refactor: Allow unambiguous windows even in unstrictly ordered sessions (#849)

Co-authored-by: Huan Chen <[email protected]>

Author: TrevorBergeron

bigframes/core/__init__.py

@@ -194,7 +194,7 @@ def promote_offsets(self, col_id: str) -> ArrayValue:
         """
         Convenience function to promote copy of column offsets to a value column. Can be used to reset index.
         """
-        if not self.session._strictly_ordered:
+        if self.node.order_ambiguous and not self.session._strictly_ordered:
             raise ValueError("Generating offsets not supported in unordered mode")
         return ArrayValue(nodes.PromoteOffsetsNode(child=self.node, col_id=col_id))
 
@@ -346,7 +346,7 @@ def project_window_op(
         """
         # TODO: Support non-deterministic windowing
         if window_spec.row_bounded or not op.order_independent:
-            if not self.session._strictly_ordered:
+            if self.node.order_ambiguous and not self.session._strictly_ordered:
                 raise ValueError(
                     "Order-dependent windowed ops not supported in unordered mode"
                 )

bigframes/core/nodes.py

@@ -127,6 +127,14 @@ def joins(self) -> bool:
         """
         return False
 
+    @property
+    @abc.abstractmethod
+    def order_ambiguous(self) -> bool:
+        """
+        Whether row ordering is potentially ambiguous. For example, ReadTable (without a primary key) could be ordered in different ways.
+        """
+        ...
+
     @functools.cached_property
     def total_variables(self) -> int:
         return self.variables_introduced + sum(
@@ -177,6 +185,10 @@ def transform_children(
     ) -> BigFrameNode:
         return replace(self, child=t(self.child))
 
+    @property
+    def order_ambiguous(self) -> bool:
+        return self.child.order_ambiguous
+
 
 @dataclass(frozen=True)
 class JoinNode(BigFrameNode):
@@ -196,6 +208,10 @@ def non_local(self) -> bool:
     def child_nodes(self) -> typing.Sequence[BigFrameNode]:
         return (self.left_child, self.right_child)
 
+    @property
+    def order_ambiguous(self) -> bool:
+        return True
+
     def __hash__(self):
         return self._node_hash
 
@@ -247,6 +263,10 @@ def __post_init__(self):
     def child_nodes(self) -> typing.Sequence[BigFrameNode]:
         return self.children
 
+    @property
+    def order_ambiguous(self) -> bool:
+        return any(child.order_ambiguous for child in self.children)
+
     def __hash__(self):
         return self._node_hash
 
@@ -293,6 +313,10 @@ def variables_introduced(self) -> int:
         """Defines the number of variables generated by the current node. Used to estimate query planning complexity."""
         return len(self.schema.items) + 1
 
+    @property
+    def order_ambiguous(self) -> bool:
+        return False
+
     def transform_children(
         self, t: Callable[[BigFrameNode], BigFrameNode]
     ) -> BigFrameNode:
@@ -350,6 +374,10 @@ def relation_ops_created(self) -> int:
         # Assume worst case, where readgbq actually has baked in analytic operation to generate index
         return 3
 
+    @property
+    def order_ambiguous(self) -> bool:
+        return len(self.total_order_cols) == 0
+
     @functools.cached_property
     def variables_introduced(self) -> int:
         return len(self.schema.items) + 1
@@ -417,6 +445,10 @@ def hidden_columns(self) -> typing.Tuple[str, ...]:
             if col not in self.schema.names
         )
 
+    @property
+    def order_ambiguous(self) -> bool:
+        return not isinstance(self.ordering, orderings.TotalOrdering)
+
     def transform_children(
         self, t: Callable[[BigFrameNode], BigFrameNode]
     ) -> BigFrameNode:
@@ -600,6 +632,10 @@ def schema(self) -> schemata.ArraySchema:
     def variables_introduced(self) -> int:
         return len(self.aggregations) + len(self.by_column_ids)
 
+    @property
+    def order_ambiguous(self) -> bool:
+        return False
+
 
 @dataclass(frozen=True)
 class WindowOpNode(UnaryNode):

tests/system/small/test_unordered.py

@@ -40,6 +40,13 @@ def test_unordered_mode_cache_aggregate(unordered_session):
     assert_pandas_df_equal(bf_result, pd_result, ignore_order=True)
 
 
+def test_unordered_mode_single_aggregate(unordered_session):
+    pd_df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}, dtype=pd.Int64Dtype())
+    bf_df = bpd.DataFrame(pd_df, session=unordered_session)
+
+    assert bf_df.a.mean() == pd_df.a.mean()
+
+
 def test_unordered_mode_print(unordered_session):
     pd_df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}, dtype=pd.Int64Dtype())
     df = bpd.DataFrame(pd_df, session=unordered_session).cache()