From 78b8560c9bf08c264b7fc72ee531624600642834 Mon Sep 17 00:00:00 2001 From: Shuowei Li Date: Wed, 5 Aug 2026 19:59:47 +0000 Subject: [PATCH 1/9] test(bigquery-storage): add dual-environment MVP validation matrix --- .../tests/unit/test_storage_mvp_validation.py | 191 ++++++++++++++++++ 1 file changed, 191 insertions(+) create mode 100644 packages/google-cloud-bigquery-storage/tests/unit/test_storage_mvp_validation.py diff --git a/packages/google-cloud-bigquery-storage/tests/unit/test_storage_mvp_validation.py b/packages/google-cloud-bigquery-storage/tests/unit/test_storage_mvp_validation.py new file mode 100644 index 000000000000..4cedda6c2584 --- /dev/null +++ b/packages/google-cloud-bigquery-storage/tests/unit/test_storage_mvp_validation.py @@ -0,0 +1,191 @@ +# -*- coding: utf-8 -*- +# +# 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 datetime +import decimal +from unittest import mock + +from google.cloud.bigquery_storage import types +from google.cloud.bigquery_storage_v1 import reader +import pandas_gbq.arrow +import pytest + +pyarrow = pytest.importorskip("pyarrow") + + +_TEST_SCHEMA = pyarrow.schema( + [ + pyarrow.field("int_col", pyarrow.int64()), + pyarrow.field("float_col", pyarrow.float64()), + pyarrow.field("num_col", pyarrow.decimal128(38, 9)), + pyarrow.field("bool_col", pyarrow.bool_()), + pyarrow.field("str_col", pyarrow.utf8()), + pyarrow.field("bytes_col", pyarrow.binary()), + pyarrow.field("date_col", pyarrow.date32()), + pyarrow.field("ts_col", pyarrow.timestamp("us", tz="UTC")), + ] +) + +_TEST_ROWS = [ + { + "int_col": 123, + "float_col": 3.14, + "num_col": decimal.Decimal("9.99"), + "bool_col": True, + "str_col": "hello world", + "bytes_col": b"ascii bytes", + "date_col": datetime.date(1998, 9, 4), + "ts_col": datetime.datetime(2000, 1, 1, 5, 0, tzinfo=datetime.timezone.utc), + }, + { + "int_col": 456, + "float_col": 2.72, + "num_col": decimal.Decimal("0.99"), + "bool_col": False, + "str_col": "hallo welt", + "bytes_col": b"\xbb\xee\xff", + "date_col": datetime.date(1995, 3, 2), + "ts_col": datetime.datetime(1965, 4, 3, 2, 1, tzinfo=datetime.timezone.utc), + }, +] + + +def _create_sample_batch(schema=_TEST_SCHEMA, rows=_TEST_ROWS): + arrays = [ + pyarrow.array([row[field.name] for row in rows], type=field.type) + for field in schema + ] + return pyarrow.RecordBatch.from_arrays(arrays, schema=schema) + + +def _create_read_rows_response(record_batch=None): + response = types.ReadRowsResponse() + if record_batch is not None: + response.row_count = record_batch.num_rows + response.arrow_record_batch.serialized_record_batch = ( + record_batch.serialize().to_pybytes() + ) + else: + response.row_count = 0 + response.arrow_record_batch.serialized_record_batch = b"" + return response + + +def _create_read_session(schema=_TEST_SCHEMA): + return types.ReadSession( + arrow_schema={"serialized_schema": schema.serialize().to_pybytes()} + ) + + +def test_storage_mvp_direct_pandas_gbq_arrow_contract(): + expected_batch = _create_sample_batch() + response = _create_read_rows_response(expected_batch) + + actual_batch = pandas_gbq.arrow.from_read_rows_response( + response, arrow_schema=_TEST_SCHEMA + ) + + assert actual_batch.equals(expected_batch) + + +def test_storage_mvp_direct_pandas_gbq_arrow_empty_response(): + response = _create_read_rows_response(record_batch=None) + + actual_batch = pandas_gbq.arrow.from_read_rows_response( + response, arrow_schema=_TEST_SCHEMA + ) + + assert actual_batch.num_rows == 0 + assert actual_batch.schema == _TEST_SCHEMA + + +def test_storage_mvp_read_rows_page_dual_environment_equivalence(): + expected_batch = _create_sample_batch() + response = _create_read_rows_response(expected_batch) + + # Environment A: With pandas-gbq installed (Delegated Pathway) + stream_parser_a = reader._ArrowStreamParser(_create_read_session(_TEST_SCHEMA)) + page_a = reader.ReadRowsPage(stream_parser_a, response) + with pytest.warns( + PendingDeprecationWarning, + match="google-cloud-bigquery-storage is deprecated", + ): + batch_delegated = page_a.to_arrow() + + # Environment B: Without pandas-gbq (Legacy Fallback Pathway) + stream_parser_b = reader._ArrowStreamParser(_create_read_session(_TEST_SCHEMA)) + page_b = reader.ReadRowsPage(stream_parser_b, response) + with ( + mock.patch.dict("sys.modules", {"pandas_gbq": None, "pandas_gbq.arrow": None}), + pytest.warns( + PendingDeprecationWarning, + match="google-cloud-bigquery-storage is deprecated", + ), + ): + batch_fallback = page_b.to_arrow() + + assert batch_delegated.equals(expected_batch) + assert batch_fallback.equals(expected_batch) + assert batch_delegated.equals(batch_fallback) + + +def test_storage_mvp_read_rows_page_empty_dual_environment_equivalence(): + empty_batch = _create_sample_batch(rows=[]) + response = _create_read_rows_response(record_batch=empty_batch) + + # Environment A: With pandas-gbq installed + stream_parser_a = reader._ArrowStreamParser(_create_read_session(_TEST_SCHEMA)) + page_a = reader.ReadRowsPage(stream_parser_a, response) + with pytest.warns(PendingDeprecationWarning): + batch_delegated = page_a.to_arrow() + + # Environment B: Without pandas-gbq installed + stream_parser_b = reader._ArrowStreamParser(_create_read_session(_TEST_SCHEMA)) + page_b = reader.ReadRowsPage(stream_parser_b, response) + with ( + mock.patch.dict("sys.modules", {"pandas_gbq": None, "pandas_gbq.arrow": None}), + pytest.warns(PendingDeprecationWarning), + ): + batch_fallback = page_b.to_arrow() + + assert batch_delegated.num_rows == 0 + assert batch_fallback.num_rows == 0 + assert batch_delegated.schema == _TEST_SCHEMA + assert batch_fallback.schema == _TEST_SCHEMA + + +def test_storage_mvp_stream_to_arrow_end_to_end(): + batch1 = _create_sample_batch() + batch2 = _create_sample_batch() + response1 = _create_read_rows_response(batch1) + response2 = _create_read_rows_response(batch2) + read_session = _create_read_session(_TEST_SCHEMA) + + gapic_client = mock.Mock() + gapic_client.read_rows.return_value = iter([response1, response2]) + + read_rows_stream = reader.ReadRowsStream( + gapic_client, + "projects/p/locations/l/sessions/s/streams/str1", + 0, + {}, + ) + + with pytest.warns(PendingDeprecationWarning): + table = read_rows_stream.to_arrow(read_session=read_session) + + assert table.num_rows == batch1.num_rows + batch2.num_rows + assert table.schema == _TEST_SCHEMA From 89493375cc07730af8436ecf27c85670ba91bad4 Mon Sep 17 00:00:00 2001 From: Shuowei Li Date: Wed, 5 Aug 2026 21:27:27 +0000 Subject: [PATCH 2/9] style(bigquery-storage): align MVP validation tests with unit_test.txt --- .../tests/unit/test_storage_mvp_validation.py | 80 +++++++++++-------- 1 file changed, 47 insertions(+), 33 deletions(-) diff --git a/packages/google-cloud-bigquery-storage/tests/unit/test_storage_mvp_validation.py b/packages/google-cloud-bigquery-storage/tests/unit/test_storage_mvp_validation.py index 4cedda6c2584..631e392d6c5f 100644 --- a/packages/google-cloud-bigquery-storage/tests/unit/test_storage_mvp_validation.py +++ b/packages/google-cloud-bigquery-storage/tests/unit/test_storage_mvp_validation.py @@ -90,7 +90,7 @@ def _create_read_session(schema=_TEST_SCHEMA): ) -def test_storage_mvp_direct_pandas_gbq_arrow_contract(): +def test_from_read_rows_response_decodes_serialized_record_batch(): expected_batch = _create_sample_batch() response = _create_read_rows_response(expected_batch) @@ -101,7 +101,7 @@ def test_storage_mvp_direct_pandas_gbq_arrow_contract(): assert actual_batch.equals(expected_batch) -def test_storage_mvp_direct_pandas_gbq_arrow_empty_response(): +def test_from_read_rows_response_handles_empty_response(): response = _create_read_rows_response(record_batch=None) actual_batch = pandas_gbq.arrow.from_read_rows_response( @@ -112,22 +112,27 @@ def test_storage_mvp_direct_pandas_gbq_arrow_empty_response(): assert actual_batch.schema == _TEST_SCHEMA -def test_storage_mvp_read_rows_page_dual_environment_equivalence(): +def test_read_rows_page_to_arrow_delegates_when_pandas_gbq_installed(): expected_batch = _create_sample_batch() response = _create_read_rows_response(expected_batch) + stream_parser = reader._ArrowStreamParser(_create_read_session(_TEST_SCHEMA)) + page = reader.ReadRowsPage(stream_parser, response) - # Environment A: With pandas-gbq installed (Delegated Pathway) - stream_parser_a = reader._ArrowStreamParser(_create_read_session(_TEST_SCHEMA)) - page_a = reader.ReadRowsPage(stream_parser_a, response) with pytest.warns( PendingDeprecationWarning, match="google-cloud-bigquery-storage is deprecated", ): - batch_delegated = page_a.to_arrow() + actual_batch = page.to_arrow() + + assert actual_batch.equals(expected_batch) + + +def test_read_rows_page_to_arrow_falls_back_when_pandas_gbq_uninstalled(): + expected_batch = _create_sample_batch() + response = _create_read_rows_response(expected_batch) + stream_parser = reader._ArrowStreamParser(_create_read_session(_TEST_SCHEMA)) + page = reader.ReadRowsPage(stream_parser, response) - # Environment B: Without pandas-gbq (Legacy Fallback Pathway) - stream_parser_b = reader._ArrowStreamParser(_create_read_session(_TEST_SCHEMA)) - page_b = reader.ReadRowsPage(stream_parser_b, response) with ( mock.patch.dict("sys.modules", {"pandas_gbq": None, "pandas_gbq.arrow": None}), pytest.warns( @@ -135,48 +140,54 @@ def test_storage_mvp_read_rows_page_dual_environment_equivalence(): match="google-cloud-bigquery-storage is deprecated", ), ): - batch_fallback = page_b.to_arrow() + actual_batch = page.to_arrow() - assert batch_delegated.equals(expected_batch) - assert batch_fallback.equals(expected_batch) - assert batch_delegated.equals(batch_fallback) + assert actual_batch.equals(expected_batch) -def test_storage_mvp_read_rows_page_empty_dual_environment_equivalence(): +def test_read_rows_page_to_arrow_empty_batch_delegates_when_pandas_gbq_installed(): empty_batch = _create_sample_batch(rows=[]) response = _create_read_rows_response(record_batch=empty_batch) + stream_parser = reader._ArrowStreamParser(_create_read_session(_TEST_SCHEMA)) + page = reader.ReadRowsPage(stream_parser, response) + + with pytest.warns( + PendingDeprecationWarning, + match="google-cloud-bigquery-storage is deprecated", + ): + actual_batch = page.to_arrow() + + assert actual_batch.num_rows == 0 + assert actual_batch.schema == _TEST_SCHEMA - # Environment A: With pandas-gbq installed - stream_parser_a = reader._ArrowStreamParser(_create_read_session(_TEST_SCHEMA)) - page_a = reader.ReadRowsPage(stream_parser_a, response) - with pytest.warns(PendingDeprecationWarning): - batch_delegated = page_a.to_arrow() - # Environment B: Without pandas-gbq installed - stream_parser_b = reader._ArrowStreamParser(_create_read_session(_TEST_SCHEMA)) - page_b = reader.ReadRowsPage(stream_parser_b, response) +def test_read_rows_page_to_arrow_empty_batch_falls_back_when_pandas_gbq_uninstalled(): + empty_batch = _create_sample_batch(rows=[]) + response = _create_read_rows_response(record_batch=empty_batch) + stream_parser = reader._ArrowStreamParser(_create_read_session(_TEST_SCHEMA)) + page = reader.ReadRowsPage(stream_parser, response) + with ( mock.patch.dict("sys.modules", {"pandas_gbq": None, "pandas_gbq.arrow": None}), - pytest.warns(PendingDeprecationWarning), + pytest.warns( + PendingDeprecationWarning, + match="google-cloud-bigquery-storage is deprecated", + ), ): - batch_fallback = page_b.to_arrow() + actual_batch = page.to_arrow() - assert batch_delegated.num_rows == 0 - assert batch_fallback.num_rows == 0 - assert batch_delegated.schema == _TEST_SCHEMA - assert batch_fallback.schema == _TEST_SCHEMA + assert actual_batch.num_rows == 0 + assert actual_batch.schema == _TEST_SCHEMA -def test_storage_mvp_stream_to_arrow_end_to_end(): +def test_read_rows_stream_to_arrow_concatenates_multiple_batches(): batch1 = _create_sample_batch() batch2 = _create_sample_batch() response1 = _create_read_rows_response(batch1) response2 = _create_read_rows_response(batch2) read_session = _create_read_session(_TEST_SCHEMA) - gapic_client = mock.Mock() gapic_client.read_rows.return_value = iter([response1, response2]) - read_rows_stream = reader.ReadRowsStream( gapic_client, "projects/p/locations/l/sessions/s/streams/str1", @@ -184,7 +195,10 @@ def test_storage_mvp_stream_to_arrow_end_to_end(): {}, ) - with pytest.warns(PendingDeprecationWarning): + with pytest.warns( + PendingDeprecationWarning, + match="google-cloud-bigquery-storage is deprecated", + ): table = read_rows_stream.to_arrow(read_session=read_session) assert table.num_rows == batch1.num_rows + batch2.num_rows From b8cde87095418256a827c74ae42e8b470109c3da Mon Sep 17 00:00:00 2001 From: Shuowei Li Date: Thu, 6 Aug 2026 20:26:14 +0000 Subject: [PATCH 3/9] test(storage): rename MVP validation test to test_reader_pandas_gbq --- ...st_storage_mvp_validation.py => test_reader_pandas_gbq.py} | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) rename packages/google-cloud-bigquery-storage/tests/unit/{test_storage_mvp_validation.py => test_reader_pandas_gbq.py} (100%) diff --git a/packages/google-cloud-bigquery-storage/tests/unit/test_storage_mvp_validation.py b/packages/google-cloud-bigquery-storage/tests/unit/test_reader_pandas_gbq.py similarity index 100% rename from packages/google-cloud-bigquery-storage/tests/unit/test_storage_mvp_validation.py rename to packages/google-cloud-bigquery-storage/tests/unit/test_reader_pandas_gbq.py index 631e392d6c5f..f997038785e0 100644 --- a/packages/google-cloud-bigquery-storage/tests/unit/test_storage_mvp_validation.py +++ b/packages/google-cloud-bigquery-storage/tests/unit/test_reader_pandas_gbq.py @@ -18,10 +18,10 @@ import decimal from unittest import mock -from google.cloud.bigquery_storage import types -from google.cloud.bigquery_storage_v1 import reader import pandas_gbq.arrow import pytest +from google.cloud.bigquery_storage import types +from google.cloud.bigquery_storage_v1 import reader pyarrow = pytest.importorskip("pyarrow") From fae8b1f4d77954a695ecf1f6c01ff2381fa9f367 Mon Sep 17 00:00:00 2001 From: Shuowei Li Date: Thu, 6 Aug 2026 13:54:55 -0700 Subject: [PATCH 4/9] Update packages/google-cloud-bigquery-storage/tests/unit/test_reader_pandas_gbq.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- .../tests/unit/test_reader_pandas_gbq.py | 1 + 1 file changed, 1 insertion(+) diff --git a/packages/google-cloud-bigquery-storage/tests/unit/test_reader_pandas_gbq.py b/packages/google-cloud-bigquery-storage/tests/unit/test_reader_pandas_gbq.py index f997038785e0..7161e74ac9fe 100644 --- a/packages/google-cloud-bigquery-storage/tests/unit/test_reader_pandas_gbq.py +++ b/packages/google-cloud-bigquery-storage/tests/unit/test_reader_pandas_gbq.py @@ -146,6 +146,7 @@ def test_read_rows_page_to_arrow_falls_back_when_pandas_gbq_uninstalled(): def test_read_rows_page_to_arrow_empty_batch_delegates_when_pandas_gbq_installed(): + pytest.importorskip("pandas_gbq") empty_batch = _create_sample_batch(rows=[]) response = _create_read_rows_response(record_batch=empty_batch) stream_parser = reader._ArrowStreamParser(_create_read_session(_TEST_SCHEMA)) From 3112825ad98d1f65cafdf658985ff965cbce6aca Mon Sep 17 00:00:00 2001 From: Shuowei Li Date: Thu, 6 Aug 2026 13:55:03 -0700 Subject: [PATCH 5/9] Update packages/google-cloud-bigquery-storage/tests/unit/test_reader_pandas_gbq.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- .../tests/unit/test_reader_pandas_gbq.py | 1 + 1 file changed, 1 insertion(+) diff --git a/packages/google-cloud-bigquery-storage/tests/unit/test_reader_pandas_gbq.py b/packages/google-cloud-bigquery-storage/tests/unit/test_reader_pandas_gbq.py index 7161e74ac9fe..d3dcce0a9c22 100644 --- a/packages/google-cloud-bigquery-storage/tests/unit/test_reader_pandas_gbq.py +++ b/packages/google-cloud-bigquery-storage/tests/unit/test_reader_pandas_gbq.py @@ -113,6 +113,7 @@ def test_from_read_rows_response_handles_empty_response(): def test_read_rows_page_to_arrow_delegates_when_pandas_gbq_installed(): + pytest.importorskip("pandas_gbq") expected_batch = _create_sample_batch() response = _create_read_rows_response(expected_batch) stream_parser = reader._ArrowStreamParser(_create_read_session(_TEST_SCHEMA)) From a4f27ea644a5bddd54766aff82662e13921cf7af Mon Sep 17 00:00:00 2001 From: Shuowei Li Date: Thu, 6 Aug 2026 13:55:09 -0700 Subject: [PATCH 6/9] Update packages/google-cloud-bigquery-storage/tests/unit/test_reader_pandas_gbq.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- .../tests/unit/test_reader_pandas_gbq.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/packages/google-cloud-bigquery-storage/tests/unit/test_reader_pandas_gbq.py b/packages/google-cloud-bigquery-storage/tests/unit/test_reader_pandas_gbq.py index d3dcce0a9c22..ae39410694c0 100644 --- a/packages/google-cloud-bigquery-storage/tests/unit/test_reader_pandas_gbq.py +++ b/packages/google-cloud-bigquery-storage/tests/unit/test_reader_pandas_gbq.py @@ -102,9 +102,10 @@ def test_from_read_rows_response_decodes_serialized_record_batch(): def test_from_read_rows_response_handles_empty_response(): + pandas_gbq_arrow = pytest.importorskip("pandas_gbq.arrow") response = _create_read_rows_response(record_batch=None) - actual_batch = pandas_gbq.arrow.from_read_rows_response( + actual_batch = pandas_gbq_arrow.from_read_rows_response( response, arrow_schema=_TEST_SCHEMA ) From 7787fdc7ced97201362a4d51fdbb4e8deb9eb8d8 Mon Sep 17 00:00:00 2001 From: Shuowei Li Date: Thu, 6 Aug 2026 13:55:15 -0700 Subject: [PATCH 7/9] Update packages/google-cloud-bigquery-storage/tests/unit/test_reader_pandas_gbq.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- .../tests/unit/test_reader_pandas_gbq.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/packages/google-cloud-bigquery-storage/tests/unit/test_reader_pandas_gbq.py b/packages/google-cloud-bigquery-storage/tests/unit/test_reader_pandas_gbq.py index ae39410694c0..663d6dc4875f 100644 --- a/packages/google-cloud-bigquery-storage/tests/unit/test_reader_pandas_gbq.py +++ b/packages/google-cloud-bigquery-storage/tests/unit/test_reader_pandas_gbq.py @@ -91,10 +91,11 @@ def _create_read_session(schema=_TEST_SCHEMA): def test_from_read_rows_response_decodes_serialized_record_batch(): + pandas_gbq_arrow = pytest.importorskip("pandas_gbq.arrow") expected_batch = _create_sample_batch() response = _create_read_rows_response(expected_batch) - actual_batch = pandas_gbq.arrow.from_read_rows_response( + actual_batch = pandas_gbq_arrow.from_read_rows_response( response, arrow_schema=_TEST_SCHEMA ) From 0322c902cd7c7e6ef1e9e45bc416a6d9970c0459 Mon Sep 17 00:00:00 2001 From: Shuowei Li Date: Thu, 6 Aug 2026 21:04:23 +0000 Subject: [PATCH 8/9] test(storage): import pandas_gbq dynamically in unit tests --- .../tests/unit/test_reader_pandas_gbq.py | 1 - 1 file changed, 1 deletion(-) diff --git a/packages/google-cloud-bigquery-storage/tests/unit/test_reader_pandas_gbq.py b/packages/google-cloud-bigquery-storage/tests/unit/test_reader_pandas_gbq.py index 663d6dc4875f..f209efb84e2e 100644 --- a/packages/google-cloud-bigquery-storage/tests/unit/test_reader_pandas_gbq.py +++ b/packages/google-cloud-bigquery-storage/tests/unit/test_reader_pandas_gbq.py @@ -18,7 +18,6 @@ import decimal from unittest import mock -import pandas_gbq.arrow import pytest from google.cloud.bigquery_storage import types from google.cloud.bigquery_storage_v1 import reader From 234c597e0d0b5412db6b9126bcce01a4c215aa7c Mon Sep 17 00:00:00 2001 From: Shuowei Li Date: Thu, 6 Aug 2026 21:14:58 +0000 Subject: [PATCH 9/9] style: remove extra lines --- .../tests/unit/test_reader_pandas_gbq.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/packages/google-cloud-bigquery-storage/tests/unit/test_reader_pandas_gbq.py b/packages/google-cloud-bigquery-storage/tests/unit/test_reader_pandas_gbq.py index f209efb84e2e..344c3218c6a9 100644 --- a/packages/google-cloud-bigquery-storage/tests/unit/test_reader_pandas_gbq.py +++ b/packages/google-cloud-bigquery-storage/tests/unit/test_reader_pandas_gbq.py @@ -1,5 +1,3 @@ -# -*- coding: utf-8 -*- -# # Copyright 2026 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License");