Repository navigation
Conversation
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?
63d04b5 to
48d916f
Compare
| started: Optional[datetime.datetime] = None, | ||
| ended: Optional[datetime.datetime] = None, | ||
| query_results_format: Optional[str] = None, | ||
| arrow_serialization_options: Optional[Dict[str, str]] = None, |
There was a problem hiding this comment.
nit: It would be nice to force this to be a kwarg too. Optional, since the RowIterator constructor isn't expected to be used directly by customers.
| arrow_serialization_options: Optional[Dict[str, str]] = None, | |
| *, | |
| arrow_serialization_options: Optional[Dict[str, str]] = None, |
| ) | ||
|
|
||
| assert rows._query_results_format == "ARROW" | ||
| assert rows._arrow_serialization_options == _DEFAULT_QUERY_OPTIONS |
There was a problem hiding this comment.
Nit: this would be a false positive if the _wait_or_cancel function somehow mutated the options dictionary. I would recommend creating an "expected" variable that is a copy of the original _DEFAULT_QUERY_OPTIONS to protect against that.
Alternatively, somehow freeze the _DEFAULT_QUERY_OPTIONS dict, but I think that might require an extra dependency, which wouldn't be worth it.
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.
48d916f to
eed63c7
Compare
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.