fix(bigquery-jdbc): prepare Arrow result set for job-stream data

java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowBatchWrapper.java

@@ -17,12 +17,16 @@
 package com.google.cloud.bigquery.jdbc;
 
 import com.google.cloud.bigquery.storage.v1.ArrowRecordBatch;
+import com.google.cloud.bigquery.storage.v1.ArrowSchema;
 import org.apache.arrow.vector.util.JsonStringArrayList;
 
 /** This class acts as a facade layer and wraps Arrow's VectorSchemaRoot & JsonStringArrayList */
 class BigQueryArrowBatchWrapper {
   // Reference to the current arrowBatch
   private final ArrowRecordBatch currentArrowBatch;
+  // Arrow schema carried with this batch, set only when the source sent one (e.g. the first
+  // ReadRowsResponse of a stream)
+  private final ArrowSchema arrowSchema;
   // Reference to the nested Records, set as null otherwise (Arrays)
   private final JsonStringArrayList nestedRecords;
 
@@ -33,34 +37,48 @@ class BigQueryArrowBatchWrapper {
 
   private BigQueryArrowBatchWrapper(
       ArrowRecordBatch currentArrowBatch,
+      ArrowSchema arrowSchema,
       JsonStringArrayList nestedRecords,
       boolean isLast,
       Throwable exception) {
     this.currentArrowBatch = currentArrowBatch;
+    this.arrowSchema = arrowSchema;
     this.nestedRecords = nestedRecords;
     this.isLast = isLast;
     this.exception = exception;
   }
 
   static BigQueryArrowBatchWrapper of(ArrowRecordBatch currentArrowBatch, boolean... isLast) {
     boolean isLastFlag = isLast != null && isLast.length == 1 && isLast[0];
-    return new BigQueryArrowBatchWrapper(currentArrowBatch, null, isLastFlag, null);
+    return new BigQueryArrowBatchWrapper(currentArrowBatch, null, null, isLastFlag, null);
+  }
+
+  /**
+   * Wraps a record batch together with the schema sent by the source. Either may be null (a
+   * schema-only message carries no batch; later messages carry no schema).
+   */
+  static BigQueryArrowBatchWrapper of(ArrowRecordBatch currentArrowBatch, ArrowSchema arrowSchema) {
+    return new BigQueryArrowBatchWrapper(currentArrowBatch, arrowSchema, null, false, null);
   }
 
   static BigQueryArrowBatchWrapper getNestedFieldValueListWrapper(
       JsonStringArrayList nestedRecords, boolean... isLast) {
     boolean isLastFlag = isLast != null && isLast.length == 1 && isLast[0];
-    return new BigQueryArrowBatchWrapper(null, nestedRecords, isLastFlag, null);
+    return new BigQueryArrowBatchWrapper(null, null, nestedRecords, isLastFlag, null);
   }
 
   static BigQueryArrowBatchWrapper ofError(Throwable exception) {
-    return new BigQueryArrowBatchWrapper(null, null, true, exception);
+    return new BigQueryArrowBatchWrapper(null, null, null, true, exception);
   }
 
   ArrowRecordBatch getCurrentArrowBatch() {
     return this.currentArrowBatch;
   }
 
+  ArrowSchema getArrowSchema() {
+    return this.arrowSchema;
+  }
+
   JsonStringArrayList getNestedRecords() {
     return this.nestedRecords;
   }

java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java

@@ -54,6 +54,8 @@
 
 /** {@link ResultSet} Implementation for Arrow datasource (Using Storage Read APIs) */
 class BigQueryArrowResultSet extends BigQueryBaseResultSet {
+  // Number of rows to return. Iteration stops once this many rows have been returned (this is
+  // also how Statement.setMaxRows is applied).
   private final long totalRows;
   // count of rows read by the current instance of ResultSet
   private long rowCount = 0;
@@ -112,12 +114,8 @@ private BigQueryArrowResultSet(
     this.nestedRowIndex = fromIndex - 1;
     this.ownedTask = ownedTask;
     this.enableTimestampPicos = enableTimestampPicos;
-    if (!isNested && arrowSchema != null) {
-      try {
-        this.arrowDeserializer = new ArrowDeserializer(arrowSchema);
-      } catch (IOException ex) {
-        throw new BigQueryJdbcException("IOException during ArrowDeserializer creation", ex);
-      }
+    if (!isNested) {
+      initDeserializerIfNeeded(arrowSchema);
     }
   }
 
@@ -269,48 +267,29 @@ public boolean next() throws SQLException {
       return false;
     } else {
       /* Non nested */
-      if (this.hasReachedEnd || this.isLast()) {
+      if (this.afterLast || this.hasReachedEnd || reachedRowLimit()) {
+        if (!this.afterLast && this.vectorSchemaRoot != null) {
+          // Free the last batch's Arrow buffers as soon as iteration ends, not at close().
+          this.vectorSchemaRoot.clear();
+        }
         this.afterLast = true;
         return false;
       }
-      try {
-        if (this.currentBatchRowIndex == -1
-            || this.currentBatchRowIndex == (this.vectorSchemaRoot.getRowCount() - 1)) {
-          /* Start of iteration or we have exhausted the current batch */
-          // Advance the cursor. Potentially blocking operation.
-          try (Scope scope = makeOriginalContextCurrent()) {
-            BigQueryArrowBatchWrapper batchWrapper = this.buffer.take();
-            checkClosed();
-            if (batchWrapper.getException() != null) {
-              throw backgroundFetchError(batchWrapper.getException());
-            }
-            if (batchWrapper.isLast()) {
-              /* Marks the end of the records */
-              if (this.vectorSchemaRoot != null) {
-                // IMP: To avoid memory leak: clear vectorSchemaRoot as it still holds
-                // the last batch
-                this.vectorSchemaRoot.clear();
-              }
-              this.hasReachedEnd = true;
-              this.rowCount++;
-              return false;
-            }
-            // Valid batch, process it
-            ArrowRecordBatch arrowBatch = batchWrapper.getCurrentArrowBatch();
-            // Populates vectorSchemaRoot
-            this.arrowDeserializer.deserializeArrowBatch(arrowBatch);
-            // Pointing to the first row in this fresh batch
-            this.currentBatchRowIndex = 0;
-            this.rowCount++;
-            return true;
-          }
-        }
-        // There are rows left in the current batch.
-        else if (this.currentBatchRowIndex < this.vectorSchemaRoot.getRowCount()) {
-          this.currentBatchRowIndex++;
-          this.rowCount++;
-          return true;
+      if (hasMoreRowsInCurrentBatch()) {
+        this.currentBatchRowIndex++;
+        this.rowCount++;
+        return true;
+      }
+      // Start of iteration or the current batch is exhausted. Potentially blocking.
+      try (Scope scope = makeOriginalContextCurrent()) {
+        if (!loadNextBatch()) {
+          this.afterLast = true;
+          return false;
         }
+        // Pointing to the first row in this fresh batch
+        this.currentBatchRowIndex = 0;
+        this.rowCount++;
+        return true;
       } catch (SQLException ex) {
         throw new BigQueryJdbcException(
             "Error occurred while advancing the cursor. This could happen when connection is closed while the next method is being called.",
@@ -322,7 +301,71 @@ else if (this.currentBatchRowIndex < this.vectorSchemaRoot.getRowCount()) {
             ex);
       }
     }
-    return false;
+  }
+
+  /** True once totalRows rows (or the row limit) have been returned. */
+  private boolean reachedRowLimit() {
+    return this.rowCount >= this.totalRows;
+  }
+
+  private boolean hasMoreRowsInCurrentBatch() {
+    return this.currentBatchRowIndex >= 0
+        && this.vectorSchemaRoot != null
+        && this.currentBatchRowIndex < this.vectorSchemaRoot.getRowCount() - 1;
+  }
+
+  private void initDeserializerIfNeeded(ArrowSchema arrowSchema) throws SQLException {
+    if (arrowSchema == null || this.arrowDeserializer != null) {
+      return;
+    }
+    try {
+      this.arrowDeserializer = new ArrowDeserializer(arrowSchema);
+    } catch (IOException ex) {
+      throw new BigQueryJdbcException("IOException during ArrowDeserializer creation", ex);
+    }
+  }
+
+  /**
+   * Loads the next batch that has rows into the vectors. The reader queues every message as-is, so
+   * this is the single place that decides what to use: messages without rows (schema, statistics or
+   * throttling updates) are skipped, and the stream must deliver all {@code totalRows} rows.
+   * Returns false at the end of the stream.
+   */
+  private boolean loadNextBatch() throws SQLException, InterruptedException {
+    while (true) {
+      BigQueryArrowBatchWrapper batchWrapper = this.buffer.take();
+      checkClosed();
+      if (batchWrapper.getException() != null) {
+        throw backgroundFetchError(batchWrapper.getException());
+      }
+      if (batchWrapper.isLast()) {
+        /* Marks the end of the records */
+        if (this.vectorSchemaRoot != null) {
+          // IMP: To avoid memory leak: clear vectorSchemaRoot as it still holds the last batch
+          this.vectorSchemaRoot.clear();
+        }
+        this.hasReachedEnd = true;
+        if (this.rowCount < this.totalRows) {
+          throw new BigQueryJdbcException(
+              String.format(
+                  "The Arrow stream ended after %d of %d rows.", this.rowCount, this.totalRows));
+        }
+        return false;
+      }
+      initDeserializerIfNeeded(batchWrapper.getArrowSchema());
+      ArrowRecordBatch arrowBatch = batchWrapper.getCurrentArrowBatch();
+      if (arrowBatch == null || arrowBatch.getSerializedRecordBatch().isEmpty()) {
+        continue; // no rows in this message
+      }
+      if (this.arrowDeserializer == null) {
+        throw new BigQueryJdbcException("Received an Arrow record batch before its schema.");
+      }
+      // Populates vectorSchemaRoot
+      this.arrowDeserializer.deserializeArrowBatch(arrowBatch);
+      if (this.vectorSchemaRoot.getRowCount() > 0) {
+        return true;
+      }
+    }
   }
 
   @Override
@@ -500,7 +543,8 @@ public boolean isBeforeFirst() throws SQLException {
     if (this.isNested) {
       return this.nestedRowIndex < this.fromIndex;
     } else {
-      return this.rowCount == 0;
+      // Per JDBC, false when the result set contains no rows.
+      return !this.afterLast && this.rowCount == 0 && this.totalRows > 0;
     }
   }
 
@@ -529,7 +573,7 @@ public boolean isLast() throws SQLException {
     if (this.isNested) {
       return !this.afterLast && this.nestedRowIndex == this.toIndexExclusive - 1;
     } else {
-      return !this.afterLast && this.rowCount == this.totalRows;
+      return !this.afterLast && this.rowCount > 0 && this.rowCount == this.totalRows;
     }
   }
 }

java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java

@@ -52,7 +52,6 @@
 import com.google.cloud.bigquery.jdbc.telemetry.v1.QueryApiType;
 import com.google.cloud.bigquery.jdbc.telemetry.v1.StatementExecution;
 import com.google.cloud.bigquery.jdbc.telemetry.v1.TelemetryManager;
-import com.google.cloud.bigquery.storage.v1.ArrowRecordBatch;
 import com.google.cloud.bigquery.storage.v1.ArrowSchema;
 import com.google.cloud.bigquery.storage.v1.ArrowSerializationOptions;
 import com.google.cloud.bigquery.storage.v1.BigQueryReadClient;
@@ -1035,9 +1034,13 @@ Future<?> populateArrowBufferedQueue(
                                 "Arrow stream reading was interrupted or shut down before completion.");
                           }
 
-                          ArrowRecordBatch currentBatch = response.getArrowRecordBatch();
+                          // Queue every response as-is; the result set decides what to use.
                           arrowBatchWrapperBlockingQueue.put(
-                              BigQueryArrowBatchWrapper.of(currentBatch));
+                              BigQueryArrowBatchWrapper.of(
+                                  response.hasArrowRecordBatch()
+                                      ? response.getArrowRecordBatch()
+                                      : null,
+                                  response.hasArrowSchema() ? response.getArrowSchema() : null));
                           rowsRead += response.getRowCount();
                         }
                         if (isWorkerShuttingDown()) {

java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSetTest.java

@@ -302,6 +302,11 @@ public void testRowCount() throws SQLException, IOException {
             QUERY_SCHEMA, arrowSchema, 1, statement, bufferWithTwoRows, workerTask, null);
 
     assertThat(resultSetRowCount(bigQueryArrowResultSet2)).isEqualTo(1);
+    // ResultSet with 3 rows buffer and 2 total rows (e.g. Statement.setMaxRows(2)).
+    assertThat(
+            resultSetRowCount(
+                idResultSet(idArrowSchema(), 2, idWrapper(1, 2, 3), arrowBatchWrapperLast)))
+        .isEqualTo(2);
   }
 
   @Test
@@ -833,4 +838,106 @@ public void testGetterWhenNotOnRowThrows() throws SQLException {
     assertThat(bigQueryArrowResultSetNested.isLast()).isFalse();
     assertThrows(SQLException.class, () -> bigQueryArrowResultSetNested.getObject(2));
   }
+
+  private static final Schema ID_SCHEMA = Schema.of(Field.of("id", StandardSQLTypeName.INT64));
+
+  private static VectorSchemaRoot idRoot(int... ids) {
+    IntVector idVector = new IntVector("id", new RootAllocator());
+    idVector.allocateNew(Math.max(ids.length, 1));
+    for (int i = 0; i < ids.length; i++) {
+      idVector.set(i, ids[i]);
+    }
+    idVector.setValueCount(ids.length);
+    VectorSchemaRoot root = new VectorSchemaRoot(ImmutableList.<FieldVector>of(idVector));
+    root.setRowCount(ids.length);
+    return root;
+  }
+
+  private static ArrowSchema idArrowSchema() throws IOException {
+    return ArrowSchema.newBuilder()
+        .setSerializedSchema(serializeSchema(idRoot().getSchema()))
+        .build();
+  }
+
+  private static BigQueryArrowBatchWrapper idWrapper(int... ids) throws IOException {
+    ArrowRecordBatch batch =
+        ArrowRecordBatch.newBuilder()
+            .setSerializedRecordBatch(serializeVectorSchemaRoot(idRoot(ids)))
+            .build();
+    return BigQueryArrowBatchWrapper.of(batch);
+  }
+
+  private BigQueryArrowResultSet idResultSet(
+      ArrowSchema arrowSchema, long totalRows, BigQueryArrowBatchWrapper... items)
+      throws SQLException {
+    BlockingQueue<BigQueryArrowBatchWrapper> queue =
+        new LinkedBlockingDeque<>(Arrays.asList(items));
+    return BigQueryArrowResultSet.of(
+        ID_SCHEMA, arrowSchema, totalRows, statement, queue, mock(Future.class), null);
+  }
+
+  @Test
+  public void testMessagesWithoutRowsAreSkipped() throws SQLException, IOException {
+    BigQueryArrowBatchWrapper noBatch = BigQueryArrowBatchWrapper.of(null);
+    BigQueryArrowBatchWrapper emptyBatch =
+        BigQueryArrowBatchWrapper.of(ArrowRecordBatch.getDefaultInstance());
+    BigQueryArrowResultSet rs =
+        idResultSet(
+            idArrowSchema(),
+            3,
+            idWrapper(1, 2),
+            noBatch,
+            idWrapper(),
+            emptyBatch,
+            idWrapper(3),
+            arrowBatchWrapperLast);
+
+    assertThat(rs.next()).isTrue();
+    assertThat(rs.next()).isTrue();
+    assertThat(rs.next()).isTrue();
+    assertThat(rs.getLong(1)).isEqualTo(3L);
+    assertThat(rs.next()).isFalse();
+  }
+
+  @Test
+  public void testStreamEndingBeforeTotalRowsThrows() throws SQLException, IOException {
+    BigQueryArrowResultSet rs =
+        idResultSet(idArrowSchema(), 3, idWrapper(1, 2), arrowBatchWrapperLast);
+
+    assertThat(rs.next()).isTrue();
+    assertThat(rs.next()).isTrue();
+    SQLException e = assertThrows(SQLException.class, rs::next);
+    assertThat(e.getCause().getMessage()).contains("ended after 2 of 3 rows");
+  }
+
+  @Test
+  public void testEmptyResultIsNotBeforeFirst() throws SQLException, IOException {
+    BigQueryArrowResultSet rs = idResultSet(idArrowSchema(), 0, arrowBatchWrapperLast);
+
+    assertThat(rs.isBeforeFirst()).isFalse();
+    assertThat(rs.isLast()).isFalse();
+    assertThat(rs.next()).isFalse();
+  }
+
+  @Test
+  public void testSchemaFromStreamMessage() throws SQLException, IOException {
+    BigQueryArrowResultSet rs =
+        idResultSet(
+            null,
+            1,
+            BigQueryArrowBatchWrapper.of(null, idArrowSchema()),
+            idWrapper(7),
+            arrowBatchWrapperLast);
+
+    assertThat(rs.next()).isTrue();
+    assertThat(rs.getLong(1)).isEqualTo(7L);
+    assertThat(rs.next()).isFalse();
+  }
+
+  @Test
+  public void testBatchBeforeSchemaThrows() throws SQLException, IOException {
+    BigQueryArrowResultSet rs = idResultSet(null, 1, idWrapper(1), arrowBatchWrapperLast);
+
+    assertThrows(SQLException.class, rs::next);
+  }
 }