diff --git a/bigframes/_config/bigquery_options.py b/bigframes/_config/bigquery_options.py index ad79543cb8..0506f1841e 100644 --- a/bigframes/_config/bigquery_options.py +++ b/bigframes/_config/bigquery_options.py @@ -16,7 +16,8 @@ from __future__ import annotations -from typing import Optional +from enum import Enum +from typing import Literal, Optional import warnings import google.api_core.exceptions @@ -26,6 +27,12 @@ import bigframes.constants import bigframes.exceptions + +class OrderingMode(Enum): + STRICT = "strict" + PARTIAL = "partial" + + SESSION_STARTED_MESSAGE = ( "Cannot change '{attribute}' once a session has started. " "Call bigframes.pandas.close_session() first, if you are using the bigframes.pandas API." @@ -57,6 +64,14 @@ def _validate_location(value: Optional[str]): ) +def _validate_ordering_mode(value: str) -> OrderingMode: + if value.casefold() == OrderingMode.STRICT.value.casefold(): + return OrderingMode.STRICT + if value.casefold() == OrderingMode.PARTIAL.value.casefold(): + return OrderingMode.PARTIAL + raise ValueError("Ordering mode must be one of 'strict' or 'partial'.") + + class BigQueryOptions: """Encapsulates configuration for working with a session.""" @@ -71,7 +86,7 @@ def __init__( kms_key_name: Optional[str] = None, skip_bq_connection_check: bool = False, *, - _strictly_ordered: bool = True, + ordering_mode: Literal["strict", "partial"] = "strict", ): self._credentials = credentials self._project = project @@ -82,8 +97,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 + # Determines the ordering strictness for the session. + self._ordering_mode = _validate_ordering_mode(ordering_mode) @property def application_name(self) -> Optional[str]: @@ -241,6 +256,10 @@ def kms_key_name(self, value: str): 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 + def ordering_mode(self) -> Literal["strict", "partial"]: + """Controls whether total row order is always maintained for DataFrame/Series.""" + return self._ordering_mode.value + + @ordering_mode.setter + def ordering_mode(self, ordering_mode: Literal["strict", "partial"]) -> None: + self._ordering_mode = _validate_ordering_mode(ordering_mode) diff --git a/bigframes/session/__init__.py b/bigframes/session/__init__.py index 9c953ee594..22ca63d25b 100644 --- a/bigframes/session/__init__.py +++ b/bigframes/session/__init__.py @@ -297,15 +297,21 @@ 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: bool = context._strictly_ordered + self._strictly_ordered: bool = context.ordering_mode != "partial" + if not self._strictly_ordered: + warnings.warn( + "Partial ordering mode is a preview feature and is subject to change.", + bigframes.exceptions.PreviewWarning, + ) + # 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 + if self._strictly_ordered else bigframes.enums.DefaultIndexKind.NULL ) self._compiler = bigframes.core.compile.SQLCompiler( - strict=context._strictly_ordered + strict=self._strictly_ordered ) self._remote_function_session = bigframes_rf._RemoteFunctionSession() diff --git a/tests/system/conftest.py b/tests/system/conftest.py index 59439c306f..55079380f4 100644 --- a/tests/system/conftest.py +++ b/tests/system/conftest.py @@ -141,9 +141,7 @@ def session() -> Generator[bigframes.Session, None, None]: @pytest.fixture(scope="session", params=["ordered", "unordered"]) def maybe_ordered_session(request) -> Generator[bigframes.Session, None, None]: - context = bigframes.BigQueryOptions( - location="US", _strictly_ordered=request.param == "ordered" - ) + context = bigframes.BigQueryOptions(location="US", ordering_mode="partial") session = bigframes.Session(context=context) yield session session.close() # close generated session at cleanup type @@ -151,7 +149,7 @@ def maybe_ordered_session(request) -> Generator[bigframes.Session, None, None]: @pytest.fixture(scope="session") def unordered_session() -> Generator[bigframes.Session, None, None]: - context = bigframes.BigQueryOptions(location="US", _strictly_ordered=False) + context = bigframes.BigQueryOptions(location="US", ordering_mode="partial") session = bigframes.Session(context=context) yield session session.close() # close generated session at cleanup type