Skip to content

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

Open
shuoweil wants to merge 5 commits into
mainfrom
shuowei-sql-cell-track
Open

test: verify read_gbq_colab label preservation in anywidget display mode#17887
shuoweil wants to merge 5 commits into
mainfrom
shuowei-sql-cell-track

Conversation

@shuoweil

@shuoweil shuoweil commented Jul 23, 2026

Copy link
Copy Markdown
Contributor

This PR ensures that read_gbq_colab telemetry labels are reliably preserved across BigQuery job configurations when evaluating DataFrames in anywidget display mode.

Fixes #<538249690> 🦕

@shuoweil
shuoweil requested review from GarrettWu, sycai and tswast July 23, 2026 23:24
@shuoweil shuoweil self-assigned this Jul 23, 2026
@shuoweil
shuoweil requested review from a team as code owners July 23, 2026 23:24

@gemini-code-assist gemini-code-assist Bot left a comment

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.

Code Review

This pull request modifies the job configuration label creation to prioritize the 'read_gbq_colab' API method as the primary 'bigframes-api' label, and adds a unit test to verify this behavior in anywidget mode. The reviewer recommends converting the 'api_methods' parameter to a local list before mutating it to prevent unintended side effects on the caller and avoid potential runtime errors if a non-list iterable is passed.

Comment thread packages/bigframes/bigframes/session/_io/bigquery/__init__.py Outdated
shuoweil and others added 2 commits July 23, 2026 16:38
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
@shuoweil

Copy link
Copy Markdown
Contributor Author

The failed importer test does not related to the change in current branch.

@parthea parthea assigned tswast and unassigned shuoweil Jul 24, 2026
(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)

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?

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.

Comment on lines +72 to +79
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)
else:
job_configs_labels["bigframes-api"] = api_methods[0]
del api_methods[0]

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?

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.

@shuoweil shuoweil changed the title fix: preserve read_gbq_colab label in anywidget mode test: verify read_gbq_colab label preservation in anywidget display mode Jul 24, 2026
@shuoweil
shuoweil requested review from sycai and tswast July 24, 2026 23:57
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants