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 3 commits
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
11 changes: 9 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,15 @@ 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]
api_methods = list(api_methods)

@sycai sycai Jul 28, 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.

Non-rhetorical question: do we need this line of change? It alters the behavior of the original code, which alters the original list that was passed in at line 73.

@shuoweil shuoweil Jul 28, 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.

Yes, this shallow copy is intentional as a safety guard. There are two reasons. (1) at line 73, del api_methods[0]
removes the primary API label. Creating a copy ensures this deletion stys local to the helper function and does not mutate the caller's original list if reused or logs elsewhere.
(2) api_methods may sometime pass in tuple. We have a unit test reate_job_configs_labels(labels, ("read_gbq_colab", "head")) will pass tuples to label. In this case, tuple is immutable. This copy is a guard to avoid errors for 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
34 changes: 34 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,40 @@ 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."""
pytest.importorskip("anywidget")
pytest.importorskip("traitlets")

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