diff --git a/datafusion/optimizer/src/push_down_filter.rs b/datafusion/optimizer/src/push_down_filter.rs index 534209c4e11b..7cb0e7c2f1f7 100644 --- a/datafusion/optimizer/src/push_down_filter.rs +++ b/datafusion/optimizer/src/push_down_filter.rs @@ -29,6 +29,7 @@ use datafusion_common::tree_node::{ use datafusion_common::{ internal_err, plan_err, qualified_name, Column, DFSchema, Result, }; +use datafusion_expr::expr::WindowFunction; use datafusion_expr::expr_rewriter::replace_col; use datafusion_expr::logical_plan::{Join, JoinType, LogicalPlan, TableScan, Union}; use datafusion_expr::utils::{ @@ -1001,23 +1002,34 @@ 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. + let extract_partition_keys = |func: &WindowFunction| { + func.partition_by + .iter() + .map(|c| Column::from_qualified_name(c.schema_name().to_string())) + .collect::>() + }; let potential_partition_keys = window .window_expr .iter() .map(|e| { - if let Expr::WindowFunction(window_expression) = e { - window_expression - .partition_by - .iter() - .map(|c| { - Column::from_qualified_name( - c.schema_name().to_string(), - ) - }) - .collect::>() - } else { - // window functions expressions are only Expr::WindowFunction - unreachable!() + match e { + Expr::WindowFunction(window_func) => { + extract_partition_keys(window_func) + } + Expr::Alias(alias) => { + if let Expr::WindowFunction(window_func) = + alias.expr.as_ref() + { + extract_partition_keys(window_func) + } else { + // window functions expressions are only Expr::WindowFunction + unreachable!() + } + } + _ => { + // window functions expressions are only Expr::WindowFunction + unreachable!() + } } }) // performs the set intersection of the partition keys of all window functions, diff --git a/datafusion/sqllogictest/test_files/subquery.slt b/datafusion/sqllogictest/test_files/subquery.slt index 25fe4c7b0390..6933514951b3 100644 --- a/datafusion/sqllogictest/test_files/subquery.slt +++ b/datafusion/sqllogictest/test_files/subquery.slt @@ -1306,7 +1306,6 @@ select struct(1, 'b') ---- {c0: 1, c1: b} - query T select (select struct(1, 'b')['c1']); ---- @@ -1330,7 +1329,6 @@ WHERE 1+2 = 3 AND column1 IN (SELECT struct(1, 'b')['c0']); ---- 1 - query I SELECT * FROM foo WHERE EXISTS (SELECT * FROM (values (1)) WHERE column1 = foo.x AND struct(1, 'b')['c0'] = 1); @@ -1344,6 +1342,47 @@ WHERE 1+2 = 3 AND EXISTS (SELECT * FROM (values (1)) WHERE column1 = foo.x AND s ---- 1 - statement ok drop table foo; + + +# Test for window alias in subquery + +# Setup source table +statement ok +CREATE TABLE source_table ( + column1 TEXT, + column2 TIMESTAMP, + column3 FLOAT +); + +statement ok +INSERT INTO source_table VALUES + ('item1', TIMESTAMP '1970-01-01 00:00:01', 50.0), + ('item2', TIMESTAMP '1970-01-01 00:00:02', 30.0), + ('item1', TIMESTAMP '1970-01-01 00:00:03', 25.0); + +# Execute the query +query TPR +WITH SubQuery AS ( + SELECT + a.column1, + a.column2 AS ts_column, + a.column3, + SUM(a.column3) OVER ( + PARTITION BY a.column1 + ORDER BY a.column2 RANGE BETWEEN INTERVAL '10 minutes' PRECEDING AND CURRENT ROW + ) AS moving_sum + FROM source_table a +) +SELECT + column1, + ts_column, + moving_sum +FROM SubQuery +WHERE moving_sum > 60; +---- +item1 1970-01-01T00:00:03 75 + +statement ok +drop table source_table; \ No newline at end of file