From de892184274897f1fdf53b3ff1a66a8ebd142f34 Mon Sep 17 00:00:00 2001 From: Keshav Dandeva Date: Fri, 13 Nov 2026 14:58:36 +0000 Subject: [PATCH 1/2] fix(bigquery-jdbc): prepare Arrow result set for job-stream data --- .../jdbc/BigQueryArrowBatchWrapper.java | 24 +++- .../bigquery/jdbc/BigQueryArrowResultSet.java | 133 +++++++++++------- .../bigquery/jdbc/BigQueryStatement.java | 34 ++++- .../jdbc/BigQueryArrowResultSetTest.java | 86 +++++++++++ .../bigquery/jdbc/BigQueryStatementTest.java | 62 +++++++- 5 files changed, 282 insertions(+), 57 deletions(-) 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..ef2ad06297b0 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,66 @@ 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 into the vectors, skipping schema-only messages. 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; + return false; + } + initDeserializerIfNeeded(batchWrapper.getArrowSchema()); + ArrowRecordBatch arrowBatch = batchWrapper.getCurrentArrowBatch(); + if (arrowBatch == null) { + continue; // schema-only 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) { + // Only batches reported as non-empty are queued, so this means corrupt data. + throw new BigQueryJdbcException("Received an Arrow record batch with no rows."); + } + return true; + } } @Override @@ -500,7 +538,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 +568,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..ede031ce9ec2 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 @@ -46,13 +46,13 @@ import com.google.cloud.bigquery.TableId; import com.google.cloud.bigquery.TableResult; import com.google.cloud.bigquery.exception.BigQueryJdbcException; +import com.google.cloud.bigquery.exception.BigQueryJdbcRuntimeException; import com.google.cloud.bigquery.exception.BigQueryJdbcSqlFeatureNotSupportedException; import com.google.cloud.bigquery.exception.BigQueryJdbcSqlSyntaxErrorException; import com.google.cloud.bigquery.jdbc.telemetry.v1.DriverFeature; 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 +1035,7 @@ Future populateArrowBufferedQueue( "Arrow stream reading was interrupted or shut down before completion."); } - ArrowRecordBatch currentBatch = response.getArrowRecordBatch(); - arrowBatchWrapperBlockingQueue.put( - BigQueryArrowBatchWrapper.of(currentBatch)); + enqueueReadRowsResponse(response, arrowBatchWrapperBlockingQueue); rowsRead += response.getRowCount(); } if (isWorkerShuttingDown()) { @@ -1114,6 +1112,34 @@ Future populateArrowBufferedQueue( return executor.submit(arrowStreamProcessor); } + /** + * Queues the rows of one {@link ReadRowsResponse}. Responses without rows (for example schema, + * statistics or throttling updates) queue no batch, but a schema they carry is still forwarded. A + * response that reports rows but carries no record batch is an error: rows are never silently + * dropped. + */ + @VisibleForTesting + static void enqueueReadRowsResponse( + ReadRowsResponse response, BlockingQueue queue) + throws InterruptedException { + long responseRows = response.getRowCount(); + ArrowSchema schema = response.hasArrowSchema() ? response.getArrowSchema() : null; + boolean hasBatch = + response.hasArrowRecordBatch() + && !response.getArrowRecordBatch().getSerializedRecordBatch().isEmpty(); + if (responseRows > 0 && !hasBatch) { + throw new BigQueryJdbcRuntimeException( + String.format( + "Read API response reported %d rows but carried no Arrow record batch.", + responseRows)); + } + if (responseRows > 0) { + queue.put(BigQueryArrowBatchWrapper.of(response.getArrowRecordBatch(), schema)); + } else if (schema != null) { + queue.put(BigQueryArrowBatchWrapper.of(null, schema)); + } + } + /** Executes SQL query using either fast query path or read API */ void processQueryResponse(String query, TableResult results) throws SQLException { processQueryResponse(query, results, null); 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..f04f514742c1 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,85 @@ 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 testZeroRowBatchIsAnError() throws SQLException, IOException { + BigQueryArrowResultSet rs = + idResultSet( + idArrowSchema(), 3, idWrapper(1, 2), idWrapper(), idWrapper(3), arrowBatchWrapperLast); + + assertThat(rs.next()).isTrue(); + assertThat(rs.next()).isTrue(); + assertThat(rs.getLong(1)).isEqualTo(2L); + // Only batches reported as non-empty are queued, so a zero-row batch means corrupt data. + assertThrows(SQLException.class, rs::next); + } + + @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); + } } diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java index e8787e72cdd1..e62a2666d10e 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java @@ -66,6 +66,7 @@ import com.google.cloud.bigquery.TableId; import com.google.cloud.bigquery.TableResult; import com.google.cloud.bigquery.exception.BigQueryJdbcException; +import com.google.cloud.bigquery.exception.BigQueryJdbcRuntimeException; import com.google.cloud.bigquery.jdbc.BigQueryStatement.JobIdWrapper; import com.google.cloud.bigquery.spi.BigQueryRpcFactory; import com.google.cloud.bigquery.storage.v1.ArrowRecordBatch; @@ -79,6 +80,7 @@ import com.google.cloud.bigquery.storage.v1.ReadStream; import com.google.common.collect.ImmutableList; import com.google.common.collect.Maps; +import com.google.protobuf.ByteString; import io.opentelemetry.api.baggage.Baggage; import io.opentelemetry.api.common.AttributeKey; import io.opentelemetry.api.trace.Span; @@ -1394,10 +1396,9 @@ public void testPopulateArrowBufferedQueueStopsAndCancelsStreamWhenResultSetClos ReadStream.newBuilder().setName("projects/p/locations/l/sessions/s/streams/0")) .build(); + // Each response must carry rows; responses without rows are dropped by the producer. ReadRowsResponse batchResponse = - ReadRowsResponse.newBuilder() - .setArrowRecordBatch(ArrowRecordBatch.getDefaultInstance()) - .build(); + ReadRowsResponse.newBuilder().setArrowRecordBatch(NON_EMPTY_BATCH).setRowCount(1).build(); CountDownLatch secondBatchReturnedLatch = new CountDownLatch(1); AtomicBoolean streamCancelled = new AtomicBoolean(false); @@ -1894,4 +1895,59 @@ private Job createScriptChildJobMock( doReturn(mock(ScriptStatistics.class)).when(stats).getScriptStatistics(); return childJob; } + + private static final ArrowRecordBatch NON_EMPTY_BATCH = + ArrowRecordBatch.newBuilder() + .setSerializedRecordBatch(ByteString.copyFromUtf8("rows")) + .build(); + + @Test + public void testEnqueueReadRowsResponseQueuesBatchWithSchemaAndRowCount() + throws InterruptedException { + BlockingQueue queue = new LinkedBlockingDeque<>(); + ReadRowsResponse response = + ReadRowsResponse.newBuilder() + .setArrowSchema(arrowSchema) + .setArrowRecordBatch(NON_EMPTY_BATCH) + .setRowCount(5) + .build(); + + BigQueryStatement.enqueueReadRowsResponse(response, queue); + + BigQueryArrowBatchWrapper wrapper = queue.poll(); + assertEquals(NON_EMPTY_BATCH, wrapper.getCurrentArrowBatch()); + assertEquals(arrowSchema, wrapper.getArrowSchema()); + assertTrue(queue.isEmpty()); + } + + @Test + public void testEnqueueReadRowsResponseWithoutRowsForwardsOnlySchema() + throws InterruptedException { + BlockingQueue queue = new LinkedBlockingDeque<>(); + + // Statistics-only response and zero-row batch: nothing to queue. + BigQueryStatement.enqueueReadRowsResponse(ReadRowsResponse.getDefaultInstance(), queue); + BigQueryStatement.enqueueReadRowsResponse( + ReadRowsResponse.newBuilder().setArrowRecordBatch(NON_EMPTY_BATCH).build(), queue); + assertTrue(queue.isEmpty()); + + // Schema-only response: the schema is forwarded without a batch. + BigQueryStatement.enqueueReadRowsResponse( + ReadRowsResponse.newBuilder().setArrowSchema(arrowSchema).build(), queue); + BigQueryArrowBatchWrapper wrapper = queue.poll(); + assertNull(wrapper.getCurrentArrowBatch()); + assertEquals(arrowSchema, wrapper.getArrowSchema()); + assertTrue(queue.isEmpty()); + } + + @Test + public void testEnqueueReadRowsResponseRejectsRowsWithoutBatch() { + BlockingQueue queue = new LinkedBlockingDeque<>(); + ReadRowsResponse response = ReadRowsResponse.newBuilder().setRowCount(3).build(); + + assertThrows( + BigQueryJdbcRuntimeException.class, + () -> BigQueryStatement.enqueueReadRowsResponse(response, queue)); + assertTrue(queue.isEmpty()); + } } From bf5767d3d5b648a2cad003de6bb25b863b9ae02b Mon Sep 17 00:00:00 2001 From: Keshav Dandeva Date: Fri, 13 Nov 2026 16:10:29 +0000 Subject: [PATCH 2/2] move Arrow batch acceptance rules into the result set --- .../bigquery/jdbc/BigQueryArrowResultSet.java | 21 ++++--- .../bigquery/jdbc/BigQueryStatement.java | 37 +++-------- .../jdbc/BigQueryArrowResultSetTest.java | 31 ++++++++-- .../bigquery/jdbc/BigQueryStatementTest.java | 62 +------------------ 4 files changed, 49 insertions(+), 102 deletions(-) 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 ef2ad06297b0..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 @@ -326,8 +326,10 @@ private void initDeserializerIfNeeded(ArrowSchema arrowSchema) throws SQLExcepti } /** - * Loads the next batch into the vectors, skipping schema-only messages. Returns false at the end - * of the stream. + * 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) { @@ -343,23 +345,26 @@ private boolean loadNextBatch() throws SQLException, InterruptedException { 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) { - continue; // schema-only message + 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) { - // Only batches reported as non-empty are queued, so this means corrupt data. - throw new BigQueryJdbcException("Received an Arrow record batch with no rows."); + if (this.vectorSchemaRoot.getRowCount() > 0) { + return true; } - return true; } } 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 ede031ce9ec2..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 @@ -46,7 +46,6 @@ import com.google.cloud.bigquery.TableId; import com.google.cloud.bigquery.TableResult; import com.google.cloud.bigquery.exception.BigQueryJdbcException; -import com.google.cloud.bigquery.exception.BigQueryJdbcRuntimeException; import com.google.cloud.bigquery.exception.BigQueryJdbcSqlFeatureNotSupportedException; import com.google.cloud.bigquery.exception.BigQueryJdbcSqlSyntaxErrorException; import com.google.cloud.bigquery.jdbc.telemetry.v1.DriverFeature; @@ -1035,7 +1034,13 @@ Future populateArrowBufferedQueue( "Arrow stream reading was interrupted or shut down before completion."); } - enqueueReadRowsResponse(response, arrowBatchWrapperBlockingQueue); + // Queue every response as-is; the result set decides what to use. + arrowBatchWrapperBlockingQueue.put( + BigQueryArrowBatchWrapper.of( + response.hasArrowRecordBatch() + ? response.getArrowRecordBatch() + : null, + response.hasArrowSchema() ? response.getArrowSchema() : null)); rowsRead += response.getRowCount(); } if (isWorkerShuttingDown()) { @@ -1112,34 +1117,6 @@ Future populateArrowBufferedQueue( return executor.submit(arrowStreamProcessor); } - /** - * Queues the rows of one {@link ReadRowsResponse}. Responses without rows (for example schema, - * statistics or throttling updates) queue no batch, but a schema they carry is still forwarded. A - * response that reports rows but carries no record batch is an error: rows are never silently - * dropped. - */ - @VisibleForTesting - static void enqueueReadRowsResponse( - ReadRowsResponse response, BlockingQueue queue) - throws InterruptedException { - long responseRows = response.getRowCount(); - ArrowSchema schema = response.hasArrowSchema() ? response.getArrowSchema() : null; - boolean hasBatch = - response.hasArrowRecordBatch() - && !response.getArrowRecordBatch().getSerializedRecordBatch().isEmpty(); - if (responseRows > 0 && !hasBatch) { - throw new BigQueryJdbcRuntimeException( - String.format( - "Read API response reported %d rows but carried no Arrow record batch.", - responseRows)); - } - if (responseRows > 0) { - queue.put(BigQueryArrowBatchWrapper.of(response.getArrowRecordBatch(), schema)); - } else if (schema != null) { - queue.put(BigQueryArrowBatchWrapper.of(null, schema)); - } - } - /** Executes SQL query using either fast query path or read API */ void processQueryResponse(String query, TableResult results) throws SQLException { processQueryResponse(query, results, null); 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 f04f514742c1..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 @@ -877,16 +877,37 @@ private BigQueryArrowResultSet idResultSet( } @Test - public void testZeroRowBatchIsAnError() throws SQLException, IOException { + 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), idWrapper(), idWrapper(3), arrowBatchWrapperLast); + idArrowSchema(), + 3, + idWrapper(1, 2), + noBatch, + idWrapper(), + emptyBatch, + idWrapper(3), + arrowBatchWrapperLast); assertThat(rs.next()).isTrue(); assertThat(rs.next()).isTrue(); - assertThat(rs.getLong(1)).isEqualTo(2L); - // Only batches reported as non-empty are queued, so a zero-row batch means corrupt data. - assertThrows(SQLException.class, rs::next); + 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 diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java index e62a2666d10e..e8787e72cdd1 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryStatementTest.java @@ -66,7 +66,6 @@ import com.google.cloud.bigquery.TableId; import com.google.cloud.bigquery.TableResult; import com.google.cloud.bigquery.exception.BigQueryJdbcException; -import com.google.cloud.bigquery.exception.BigQueryJdbcRuntimeException; import com.google.cloud.bigquery.jdbc.BigQueryStatement.JobIdWrapper; import com.google.cloud.bigquery.spi.BigQueryRpcFactory; import com.google.cloud.bigquery.storage.v1.ArrowRecordBatch; @@ -80,7 +79,6 @@ import com.google.cloud.bigquery.storage.v1.ReadStream; import com.google.common.collect.ImmutableList; import com.google.common.collect.Maps; -import com.google.protobuf.ByteString; import io.opentelemetry.api.baggage.Baggage; import io.opentelemetry.api.common.AttributeKey; import io.opentelemetry.api.trace.Span; @@ -1396,9 +1394,10 @@ public void testPopulateArrowBufferedQueueStopsAndCancelsStreamWhenResultSetClos ReadStream.newBuilder().setName("projects/p/locations/l/sessions/s/streams/0")) .build(); - // Each response must carry rows; responses without rows are dropped by the producer. ReadRowsResponse batchResponse = - ReadRowsResponse.newBuilder().setArrowRecordBatch(NON_EMPTY_BATCH).setRowCount(1).build(); + ReadRowsResponse.newBuilder() + .setArrowRecordBatch(ArrowRecordBatch.getDefaultInstance()) + .build(); CountDownLatch secondBatchReturnedLatch = new CountDownLatch(1); AtomicBoolean streamCancelled = new AtomicBoolean(false); @@ -1895,59 +1894,4 @@ private Job createScriptChildJobMock( doReturn(mock(ScriptStatistics.class)).when(stats).getScriptStatistics(); return childJob; } - - private static final ArrowRecordBatch NON_EMPTY_BATCH = - ArrowRecordBatch.newBuilder() - .setSerializedRecordBatch(ByteString.copyFromUtf8("rows")) - .build(); - - @Test - public void testEnqueueReadRowsResponseQueuesBatchWithSchemaAndRowCount() - throws InterruptedException { - BlockingQueue queue = new LinkedBlockingDeque<>(); - ReadRowsResponse response = - ReadRowsResponse.newBuilder() - .setArrowSchema(arrowSchema) - .setArrowRecordBatch(NON_EMPTY_BATCH) - .setRowCount(5) - .build(); - - BigQueryStatement.enqueueReadRowsResponse(response, queue); - - BigQueryArrowBatchWrapper wrapper = queue.poll(); - assertEquals(NON_EMPTY_BATCH, wrapper.getCurrentArrowBatch()); - assertEquals(arrowSchema, wrapper.getArrowSchema()); - assertTrue(queue.isEmpty()); - } - - @Test - public void testEnqueueReadRowsResponseWithoutRowsForwardsOnlySchema() - throws InterruptedException { - BlockingQueue queue = new LinkedBlockingDeque<>(); - - // Statistics-only response and zero-row batch: nothing to queue. - BigQueryStatement.enqueueReadRowsResponse(ReadRowsResponse.getDefaultInstance(), queue); - BigQueryStatement.enqueueReadRowsResponse( - ReadRowsResponse.newBuilder().setArrowRecordBatch(NON_EMPTY_BATCH).build(), queue); - assertTrue(queue.isEmpty()); - - // Schema-only response: the schema is forwarded without a batch. - BigQueryStatement.enqueueReadRowsResponse( - ReadRowsResponse.newBuilder().setArrowSchema(arrowSchema).build(), queue); - BigQueryArrowBatchWrapper wrapper = queue.poll(); - assertNull(wrapper.getCurrentArrowBatch()); - assertEquals(arrowSchema, wrapper.getArrowSchema()); - assertTrue(queue.isEmpty()); - } - - @Test - public void testEnqueueReadRowsResponseRejectsRowsWithoutBatch() { - BlockingQueue queue = new LinkedBlockingDeque<>(); - ReadRowsResponse response = ReadRowsResponse.newBuilder().setRowCount(3).build(); - - assertThrows( - BigQueryJdbcRuntimeException.class, - () -> BigQueryStatement.enqueueReadRowsResponse(response, queue)); - assertTrue(queue.isEmpty()); - } }