Skip to content

test: verify read_gbq_colab label preservation in anywidget display mode #17887

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
shuoweil merged 5 commits into from
Jul 28, 2026
+33 −1
Merged

test: verify read_gbq_colab label preservation in anywidget display mode #17887

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
d0b590e
fix: preserve read_gbq_colab label in anywidget mode
shuoweil Jul 23, 2026
dc53247
Update packages/bigframes/bigframes/session/_io/bigquery/__init__.py
shuoweil Jul 23, 2026
a8a07ba
test: skip anywidget test when dependencies missing
shuoweil Jul 24, 2026
c563483
Merge branch 'main' into shuowei-sql-cell-track
shuoweil Jul 24, 2026
9185cdd
test(bigframes): verify read_gbq_colab label preservation in anywidge…
shuoweil Jul 24, 2026
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: 8 additions & 2 deletions packages/bigframes/bigframes/session/_io/bigquery/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -68,8 +68,14 @@ def create_job_configs_labels(
job_configs_labels = dict(job_configs_labels)

if api_methods and "bigframes-api" not in job_configs_labels:
job_configs_labels["bigframes-api"] = api_methods[0]
del api_methods[0]
colab_idx = next(
(i for i, m in enumerate(api_methods) if "read_gbq_colab" in m), None
)
if colab_idx is not None:
job_configs_labels["bigframes-api"] = api_methods.pop(colab_idx)

@sycai sycai Jul 24, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

It's not immediately clear to me why we need to pop at the colab_idx instead of at head. Could you add some explanatory comments here?

@shuoweil shuoweil Jul 24, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

We have reverted the method-specific popping logic. We used to want to override the label precedence in Python due to the tracking mechanism of dashboard. However, we toss this approach for now, and I will update the change in dashboard instead.

else:
job_configs_labels["bigframes-api"] = api_methods[0]
del api_methods[0]
Comment thread
shuoweil marked this conversation as resolved.
Outdated

@tswast tswast Jul 24, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

I don't understand why we need a special case for read_gbq_colab. It should be the most recent API call before we execute a requests, so it should always be preserved. Could you double-check the tests to make sure we drop only the oldest labels when we have too many?

@shuoweil shuoweil Jul 24, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Agreed. We have removed this special case. In this PR, we just make sure session-read_gbq_colab is anyiwdget mode is preserved. I will make the changes from dashboard to loop through all labels. This approach would avoid hardcoded specials cases.


# Make sure we always populate bigframes-api with _something_, even if we
# have a code path which doesn't populate the list of api_methods. See
Expand Down
31 changes: 31 additions & 0 deletions packages/bigframes/tests/unit/session/test_read_gbq_colab.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,37 @@ def test_read_gbq_colab_includes_label():
assert "session-read_gbq_colab" in label_values


def test_read_gbq_colab_includes_label_in_anywidget_mode():
"""Make sure read_gbq_colab label is preserved as the primary bigframes-api label in anywidget mode."""
import bigframes
import bigframes.display.html as bf_html

bqclient = mock.create_autospec(bigquery.Client, instance=True)
bqclient.project = "proj"
session = mocks.create_bigquery_session(bqclient=bqclient)
df = session._read_gbq_colab("SELECT 'read-gbq-colab-test'")

with bigframes.option_context("display.render_mode", "anywidget"):
_ = bf_html.get_anywidget_bundle(df)

label_values = []
bigframes_api_labels = []
for kall in itertools.chain(
bqclient.query_and_wait.call_args_list,
bqclient._query_and_wait_bigframes.call_args_list,
bqclient.query.call_args_list,
):
job_config = kall.kwargs.get("job_config")
if job_config is None:
continue
label_values.extend(job_config.labels.values())
if "bigframes-api" in job_config.labels:
bigframes_api_labels.append(job_config.labels["bigframes-api"])

assert "session-read_gbq_colab" in label_values
assert "session-read_gbq_colab" in bigframes_api_labels


@pytest.mark.parametrize("dry_run", [True, False])
def test_read_gbq_colab_includes_formatted_values_in_dry_run(monkeypatch, dry_run):
bqclient = mock.create_autospec(bigquery.Client, instance=True)
Expand Down
Loading