feat(bigquery): add zero-copy queryArrow API for Arrow VectorSchemaRoot streaming - #14402
jinseopkim0 wants to merge 12 commits into
Conversation
![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 introduces support for executing queries in Apache Arrow format via a new queryArrow API, which utilizes the BigQuery Storage Read API for streaming results. Key changes include adding necessary gRPC dependencies, implementing queryArrow in BigQueryImpl with lazy initialization of a BigQueryReadClient, and updating QueryRequestInfo and tests. The review feedback highlights a critical improvement: instead of caching a single global BigQueryReadClient, the client should be cached by region/location to prevent failures when querying regional datasets, with corresponding updates to endpoint configuration and client retrieval.
|
@gemini-code-assist review |
There was a problem hiding this comment.
Code Review
This pull request introduces support for executing queries in Apache Arrow format by adding queryArrow methods to the BigQuery interface and implementing them in BigQueryImpl. It integrates the BigQuery Storage Read API, lazily initializing and caching BigQueryReadClient instances. Feedback on the implementation highlights two main issues: first, reusing an auto-created 'global' client for regional queries bypasses regional routing and data residency requirements, which can be resolved by tracking whether the global client was explicitly user-provided; second, a bug in the location resolution logic within createArrowQueryResultFromTable incorrectly evaluates to null when a jobId is present but lacks a location, ignoring the configured default location.
| if (bqReadClients == null) { | ||
| bqReadClients = Maps.newHashMap(); | ||
| } | ||
| BigQueryReadClient client = bqReadClients.get(cacheKey); | ||
| if (client == null && bqReadClients.containsKey("global")) { | ||
| client = bqReadClients.get("global"); | ||
| } | ||
| if (client == null) { |
There was a problem hiding this comment.
do we need caching for this use case? Would there be any bottleneck within a single read client if a user has high throughput for a single region?
There was a problem hiding this comment.
Thanks for the questions.
-
Creating clients is expensive, so caching is recommended. In https://grpc.io/docs/guides/performance/#java, it says "Always re-use stubs and channels when possible." Also, the
BigQueryReadClientJavadoc notes thatclose()is required to clean up allocated thread and channel resources. Creating clients per-query would introduce cold connection latency, thread churn, and channel lifecycle management issues. -
A single read client should be able to handle high throughput for a single region. BigQueryReadClient is threadsafe, and single client usage can be seen in ConnectionImpl.java.
|
@gemini-code-assist review |
There was a problem hiding this comment.
Code Review
This pull request introduces support for querying BigQuery and returning results in Apache Arrow format via the new queryArrow API. It adds necessary gRPC dependencies, implements the queryArrow methods in BigQueryImpl with both fast-path and fallback execution paths, manages regional BigQueryReadClient caching, and includes corresponding unit tests. Key feedback includes ensuring that fast-path execution is bypassed when JobOptions are provided to avoid silently ignoring them, implementing resource cleanup for cached BigQueryReadClients to prevent leaks, checking that the error list is not empty before throwing a BigQueryException, and properly catching runtime exceptions like IllegalArgumentException during Arrow schema deserialization and record batch decoding.
|
@gemini-code-assist review |
There was a problem hiding this comment.
Code Review
This pull request introduces beta support for running queries that return Apache Arrow results directly via a new queryArrow method. It adds necessary dependencies, updates the BigQuery interface and its implementation BigQueryImpl, and handles both fast-path and fallback query execution using the BigQuery Storage Read API. Additionally, BigQueryReadClient instances are lazily created and cached by location. A critical issue was identified in BigQueryImpl where the newly introduced readClientLock is declared as final and initialized inline; since BigQueryImpl is serializable, this will lead to a NullPointerException after deserialization. It is recommended to make the lock transient and initialize it inside a custom readObject method.
|
@gemini-code-assist review |
There was a problem hiding this comment.
Code Review
This pull request implements support for streaming query results in Apache Arrow format by introducing the queryArrow API and integrating the BigQuery Storage Read API. Key changes include updating the BigQuery interface to extend AutoCloseable, implementing lazy initialization and caching of BigQueryReadClient in BigQueryImpl, and adding corresponding unit tests. Feedback suggests improving resource cleanup in BigQueryImpl.close() by closing clients in reverse order with exception safety, and enhancing emulator detection in configureReadSettings to support containerized environments.
|
@gemini-code-assist review |
There was a problem hiding this comment.
Code Review
This pull request introduces support for querying and streaming Apache Arrow results directly via the BigQuery Storage Read API. It adds queryArrow methods to the BigQuery interface and implements them in BigQueryImpl, supporting both a fast-path execution and a fallback path via a Storage Read API session. It also implements caching of BigQueryReadClient instances by location. Feedback on the changes highlights a critical issue in configureReadSettings where setting NoCredentialsProvider.create() when credentials are null overrides the default credentials provider chain, breaking the standard Google Application Default Credentials (ADC) flow. It is recommended to omit this fallback so that BigQueryReadSettings can resolve ADC normally.
|
@gemini-code-assist review |
There was a problem hiding this comment.
Code Review
This pull request introduces support for querying BigQuery results in Apache Arrow format via the new queryArrow API. It updates the BigQuery interface to extend AutoCloseable and adds queryArrow methods. In BigQueryImpl, these methods are implemented to support both a fast-path execution using the jobs.query REST RPC and a fallback path using jobs.insert combined with the BigQuery Storage Read API. Additionally, BigQueryImpl now manages and caches BigQueryReadClient instances by location, ensuring they are properly closed. Unit tests have been added to verify the new functionality, including client caching and format validation. I have no feedback to provide.
|
@gemini-code-assist review |
There was a problem hiding this comment.
Code Review
This pull request introduces support for executing queries and retrieving results in Apache Arrow format via the new queryArrow methods in the BigQuery interface and its implementation BigQueryImpl. It adds necessary gRPC dependencies, implements AutoCloseable on BigQuery to manage the lifecycle of cached BigQueryReadClient instances, and updates QueryRequestInfo to serialize Arrow-related options. Additionally, comprehensive unit tests have been added to verify the new functionality, including caching and fallback behaviors. There are no review comments, so no further feedback is provided.
2 participants
This PR introduces the
queryArrowAPI on theBigQueryclient veneer.It enables applications to stream query results as Apache Arrow
VectorSchemaRootbatches with zero memory copies, supporting fast-path query execution and automatic fallback to BigQuery Storage Read API sessions for larger queries.