perf: Utilize ORDER BY LIMIT over ROW_NUMBER where possible (#1077)

Author: TrevorBergeron

bigframes/core/compile/api.py

@@ -18,6 +18,7 @@
 import google.cloud.bigquery as bigquery
 
 import bigframes.core.compile.compiler as compiler
+import bigframes.core.rewrite as rewrites
 
 if TYPE_CHECKING:
     import bigframes.core.nodes
@@ -42,6 +43,7 @@ def compile_unordered(
         col_id_overrides: Mapping[str, str] = {},
     ) -> str:
         """Compile node into sql where rows are unsorted, and no ordering information is preserved."""
+        # TODO: Enable limit pullup, but only if not being used to write to clustered table.
         return self._compiler.compile_unordered_ir(node).to_sql(
             col_id_overrides=col_id_overrides
         )
@@ -53,8 +55,10 @@ def compile_ordered(
         col_id_overrides: Mapping[str, str] = {},
     ) -> str:
         """Compile node into sql where rows are sorted with ORDER BY."""
-        return self._compiler.compile_ordered_ir(node).to_sql(
-            col_id_overrides=col_id_overrides, ordered=True
+        # If we are ordering the query anyways, compiling the slice as a limit is probably a good idea.
+        new_node, limit = rewrites.pullup_limit_from_slice(node)
+        return self._compiler.compile_ordered_ir(new_node).to_sql(
+            col_id_overrides=col_id_overrides, ordered=True, limit=limit
         )
 
     def compile_raw(

bigframes/core/compile/compiled.py

@@ -943,8 +943,9 @@ def to_sql(
         self,
         col_id_overrides: typing.Mapping[str, str] = {},
         ordered: bool = False,
+        limit: Optional[int] = None,
     ) -> str:
-        if ordered:
+        if ordered or limit:
             # Need to bake ordering expressions into the selected column in order for our ordering clause builder to work.
             baked_ir = self._bake_ordering()
             sql = ibis_bigquery.Backend().compile(
@@ -969,7 +970,11 @@ def to_sql(
                 order_by_clause = bigframes.core.sql.ordering_clause(
                     baked_ir._ordering.all_ordering_columns
                 )
-                sql += f"{order_by_clause}\n"
+                sql += f"\n{order_by_clause}"
+            if limit is not None:
+                if not isinstance(limit, int):
+                    raise TypeError(f"Limit param: {limit} must be an int.")
+                sql += f"\nLIMIT {limit}"
         else:
             sql = ibis_bigquery.Backend().compile(
                 self._to_ibis_expr(

bigframes/core/nodes.py

@@ -20,7 +20,7 @@
 import functools
 import itertools
 import typing
-from typing import Callable, Iterable, Optional, Sequence, Tuple
+from typing import Callable, cast, Iterable, Optional, Sequence, Tuple
 
 import google.cloud.bigquery as bq
 
@@ -30,6 +30,7 @@
 import bigframes.core.identifiers as bfet_ids
 from bigframes.core.ordering import OrderingExpression
 import bigframes.core.schema as schemata
+import bigframes.core.slices as slices
 import bigframes.core.window_spec as window
 import bigframes.dtypes
 import bigframes.operations.aggregations as agg_ops
@@ -82,6 +83,11 @@ def child_nodes(self) -> typing.Sequence[BigFrameNode]:
         """Direct children of this node"""
         return tuple([])
 
+    @property
+    @abc.abstractmethod
+    def row_count(self) -> typing.Optional[int]:
+        return None
+
     @functools.cached_property
     def session(self):
         sessions = []
@@ -304,6 +310,26 @@ def variables_introduced(self) -> int:
     def relation_ops_created(self) -> int:
         return 2
 
+    @property
+    def is_limit(self) -> bool:
+        """Returns whether this is equivalent to a ORDER BY ... LIMIT N."""
+        # TODO: Handle tail case.
+        return (
+            (not self.start)
+            and (self.step == 1)
+            and (self.stop is not None)
+            and (self.stop > 0)
+        )
+
+    @property
+    def row_count(self) -> typing.Optional[int]:
+        child_length = self.child.row_count
+        if child_length is None:
+            return None
+        return slices.slice_output_rows(
+            (self.start, self.stop, self.step), child_length
+        )
+
 
 @dataclass(frozen=True, eq=False)
 class JoinNode(BigFrameNode):
@@ -351,6 +377,15 @@ def variables_introduced(self) -> int:
     def joins(self) -> bool:
         return True
 
+    @property
+    def row_count(self) -> Optional[int]:
+        if self.type == "cross":
+            if self.left_child.row_count is None or self.right_child.row_count is None:
+                return None
+            return self.left_child.row_count * self.right_child.row_count
+
+        return None
+
     def transform_children(
         self, t: Callable[[BigFrameNode], BigFrameNode]
     ) -> BigFrameNode:
@@ -412,6 +447,16 @@ 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) + OVERHEAD_VARIABLES
 
+    @property
+    def row_count(self) -> Optional[int]:
+        sub_counts = [node.row_count for node in self.child_nodes]
+        total = 0
+        for count in sub_counts:
+            if count is None:
+                return None
+            total += count
+        return total
+
     def transform_children(
         self, t: Callable[[BigFrameNode], BigFrameNode]
     ) -> BigFrameNode:
@@ -460,6 +505,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) + OVERHEAD_VARIABLES
 
+    @property
+    def row_count(self) -> Optional[int]:
+        return None
+
     def transform_children(
         self, t: Callable[[BigFrameNode], BigFrameNode]
     ) -> BigFrameNode:
@@ -484,19 +533,18 @@ def roots(self) -> typing.Set[BigFrameNode]:
         return {self}
 
     @property
-    def supports_fast_head(self) -> bool:
+    def fast_offsets(self) -> bool:
+        return False
+
+    @property
+    def fast_ordered_limit(self) -> bool:
         return False
 
     def transform_children(
         self, t: Callable[[BigFrameNode], BigFrameNode]
     ) -> BigFrameNode:
         return self
 
-    @property
-    def row_count(self) -> typing.Optional[int]:
-        """How many rows are in the data source. None means unknown."""
-        return None
-
 
 class ScanItem(typing.NamedTuple):
     id: bfet_ids.ColumnId
@@ -528,7 +576,11 @@ def variables_introduced(self) -> int:
         return len(self.scan_list.items) + 1
 
     @property
-    def supports_fast_head(self) -> bool:
+    def fast_offsets(self) -> bool:
+        return True
+
+    @property
+    def fast_ordered_limit(self) -> bool:
         return True
 
     @property
@@ -635,12 +687,27 @@ def relation_ops_created(self) -> int:
         return 3
 
     @property
-    def supports_fast_head(self) -> bool:
-        # Fast head is only supported when row offsets are available.
-        # In the future, ORDER BY+LIMIT optimizations may allow fast head when
-        # clustered and/or partitioned on ordering key
+    def fast_offsets(self) -> bool:
+        # Fast head is only supported when row offsets are available or data is clustered over ordering key.
         return (self.source.ordering is not None) and self.source.ordering.is_sequential
 
+    @property
+    def fast_ordered_limit(self) -> bool:
+        if self.source.ordering is None:
+            return False
+        order_cols = self.source.ordering.all_ordering_columns
+        # monotonicity would probably be fine
+        if not all(col.scalar_expression.is_identity for col in order_cols):
+            return False
+        order_col_ids = tuple(
+            cast(ex.DerefOp, col.scalar_expression).id.name for col in order_cols
+        )
+        cluster_col_ids = self.source.table.cluster_cols
+        if cluster_col_ids is None:
+            return False
+
+        return order_col_ids == cluster_col_ids[: len(order_col_ids)]
+
     @property
     def order_ambiguous(self) -> bool:
         return (
@@ -706,6 +773,10 @@ def relation_ops_created(self) -> int:
     def variables_introduced(self) -> int:
         return 1
 
+    @property
+    def row_count(self) -> Optional[int]:
+        return self.child.row_count
+
     def prune(self, used_cols: COLUMN_SET) -> BigFrameNode:
         if self.col_id not in used_cols:
             return self.child.prune(used_cols)
@@ -726,6 +797,10 @@ def row_preserving(self) -> bool:
     def variables_introduced(self) -> int:
         return 1
 
+    @property
+    def row_count(self) -> Optional[int]:
+        return None
+
     def prune(self, used_cols: COLUMN_SET) -> BigFrameNode:
         consumed_ids = used_cols.union(self.predicate.column_references)
         pruned_child = self.child.prune(consumed_ids)
@@ -749,6 +824,10 @@ def relation_ops_created(self) -> int:
     def explicitly_ordered(self) -> bool:
         return True
 
+    @property
+    def row_count(self) -> Optional[int]:
+        return self.child.row_count
+
     def prune(self, used_cols: COLUMN_SET) -> BigFrameNode:
         ordering_cols = itertools.chain.from_iterable(
             map(lambda x: x.referenced_columns, self.by)
@@ -772,6 +851,10 @@ def relation_ops_created(self) -> int:
         # Doesnt directly create any relational operations
         return 0
 
+    @property
+    def row_count(self) -> Optional[int]:
+        return self.child.row_count
+
 
 @dataclass(frozen=True, eq=False)
 class SelectionNode(UnaryNode):
@@ -798,6 +881,10 @@ def variables_introduced(self) -> int:
     def defines_namespace(self) -> bool:
         return True
 
+    @property
+    def row_count(self) -> Optional[int]:
+        return self.child.row_count
+
     def prune(self, used_cols: COLUMN_SET) -> BigFrameNode:
         pruned_selections = tuple(
             select for select in self.input_output_pairs if select[1] in used_cols
@@ -842,6 +929,10 @@ def variables_introduced(self) -> int:
         new_vars = sum(1 for i in self.assignments if not i[0].is_identity)
         return new_vars
 
+    @property
+    def row_count(self) -> Optional[int]:
+        return self.child.row_count
+
     def prune(self, used_cols: COLUMN_SET) -> BigFrameNode:
         pruned_assignments = tuple(i for i in self.assignments if i[1] in used_cols)
         if len(pruned_assignments) == 0:
@@ -877,6 +968,10 @@ def variables_introduced(self) -> int:
     def defines_namespace(self) -> bool:
         return True
 
+    @property
+    def row_count(self) -> Optional[int]:
+        return 1
+
 
 @dataclass(frozen=True, eq=False)
 class AggregateNode(UnaryNode):
@@ -926,6 +1021,12 @@ def explicitly_ordered(self) -> bool:
     def defines_namespace(self) -> bool:
         return True
 
+    @property
+    def row_count(self) -> Optional[int]:
+        if not self.by_column_ids:
+            return 1
+        return None
+
     def prune(self, used_cols: COLUMN_SET) -> BigFrameNode:
         by_ids = (ref.id for ref in self.by_column_ids)
         pruned_aggs = tuple(agg for agg in self.aggregations if agg[1] in used_cols)
@@ -963,6 +1064,10 @@ def relation_ops_created(self) -> int:
         # Assume that if not reprojecting, that there is a sequence of window operations sharing the same window
         return 0 if self.skip_reproject_unsafe else 4
 
+    @property
+    def row_count(self) -> Optional[int]:
+        return self.child.row_count
+
     @functools.cached_property
     def added_field(self) -> Field:
         input_type = self.child.get_type(self.column_name.id)
@@ -994,6 +1099,10 @@ def row_preserving(self) -> bool:
     def variables_introduced(self) -> int:
         return 1
 
+    @property
+    def row_count(self) -> Optional[int]:
+        return None
+
 
 # TODO: Explode should create a new column instead of overriding the existing one
 @dataclass(frozen=True, eq=False)
@@ -1030,6 +1139,10 @@ def variables_introduced(self) -> int:
     def defines_namespace(self) -> bool:
         return True
 
+    @property
+    def row_count(self) -> Optional[int]:
+        return None
+
     def prune(self, used_cols: COLUMN_SET) -> BigFrameNode:
         # Cannot prune explode op
         return self.transform_children(