chore: create private _read_gbq_colab which will enable partial ordering mode even when session is in strict mode (#1688)

* chore: add private _read_gbq_colab method that uses partial ordering
mode, disables progress bars, disables default index, and communicates
via callbacks

* add colab read gbq

* add test for ordering

* add ordered argument to to_pandas_batches

* add unit test looking for job labels

* remove ordered option for to_pandas_batches

* ignore type for mock job configs

Author: tswast

bigframes/session/__init__.py

@@ -476,6 +476,27 @@ def _register_object(
     ):
         self._objects.append(weakref.ref(object))
 
+    def _read_gbq_colab(
+        self,
+        query: str,
+        # TODO: Add a callback parameter that takes some kind of Event object.
+        # TODO: Add parameter for variables for string formatting.
+        # TODO: Add dry_run parameter.
+    ) -> dataframe.DataFrame:
+        """A version of read_gbq that has the necessary default values for use in colab integrations.
+
+        This includes, no ordering, no index, no progress bar, always use string
+        formatting for embedding local variables / dataframes.
+        """
+
+        # TODO: Allow for a table ID to avoid queries like read_gbq?
+        return self._loader.read_gbq_query(
+            query=query,
+            index_col=bigframes.enums.DefaultIndexKind.NULL,
+            api_name="read_gbq_colab",
+            force_total_order=False,
+        )
+
     @overload
     def read_gbq_query(  # type: ignore[overload-overlap]
         self,

bigframes/session/loader.py

@@ -373,6 +373,7 @@ def read_gbq_table(  # type: ignore[overload-overlap]
         filters: third_party_pandas_gbq.FiltersType = ...,
         enable_snapshot: bool = ...,
         dry_run: Literal[False] = ...,
+        force_total_order: Optional[bool] = ...,
     ) -> dataframe.DataFrame:
         ...
 
@@ -394,6 +395,7 @@ def read_gbq_table(
         filters: third_party_pandas_gbq.FiltersType = ...,
         enable_snapshot: bool = ...,
         dry_run: Literal[True] = ...,
+        force_total_order: Optional[bool] = ...,
     ) -> pandas.Series:
         ...
 
@@ -414,6 +416,7 @@ def read_gbq_table(
         filters: third_party_pandas_gbq.FiltersType = (),
         enable_snapshot: bool = True,
         dry_run: bool = False,
+        force_total_order: Optional[bool] = None,
     ) -> dataframe.DataFrame | pandas.Series:
         import bigframes._tools.strings
         import bigframes.dataframe as dataframe
@@ -608,7 +611,14 @@ def read_gbq_table(
             session=self._session,
         )
         # if we don't have a unique index, we order by row hash if we are in strict mode
-        if self._force_total_order:
+        if (
+            # If the user has explicitly selected or disabled total ordering for
+            # this API call, respect that choice.
+            (force_total_order is not None and force_total_order)
+            # If the user has not explicitly selected or disabled total ordering
+            # for this API call, respect the default choice for the session.
+            or (force_total_order is None and self._force_total_order)
+        ):
             if not primary_key:
                 array_value = array_value.order_by(
                     [
@@ -712,6 +722,7 @@ def read_gbq_query(  # type: ignore[overload-overlap]
         use_cache: Optional[bool] = ...,
         filters: third_party_pandas_gbq.FiltersType = ...,
         dry_run: Literal[False] = ...,
+        force_total_order: Optional[bool] = ...,
     ) -> dataframe.DataFrame:
         ...
 
@@ -728,6 +739,7 @@ def read_gbq_query(
         use_cache: Optional[bool] = ...,
         filters: third_party_pandas_gbq.FiltersType = ...,
         dry_run: Literal[True] = ...,
+        force_total_order: Optional[bool] = ...,
     ) -> pandas.Series:
         ...
 
@@ -743,6 +755,7 @@ def read_gbq_query(
         use_cache: Optional[bool] = None,
         filters: third_party_pandas_gbq.FiltersType = (),
         dry_run: bool = False,
+        force_total_order: Optional[bool] = None,
     ) -> dataframe.DataFrame | pandas.Series:
         import bigframes.dataframe as dataframe
 
@@ -833,6 +846,7 @@ def read_gbq_query(
             columns=columns,
             use_cache=configuration["query"]["useQueryCache"],
             api_name=api_name,
+            force_total_order=force_total_order,
             # max_results and filters are omitted because they are already
             # handled by to_query(), above.
         )

bigframes/testing/mocks.py

@@ -51,6 +51,9 @@ def create_bigquery_session(
         google.auth.credentials.Credentials, instance=True
     )
 
+    bq_time = datetime.datetime.now()
+    table_time = bq_time + datetime.timedelta(minutes=1)
+
     if anonymous_dataset is None:
         anonymous_dataset = google.cloud.bigquery.DatasetReference(
             "test-project",
@@ -65,6 +68,8 @@ def create_bigquery_session(
         # Mock the location.
         table = mock.create_autospec(google.cloud.bigquery.Table, instance=True)
         table._properties = {}
+        # TODO(tswast): support tables created before and after the session started.
+        type(table).created = mock.PropertyMock(return_value=table_time)
         type(table).location = mock.PropertyMock(return_value=location)
         type(table).schema = mock.PropertyMock(return_value=table_schema)
         type(table).reference = mock.PropertyMock(
@@ -73,7 +78,10 @@ def create_bigquery_session(
         type(table).num_rows = mock.PropertyMock(return_value=1000000000)
         bqclient.get_table.return_value = table
 
-    def query_mock(query, *args, **kwargs):
+    job_configs = []
+
+    def query_mock(query, *args, job_config=None, **kwargs):
+        job_configs.append(job_config)
         query_job = mock.create_autospec(google.cloud.bigquery.QueryJob)
         type(query_job).destination = mock.PropertyMock(
             return_value=anonymous_dataset.table("test_table"),
@@ -83,15 +91,16 @@ def query_mock(query, *args, **kwargs):
         )
 
         if query.startswith("SELECT CURRENT_TIMESTAMP()"):
-            query_job.result = mock.MagicMock(return_value=[[datetime.datetime.now()]])
+            query_job.result = mock.MagicMock(return_value=[[bq_time]])
         else:
             type(query_job).schema = mock.PropertyMock(return_value=table_schema)
 
         return query_job
 
     existing_query_and_wait = bqclient.query_and_wait
 
-    def query_and_wait_mock(query, *args, **kwargs):
+    def query_and_wait_mock(query, *args, job_config=None, **kwargs):
+        job_configs.append(job_config)
         if query.startswith("SELECT CURRENT_TIMESTAMP()"):
             return iter([[datetime.datetime.now()]])
         else:
@@ -109,6 +118,7 @@ def query_and_wait_mock(query, *args, **kwargs):
     session._bq_connection_manager = mock.create_autospec(
         bigframes.clients.BqConnectionManager, instance=True
     )
+    session._job_configs = job_configs  # type: ignore
     return session
 
 

tests/system/small/session/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

tests/system/small/session/test_read_gbq_colab.py

@@ -0,0 +1,41 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""System tests for read_gbq_colab helper functions."""
+
+
+def test_read_gbq_colab_to_pandas_batches_preserves_order_by(maybe_ordered_session):
+    df = maybe_ordered_session._read_gbq_colab(
+        """
+        SELECT
+            name,
+            SUM(number) AS total
+        FROM
+            `bigquery-public-data.usa_names.usa_1910_2013`
+        WHERE state LIKE 'W%'
+        GROUP BY name
+        ORDER BY total DESC
+        LIMIT 300
+        """
+    )
+    batches = df.to_pandas_batches(
+        page_size=100,
+    )
+
+    total_rows = 0
+    for batch in batches:
+        assert batch["total"].is_monotonic_decreasing
+        total_rows += len(batch.index)
+
+    assert total_rows > 0

tests/unit/session/test_read_gbq_colab.py

@@ -0,0 +1,32 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit tests for read_gbq_colab helper functions."""
+
+from bigframes.testing import mocks
+
+
+def test_read_gbq_colab_includes_label():
+    """Make sure we can tell direct colab usage apart from regular read_gbq usage."""
+    session = mocks.create_bigquery_session()
+    _ = session._read_gbq_colab("SELECT 'read-gbq-colab-test'")
+    configs = session._job_configs  # type: ignore
+
+    label_values = []
+    for config in configs:
+        if config is None:
+            continue
+        label_values.extend(config.labels.values())
+
+    assert "read_gbq_colab" in label_values