test: verify read_gbq_colab label preservation in anywidget display mode#17887
test: verify read_gbq_colab label preservation in anywidget display mode#17887shuoweil wants to merge 5 commits into
Conversation
There was a problem hiding this comment.
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.
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
|
The failed importer test does not related to the change in current branch. |
| (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) |
There was a problem hiding this comment.
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?
There was a problem hiding this comment.
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.
| 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] |
There was a problem hiding this comment.
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?
There was a problem hiding this comment.
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.
This PR ensures that
read_gbq_colabtelemetry labels are reliably preserved across BigQuery job configurations when evaluating DataFrames inanywidgetdisplay mode.Fixes #<538249690> 🦕