diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowBatchWrapper.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowBatchWrapper.java index 62588c58ab42..bd6041a6f7f4 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowBatchWrapper.java +++ b/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,10 +37,12 @@ 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; @@ -44,23 +50,35 @@ private BigQueryArrowBatchWrapper( 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; } diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java index 620f91641933..afac21f9c456 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSet.java +++ b/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; } } } diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java index 4905a6d5aa46..41f8d59c30b7 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java +++ b/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()) { diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSetTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSetTest.java index 72a91dc42145..33c5084a1cf1 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryArrowResultSetTest.java +++ b/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.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 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); + } }