From fe2f752d1fbed1bd803051e6016ce4aaf3ee5dc3 Mon Sep 17 00:00:00 2001 From: ohmayr Date: Tue, 13 Oct 2026 09:32:49 +0000 Subject: [PATCH] feat(bigquery): add timestamp_precision for Arrow query results ARROW results always truncated picosecond TIMESTAMP columns to microseconds, because query_and_wait() had no way to ask for another precision. Client.query() and Client.list_rows() already take timestamp_precision for STRUCT_ENCODING results. Add the same timestamp_precision parameter, with the existing TimestampPrecision enum, to query_and_wait(). PICOSECOND maps to TIMESTAMP_PRECISION_PICOS, which returns picosecond columns at full precision as ISO 8601 strings, and anything else maps to TIMESTAMP_PRECISION_MICROS. Nanoseconds aren't offered, since the picosecond strings already keep full precision. The value goes into the arrowSerializationOptions sent on jobs.query and ReadRows, and setting it without ARROW raises ValueError, like compression_codec. It's labeled [Private Preview], as on Client.query(). --- .../google/cloud/bigquery/_arrow_query.py | 21 +++- .../google/cloud/bigquery/_job_helpers.py | 15 ++- .../google/cloud/bigquery/client.py | 12 ++- .../tests/system/test_arrow.py | 102 ++++++++++++++++++ .../unit/test_query_results_format_arrow.py | 48 +++++++++ 5 files changed, 191 insertions(+), 7 deletions(-) diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py b/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py index 4a235ed66f85..84c9fc853758 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py @@ -42,6 +42,7 @@ def validate_query( query_results_format: Optional[str], *, compression_codec: Optional[str] = None, + timestamp_precision: Optional[enums.TimestampPrecision] = None, ) -> None: """Raise before sending a query whose ``ARROW`` results can't be read. @@ -50,10 +51,13 @@ def validate_query( ``query_and_wait()``. compression_codec: The ``compression_codec`` passed to ``query_and_wait()``. + timestamp_precision: The ``timestamp_precision`` passed to + ``query_and_wait()``. Raises: ValueError: - If ``compression_codec`` is set without ``ARROW``. + If ``compression_codec`` or ``timestamp_precision`` is set + without ``ARROW``. google.cloud.bigquery.exceptions.BigQueryError: If pyarrow or google-cloud-bigquery-storage is missing or too old to read ``ARROW`` results. @@ -61,6 +65,8 @@ def validate_query( if query_results_format != enums.QueryResultsFormat.ARROW: if compression_codec is not None: raise ValueError("compression_codec requires query_results_format=ARROW.") + if timestamp_precision is not None: + raise ValueError("timestamp_precision requires query_results_format=ARROW.") return # Check before the query runs, so a query doesn't succeed while its @@ -84,7 +90,10 @@ def _check_dependencies() -> None: ) -def serialization_options(compression_codec: Optional[str]) -> Dict[str, str]: +def serialization_options( + compression_codec: Optional[str], + timestamp_precision: Optional[enums.TimestampPrecision] = None, +) -> Dict[str, str]: """Build the ``arrowSerializationOptions`` sent on ``jobs.query``. ``ReadRows`` gets the same options. Options the caller didn't set are sent @@ -94,7 +103,13 @@ def serialization_options(compression_codec: Optional[str]) -> Dict[str, str]: return { "bufferCompression": getattr(compression_codec, "value", compression_codec) or "COMPRESSION_UNSPECIFIED", - "picosTimestampPrecision": "TIMESTAMP_PRECISION_MICROS", + # Picosecond values come back as ISO 8601 strings at full precision, + # so there's no need to offer nanoseconds. + "picosTimestampPrecision": ( + "TIMESTAMP_PRECISION_PICOS" + if timestamp_precision == enums.TimestampPrecision.PICOSECOND + else "TIMESTAMP_PRECISION_MICROS" + ), } diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py b/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py index 9a458374f83d..8470c19becd2 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py @@ -433,6 +433,7 @@ def query_and_wait( max_results: Optional[int] = None, query_results_format: Optional[str] = None, compression_codec: Optional[str] = None, + timestamp_precision: Optional[enums.TimestampPrecision] = None, callback: Callable = lambda _: None, ) -> table.RowIterator: """Run the query, wait for it to finish, and return the results. @@ -482,6 +483,11 @@ def query_and_wait( [Beta] The format for query results (e.g. "ARROW" or :class:`~google.cloud.bigquery.enums.QueryResultsFormat.ARROW`). compression_codec (Optional[Union[str, google.cloud.bigquery.enums.QueryResultsCompressionCodec]]): [Beta] Compression codec for Arrow serialization (e.g. "LZ4_FRAME" or :class:`~google.cloud.bigquery.enums.QueryResultsCompressionCodec.LZ4_FRAME`). + timestamp_precision (Optional[enums.TimestampPrecision]): + [Private Preview] Precision of picosecond ``TIMESTAMP`` columns in + ``ARROW`` results. With ``PICOSECOND``, they arrive at full + precision as ISO 8601 strings. Otherwise, they're truncated to + microseconds. Requires ``query_results_format=ARROW``. callback (Callable): A callback function used by bigframes to report query progress. @@ -503,13 +509,16 @@ def query_and_wait( :class:`~google.cloud.bigquery.job.QueryJobConfig` class. ValueError: - If ``compression_codec`` is set without ``ARROW``. + If ``compression_codec`` or ``timestamp_precision`` is set + without ``ARROW``. google.cloud.bigquery.exceptions.BigQueryError: If pyarrow or google-cloud-bigquery-storage is missing or too old to read ``ARROW`` results. The query isn't sent. """ _arrow_query.validate_query( - query_results_format, compression_codec=compression_codec + query_results_format, + compression_codec=compression_codec, + timestamp_precision=timestamp_precision, ) request_body = _to_query_request( @@ -520,7 +529,7 @@ def query_and_wait( arrow_serialization_options = None if query_results_format == enums.QueryResultsFormat.ARROW: arrow_serialization_options = _arrow_query.serialization_options( - compression_codec + compression_codec, timestamp_precision ) request_body["arrowSerializationOptions"] = arrow_serialization_options diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/client.py b/packages/google-cloud-bigquery/google/cloud/bigquery/client.py index 214982dc2a23..b78438c5bc23 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/client.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/client.py @@ -3677,6 +3677,7 @@ def query_and_wait( max_results: Optional[int] = None, query_results_format: Optional[str] = None, compression_codec: Optional[str] = None, + timestamp_precision: Optional[enums.TimestampPrecision] = None, ) -> RowIterator: """Run the query, wait for it to finish, and return the results. @@ -3728,6 +3729,11 @@ def query_and_wait( [Beta] The format for query results (e.g. "ARROW" or :class:`~google.cloud.bigquery.enums.QueryResultsFormat.ARROW`). compression_codec (Optional[Union[str, google.cloud.bigquery.enums.QueryResultsCompressionCodec]]): [Beta] Compression codec for Arrow serialization (e.g. "LZ4_FRAME" or :class:`~google.cloud.bigquery.enums.QueryResultsCompressionCodec.LZ4_FRAME`). + timestamp_precision (Optional[enums.TimestampPrecision]): + [Private Preview] Precision of picosecond ``TIMESTAMP`` columns + in ``ARROW`` results. With ``PICOSECOND``, they arrive at full + precision as ISO 8601 strings. Otherwise, they're truncated to + microseconds. Requires ``query_results_format=ARROW``. Returns: google.cloud.bigquery.table.RowIterator: @@ -3747,7 +3753,8 @@ def query_and_wait( :class:`~google.cloud.bigquery.job.QueryJobConfig` class. ValueError: - If ``compression_codec`` is set without ``ARROW``. + If ``compression_codec`` or ``timestamp_precision`` is set + without ``ARROW``. google.cloud.bigquery.exceptions.BigQueryError: If pyarrow or google-cloud-bigquery-storage is missing or too old to read ``ARROW`` results. The query isn't sent. @@ -3765,6 +3772,7 @@ def query_and_wait( max_results=max_results, query_results_format=query_results_format, compression_codec=compression_codec, + timestamp_precision=timestamp_precision, ) def _query_and_wait_bigframes( @@ -3782,6 +3790,7 @@ def _query_and_wait_bigframes( max_results: Optional[int] = None, query_results_format: Optional[str] = None, compression_codec: Optional[str] = None, + timestamp_precision: Optional[enums.TimestampPrecision] = None, callback: Callable = lambda _: None, ) -> RowIterator: """See query_and_wait. @@ -3816,6 +3825,7 @@ def _query_and_wait_bigframes( max_results=max_results, query_results_format=query_results_format, compression_codec=compression_codec, + timestamp_precision=timestamp_precision, callback=callback, ) diff --git a/packages/google-cloud-bigquery/tests/system/test_arrow.py b/packages/google-cloud-bigquery/tests/system/test_arrow.py index b6e4c4a1d079..4c4bf4fcf583 100644 --- a/packages/google-cloud-bigquery/tests/system/test_arrow.py +++ b/packages/google-cloud-bigquery/tests/system/test_arrow.py @@ -14,11 +14,13 @@ """System tests for Arrow connector.""" +import datetime from typing import Optional import pyarrow import pytest +from google.api_core import exceptions from google.cloud import bigquery from google.cloud.bigquery import enums @@ -310,3 +312,103 @@ def read(**kwargs): assert compressed.num_rows == 1000 assert compressed.equals(read()) + + +_MICROSECOND_VALUES = [ + datetime.datetime(2025, 1, day, 0, 0, 0, 123456, tzinfo=datetime.timezone.utc) + for day in (1, 2, 3) +] +_TIMESTAMP_PRECISIONS = pytest.mark.parametrize( + "timestamp_precision", + [ + None, + enums.TimestampPrecision.MICROSECOND, + enums.TimestampPrecision.PICOSECOND, + ], + ids=["default", "microsecond", "picosecond"], +) +_PAGE_SIZES = pytest.mark.parametrize( + "page_size", [None, 1], ids=["first_page_only", "with_stream"] +) + + +def _read_timestamps(bigquery_client, query, timestamp_precision, page_size): + rows = bigquery_client.query_and_wait( + query, + job_config=bigquery.QueryJobConfig(use_query_cache=False), + page_size=page_size, + query_results_format=enums.QueryResultsFormat.ARROW, + timestamp_precision=timestamp_precision, + ) + table = pyarrow.Table.from_batches(list(rows.to_arrow_iterable())) + return table.column("ts").to_pylist() + + +@_PAGE_SIZES +@_TIMESTAMP_PRECISIONS +def test_query_and_wait_arrow_timestamp_precision_microsecond_column( + bigquery_client, timestamp_precision, page_size +): + values = _read_timestamps( + bigquery_client, + "SELECT TIMESTAMP_ADD(TIMESTAMP '2025-01-01 00:00:00.123456 UTC', " + "INTERVAL x DAY) AS ts FROM UNNEST([0, 1, 2]) AS x ORDER BY x", + timestamp_precision, + page_size, + ) + + assert values == _MICROSECOND_VALUES + + +@_PAGE_SIZES +@_TIMESTAMP_PRECISIONS +def test_query_and_wait_arrow_timestamp_precision_picosecond_column( + bigquery_client, timestamp_precision, page_size +): + try: + values = _read_timestamps( + bigquery_client, + "SELECT CAST(FORMAT('2025-01-0%dT00:00:00.123456789012Z', x) " + "AS TIMESTAMP(12)) AS ts FROM UNNEST([1, 2, 3]) AS x ORDER BY x", + timestamp_precision, + page_size, + ) + except exceptions.BadRequest as exc: + # TIMESTAMP(12) is enabled per project. The test project has it, but + # a developer's own project may not. + if "Timestamp precision type parameter is not supported" not in str(exc): + raise + pytest.xfail(f"TIMESTAMP(12) isn't enabled in this project: {exc.message}") + + if timestamp_precision == enums.TimestampPrecision.PICOSECOND: + assert [value[:10] for value in values] == [ + "2025-01-01", + "2025-01-02", + "2025-01-03", + ] + assert all(".123456789012" in value for value in values) + else: + assert values == _MICROSECOND_VALUES + + +@_PAGE_SIZES +@_TIMESTAMP_PRECISIONS +def test_query_and_wait_arrow_timestamp_precision_picosecond_table( + bigquery_client, scalars_table_pico, timestamp_precision, page_size +): + values = _read_timestamps( + bigquery_client, + f"SELECT pico_col AS ts FROM `{scalars_table_pico}` ORDER BY pico_col", + timestamp_precision, + page_size, + ) + + if timestamp_precision == enums.TimestampPrecision.PICOSECOND: + assert [value[:10] for value in values] == [ + "2025-01-01", + "2025-01-02", + "2025-01-03", + ] + assert all(".123456789012" in value for value in values) + else: + assert values == _MICROSECOND_VALUES diff --git a/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py b/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py index b06151b9a6a3..4a576007e3a5 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py +++ b/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py @@ -25,6 +25,7 @@ from google.cloud.bigquery import exceptions from google.cloud.bigquery.client import Client from google.cloud.bigquery.enums import QueryResultsCompressionCodec +from google.cloud.bigquery.enums import TimestampPrecision from google.cloud.bigquery.enums import WriteDisposition from google.cloud.bigquery.job import QueryJobConfig from google.cloud.bigquery.table import RowIterator, _EmptyRowIterator @@ -909,5 +910,52 @@ def test_download_arrow_from_job_id_sends_query_options_on_read_rows(): ) +@pytest.mark.parametrize( + "timestamp_precision", + [TimestampPrecision.MICROSECOND, TimestampPrecision.PICOSECOND], +) +def test_query_and_wait_timestamp_precision_requires_arrow( + check_dependencies, timestamp_precision +): + client = mock.MagicMock(spec=Client) + + with pytest.raises(ValueError, match="timestamp_precision requires"): + _query_and_wait(client, timestamp_precision=timestamp_precision) + + client._call_api.assert_not_called() + check_dependencies.assert_not_called() + + +@pytest.mark.parametrize( + ("timestamp_precision", "expected"), + [ + (None, "TIMESTAMP_PRECISION_MICROS"), + (TimestampPrecision.MICROSECOND, "TIMESTAMP_PRECISION_MICROS"), + (TimestampPrecision.PICOSECOND, "TIMESTAMP_PRECISION_PICOS"), + ], +) +def test_client_query_and_wait_arrow_timestamp_precision(timestamp_precision, expected): + client = make_client() + client._call_api = mock.Mock( + return_value={ + "jobReference": {"projectId": "PROJECT", "jobId": "j", "location": "US"}, + "jobComplete": True, + "schema": {"fields": []}, + } + ) + + rows = client.query_and_wait( + "SELECT 1", + query_results_format="ARROW", + timestamp_precision=timestamp_precision, + ) + + data = client._call_api.call_args.kwargs["data"] + assert data["arrowSerializationOptions"]["picosTimestampPrecision"] == expected + assert rows._arrow_serialization_options["picosTimestampPrecision"] == expected + # formatOptions only shapes STRUCT_ENCODING rows, which ARROW doesn't return. + assert "timestampOutputFormat" not in data["formatOptions"] + + if __name__ == "__main__": unittest.main()