refactor: Create ordering base class to represent partial ordering (#814)

Author: TrevorBergeron

bigframes/core/__init__.py

@@ -78,7 +78,7 @@ def from_cached(
         cls,
         original: ArrayValue,
         table: google.cloud.bigquery.Table,
-        ordering: orderings.ExpressionOrdering,
+        ordering: orderings.TotalOrdering,
     ):
         node = nodes.CachedTableNode(
             original_node=original.node,
@@ -147,7 +147,7 @@ def _compiled_schema(self) -> schemata.ArraySchema:
     def as_cached(
         self: ArrayValue,
         cache_table: google.cloud.bigquery.Table,
-        ordering: Optional[orderings.ExpressionOrdering],
+        ordering: Optional[orderings.TotalOrdering],
     ) -> ArrayValue:
         """
         Replace the node with an equivalent one that references a tabel where the value has been materialized to.

bigframes/core/compile/api.py

@@ -46,7 +46,7 @@ def compile_ordered(
 
 def compile_raw(
     node: bigframes.core.nodes.BigFrameNode,
-) -> Tuple[str, bigframes.core.ordering.ExpressionOrdering]:
+) -> Tuple[str, bigframes.core.ordering.TotalOrdering]:
     """Compile node into sql that exposes all columns, including hidden ordering-only columns."""
     ir = compiler.compile_ordered_ir(node)
     sql = ir.raw_sql()

bigframes/core/compile/compiled.py

@@ -36,9 +36,9 @@
 from bigframes.core.ordering import (
     ascending_over,
     encode_order_string,
-    ExpressionOrdering,
     IntegerEncoding,
     OrderingExpression,
+    TotalOrdering,
 )
 import bigframes.core.schema as schemata
 import bigframes.core.sql
@@ -194,7 +194,7 @@ def _aggregate_base(
         if by_column_ids:
             result = table.group_by(by_column_ids).aggregate(**stats)
             # Must have deterministic ordering, so order by the unique "by" column
-            ordering = ExpressionOrdering(
+            ordering = TotalOrdering(
                 tuple([ascending_over(column_id) for column_id in by_column_ids]),
                 total_ordering_columns=frozenset(by_column_ids),
             )
@@ -210,7 +210,7 @@ def _aggregate_base(
             # Ordering is irrelevant for single-row output, but set ordering id regardless
             # as other ops(join etc.) expect it.
             # TODO: Maybe can make completely empty
-            ordering = ExpressionOrdering(
+            ordering = TotalOrdering(
                 ordering_value_columns=tuple([]),
                 total_ordering_columns=frozenset([]),
             )
@@ -279,7 +279,7 @@ def row_count(self) -> OrderedIR:
         return OrderedIR(
             ibis_table,
             (ibis_table["count"],),
-            ordering=ExpressionOrdering(
+            ordering=TotalOrdering(
                 ordering_value_columns=(ascending_over("count"),),
                 total_ordering_columns=frozenset(["count"]),
             ),
@@ -519,7 +519,7 @@ def __init__(
         table: ibis_types.Table,
         columns: Sequence[ibis_types.Value],
         hidden_ordering_columns: Optional[Sequence[ibis_types.Value]] = None,
-        ordering: ExpressionOrdering = ExpressionOrdering(),
+        ordering: TotalOrdering = TotalOrdering(),
         predicates: Optional[Collection[ibis_types.BooleanValue]] = None,
     ):
         super().__init__(table, columns, predicates)
@@ -598,10 +598,7 @@ def from_pandas(
         return cls(
             keys_memtable,
             columns=[keys_memtable[column].name(column) for column in pd_df.columns],
-            ordering=ExpressionOrdering(
-                ordering_value_columns=tuple([ascending_over(ORDER_ID_COLUMN)]),
-                total_ordering_columns=frozenset([ORDER_ID_COLUMN]),
-            ),
+            ordering=TotalOrdering.from_offset_col(ORDER_ID_COLUMN),
             hidden_ordering_columns=(keys_memtable[ORDER_ID_COLUMN],),
         )
 
@@ -760,7 +757,7 @@ def explode(self, column_ids: typing.Sequence[str]) -> OrderedIR:
             ],
             table_w_unnest[unnest_offset_id],
         ]
-        ordering = ExpressionOrdering(
+        ordering = TotalOrdering(
             ordering_value_columns=tuple(
                 [
                     *self._ordering.ordering_value_columns,
@@ -1153,7 +1150,7 @@ def _bake_ordering(self) -> OrderedIR:
                         self._ibis_bindings[expr.scalar_expression.id]
                     )
 
-        new_ordering = ExpressionOrdering(
+        new_ordering = TotalOrdering(
             tuple(new_exprs),
             self._ordering.integer_encoding,
             self._ordering.string_encoding,
@@ -1176,7 +1173,7 @@ def _project_offsets(self) -> OrderedIR:
             ordering_mode="offset_col", order_col_name=ORDER_ID_COLUMN
         )
         columns = [table[column_name] for column_name in self._column_names]
-        ordering = ExpressionOrdering(
+        ordering = TotalOrdering(
             ordering_value_columns=tuple([ascending_over(ORDER_ID_COLUMN)]),
             total_ordering_columns=frozenset([ORDER_ID_COLUMN]),
             integer_encoding=IntegerEncoding(True, is_sequential=True),
@@ -1300,7 +1297,7 @@ class Builder:
         def __init__(
             self,
             table: ibis_types.Table,
-            ordering: ExpressionOrdering,
+            ordering: TotalOrdering,
             columns: Collection[ibis_types.Value] = (),
             hidden_ordering_columns: Collection[ibis_types.Value] = (),
             predicates: Optional[Collection[ibis_types.BooleanValue]] = None,

bigframes/core/compile/compiler.py

@@ -195,7 +195,7 @@ def compile_read_table_ordered(node: nodes.ReadTableNode):
             )
         else:
             integer_encoding = bf_ordering.IntegerEncoding()
-        ordering = bf_ordering.ExpressionOrdering(
+        ordering = bf_ordering.TotalOrdering(
             ordering_value_columns,
             integer_encoding=integer_encoding,
             total_ordering_columns=frozenset(node.total_order_cols),

bigframes/core/compile/concat.py

@@ -21,9 +21,9 @@
 import bigframes.core.compile.compiled as compiled
 from bigframes.core.ordering import (
     ascending_over,
-    ExpressionOrdering,
     reencode_order_string,
     StringEncoding,
+    TotalOrdering,
 )
 
 ORDER_ID_COLUMN = "bigframes_ordering_id"
@@ -83,7 +83,7 @@ def concat_ordered(
         )
         tables.append(table)
     combined_table = ibis.union(*tables)
-    ordering = ExpressionOrdering(
+    ordering = TotalOrdering(
         ordering_value_columns=tuple([ascending_over(ORDER_ID_COLUMN)]),
         total_ordering_columns=frozenset([ORDER_ID_COLUMN]),
         string_encoding=StringEncoding(True, prefix_size + max_encoding_size),

bigframes/core/compile/default_ordering.py

@@ -82,7 +82,7 @@ def gen_default_ordering(table: ibis.table, use_double_hash: bool = True):
         itertools.chain(original_column_ids, order_values)
     )
 
-    ordering = order.ExpressionOrdering(
+    ordering = order.TotalOrdering(
         ordering_value_columns=tuple(
             order.ascending_over(col.get_name()) for col in order_values
         ),

bigframes/core/compile/single_column.py

@@ -16,8 +16,6 @@
 
 from __future__ import annotations
 
-from typing import Mapping
-
 import ibis
 import ibis.expr.datatypes as ibis_dtypes
 import ibis.expr.types as ibis_types
@@ -84,7 +82,7 @@ def join_by_column_ordered(
     )
 
     # Preserve ordering accross joins.
-    ordering = join_orderings(
+    ordering = orderings.join_orderings(
         left._ordering,
         right._ordering,
         l_mapping,
@@ -173,33 +171,3 @@ def value_to_join_key(value: ibis_types.Value):
     if not value.type().is_string():
         value = value.cast(ibis_dtypes.str)
     return value.fillna(ibis_types.literal("$NULL_SENTINEL$"))
-
-
-def join_orderings(
-    left: orderings.ExpressionOrdering,
-    right: orderings.ExpressionOrdering,
-    left_id_mapping: Mapping[str, str],
-    right_id_mapping: Mapping[str, str],
-    left_order_dominates: bool = True,
-) -> orderings.ExpressionOrdering:
-    left_ordering_refs = [
-        ref.remap_names(left_id_mapping) for ref in left.all_ordering_columns
-    ]
-    right_ordering_refs = [
-        ref.remap_names(right_id_mapping) for ref in right.all_ordering_columns
-    ]
-    if left_order_dominates:
-        joined_refs = [*left_ordering_refs, *right_ordering_refs]
-    else:
-        joined_refs = [*right_ordering_refs, *left_ordering_refs]
-
-    left_total_order_cols = frozenset(
-        [left_id_mapping[id] for id in left.total_ordering_columns]
-    )
-    right_total_order_cols = frozenset(
-        [right_id_mapping[id] for id in right.total_ordering_columns]
-    )
-    return orderings.ExpressionOrdering(
-        ordering_value_columns=tuple(joined_refs),
-        total_ordering_columns=left_total_order_cols | right_total_order_cols,
-    )

bigframes/core/nodes.py

@@ -373,7 +373,7 @@ class CachedTableNode(BigFrameNode):
     table_id: str = field()
     physical_schema: Tuple[bq.SchemaField, ...] = field()
 
-    ordering: typing.Optional[orderings.ExpressionOrdering] = field()
+    ordering: typing.Optional[orderings.TotalOrdering] = field()
 
     def __post_init__(self):
         # enforce invariants

bigframes/core/ordering.py

@@ -94,10 +94,42 @@ class IntegerEncoding:
 
 
 @dataclass(frozen=True)
-class ExpressionOrdering:
-    """Immutable object that holds information about the ordering of rows in a ArrayValue object."""
+class RowOrdering:
+    """Immutable object that holds information about the ordering of rows in a ArrayValue object. May not be unambiguous."""
 
     ordering_value_columns: typing.Tuple[OrderingExpression, ...] = ()
+
+    @property
+    def all_ordering_columns(self) -> Sequence[OrderingExpression]:
+        return list(self.ordering_value_columns)
+
+    @property
+    def referenced_columns(self) -> Set[str]:
+        return set(
+            col
+            for part in self.ordering_value_columns
+            for col in part.scalar_expression.unbound_variables
+        )
+
+    def with_reverse(self) -> RowOrdering:
+        """Reverses the ordering."""
+        return RowOrdering(
+            tuple([col.with_reverse() for col in self.ordering_value_columns]),
+        )
+
+    def with_column_remap(self, mapping: typing.Mapping[str, str]) -> RowOrdering:
+        new_value_columns = [
+            col.remap_names(mapping) for col in self.all_ordering_columns
+        ]
+        return TotalOrdering(
+            tuple(new_value_columns),
+        )
+
+
+@dataclass(frozen=True)
+class TotalOrdering(RowOrdering):
+    """Immutable object that holds information about the ordering of rows in a ArrayValue object. Guaranteed to be unambiguous."""
+
     integer_encoding: IntegerEncoding = IntegerEncoding(False)
     string_encoding: StringEncoding = StringEncoding(False)
     # A table has a total ordering defined by the identities of a set of 1 or more columns.
@@ -106,8 +138,8 @@ class ExpressionOrdering:
     total_ordering_columns: frozenset[str] = field(default_factory=frozenset)
 
     @classmethod
-    def from_offset_col(cls, col: str) -> ExpressionOrdering:
-        return ExpressionOrdering(
+    def from_offset_col(cls, col: str) -> TotalOrdering:
+        return TotalOrdering(
             (ascending_over(col),),
             integer_encoding=IntegerEncoding(True, is_sequential=True),
             total_ordering_columns=frozenset({col}),
@@ -119,7 +151,7 @@ def with_non_sequential(self):
         This is useful when filtering, but not sorting, an expression.
         """
         if self.integer_encoding.is_sequential:
-            return ExpressionOrdering(
+            return TotalOrdering(
                 self.ordering_value_columns,
                 integer_encoding=IntegerEncoding(
                     self.integer_encoding.is_encoded, is_sequential=False
@@ -132,7 +164,7 @@ def with_non_sequential(self):
     def with_ordering_columns(
         self,
         ordering_value_columns: Sequence[OrderingExpression] = (),
-    ) -> ExpressionOrdering:
+    ) -> TotalOrdering:
         """Creates a new ordering that reorders by the given columns.
 
         Args:
@@ -147,7 +179,7 @@ def with_ordering_columns(
         new_ordering = self._truncate_ordering(
             (*ordering_value_columns, *self.ordering_value_columns)
         )
-        return ExpressionOrdering(
+        return TotalOrdering(
             new_ordering,
             total_ordering_columns=self.total_ordering_columns,
         )
@@ -173,7 +205,7 @@ def _truncate_ordering(
 
     def with_reverse(self):
         """Reverses the ordering."""
-        return ExpressionOrdering(
+        return TotalOrdering(
             tuple([col.with_reverse() for col in self.ordering_value_columns]),
             total_ordering_columns=self.total_ordering_columns,
         )
@@ -185,7 +217,7 @@ def with_column_remap(self, mapping: typing.Mapping[str, str]):
         new_total_order = frozenset(
             mapping.get(col_id, col_id) for col_id in self.total_ordering_columns
         )
-        return ExpressionOrdering(
+        return TotalOrdering(
             tuple(new_value_columns),
             integer_encoding=self.integer_encoding,
             string_encoding=self.string_encoding,
@@ -211,18 +243,6 @@ def is_string_encoded(self) -> bool:
     def is_sequential(self) -> bool:
         return self.integer_encoding.is_encoded and self.integer_encoding.is_sequential
 
-    @property
-    def all_ordering_columns(self) -> Sequence[OrderingExpression]:
-        return list(self.ordering_value_columns)
-
-    @property
-    def referenced_columns(self) -> Set[str]:
-        return set(
-            col
-            for part in self.ordering_value_columns
-            for col in part.scalar_expression.unbound_variables
-        )
-
 
 def encode_order_string(
     order_id: ibis_types.IntegerColumn, length: int = DEFAULT_ORDERING_ID_LENGTH
@@ -257,3 +277,58 @@ def descending_over(id: str, nulls_last: bool = True) -> OrderingExpression:
     return OrderingExpression(
         expression.free_var(id), direction=OrderingDirection.DESC, na_last=nulls_last
     )
+
+
+@typing.overload
+def join_orderings(
+    left: TotalOrdering,
+    right: TotalOrdering,
+    left_id_mapping: Mapping[str, str],
+    right_id_mapping: Mapping[str, str],
+    left_order_dominates: bool = True,
+) -> TotalOrdering:
+    ...
+
+
+@typing.overload
+def join_orderings(
+    left: RowOrdering,
+    right: RowOrdering,
+    left_id_mapping: Mapping[str, str],
+    right_id_mapping: Mapping[str, str],
+    left_order_dominates: bool = True,
+) -> RowOrdering:
+    ...
+
+
+def join_orderings(
+    left: RowOrdering,
+    right: RowOrdering,
+    left_id_mapping: Mapping[str, str],
+    right_id_mapping: Mapping[str, str],
+    left_order_dominates: bool = True,
+) -> RowOrdering:
+    left_ordering_refs = [
+        ref.remap_names(left_id_mapping) for ref in left.all_ordering_columns
+    ]
+    right_ordering_refs = [
+        ref.remap_names(right_id_mapping) for ref in right.all_ordering_columns
+    ]
+    if left_order_dominates:
+        joined_refs = [*left_ordering_refs, *right_ordering_refs]
+    else:
+        joined_refs = [*right_ordering_refs, *left_ordering_refs]
+
+    if isinstance(left, TotalOrdering) and isinstance(right, TotalOrdering):
+        left_total_order_cols = frozenset(
+            [left_id_mapping[id] for id in left.total_ordering_columns]
+        )
+        right_total_order_cols = frozenset(
+            [right_id_mapping[id] for id in right.total_ordering_columns]
+        )
+        return TotalOrdering(
+            ordering_value_columns=tuple(joined_refs),
+            total_ordering_columns=left_total_order_cols | right_total_order_cols,
+        )
+    else:
+        return RowOrdering(tuple(joined_refs))