Skip to content

chore: use bigframes.pandas._read_gbq_colab() in benchmarks #1942

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
tswast merged 1 commit into from
Jul 29, 2025
Merged

chore: use bigframes.pandas._read_gbq_colab() in benchmarks #1942

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
10 changes: 3 additions & 7 deletions tests/benchmark/read_gbq_colab/aggregate_output.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,19 +15,15 @@

import benchmark.utils as utils

import bigframes.session
import bigframes.pandas as bpd

PAGE_SIZE = utils.READ_GBQ_COLAB_PAGE_SIZE


def aggregate_output(
*, project_id, dataset_id, table_id, session: bigframes.session.Session
):
def aggregate_output(*, project_id, dataset_id, table_id):
# TODO(tswast): Support alternative query if table_id is a local DataFrame,
# e.g. "{local_inline}" or "{local_large}"
df = session._read_gbq_colab(
f"SELECT * FROM `{project_id}`.{dataset_id}.{table_id}"
)
df = bpd._read_gbq_colab(f"SELECT * FROM `{project_id}`.{dataset_id}.{table_id}")

# Simulate getting the first page, since we'll always do that first in the UI.
df.shape
Expand Down
9 changes: 4 additions & 5 deletions tests/benchmark/read_gbq_colab/dry_run.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,20 +15,20 @@

import benchmark.utils as utils

import bigframes.session
import bigframes.pandas


def dry_run(*, project_id, dataset_id, table_id, session: bigframes.session.Session):
def dry_run(*, project_id, dataset_id, table_id):
# TODO(tswast): Support alternative query if table_id is a local DataFrame,
# e.g. "{local_inline}" or "{local_large}"
session._read_gbq_colab(
bigframes.pandas._read_gbq_colab(
f"SELECT * FROM `{project_id}`.{dataset_id}.{table_id}",
dry_run=True,
)


if __name__ == "__main__":
config = utils.get_configuration(include_table_id=True)
config = utils.get_configuration(include_table_id=True, start_session=False)
current_path = pathlib.Path(__file__).absolute()

utils.get_execution_time(
Expand All @@ -38,5 +38,4 @@ def dry_run(*, project_id, dataset_id, table_id, session: bigframes.session.Sess
project_id=config.project_id,
dataset_id=config.dataset_id,
table_id=config.table_id,
session=config.session,
)
12 changes: 6 additions & 6 deletions tests/benchmark/read_gbq_colab/filter_output.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,19 +15,20 @@

import benchmark.utils as utils

import bigframes.session
import bigframes.pandas as bpd

PAGE_SIZE = utils.READ_GBQ_COLAB_PAGE_SIZE


def filter_output(
*, project_id, dataset_id, table_id, session: bigframes.session.Session
*,
project_id,
dataset_id,
table_id,
):
# TODO(tswast): Support alternative query if table_id is a local DataFrame,
# e.g. "{local_inline}" or "{local_large}"
df = session._read_gbq_colab(
f"SELECT * FROM `{project_id}`.{dataset_id}.{table_id}"
)
df = bpd._read_gbq_colab(f"SELECT * FROM `{project_id}`.{dataset_id}.{table_id}")

# Simulate getting the first page, since we'll always do that first in the UI.
df.shape
Expand All @@ -54,5 +55,4 @@ def filter_output(
project_id=config.project_id,
dataset_id=config.dataset_id,
table_id=config.table_id,
session=config.session,
)
9 changes: 4 additions & 5 deletions tests/benchmark/read_gbq_colab/first_page.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,15 +15,15 @@

import benchmark.utils as utils

import bigframes.session
import bigframes.pandas

PAGE_SIZE = utils.READ_GBQ_COLAB_PAGE_SIZE


def first_page(*, project_id, dataset_id, table_id, session: bigframes.session.Session):
def first_page(*, project_id, dataset_id, table_id):
# TODO(tswast): Support alternative query if table_id is a local DataFrame,
# e.g. "{local_inline}" or "{local_large}"
df = session._read_gbq_colab(
df = bigframes.pandas._read_gbq_colab(
f"SELECT * FROM `{project_id}`.{dataset_id}.{table_id}"
)

Expand All @@ -33,7 +33,7 @@ def first_page(*, project_id, dataset_id, table_id, session: bigframes.session.S


if __name__ == "__main__":
config = utils.get_configuration(include_table_id=True)
config = utils.get_configuration(include_table_id=True, start_session=False)
current_path = pathlib.Path(__file__).absolute()

utils.get_execution_time(
Expand All @@ -43,5 +43,4 @@ def first_page(*, project_id, dataset_id, table_id, session: bigframes.session.S
project_id=config.project_id,
dataset_id=config.dataset_id,
table_id=config.table_id,
session=config.session,
)
9 changes: 4 additions & 5 deletions tests/benchmark/read_gbq_colab/last_page.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,15 +15,15 @@

import benchmark.utils as utils

import bigframes.session
import bigframes.pandas

PAGE_SIZE = utils.READ_GBQ_COLAB_PAGE_SIZE


def last_page(*, project_id, dataset_id, table_id, session: bigframes.session.Session):
def last_page(*, project_id, dataset_id, table_id):
# TODO(tswast): Support alternative query if table_id is a local DataFrame,
# e.g. "{local_inline}" or "{local_large}"
df = session._read_gbq_colab(
df = bigframes.pandas._read_gbq_colab(
f"SELECT * FROM `{project_id}`.{dataset_id}.{table_id}"
)

Expand All @@ -34,7 +34,7 @@ def last_page(*, project_id, dataset_id, table_id, session: bigframes.session.Se


if __name__ == "__main__":
config = utils.get_configuration(include_table_id=True)
config = utils.get_configuration(include_table_id=True, start_session=False)
current_path = pathlib.Path(__file__).absolute()

utils.get_execution_time(
Expand All @@ -44,5 +44,4 @@ def last_page(*, project_id, dataset_id, table_id, session: bigframes.session.Se
project_id=config.project_id,
dataset_id=config.dataset_id,
table_id=config.table_id,
session=config.session,
)
11 changes: 4 additions & 7 deletions tests/benchmark/read_gbq_colab/sort_output.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,17 +15,15 @@

import benchmark.utils as utils

import bigframes.session
import bigframes.pandas

PAGE_SIZE = utils.READ_GBQ_COLAB_PAGE_SIZE


def sort_output(
*, project_id, dataset_id, table_id, session: bigframes.session.Session
):
def sort_output(*, project_id, dataset_id, table_id):
# TODO(tswast): Support alternative query if table_id is a local DataFrame,
# e.g. "{local_inline}" or "{local_large}"
df = session._read_gbq_colab(
df = bigframes.pandas._read_gbq_colab(
f"SELECT * FROM `{project_id}`.{dataset_id}.{table_id}"
)

Expand All @@ -44,7 +42,7 @@ def sort_output(


if __name__ == "__main__":
config = utils.get_configuration(include_table_id=True)
config = utils.get_configuration(include_table_id=True, start_session=False)
current_path = pathlib.Path(__file__).absolute()

utils.get_execution_time(
Expand All @@ -54,5 +52,4 @@ def sort_output(
project_id=config.project_id,
dataset_id=config.dataset_id,
table_id=config.table_id,
session=config.session,
)
6 changes: 3 additions & 3 deletions tests/benchmark/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,12 +25,12 @@
class BenchmarkConfig:
project_id: str
dataset_id: str
session: bigframes.Session
session: bigframes.Session | None
benchmark_suffix: str | None
table_id: str | None = None


def get_configuration(include_table_id=False) -> BenchmarkConfig:
def get_configuration(include_table_id=False, start_session=True) -> BenchmarkConfig:
parser = argparse.ArgumentParser()
parser.add_argument(
"--project_id",
Expand Down Expand Up @@ -65,7 +65,7 @@ def get_configuration(include_table_id=False) -> BenchmarkConfig:
)

args = parser.parse_args()
session = _initialize_session(_str_to_bool(args.ordered))
session = _initialize_session(_str_to_bool(args.ordered)) if start_session else None

return BenchmarkConfig(
project_id=args.project_id,
Expand Down