From d0b590eb13a65fc86f8795dd6d9a712c01082c37 Mon Sep 17 00:00:00 2001 From: Shuowei Li Date: Thu, 23 Jul 2026 23:21:55 +0000 Subject: [PATCH 1/4] fix: preserve read_gbq_colab label in anywidget mode --- .../session/_io/bigquery/__init__.py | 10 ++++-- .../tests/unit/session/test_read_gbq_colab.py | 31 +++++++++++++++++++ 2 files changed, 39 insertions(+), 2 deletions(-) diff --git a/packages/bigframes/bigframes/session/_io/bigquery/__init__.py b/packages/bigframes/bigframes/session/_io/bigquery/__init__.py index 58bd5bd61748..47f677d6a97f 100644 --- a/packages/bigframes/bigframes/session/_io/bigquery/__init__.py +++ b/packages/bigframes/bigframes/session/_io/bigquery/__init__.py @@ -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) + else: + job_configs_labels["bigframes-api"] = api_methods[0] + del api_methods[0] # 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 diff --git a/packages/bigframes/tests/unit/session/test_read_gbq_colab.py b/packages/bigframes/tests/unit/session/test_read_gbq_colab.py index a168ccbad5e6..c2e696927434 100644 --- a/packages/bigframes/tests/unit/session/test_read_gbq_colab.py +++ b/packages/bigframes/tests/unit/session/test_read_gbq_colab.py @@ -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) From dc5324724e8c7d7b9d11d59564923f3c02c58f20 Mon Sep 17 00:00:00 2001 From: Shuowei Li Date: Thu, 23 Jul 2026 16:38:43 -0700 Subject: [PATCH 2/4] Update packages/bigframes/bigframes/session/_io/bigquery/__init__.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- packages/bigframes/bigframes/session/_io/bigquery/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/packages/bigframes/bigframes/session/_io/bigquery/__init__.py b/packages/bigframes/bigframes/session/_io/bigquery/__init__.py index 47f677d6a97f..a5b200296732 100644 --- a/packages/bigframes/bigframes/session/_io/bigquery/__init__.py +++ b/packages/bigframes/bigframes/session/_io/bigquery/__init__.py @@ -68,6 +68,7 @@ def create_job_configs_labels( job_configs_labels = dict(job_configs_labels) if api_methods and "bigframes-api" not in job_configs_labels: + api_methods = list(api_methods) colab_idx = next( (i for i, m in enumerate(api_methods) if "read_gbq_colab" in m), None ) From a8a07ba391b03622d0fe2248e1b5ef2170093585 Mon Sep 17 00:00:00 2001 From: Shuowei Li Date: Fri, 24 Jul 2026 00:06:42 +0000 Subject: [PATCH 3/4] test: skip anywidget test when dependencies missing --- packages/bigframes/tests/unit/session/test_read_gbq_colab.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/packages/bigframes/tests/unit/session/test_read_gbq_colab.py b/packages/bigframes/tests/unit/session/test_read_gbq_colab.py index c2e696927434..0916d2d918b5 100644 --- a/packages/bigframes/tests/unit/session/test_read_gbq_colab.py +++ b/packages/bigframes/tests/unit/session/test_read_gbq_colab.py @@ -49,6 +49,9 @@ def test_read_gbq_colab_includes_label(): 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 From 9185cdd1836aaa3e94b0d0c4e357aba4abd6de9c Mon Sep 17 00:00:00 2001 From: Shuowei Li Date: Fri, 24 Jul 2026 22:28:10 +0000 Subject: [PATCH 4/4] test(bigframes): verify read_gbq_colab label preservation in anywidget display mode --- .../bigframes/session/_io/bigquery/__init__.py | 10 ++-------- .../bigframes/tests/unit/session/test_io_bigquery.py | 3 ++- .../tests/unit/session/test_read_gbq_colab.py | 6 +----- 3 files changed, 5 insertions(+), 14 deletions(-) diff --git a/packages/bigframes/bigframes/session/_io/bigquery/__init__.py b/packages/bigframes/bigframes/session/_io/bigquery/__init__.py index a5b200296732..3d60bcc80744 100644 --- a/packages/bigframes/bigframes/session/_io/bigquery/__init__.py +++ b/packages/bigframes/bigframes/session/_io/bigquery/__init__.py @@ -69,14 +69,8 @@ def create_job_configs_labels( if api_methods and "bigframes-api" not in job_configs_labels: api_methods = list(api_methods) - 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] + job_configs_labels["bigframes-api"] = api_methods[0] + del api_methods[0] # 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 diff --git a/packages/bigframes/tests/unit/session/test_io_bigquery.py b/packages/bigframes/tests/unit/session/test_io_bigquery.py index 9f44a02ef49b..e6fa7a901eca 100644 --- a/packages/bigframes/tests/unit/session/test_io_bigquery.py +++ b/packages/bigframes/tests/unit/session/test_io_bigquery.py @@ -172,7 +172,7 @@ def test_add_and_trim_labels_length_limit_met(): for _ in range(52): df.head() - io_bq.add_and_trim_labels(job_config=job_config) + io_bq.add_and_trim_labels(job_config=job_config, session=df._session) assert job_config.labels is not None assert len(job_config.labels) == 56 assert "dataframe-max" not in job_config.labels.values() @@ -218,6 +218,7 @@ def test_start_query_with_job_labels_length_limit_met( timeout=timeout, metrics=None, publisher=bigframes.core.events.Publisher(), + session=df._session, ) assert job_config.labels is not None diff --git a/packages/bigframes/tests/unit/session/test_read_gbq_colab.py b/packages/bigframes/tests/unit/session/test_read_gbq_colab.py index 0916d2d918b5..fc4181b6a2b4 100644 --- a/packages/bigframes/tests/unit/session/test_read_gbq_colab.py +++ b/packages/bigframes/tests/unit/session/test_read_gbq_colab.py @@ -48,7 +48,7 @@ def test_read_gbq_colab_includes_label(): 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.""" + """Make sure read_gbq_colab label is preserved in recent-bigframes-api labels in anywidget mode.""" pytest.importorskip("anywidget") pytest.importorskip("traitlets") @@ -64,7 +64,6 @@ def test_read_gbq_colab_includes_label_in_anywidget_mode(): _ = 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, @@ -74,11 +73,8 @@ def test_read_gbq_colab_includes_label_in_anywidget_mode(): 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])