Skip to content

Commit 8f4c51f

Browse files
authored
fix: keep filters with correlated subqueries above extension nodes (#25294)
## Which issue does this PR close? - Closes #15046. ## Rationale for this change An extension node can name columns whose predicates must not be pushed past it, via `UserDefinedLogicalNode::prevent_predicate_push_down_columns`. A predicate whose only dependency on those columns runs through a correlated subquery is pushed down anyway, and the resulting plan does not just perform worse — it fails to execute. Using the `NoopPlan` node already present in the optimizer's own tests, which refuses push-down of column `c`: ``` Filter: EXISTS (SELECT sq.a FROM sq WHERE outer_ref(test.c) = sq.a) NoopPlan TableScan: test ``` optimizing that plan today gives ``` Optimizer rule 'push_down_filter' failed Invalid (non-executable) plan after Optimizer rule: push_down_filter In/Exist/SetComparison subquery can only be used in Projection, Filter, TableScan, Window functions, Aggregate and Join plan nodes, but was used in [NoopPlan] ``` The predicate is moved below the node that asked to keep it, lands somewhere subqueries are not permitted, and the invariant check turns that into a query failure. The cause is the one #15046 describes. The check is ```rust !expr.column_refs().iter().any(|c| prevent_cols.contains(&c.name)) ``` and `column_refs` collects only `Expr::Column`. A subquery records the outer columns it correlates on in `Subquery::outer_ref_columns`, which `Expr`'s traversal does not descend into, so `EXISTS (... WHERE outer.c = ...)` reports no column references at all and looks unconditionally safe to push. ## What changes are included in this PR? A `references_any_column` helper in `push_down_filter.rs` that asks `column_refs` first and then, for `Expr::Exists`, `Expr::InSubquery` and `Expr::ScalarSubquery`, also looks at the `outer_ref_columns` those carry. The extension-node branch uses it instead of `column_refs` alone. Nothing else changes. **I deliberately did not change `column_refs` itself**, which is what the issue title asks for. It has many callers whose behaviour would shift — `optimize_projections`, CSE, window and join handling all use it, and several of them want "columns produced by this plan's own schema", not outer references belonging to a parent. Widening it looks like a separate change with its own risk, so this PR fixes the consumer that has a demonstrable failure. Happy to look at the broader change instead if you would rather that were the fix. ## What is the testing strategy for this PR? New unit test `push_down_filter::tests::user_defined_plan_outer_referenced_column`, next to the existing `user_defined_plan` test and reusing its `NoopPlan`. It asserts the `EXISTS` predicate stays above the node. The test discriminates: with the outer-reference collection short-circuited out, it fails with the invalid-plan error quoted above; with the fix it passes. `datafusion-optimizer` (808) and `datafusion-expr` (262) pass, as does `sqllogictest`. ## Are there any user-facing changes? No public API changes. Plans change only where an extension node declares `prevent_predicate_push_down_columns` and a predicate correlates a subquery on one of those columns — previously an invalid plan, now a correct one. No in-tree node declares those columns outside tests, so no bundled plan snapshots move. Two adjacent things found while working on this, neither touched here: - `find_out_reference_exprs` does not see these outer references either, for the same reason, so it cannot be used as the collection mechanism. That is the same traversal gap as #16147, which @AdamGS has taken. - The equivalent guard for correlated subqueries in other push-down branches was not audited; this PR only covers the extension-node branch that #15046 reports.
1 parent c0e872f commit 8f4c51f

1 file changed

Lines changed: 132 additions & 12 deletions

File tree

‎datafusion/optimizer/src/push_down_filter.rs‎

Lines changed: 132 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,9 @@ use itertools::Itertools;
2626
use log::{Level, debug, log_enabled};
2727

2828
use datafusion_common::instant::Instant;
29-
use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
29+
use datafusion_common::tree_node::{
30+
Transformed, TransformedResult, TreeNode, TreeNodeRecursion,
31+
};
3032
use datafusion_common::{
3133
Column, DFSchema, Result, assert_eq_or_internal_err, internal_err, plan_err,
3234
qualified_name,
@@ -783,6 +785,45 @@ fn infer_join_predicates_impl<
783785
Ok(())
784786
}
785787

788+
/// Whether `expr` depends on any of the columns named in `names`.
789+
///
790+
/// This is the columns `Expr::column_refs` would collect plus the outer columns
791+
/// that any subquery inside `expr` correlates on. A subquery records those in
792+
/// `Subquery::outer_ref_columns` rather than as an `Expr::Column` in the
793+
/// predicate, and `Expr`'s own traversal does not descend into that field, so
794+
/// looking only at `column_refs` would report such a predicate as depending on
795+
/// nothing and let it be pushed past a node that asked to keep those columns.
796+
fn references_any_column(expr: &Expr, names: &HashSet<String>) -> bool {
797+
let mut found = false;
798+
expr.apply(|e| {
799+
let outer_refs = match e {
800+
Expr::Column(col) => {
801+
if names.contains(&col.name) {
802+
found = true;
803+
return Ok(TreeNodeRecursion::Stop);
804+
}
805+
return Ok(TreeNodeRecursion::Continue);
806+
}
807+
Expr::Exists(exists) => &exists.subquery.outer_ref_columns,
808+
Expr::InSubquery(in_subquery) => &in_subquery.subquery.outer_ref_columns,
809+
Expr::ScalarSubquery(subquery) => &subquery.outer_ref_columns,
810+
Expr::SetComparison(set_comparison) => {
811+
&set_comparison.subquery.outer_ref_columns
812+
}
813+
_ => return Ok(TreeNodeRecursion::Continue),
814+
};
815+
if outer_refs.iter().any(|outer_ref| {
816+
matches!(outer_ref, Expr::OuterReferenceColumn(_, c) if names.contains(&c.name))
817+
}) {
818+
found = true;
819+
return Ok(TreeNodeRecursion::Stop);
820+
}
821+
Ok(TreeNodeRecursion::Continue)
822+
})
823+
.expect("traversal is infallible");
824+
found
825+
}
826+
786827
impl OptimizerRule for PushDownFilter {
787828
fn name(&self) -> &str {
788829
"push_down_filter"
@@ -1301,12 +1342,7 @@ impl OptimizerRule for PushDownFilter {
13011342
let predicate_push_or_keep: Vec<bool> =
13021343
split_conjunction(&filter.predicate)
13031344
.iter()
1304-
.map(|expr| {
1305-
!expr
1306-
.column_refs()
1307-
.iter()
1308-
.any(|c| prevent_cols.contains(&c.name))
1309-
})
1345+
.map(|expr| !references_any_column(expr, &prevent_cols))
13101346
.collect();
13111347

13121348
// all predicates are kept, no changes needed
@@ -1523,14 +1559,14 @@ mod tests {
15231559
use arrow::datatypes::{Field, Metadata, Schema, SchemaRef};
15241560
use async_trait::async_trait;
15251561

1526-
use datafusion_common::{DFSchemaRef, DataFusionError, ScalarValue};
1527-
use datafusion_expr::expr::ScalarFunction;
1562+
use datafusion_common::{DFSchemaRef, DataFusionError, ScalarValue, Spans};
1563+
use datafusion_expr::expr::{ScalarFunction, SetComparison, SetQuantifier};
15281564
use datafusion_expr::logical_plan::table_scan;
15291565
use datafusion_expr::{
15301566
ColumnarValue, ExprFunctionExt, Extension, LogicalPlanBuilder,
1531-
ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, TableScan, TableSource,
1532-
TableType, UserDefinedLogicalNodeCore, Volatility, WindowFunctionDefinition, col,
1533-
in_list, in_subquery, lit,
1567+
ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Subquery, TableScan,
1568+
TableSource, TableType, UserDefinedLogicalNodeCore, Volatility,
1569+
WindowFunctionDefinition, col, exists, in_list, in_subquery, lit, out_ref_col,
15341570
};
15351571

15361572
use crate::OptimizerContext;
@@ -2185,6 +2221,90 @@ mod tests {
21852221
}
21862222
}
21872223

2224+
#[test]
2225+
fn user_defined_plan_outer_referenced_column() -> Result<()> {
2226+
let table_scan = test_table_scan()?;
2227+
2228+
// A subquery correlated on `test.c` — the column `NoopPlan` refuses to
2229+
// have predicates pushed past. The correlation is carried by the
2230+
// subquery's `outer_ref_columns`, not by an `Expr::Column` in the
2231+
// predicate itself.
2232+
let subquery = LogicalPlanBuilder::from(test_table_scan_with_name("sq")?)
2233+
.filter(out_ref_col(DataType::UInt32, "test.c").eq(col("sq.a")))?
2234+
.project(vec![col("sq.a")])?
2235+
.build()?;
2236+
2237+
let custom_plan = LogicalPlan::Extension(Extension {
2238+
node: Arc::new(NoopPlan {
2239+
input: vec![table_scan.clone()],
2240+
schema: Arc::clone(table_scan.schema()),
2241+
}),
2242+
});
2243+
let plan = LogicalPlanBuilder::from(custom_plan)
2244+
.filter(exists(Arc::new(subquery)))?
2245+
.build()?;
2246+
2247+
// The predicate depends on `test.c`, so it must stay above NoopPlan.
2248+
assert_optimized_plan_equal!(
2249+
plan,
2250+
@r"
2251+
Filter: EXISTS (<subquery>)
2252+
Subquery:
2253+
Projection: sq.a
2254+
TableScan: sq, full_filters=[outer_ref(test.c) = sq.a]
2255+
NoopPlan
2256+
TableScan: test
2257+
"
2258+
)
2259+
}
2260+
2261+
#[test]
2262+
fn user_defined_plan_outer_referenced_column_set_comparison() -> Result<()> {
2263+
let table_scan = test_table_scan()?;
2264+
2265+
// `test.a > ANY (SELECT sq.a FROM sq WHERE test.c = sq.a)`: the
2266+
// comparison expression names only `test.a`, so the dependency on
2267+
// `test.c` exists solely in the subquery's `outer_ref_columns`.
2268+
let subquery = LogicalPlanBuilder::from(test_table_scan_with_name("sq")?)
2269+
.filter(out_ref_col(DataType::UInt32, "test.c").eq(col("sq.a")))?
2270+
.project(vec![col("sq.a")])?
2271+
.build()?;
2272+
let outer_ref_columns = subquery.all_out_ref_exprs();
2273+
let set_comparison = Expr::SetComparison(SetComparison::new(
2274+
Box::new(col("test.a")),
2275+
Subquery {
2276+
subquery: Arc::new(subquery),
2277+
outer_ref_columns,
2278+
spans: Spans::new(),
2279+
},
2280+
Operator::Gt,
2281+
SetQuantifier::Any,
2282+
));
2283+
2284+
let custom_plan = LogicalPlan::Extension(Extension {
2285+
node: Arc::new(NoopPlan {
2286+
input: vec![table_scan.clone()],
2287+
schema: Arc::clone(table_scan.schema()),
2288+
}),
2289+
});
2290+
let plan = LogicalPlanBuilder::from(custom_plan)
2291+
.filter(set_comparison)?
2292+
.build()?;
2293+
2294+
// The predicate depends on `test.c`, so it must stay above NoopPlan.
2295+
assert_optimized_plan_equal!(
2296+
plan,
2297+
@r"
2298+
Filter: test.a > ANY (<subquery>)
2299+
Subquery:
2300+
Projection: sq.a
2301+
TableScan: sq, full_filters=[outer_ref(test.c) = sq.a]
2302+
NoopPlan
2303+
TableScan: test
2304+
"
2305+
)
2306+
}
2307+
21882308
#[test]
21892309
fn user_defined_plan() -> Result<()> {
21902310
let table_scan = test_table_scan()?;

0 commit comments

Comments
 (0)