refactor: create internal session ordering mode flag (#772)

Author: TrevorBergeron

bigframes/_config/bigquery_options.py

@@ -70,6 +70,8 @@ def __init__(
         application_name: Optional[str] = None,
         kms_key_name: Optional[str] = None,
         skip_bq_connection_check: bool = False,
+        *,
+        _strictly_ordered: bool = True,
     ):
         self._credentials = credentials
         self._project = project
@@ -80,6 +82,8 @@ def __init__(
         self._kms_key_name = kms_key_name
         self._skip_bq_connection_check = skip_bq_connection_check
         self._session_started = False
+        # Determines the ordering strictness for the session. For internal use only.
+        self._strictly_ordered_internal = _strictly_ordered
 
     @property
     def application_name(self) -> Optional[str]:
@@ -235,3 +239,8 @@ def kms_key_name(self, value: str):
             raise ValueError(SESSION_STARTED_MESSAGE.format(attribute="kms_key_name"))
 
         self._kms_key_name = value
+
+    @property
+    def _strictly_ordered(self) -> bool:
+        """Internal use only. Controls whether total row order is always maintained for DataFrame/Series."""
+        return self._strictly_ordered_internal

bigframes/core/blocks.py

@@ -209,7 +209,7 @@ def shape(self) -> typing.Tuple[int, int]:
             except Exception:
                 pass
 
-        iter, _ = self.session._execute(row_count_expr, sorted=False)
+        iter, _ = self.session._execute(row_count_expr, ordered=False)
         row_count = next(iter)[0]
         return (row_count, len(self.value_columns))
 
@@ -518,7 +518,7 @@ def to_pandas_batches(
         dtypes = dict(zip(self.index_columns, self.index.dtypes))
         dtypes.update(zip(self.value_columns, self.dtypes))
         _, query_job = self.session._query_to_destination(
-            self.session._to_sql(self.expr, sorted=True),
+            self.session._to_sql(self.expr, ordered=self.session._strictly_ordered),
             list(self.index_columns),
             api_name="cached",
             do_clustering=False,
@@ -553,7 +553,7 @@ def _materialize_local(
         """Run query and download results as a pandas DataFrame. Return the total number of results as well."""
         # TODO(swast): Allow for dry run and timeout.
         _, query_job = self.session._query_to_destination(
-            self.session._to_sql(self.expr, sorted=materialize_options.ordered),
+            self.session._to_sql(self.expr, ordered=materialize_options.ordered),
             list(self.index_columns),
             api_name="cached",
             do_clustering=False,
@@ -1736,7 +1736,7 @@ def transpose(
         original_row_index = (
             original_row_index
             if original_row_index is not None
-            else self.index.to_pandas()
+            else self.index.to_pandas(ordered=True)
         )
         original_row_count = len(original_row_index)
         if original_row_count > bigframes.constants.MAX_COLUMNS:
@@ -2507,7 +2507,7 @@ def column_ids(self) -> Sequence[str]:
         """Column(s) to use as row labels."""
         return self._block._index_columns
 
-    def to_pandas(self) -> pd.Index:
+    def to_pandas(self, *, ordered: Optional[bool] = None) -> pd.Index:
         """Executes deferred operations and downloads the results."""
         if len(self.column_ids) == 0:
             raise bigframes.exceptions.NullIndexError(
@@ -2517,7 +2517,12 @@ def to_pandas(self) -> pd.Index:
         index_columns = list(self._block.index_columns)
         dtypes = dict(zip(index_columns, self.dtypes))
         expr = self._expr.select_columns(index_columns)
-        results, _ = self.session._execute(expr)
+        results, _ = self.session._execute(
+            expr,
+            ordered=ordered
+            if (ordered is not None)
+            else self.session._strictly_ordered,
+        )
         df = expr.session._rows_to_dataframe(results, dtypes)
         df = df.set_index(index_columns)
         index = df.index

bigframes/core/compile/api.py

@@ -40,7 +40,7 @@ def compile_ordered(
 ) -> str:
     """Compile node into sql where rows are sorted with ORDER BY."""
     return compiler.compile_ordered_ir(node).to_sql(
-        col_id_overrides=col_id_overrides, sorted=True
+        col_id_overrides=col_id_overrides, ordered=True
     )
 
 

bigframes/core/compile/compiled.py

@@ -257,9 +257,9 @@ def to_sql(
         self,
         offset_column: typing.Optional[str] = None,
         col_id_overrides: typing.Mapping[str, str] = {},
-        sorted: bool = False,
+        ordered: bool = False,
     ) -> str:
-        if offset_column or sorted:
+        if offset_column or ordered:
             raise ValueError("Cannot produce sorted sql in unordered mode")
         sql = ibis_bigquery.Backend().compile(
             self._to_ibis_expr(
@@ -890,9 +890,9 @@ def _reproject_to_table(self) -> OrderedIR:
     def to_sql(
         self,
         col_id_overrides: typing.Mapping[str, str] = {},
-        sorted: bool = False,
+        ordered: bool = False,
     ) -> str:
-        if sorted:
+        if ordered:
             # Need to bake ordering expressions into the selected column in order for our ordering clause builder to work.
             baked_ir = self._bake_ordering()
             sql = ibis_bigquery.Backend().compile(

bigframes/core/indexes/base.py

@@ -481,7 +481,9 @@ def to_pandas(self) -> pandas.Index:
             pandas.Index:
                 A pandas Index with all of the labels from this Index.
         """
-        return self._block.index.to_pandas()
+        return self._block.index.to_pandas(
+            ordered=self._block.session._strictly_ordered
+        )
 
     def to_numpy(self, dtype=None, **kwargs) -> np.ndarray:
         return self.to_pandas().to_numpy(dtype, **kwargs)

bigframes/dataframe.py

@@ -1189,7 +1189,7 @@ def to_pandas(
         sampling_method: Optional[str] = None,
         random_state: Optional[int] = None,
         *,
-        ordered: bool = True,
+        ordered: Optional[bool] = None,
     ) -> pandas.DataFrame:
         """Write DataFrame to pandas DataFrame.
 
@@ -1209,9 +1209,10 @@ def to_pandas(
                 The seed for the uniform downsampling algorithm. If provided, the uniform method may
                 take longer to execute and require more computation. If set to a value other than
                 None, this will supersede the global config.
-            ordered (bool, default True):
+            ordered (bool, default None):
                 Determines whether the resulting pandas dataframe will be deterministically ordered.
-                In some cases, unordered may result in a faster-executing query.
+                In some cases, unordered may result in a faster-executing query. If set to a value
+                other than None, will override Session default.
 
         Returns:
             pandas.DataFrame: A pandas DataFrame with all rows and columns of this DataFrame if the
@@ -1224,7 +1225,7 @@ def to_pandas(
             max_download_size=max_download_size,
             sampling_method=sampling_method,
             random_state=random_state,
-            ordered=ordered,
+            ordered=ordered if ordered is not None else self._session._strictly_ordered,
         )
         self._set_internal_query_job(query_job)
         return df.set_axis(self._block.column_labels, axis=1, copy=False)
@@ -3339,7 +3340,7 @@ def _run_io_query(
         _, query_job = session._execute(
             export_array,
             job_config=job_config,
-            sorted=False,
+            ordered=False,
             col_id_overrides=id_overrides,
         )
         self._set_internal_query_job(query_job)

bigframes/series.py

@@ -323,7 +323,7 @@ def to_pandas(
         sampling_method: Optional[str] = None,
         random_state: Optional[int] = None,
         *,
-        ordered: bool = True,
+        ordered: Optional[bool] = None,
     ) -> pandas.Series:
         """Writes Series to pandas Series.
 
@@ -343,9 +343,10 @@ def to_pandas(
                 The seed for the uniform downsampling algorithm. If provided, the uniform method may
                 take longer to execute and require more computation. If set to a value other than
                 None, this will supersede the global config.
-            ordered (bool, default True):
+            ordered (bool, default None):
                 Determines whether the resulting pandas series will be deterministically ordered.
-                In some cases, unordered may result in a faster-executing query.
+                In some cases, unordered may result in a faster-executing query. If set to a value
+                other than None, will override Session default.
 
 
         Returns:
@@ -357,7 +358,7 @@ def to_pandas(
             max_download_size=max_download_size,
             sampling_method=sampling_method,
             random_state=random_state,
-            ordered=ordered,
+            ordered=ordered if ordered is not None else self._session._strictly_ordered,
         )
         self._set_internal_query_job(query_job)
         series = df.squeeze(axis=1)

bigframes/session/__init__.py

@@ -297,7 +297,13 @@ def __init__(
         self._execution_count = 0
         # Whether this session treats objects as totally ordered.
         # Will expose as feature later, only False for internal testing
-        self._strictly_ordered = True
+        self._strictly_ordered: bool = context._strictly_ordered
+        # Sequential index needs total ordering to generate, so use null index with unstrict ordering.
+        self._default_index_type: bigframes.enums.DefaultIndexKind = (
+            bigframes.enums.DefaultIndexKind.SEQUENTIAL_INT64
+            if context._strictly_ordered
+            else bigframes.enums.DefaultIndexKind.NULL
+        )
 
     @property
     def bqclient(self):
@@ -882,11 +888,11 @@ def _read_gbq_table(
         # Create Default Sequential Index if still have no index
         # ----------------------------------------------------
 
-        # If no index columns provided or found, fall back to sequential index
+        # If no index columns provided or found, fall back to session default
         if (index_col != bigframes.enums.DefaultIndexKind.NULL) and len(
             index_cols
         ) == 0:
-            index_col = bigframes.enums.DefaultIndexKind.SEQUENTIAL_INT64
+            index_col = self._default_index_type
 
         index_names: Sequence[Hashable] = index_cols
         if index_col == bigframes.enums.DefaultIndexKind.SEQUENTIAL_INT64:
@@ -1935,14 +1941,14 @@ def _execute(
         array_value: core.ArrayValue,
         job_config: Optional[bigquery.job.QueryJobConfig] = None,
         *,
-        sorted: bool = True,
+        ordered: bool = True,
         dry_run=False,
         col_id_overrides: Mapping[str, str] = {},
     ) -> tuple[bigquery.table.RowIterator, bigquery.QueryJob]:
         if not dry_run:
             self._add_execution(1)
         sql = self._to_sql(
-            array_value, sorted=sorted, col_id_overrides=col_id_overrides
+            array_value, ordered=ordered, col_id_overrides=col_id_overrides
         )  # type:ignore
         if job_config is None:
             job_config = bigquery.QueryJobConfig(dry_run=dry_run)
@@ -1977,12 +1983,12 @@ def _to_sql(
         array_value: core.ArrayValue,
         offset_column: typing.Optional[str] = None,
         col_id_overrides: typing.Mapping[str, str] = {},
-        sorted: bool = False,
+        ordered: bool = False,
     ) -> str:
         if offset_column:
             array_value = array_value.promote_offsets(offset_column)
         node_w_cached = self._with_cached_executions(array_value.node)
-        if sorted:
+        if ordered:
             return bigframes.core.compile.compile_ordered(
                 node_w_cached, col_id_overrides=col_id_overrides
             )

tests/system/conftest.py

@@ -141,11 +141,8 @@ def session() -> Generator[bigframes.Session, None, None]:
 
 @pytest.fixture(scope="session")
 def unordered_session() -> Generator[bigframes.Session, None, None]:
-    context = bigframes.BigQueryOptions(
-        location="US",
-    )
+    context = bigframes.BigQueryOptions(location="US", _strictly_ordered=False)
     session = bigframes.Session(context=context)
-    session._strictly_ordered = False
     yield session
     session.close()  # close generated session at cleanup type
 

tests/system/small/test_unordered.py

@@ -12,9 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import pandas as pd
+import pyarrow as pa
 
 import bigframes.pandas as bpd
-from tests.system.utils import assert_pandas_df_equal
+from tests.system.utils import assert_pandas_df_equal, skip_legacy_pandas
 
 
 def test_unordered_mode_cache_aggregate(unordered_session):
@@ -26,3 +27,35 @@ def test_unordered_mode_cache_aggregate(unordered_session):
     pd_result = pd_df - pd_df.mean()
 
     assert_pandas_df_equal(bf_result, pd_result, ignore_order=True)
+
+
+@skip_legacy_pandas
+def test_unordered_mode_read_gbq(unordered_session):
+    df = unordered_session.read_gbq(
+        """SELECT
+        [1, 3, 2] AS array_column,
+        STRUCT(
+            "a" AS string_field,
+            1.2 AS float_field) AS struct_column"""
+    )
+    expected = pd.DataFrame(
+        {
+            "array_column": pd.Series(
+                [[1, 3, 2]],
+                dtype=(pd.ArrowDtype(pa.list_(pa.int64()))),
+            ),
+            "struct_column": pd.Series(
+                [{"string_field": "a", "float_field": 1.2}],
+                dtype=pd.ArrowDtype(
+                    pa.struct(
+                        [
+                            ("string_field", pa.string()),
+                            ("float_field", pa.float64()),
+                        ]
+                    )
+                ),
+            ),
+        }
+    )
+    # Don't need ignore_order as there is only 1 row
+    assert_pandas_df_equal(df.to_pandas(), expected)