Skip to content

Commit 93052cd

Browse files
ygf11ygf11
authored
Refactor extract_join_keys and move the ExtractEquijoinPredicate rule (#4760)
* Refactor extract_join_keys with split_conjunction * reorder the ExtractEquijoinPredicate rule * fix cargo test * trigger ci aggain * trigger ci aggain * add integration test * fix test
1 parent 0d6d371 commit 93052cd

3 files changed

Lines changed: 163 additions & 70 deletions

File tree

‎datafusion/core/tests/dataframe.rs‎

Lines changed: 98 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -17,21 +17,20 @@
1717

1818
use arrow::datatypes::{DataType, Field, Schema};
1919
use arrow::{
20-
array::{Int32Array, StringArray},
20+
array::{Int32Array, StringArray, UInt32Array},
2121
record_batch::RecordBatch,
2222
};
2323
use datafusion::from_slice::FromSlice;
2424
use std::sync::Arc;
2525

26-
use datafusion::assert_batches_eq;
2726
use datafusion::dataframe::DataFrame;
2827
use datafusion::error::Result;
2928
use datafusion::execution::context::SessionContext;
3029
use datafusion::prelude::CsvReadOptions;
3130
use datafusion::prelude::JoinType;
31+
use datafusion::{assert_batches_eq, assert_batches_sorted_eq};
3232
use datafusion_expr::expr::{GroupingSet, Sort};
33-
use datafusion_expr::{avg, count, lit, sum};
34-
use datafusion_expr::{col, Expr};
33+
use datafusion_expr::{avg, col, count, lit, sum, Expr, ExprSchemable};
3534

3635
#[tokio::test]
3736
async fn join() -> Result<()> {
@@ -352,6 +351,62 @@ async fn test_grouping_set_array_agg_with_overflow() -> Result<()> {
352351
Ok(())
353352
}
354353

354+
#[tokio::test]
355+
async fn join_with_alias_filter() -> Result<()> {
356+
let join_ctx = create_join_context()?;
357+
let t1 = join_ctx.table("t1")?;
358+
let t2 = join_ctx.table("t2")?;
359+
let t1_schema = t1.schema().clone();
360+
let t2_schema = t2.schema().clone();
361+
362+
// filter: t1.a + CAST(Int64(1), UInt32) = t2.a + CAST(Int64(2), UInt32) as t1.a + 1 = t2.a + 2
363+
let filter = Expr::eq(
364+
col("t1.a") + lit(3i64).cast_to(&DataType::UInt32, &t1_schema)?,
365+
col("t2.a") + lit(1i32).cast_to(&DataType::UInt32, &t2_schema)?,
366+
)
367+
.alias("t1.b + 1 = t2.a + 2");
368+
369+
let df = t1
370+
.join(t2, JoinType::Inner, &[], &[], Some(filter))?
371+
.select(vec![
372+
col("t1.a"),
373+
col("t2.a"),
374+
col("t1.b"),
375+
col("t1.c"),
376+
col("t2.b"),
377+
col("t2.c"),
378+
])?;
379+
let optimized_plan = df.clone().into_optimized_plan()?;
380+
381+
let expected = vec![
382+
"Projection: t1.a, t2.a, t1.b, t1.c, t2.b, t2.c [a:UInt32, a:UInt32, b:Utf8, c:Int32, b:Utf8, c:Int32]",
383+
" Inner Join: t1.a + UInt32(3) = t2.a + UInt32(1) [a:UInt32, b:Utf8, c:Int32, a:UInt32, b:Utf8, c:Int32]",
384+
" TableScan: t1 projection=[a, b, c] [a:UInt32, b:Utf8, c:Int32]",
385+
" TableScan: t2 projection=[a, b, c] [a:UInt32, b:Utf8, c:Int32]",
386+
];
387+
388+
let formatted = optimized_plan.display_indent_schema().to_string();
389+
let actual: Vec<&str> = formatted.trim().lines().collect();
390+
assert_eq!(
391+
expected, actual,
392+
"\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n"
393+
);
394+
395+
let results = df.collect().await?;
396+
let expected: Vec<&str> = vec![
397+
"+----+----+---+----+---+---+",
398+
"| a | a | b | c | b | c |",
399+
"+----+----+---+----+---+---+",
400+
"| 11 | 13 | c | 30 | c | 3 |",
401+
"| 1 | 3 | a | 10 | a | 1 |",
402+
"+----+----+---+----+---+---+",
403+
];
404+
405+
assert_batches_sorted_eq!(expected, &results);
406+
407+
Ok(())
408+
}
409+
355410
fn create_test_table() -> Result<DataFrame> {
356411
let schema = Arc::new(Schema::new(vec![
357412
Field::new("a", DataType::Utf8, false),
@@ -388,3 +443,42 @@ async fn aggregates_table(ctx: &SessionContext) -> Result<DataFrame> {
388443
)
389444
.await
390445
}
446+
447+
fn create_join_context() -> Result<SessionContext> {
448+
let t1 = Arc::new(Schema::new(vec![
449+
Field::new("a", DataType::UInt32, false),
450+
Field::new("b", DataType::Utf8, false),
451+
Field::new("c", DataType::Int32, false),
452+
]));
453+
let t2 = Arc::new(Schema::new(vec![
454+
Field::new("a", DataType::UInt32, false),
455+
Field::new("b", DataType::Utf8, false),
456+
Field::new("c", DataType::Int32, false),
457+
]));
458+
459+
// define data.
460+
let batch1 = RecordBatch::try_new(
461+
t1,
462+
vec![
463+
Arc::new(UInt32Array::from_slice([1, 10, 11, 100])),
464+
Arc::new(StringArray::from_slice(["a", "b", "c", "d"])),
465+
Arc::new(Int32Array::from_slice([10, 20, 30, 40])),
466+
],
467+
)?;
468+
// define data.
469+
let batch2 = RecordBatch::try_new(
470+
t2,
471+
vec![
472+
Arc::new(UInt32Array::from_slice([3, 10, 13, 100])),
473+
Arc::new(StringArray::from_slice(["a", "b", "c", "d"])),
474+
Arc::new(Int32Array::from_slice([1, 2, 3, 4])),
475+
],
476+
)?;
477+
478+
let ctx = SessionContext::new();
479+
480+
ctx.register_batch("t1", batch1)?;
481+
ctx.register_batch("t2", batch2)?;
482+
483+
Ok(ctx)
484+
}

‎datafusion/optimizer/src/extract_equijoin_predicate.rs‎

Lines changed: 64 additions & 65 deletions
Original file line numberDiff line numberDiff line change
@@ -17,13 +17,17 @@
1717

1818
//! Optimizer rule to extract equijoin expr from filter
1919
use crate::optimizer::ApplyOrder;
20+
use crate::utils::split_conjunction;
2021
use crate::{OptimizerConfig, OptimizerRule};
2122
use datafusion_common::DFSchema;
2223
use datafusion_common::Result;
2324
use datafusion_expr::utils::{can_hash, find_valid_equijoin_key_pair};
2425
use datafusion_expr::{BinaryExpr, Expr, ExprSchemable, Join, LogicalPlan, Operator};
2526
use std::sync::Arc;
2627

28+
// equijoin predicate
29+
type EquijoinPredicate = (Expr, Expr);
30+
2731
/// Optimization rule that extract equijoin expr from the filter
2832
#[derive(Default)]
2933
pub struct ExtractEquijoinPredicate;
@@ -56,27 +60,22 @@ impl OptimizerRule for ExtractEquijoinPredicate {
5660
let right_schema = right.schema();
5761

5862
filter.as_ref().map_or(Result::Ok(None), |expr| {
59-
let mut accum: Vec<(Expr, Expr)> = vec![];
60-
let mut accum_filter: Vec<Expr> = vec![];
61-
// TODO: avoding clone with split_conjunction
62-
extract_join_keys(
63-
expr.clone(),
64-
&mut accum,
65-
&mut accum_filter,
66-
left_schema,
67-
right_schema,
68-
)?;
69-
70-
let optimized_plan = (!accum.is_empty()).then(|| {
63+
let (equijoin_predicates, non_equijoin_expr) =
64+
split_eq_and_noneq_join_predicate(
65+
expr,
66+
left_schema,
67+
right_schema,
68+
)?;
69+
70+
let optimized_plan = (!equijoin_predicates.is_empty()).then(|| {
7171
let mut new_on = on.clone();
72-
new_on.extend(accum);
72+
new_on.extend(equijoin_predicates);
7373

74-
let new_filter = accum_filter.into_iter().reduce(Expr::and);
7574
LogicalPlan::Join(Join {
7675
left: left.clone(),
7776
right: right.clone(),
7877
on: new_on,
79-
filter: new_filter,
78+
filter: non_equijoin_expr,
8079
join_type: *join_type,
8180
join_constraint: *join_constraint,
8281
schema: schema.clone(),
@@ -100,30 +99,22 @@ impl OptimizerRule for ExtractEquijoinPredicate {
10099
}
101100
}
102101

103-
/// Extracts equijoin ON condition be a single Eq or multiple conjunctive Eqs
104-
/// Filters matching this pattern are added to `accum`
105-
/// Filters that don't match this pattern are added to `accum_filter`
106-
/// Examples:
107-
/// ```text
108-
/// foo = bar => accum=[(foo, bar)] accum_filter=[]
109-
/// foo = bar AND bar = baz => accum=[(foo, bar), (bar, baz)] accum_filter=[]
110-
/// foo = bar AND baz > 1 => accum=[(foo, bar)] accum_filter=[baz > 1]
111-
///
112-
/// For equijoin join key, assume we have tables -- a(c0, c1 c2) and b(c0, c1, c2):
113-
/// (a.c0 = 10) => accum=[], accum_filter=[a.c0 = 10]
114-
/// (a.c0 + 1 = b.c0 * 2) => accum=[(a.c0 + 1, b.c0 * 2)], accum_filter=[]
115-
/// (a.c0 + b.c0 = 10) => accum=[], accum_filter=[a.c0 + b.c0 = 10]
116-
/// ```
117-
fn extract_join_keys(
118-
expr: Expr,
119-
accum: &mut Vec<(Expr, Expr)>,
120-
accum_filter: &mut Vec<Expr>,
102+
fn split_eq_and_noneq_join_predicate(
103+
filter: &Expr,
121104
left_schema: &Arc<DFSchema>,
122105
right_schema: &Arc<DFSchema>,
123-
) -> Result<()> {
124-
match &expr {
125-
Expr::BinaryExpr(BinaryExpr { left, op, right }) => match op {
126-
Operator::Eq => {
106+
) -> Result<(Vec<EquijoinPredicate>, Option<Expr>)> {
107+
let exprs = split_conjunction(filter);
108+
109+
let mut accum_join_keys: Vec<(Expr, Expr)> = vec![];
110+
let mut accum_filters: Vec<Expr> = vec![];
111+
for expr in exprs {
112+
match expr {
113+
Expr::BinaryExpr(BinaryExpr {
114+
left,
115+
op: Operator::Eq,
116+
right,
117+
}) => {
127118
let left = left.as_ref();
128119
let right = right.as_ref();
129120

@@ -139,48 +130,27 @@ fn extract_join_keys(
139130
let right_expr_type = right_expr.get_type(right_schema)?;
140131

141132
if can_hash(&left_expr_type) && can_hash(&right_expr_type) {
142-
accum.push((left_expr, right_expr));
133+
accum_join_keys.push((left_expr, right_expr));
143134
} else {
144-
accum_filter.push(expr);
135+
accum_filters.push(expr.clone());
145136
}
146137
} else {
147-
accum_filter.push(expr);
148-
}
149-
}
150-
Operator::And => {
151-
if let Expr::BinaryExpr(BinaryExpr { left, op: _, right }) = expr {
152-
extract_join_keys(
153-
*left,
154-
accum,
155-
accum_filter,
156-
left_schema,
157-
right_schema,
158-
)?;
159-
extract_join_keys(
160-
*right,
161-
accum,
162-
accum_filter,
163-
left_schema,
164-
right_schema,
165-
)?;
138+
accum_filters.push(expr.clone());
166139
}
167140
}
168-
_other => {
169-
accum_filter.push(expr);
170-
}
171-
},
172-
_other => {
173-
accum_filter.push(expr);
141+
_ => accum_filters.push(expr.clone()),
174142
}
175143
}
176144

177-
Ok(())
145+
let result_filter = accum_filters.into_iter().reduce(Expr::and);
146+
Ok((accum_join_keys, result_filter))
178147
}
179148

180149
#[cfg(test)]
181150
mod tests {
182151
use super::*;
183152
use crate::test::*;
153+
use arrow::datatypes::DataType;
184154
use datafusion_common::Column;
185155
use datafusion_expr::{
186156
col, lit, logical_plan::builder::LogicalPlanBuilder, JoinType,
@@ -387,4 +357,33 @@ mod tests {
387357

388358
assert_plan_eq(&plan, expected)
389359
}
360+
361+
#[test]
362+
fn join_with_alias_filter() -> Result<()> {
363+
let t1 = test_table_scan_with_name("t1")?;
364+
let t2 = test_table_scan_with_name("t2")?;
365+
366+
let t1_schema = t1.schema().clone();
367+
let t2_schema = t2.schema().clone();
368+
369+
// filter: t1.a + CAST(Int64(1), UInt32) = t2.a + CAST(Int64(2), UInt32) as t1.a + 1 = t2.a + 2
370+
let filter = Expr::eq(
371+
col("t1.a") + lit(1i64).cast_to(&DataType::UInt32, &t1_schema)?,
372+
col("t2.a") + lit(2i32).cast_to(&DataType::UInt32, &t2_schema)?,
373+
)
374+
.alias("t1.a + 1 = t2.a + 2");
375+
let plan = LogicalPlanBuilder::from(t1)
376+
.join(
377+
t2,
378+
JoinType::Left,
379+
(Vec::<Column>::new(), Vec::<Column>::new()),
380+
Some(filter),
381+
)?
382+
.build()?;
383+
let expected = "Left Join: t1.a + CAST(Int64(1) AS UInt32) = t2.a + CAST(Int32(2) AS UInt32) [a:UInt32, b:UInt32, c:UInt32, a:UInt32, b:UInt32, c:UInt32]\
384+
\n TableScan: t1 [a:UInt32, b:UInt32, c:UInt32]\
385+
\n TableScan: t2 [a:UInt32, b:UInt32, c:UInt32]";
386+
387+
assert_plan_eq(&plan, expected)
388+
}
390389
}

‎datafusion/optimizer/src/optimizer.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -222,12 +222,12 @@ impl Optimizer {
222222
let rules: Vec<Arc<dyn OptimizerRule + Sync + Send>> = vec![
223223
Arc::new(InlineTableScan::new()),
224224
Arc::new(TypeCoercion::new()),
225-
Arc::new(ExtractEquijoinPredicate::new()),
226225
Arc::new(SimplifyExpressions::new()),
227226
Arc::new(UnwrapCastInComparison::new()),
228227
Arc::new(DecorrelateWhereExists::new()),
229228
Arc::new(DecorrelateWhereIn::new()),
230229
Arc::new(ScalarSubqueryToJoin::new()),
230+
Arc::new(ExtractEquijoinPredicate::new()),
231231
// simplify expressions does not simplify expressions in subqueries, so we
232232
// run it again after running the optimizations that potentially converted
233233
// subqueries to joins

0 commit comments

Comments
 (0)