Avoid Aliased Window Expr Enter Unreachable Code (#14109)

apache · Jan 14, 2025 · fda500a · fda500a
1 parent 722307f
commit fda500a
Show file tree

Hide file tree

Showing 2 changed files with 67 additions and 16 deletions.
diff --git 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::<HashSet<_>>()
+                };
                 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::<HashSet<_>>()
-                        } 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
@@ -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;