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,
  )