Skip to content

Commit 499119b

Browse files
Retain short IN LISTs with specialized filters
1 parent bec8193 commit 499119b

24 files changed

Lines changed: 679 additions & 155 deletions

‎datafusion/core/tests/sql/explain_analyze.rs‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -519,7 +519,6 @@ async fn csv_explain_verbose() {
519519
async fn csv_explain_inlist_verbose() {
520520
let ctx = SessionContext::new();
521521
register_aggregate_csv_by_sql(&ctx).await;
522-
// Inlist len <=3 case will be transformed to OR List so we test with len=4
523522
let sql = "EXPLAIN VERBOSE SELECT c1 FROM aggregate_test_100 where c2 in (1,2,4,5)";
524523
let actual = execute(&ctx, sql).await;
525524

‎datafusion/optimizer/src/push_down_filter.rs‎

Lines changed: 117 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,9 @@ use datafusion_expr::{
4343
};
4444

4545
use crate::optimizer::ApplyOrder;
46-
use crate::simplify_expressions::{reorder_predicates, simplify_predicates};
46+
use crate::simplify_expressions::{
47+
reorder_predicates, simplify_inlist_disjunction, simplify_predicates,
48+
};
4749
use crate::utils::{
4850
ColumnReference, has_all_column_refs, is_restrict_null_predicate, schema_columns,
4951
};
@@ -307,28 +309,35 @@ fn can_evaluate_as_join_condition(predicate: &Expr) -> Result<bool> {
307309
/// * `((A AND B) OR (C AND D)) AND (A OR C)`
308310
/// * `((A AND B) OR (C AND D)) AND ((A AND B) OR C)`
309311
/// * do nothing.
310-
fn extract_or_clauses_for_join<'a>(
311-
filters: &'a [Expr],
312-
schema_cols: &'a HashSet<ColumnReference>,
313-
) -> impl Iterator<Item = Expr> + 'a {
312+
fn extract_or_clauses_for_join(
313+
filters: &[Expr],
314+
schema_cols: &HashSet<ColumnReference>,
315+
schema: &DFSchema,
316+
) -> Result<Vec<Expr>> {
314317
// new formed OR clauses and their column references
315-
filters.iter().filter_map(move |expr| {
316-
if let Expr::BinaryExpr(BinaryExpr {
317-
left,
318-
op: Operator::Or,
319-
right,
320-
}) = expr
321-
{
322-
let left_expr = extract_or_clause(left.as_ref(), schema_cols);
323-
let right_expr = extract_or_clause(right.as_ref(), schema_cols);
324-
325-
// If nothing can be extracted from any sub clauses, do nothing for this OR clause.
326-
if let (Some(left_expr), Some(right_expr)) = (left_expr, right_expr) {
327-
return Some(or(left_expr, right_expr));
318+
filters
319+
.iter()
320+
.filter_map(move |expr| {
321+
if let Expr::BinaryExpr(BinaryExpr {
322+
left,
323+
op: Operator::Or,
324+
right,
325+
}) = expr
326+
{
327+
let left_expr = extract_or_clause(left.as_ref(), schema_cols);
328+
let right_expr = extract_or_clause(right.as_ref(), schema_cols);
329+
330+
// If nothing can be extracted from any sub clauses, do nothing for this OR clause.
331+
if let (Some(left_expr), Some(right_expr)) = (left_expr, right_expr) {
332+
return Some(simplify_inlist_disjunction(
333+
or(left_expr, right_expr),
334+
schema,
335+
));
336+
}
328337
}
329-
}
330-
None
331-
})
338+
None
339+
})
340+
.collect()
332341
}
333342

334343
/// extract qual from OR sub-clause.
@@ -478,21 +487,25 @@ fn push_down_all_join(
478487
left_push.extend(extract_or_clauses_for_join(
479488
&keep_predicates,
480489
&left_schema_columns,
481-
));
490+
left_schema,
491+
)?);
482492
left_push.extend(extract_or_clauses_for_join(
483493
&join_conditions,
484494
&left_schema_columns,
485-
));
495+
left_schema,
496+
)?);
486497
}
487498
if right_preserved {
488499
right_push.extend(extract_or_clauses_for_join(
489500
&keep_predicates,
490501
&right_schema_columns,
491-
));
502+
right_schema,
503+
)?);
492504
right_push.extend(extract_or_clauses_for_join(
493505
&join_conditions,
494506
&right_schema_columns,
495-
));
507+
right_schema,
508+
)?);
496509
}
497510

498511
// For predicates from join filter, we should check with if a join side is preserved
@@ -501,13 +514,15 @@ fn push_down_all_join(
501514
left_push.extend(extract_or_clauses_for_join(
502515
&on_filter_join_conditions,
503516
&left_schema_columns,
504-
));
517+
left_schema,
518+
)?);
505519
}
506520
if on_right_preserved {
507521
right_push.extend(extract_or_clauses_for_join(
508522
&on_filter_join_conditions,
509523
&right_schema_columns,
510-
));
524+
right_schema,
525+
)?);
511526
}
512527

513528
// Add any new join conditions as the non join predicates
@@ -3699,6 +3714,81 @@ mod tests {
36993714
)
37003715
}
37013716

3717+
#[test]
3718+
fn join_or_pushdown_deduplicates_retained_inlist_across_passes() -> Result<()> {
3719+
let optimizer = Optimizer::with_rules(vec![
3720+
Arc::new(SimplifyExpressions::new()),
3721+
Arc::new(PushDownFilter::new()),
3722+
]);
3723+
let context = OptimizerContext::new().with_max_passes(1);
3724+
let expected = col("b").in_list(vec![lit(1_i32), lit(2_i32)], false);
3725+
3726+
for support in [
3727+
TableProviderFilterPushDown::Unsupported,
3728+
TableProviderFilterPushDown::Inexact,
3729+
TableProviderFilterPushDown::Exact,
3730+
] {
3731+
let left =
3732+
table_scan_with_pushdown_provider_builder(support.clone(), vec![], None)?
3733+
.alias("left")?
3734+
.build()?;
3735+
let right =
3736+
table_scan_with_pushdown_provider_builder(support.clone(), vec![], None)?
3737+
.alias("right")?
3738+
.build()?;
3739+
let mut plan = LogicalPlanBuilder::from(left)
3740+
.join(
3741+
right,
3742+
JoinType::Left,
3743+
(vec!["left.a"], vec!["right.a"]),
3744+
None,
3745+
)?
3746+
.filter(
3747+
col("right.b")
3748+
.not_eq(lit(10_i32))
3749+
.and(col("left.b").eq(lit(1_i32)))
3750+
.or(col("left.b").eq(lit(2_i32))),
3751+
)?
3752+
.build()?;
3753+
3754+
// The predicate above the join remains in place, so every pass
3755+
// extracts the same left-side disjunction again. Each destination
3756+
// must retain only one copy, including filters stored in a scan.
3757+
for _ in 0..3 {
3758+
plan = optimizer.optimize(plan, &context, observe)?;
3759+
let LogicalPlan::Filter(filter) = &plan else {
3760+
panic!("expected the unpushable filter: {plan}");
3761+
};
3762+
let LogicalPlan::Join(join) = filter.input.as_ref() else {
3763+
panic!("expected a join: {plan}");
3764+
};
3765+
let LogicalPlan::SubqueryAlias(alias) = join.left.as_ref() else {
3766+
panic!("expected the left alias: {plan}");
3767+
};
3768+
let scan = if support == TableProviderFilterPushDown::Exact {
3769+
alias.input.as_ref()
3770+
} else {
3771+
let LogicalPlan::Filter(filter) = alias.input.as_ref() else {
3772+
panic!("expected the pushed filter: {plan}");
3773+
};
3774+
assert_eq!(filter.predicate, expected, "{plan}");
3775+
filter.input.as_ref()
3776+
};
3777+
let LogicalPlan::TableScan(scan) = scan else {
3778+
panic!("expected the left scan: {plan}");
3779+
};
3780+
let expected_scan_filters =
3781+
if support == TableProviderFilterPushDown::Unsupported {
3782+
vec![]
3783+
} else {
3784+
vec![expected.clone()]
3785+
};
3786+
assert_eq!(scan.filters, expected_scan_filters, "{plan}");
3787+
}
3788+
}
3789+
Ok(())
3790+
}
3791+
37023792
#[test]
37033793
fn left_semi_join() -> Result<()> {
37043794
let left = test_table_scan_with_name("test1")?;

0 commit comments

Comments
 (0)