refactor: SQL builder supports partial ordering (#838)

Author: TrevorBergeron

bigframes/core/compile/api.py

@@ -51,7 +51,8 @@ def compile_raw(
     ir = compiler.compile_ordered_ir(node)
     sql = ir.raw_sql()
     ordering_info = ir._ordering
-    return sql, ordering_info
+    assert ir.has_total_order
+    return sql, ordering_info  # type: ignore
 
 
 def test_only_try_evaluate(node: bigframes.core.nodes.BigFrameNode):

bigframes/core/compile/compiled.py

@@ -37,7 +37,9 @@
     ascending_over,
     encode_order_string,
     IntegerEncoding,
+    join_orderings,
     OrderingExpression,
+    RowOrdering,
     TotalOrdering,
 )
 import bigframes.core.schema as schemata
@@ -519,7 +521,7 @@ def __init__(
         table: ibis_types.Table,
         columns: Sequence[ibis_types.Value],
         hidden_ordering_columns: Optional[Sequence[ibis_types.Value]] = None,
-        ordering: TotalOrdering = TotalOrdering(),
+        ordering: RowOrdering = RowOrdering(),
         predicates: Optional[Collection[ibis_types.BooleanValue]] = None,
     ):
         super().__init__(table, columns, predicates)
@@ -566,6 +568,10 @@ def __init__(
     def is_ordered_ir(self) -> bool:
         return True
 
+    @property
+    def has_total_order(self) -> bool:
+        return isinstance(self._ordering, TotalOrdering)
+
     @classmethod
     def from_pandas(
         cls,
@@ -757,16 +763,13 @@ def explode(self, column_ids: typing.Sequence[str]) -> OrderedIR:
             ],
             table_w_unnest[unnest_offset_id],
         ]
-        ordering = TotalOrdering(
-            ordering_value_columns=tuple(
-                [
-                    *self._ordering.ordering_value_columns,
-                    ascending_over(unnest_offset_id),
-                ]
-            ),
-            total_ordering_columns=frozenset(
-                [*self._ordering.total_ordering_columns, unnest_offset_id]
-            ),
+        l_mappings = {id: id for id in self._ordering.referenced_columns}
+        r_mappings = {unnest_offset_id: unnest_offset_id}
+        ordering = join_orderings(
+            self._ordering,
+            TotalOrdering.from_offset_col(unnest_offset_id),
+            l_mappings,
+            r_mappings,
         )
 
         return OrderedIR(
@@ -1150,12 +1153,19 @@ def _bake_ordering(self) -> OrderedIR:
                         self._ibis_bindings[expr.scalar_expression.id]
                     )
 
-        new_ordering = TotalOrdering(
-            tuple(new_exprs),
-            self._ordering.integer_encoding,
-            self._ordering.string_encoding,
-            self._ordering.total_ordering_columns,
-        )
+        if isinstance(self._ordering, TotalOrdering):
+            new_ordering: RowOrdering = TotalOrdering(
+                tuple(new_exprs),
+                self._ordering.integer_encoding,
+                self._ordering.string_encoding,
+                self._ordering.total_ordering_columns,
+            )
+        else:
+            new_ordering = RowOrdering(
+                tuple(new_exprs),
+                self._ordering.integer_encoding,
+                self._ordering.string_encoding,
+            )
         return OrderedIR(
             self._table,
             columns=self.columns,
@@ -1297,7 +1307,7 @@ class Builder:
         def __init__(
             self,
             table: ibis_types.Table,
-            ordering: TotalOrdering,
+            ordering: RowOrdering,
             columns: Collection[ibis_types.Value] = (),
             hidden_ordering_columns: Collection[ibis_types.Value] = (),
             predicates: Optional[Collection[ibis_types.BooleanValue]] = None,

bigframes/core/compile/compiler.py

@@ -38,7 +38,9 @@
 
 
 def compile_ordered_ir(node: nodes.BigFrameNode) -> compiled.OrderedIR:
-    return typing.cast(compiled.OrderedIR, compile_node(node, True))
+    ir = typing.cast(compiled.OrderedIR, compile_node(node, True))
+    assert ir.has_total_order
+    return ir
 
 
 def compile_unordered_ir(node: nodes.BigFrameNode) -> compiled.UnorderedIR:

bigframes/core/ordering.py

@@ -98,6 +98,8 @@ 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, ...] = ()
+    integer_encoding: IntegerEncoding = IntegerEncoding(False)
+    string_encoding: StringEncoding = StringEncoding(False)
 
     @property
     def all_ordering_columns(self) -> Sequence[OrderingExpression]:
@@ -111,6 +113,20 @@ def referenced_columns(self) -> Set[str]:
             for col in part.scalar_expression.unbound_variables
         )
 
+    @property
+    def is_string_encoded(self) -> bool:
+        """True if ordering is fully defined by a fixed length string column."""
+        return self.string_encoding.is_encoded
+
+    @property
+    def is_sequential(self) -> bool:
+        return self.integer_encoding.is_encoded and self.integer_encoding.is_sequential
+
+    @property
+    def total_order_col(self) -> Optional[OrderingExpression]:
+        """Returns column id of columns that defines total ordering, if such as column exists"""
+        return None
+
     def with_reverse(self) -> RowOrdering:
         """Reverses the ordering."""
         return RowOrdering(
@@ -121,17 +137,66 @@ 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(
+        return RowOrdering(
             tuple(new_value_columns),
         )
 
+    def with_non_sequential(self):
+        """Create a copy that is marked as non-sequential.
+
+        This is useful when filtering, but not sorting, an expression.
+        """
+        if self.integer_encoding.is_sequential:
+            return RowOrdering(
+                self.ordering_value_columns,
+                integer_encoding=IntegerEncoding(
+                    self.integer_encoding.is_encoded, is_sequential=False
+                ),
+            )
+
+        return self
+
+    def with_ordering_columns(
+        self,
+        ordering_value_columns: Sequence[OrderingExpression] = (),
+    ) -> RowOrdering:
+        """Creates a new ordering that reorders by the given columns.
+
+        Args:
+            ordering_value_columns:
+                In decreasing precedence order, the values used to sort the ordering
+
+        Returns:
+            Modified ExpressionOrdering
+        """
+
+        # Truncate to remove any unneded col references after all total order cols included
+        new_ordering = self._truncate_ordering(
+            (*ordering_value_columns, *self.ordering_value_columns)
+        )
+        return RowOrdering(
+            new_ordering,
+        )
+
+    def _truncate_ordering(
+        self, order_refs: tuple[OrderingExpression, ...]
+    ) -> tuple[OrderingExpression, ...]:
+        # Truncate once we refer to a full key in bijective operations
+        columns_seen: Set[str] = set()
+        truncated_refs = []
+        for order_part in order_refs:
+            expr = order_part.scalar_expression
+            if not set(expr.unbound_variables).issubset(columns_seen):
+                if expr.is_bijective:
+                    columns_seen.update(expr.unbound_variables)
+                truncated_refs.append(order_part)
+        return tuple(truncated_refs)
+
 
 @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.
     # These columns must always be part of the ordering, in order to guarantee that the ordering is total.
     # Therefore, any modifications(or drops) done to these columns must result in hidden copies being made.
@@ -234,15 +299,6 @@ def total_order_col(self) -> Optional[OrderingExpression]:
             return None
         return order_ref
 
-    @property
-    def is_string_encoded(self) -> bool:
-        """True if ordering is fully defined by a fixed length string column."""
-        return self.string_encoding.is_encoded
-
-    @property
-    def is_sequential(self) -> bool:
-        return self.integer_encoding.is_encoded and self.integer_encoding.is_sequential
-
 
 def encode_order_string(
     order_id: ibis_types.IntegerColumn, length: int = DEFAULT_ORDERING_ID_LENGTH