diff --git a/bigquery/bigframes/query_dry_run.py b/bigquery/bigframes/query_dry_run.py new file mode 100644 index 0000000000..ee1c1ab834 --- /dev/null +++ b/bigquery/bigframes/query_dry_run.py @@ -0,0 +1,60 @@ +# 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 +# +# http://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 bigquery_bigframes_query_dry_run] +from typing import Tuple + +import bigframes.pandas as bpd +import pandas as pd + +# Set partial ordering mode as the default configuration for BigQuery +# DataFrames. +bpd.options.bigquery.ordering_mode = "partial" + + +def query_dry_run( + project_id: str = "your-project-id", +) -> Tuple[pd.Series, pd.Series]: + """Estimates the bytes processed by a SQL query and a DataFrame operation.""" + bpd.options.bigquery.project = project_id + + sql = """ + SELECT name, SUM(number) AS total_people + FROM `bigquery-public-data.usa_names.usa_1910_current` + WHERE state = 'TX' + GROUP BY name + """ + + # Perform a dry run on a SQL query without executing it. + sql_stats = bpd.read_gbq(sql, use_cache=False, dry_run=True) + print(f"SQL query will process {sql_stats['totalBytesProcessed']} bytes.") + + # Perform a dry run on BigQuery DataFrames operations before downloading + # results to an in-memory pandas DataFrame. + df = bpd.read_gbq("bigquery-public-data.usa_names.usa_1910_current") + df_tx = df[df["state"] == "TX"].groupby("name", as_index=False)["number"].sum() + df_stats = df_tx.to_pandas(dry_run=True) + print(f"DataFrame operation will process {df_stats['totalBytesProcessed']} bytes.") + + return sql_stats, df_stats + + +# [END bigquery_bigframes_query_dry_run] + + +if __name__ == "__main__": + import os + + project = os.environ.get("GOOGLE_CLOUD_PROJECT", "your-project-id") + query_dry_run(project_id=project) diff --git a/bigquery/bigframes/query_dry_run_test.py b/bigquery/bigframes/query_dry_run_test.py new file mode 100644 index 0000000000..0755b678fb --- /dev/null +++ b/bigquery/bigframes/query_dry_run_test.py @@ -0,0 +1,29 @@ +# 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 +# +# http://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 bigframes.pandas as bpd +import pytest + +import query_dry_run + + +def test_query_dry_run(project_id: str, capsys: pytest.CaptureFixture) -> None: + bpd.close_session() + + sql_stats, df_stats = query_dry_run.query_dry_run(project_id=project_id) + out, _ = capsys.readouterr() + + assert sql_stats["totalBytesProcessed"] > 0 + assert df_stats["totalBytesProcessed"] > 0 + assert "bytes" in out diff --git a/bigquery/pandas-gbq-migration/query_dry_run.py b/bigquery/pandas-gbq-migration/query_dry_run.py new file mode 100644 index 0000000000..15889f52e6 --- /dev/null +++ b/bigquery/pandas-gbq-migration/query_dry_run.py @@ -0,0 +1,52 @@ +# 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 +# +# http://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 bigquery_migration_pandas_gbq_query_dry_run] +import pandas as pd +import pandas_gbq + + +def query_dry_run(project_id: str = "your-project-id") -> pd.Series: + """Estimates the bytes processed by a SQL query using a dry run.""" + sql = """ + SELECT name, SUM(number) AS total_people + FROM `bigquery-public-data.usa_names.usa_1910_current` + WHERE state = 'TX' + GROUP BY name + """ + + # Perform a dry run to estimate bytes processed without running the query. + query_config = { + "query": { + "useQueryCache": False, + } + } + stats = pandas_gbq.read_gbq( + sql, + project_id=project_id, + configuration=query_config, + dry_run=True, + ) + print(f"This query will process {stats['totalBytesProcessed']} bytes.") + return stats + + +# [END bigquery_migration_pandas_gbq_query_dry_run] + + +if __name__ == "__main__": + import os + + project = os.environ.get("GOOGLE_CLOUD_PROJECT", "your-project-id") + query_dry_run(project_id=project) diff --git a/bigquery/pandas-gbq-migration/query_dry_run_test.py b/bigquery/pandas-gbq-migration/query_dry_run_test.py new file mode 100644 index 0000000000..d06ceada95 --- /dev/null +++ b/bigquery/pandas-gbq-migration/query_dry_run_test.py @@ -0,0 +1,32 @@ +# 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 +# +# http://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 os + +import pytest + +import query_dry_run + + +@pytest.fixture(scope="session") +def project_id() -> str: + return os.environ["GOOGLE_CLOUD_PROJECT"] + + +def test_query_dry_run(project_id: str, capsys: pytest.CaptureFixture) -> None: + stats = query_dry_run.query_dry_run(project_id=project_id) + out, _ = capsys.readouterr() + + assert stats["totalBytesProcessed"] > 0 + assert "bytes" in out diff --git a/bigquery/pandas-gbq-migration/requirements.txt b/bigquery/pandas-gbq-migration/requirements.txt index cec3110949..dd6bdabfdb 100644 --- a/bigquery/pandas-gbq-migration/requirements.txt +++ b/bigquery/pandas-gbq-migration/requirements.txt @@ -1,8 +1,10 @@ google-cloud-bigquery==3.27.0 google-cloud-bigquery-storage==2.27.0 pandas==2.0.3; python_version == '3.8' -pandas==2.2.3; python_version > '3.8' -pandas-gbq==0.24.0 +pandas==2.2.3; python_version < '3.11' +pandas==3.1.0; python_version >= '3.11' +pandas-gbq==0.35.3 grpcio==1.70.0; python_version == '3.8' -grpcio==1.74.0; python_version > '3.8' +grpcio==1.74.0; python_version == '3.9' +grpcio==1.85.1; python_version >= '3.10' pyarrow==24.0.0 diff --git a/bigquery/polars/README.md b/bigquery/polars/README.md new file mode 100644 index 0000000000..14de2f0ebc --- /dev/null +++ b/bigquery/polars/README.md @@ -0,0 +1,4 @@ +# polars-bigquery samples + +This directory contains samples used to demonstrate use of the `polars-bigquery` +package. diff --git a/bigquery/polars/conftest.py b/bigquery/polars/conftest.py new file mode 100644 index 0000000000..731ebc3e30 --- /dev/null +++ b/bigquery/polars/conftest.py @@ -0,0 +1,22 @@ +# Copyright 2021 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 os + +import pytest + + +@pytest.fixture(scope="session") +def project_id() -> str: + return os.environ["GOOGLE_CLOUD_PROJECT"] diff --git a/bigquery/polars/noxfile_config.py b/bigquery/polars/noxfile_config.py new file mode 100644 index 0000000000..a4cb0cac19 --- /dev/null +++ b/bigquery/polars/noxfile_config.py @@ -0,0 +1,39 @@ +# Copyright 2021 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 +# +# http://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. + +# Default TEST_CONFIG_OVERRIDE for python repos. + +# You can copy this file into your directory, then it will be imported from +# the noxfile.py. + +# The source of truth: +# https://github.com/GoogleCloudPlatform/python-docs-samples/blob/main/noxfile_config.py + +TEST_CONFIG_OVERRIDE = { + # You can opt out from the test for specific Python versions. + # Skipping for Python 3.9 due to pyarrow compilation failure. + "ignored_versions": ["3.8", "3.9", "3.11", "3.12", "3.13"], + # Old samples are opted out of enforcing Python type hints + # All new samples should feature them + "enforce_type_hints": False, + # An envvar key for determining the project id to use. Change it + # to 'BUILD_SPECIFIC_GCLOUD_PROJECT' if you want to opt in using a + # build specific Cloud project. You can also use your own string + # to use your own Cloud project. + "gcloud_project_env": "GOOGLE_CLOUD_PROJECT", + # 'gcloud_project_env': 'BUILD_SPECIFIC_GCLOUD_PROJECT', + # A dictionary you want to inject into your test. Don't put any + # secrets here. These values will override predefined values. + "envs": {}, +} diff --git a/bigquery/polars/query_bqstorage.py b/bigquery/polars/query_bqstorage.py new file mode 100644 index 0000000000..626b5fa687 --- /dev/null +++ b/bigquery/polars/query_bqstorage.py @@ -0,0 +1,51 @@ +# 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 +# +# http://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 bigquery_polars_query_bqstorage] +import polars +import polars_bigquery + + +def construct_client(project_id: str = "your-project-id") -> polars_bigquery.Client: + return polars_bigquery.Client( + quota_project_id=project_id, + ) + + +def query_bqstorage(plbq: polars_bigquery.Client) -> polars.DataFrame: + sql = """ + SELECT year, name, SUM(number) + FROM `bigquery-public-data.usa_names.usa_1910_current` + GROUP BY year, name + """ + + # Run a query and immediately download the results. polars-biqquery + # automatically uses the BigQuery Storage Read API to download results that + # don't fit in a single response message. + df = plbq.read_query( + sql, + # For the best performance, set maintain_order=False if you don't care + # about the order of the rows in the result. This allows polars-bigquery + # to create multiple BigQuery Storage Read API read streams if needed. + maintain_order=False, + ) + return df +# [END bigquery_polars_query_bqstorage] + + +if __name__ == "__main__": + import os + + project = os.environ.get("GOOGLE_CLOUD_PROJECT", "your-project-id") + print(query_bqstorage(construct_client(project))) diff --git a/bigquery/polars/query_bqstorage_test.py b/bigquery/polars/query_bqstorage_test.py new file mode 100644 index 0000000000..3d884d125c --- /dev/null +++ b/bigquery/polars/query_bqstorage_test.py @@ -0,0 +1,21 @@ +# 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 +# +# http://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 query_bqstorage + + +def test_query_bqstorage(project_id: str) -> None: + plbq = query_bqstorage.construct_client(project_id=project_id) + pandas_df = query_bqstorage.query_bqstorage(plbq) + assert pandas_df is not None diff --git a/bigquery/polars/query_dry_run.py b/bigquery/polars/query_dry_run.py new file mode 100644 index 0000000000..ebf1e1709f --- /dev/null +++ b/bigquery/polars/query_dry_run.py @@ -0,0 +1,57 @@ +# 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 +# +# http://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 bigquery_polars_query_dry_run] +import polars +import polars_bigquery + + +def construct_client( + project_id: str = "your-project-id", +) -> polars_bigquery.Client: + """Constructs a Polars BigQuery client.""" + return polars_bigquery.Client( + quota_project_id=project_id, + ) + + +def query_dry_run(plbq: polars_bigquery.Client) -> polars.DataFrame: + """Estimates the bytes processed by a SQL query using a dry run.""" + sql = """ + SELECT name, SUM(number) AS total_people + FROM `bigquery-public-data.usa_names.usa_1910_current` + WHERE state = 'TX' + GROUP BY name + """ + + # Perform a dry run to validate the query and estimate bytes processed + # without executing the query. + stats = plbq.read_query( + sql, + dry_run=True, + use_query_cache=False, + ) + total_bytes = stats.item(0, "totalBytesProcessed") + print(f"This query will process {total_bytes} bytes.") + return stats + + +# [END bigquery_polars_query_dry_run] + + +if __name__ == "__main__": + import os + + project = os.environ.get("GOOGLE_CLOUD_PROJECT", "your-project-id") + query_dry_run(construct_client(project_id=project)) diff --git a/bigquery/polars/query_dry_run_test.py b/bigquery/polars/query_dry_run_test.py new file mode 100644 index 0000000000..26767993ba --- /dev/null +++ b/bigquery/polars/query_dry_run_test.py @@ -0,0 +1,27 @@ +# 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 +# +# http://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 pytest + +import query_dry_run + + +def test_query_dry_run(project_id: str, capsys: pytest.CaptureFixture) -> None: + plbq = query_dry_run.construct_client(project_id=project_id) + + stats = query_dry_run.query_dry_run(plbq) + out, _ = capsys.readouterr() + + assert stats.item(0, "totalBytesProcessed") > 0 + assert "bytes" in out diff --git a/bigquery/polars/query_standard_sql.py b/bigquery/polars/query_standard_sql.py new file mode 100644 index 0000000000..07c54e0fc2 --- /dev/null +++ b/bigquery/polars/query_standard_sql.py @@ -0,0 +1,51 @@ +# 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 +# +# http://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 bigquery_polars_query] +import polars +import polars_bigquery + + +def construct_client(project_id: str = "your-project-id") -> polars_bigquery.Client: + return polars_bigquery.Client( + quota_project_id=project_id, + ) + + +def query_standard_sql(plbq: polars_bigquery.Client) -> polars.DataFrame: + sql = """ + SELECT name, sum(number) as total_people + FROM `bigquery-public-data.usa_names.usa_1910_current` + WHERE state = 'TX' + GROUP BY name + ORDER BY total_people DESC + LIMIT 100 + """ + + # Run a query and immediately download the results. + df = plbq.read_query( + sql, + # For the best performance, set maintain_order=False if you don't care + # about the order of the rows in the result. + maintain_order=True, + ) + return df +# [END bigquery_polars_query] + + +if __name__ == "__main__": + import os + + project = os.environ.get("GOOGLE_CLOUD_PROJECT", "your-project-id") + print(query_standard_sql(construct_client(project_id=project))) diff --git a/bigquery/polars/query_standard_sql_test.py b/bigquery/polars/query_standard_sql_test.py new file mode 100644 index 0000000000..86e87cf589 --- /dev/null +++ b/bigquery/polars/query_standard_sql_test.py @@ -0,0 +1,21 @@ +# 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 +# +# http://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 query_standard_sql + + +def test_query_standard_sql(project_id: str) -> None: + plbq = query_standard_sql.construct_client(project_id=project_id) + df = query_standard_sql.query_standard_sql(plbq) + assert df is not None diff --git a/bigquery/polars/requirements-test.txt b/bigquery/polars/requirements-test.txt new file mode 100644 index 0000000000..c9e154ba44 --- /dev/null +++ b/bigquery/polars/requirements-test.txt @@ -0,0 +1 @@ +pytest==9.0.3; python_version >= "3.10" diff --git a/bigquery/polars/requirements.txt b/bigquery/polars/requirements.txt new file mode 100644 index 0000000000..829e201beb --- /dev/null +++ b/bigquery/polars/requirements.txt @@ -0,0 +1,2 @@ +polars-bigquery==0.3.0 +polars==2.1.0