zhuqi-lucas commented on code in PR #25865:
URL: https://github.com/apache/datafusion/pull/25865#discussion_r4227473649
##########
datafusion/optimizer/src/push_down_filter.rs:
##########
@@ -1975,21 +2044,449 @@ mod tests {
let plan = LogicalPlanBuilder::from(table_scan)
.window(vec![window])?
- // unlike with aggregations, single partition column "test.a +
test.b" is not available
- // to the plan, so we use multiple columns when filtering
+ // the single partition column "test.a + test.b" is not available
to the
+ // plan, so the predicate is written over the underlying columns
.filter(add(col("a"), col("b")).gt(lit(10i64)))?
.build()?;
+ assert_optimized_plan_equal!(
+ plan,
+ @r"
+ WindowAggr: windowExpr=[[rank() PARTITION BY [test.a + test.b] ORDER
BY [test.c ASC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]]
+ TableScan: test, full_filters=[test.a + test.b > Int64(10)]
+ "
+ )
+ }
+
+ /// verifies that a single predicate spanning a column key and an
expression
+ /// key is pushed: it is constant within every partition, but neither the
old
Review Comment:
Gone with that test.
##########
datafusion/optimizer/src/push_down_filter.rs:
##########
@@ -1975,21 +2044,449 @@ mod tests {
let plan = LogicalPlanBuilder::from(table_scan)
.window(vec![window])?
- // unlike with aggregations, single partition column "test.a +
test.b" is not available
- // to the plan, so we use multiple columns when filtering
+ // the single partition column "test.a + test.b" is not available
to the
+ // plan, so the predicate is written over the underlying columns
.filter(add(col("a"), col("b")).gt(lit(10i64)))?
.build()?;
+ assert_optimized_plan_equal!(
+ plan,
+ @r"
+ WindowAggr: windowExpr=[[rank() PARTITION BY [test.a + test.b] ORDER
BY [test.c ASC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]]
+ TableScan: test, full_filters=[test.a + test.b > Int64(10)]
+ "
+ )
+ }
+
+ /// verifies that a single predicate spanning a column key and an
expression
+ /// key is pushed: it is constant within every partition, but neither the
old
+ /// column-name matching nor a rule that looked only at expression keys
could
+ /// see that
+ #[test]
+ fn filter_move_window_mixed_column_and_expression_keys() -> Result<()> {
+ let table_scan = test_table_scan()?;
+
+ let window = Expr::from(WindowFunction::new(
+ WindowFunctionDefinition::WindowUDF(
+ datafusion_functions_window::rank::rank_udwf(),
+ ),
+ vec![],
+ ))
+ .partition_by(vec![col("a"), add(col("a"), col("b"))]) // PARTITION BY
a, a + b
+ .order_by(vec![col("c").sort(true, true)])
+ .build()
+ .unwrap();
+
+ let plan = LogicalPlanBuilder::from(table_scan)
+ .window(vec![window])?
+ .filter(
+ col("a")
+ .gt(lit(1i64))
+ .or(add(col("a"), col("b")).gt(lit(10i64))),
+ )?
+ .build()?;
+
+ assert_optimized_plan_equal!(
+ plan,
+ @r"
+ WindowAggr: windowExpr=[[rank() PARTITION BY [test.a, test.a + test.b]
ORDER BY [test.c ASC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT
ROW]]
+ TableScan: test, full_filters=[test.a > Int64(1) OR test.a + test.b
> Int64(10)]
+ "
+ )
+ }
+
+ /// verifies that an expression partition key does not make the columns it
+ /// reads pushable on their own: `a` alone is not constant within a
partition
+ /// of `a + b`
+ #[test]
+ fn filter_keep_window_column_underlying_expression_key() -> Result<()> {
+ let table_scan = test_table_scan()?;
+
+ let window = Expr::from(WindowFunction::new(
+ WindowFunctionDefinition::WindowUDF(
+ datafusion_functions_window::rank::rank_udwf(),
+ ),
+ vec![],
+ ))
+ .partition_by(vec![add(col("a"), col("b"))]) // PARTITION BY a + b
+ .order_by(vec![col("c").sort(true, true)])
+ .build()
+ .unwrap();
+
+ let plan = LogicalPlanBuilder::from(table_scan)
+ .window(vec![window])?
+ .filter(col("a").gt(lit(10i64)))?
+ .build()?;
+ assert_plan_not_transformed!(plan.clone());
+
+ assert_optimized_plan_equal!(
+ plan,
+ @r"
+ Filter: test.a > Int64(10)
+ WindowAggr: windowExpr=[[rank() PARTITION BY [test.a + test.b] ORDER
BY [test.c ASC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]]
+ TableScan: test
+ "
+ )
+ }
+
+ /// verifies that a predicate on a *function* expression key is pushed.
This is
+ /// the shape the rule exists for: `PARTITION BY NULLIF(c, '')` with a
predicate
+ /// on `NULLIF(c, '')`, where the key is not an arithmetic expression and
the
+ /// column it reads is not a key on its own.
+ #[test]
+ fn filter_move_window_function_expression_key() -> Result<()> {
+ let table_scan = test_table_scan()?;
+
+ let window = Expr::from(WindowFunction::new(
+ WindowFunctionDefinition::WindowUDF(
+ datafusion_functions_window::rank::rank_udwf(),
+ ),
+ vec![],
+ ))
+ .partition_by(vec![immutable_udf_call(col("c"))])
+ .order_by(vec![col("c").sort(true, true)])
+ .build()
+ .unwrap();
+
+ let plan = LogicalPlanBuilder::from(table_scan)
+ .window(vec![window])?
+ .filter(immutable_udf_call(col("c")).is_not_null())?
+ .build()?;
+
+ assert_optimized_plan_equal!(
+ plan,
+ @r"
+ WindowAggr: windowExpr=[[rank() PARTITION BY [TestScalarUDF(test.c)]
ORDER BY [test.c ASC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT
ROW]]
+ TableScan: test, full_filters=[TestScalarUDF(test.c) IS NOT NULL]
+ "
+ )
+ }
+
+ /// verifies that matching is structural, so a predicate written with the
+ /// operands in the other order does not match the key and stays above the
+ /// window. Conservative rather than wrong.
+ #[test]
+ fn filter_keep_window_expression_key_operand_order() -> Result<()> {
+ let table_scan = test_table_scan()?;
+
+ let window = Expr::from(WindowFunction::new(
+ WindowFunctionDefinition::WindowUDF(
+ datafusion_functions_window::rank::rank_udwf(),
+ ),
+ vec![],
+ ))
+ .partition_by(vec![add(col("a"), col("b"))]) // PARTITION BY a + b
+ .order_by(vec![col("c").sort(true, true)])
+ .build()
+ .unwrap();
+
+ let plan = LogicalPlanBuilder::from(table_scan)
+ .window(vec![window])?
+ .filter(add(col("b"), col("a")).gt(lit(10i64)))? // b + a, not a +
b
+ .build()?;
+ assert_plan_not_transformed!(plan.clone());
+
+ assert_optimized_plan_equal!(
+ plan,
+ @r"
+ Filter: test.b + test.a > Int64(10)
+ WindowAggr: windowExpr=[[rank() PARTITION BY [test.a + test.b] ORDER
BY [test.c ASC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]]
+ TableScan: test
+ "
+ )
+ }
+
+ /// verifies that a predicate built *on top of* an expression key is
pushed:
+ /// the key subtree accounts for everything it reads, and the literal
around it
+ /// reads nothing, so the predicate is still constant within each partition
+ #[test]
+ fn filter_move_window_expression_key_subtree() -> Result<()> {
+ let table_scan = test_table_scan()?;
+
+ let window = Expr::from(WindowFunction::new(
+ WindowFunctionDefinition::WindowUDF(
+ datafusion_functions_window::rank::rank_udwf(),
+ ),
+ vec![],
+ ))
+ .partition_by(vec![add(col("a"), col("b"))]) // PARTITION BY a + b
+ .order_by(vec![col("c").sort(true, true)])
+ .build()
+ .unwrap();
+
+ let plan = LogicalPlanBuilder::from(table_scan)
+ .window(vec![window])?
+ .filter(add(add(col("a"), col("b")), lit(1i64)).gt(lit(10i64)))?
Review Comment:
Now in the helper doc and as the IS NOT NULL case in the slt.
##########
datafusion/optimizer/src/push_down_filter.rs:
##########
@@ -1975,21 +2044,449 @@ mod tests {
let plan = LogicalPlanBuilder::from(table_scan)
.window(vec![window])?
- // unlike with aggregations, single partition column "test.a +
test.b" is not available
- // to the plan, so we use multiple columns when filtering
+ // the single partition column "test.a + test.b" is not available
to the
+ // plan, so the predicate is written over the underlying columns
.filter(add(col("a"), col("b")).gt(lit(10i64)))?
.build()?;
+ assert_optimized_plan_equal!(
+ plan,
+ @r"
+ WindowAggr: windowExpr=[[rank() PARTITION BY [test.a + test.b] ORDER
BY [test.c ASC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]]
+ TableScan: test, full_filters=[test.a + test.b > Int64(10)]
+ "
+ )
+ }
+
+ /// verifies that a single predicate spanning a column key and an
expression
+ /// key is pushed: it is constant within every partition, but neither the
old
+ /// column-name matching nor a rule that looked only at expression keys
could
+ /// see that
+ #[test]
+ fn filter_move_window_mixed_column_and_expression_keys() -> Result<()> {
+ let table_scan = test_table_scan()?;
+
+ let window = Expr::from(WindowFunction::new(
+ WindowFunctionDefinition::WindowUDF(
+ datafusion_functions_window::rank::rank_udwf(),
+ ),
+ vec![],
+ ))
+ .partition_by(vec![col("a"), add(col("a"), col("b"))]) // PARTITION BY
a, a + b
+ .order_by(vec![col("c").sort(true, true)])
+ .build()
+ .unwrap();
+
+ let plan = LogicalPlanBuilder::from(table_scan)
+ .window(vec![window])?
+ .filter(
+ col("a")
+ .gt(lit(1i64))
+ .or(add(col("a"), col("b")).gt(lit(10i64))),
+ )?
+ .build()?;
+
+ assert_optimized_plan_equal!(
+ plan,
+ @r"
+ WindowAggr: windowExpr=[[rank() PARTITION BY [test.a, test.a + test.b]
ORDER BY [test.c ASC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT
ROW]]
+ TableScan: test, full_filters=[test.a > Int64(1) OR test.a + test.b
> Int64(10)]
+ "
+ )
+ }
+
+ /// verifies that an expression partition key does not make the columns it
+ /// reads pushable on their own: `a` alone is not constant within a
partition
+ /// of `a + b`
+ #[test]
+ fn filter_keep_window_column_underlying_expression_key() -> Result<()> {
+ let table_scan = test_table_scan()?;
+
+ let window = Expr::from(WindowFunction::new(
+ WindowFunctionDefinition::WindowUDF(
+ datafusion_functions_window::rank::rank_udwf(),
+ ),
+ vec![],
+ ))
+ .partition_by(vec![add(col("a"), col("b"))]) // PARTITION BY a + b
+ .order_by(vec![col("c").sort(true, true)])
+ .build()
+ .unwrap();
+
+ let plan = LogicalPlanBuilder::from(table_scan)
+ .window(vec![window])?
+ .filter(col("a").gt(lit(10i64)))?
+ .build()?;
+ assert_plan_not_transformed!(plan.clone());
+
+ assert_optimized_plan_equal!(
+ plan,
+ @r"
+ Filter: test.a > Int64(10)
+ WindowAggr: windowExpr=[[rank() PARTITION BY [test.a + test.b] ORDER
BY [test.c ASC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]]
+ TableScan: test
+ "
+ )
+ }
+
+ /// verifies that a predicate on a *function* expression key is pushed.
This is
+ /// the shape the rule exists for: `PARTITION BY NULLIF(c, '')` with a
predicate
+ /// on `NULLIF(c, '')`, where the key is not an arithmetic expression and
the
+ /// column it reads is not a key on its own.
+ #[test]
+ fn filter_move_window_function_expression_key() -> Result<()> {
+ let table_scan = test_table_scan()?;
+
+ let window = Expr::from(WindowFunction::new(
+ WindowFunctionDefinition::WindowUDF(
+ datafusion_functions_window::rank::rank_udwf(),
+ ),
+ vec![],
+ ))
+ .partition_by(vec![immutable_udf_call(col("c"))])
+ .order_by(vec![col("c").sort(true, true)])
+ .build()
+ .unwrap();
+
+ let plan = LogicalPlanBuilder::from(table_scan)
+ .window(vec![window])?
+ .filter(immutable_udf_call(col("c")).is_not_null())?
+ .build()?;
+
+ assert_optimized_plan_equal!(
+ plan,
+ @r"
+ WindowAggr: windowExpr=[[rank() PARTITION BY [TestScalarUDF(test.c)]
ORDER BY [test.c ASC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT
ROW]]
+ TableScan: test, full_filters=[TestScalarUDF(test.c) IS NOT NULL]
+ "
+ )
+ }
+
+ /// verifies that matching is structural, so a predicate written with the
+ /// operands in the other order does not match the key and stays above the
+ /// window. Conservative rather than wrong.
+ #[test]
+ fn filter_keep_window_expression_key_operand_order() -> Result<()> {
+ let table_scan = test_table_scan()?;
+
+ let window = Expr::from(WindowFunction::new(
+ WindowFunctionDefinition::WindowUDF(
+ datafusion_functions_window::rank::rank_udwf(),
+ ),
+ vec![],
+ ))
+ .partition_by(vec![add(col("a"), col("b"))]) // PARTITION BY a + b
+ .order_by(vec![col("c").sort(true, true)])
+ .build()
+ .unwrap();
+
+ let plan = LogicalPlanBuilder::from(table_scan)
+ .window(vec![window])?
+ .filter(add(col("b"), col("a")).gt(lit(10i64)))? // b + a, not a +
b
+ .build()?;
+ assert_plan_not_transformed!(plan.clone());
+
+ assert_optimized_plan_equal!(
+ plan,
+ @r"
+ Filter: test.b + test.a > Int64(10)
+ WindowAggr: windowExpr=[[rank() PARTITION BY [test.a + test.b] ORDER
BY [test.c ASC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]]
+ TableScan: test
+ "
+ )
+ }
+
+ /// verifies that a predicate built *on top of* an expression key is
pushed:
+ /// the key subtree accounts for everything it reads, and the literal
around it
+ /// reads nothing, so the predicate is still constant within each partition
+ #[test]
+ fn filter_move_window_expression_key_subtree() -> Result<()> {
+ let table_scan = test_table_scan()?;
+
+ let window = Expr::from(WindowFunction::new(
+ WindowFunctionDefinition::WindowUDF(
+ datafusion_functions_window::rank::rank_udwf(),
+ ),
+ vec![],
+ ))
+ .partition_by(vec![add(col("a"), col("b"))]) // PARTITION BY a + b
+ .order_by(vec![col("c").sort(true, true)])
+ .build()
+ .unwrap();
+
+ let plan = LogicalPlanBuilder::from(table_scan)
+ .window(vec![window])?
+ .filter(add(add(col("a"), col("b")), lit(1i64)).gt(lit(10i64)))?
+ .build()?;
+
+ assert_optimized_plan_equal!(
+ plan,
+ @r"
+ WindowAggr: windowExpr=[[rank() PARTITION BY [test.a + test.b] ORDER
BY [test.c ASC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]]
+ TableScan: test, full_filters=[test.a + test.b + Int64(1) >
Int64(10)]
+ "
+ )
+ }
+
+ /// verifies that an expression key shared by every window function is
pushed,
Review Comment:
Added: a + b > 5 OR a + b < 3 under PARTITION BY a + b, pushed.
##########
datafusion/optimizer/src/push_down_filter.rs:
##########
@@ -1135,8 +1135,13 @@ impl OptimizerRule for PushDownFilter {
// multiple window functions, each with potentially different
partition keys.
// Therefore, we need to ensure that any potential partition
key returned is used in
// ALL window functions. Otherwise, filters cannot be pushed
by through that column.
- fn extract_partition_keys(func: &WindowFunction) ->
HashSet<Column> {
- expr_columns(&func.params.partition_by)
+ // Keyed by the partition *expression*, not by a name
synthesised
+ // from it. `PARTITION BY a + b` used to be mapped through
+ // `qualified_name()` to a column literally called "a + b",
which no
+ // predicate's real column refs could ever match, so such a
key was
+ // dead weight in this set.
+ fn extract_partition_keys(func: &WindowFunction) ->
HashSet<Expr> {
+ func.params.partition_by.iter().cloned().collect()
Review Comment:
Done, HashSet<&Expr>.
##########
datafusion/optimizer/src/push_down_filter.rs:
##########
@@ -1108,21 +1115,27 @@ impl OptimizerRule for PushDownFilter {
let mut keep_predicates = vec![];
let mut push_predicates = vec![];
for expr in predicates {
- let cols = expr.column_refs();
- if cols.iter().all(|c|
potential_partition_keys.contains(c)) {
+ // A volatile predicate has to stay above the window.
Pushing it
+ // changes which rows the window function sees, and so the
value
+ // it computes for the rows that do survive. Checking this
first
+ // also covers a volatile predicate that reads no columns
at all,
+ // such as `random() < 0.5`, which would otherwise satisfy
the
+ // partition-key test vacuously.
+ if !expr.is_volatile()
+ && reads_only_partition_keys(&expr,
&potential_partition_keys)?
+ {
Review Comment:
Addressed in eb779972d; the IN and scalar subquery cases in the slt stay
above the window.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
