feat: prepare client-side delegation and telemetry for pandas-gbq

packages/google-cloud-bigquery/.coveragerc

@@ -12,3 +12,5 @@ exclude_lines =
     pragma: (no cover|NO COVER)
     # Ignore debug-only repr
     def __repr__
+    except ImportError:
+    except ImportError as .*:

packages/google-cloud-bigquery/google/cloud/bigquery/_versions_helpers.py

@@ -16,10 +16,8 @@
 from typing import Any
 
 import packaging.version
-
 from google.cloud.bigquery import exceptions
 
-
 _MIN_PYARROW_VERSION = packaging.version.Version("3.0.0")
 _MIN_BQ_STORAGE_VERSION = packaging.version.Version("2.0.0")
 _BQ_STORAGE_OPTIONAL_READ_SESSION_VERSION = packaging.version.Version("2.6.0")
@@ -247,3 +245,49 @@ def try_import(self, raise_if_error: bool = False) -> Any:
     and PYARROW_VERSIONS.try_import() is not None
     and PYARROW_VERSIONS.installed_version >= _MIN_PYARROW_VERSION_RANGE
 )
+
+
+class PandasGBQVersions:
+    """Version and delegation comparisons for pandas-gbq package."""
+
+    def __init__(self):
+        self._installed_version = None
+        self._delegation_api_version = None
+
+    @property
+    def installed_version(self) -> packaging.version.Version:
+        """Return the parsed version of pandas-gbq"""
+        if self._installed_version is None:
+            try:
+                import pandas_gbq  # type: ignore
+
+                self._installed_version = packaging.version.parse(
+                    getattr(pandas_gbq, "__version__", "0.0.0")
+                )
+            except Exception:
+                self._installed_version = packaging.version.parse("0.0.0")
+
+        return self._installed_version
+
+    @property
+    def delegation_api_version(self) -> int:
+        """Return the delegation API version of pandas-gbq if installed, otherwise 0."""
+        if self._delegation_api_version is None:
+            try:
+                import pandas_gbq  # type: ignore
+
+                self._delegation_api_version = getattr(
+                    pandas_gbq, "_internal_delegation_api_version", 0
+                )
+            except Exception:
+                self._delegation_api_version = 0
+
+        return self._delegation_api_version
+
+    @property
+    def is_delegation_supported(self) -> bool:
+        """True if the installed pandas-gbq version supports query delegation API (version >= 1)."""
+        return self.delegation_api_version >= 1
+
+
+PANDAS_GBQ_VERSIONS = PandasGBQVersions()

packages/google-cloud-bigquery/google/cloud/bigquery/magics/magics.py

@@ -23,10 +23,10 @@
 
 from __future__ import print_function
 
-import re
 import ast
 import copy
 import functools
+import re
 import sys
 import time
 import warnings
@@ -39,14 +39,12 @@
 except ImportError:
     raise ImportError("This module can only be loaded in IPython.")
 
-from google.api_core import client_info
-from google.api_core import client_options
-from google.api_core.exceptions import NotFound
 import google.auth  # type: ignore
-from google.cloud import bigquery
 import google.cloud.bigquery.dataset
-from google.cloud.bigquery import _versions_helpers
-from google.cloud.bigquery import exceptions
+from google.api_core import client_info, client_options
+from google.api_core.exceptions import NotFound
+from google.cloud import bigquery
+from google.cloud.bigquery import _versions_helpers, exceptions
 from google.cloud.bigquery.dbapi import _helpers
 from google.cloud.bigquery.magics import line_arg_parser as lap
 
@@ -231,7 +229,7 @@ def progress_bar_type(self, value):
 # their code.
 if bigquery_magics is not None:
     context = bigquery_magics.context
-else:
+else:  # pragma: NO COVER
     context = Context()
 
 

packages/google-cloud-bigquery/google/cloud/bigquery/table.py

@@ -21,9 +21,8 @@
 import functools
 import operator
 import typing
-from typing import Any, Dict, Iterable, Iterator, List, Optional, Tuple, Union, Sequence
-
 import warnings
+from typing import Any, Dict, Iterable, Iterator, List, Optional, Sequence, Tuple, Union
 
 try:
     import pandas  # type: ignore
@@ -56,30 +55,33 @@
     _read_wkt = wkt.loads
 
 import google.api_core.exceptions
-from google.api_core.page_iterator import HTTPIterator
-
 import google.cloud._helpers  # type: ignore
-from google.cloud.bigquery import _helpers
-from google.cloud.bigquery import _pandas_helpers
-from google.cloud.bigquery import _versions_helpers
+from google.api_core.page_iterator import HTTPIterator
+from google.cloud.bigquery import (
+    _helpers,
+    _pandas_helpers,
+    _string_references,
+    _versions_helpers,
+    external_config,
+)
 from google.cloud.bigquery import exceptions as bq_exceptions
+from google.cloud.bigquery import schema as _schema
 from google.cloud.bigquery._tqdm_helpers import get_progress_bar
 from google.cloud.bigquery.encryption_configuration import EncryptionConfiguration
 from google.cloud.bigquery.enums import DefaultPandasDTypes
 from google.cloud.bigquery.external_config import ExternalConfig
-from google.cloud.bigquery import schema as _schema
-from google.cloud.bigquery.schema import _build_schema_resource
-from google.cloud.bigquery.schema import _parse_schema_resource
-from google.cloud.bigquery.schema import _to_schema_fields
-from google.cloud.bigquery import external_config
-from google.cloud.bigquery import _string_references
+from google.cloud.bigquery.schema import (
+    _build_schema_resource,
+    _parse_schema_resource,
+    _to_schema_fields,
+)
 
 if typing.TYPE_CHECKING:  # pragma: NO COVER
     # Unconditionally import optional dependencies again to tell pytype that
     # they are not None, avoiding false "no attribute" errors.
+    import geopandas  # type: ignore
     import pandas
     import pyarrow
-    import geopandas  # type: ignore
     from google.cloud import bigquery_storage  # type: ignore
     from google.cloud.bigquery.dataset import DatasetReference
 
@@ -2709,6 +2711,14 @@ def to_dataframe(
                 is not supported dtype.
 
         """
+        if not _versions_helpers.PANDAS_GBQ_VERSIONS.is_delegation_supported:
+            warnings.warn(
+                "Retrieving dataframes via the core client is deprecated. "
+                "Please install 'pandas-gbq' for the new high-performance backend.",
+                PendingDeprecationWarning,
+                stacklevel=2,
+            )
+
         _pandas_helpers.verify_pandas_imports()
 
         if geography_as_object and shapely is None:
@@ -2801,6 +2811,44 @@ def to_dataframe(
             create_bqstorage_client = False
             bqstorage_client = None
 
+        if _versions_helpers.PANDAS_GBQ_VERSIONS.is_delegation_supported:
+            import pandas_gbq  # type: ignore
+
+            if (
+                self.client
+                and hasattr(self.client, "_connection")
+                and hasattr(self.client._connection, "_client_info")
+            ):
+                client_info = self.client._connection._client_info
+                if client_info:
+                    ua = client_info.user_agent or ""
+                    if "pandas-gbq" not in ua:
+                        pandas_gbq_version = getattr(pandas_gbq, "__version__", "0.0.0")
+                        client_info.user_agent = (
+                            f"{ua} pandas-gbq/{pandas_gbq_version}".strip()
+                        )
+
+            return pandas_gbq.pandas.from_row_iterator(
+                self,
+                bqstorage_client=bqstorage_client,
+                dtypes=dtypes,
+                progress_bar_type=progress_bar_type,
+                create_bqstorage_client=create_bqstorage_client,
+                geography_as_object=geography_as_object,
+                bool_dtype=bool_dtype,
+                int_dtype=int_dtype,
+                float_dtype=float_dtype,
+                string_dtype=string_dtype,
+                date_dtype=date_dtype,
+                datetime_dtype=datetime_dtype,
+                time_dtype=time_dtype,
+                timestamp_dtype=timestamp_dtype,
+                range_date_dtype=range_date_dtype,
+                range_datetime_dtype=range_datetime_dtype,
+                range_timestamp_dtype=range_timestamp_dtype,
+                timeout=timeout,
+            )
+
         record_batch = self.to_arrow(
             progress_bar_type=progress_bar_type,
             bqstorage_client=bqstorage_client,

packages/google-cloud-bigquery/tests/unit/test__pandas_helpers.py

@@ -18,14 +18,13 @@
 import decimal
 import functools
 import gc
+import importlib.metadata as metadata
 import operator
 import queue
 import time
+import warnings
 from typing import Union
 from unittest import mock
-import warnings
-
-import importlib.metadata as metadata
 
 try:
     import pandas
@@ -45,13 +44,13 @@
     geopandas = None
 
 import pytest
-
 from google import api_core
-
-from google.cloud.bigquery import exceptions
-from google.cloud.bigquery import _pyarrow_helpers
-from google.cloud.bigquery import _versions_helpers
-from google.cloud.bigquery import schema
+from google.cloud.bigquery import (
+    _pyarrow_helpers,
+    _versions_helpers,
+    exceptions,
+    schema,
+)
 from google.cloud.bigquery._pandas_helpers import determine_requested_streams
 
 pyarrow = _versions_helpers.PYARROW_VERSIONS.try_import()
@@ -1831,8 +1830,7 @@ def test__download_table_bqstorage(
     expected_call_count,
     expected_maxsize,
 ):
-    from google.cloud.bigquery import dataset
-    from google.cloud.bigquery import table
+    from google.cloud.bigquery import dataset, table
 
     queue_used = None  # A reference to the queue used by code under test.
 
@@ -1885,10 +1883,9 @@ def test__download_table_bqstorage_shuts_down_workers(
     the child threads are also stopped.
     """
     pytest.importorskip("google.cloud.bigquery_storage_v1")
-    from google.cloud.bigquery import dataset
-    from google.cloud.bigquery import table
     import google.cloud.bigquery_storage_v1.reader
     import google.cloud.bigquery_storage_v1.types
+    from google.cloud.bigquery import dataset, table
 
     monkeypatch.setattr(
         _versions_helpers.BQ_STORAGE_VERSIONS, "_installed_version", None
@@ -2211,10 +2208,10 @@ def test_determine_requested_streams_invalid_max_stream_count():
     bigquery_storage is None, reason="Requires google-cloud-bigquery-storage"
 )
 def test__download_table_bqstorage_w_timeout_error(module_under_test):
-    from google.cloud.bigquery import dataset
-    from google.cloud.bigquery import table
     from unittest import mock
 
+    from google.cloud.bigquery import dataset, table
+
     mock_bqstorage_client = mock.create_autospec(
         bigquery_storage.BigQueryReadClient, instance=True
     )
@@ -2248,10 +2245,10 @@ def slow_download_stream(
     bigquery_storage is None, reason="Requires google-cloud-bigquery-storage"
 )
 def test__download_table_bqstorage_w_timeout_success(module_under_test):
-    from google.cloud.bigquery import dataset
-    from google.cloud.bigquery import table
     from unittest import mock
 
+    from google.cloud.bigquery import dataset, table
+
     mock_bqstorage_client = mock.create_autospec(
         bigquery_storage.BigQueryReadClient, instance=True
     )
@@ -2409,3 +2406,38 @@ def test_download_arrow_bqstorage_passes_timeout_to_create_read_session(
     assert retry_policy is not None
     # Check if deadline is set correctly in the retry policy
     assert retry_policy._deadline == timeout
+
+
+@pytest.mark.skipif(pandas is None, reason="Requires `pandas`")
+def test_dataframe_to_bq_schema_w_unused_schema_field(module_under_test):
+    with mock.patch.object(module_under_test, "pandas_gbq", None):
+        with pytest.raises(
+            ValueError, match="bq_schema contains fields not present in dataframe"
+        ):
+            module_under_test.dataframe_to_bq_schema(
+                pandas.DataFrame(), (schema.SchemaField("not_in_df", "STRING"),)
+            )
+
+
+@pytest.mark.skipif(pandas is None, reason="Requires `pandas`")
+@pytest.mark.skipif(isinstance(pyarrow, mock.Mock), reason="Requires `pyarrow`")
+def test_get_schema_by_pyarrow_bignumeric(module_under_test):
+    series = pandas.Series([decimal.Decimal("1.12345678901")])
+    result = module_under_test._get_schema_by_pyarrow("col", series)
+    assert result is not None
+    assert result.field_type == "BIGNUMERIC"
+
+
+@pytest.mark.skipif(pandas is None, reason="Requires `pandas`")
+@pytest.mark.skipif(isinstance(pyarrow, mock.Mock), reason="Requires `pyarrow`")
+def test_get_types_mapper_range_timestamp_mismatch(module_under_test):
+    if not hasattr(pandas, "ArrowDtype"):
+        return
+    range_ts = pandas.ArrowDtype(
+        pyarrow.struct(
+            [("start", pyarrow.timestamp("us")), ("end", pyarrow.timestamp("us"))]
+        )
+    )
+    mapper = module_under_test.default_types_mapper(range_timestamp_dtype=range_ts)
+    unmatched_struct = pyarrow.struct([("other", pyarrow.int64())])
+    assert mapper(unmatched_struct) is None

packages/google-cloud-bigquery/tests/unit/test__pyarrow_helpers.py

@@ -44,3 +44,16 @@ def test_bq_to_arrow_scalars(module_under_test):
 def test_arrow_scalar_ids_to_bq(module_under_test):
     assert module_under_test.arrow_scalar_ids_to_bq(pyarrow.bool_().id) == "BOOL"
     assert module_under_test.arrow_scalar_ids_to_bq("UNKNOWN_TYPE") is None
+
+
+def test_pyarrow_helpers_when_pyarrow_none(module_under_test):
+    import importlib
+    import sys
+    from unittest import mock
+
+    with mock.patch.dict(sys.modules, {"pyarrow": None}):
+        importlib.reload(module_under_test)
+        assert module_under_test.pyarrow is None
+        assert module_under_test.arrow_scalar_ids_to_bq(1) is None
+
+    importlib.reload(module_under_test)