Thanks to visit codestin.com
Credit goes to github.com

Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 17 commits into from
May 5, 2025
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
17 commits
Select commit Hold shift + click to select a range
1bc248e
feat: add dry_run parameter to read_gbq() and read_gbq_query()
sycai Apr 30, 2025
073c341
fix lint
sycai Apr 30, 2025
c0f5163
Merge branch 'main' into sycai_dry_run
sycai Apr 30, 2025
8c8a723
Merge branch 'main' into sycai_dry_run
sycai Apr 30, 2025
7b30eb1
Merge branch 'main' into sycai_dry_run
sycai Apr 30, 2025
2bceff7
chore(main): release 2.2.0 (#1643)
release-please[bot] Apr 30, 2025
42edce4
create a different stats report for reading gbq tables
sycai May 1, 2025
d59c4bd
Merge branch 'main' into sycai_dry_run
sycai May 1, 2025
4d6a59a
fix lint
sycai May 1, 2025
3624795
🦉 Updates from OwlBot post-processor
gcf-owl-bot[bot] May 1, 2025
089d586
Merge branch 'sycai_dry_run' of https://github.com/googleapis/python-…
gcf-owl-bot[bot] May 1, 2025
669c041
rename column count and column dtypes
sycai May 1, 2025
ac8c05c
Merge branch 'main' into sycai_dry_run
sycai May 1, 2025
7dc90d9
fix typo
sycai May 1, 2025
698b409
Merge branch 'main' into sycai_dry_run
sycai May 1, 2025
80e08bf
format code
sycai May 1, 2025
3fe8b8e
Merge branch 'main' into sycai_dry_run
sycai May 5, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
71 changes: 10 additions & 61 deletions bigframes/core/blocks.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
107 changes: 104 additions & 3 deletions bigframes/pandas/io/api.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
Literal,
MutableSequence,
Optional,
overload,
Sequence,
Tuple,
Union,
Expand Down Expand Up @@ -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,
*,
Expand All @@ -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,
Expand All @@ -177,6 +211,7 @@ def read_gbq(
filters=filters,
use_cache=use_cache,
col_order=col_order,
dry_run=dry_run,
)


Expand Down Expand Up @@ -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,
*,
Expand All @@ -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,
Expand All @@ -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,
*,
Expand All @@ -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,
Expand All @@ -256,6 +356,7 @@ def read_gbq_table(
filters=filters,
use_cache=use_cache,
col_order=col_order,
dry_run=dry_run,
)


Expand Down
Loading