test: Cross validate sort execution between engines (#1823)

Author: TrevorBergeron

bigframes/core/bigframe_node.py

@@ -20,9 +20,19 @@
 import functools
 import itertools
 import typing
-from typing import Callable, Dict, Generator, Iterable, Mapping, Sequence, Set, Tuple
-
-from bigframes.core import field, identifiers
+from typing import (
+    Callable,
+    Dict,
+    Generator,
+    Iterable,
+    Mapping,
+    Sequence,
+    Set,
+    Tuple,
+    Union,
+)
+
+from bigframes.core import expression, field, identifiers
 import bigframes.core.schema as schemata
 import bigframes.dtypes
 
@@ -278,6 +288,13 @@ def _dtype_lookup(self) -> dict[identifiers.ColumnId, bigframes.dtypes.Dtype]:
     def field_by_id(self) -> Mapping[identifiers.ColumnId, field.Field]:
         return {field.id: field for field in self.fields}
 
+    @property
+    def _node_expressions(
+        self,
+    ) -> Sequence[Union[expression.Expression, expression.Aggregation]]:
+        """List of scalar expressions. Intended for checking engine compatibility with used ops."""
+        return ()
+
     # Plan algorithms
     def unique_nodes(
         self: BigFrameNode,

bigframes/core/nodes.py

@@ -274,6 +274,10 @@ def joins_nulls(self) -> bool:
         right_nullable = self.right_child.field_by_id[self.right_col.id].nullable
         return left_nullable or right_nullable
 
+    @property
+    def _node_expressions(self):
+        return (self.left_col, self.right_col)
+
     def replace_additive_base(self, node: BigFrameNode):
         return dataclasses.replace(self, left_child=node)
 
@@ -387,6 +391,10 @@ def referenced_ids(self) -> COLUMN_SET:
     def consumed_ids(self) -> COLUMN_SET:
         return frozenset(*self.ids, *self.referenced_ids)
 
+    @property
+    def _node_expressions(self):
+        return tuple(itertools.chain.from_iterable(self.conditions))
+
     def transform_children(self, t: Callable[[BigFrameNode], BigFrameNode]) -> JoinNode:
         transformed = dataclasses.replace(
             self, left_child=t(self.left_child), right_child=t(self.right_child)
@@ -996,6 +1004,10 @@ def consumed_ids(self) -> COLUMN_SET:
     def referenced_ids(self) -> COLUMN_SET:
         return frozenset(self.predicate.column_references)
 
+    @property
+    def _node_expressions(self):
+        return (self.predicate,)
+
     def remap_vars(
         self, mappings: Mapping[identifiers.ColumnId, identifiers.ColumnId]
     ) -> FilterNode:
@@ -1050,6 +1062,10 @@ def referenced_ids(self) -> COLUMN_SET:
             itertools.chain.from_iterable(map(lambda x: x.referenced_columns, self.by))
         )
 
+    @property
+    def _node_expressions(self):
+        return tuple(map(lambda x: x.scalar_expression, self.by))
+
     def remap_vars(
         self, mappings: Mapping[identifiers.ColumnId, identifiers.ColumnId]
     ) -> OrderByNode:
@@ -1178,6 +1194,10 @@ def node_defined_ids(self) -> Tuple[identifiers.ColumnId, ...]:
     def consumed_ids(self) -> COLUMN_SET:
         return frozenset(ref.id for ref, id in self.input_output_pairs)
 
+    @property
+    def _node_expressions(self):
+        return tuple(ref for ref, id in self.input_output_pairs)
+
     def get_id_mapping(self) -> dict[identifiers.ColumnId, identifiers.ColumnId]:
         return {ref.id: id for ref, id in self.input_output_pairs}
 
@@ -1265,6 +1285,10 @@ def referenced_ids(self) -> COLUMN_SET:
             )
         )
 
+    @property
+    def _node_expressions(self):
+        return tuple(ex for ex, id in self.assignments)
+
     @property
     def additive_base(self) -> BigFrameNode:
         return self.child
@@ -1361,6 +1385,13 @@ def has_ordered_ops(self) -> bool:
             aggregate.op.order_independent for aggregate, _ in self.aggregations
         )
 
+    @property
+    def _node_expressions(self):
+        by_ids = (ref for ref in self.by_column_ids)
+        aggs = tuple(agg for agg, _ in self.aggregations)
+        order_ids = tuple(part.scalar_expression for part in self.order_by)
+        return (*by_ids, *aggs, *order_ids)
+
     def remap_vars(
         self, mappings: Mapping[identifiers.ColumnId, identifiers.ColumnId]
     ) -> AggregateNode:
@@ -1463,6 +1494,10 @@ def inherits_order(self) -> bool:
     def additive_base(self) -> BigFrameNode:
         return self.child
 
+    @property
+    def _node_expressions(self):
+        return (self.expression, *self.window_spec.expressions)
+
     def replace_additive_base(self, node: BigFrameNode) -> WindowOpNode:
         return dataclasses.replace(self, child=node)
 
@@ -1584,6 +1619,10 @@ def node_defined_ids(self) -> Tuple[identifiers.ColumnId, ...]:
     def referenced_ids(self) -> COLUMN_SET:
         return frozenset(ref.id for ref in self.column_ids)
 
+    @property
+    def _node_expressions(self):
+        return self.column_ids
+
     def remap_vars(
         self, mappings: Mapping[identifiers.ColumnId, identifiers.ColumnId]
     ) -> ExplodeNode:
@@ -1657,6 +1696,10 @@ def row_count(self) -> Optional[int]:
     def variables_introduced(self) -> int:
         return 0
 
+    @property
+    def _node_expressions(self):
+        return tuple(ref for ref, _ in self.output_cols)
+
 
 # Tree operators
 def top_down(

bigframes/core/window_spec.py

@@ -16,7 +16,7 @@
 from dataclasses import dataclass, replace
 import datetime
 import itertools
-from typing import Literal, Mapping, Optional, Set, Tuple, Union
+from typing import Literal, Mapping, Optional, Sequence, Set, Tuple, Union
 
 import numpy as np
 import pandas as pd
@@ -260,6 +260,11 @@ def is_unbounded(self):
             self.bounds.start is None and self.bounds.end is None
         )
 
+    @property
+    def expressions(self) -> Sequence[ex.Expression]:
+        ordering_exprs = (item.scalar_expression for item in self.ordering)
+        return (*self.grouping_keys, *ordering_exprs)
+
     @property
     def all_referenced_columns(self) -> Set[ids.ColumnId]:
         """

bigframes/session/polars_executor.py

@@ -13,26 +13,46 @@
 # limitations under the License.
 from __future__ import annotations
 
+import itertools
 from typing import Optional, TYPE_CHECKING
 
 import pyarrow as pa
 
-from bigframes.core import array_value, bigframe_node, local_data, nodes
+from bigframes.core import array_value, bigframe_node, expression, local_data, nodes
+import bigframes.operations
 from bigframes.session import executor, semi_executor
 
 if TYPE_CHECKING:
     import polars as pl
 
-
+# Polars executor can execute more node types, but these are the validated ones
 _COMPATIBLE_NODES = (
     nodes.ReadLocalNode,
     nodes.OrderByNode,
     nodes.ReversedNode,
     nodes.SelectionNode,
-    nodes.FilterNode,  # partial support
-    nodes.ProjectionNode,  # partial support
 )
 
+_COMPATIBLE_SCALAR_OPS = ()
+
+
+def _get_expr_ops(expr: expression.Expression) -> set[bigframes.operations.ScalarOp]:
+    if isinstance(expr, expression.OpExpression):
+        return set(itertools.chain.from_iterable(map(_get_expr_ops, expr.children)))
+    return set()
+
+
+def _is_node_polars_executable(node: nodes.BigFrameNode):
+    if not isinstance(node, _COMPATIBLE_NODES):
+        return False
+    for expr in node._node_expressions:
+        if isinstance(expr, expression.Aggregation):
+            return False
+        if isinstance(expr, expression.Expression):
+            if not _get_expr_ops(expr).issubset(_COMPATIBLE_SCALAR_OPS):
+                return False
+    return True
+
 
 class PolarsExecutor(semi_executor.SemiExecutor):
     def __init__(self):
@@ -67,7 +87,7 @@ def execute(
         )
 
     def _can_execute(self, plan: bigframe_node.BigFrameNode):
-        return all(isinstance(node, _COMPATIBLE_NODES) for node in plan.unique_nodes())
+        return all(_is_node_polars_executable(node) for node in plan.unique_nodes())
 
     def _adapt_array(self, array: pa.Array) -> pa.Array:
         target_type = local_data.logical_type_replacements(array.type)

tests/system/small/engines/test_sorting.py

@@ -0,0 +1,103 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+
+from bigframes.core import array_value, nodes, ordering
+import bigframes.operations as bf_ops
+from bigframes.session import polars_executor
+from bigframes.testing.engine_utils import assert_equivalence_execution
+
+pytest.importorskip("polars")
+
+# Polars used as reference as its fast and local. Generally though, prefer gbq engine where they disagree.
+REFERENCE_ENGINE = polars_executor.PolarsExecutor()
+
+
+@pytest.mark.parametrize("engine", ["polars", "bq"], indirect=True)
+def test_engines_reverse(
+    scalars_array_value: array_value.ArrayValue,
+    engine,
+):
+    node = apply_reverse(scalars_array_value.node)
+    assert_equivalence_execution(node, REFERENCE_ENGINE, engine)
+
+
+@pytest.mark.parametrize("engine", ["polars", "bq"], indirect=True)
+def test_engines_double_reverse(
+    scalars_array_value: array_value.ArrayValue,
+    engine,
+):
+    node = apply_reverse(scalars_array_value.node)
+    assert_equivalence_execution(node, REFERENCE_ENGINE, engine)
+
+
+@pytest.mark.parametrize("engine", ["polars", "bq"], indirect=True)
+@pytest.mark.parametrize(
+    "sort_col",
+    [
+        "bool_col",
+        "int64_col",
+        "bytes_col",
+        "date_col",
+        "datetime_col",
+        "int64_col",
+        "int64_too",
+        "numeric_col",
+        "float64_col",
+        "string_col",
+        "time_col",
+        "timestamp_col",
+    ],
+)
+def test_engines_sort_over_column(
+    scalars_array_value: array_value.ArrayValue, engine, sort_col
+):
+    node = apply_reverse(scalars_array_value.node)
+    ORDER_EXPRESSIONS = (ordering.descending_over(sort_col, nulls_last=False),)
+    node = nodes.OrderByNode(node, ORDER_EXPRESSIONS)
+    assert_equivalence_execution(node, REFERENCE_ENGINE, engine)
+
+
+@pytest.mark.parametrize("engine", ["polars", "bq"], indirect=True)
+def test_engines_sort_multi_column_refs(
+    scalars_array_value: array_value.ArrayValue,
+    engine,
+):
+    node = scalars_array_value.node
+    ORDER_EXPRESSIONS = (
+        ordering.ascending_over("bool_col", nulls_last=False),
+        ordering.descending_over("int64_col"),
+    )
+    node = nodes.OrderByNode(node, ORDER_EXPRESSIONS)
+    assert_equivalence_execution(node, REFERENCE_ENGINE, engine)
+
+
+@pytest.mark.parametrize("engine", ["polars"], indirect=True)
+def test_polars_engines_skips_unrecognized_order_expr(
+    scalars_array_value: array_value.ArrayValue,
+    engine,
+):
+    node = scalars_array_value.node
+    ORDER_EXPRESSIONS = (
+        ordering.OrderingExpression(
+            scalar_expression=bf_ops.sin_op.as_expr("float_col")
+        ),
+    )
+    node = nodes.OrderByNode(node, ORDER_EXPRESSIONS)
+    assert engine.execute(node, ordered=True) is None
+
+
+def apply_reverse(node: nodes.BigFrameNode) -> nodes.BigFrameNode:
+    return nodes.ReversedNode(node)