google-cloud-bigquery 3.31.0__py3-none-any.whl

google/cloud/bigquery/_pyarrow_helpers.py

@@ -0,0 +1,147 @@
+# Copyright 2023 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.
+
+"""Shared helper functions for connecting BigQuery and pyarrow.
+
+NOTE: This module is DEPRECATED. Please make updates in the pandas-gbq package,
+instead. See: go/pandas-gbq-and-bigframes-redundancy,
+https://github.com/googleapis/python-bigquery-pandas/blob/main/pandas_gbq/schema/bigquery_to_pyarrow.py
+and
+https://github.com/googleapis/python-bigquery-pandas/blob/main/pandas_gbq/schema/pyarrow_to_bigquery.py
+"""
+
+from typing import Any
+
+try:
+    import pyarrow  # type: ignore
+except ImportError:
+    pyarrow = None
+
+try:
+    import db_dtypes  # type: ignore
+
+    db_dtypes_import_exception = None
+except ImportError as exc:
+    db_dtypes = None
+    db_dtypes_import_exception = exc
+
+
+def pyarrow_datetime():
+    return pyarrow.timestamp("us", tz=None)
+
+
+def pyarrow_numeric():
+    return pyarrow.decimal128(38, 9)
+
+
+def pyarrow_bignumeric():
+    # 77th digit is partial.
+    # https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types#decimal_types
+    return pyarrow.decimal256(76, 38)
+
+
+def pyarrow_time():
+    return pyarrow.time64("us")
+
+
+def pyarrow_timestamp():
+    return pyarrow.timestamp("us", tz="UTC")
+
+
+_BQ_TO_ARROW_SCALARS = {}
+_ARROW_SCALAR_IDS_TO_BQ = {}
+
+if pyarrow:
+    # This dictionary is duplicated in bigquery_storage/test/unite/test_reader.py
+    # When modifying it be sure to update it there as well.
+    # Note(todo!!): type "BIGNUMERIC"'s matching pyarrow type is added in _pandas_helpers.py
+    _BQ_TO_ARROW_SCALARS = {
+        "BOOL": pyarrow.bool_,
+        "BOOLEAN": pyarrow.bool_,
+        "BYTES": pyarrow.binary,
+        "DATE": pyarrow.date32,
+        "DATETIME": pyarrow_datetime,
+        "FLOAT": pyarrow.float64,
+        "FLOAT64": pyarrow.float64,
+        "GEOGRAPHY": pyarrow.string,
+        "INT64": pyarrow.int64,
+        "INTEGER": pyarrow.int64,
+        # Normally, we'd prefer JSON type built-in to pyarrow (added in 19.0.0),
+        # but we'd like this to map as closely to the BQ Storage API as
+        # possible, which uses the string() dtype, as JSON support in Arrow
+        # predates JSON support in BigQuery by several years.
+        "JSON": pyarrow.string,
+        "NUMERIC": pyarrow_numeric,
+        "STRING": pyarrow.string,
+        "TIME": pyarrow_time,
+        "TIMESTAMP": pyarrow_timestamp,
+    }
+
+    # DEPRECATED: update pandas_gbq.schema.pyarrow_to_bigquery, instead.
+    _ARROW_SCALAR_IDS_TO_BQ = {
+        # https://arrow.apache.org/docs/python/api/datatypes.html#type-classes
+        pyarrow.bool_().id: "BOOL",
+        pyarrow.int8().id: "INT64",
+        pyarrow.int16().id: "INT64",
+        pyarrow.int32().id: "INT64",
+        pyarrow.int64().id: "INT64",
+        pyarrow.uint8().id: "INT64",
+        pyarrow.uint16().id: "INT64",
+        pyarrow.uint32().id: "INT64",
+        pyarrow.uint64().id: "INT64",
+        pyarrow.float16().id: "FLOAT64",
+        pyarrow.float32().id: "FLOAT64",
+        pyarrow.float64().id: "FLOAT64",
+        pyarrow.time32("ms").id: "TIME",
+        pyarrow.time64("ns").id: "TIME",
+        pyarrow.timestamp("ns").id: "TIMESTAMP",
+        pyarrow.date32().id: "DATE",
+        pyarrow.date64().id: "DATETIME",  # because millisecond resolution
+        pyarrow.binary().id: "BYTES",
+        pyarrow.string().id: "STRING",  # also alias for pyarrow.utf8()
+        pyarrow.large_string().id: "STRING",
+        # The exact scale and precision don't matter, see below.
+        pyarrow.decimal128(38, scale=9).id: "NUMERIC",
+        # NOTE: all extension types (e.g. json_, uuid, db_dtypes.JSONArrowType)
+        # have the same id (31 as of version 19.0.1), so these should not be
+        # matched by id.
+    }
+
+    _BQ_TO_ARROW_SCALARS["BIGNUMERIC"] = pyarrow_bignumeric
+    # The exact decimal's scale and precision are not important, as only
+    # the type ID matters, and it's the same for all decimal256 instances.
+    _ARROW_SCALAR_IDS_TO_BQ[pyarrow.decimal256(76, scale=38).id] = "BIGNUMERIC"
+
+
+def bq_to_arrow_scalars(bq_scalar: str):
+    """
+    DEPRECATED: update pandas_gbq.schema.bigquery_to_pyarrow, instead, which is
+    to be added in https://github.com/googleapis/python-bigquery-pandas/pull/893.
+
+    Returns:
+        The Arrow scalar type that the input BigQuery scalar type maps to.
+        If it cannot find the BigQuery scalar, return None.
+    """
+    return _BQ_TO_ARROW_SCALARS.get(bq_scalar)
+
+
+def arrow_scalar_ids_to_bq(arrow_scalar: Any):
+    """
+    DEPRECATED: update pandas_gbq.schema.pyarrow_to_bigquery, instead.
+
+    Returns:
+        The BigQuery scalar type that the input arrow scalar type maps to.
+        If it cannot find the arrow scalar, return None.
+    """
+    return _ARROW_SCALAR_IDS_TO_BQ.get(arrow_scalar)

google/cloud/bigquery/_tqdm_helpers.py

@@ -0,0 +1,137 @@
+# Copyright 2019 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.
+
+"""Shared helper functions for tqdm progress bar."""
+
+import concurrent.futures
+import sys
+import time
+import typing
+from typing import Optional
+import warnings
+
+try:
+    import tqdm  # type: ignore
+except ImportError:
+    tqdm = None
+
+try:
+    import tqdm.notebook as tqdm_notebook  # type: ignore
+except ImportError:
+    tqdm_notebook = None
+
+if typing.TYPE_CHECKING:  # pragma: NO COVER
+    from google.cloud.bigquery import QueryJob
+    from google.cloud.bigquery.table import RowIterator
+
+_NO_TQDM_ERROR = (
+    "A progress bar was requested, but there was an error loading the tqdm "
+    "library. Please install tqdm to use the progress bar functionality."
+)
+
+_PROGRESS_BAR_UPDATE_INTERVAL = 0.5
+
+
+def get_progress_bar(progress_bar_type, description, total, unit):
+    """Construct a tqdm progress bar object, if tqdm is installed."""
+    if tqdm is None or tqdm_notebook is None and progress_bar_type == "tqdm_notebook":
+        if progress_bar_type is not None:
+            warnings.warn(_NO_TQDM_ERROR, UserWarning, stacklevel=3)
+        return None
+
+    try:
+        if progress_bar_type == "tqdm":
+            return tqdm.tqdm(
+                bar_format="{l_bar}{bar}|",
+                colour="green",
+                desc=description,
+                file=sys.stdout,
+                total=total,
+                unit=unit,
+            )
+        elif progress_bar_type == "tqdm_notebook":
+            return tqdm_notebook.tqdm(
+                bar_format="{l_bar}{bar}|",
+                desc=description,
+                file=sys.stdout,
+                total=total,
+                unit=unit,
+            )
+        elif progress_bar_type == "tqdm_gui":
+            return tqdm.tqdm_gui(desc=description, total=total, unit=unit)
+    except (KeyError, TypeError):  # pragma: NO COVER
+        # Protect ourselves from any tqdm errors. In case of
+        # unexpected tqdm behavior, just fall back to showing
+        # no progress bar.
+        warnings.warn(_NO_TQDM_ERROR, UserWarning, stacklevel=3)
+    return None
+
+
+def wait_for_query(
+    query_job: "QueryJob",
+    progress_bar_type: Optional[str] = None,
+    max_results: Optional[int] = None,
+) -> "RowIterator":
+    """Return query result and display a progress bar while the query running, if tqdm is installed.
+
+    Args:
+        query_job:
+            The job representing the execution of the query on the server.
+        progress_bar_type:
+            The type of progress bar to use to show query progress.
+        max_results:
+            The maximum number of rows the row iterator should return.
+
+    Returns:
+        A row iterator over the query results.
+    """
+    default_total = 1
+    current_stage = None
+    start_time = time.perf_counter()
+
+    progress_bar = get_progress_bar(
+        progress_bar_type, "Query is running", default_total, "query"
+    )
+    if progress_bar is None:
+        return query_job.result(max_results=max_results)
+
+    i = 0
+    while True:
+        if query_job.query_plan:
+            default_total = len(query_job.query_plan)
+            current_stage = query_job.query_plan[i]
+            progress_bar.total = len(query_job.query_plan)
+            progress_bar.set_description(
+                f"Query executing stage {current_stage.name} and status {current_stage.status} : {time.perf_counter() - start_time:.2f}s"
+            )
+        try:
+            query_result = query_job.result(
+                timeout=_PROGRESS_BAR_UPDATE_INTERVAL, max_results=max_results
+            )
+            progress_bar.update(default_total)
+            progress_bar.set_description(
+                f"Job ID {query_job.job_id} successfully executed",
+            )
+            break
+        except concurrent.futures.TimeoutError:
+            query_job.reload()  # Refreshes the state via a GET request.
+            if current_stage:
+                if current_stage.status == "COMPLETE":
+                    if i < default_total - 1:
+                        progress_bar.update(i + 1)
+                        i += 1
+            continue
+
+    progress_bar.close()
+    return query_result

google/cloud/bigquery/_versions_helpers.py

@@ -0,0 +1,264 @@
+# Copyright 2023 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.
+
+"""Shared helper functions for verifying versions of installed modules."""
+
+import sys
+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")
+_MIN_PANDAS_VERSION = packaging.version.Version("1.1.0")
+
+_MIN_PANDAS_VERSION_RANGE = packaging.version.Version("1.5.0")
+_MIN_PYARROW_VERSION_RANGE = packaging.version.Version("10.0.1")
+
+
+class PyarrowVersions:
+    """Version comparisons for pyarrow package."""
+
+    def __init__(self):
+        self._installed_version = None
+
+    @property
+    def installed_version(self) -> packaging.version.Version:
+        """Return the parsed version of pyarrow."""
+        if self._installed_version is None:
+            import pyarrow  # type: ignore
+
+            self._installed_version = packaging.version.parse(
+                # Use 0.0.0, since it is earlier than any released version.
+                # Legacy versions also have the same property, but
+                # creating a LegacyVersion has been deprecated.
+                # https://github.com/pypa/packaging/issues/321
+                getattr(pyarrow, "__version__", "0.0.0")
+            )
+
+        return self._installed_version
+
+    @property
+    def use_compliant_nested_type(self) -> bool:
+        return self.installed_version.major >= 4
+
+    def try_import(self, raise_if_error: bool = False) -> Any:
+        """Verifies that a recent enough version of pyarrow extra is installed.
+
+        The function assumes that pyarrow extra is installed, and should thus
+        be used in places where this assumption holds.
+
+        Because `pip` can install an outdated version of this extra despite
+        the constraints in `setup.py`, the calling code can use this helper
+        to verify the version compatibility at runtime.
+
+        Returns:
+            The ``pyarrow`` module or ``None``.
+
+        Raises:
+            exceptions.LegacyPyarrowError:
+                If the pyarrow package is outdated and ``raise_if_error`` is
+                ``True``.
+        """
+        try:
+            import pyarrow
+        except ImportError as exc:
+            if raise_if_error:
+                raise exceptions.LegacyPyarrowError(
+                    "pyarrow package not found. Install pyarrow version >="
+                    f" {_MIN_PYARROW_VERSION}."
+                ) from exc
+            return None
+
+        if self.installed_version < _MIN_PYARROW_VERSION:
+            if raise_if_error:
+                msg = (
+                    "Dependency pyarrow is outdated, please upgrade"
+                    f" it to version >= {_MIN_PYARROW_VERSION}"
+                    f" (version found: {self.installed_version})."
+                )
+                raise exceptions.LegacyPyarrowError(msg)
+            return None
+
+        return pyarrow
+
+
+PYARROW_VERSIONS = PyarrowVersions()
+
+
+class BQStorageVersions:
+    """Version comparisons for google-cloud-bigqueyr-storage package."""
+
+    def __init__(self):
+        self._installed_version = None
+
+    @property
+    def installed_version(self) -> packaging.version.Version:
+        """Return the parsed version of google-cloud-bigquery-storage."""
+        if self._installed_version is None:
+            from google.cloud import bigquery_storage
+
+            self._installed_version = packaging.version.parse(
+                # Use 0.0.0, since it is earlier than any released version.
+                # Legacy versions also have the same property, but
+                # creating a LegacyVersion has been deprecated.
+                # https://github.com/pypa/packaging/issues/321
+                getattr(bigquery_storage, "__version__", "0.0.0")
+            )
+
+        return self._installed_version  # type: ignore
+
+    @property
+    def is_read_session_optional(self) -> bool:
+        """True if read_session is optional to rows().
+
+        See: https://github.com/googleapis/python-bigquery-storage/pull/228
+        """
+        return self.installed_version >= _BQ_STORAGE_OPTIONAL_READ_SESSION_VERSION
+
+    def try_import(self, raise_if_error: bool = False) -> Any:
+        """Tries to import the bigquery_storage module, and returns results
+        accordingly. It also verifies the module version is recent enough.
+
+        If the import succeeds, returns the ``bigquery_storage`` module.
+
+        If the import fails,
+        returns ``None`` when ``raise_if_error == False``,
+        raises Error when ``raise_if_error == True``.
+
+        Returns:
+            The ``bigquery_storage`` module or ``None``.
+
+        Raises:
+            exceptions.BigQueryStorageNotFoundError:
+                If google-cloud-bigquery-storage is not installed
+            exceptions.LegacyBigQueryStorageError:
+                If google-cloud-bigquery-storage package is outdated
+        """
+        try:
+            from google.cloud import bigquery_storage  # type: ignore
+        except ImportError:
+            if raise_if_error:
+                msg = (
+                    "Package google-cloud-bigquery-storage not found. "
+                    "Install google-cloud-bigquery-storage version >= "
+                    f"{_MIN_BQ_STORAGE_VERSION}."
+                )
+                raise exceptions.BigQueryStorageNotFoundError(msg)
+            return None
+
+        if self.installed_version < _MIN_BQ_STORAGE_VERSION:
+            if raise_if_error:
+                msg = (
+                    "Dependency google-cloud-bigquery-storage is outdated, "
+                    f"please upgrade it to version >= {_MIN_BQ_STORAGE_VERSION} "
+                    f"(version found: {self.installed_version})."
+                )
+                raise exceptions.LegacyBigQueryStorageError(msg)
+            return None
+
+        return bigquery_storage
+
+
+BQ_STORAGE_VERSIONS = BQStorageVersions()
+
+
+class PandasVersions:
+    """Version comparisons for pandas package."""
+
+    def __init__(self):
+        self._installed_version = None
+
+    @property
+    def installed_version(self) -> packaging.version.Version:
+        """Return the parsed version of pandas"""
+        if self._installed_version is None:
+            import pandas  # type: ignore
+
+            self._installed_version = packaging.version.parse(
+                # Use 0.0.0, since it is earlier than any released version.
+                # Legacy versions also have the same property, but
+                # creating a LegacyVersion has been deprecated.
+                # https://github.com/pypa/packaging/issues/321
+                getattr(pandas, "__version__", "0.0.0")
+            )
+
+        return self._installed_version
+
+    def try_import(self, raise_if_error: bool = False) -> Any:
+        """Verify that a recent enough version of pandas extra is installed.
+        The function assumes that pandas extra is installed, and should thus
+        be used in places where this assumption holds.
+        Because `pip` can install an outdated version of this extra despite
+        the constraints in `setup.py`, the calling code can use this helper
+        to verify the version compatibility at runtime.
+        Returns:
+            The ``pandas`` module or ``None``.
+        Raises:
+            exceptions.LegacyPandasError:
+                If the pandas package is outdated and ``raise_if_error`` is
+                ``True``.
+        """
+        try:
+            import pandas
+        except ImportError as exc:
+            if raise_if_error:
+                raise exceptions.LegacyPandasError(
+                    "pandas package not found. Install pandas version >="
+                    f" {_MIN_PANDAS_VERSION}"
+                ) from exc
+            return None
+
+        if self.installed_version < _MIN_PANDAS_VERSION:
+            if raise_if_error:
+                msg = (
+                    "Dependency pandas is outdated, please upgrade"
+                    f" it to version >= {_MIN_PANDAS_VERSION}"
+                    f" (version found: {self.installed_version})."
+                )
+                raise exceptions.LegacyPandasError(msg)
+            return None
+
+        return pandas
+
+
+PANDAS_VERSIONS = PandasVersions()
+
+# Since RANGE support in pandas requires specific versions
+# of both pyarrow and pandas, we make this a separate
+# constant instead of as a property of PANDAS_VERSIONS
+# or PYARROW_VERSIONS.
+SUPPORTS_RANGE_PYARROW = (
+    PANDAS_VERSIONS.try_import() is not None
+    and PANDAS_VERSIONS.installed_version >= _MIN_PANDAS_VERSION_RANGE
+    and PYARROW_VERSIONS.try_import() is not None
+    and PYARROW_VERSIONS.installed_version >= _MIN_PYARROW_VERSION_RANGE
+)
+
+
+def extract_runtime_version():
+    # Retrieve the version information
+    version_info = sys.version_info
+
+    # Extract the major, minor, and micro components
+    major = version_info.major
+    minor = version_info.minor
+    micro = version_info.micro
+
+    # Display the version number in a clear format
+    return major, minor, micro