refactor: Make row hash ordering explicit in tree (#1284)

TrevorBergeron · web-flow · commit cf852a2b8eb2 · 2025-01-15T16:06:15.000-08:00
diff --git a/bigframes/bigquery/_operations/search.py b/bigframes/bigquery/_operations/search.py
@@ -122,12 +122,12 @@ def vector_search(
         ...             base_table="bigframes-dev.bigframes_tests_sys.base_table",
         ...             column_to_search="my_embedding",
         ...             query=search_query,
-        ...             top_k=2)
+        ...             top_k=2).sort_values("id")
           query_id  embedding  id my_embedding  distance
-        1      cat  [3.  5.2]   5    [5.  5.4]  2.009975
         0      dog    [1. 2.]   1      [1. 2.]       0.0
-        0      dog    [1. 2.]   4    [1.  3.2]       1.2
         1      cat  [3.  5.2]   2      [2. 4.]   1.56205
+        0      dog    [1. 2.]   4    [1.  3.2]       1.2
+        1      cat  [3.  5.2]   5    [5.  5.4]  2.009975
         <BLANKLINE>
         [4 rows x 5 columns]
 
@@ -141,12 +141,12 @@ def vector_search(
         ...             column_to_search="my_embedding",
         ...             query=search_query,
         ...             top_k=2,
-        ...             use_brute_force=True)
+        ...             use_brute_force=True).sort_values("id")
              embedding  id my_embedding  distance
         dog    [1. 2.]   1      [1. 2.]       0.0
-        cat  [3.  5.2]   5    [5.  5.4]  2.009975
-        dog    [1. 2.]   4    [1.  3.2]       1.2
         cat  [3.  5.2]   2      [2. 4.]   1.56205
+        dog    [1. 2.]   4    [1.  3.2]       1.2
+        cat  [3.  5.2]   5    [5.  5.4]  2.009975
         <BLANKLINE>
         [4 rows x 4 columns]
 
diff --git a/bigframes/core/__init__.py b/bigframes/core/__init__.py
@@ -220,8 +220,14 @@ def filter_by_id(self, predicate_id: str, keep_null: bool = False) -> ArrayValue
     def filter(self, predicate: ex.Expression):
         return ArrayValue(nodes.FilterNode(child=self.node, predicate=predicate))
 
-    def order_by(self, by: Sequence[OrderingExpression]) -> ArrayValue:
-        return ArrayValue(nodes.OrderByNode(child=self.node, by=tuple(by)))
+    def order_by(
+        self, by: Sequence[OrderingExpression], is_total_order: bool = False
+    ) -> ArrayValue:
+        return ArrayValue(
+            nodes.OrderByNode(
+                child=self.node, by=tuple(by), is_total_order=is_total_order
+            )
+        )
 
     def reversed(self) -> ArrayValue:
         return ArrayValue(nodes.ReversedNode(child=self.node))
diff --git a/bigframes/core/compile/compiled.py b/bigframes/core/compile/compiled.py
@@ -278,6 +278,24 @@ def to_sql(
         sql = ibis_bigquery.Backend().compile(self._to_ibis_expr())
         return typing.cast(str, sql)
 
+    def with_total_order(self, by: Sequence[OrderingExpression]) -> OrderedIR:
+        return OrderedIR(
+            table=self._table,
+            columns=self._columns,
+            predicates=self._predicates,
+            ordering=TotalOrdering(
+                ordering_value_columns=tuple(by),
+                total_ordering_columns=frozenset(
+                    map(
+                        ex.DerefOp,
+                        itertools.chain.from_iterable(
+                            col.referenced_columns for col in by
+                        ),
+                    )
+                ),
+            ),
+        )
+
     def row_count(self, name: str) -> OrderedIR:
         original_table = self._to_ibis_expr()
         ibis_table = original_table.agg(
@@ -576,6 +594,13 @@ def __init__(
     def is_ordered_ir(self) -> bool:
         return True
 
+    @property
+    def order_non_deterministic(self) -> bool:
+        # ordering suffix non-determinism is ok, as rand() is used as suffix for auto-generated order keys.
+        # but must be resolved before or explode, otherwise the engine might pull the rand() evaluation above the join,
+        # creating inconsistencies
+        return not all(col.deterministic for col in self._ordering.all_ordering_columns)
+
     @property
     def has_total_order(self) -> bool:
         return isinstance(self._ordering, TotalOrdering)
@@ -722,6 +747,9 @@ def _uniform_sampling(self, fraction: float) -> OrderedIR:
         )
 
     def explode(self, columns: typing.Sequence[ex.DerefOp]) -> OrderedIR:
+        if self.order_non_deterministic:
+            id = bigframes.core.guid.generate_guid()
+            return self.promote_offsets(id)
         table = self._to_ibis_expr(ordering_mode="unordered", expose_hidden_cols=True)
         column_ids = tuple(ref.id.sql for ref in columns)
 
@@ -1229,7 +1257,14 @@ def _bake_ordering(self) -> OrderedIR:
                 tuple(new_exprs),
                 self._ordering.integer_encoding,
                 self._ordering.string_encoding,
-                self._ordering.total_ordering_columns,
+                total_ordering_columns=frozenset(
+                    map(
+                        ex.DerefOp,
+                        itertools.chain.from_iterable(
+                            col.referenced_columns for col in new_exprs
+                        ),
+                    )
+                ),
             )
         else:
             new_ordering = RowOrdering(
diff --git a/bigframes/core/compile/compiler.py b/bigframes/core/compile/compiler.py
@@ -26,7 +26,6 @@
 
 import bigframes.core.compile.compiled as compiled
 import bigframes.core.compile.concat as concat_impl
-import bigframes.core.compile.default_ordering as default_ordering
 import bigframes.core.compile.ibis_types
 import bigframes.core.compile.scalar_op_compiler
 import bigframes.core.compile.scalar_op_compiler as compile_scalar
@@ -104,10 +103,7 @@ def set_output_names(
         )
 
     def compile_ordered_ir(self, node: nodes.BigFrameNode) -> compiled.OrderedIR:
-        ir = typing.cast(compiled.OrderedIR, self.compile_node(node, True))
-        if self.strict:
-            assert ir.has_total_order
-        return ir
+        return typing.cast(compiled.OrderedIR, self.compile_node(node, True))
 
     def compile_unordered_ir(self, node: nodes.BigFrameNode) -> compiled.UnorderedIR:
         return typing.cast(compiled.UnorderedIR, self.compile_node(node, False))
@@ -274,17 +270,6 @@ def compile_read_table_ordered(
                 for source_id, out_id in full_mapping.items()
                 if source_id not in visible_column_mapping
             )
-        elif self.strict:  # In strict mode, we fallback to ordering by row hash
-            order_values = [
-                col.name(guids.generate_guid())
-                for col in default_ordering.gen_default_ordering(
-                    ibis_table, use_double_hash=True
-                )
-            ]
-            ordering = bf_ordering.TotalOrdering.from_primary_key(
-                [value.get_name() for value in order_values]
-            )
-            hidden_columns = tuple(order_values)
         else:
             # In unstrict mode, don't generate total ordering from hashing as this is
             # expensive (prevent removing any columns from table scan)
@@ -316,7 +301,11 @@ def compile_filter(self, node: nodes.FilterNode, ordered: bool = True):
     @_compile_node.register
     def compile_orderby(self, node: nodes.OrderByNode, ordered: bool = True):
         if ordered:
-            return self.compile_ordered_ir(node.child).order_by(node.by)
+            if node.is_total_order:
+                # more efficient, can just discard any previous ordering and get same result
+                return self.compile_unordered_ir(node.child).with_total_order(node.by)
+            else:
+                return self.compile_ordered_ir(node.child).order_by(node.by)
         else:
             return self.compile_unordered_ir(node.child)
 
diff --git a/bigframes/core/compile/default_ordering.py b/bigframes/core/compile/default_ordering.py
@@ -18,7 +18,7 @@
 
 from __future__ import annotations
 
-from typing import cast
+from typing import cast, Sequence
 
 import bigframes_vendored.ibis
 import bigframes_vendored.ibis.expr.datatypes as ibis_dtypes
@@ -28,7 +28,7 @@
 import bigframes.core.guid as guid
 
 
-def _convert_to_nonnull_string(column: ibis_types.Column) -> ibis_types.StringValue:
+def _convert_to_nonnull_string(column: ibis_types.Value) -> ibis_types.StringValue:
     col_type = column.type()
     if (
         col_type.is_numeric()
@@ -60,29 +60,35 @@ def _convert_to_nonnull_string(column: ibis_types.Column) -> ibis_types.StringVa
     )
 
 
-def gen_default_ordering(
-    table: ibis_types.Table, use_double_hash: bool = True
-) -> list[bigframes_vendored.ibis.Value]:
+def gen_row_key(
+    columns: Sequence[ibis_types.Value],
+) -> bigframes_vendored.ibis.Value:
     ordering_hash_part = guid.generate_guid("bigframes_ordering_")
     ordering_hash_part2 = guid.generate_guid("bigframes_ordering_")
     ordering_rand_part = guid.generate_guid("bigframes_ordering_")
 
     # All inputs into hash must be non-null or resulting hash will be null
-    str_values = list(
-        map(lambda col: _convert_to_nonnull_string(table[col]), table.columns)
-    )
+    str_values = list(map(_convert_to_nonnull_string, columns))
     full_row_str = (
         str_values[0].concat(*str_values[1:]) if len(str_values) > 1 else str_values[0]
     )
-    full_row_hash = full_row_str.hash().name(ordering_hash_part)
+    full_row_hash = (
+        full_row_str.hash()
+        .name(ordering_hash_part)
+        .cast(ibis_dtypes.String(nullable=True))
+    )
     # By modifying value slightly, we get another hash uncorrelated with the first
-    full_row_hash_p2 = (full_row_str + "_").hash().name(ordering_hash_part2)
+    full_row_hash_p2 = (
+        (full_row_str + "_")
+        .hash()
+        .name(ordering_hash_part2)
+        .cast(ibis_dtypes.String(nullable=True))
+    )
     # Used to disambiguate between identical rows (which will have identical hash)
-    random_value = bigframes_vendored.ibis.random().name(ordering_rand_part)
-
-    order_values = (
-        [full_row_hash, full_row_hash_p2, random_value]
-        if use_double_hash
-        else [full_row_hash, random_value]
+    random_value = (
+        bigframes_vendored.ibis.random()
+        .name(ordering_rand_part)
+        .cast(ibis_dtypes.String(nullable=True))
     )
-    return order_values
+
+    return full_row_hash.concat(full_row_hash_p2, random_value)
diff --git a/bigframes/core/compile/scalar_op_compiler.py b/bigframes/core/compile/scalar_op_compiler.py
@@ -26,6 +26,7 @@
 import numpy as np
 import pandas as pd
 
+import bigframes.core.compile.default_ordering
 import bigframes.core.compile.ibis_types
 import bigframes.core.expression as ex
 import bigframes.dtypes
@@ -1850,6 +1851,11 @@ def struct_op_impl(
     return ibis_types.struct(data)
 
 
+@scalar_op_compiler.register_nary_op(ops.RowKey, pass_op=True)
+def rowkey_op_impl(*values: ibis_types.Value, op: ops.RowKey) -> ibis_types.Value:
+    return bigframes.core.compile.default_ordering.gen_row_key(values)
+
+
 # Helpers
 def is_null(value) -> bool:
     # float NaN/inf should be treated as distinct from 'true' null values
diff --git a/bigframes/core/compile/single_column.py b/bigframes/core/compile/single_column.py
@@ -50,6 +50,12 @@ def join_by_column_ordered(
         first the coalesced join keys, then, all the left columns, and
         finally, all the right columns.
     """
+    if type == "right":
+        if left.order_non_deterministic:
+            right = right._bake_ordering()
+    else:
+        if left.order_non_deterministic:
+            left = left._bake_ordering()
 
     # Do not reset the generator
     l_value_mapping = dict(zip(left.column_ids, left.column_ids))
diff --git a/bigframes/core/expression.py b/bigframes/core/expression.py
@@ -205,6 +205,10 @@ def bind_variables(
     def is_bijective(self) -> bool:
         return False
 
+    @property
+    def deterministic(self) -> bool:
+        return True
+
     @property
     def is_identity(self) -> bool:
         """True for identity operation that does not transform input."""
@@ -409,4 +413,10 @@ def bind_refs(
     @property
     def is_bijective(self) -> bool:
         # TODO: Mark individual functions as bijective?
-        return False
+        return all(input.is_bijective for input in self.inputs) and self.op.is_bijective
+
+    @property
+    def deterministic(self) -> bool:
+        return (
+            all(input.deterministic for input in self.inputs) and self.op.deterministic
+        )
diff --git a/bigframes/core/nodes.py b/bigframes/core/nodes.py
@@ -976,6 +976,9 @@ def remap_refs(self, mappings: Mapping[bfet_ids.ColumnId, bfet_ids.ColumnId]):
 @dataclasses.dataclass(frozen=True, eq=False)
 class OrderByNode(UnaryNode):
     by: Tuple[OrderingExpression, ...]
+    # This is an optimization, if true, can discard previous orderings.
+    # might be a total ordering even if false
+    is_total_order: bool = False
 
     @property
     def variables_introduced(self) -> int:
diff --git a/bigframes/core/ordering.py b/bigframes/core/ordering.py
@@ -59,6 +59,10 @@ class OrderingExpression:
     def referenced_columns(self) -> Set[ids.ColumnId]:
         return set(self.scalar_expression.column_references)
 
+    @property
+    def deterministic(self) -> bool:
+        return self.scalar_expression.deterministic
+
     def remap_column_refs(
         self,
         mapping: Mapping[ids.ColumnId, ids.ColumnId],
diff --git a/bigframes/operations/__init__.py b/bigframes/operations/__init__.py
@@ -25,6 +25,7 @@
 from pandas.tseries.offsets import DateOffset
 import pyarrow as pa
 
+import bigframes.dtypes
 import bigframes.dtypes as dtypes
 import bigframes.operations.type as op_typing
 
@@ -42,10 +43,20 @@ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionT
         ...
 
     @property
-    def order_preserving(self) -> bool:
+    def is_monotonic(self) -> bool:
         """Whether the row operation preserves total ordering. Can be pruned from ordering expressions."""
         ...
 
+    @property
+    def is_bijective(self) -> bool:
+        """Whether the operation has a 1:1 mapping between inputs and outputs"""
+        ...
+
+    @property
+    def deterministic(self) -> bool:
+        """Whether the operation is deterministic" (given deterministic inputs)"""
+        ...
+
 
 @dataclasses.dataclass(frozen=True)
 class ScalarOp:
@@ -57,10 +68,20 @@ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionT
         raise NotImplementedError("Abstract operation has no output type")
 
     @property
-    def order_preserving(self) -> bool:
+    def is_monotonic(self) -> bool:
         """Whether the row operation preserves total ordering. Can be pruned from ordering expressions."""
         return False
 
+    @property
+    def is_bijective(self) -> bool:
+        """Whether the operation has a 1:1 mapping between inputs and outputs"""
+        return False
+
+    @property
+    def deterministic(self) -> bool:
+        """Whether the operation is deterministic" (given deterministic inputs)"""
+        return True
+
 
 @dataclasses.dataclass(frozen=True)
 class NaryOp(ScalarOp):
@@ -1030,6 +1051,24 @@ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionT
         )  # [(name1, value1), (name2, value2), ...]
 
 
+# Really doesn't need to be its own op, but allows us to try to get the most compact representation
+@dataclasses.dataclass(frozen=True)
+class RowKey(NaryOp):
+    name: typing.ClassVar[str] = "rowkey"
+
+    def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
+        return bigframes.dtypes.STRING_DTYPE
+
+    @property
+    def is_bijective(self) -> bool:
+        """Whether the operation has a 1:1 mapping between inputs and outputs"""
+        return True
+
+    @property
+    def deterministic(self) -> bool:
+        return False
+
+
 # Just parameterless unary ops for now
 # TODO: Parameter mappings
 NUMPY_TO_OP: dict[np.ufunc, UnaryOp] = {
diff --git a/bigframes/session/__init__.py b/bigframes/session/__init__.py
@@ -269,6 +269,7 @@ def __init__(
             storage_manager=self._temp_storage_manager,
             default_index_type=self._default_index_type,
             scan_index_uniqueness=self._strictly_ordered,
+            force_total_order=self._strictly_ordered,
             metrics=self._metrics,
         )
 
diff --git a/bigframes/session/loader.py b/bigframes/session/loader.py
diff --git a/tests/system/small/bigquery/test_vector_search.py b/tests/system/small/bigquery/test_vector_search.py
diff --git a/tests/system/small/test_null_index.py b/tests/system/small/test_null_index.py
diff --git a/tests/system/small/test_session.py b/tests/system/small/test_session.py
