Repository navigation
fix(bigquery-jdbc): prepare Arrow result set for job-stream data - #14941
keshavdandeva wants to merge 3 commits 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
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
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.
| 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?
There was a problem hiding this comment.
For today's read-session path there's no observable change: every message we've seen has rows and a batch, and those are queued exactly as before. It only changes behaviour for messages we haven't seen yet. Messages without rows (allowed by ReadRowsResponse) are no longer queued, so the reader can't turn them into a phantom row. A message that reports rows but has no batch now fails instead of silently dropping those rows. It also passes on a schema sent within the stream, which the job stream needs.
| /** 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?
There was a problem hiding this comment.
I added it originally because in early testing totalRows was missing from some jobs.query responses, and I assumed we might have to stream without it. Double-checking all the cases, it's only missing for DML/DDL (nothing to stream) and when the job hasn't finished yet, and in that case we wait first and the wait call always returns it. So there's no real scenario where we stream without a known row count. I've removed it, and isLast()/isBeforeFirst() are plain comparisons again.
keshavdandeva
changed the title
b/582676539
Prepares the Arrow reader for the job stream (b/581544197).
isBeforeFirst()/isLast()return false for an empty result.