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..aca875e7f5e3 --- /dev/null +++ b/packages/google-cloud-bigquery-storage/samples/snippets/read_rows_query_job.py @@ -0,0 +1,77 @@ +# 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, 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 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. + + 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() + read_client = bigquery_storage_v1.BigQueryReadClient() + + query = """ + SELECT name, number, state + FROM `bigquery-public-data.usa_names.usa_1910_current` + LIMIT 20000 + """ + + # 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 + + 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 and yield each record batch using the schema. + if ( + chunk.arrow_record_batch + and chunk.arrow_record_batch.serialized_record_batch + ): + yield pyarrow.ipc.read_record_batch( + pyarrow.py_buffer(chunk.arrow_record_batch.serialized_record_batch), + schema, + ) + + +# [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..cb822e30b046 --- /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 == 20000 + 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'