Skip to content

Commit 1be6e3c

Browse files
authored
samples(bigquery-storage): add Arrow query results samples for query_and_wait and read_rows (#18126)
## Description Adds documentation code snippets and system tests demonstrating high-performance query result retrieval in Apache Arrow format with LZ4 compression using the BigQuery Storage API: 1. **`query_and_wait()` with Arrow format & LZ4 frame compression** (`query_and_wait_arrow.py`): - Demonstrates calling `client.query_and_wait()` with `query_results_format=enums.QueryResultsFormat.ARROW` and `compression_codec=enums.QueryResultsCompressionCodec.LZ4_FRAME`. - Returns results as an iterable of `pyarrow.RecordBatch` via `results.to_arrow_iterable()`. - Wrapped in region tag: `[START bigquerystorage_query_and_wait_arrow]`. 2. **Direct `read_rows` on query job default stream** (`read_rows_query_job.py`): - Demonstrates directly reading query results using `BigQueryReadClient.read_rows` against the job stream `projects/{project}/locations/{location}/jobs/{job_id}/streams/_default`. - Deserializes schema and record batches safely via `pyarrow.ipc`. - Wrapped in region tag: `[START bigquerystorage_read_rows_query_job]`. 3. **Tests & Dependencies**: - Added `query_and_wait_arrow_test.py` and `read_rows_query_job_test.py` verifying batch iteration and schema types. - Added `pyarrow` dependency pins to `samples/snippets/requirements.txt`. Follow-up to #18027 Related to #18047 ## Checklist - [x] Ensure the tests and linter pass - [x] Code coverage does not decrease (if any source code was changed) - [x] Appropriate docs were updated (if necessary)
1 parent da3a9ed commit 1be6e3c

5 files changed

Lines changed: 198 additions & 6 deletions

File tree

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
# [START bigquerystorage_query_and_wait_arrow]
16+
from typing import Iterable
17+
18+
import pyarrow
19+
from google.cloud.bigquery import enums
20+
21+
from google.cloud import bigquery
22+
23+
24+
def query_and_wait_arrow() -> Iterable[pyarrow.RecordBatch]:
25+
"""Queries BigQuery and returns results as an iterable of Apache Arrow RecordBatches.
26+
27+
Returns:
28+
Iterable[pyarrow.RecordBatch]: An iterable of Apache Arrow RecordBatch objects.
29+
"""
30+
# Initialize a BigQuery client.
31+
client = bigquery.Client()
32+
33+
query = """
34+
SELECT name, number, state
35+
FROM `bigquery-public-data.usa_names.usa_1910_current`
36+
LIMIT 100000
37+
"""
38+
39+
# Run the query and wait for results returned directly in Arrow format
40+
# compressed with LZ4_FRAME.
41+
results = client.query_and_wait(
42+
query,
43+
query_results_format=enums.QueryResultsFormat.ARROW,
44+
compression_codec=enums.QueryResultsCompressionCodec.LZ4_FRAME,
45+
)
46+
47+
# Return results as an iterable of pyarrow.RecordBatch objects.
48+
# Each batch contains a slice of the rows in Apache Arrow format.
49+
batches = results.to_arrow_iterable()
50+
return batches
51+
52+
53+
# [END bigquerystorage_query_and_wait_arrow]
Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import pyarrow
16+
17+
from . import query_and_wait_arrow
18+
19+
20+
def test_query_and_wait_arrow():
21+
batches = query_and_wait_arrow.query_and_wait_arrow()
22+
23+
total_rows = 0
24+
batch_count = 0
25+
for batch in batches:
26+
assert isinstance(batch, pyarrow.RecordBatch)
27+
assert batch.schema.names == ["name", "number", "state"]
28+
assert batch.schema.field("name").type == pyarrow.string()
29+
assert batch.schema.field("number").type == pyarrow.int64()
30+
assert batch.schema.field("state").type == pyarrow.string()
31+
total_rows += batch.num_rows
32+
batch_count += 1
33+
34+
assert total_rows == 100000
35+
assert batch_count > 0
Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
# [START bigquerystorage_read_rows_query_job]
16+
from typing import Iterable, Optional
17+
18+
import pyarrow
19+
20+
from google.cloud import bigquery, bigquery_storage_v1
21+
22+
23+
def read_rows_query_job() -> Iterable[pyarrow.RecordBatch]:
24+
"""Queries BigQuery and yields batches directly via BigQueryReadClient using a job stream.
25+
26+
Yields:
27+
pyarrow.RecordBatch: Apache Arrow RecordBatch objects streamed from BigQuery.
28+
"""
29+
# Initialize BigQuery and BigQuery Storage clients.
30+
client = bigquery.Client()
31+
read_client = bigquery_storage_v1.BigQueryReadClient()
32+
33+
query = """
34+
SELECT name, number, state
35+
FROM `bigquery-public-data.usa_names.usa_1910_current`
36+
LIMIT 20000
37+
"""
38+
39+
# Start the query job.
40+
job = client.query(query)
41+
42+
# Construct the job default stream name.
43+
# Format: projects/{project_id}/locations/{location}/jobs/{job_id}/streams/_default
44+
stream = f"projects/{job.project}/locations/{job.location}/jobs/{job.job_id}/streams/_default"
45+
46+
# Read rows directly from the stream using the Storage Read API.
47+
schema: Optional[pyarrow.Schema] = None
48+
49+
for chunk in read_client.read_rows(name=stream, offset=0):
50+
# Extract the schema from the first chunk that provides it.
51+
if (
52+
schema is None
53+
and chunk.arrow_schema
54+
and chunk.arrow_schema.serialized_schema
55+
):
56+
schema = pyarrow.ipc.read_schema(
57+
pyarrow.py_buffer(chunk.arrow_schema.serialized_schema)
58+
)
59+
60+
# Deserialize and yield each record batch using the schema.
61+
if (
62+
chunk.arrow_record_batch
63+
and chunk.arrow_record_batch.serialized_record_batch
64+
):
65+
yield pyarrow.ipc.read_record_batch(
66+
pyarrow.py_buffer(chunk.arrow_record_batch.serialized_record_batch),
67+
schema,
68+
)
69+
70+
71+
# [END bigquerystorage_read_rows_query_job]
Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import pyarrow
16+
17+
from . import read_rows_query_job
18+
19+
20+
def test_read_rows_query_job():
21+
batches = read_rows_query_job.read_rows_query_job()
22+
23+
total_rows = 0
24+
batch_count = 0
25+
for batch in batches:
26+
assert isinstance(batch, pyarrow.RecordBatch)
27+
assert batch.schema.names == ["name", "number", "state"]
28+
assert batch.schema.field("name").type == pyarrow.string()
29+
assert batch.schema.field("number").type == pyarrow.int64()
30+
assert batch.schema.field("state").type == pyarrow.string()
31+
total_rows += batch.num_rows
32+
batch_count += 1
33+
34+
assert total_rows == 20000
35+
assert batch_count > 0
Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,4 @@
1-
google-cloud-bigquery-storage==2.38.0
2-
google-cloud-bigquery===3.30.0; python_version <= '3.8'
3-
google-cloud-bigquery==3.41.0; python_version >= '3.9'
4-
pytest===7.4.3; python_version == '3.7'
5-
pytest===8.3.5; python_version == '3.8'
6-
pytest==9.0.3; python_version >= '3.9'
1+
google-cloud-bigquery-storage==2.41.0
2+
google-cloud-bigquery==3.44.0
3+
pyarrow==24.0.0
4+
pytest==9.0.3

0 commit comments

Comments
 (0)