[CALCITE-2322] Add fetch row count support

core/src/main/java/org/apache/calcite/avatica/AvaticaStatement.java

@@ -66,7 +66,7 @@ public abstract class AvaticaStatement
   final int resultSetType;
   final int resultSetConcurrency;
   final int resultSetHoldability;
-  private int fetchSize = DEFAULT_FETCH_SIZE;
+  private int fetchSizeRows;
   private int fetchDirection;
   protected long maxRowCount = 0;
 
@@ -108,6 +108,7 @@ protected AvaticaStatement(AvaticaConnection connection,
     this.resultSetType = resultSetType;
     this.resultSetConcurrency = resultSetConcurrency;
     this.resultSetHoldability = resultSetHoldability;
+    this.fetchSizeRows = connection.config().fetchSizeRows(); // Default to connection config value
     this.signature = signature;
     this.closed = false;
     if (h == null) {
@@ -407,12 +408,12 @@ public int getFetchDirection() throws SQLException {
 
   public void setFetchSize(int rows) throws SQLException {
     checkOpen();
-    this.fetchSize = rows;
+    this.fetchSizeRows = rows;
   }
 
   public int getFetchSize() throws SQLException {
     checkOpen();
-    return fetchSize;
+    return fetchSizeRows;
   }
 
   public int getResultSetConcurrency() throws SQLException {

core/src/main/java/org/apache/calcite/avatica/BuiltInConnectionProperty.java

@@ -86,7 +86,10 @@ public enum BuiltInConnectionProperty implements ConnectionProperty {
   KEY_PASSWORD("key_password", Type.STRING, "", false),
 
   HOSTNAME_VERIFICATION("hostname_verification", Type.ENUM, HostnameVerification.STRICT,
-      HostnameVerification.class, false);
+      HostnameVerification.class, false),
+
+  /** The number of rows to fetch per call, default is 100 rows. */
+  FETCH_SIZE_ROWS("fetch_size_rows", Type.NUMBER, AvaticaStatement.DEFAULT_FETCH_SIZE, false);
 
   private final String camelName;
   private final Type type;

core/src/main/java/org/apache/calcite/avatica/ConnectionConfig.java

@@ -62,6 +62,8 @@ public interface ConnectionConfig {
   String keyPassword();
   /** @see BuiltInConnectionProperty#HOSTNAME_VERIFICATION */
   HostnameVerification hostnameVerification();
+  /** @see BuiltInConnectionProperty#FETCH_SIZE_ROWS */
+  int fetchSizeRows();
 }
 
 // End ConnectionConfig.java

core/src/main/java/org/apache/calcite/avatica/ConnectionConfigImpl.java

@@ -128,6 +128,10 @@ public HostnameVerification hostnameVerification() {
         .getEnum(HostnameVerification.class);
   }
 
+  public int fetchSizeRows() {
+    return BuiltInConnectionProperty.FETCH_SIZE_ROWS.wrap(properties).getInt();
+  }
+
   /** Converts a {@link Properties} object containing (name, value)
    * pairs into a map whose keys are
    * {@link org.apache.calcite.avatica.InternalProperty} objects.

core/src/main/java/org/apache/calcite/avatica/MetaImpl.java

@@ -1543,13 +1543,21 @@ public Iterator<Object> iterator() {
   private class FetchIterator implements Iterator<Object> {
     private final AvaticaStatement stmt;
     private final QueryState state;
+    private final int fetchRowCount;
     private Frame frame;
     private Iterator<Object> rows;
     private long currentOffset = 0;
 
     private FetchIterator(AvaticaStatement stmt, QueryState state, Frame firstFrame) {
       this.stmt = stmt;
       this.state = state;
+      int fetchRowCount;
+      try {
+        fetchRowCount = stmt.getFetchSize();
+      } catch (SQLException e) {
+        fetchRowCount = AvaticaStatement.DEFAULT_FETCH_SIZE;
+      }
+      this.fetchRowCount = fetchRowCount;
       if (firstFrame == null) {
         frame = Frame.MORE;
         rows = EmptyIterator.INSTANCE;
@@ -1589,7 +1597,7 @@ private void moveNext() {
         }
         try {
           // currentOffset updated after element is read from `rows` iterator
-          frame = fetch(stmt.handle, currentOffset, AvaticaStatement.DEFAULT_FETCH_SIZE);
+          frame = fetch(stmt.handle, currentOffset, fetchRowCount);
         } catch (NoSuchStatementException e) {
           resetStatement();
           // re-fetch the batch where we left off

site/_docs/client_reference.md

@@ -172,3 +172,13 @@ on-hover images for the permalink, but oh well.
 : _Default_: `null`.
 
 : _Required_: Only if `truststore` was provided.
+
+<strong><a name="fetch_size_rows" href="#fetch_size_rows">fetch_size_rows</a></strong>
+
+: _Description_: The number of rows to fetch. If
+    <a href="https://docs.oracle.com/javase/8/docs/api/java/sql/Statement.html#setFetchSize-int-">
+    Statement:setFetchSize</a> is set, that value overrides fetch_size_rows.
+
+: _Default_: `100`.
+
+: _Required_: No.