Ted-Jiang commented on code in PR #2729:
URL: https://github.com/apache/arrow-datafusion/pull/2729#discussion_r897613300
##########
datafusion/optimizer/src/filter_push_down.rs:
##########
@@ -1831,4 +1831,218 @@ mod tests {
Ok(())
}
+
+ #[test]
+ fn test_filter_with_alias() -> Result<()> {
+ // in table scan the true col name is 'test.a',
+ // but we rename it as 'b', and use col 'b' in filter
+ // we need rewrite filter col before push down.
+ let table_scan = test_table_scan()?;
+ let plan = LogicalPlanBuilder::from(table_scan)
+ .project(vec![col("a").alias("b"), col("c")])?
+ .filter(and(col("b").gt(lit(10i64)), col("c").gt(lit(10i64))))?
+ .build()?;
+
+ // filter on col b
+ assert_eq!(
+ format!("{:?}", plan),
+ "\
+ Filter: #b > Int64(10) AND #test.c > Int64(10)\
+ \n Projection: #test.a AS b, #test.c\
+ \n TableScan: test projection=None\
+ "
+ );
+
+ // rewrite filter col b to test.a
+ let expected = "\
+ Projection: #test.a AS b, #test.c\
+ \n Filter: #test.a > Int64(10) AND #test.c > Int64(10)\
+ \n TableScan: test projection=None\
+ ";
+
+ assert_optimized_plan_eq(&plan, expected);
+
+ Ok(())
+ }
+
+ #[test]
+ fn test_filter_with_alias_2() -> Result<()> {
+ // in table scan the true col name is 'test.a',
+ // but we rename it as 'b', and use col 'b' in filter
+ // we need rewrite filter col before push down.
+ let table_scan = test_table_scan()?;
+ let plan = LogicalPlanBuilder::from(table_scan)
+ .project(vec![col("a").alias("b"), col("c")])?
+ .project(vec![col("b"), col("c")])?
+ .filter(and(col("b").gt(lit(10i64)), col("c").gt(lit(10i64))))?
+ .build()?;
+
+ // filter on col b
+ assert_eq!(
+ format!("{:?}", plan),
+ "\
+ Filter: #b > Int64(10) AND #test.c > Int64(10)\
+ \n Projection: #b, #test.c\
+ \n Projection: #test.a AS b, #test.c\
+ \n TableScan: test projection=None\
+ "
+ );
+
+ // rewrite filter col b to test.a
+ let expected = "\
+ Projection: #b, #test.c\
+ \n Projection: #test.a AS b, #test.c\
+ \n Filter: #test.a > Int64(10) AND #test.c > Int64(10)\
+ \n TableScan: test projection=None\
+ ";
+
+ assert_optimized_plan_eq(&plan, expected);
+
+ Ok(())
+ }
+
+ #[test]
+ fn test_filter_with_multi_alias() -> Result<()> {
+ let table_scan = test_table_scan()?;
+ let plan = LogicalPlanBuilder::from(table_scan)
+ .project(vec![col("a").alias("b"), col("c").alias("d")])?
+ .filter(and(col("b").gt(lit(10i64)), col("d").gt(lit(10i64))))?
+ .build()?;
+
+ // filter on col b and d
+ assert_eq!(
+ format!("{:?}", plan),
+ "\
+ Filter: #b > Int64(10) AND #d > Int64(10)\
+ \n Projection: #test.a AS b, #test.c AS d\
+ \n TableScan: test projection=None\
+ "
+ );
+
+ // rewrite filter col b to test.a, col d to test.c
+ let expected = "\
+ Projection: #test.a AS b, #test.c AS d\
+ \n Filter: #test.a > Int64(10) AND #test.c > Int64(10)\
+ \n TableScan: test projection=None\
+ ";
+
+ assert_optimized_plan_eq(&plan, expected);
+
+ Ok(())
+ }
+
+ /// predicate on join key in filter expression should be pushed down to
both inputs
+ #[test]
+ fn join_filter_with_alias() -> Result<()> {
+ let table_scan = test_table_scan()?;
+ let left = LogicalPlanBuilder::from(table_scan)
+ .project(vec![col("a").alias("c")])?
+ .build()?;
+ let right_table_scan = test_table_scan_with_name("test2")?;
+ let right = LogicalPlanBuilder::from(right_table_scan)
+ .project(vec![col("b").alias("d")])?
+ .build()?;
+ let filter = col("c").gt(lit(1u32));
+ let plan = LogicalPlanBuilder::from(left)
+ .join(
+ &right,
+ JoinType::Inner,
+ (vec![Column::from_name("c")], vec![Column::from_name("d")]),
+ Some(filter),
+ )?
+ .build()?;
+
+ assert_eq!(
+ format!("{:?}", plan),
+ "\
+ Inner Join: #c = #d Filter: #c > UInt32(1)\
+ \n Projection: #test.a AS c\
+ \n TableScan: test projection=None\
+ \n Projection: #test2.b AS d\
+ \n TableScan: test2 projection=None"
+ );
+
+ // Change filter on col `c`, 'd' to `test.a`, 'test.b'
+ let expected = "\
+ Inner Join: #c = #d\
+ \n Projection: #test.a AS c\
+ \n Filter: #test.a > UInt32(1)\
+ \n TableScan: test projection=None\
+ \n Projection: #test2.b AS d\
+ \n Filter: #test2.b > UInt32(1)\
+ \n TableScan: test2 projection=None";
+ assert_optimized_plan_eq(&plan, expected);
+ Ok(())
+ }
+
+ #[test]
+ fn test_in_filter_with_alias() -> Result<()> {
Review Comment:
This test will fail without code change.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscr...@arrow.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
