feat: add dry_run parameter to read_gbq(), read_gbq_table() and read_gbq_query() (#1674)

* feat: add dry_run parameter to read_gbq() and read_gbq_query()

* fix lint

* chore(main): release 2.2.0 (#1643)

Co-authored-by: release-please[bot] <55107282+release-please[bot]@users.noreply.github.com>

* create a different stats report for reading gbq tables

* fix lint

* 🦉 Updates from OwlBot post-processor

See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md

* rename column count and column dtypes

* fix typo

* format code

---------

Co-authored-by: release-please[bot] <55107282+release-please[bot]@users.noreply.github.com>
Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>

Author: 3 people

bigframes/core/blocks.py

@@ -22,25 +22,14 @@
 from __future__ import annotations
 
 import ast
-import copy
 import dataclasses
 import datetime
 import functools
 import itertools
 import random
 import textwrap
 import typing
-from typing import (
-    Any,
-    Iterable,
-    List,
-    Literal,
-    Mapping,
-    Optional,
-    Sequence,
-    Tuple,
-    Union,
-)
+from typing import Iterable, List, Literal, Mapping, Optional, Sequence, Tuple, Union
 import warnings
 
 import bigframes_vendored.constants as constants
@@ -69,6 +58,7 @@
 import bigframes.exceptions as bfe
 import bigframes.operations as ops
 import bigframes.operations.aggregations as agg_ops
+from bigframes.session import dry_runs
 from bigframes.session import executor as executors
 
 # Type constraint for wherever column labels are used
@@ -822,59 +812,18 @@ def _compute_dry_run(
         if sampling.enable_downsampling:
             raise NotImplementedError("Dry run with sampling is not supported")
 
-        index: List[Any] = []
-        values: List[Any] = []
-
-        index.append("columnCount")
-        values.append(len(self.value_columns))
-        index.append("columnDtypes")
-        values.append(
-            {
-                col: self.expr.get_column_type(self.resolve_label_exact_or_error(col))
-                for col in self.column_labels
-            }
-        )
-
-        index.append("indexLevel")
-        values.append(self.index.nlevels)
-        index.append("indexDtypes")
-        values.append(self.index.dtypes)
-
         expr = self._apply_value_keys_to_expr(value_keys=value_keys)
         query_job = self.session._executor.dry_run(expr, ordered)
-        job_api_repr = copy.deepcopy(query_job._properties)
-
-        job_ref = job_api_repr["jobReference"]
-        for key, val in job_ref.items():
-            index.append(key)
-            values.append(val)
-
-        index.append("jobType")
-        values.append(job_api_repr["configuration"]["jobType"])
-
-        query_config = job_api_repr["configuration"]["query"]
-        for key in ("destinationTable", "useLegacySql"):
-            index.append(key)
-            values.append(query_config.get(key))
-
-        query_stats = job_api_repr["statistics"]["query"]
-        for key in (
-            "referencedTables",
-            "totalBytesProcessed",
-            "cacheHit",
-            "statementType",
-        ):
-            index.append(key)
-            values.append(query_stats.get(key))
 
-        index.append("creationTime")
-        values.append(
-            pd.Timestamp(
-                job_api_repr["statistics"]["creationTime"], unit="ms", tz="UTC"
-            )
-        )
+        column_dtypes = {
+            col: self.expr.get_column_type(self.resolve_label_exact_or_error(col))
+            for col in self.column_labels
+        }
 
-        return pd.Series(values, index=index), query_job
+        dry_run_stats = dry_runs.get_query_stats_with_dtypes(
+            query_job, column_dtypes, self.index.dtypes
+        )
+        return dry_run_stats, query_job
 
     def _apply_value_keys_to_expr(self, value_keys: Optional[Iterable[str]] = None):
         expr = self._expr

bigframes/pandas/io/api.py

@@ -25,6 +25,7 @@
     Literal,
     MutableSequence,
     Optional,
+    overload,
     Sequence,
     Tuple,
     Union,
@@ -155,6 +156,38 @@ def read_json(
 read_json.__doc__ = inspect.getdoc(bigframes.session.Session.read_json)
 
 
+@overload
+def read_gbq(  # type: ignore[overload-overlap]
+    query_or_table: str,
+    *,
+    index_col: Iterable[str] | str | bigframes.enums.DefaultIndexKind = ...,
+    columns: Iterable[str] = ...,
+    configuration: Optional[Dict] = ...,
+    max_results: Optional[int] = ...,
+    filters: vendored_pandas_gbq.FiltersType = ...,
+    use_cache: Optional[bool] = ...,
+    col_order: Iterable[str] = ...,
+    dry_run: Literal[False] = ...,
+) -> bigframes.dataframe.DataFrame:
+    ...
+
+
+@overload
+def read_gbq(
+    query_or_table: str,
+    *,
+    index_col: Iterable[str] | str | bigframes.enums.DefaultIndexKind = ...,
+    columns: Iterable[str] = ...,
+    configuration: Optional[Dict] = ...,
+    max_results: Optional[int] = ...,
+    filters: vendored_pandas_gbq.FiltersType = ...,
+    use_cache: Optional[bool] = ...,
+    col_order: Iterable[str] = ...,
+    dry_run: Literal[True] = ...,
+) -> pandas.Series:
+    ...
+
+
 def read_gbq(
     query_or_table: str,
     *,
@@ -165,7 +198,8 @@ def read_gbq(
     filters: vendored_pandas_gbq.FiltersType = (),
     use_cache: Optional[bool] = None,
     col_order: Iterable[str] = (),
-) -> bigframes.dataframe.DataFrame:
+    dry_run: bool = False,
+) -> bigframes.dataframe.DataFrame | pandas.Series:
     _set_default_session_location_if_possible(query_or_table)
     return global_session.with_default_session(
         bigframes.session.Session.read_gbq,
@@ -177,6 +211,7 @@ def read_gbq(
         filters=filters,
         use_cache=use_cache,
         col_order=col_order,
+        dry_run=dry_run,
     )
 
 
@@ -208,6 +243,38 @@ def read_gbq_object_table(
 )
 
 
+@overload
+def read_gbq_query(  # type: ignore[overload-overlap]
+    query: str,
+    *,
+    index_col: Iterable[str] | str | bigframes.enums.DefaultIndexKind = ...,
+    columns: Iterable[str] = ...,
+    configuration: Optional[Dict] = ...,
+    max_results: Optional[int] = ...,
+    use_cache: Optional[bool] = ...,
+    col_order: Iterable[str] = ...,
+    filters: vendored_pandas_gbq.FiltersType = ...,
+    dry_run: Literal[False] = ...,
+) -> bigframes.dataframe.DataFrame:
+    ...
+
+
+@overload
+def read_gbq_query(
+    query: str,
+    *,
+    index_col: Iterable[str] | str | bigframes.enums.DefaultIndexKind = ...,
+    columns: Iterable[str] = ...,
+    configuration: Optional[Dict] = ...,
+    max_results: Optional[int] = ...,
+    use_cache: Optional[bool] = ...,
+    col_order: Iterable[str] = ...,
+    filters: vendored_pandas_gbq.FiltersType = ...,
+    dry_run: Literal[True] = ...,
+) -> pandas.Series:
+    ...
+
+
 def read_gbq_query(
     query: str,
     *,
@@ -218,7 +285,8 @@ def read_gbq_query(
     use_cache: Optional[bool] = None,
     col_order: Iterable[str] = (),
     filters: vendored_pandas_gbq.FiltersType = (),
-) -> bigframes.dataframe.DataFrame:
+    dry_run: bool = False,
+) -> bigframes.dataframe.DataFrame | pandas.Series:
     _set_default_session_location_if_possible(query)
     return global_session.with_default_session(
         bigframes.session.Session.read_gbq_query,
@@ -230,12 +298,43 @@ def read_gbq_query(
         use_cache=use_cache,
         col_order=col_order,
         filters=filters,
+        dry_run=dry_run,
     )
 
 
 read_gbq_query.__doc__ = inspect.getdoc(bigframes.session.Session.read_gbq_query)
 
 
+@overload
+def read_gbq_table(  # type: ignore[overload-overlap]
+    query: str,
+    *,
+    index_col: Iterable[str] | str | bigframes.enums.DefaultIndexKind = ...,
+    columns: Iterable[str] = ...,
+    max_results: Optional[int] = ...,
+    filters: vendored_pandas_gbq.FiltersType = ...,
+    use_cache: bool = ...,
+    col_order: Iterable[str] = ...,
+    dry_run: Literal[False] = ...,
+) -> bigframes.dataframe.DataFrame:
+    ...
+
+
+@overload
+def read_gbq_table(
+    query: str,
+    *,
+    index_col: Iterable[str] | str | bigframes.enums.DefaultIndexKind = ...,
+    columns: Iterable[str] = ...,
+    max_results: Optional[int] = ...,
+    filters: vendored_pandas_gbq.FiltersType = ...,
+    use_cache: bool = ...,
+    col_order: Iterable[str] = ...,
+    dry_run: Literal[True] = ...,
+) -> pandas.Series:
+    ...
+
+
 def read_gbq_table(
     query: str,
     *,
@@ -245,7 +344,8 @@ def read_gbq_table(
     filters: vendored_pandas_gbq.FiltersType = (),
     use_cache: bool = True,
     col_order: Iterable[str] = (),
-) -> bigframes.dataframe.DataFrame:
+    dry_run: bool = False,
+) -> bigframes.dataframe.DataFrame | pandas.Series:
     _set_default_session_location_if_possible(query)
     return global_session.with_default_session(
         bigframes.session.Session.read_gbq_table,
@@ -256,6 +356,7 @@ def read_gbq_table(
         filters=filters,
         use_cache=use_cache,
         col_order=col_order,
+        dry_run=dry_run,
     )