Repository navigation
fix(bigquery-jdbc): harden Arrow result set for empty batches, unknown row counts and lazy schema - #14941
fix(bigquery-jdbc): harden Arrow result set for empty batches, unknown row counts and lazy schema#14941keshavdandeva wants to merge 1 commit into
Conversation
…n row counts and lazy schema - Producer skips Read API responses without rows and fails fast when a response reports rows but carries no record batch. - Result set skips zero-row batches instead of returning a phantom row. - Support an unknown total row count: iterate until the end marker; isLast()/isBeforeFirst() look ahead without skipping rows. - Accept the Arrow schema from the first stream message. - Fix isLast() when Statement.setMaxRows exceeds the result size. - isBeforeFirst()/isLast() return false for an empty result (JDBC spec). - Getters raise SQLException instead of IndexOutOfBoundsException or a stale value when the cursor is not on a row. - Deserialize record batches without copying them to a heap byte array. Bug: b/582676539
There was a problem hiding this comment.
Code Review
This pull request refactors the BigQuery JDBC driver's Arrow result set processing to support streams with an unknown total row count, skip empty batches, and dynamically forward schemas. It also optimizes memory usage during Arrow batch deserialization by reading directly from the ByteString input stream instead of copying it to a heap byte array. Feedback is provided to wrap the newly introduced InputStream and ReadableByteChannel in a try-with-resources block to prevent potential static analysis warnings regarding unclosed resources.
| // set the resultset | ||
| long totalRows = (getMaxRows() > 0) ? getMaxRows() : results.getTotalRows(); | ||
| long totalRows = | ||
| (getMaxRows() > 0) |
There was a problem hiding this comment.
I'd move to helper method; Same applies to JsonResultSet
(maybe worth moving to a standalone PR & merge to main directly)
| ArrowRecordBatch currentBatch = response.getArrowRecordBatch(); | ||
| arrowBatchWrapperBlockingQueue.put( | ||
| BigQueryArrowBatchWrapper.of(currentBatch)); | ||
| enqueueReadRowsResponse(response, arrowBatchWrapperBlockingQueue); |
There was a problem hiding this comment.
What is the impact of this change?
| /** This class acts as a facade layer and wraps Arrow's VectorSchemaRoot & JsonStringArrayList */ | ||
| class BigQueryArrowBatchWrapper { | ||
| /** Row count used when the producer did not report one. */ | ||
| static final long UNKNOWN_ROW_COUNT = -1; |
There was a problem hiding this comment.
What is a scenario for unknown # of rows?
b/582676539
Hardens
BigQueryArrowResultSetand the Read API producer ahead of the Arrow job-stream work (b/581544197). Targets thejdbc/arrow-support-featurebranch.Changes
IndexOutOfBoundsExceptionand the following row was lost).isLast()/isBeforeFirst()look ahead without skipping rows.isLast()never returning true whenStatement.setMaxRowsexceeds the result size.isBeforeFirst()/isLast()return false for an empty result, per the JDBC spec.SQLExceptioninstead ofIndexOutOfBoundsExceptionor a stale value when the cursor is not on a row.ByteStringstream instead of being copied to a heapbyte[].Testing
BigQueryArrowResultSetTestandBigQueryStatementTest;testRowCountnow exercises a real row limit.