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()