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(