alamb commented on code in PR #25865:
URL: https://github.com/apache/datafusion/pull/25865#discussion_r4221307941
##########
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!(
Review Comment:
Is there a reason we can't use .slt for these tests? The setup code for
buulding the dataframe is quite verbose compare to the same thing as SQL, and I
found it hard to map from the dataframe API expressions to the SQL level
expressions (e.g. `add(col(a), col(b))` in the setup code, but `a + b` in the
explain)
I am wondering if there is some way to reduce the size of this PR
##########
datafusion/sqllogictest/test_files/push_down_filter_regression.slt:
##########
@@ -723,3 +723,92 @@ query I
SELECT sum(c) FROM (SELECT random() < 0.5 AS k, count(*) AS c FROM
generate_series(1, 10000) GROUP BY random() < 0.5) WHERE k OR NOT k;
----
10000
+
+# Window filter pushdown over an expression PARTITION BY key.
+#
+# A predicate that reads only the partition keys is constant within a
partition,
+# so pushing it below the window drops whole partitions and leaves every
+# surviving row's window value alone. The unit tests pin where the filter
lands;
Review Comment:
> The unit tests pin where the filter lands;
This (and the stuff below it) I think is pretty obvious (aka these are
explain tests so changes to where the filter is will cause the filters not to
push down)
##########
datafusion/optimizer/src/push_down_filter.rs:
##########
@@ -1585,6 +1596,63 @@ fn with_filters(predicates: Vec<Expr>, plan:
LogicalPlan) -> LogicalPlan {
}
}
+/// Does `expr` read nothing beyond the given window partition keys?
+///
+/// A subtree that is exactly one of the keys counts as read in full, so a
+/// predicate on an *expression* key, say `NULLIF(c, '') IS NOT NULL` against
+/// `PARTITION BY NULLIF(c, '')`, qualifies even though the column it
ultimately
+/// reads (`c`) is not a key on its own. Such a predicate is constant within
each
+/// partition, so applying it below the window drops whole partitions and
leaves
+/// every surviving row's window value unchanged.
+///
+/// Matching is structural, which makes this conservative rather than wrong: a
+/// predicate written `b + a` does not match a key written `a + b`, and is
simply
+/// left above the window.
+///
+/// A node carrying a subquery counts as reading something else, whatever the
Review Comment:
I found this paragraph very hard to understand -- I don't think we need to
explain the intricate details of why subqueries can't be pushed down
##########
datafusion/optimizer/src/push_down_filter.rs:
##########
@@ -1585,6 +1596,63 @@ fn with_filters(predicates: Vec<Expr>, plan:
LogicalPlan) -> LogicalPlan {
}
}
+/// Does `expr` read nothing beyond the given window partition keys?
+///
+/// A subtree that is exactly one of the keys counts as read in full, so a
Review Comment:
I found "is exactly one of the keys counts as read in full" very hard to
understand
I think this is trying to explain that we can pass any expression down that
only refers to partition columns (as in is an expression of columns that only
appear in the expression)?
##########
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
Review Comment:
Is it important to explain here what used to happen? This comment on the old
behavior seems like it will be irrelevant once the PR will merge (it would be
better as a comment on the PR I think, not in the code)
##########
datafusion/optimizer/src/push_down_filter.rs:
##########
@@ -1172,21 +1177,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
Review Comment:
I am not sure "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."
is needed as it explain what the code does, doesn't it?
##########
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:
do you need to clone here? Can we just use `HashSet<&Expr>` to keep
references rather tahn deeply clone the exprs?
##########
datafusion/sqllogictest/test_files/push_down_filter_regression.slt:
##########
@@ -723,3 +723,92 @@ query I
SELECT sum(c) FROM (SELECT random() < 0.5 AS k, count(*) AS c FROM
generate_series(1, 10000) GROUP BY random() < 0.5) WHERE k OR NOT k;
----
10000
+
+# Window filter pushdown over an expression PARTITION BY key.
+#
+# A predicate that reads only the partition keys is constant within a
partition,
+# so pushing it below the window drops whole partitions and leaves every
+# surviving row's window value alone. The unit tests pin where the filter
lands;
+# these pin that the answers do not move, which is what a wrong push breaks.
+
+statement ok
+create table window_expr_key(k varchar, v int) as values
+ ('a', 1), ('a', 2), ('a', 3),
+ ('', 4), ('', 5),
+ ('b', 6);
+
+# Partitions under NULLIF(k, '') are 'a' => {1,2,3}, NULL => {4,5}, 'b' => {6}.
Review Comment:
can you also please add EXPLAIN to these tests so we can see the shape?
##########
datafusion/optimizer/src/push_down_filter.rs:
##########
@@ -1585,6 +1596,63 @@ fn with_filters(predicates: Vec<Expr>, plan:
LogicalPlan) -> LogicalPlan {
}
}
+/// Does `expr` read nothing beyond the given window partition keys?
+///
+/// A subtree that is exactly one of the keys counts as read in full, so a
Review Comment:
I think this comment would be much more helpful with some examples
Something like
given `PARTITION BY (a, b)`
Can push down filters: `a < 5`, `a+b = 4` etc
Can not push down filters like `c < 5`
##########
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:
would help here to note with `((a + b) + 1) > 10`
##########
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:
I also recommend a test that has multiple instances of expressions in
partition by
```sql
PARTITION BY (a+b)
...
WHERE ((a +b) > 5 OR (a + b) < 10)
```
##########
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:
again, a reference to the old code is not helpful in code comments -- it is
helpful in the context of a PR and I think should be a comment on the PR
--
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]
