From 509b0c453cf64f6a16018ca8c92f22d2d5954b32 Mon Sep 17 00:00:00 2001 From: Alex Tolpin Date: Sun, 16 Aug 2026 17:48:26 +0000 Subject: [PATCH 1/2] arrow snippets --- .../samples/snippets/query_and_wait_arrow.py | 58 +++++++++++++ .../snippets/query_and_wait_arrow_test.py | 35 ++++++++ .../samples/snippets/read_rows_query_job.py | 82 +++++++++++++++++++ .../snippets/read_rows_query_job_test.py | 35 ++++++++ .../samples/snippets/requirements.txt | 3 + 5 files changed, 213 insertions(+) create mode 100644 packages/google-cloud-bigquery-storage/samples/snippets/query_and_wait_arrow.py create mode 100644 packages/google-cloud-bigquery-storage/samples/snippets/query_and_wait_arrow_test.py create mode 100644 packages/google-cloud-bigquery-storage/samples/snippets/read_rows_query_job.py create mode 100644 packages/google-cloud-bigquery-storage/samples/snippets/read_rows_query_job_test.py diff --git a/packages/google-cloud-bigquery-storage/samples/snippets/query_and_wait_arrow.py b/packages/google-cloud-bigquery-storage/samples/snippets/query_and_wait_arrow.py new file mode 100644 index 000000000000..c19c97f3eff3 --- /dev/null +++ b/packages/google-cloud-bigquery-storage/samples/snippets/query_and_wait_arrow.py @@ -0,0 +1,58 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# [START bigquerystorage_query_and_wait_arrow] +from typing import Iterable, Optional + +from google.cloud import bigquery +from google.cloud.bigquery import enums +import pyarrow + + +def query_and_wait_arrow( + project_id: Optional[str] = None, +) -> Iterable[pyarrow.RecordBatch]: + """Queries BigQuery and returns results as an iterable of Apache Arrow RecordBatches. + + Args: + project_id (Optional[str]): The Google Cloud project ID to bill for the query. + If not specified, the project is inferred from the environment. + + Returns: + Iterable[pyarrow.RecordBatch]: An iterable of Apache Arrow RecordBatch objects. + """ + # Initialize a BigQuery client. + client = bigquery.Client(project=project_id) if project_id else bigquery.Client() + + query = """ + SELECT name, number, state + FROM `bigquery-public-data.usa_names.usa_1910_current` + LIMIT 100000 + """ + + # Run the query and wait for results returned directly in Arrow format + # compressed with LZ4_FRAME. + results = client.query_and_wait( + query, + query_results_format=enums.QueryResultsFormat.ARROW, + compression_codec=enums.QueryResultsCompressionCodec.LZ4_FRAME, + ) + + # Return results as an iterable of pyarrow.RecordBatch objects. + # Each batch contains a slice of the rows in Apache Arrow format. + batches = results.to_arrow_iterable() + return batches + + +# [END bigquerystorage_query_and_wait_arrow] diff --git a/packages/google-cloud-bigquery-storage/samples/snippets/query_and_wait_arrow_test.py b/packages/google-cloud-bigquery-storage/samples/snippets/query_and_wait_arrow_test.py new file mode 100644 index 000000000000..fc8285f5f78d --- /dev/null +++ b/packages/google-cloud-bigquery-storage/samples/snippets/query_and_wait_arrow_test.py @@ -0,0 +1,35 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pyarrow + +from . import query_and_wait_arrow + + +def test_query_and_wait_arrow(project_id: str): + batches = query_and_wait_arrow.query_and_wait_arrow(project_id=project_id) + + total_rows = 0 + batch_count = 0 + for batch in batches: + assert isinstance(batch, pyarrow.RecordBatch) + assert batch.schema.names == ["name", "number", "state"] + assert batch.schema.field("name").type == pyarrow.string() + assert batch.schema.field("number").type == pyarrow.int64() + assert batch.schema.field("state").type == pyarrow.string() + total_rows += batch.num_rows + batch_count += 1 + + assert total_rows == 100000 + assert batch_count > 0 diff --git a/packages/google-cloud-bigquery-storage/samples/snippets/read_rows_query_job.py b/packages/google-cloud-bigquery-storage/samples/snippets/read_rows_query_job.py new file mode 100644 index 000000000000..070cde795217 --- /dev/null +++ b/packages/google-cloud-bigquery-storage/samples/snippets/read_rows_query_job.py @@ -0,0 +1,82 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# [START bigquerystorage_read_rows_query_job] +from typing import Iterable, List, Optional + +from google.cloud import bigquery +from google.cloud import bigquery_storage_v1 +import pyarrow + + +def read_rows_query_job( + project_id: Optional[str] = None, +) -> Iterable[pyarrow.RecordBatch]: + """Queries BigQuery and reads results directly via BigQueryReadClient using a job stream. + + Args: + project_id (Optional[str]): The Google Cloud project ID to bill for the query. + If not specified, the project is inferred from the environment. + + Returns: + Iterable[pyarrow.RecordBatch]: An iterable of Apache Arrow RecordBatch objects. + """ + # Initialize BigQuery and BigQuery Storage clients. + client = bigquery.Client(project=project_id) if project_id else bigquery.Client() + read_client = bigquery_storage_v1.BigQueryReadClient() + + query = """ + SELECT name, number, state + FROM `bigquery-public-data.usa_names.usa_1910_current` + LIMIT 100 + """ + + # Start the query job. + job = client.query(query) + + # Construct the job default stream name. + # Format: projects/{project_id}/locations/{location}/jobs/{job_id}/streams/_default + stream = f"projects/{job.project}/locations/{job.location}/jobs/{job.job_id}/streams/_default" + + # Read rows directly from the stream using the Storage Read API. + schema: Optional[pyarrow.Schema] = None + batches: List[pyarrow.RecordBatch] = [] + + for chunk in read_client.read_rows(name=stream, offset=0): + # Extract the schema from the first chunk that provides it. + if ( + schema is None + and chunk.arrow_schema + and chunk.arrow_schema.serialized_schema + ): + schema = pyarrow.ipc.read_schema( + pyarrow.py_buffer(chunk.arrow_schema.serialized_schema) + ) + + # Deserialize each record batch using the schema. + if ( + chunk.arrow_record_batch + and chunk.arrow_record_batch.serialized_record_batch + ): + batches.append( + pyarrow.ipc.read_record_batch( + pyarrow.py_buffer(chunk.arrow_record_batch.serialized_record_batch), + schema, + ) + ) + + return batches + + +# [END bigquerystorage_read_rows_query_job] diff --git a/packages/google-cloud-bigquery-storage/samples/snippets/read_rows_query_job_test.py b/packages/google-cloud-bigquery-storage/samples/snippets/read_rows_query_job_test.py new file mode 100644 index 000000000000..b648214dba9d --- /dev/null +++ b/packages/google-cloud-bigquery-storage/samples/snippets/read_rows_query_job_test.py @@ -0,0 +1,35 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pyarrow + +from . import read_rows_query_job + + +def test_read_rows_query_job(project_id: str): + batches = read_rows_query_job.read_rows_query_job(project_id=project_id) + + total_rows = 0 + batch_count = 0 + for batch in batches: + assert isinstance(batch, pyarrow.RecordBatch) + assert batch.schema.names == ["name", "number", "state"] + assert batch.schema.field("name").type == pyarrow.string() + assert batch.schema.field("number").type == pyarrow.int64() + assert batch.schema.field("state").type == pyarrow.string() + total_rows += batch.num_rows + batch_count += 1 + + assert total_rows == 100 + assert batch_count > 0 diff --git a/packages/google-cloud-bigquery-storage/samples/snippets/requirements.txt b/packages/google-cloud-bigquery-storage/samples/snippets/requirements.txt index 7927a9f7c8fa..4896a882a004 100644 --- a/packages/google-cloud-bigquery-storage/samples/snippets/requirements.txt +++ b/packages/google-cloud-bigquery-storage/samples/snippets/requirements.txt @@ -1,6 +1,9 @@ google-cloud-bigquery-storage==2.38.0 google-cloud-bigquery===3.30.0; python_version <= '3.8' google-cloud-bigquery==3.41.0; python_version >= '3.9' +pyarrow===12.0.1; python_version == '3.7' +pyarrow===17.0.0; python_version == '3.8' +pyarrow==24.0.0; python_version >= '3.9' pytest===7.4.3; python_version == '3.7' pytest===8.3.5; python_version == '3.8' pytest==9.0.3; python_version >= '3.9' From 7a151fb0a3c3067b2a0babfe6807c0f5ff6140a0 Mon Sep 17 00:00:00 2001 From: Alex Tolpin Date: Sun, 16 Aug 2026 17:58:38 +0000 Subject: [PATCH 2/2] yield batches --- .../samples/snippets/read_rows_query_job.py | 23 ++++++++----------- .../snippets/read_rows_query_job_test.py | 2 +- 2 files changed, 10 insertions(+), 15 deletions(-) diff --git a/packages/google-cloud-bigquery-storage/samples/snippets/read_rows_query_job.py b/packages/google-cloud-bigquery-storage/samples/snippets/read_rows_query_job.py index 070cde795217..aca875e7f5e3 100644 --- a/packages/google-cloud-bigquery-storage/samples/snippets/read_rows_query_job.py +++ b/packages/google-cloud-bigquery-storage/samples/snippets/read_rows_query_job.py @@ -13,7 +13,7 @@ # limitations under the License. # [START bigquerystorage_read_rows_query_job] -from typing import Iterable, List, Optional +from typing import Iterable, Optional from google.cloud import bigquery from google.cloud import bigquery_storage_v1 @@ -23,14 +23,14 @@ def read_rows_query_job( project_id: Optional[str] = None, ) -> Iterable[pyarrow.RecordBatch]: - """Queries BigQuery and reads results directly via BigQueryReadClient using a job stream. + """Queries BigQuery and yields batches directly via BigQueryReadClient using a job stream. Args: project_id (Optional[str]): The Google Cloud project ID to bill for the query. If not specified, the project is inferred from the environment. - Returns: - Iterable[pyarrow.RecordBatch]: An iterable of Apache Arrow RecordBatch objects. + Yields: + pyarrow.RecordBatch: Apache Arrow RecordBatch objects streamed from BigQuery. """ # Initialize BigQuery and BigQuery Storage clients. client = bigquery.Client(project=project_id) if project_id else bigquery.Client() @@ -39,7 +39,7 @@ def read_rows_query_job( query = """ SELECT name, number, state FROM `bigquery-public-data.usa_names.usa_1910_current` - LIMIT 100 + LIMIT 20000 """ # Start the query job. @@ -51,7 +51,6 @@ def read_rows_query_job( # Read rows directly from the stream using the Storage Read API. schema: Optional[pyarrow.Schema] = None - batches: List[pyarrow.RecordBatch] = [] for chunk in read_client.read_rows(name=stream, offset=0): # Extract the schema from the first chunk that provides it. @@ -64,19 +63,15 @@ def read_rows_query_job( pyarrow.py_buffer(chunk.arrow_schema.serialized_schema) ) - # Deserialize each record batch using the schema. + # Deserialize and yield each record batch using the schema. if ( chunk.arrow_record_batch and chunk.arrow_record_batch.serialized_record_batch ): - batches.append( - pyarrow.ipc.read_record_batch( - pyarrow.py_buffer(chunk.arrow_record_batch.serialized_record_batch), - schema, - ) + yield pyarrow.ipc.read_record_batch( + pyarrow.py_buffer(chunk.arrow_record_batch.serialized_record_batch), + schema, ) - return batches - # [END bigquerystorage_read_rows_query_job] diff --git a/packages/google-cloud-bigquery-storage/samples/snippets/read_rows_query_job_test.py b/packages/google-cloud-bigquery-storage/samples/snippets/read_rows_query_job_test.py index b648214dba9d..cb822e30b046 100644 --- a/packages/google-cloud-bigquery-storage/samples/snippets/read_rows_query_job_test.py +++ b/packages/google-cloud-bigquery-storage/samples/snippets/read_rows_query_job_test.py @@ -31,5 +31,5 @@ def test_read_rows_query_job(project_id: str): total_rows += batch.num_rows batch_count += 1 - assert total_rows == 100 + assert total_rows == 20000 assert batch_count > 0