Repository navigation
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 passing arrowSerializationOptions (such as compression codecs) from jobs.query down to ReadRows in the BigQuery Storage API when using the ARROW format. This involves updating RowIterator to store these options, modifying query_and_wait to set them, and updating download_arrow_from_job to pass them to read_rows. Feedback was provided to make _read_rows_options more robust against potential KeyErrors by using .get() with default fallbacks instead of direct dictionary lookups.
| def _read_rows_options(options: Optional[Dict[str, str]]) -> Dict[str, str]: | ||
| """Convert ``jobs.query`` ``arrowSerializationOptions`` to the | ||
| ``ArrowSerializationOptions`` fields that ``ReadRows`` takes.""" | ||
| options = options or serialization_options(None) | ||
| return { | ||
| "buffer_compression": options["bufferCompression"], | ||
| "picos_timestamp_precision": options["picosTimestampPrecision"], | ||
| } |
There was a problem hiding this comment.
Using direct dictionary lookups like options["bufferCompression"] can raise a KeyError if the dictionary is passed with missing keys (for example, if a user instantiates RowIterator with custom options). It is safer and more robust to use .get() with fallback defaults, which also simplifies the fallback logic when options is None.
| def _read_rows_options(options: Optional[Dict[str, str]]) -> Dict[str, str]: | |
| """Convert ``jobs.query`` ``arrowSerializationOptions`` to the | |
| ``ArrowSerializationOptions`` fields that ``ReadRows`` takes.""" | |
| options = options or serialization_options(None) | |
| return { | |
| "buffer_compression": options["bufferCompression"], | |
| "picos_timestamp_precision": options["picosTimestampPrecision"], | |
| } | |
| def _read_rows_options(options: Optional[Dict[str, str]]) -> Dict[str, str]: | |
| """Convert jobs.query arrowSerializationOptions to the | |
| ArrowSerializationOptions fields that ReadRows takes.""" | |
| options = options or {} | |
| return { | |
| "buffer_compression": options.get("bufferCompression", "COMPRESSION_UNSPECIFIED"), | |
| "picos_timestamp_precision": options.get("picosTimestampPrecision", "TIMESTAMP_PRECISION_MICROS"), | |
| } |
There was a problem hiding this comment.
These options always come from serialization_options(), which sets both keys, so a missing key would be a bug I'd want to see as a KeyError rather than paper over with defaults. RowIterator's constructor is internal to query_and_wait() here.
a8b7900 to
6543176
Compare
8e73bb3 to
63d04b5
Compare
| return { | ||
| "bufferCompression": getattr(compression_codec, "value", compression_codec) | ||
| or "COMPRESSION_UNSPECIFIED", | ||
| "picosTimestampPrecision": "TIMESTAMP_PRECISION_MICROS", |
There was a problem hiding this comment.
I assume it'll be a separate project to plumb through all of the TIMESTAMP_PRECISION_PICOS features at some point in the future?
query_and_wait() sent compression_codec inside formatOptions, which only has useInt64Timestamp and timestampOutputFormat, so jobs.query ignored it, and ReadRows never got it at all. The first page and the stream each used their own endpoint's defaults, so one result could mix encodings. Build one arrowSerializationOptions value for ARROW queries and send it at the top level of the jobs.query request, where it also joins the jobs.query allowlist. Options the caller didn't set are sent with their defaults (COMPRESSION_UNSPECIFIED, TIMESTAMP_PRECISION_MICROS), so the encoding never depends on server defaults. RowIterator keeps the value on both the jobs.query and the waited-job paths and sends it on read_rows(). jobs.insert never gets these fields, since they come from query_and_wait() arguments rather than QueryJobConfig. The mypy session now installs the in-repo google-cloud-bigquery-storage too, so the new read_rows() argument is type-checked.
google-cloud-bigquery-storage 2.43.0 has the read_rows() arrow_serialization_options argument, so mypy can type-check against the released package instead of the in-repo copy.
63d04b5 to
48d916f
Compare
This branch has not been deployed
2 participants
Sends the same
arrowSerializationOptionsonjobs.queryandReadRows, so the first page and the stream use the same Arrow encoding. Previouslycompression_codecwas sent informatOptions, whichjobs.queryignores, and never reachedReadRows.arrowSerializationOptionsat the top level ofjobs.query, with explicit defaults for anything the caller didn't set.read_rows()(feat(bigquery-storage): add Arrow options and cancel to read_rows #18643).LZ4_FRAMEandZSTDwith uncompressed results.