refactor: Block .head consistently in unordered mode (#854)

Author: TrevorBergeron

bigframes/constants.py

@@ -99,3 +99,5 @@
 
 # BigQuery default is 10000, leave 100 for overhead
 MAX_COLUMNS = 9900
+
+SUGGEST_PEEK_PREVIEW = "Use .peek(n) to preview n arbitrary rows."

bigframes/core/groupby/__init__.py

@@ -109,6 +109,7 @@ def __getitem__(
                 dropna=self._dropna,
             )
 
+    @validations.requires_strict_ordering()
     def head(self, n: int = 5) -> df.DataFrame:
         block = self._block
         if self._dropna:
@@ -531,6 +532,7 @@ def __init__(
     def _session(self) -> core.Session:
         return self._block.session
 
+    @validations.requires_strict_ordering()
     def head(self, n: int = 5) -> series.Series:
         block = self._block
         if self._dropna:

bigframes/core/validations.py

@@ -17,7 +17,7 @@
 from __future__ import annotations
 
 import functools
-from typing import Protocol, TYPE_CHECKING
+from typing import Optional, Protocol, TYPE_CHECKING
 
 import bigframes.constants
 import bigframes.exceptions
@@ -32,20 +32,23 @@ def _session(self) -> Session:
         ...
 
 
-def requires_strict_ordering():
+def requires_strict_ordering(suggestion: Optional[str] = None):
     def decorator(meth):
         @functools.wraps(meth)
         def guarded_meth(object: HasSession, *args, **kwargs):
-            enforce_ordered(object, meth.__name__)
+            enforce_ordered(object, meth.__name__, suggestion)
             return meth(object, *args, **kwargs)
 
         return guarded_meth
 
     return decorator
 
 
-def enforce_ordered(object: HasSession, opname: str) -> None:
+def enforce_ordered(
+    object: HasSession, opname: str, suggestion: Optional[str] = None
+) -> None:
     if not object._session._strictly_ordered:
+        suggestion_substr = suggestion + " " if suggestion else ""
         raise bigframes.exceptions.OrderRequiredError(
-            f"Op {opname} not supported when strict ordering is disabled. {bigframes.constants.FEEDBACK_LINK}"
+            f"Op {opname} not supported when strict ordering is disabled. {suggestion_substr}{bigframes.constants.FEEDBACK_LINK}"
         )

bigframes/dataframe.py

@@ -49,6 +49,7 @@
 
 import bigframes
 import bigframes._config.display_options as display_options
+import bigframes.constants
 import bigframes.constants as constants
 import bigframes.core
 from bigframes.core import log_adapter
@@ -1293,6 +1294,7 @@ def _compute_dry_run(self) -> bigquery.QueryJob:
     def copy(self) -> DataFrame:
         return DataFrame(self._block)
 
+    @validations.requires_strict_ordering(bigframes.constants.SUGGEST_PEEK_PREVIEW)
     def head(self, n: int = 5) -> DataFrame:
         return typing.cast(DataFrame, self.iloc[:n])
 

bigframes/series.py

@@ -633,7 +633,7 @@ def dropna(
             result = result.reset_index()
         return Series(result)
 
-    @validations.requires_strict_ordering()
+    @validations.requires_strict_ordering(bigframes.constants.SUGGEST_PEEK_PREVIEW)
     def head(self, n: int = 5) -> Series:
         return typing.cast(Series, self.iloc[0:n])
 

tests/system/small/test_unordered.py

@@ -131,6 +131,10 @@ def test_unordered_drop_duplicates(unordered_session, keep):
             lambda x: x.a.iloc[1::2],
             id="series_iloc",
         ),
+        pytest.param(
+            lambda x: x.head(3),
+            id="head",
+        ),
     ],
 )
 def test_unordered_mode_blocks_windowing(unordered_session, function):