This is an automated email from the ASF dual-hosted git repository.
agrove pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git
The following commit(s) were added to refs/heads/master by this push:
new 748b6a65a Rewrite subexpressions of InSubquery in rewrite_expression
(#2765)
748b6a65a is described below
commit 748b6a65a5fa801595fd80a3c7b2728be3c9cdaa
Author: Mike Roberts <[email protected]>
AuthorDate: Wed Jun 22 15:02:36 2022 +0100
Rewrite subexpressions of InSubquery in rewrite_expression (#2765)
* Visit sub-expressions of Expr::InSubquery in ExpressionVisitor
* Rewrite subexpression of Expr::InSubquery in rewrite_expression
* Add test case for predicate pushdown of InSubquery filters, with a column
alias
---
datafusion/expr/src/expr_visitor.rs | 2 +-
datafusion/optimizer/src/filter_push_down.rs | 35 +++++++++++++++++++++++++++-
datafusion/optimizer/src/utils.rs | 8 ++++++-
3 files changed, 42 insertions(+), 3 deletions(-)
diff --git a/datafusion/expr/src/expr_visitor.rs
b/datafusion/expr/src/expr_visitor.rs
index 05bc3bbfa..162db60a0 100644
--- a/datafusion/expr/src/expr_visitor.rs
+++ b/datafusion/expr/src/expr_visitor.rs
@@ -102,6 +102,7 @@ impl ExprVisitable for Expr {
| Expr::Cast { expr, .. }
| Expr::TryCast { expr, .. }
| Expr::Sort { expr, .. }
+ | Expr::InSubquery { expr, .. }
| Expr::GetIndexedField { expr, .. } => expr.accept(visitor),
Expr::GroupingSet(GroupingSet::Rollup(exprs)) => exprs
.iter()
@@ -120,7 +121,6 @@ impl ExprVisitable for Expr {
| Expr::ScalarVariable(_, _)
| Expr::Literal(_)
| Expr::Exists { .. }
- | Expr::InSubquery { .. }
| Expr::ScalarSubquery(_)
| Expr::Wildcard
| Expr::QualifiedWildcard { .. } => Ok(visitor),
diff --git a/datafusion/optimizer/src/filter_push_down.rs
b/datafusion/optimizer/src/filter_push_down.rs
index 561b6b36b..eded1af3e 100644
--- a/datafusion/optimizer/src/filter_push_down.rs
+++ b/datafusion/optimizer/src/filter_push_down.rs
@@ -654,7 +654,7 @@ mod tests {
use async_trait::async_trait;
use datafusion_common::DFSchema;
use datafusion_expr::{
- and, col, in_list, lit,
+ and, col, in_list, in_subquery, lit,
logical_plan::{builder::union_with_alias, JoinType},
sum, Expr, LogicalPlanBuilder, Operator, TableSource, TableType,
};
@@ -2045,4 +2045,37 @@ mod tests {
Ok(())
}
+
+ #[test]
+ fn test_in_subquery_with_alias() -> Result<()> {
+ // in table scan the true col name is 'test.a',
+ // but we rename it as 'b', and use col 'b' in subquery filter
+ let table_scan = test_table_scan()?;
+ let table_scan_sq = test_table_scan_with_name("sq")?;
+ let subplan = Arc::new(
+ LogicalPlanBuilder::from(table_scan_sq)
+ .project(vec![col("c")])?
+ .build()?,
+ );
+ let plan = LogicalPlanBuilder::from(table_scan)
+ .project(vec![col("a").alias("b"), col("c")])?
+ .filter(in_subquery(col("b"), subplan))?
+ .build()?;
+
+ // filter on col b in subquery
+ let expected_before = "\
+ Filter: #b IN (Subquery: Projection: #sq.c\n TableScan: sq
projection=None)\
+ \n Projection: #test.a AS b, #test.c\
+ \n TableScan: test projection=None";
+ assert_eq!(format!("{:?}", plan), expected_before);
+
+ // rewrite filter col b to test.a
+ let expected_after = "\
+ Projection: #test.a AS b, #test.c\
+ \n Filter: #test.a IN (Subquery: Projection: #sq.c\n TableScan: sq
projection=None)\
+ \n TableScan: test projection=None";
+ assert_optimized_plan_eq(&plan, expected_after);
+
+ Ok(())
+ }
}
diff --git a/datafusion/optimizer/src/utils.rs
b/datafusion/optimizer/src/utils.rs
index 691d8b663..dfde370a0 100644
--- a/datafusion/optimizer/src/utils.rs
+++ b/datafusion/optimizer/src/utils.rs
@@ -281,10 +281,16 @@ pub fn rewrite_expression(expr: &Expr, expressions:
&[Expr]) -> Result<Expr> {
list: list.clone(),
negated: *negated,
}),
+ Expr::InSubquery {
+ subquery, negated, ..
+ } => Ok(Expr::InSubquery {
+ expr: Box::new(expressions[0].clone()),
+ subquery: subquery.clone(),
+ negated: *negated,
+ }),
Expr::Column(_)
| Expr::Literal(_)
| Expr::Exists { .. }
- | Expr::InSubquery { .. }
| Expr::ScalarSubquery(_)
| Expr::ScalarVariable(_, _) => Ok(expr.clone()),
Expr::Sort {
