diff --git a/bigframes/core/expression.py b/bigframes/core/expression.py index 7b20e430ff..0e94193bd3 100644 --- a/bigframes/core/expression.py +++ b/bigframes/core/expression.py @@ -253,6 +253,11 @@ def is_identity(self) -> bool: def transform_children(self, t: Callable[[Expression], Expression]) -> Expression: ... + def bottom_up(self, t: Callable[[Expression], Expression]) -> Expression: + expr = self.transform_children(lambda child: child.bottom_up(t)) + expr = t(expr) + return expr + def walk(self) -> Generator[Expression, None, None]: yield self for child in self.children: diff --git a/bigframes/core/rewrite/op_lowering.py b/bigframes/core/rewrite/op_lowering.py index a64a4cc8c4..6473c3bf8a 100644 --- a/bigframes/core/rewrite/op_lowering.py +++ b/bigframes/core/rewrite/op_lowering.py @@ -44,7 +44,7 @@ def lower_expr_step(expr: expression.Expression) -> expression.Expression: return maybe_rule.lower(expr) return expr - return lower_expr_step(expr.transform_children(lower_expr_step)) + return expr.bottom_up(lower_expr_step) def lower_node(node: bigframe_node.BigFrameNode) -> bigframe_node.BigFrameNode: if isinstance( diff --git a/tests/system/small/engines/test_generic_ops.py b/tests/system/small/engines/test_generic_ops.py index 1d28c335a6..8deef3638e 100644 --- a/tests/system/small/engines/test_generic_ops.py +++ b/tests/system/small/engines/test_generic_ops.py @@ -423,3 +423,18 @@ def test_engines_isin_op(scalars_array_value: array_value.ArrayValue, engine): ) assert_equivalence_execution(arr.node, REFERENCE_ENGINE, engine) + + +@pytest.mark.parametrize("engine", ["polars", "bq"], indirect=True) +def test_engines_isin_op_nested_filter( + scalars_array_value: array_value.ArrayValue, engine +): + isin_clause = ops.IsInOp((1, 2, 3)).as_expr(expression.deref("int64_col")) + filter_clause = ops.invert_op.as_expr( + ops.or_op.as_expr( + expression.deref("bool_col"), ops.invert_op.as_expr(isin_clause) + ) + ) + arr = scalars_array_value.filter(filter_clause) + + assert_equivalence_execution(arr.node, REFERENCE_ENGINE, engine)