Skip to content

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

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

Closed
tswast wants to merge 1 commit into from
+24 −42
Closed

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

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
12 changes: 4 additions & 8 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,17 +15,15 @@

import benchmark.utils as utils

import bigframes.session
import bigframes.pandas

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(
df = bigframes.pandas._read_gbq_colab(
f"SELECT * FROM `{project_id}`.{dataset_id}.{table_id}"
)

Expand Down Expand Up @@ -56,9 +54,8 @@ def aggregate_output(
project_id,
dataset_id,
table_id,
session,
suffix,
) = utils.get_configuration(include_table_id=True)
) = utils.get_configuration(include_table_id=True, include_session=False)
Copy link
Collaborator Author

@tswast tswast Jul 25, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This parameter doesn't actually exist. Might be nice to have a start_session=False argument, though.

current_path = pathlib.Path(__file__).absolute()

utils.get_execution_time(
Expand All @@ -68,5 +65,4 @@ def aggregate_output(
project_id=project_id,
dataset_id=dataset_id,
table_id=table_id,
session=session,
)
10 changes: 4 additions & 6 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,13 +15,13 @@

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,
)
Expand All @@ -32,9 +32,8 @@ def dry_run(*, project_id, dataset_id, table_id, session: bigframes.session.Sess
project_id,
dataset_id,
table_id,
session,
suffix,
) = utils.get_configuration(include_table_id=True)
) = utils.get_configuration(include_table_id=True, include_session=False)
current_path = pathlib.Path(__file__).absolute()

utils.get_execution_time(
Expand All @@ -44,5 +43,4 @@ def dry_run(*, project_id, dataset_id, table_id, session: bigframes.session.Sess
project_id=project_id,
dataset_id=dataset_id,
table_id=table_id,
session=session,
)
12 changes: 4 additions & 8 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,17 +15,15 @@

import benchmark.utils as utils

import bigframes.session
import bigframes.pandas

PAGE_SIZE = utils.READ_GBQ_COLAB_PAGE_SIZE


def filter_output(
*, project_id, dataset_id, table_id, session: bigframes.session.Session
):
def filter_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 @@ -48,9 +46,8 @@ def filter_output(
project_id,
dataset_id,
table_id,
session,
suffix,
) = utils.get_configuration(include_table_id=True)
) = utils.get_configuration(include_table_id=True, include_session=False)
current_path = pathlib.Path(__file__).absolute()

utils.get_execution_time(
Expand All @@ -60,5 +57,4 @@ def filter_output(
project_id=project_id,
dataset_id=dataset_id,
table_id=table_id,
session=session,
)
10 changes: 4 additions & 6 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 @@ -37,9 +37,8 @@ def first_page(*, project_id, dataset_id, table_id, session: bigframes.session.S
project_id,
dataset_id,
table_id,
session,
suffix,
) = utils.get_configuration(include_table_id=True)
) = utils.get_configuration(include_table_id=True, include_session=False)
current_path = pathlib.Path(__file__).absolute()

utils.get_execution_time(
Expand All @@ -49,5 +48,4 @@ def first_page(*, project_id, dataset_id, table_id, session: bigframes.session.S
project_id=project_id,
dataset_id=dataset_id,
table_id=table_id,
session=session,
)
10 changes: 4 additions & 6 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 @@ -38,9 +38,8 @@ def last_page(*, project_id, dataset_id, table_id, session: bigframes.session.Se
project_id,
dataset_id,
table_id,
session,
suffix,
) = utils.get_configuration(include_table_id=True)
) = utils.get_configuration(include_table_id=True, include_session=False)
current_path = pathlib.Path(__file__).absolute()

utils.get_execution_time(
Expand All @@ -50,5 +49,4 @@ def last_page(*, project_id, dataset_id, table_id, session: bigframes.session.Se
project_id=project_id,
dataset_id=dataset_id,
table_id=table_id,
session=session,
)
12 changes: 4 additions & 8 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 @@ -48,9 +46,8 @@ def sort_output(
project_id,
dataset_id,
table_id,
session,
suffix,
) = utils.get_configuration(include_table_id=True)
) = utils.get_configuration(include_table_id=True, include_session=False)
current_path = pathlib.Path(__file__).absolute()

utils.get_execution_time(
Expand All @@ -60,5 +57,4 @@ def sort_output(
project_id=project_id,
dataset_id=dataset_id,
table_id=table_id,
session=session,
)