Repository: apache/datafusion (Rust)

Issue:
Unused Window functions experssion is wrongly removed from LogicalPlan during optimalization
**Describe the bug**
Unused window function is wrongly removed from Logical Plan during optimization. It leaves WindowAgg empty, which causes `Impossibly got empty window expression` error.

**To Reproduce**

```sql
RUST_LOG=debug datafusion-cli

WITH _sample_data AS (
  SELECT 1 as a, 'aa' AS b
  UNION ALL
  SELECT 3 as a, 'aa' AS b
  UNION ALL
  SELECT 5 as a, 'bb' AS b
  UNION ALL
  SELECT 7 as a, 'bb' AS b
  
), _data2 AS (
  SELECT
    row_number() OVER (PARTITION BY s.b ORDER BY s.a) AS seq,
    s.a,
    s.b
  FROM _sample_data s
)
SELECT d.b, MAX(d.a) AS max_a
FROM _data2 d
GROUP BY d.b
ORDER BY d.b;

-- Result
-- Error: Impossibly got empty window expression
```

**Expected behavior**

With above SQL i should got:
```sql
-- Result
| b  | max_a |
+----+-------+
| aa | 3     | 
| bb | 7     | 
+----+-------+

```

**Additional context**

Input logical plan:
```
    Sort: #d.b ASC NULLS LAST
      Projection: #d.b, #MAX(d.a) AS max_a
        Aggregate: groupBy=[[#d.b]], aggr=[[MAX(#d.a)]]
          Projection: #_data2.seq, #_data2.a, #_data2.b, alias=d
            Projection: #ROW_NUMBER() PARTITION BY [#s.b] ORDER BY [#s.a ASC NULLS LAST] AS seq, #s.a, #s.b, alias=_data2
              WindowAggr: windowExpr=[[ROW_NUMBER() PARTITION BY [#s.b] ORDER BY [#s.a ASC NULLS LAST]]]
                Projection: #_sample_data.a, #_sample_data.b, alias=s
                  Union
                    Projection: Int64(1) AS a, Utf8("aa") AS b
                      EmptyRelation
                    Projection: Int64(3) AS a, Utf8("aa") AS b
                      EmptyRelation
                    Projection: Int64(5) AS a, Utf8("bb") AS b
                      EmptyRelation
                    Projection: Int64(7) AS a, Utf8("bb") AS b
                      EmptyRelation
```

Optimized logical plan:
```
    Sort: #d.b ASC NULLS LAST
      Projection: #d.b, #MAX(d.a) AS max_a
        Aggregate: groupBy=[[#d.b]], aggr=[[MAX(#d.a)]]
          Projection: #_data2.a, #_data2.b, alias=d
            Projection: #s.a, #s.b, alias=_data2
              WindowAggr: windowExpr=[[]]  <-------- Here is empty expression
                Projection: #_sample_data.a, #_sample_data.b, alias=s
                  Union
                    Projection: Int64(1) AS a, Utf8("aa") AS b
                      EmptyRelation
                    Projection: Int64(3) AS a, Utf8("aa") AS b
                      EmptyRelation
                    Projection: Int64(5) AS a, Utf8("bb") AS b
                      EmptyRelation
                    Projection: Int64(7) AS a, Utf8("bb") AS b
                      EmptyRelation
```

We can see *windowExpr* is empty after optimization. I have tested this case on master ('[a825891](https://github.com/apache/arrow-datafusion/commit/a82589100d7b07cb629476a2df7eac002f48b8a0)').

I don't use `seq` column in last CTE and this causes error. WIth below SQL it's work
```sql
WITH _sample_data AS (
  SELECT 1 as a, 'aa' AS b
  UNION ALL
  SELECT 3 as a, 'aa' AS b
  UNION ALL
  SELECT 5 as a, 'bb' AS b
  UNION ALL
  SELECT 7 as a, 'bb' AS b
  
), _data2 AS (
  SELECT
    row_number() OVER (PARTITION BY s.b ORDER BY s.a) AS seq,
    s.a,
    s.b
  FROM _sample_data s
)
SELECT d.b, MAX(d.a) AS max_a, max(d.seq)
FROM _data2 d
GROUP BY d.b
ORDER BY d.b;

-- Resut:
| b  | max_a | MAX(d.seq) |
+----+-------+------------+
| aa | 3     | 2          |
| bb | 7     | 2          |
+----+-------+------------+
```



Discussion:
I think this is because the window exprs are not referenced by any operators,  and projection_push_down make empty new_window_expr for window logical plan.

If the window exprs are not referenced by any operators, we can remove the window exprs.  cc @mateuszkj 



Patch:
diff --git a/datafusion/core/src/optimizer/projection_push_down.rs b/datafusion/core/src/optimizer/projection_push_down.rs
index 494b7dd64044..4feef4f99057 100644
--- a/datafusion/core/src/optimizer/projection_push_down.rs
+++ b/datafusion/core/src/optimizer/projection_push_down.rs
@@ -277,6 +277,18 @@ fn optimize_plan(
                 })?;
             }
 
+            // none columns in window expr are needed, remove the window expr
+            if new_window_expr.is_empty() {
+                return LogicalPlanBuilder::from(optimize_plan(
+                    _optimizer,
+                    input,
+                    required_columns,
+                    true,
+                    _execution_props,
+                )?)
+                .build();
+            };
+
             // for all the retained window expr, find their sort expressions if any, and retain these
             exprlist_to_columns(
                 &find_sort_exprs(&new_window_expr),


Test patch:
diff --git a/datafusion/core/tests/sql/window.rs b/datafusion/core/tests/sql/window.rs
index 7e36177fde9a..bdbc77067ebd 100644
--- a/datafusion/core/tests/sql/window.rs
+++ b/datafusion/core/tests/sql/window.rs
@@ -298,3 +298,138 @@ async fn window_partition_by_order_by() -> Result<()> {
     assert_batches_eq!(expected, &results);
     Ok(())
 }
+
+#[tokio::test]
+async fn window_expr_eliminate() -> Result<()> {
+    let ctx = SessionContext::new();
+
+    // window expr is not referenced anywhere, eliminate it.
+    let sql = "WITH _sample_data AS (
+                SELECT 1 as a, 'aa' AS b
+                UNION ALL
+                SELECT 3 as a, 'aa' AS b
+                UNION ALL
+                SELECT 5 as a, 'bb' AS b
+                UNION ALL
+                SELECT 7 as a, 'bb' AS b
+            ), _data2 AS (
+                SELECT
+                row_number() OVER (PARTITION BY s.b ORDER BY s.a) AS seq,
+                s.a,
+                s.b
+                FROM _sample_data s
+            )
+            SELECT d.b, MAX(d.a) AS max_a
+            FROM _data2 d
+            GROUP BY d.b
+            ORDER BY d.b;";
+
+    let msg = format!("Creating logical plan for '{}'", sql);
+    let plan = ctx
+        .create_logical_plan(&("explain ".to_owned() + sql))
+        .expect(&msg);
+    let state = ctx.state.read().clone();
+    let plan = state.optimize(&plan)?;
+    let expected = vec![
+        "Explain [plan_type:Utf8, plan:Utf8]",
+        "  Sort: #d.b ASC NULLS LAST [b:Utf8, max_a:Int64;N]",
+        "    Projection: #d.b, #MAX(d.a) AS max_a [b:Utf8, max_a:Int64;N]",
+        "      Aggregate: groupBy=[[#d.b]], aggr=[[MAX(#d.a)]] [b:Utf8, MAX(d.a):Int64;N]",
+        "        Projection: #_data2.a, #_data2.b, alias=d [a:Int64, b:Utf8]",
+        "          Projection: #s.a, #s.b, alias=_data2 [a:Int64, b:Utf8]",
+        "            Projection: #a, #b, alias=s [a:Int64, b:Utf8]",
+        "              Union [a:Int64, b:Utf8]",
+        "                Projection: Int64(1) AS a, Utf8(\"aa\") AS b [a:Int64, b:Utf8]",
+        "                  EmptyRelation []",
+        "                Projection: Int64(3) AS a, Utf8(\"aa\") AS b [a:Int64, b:Utf8]",
+        "                  EmptyRelation []",
+        "                Projection: Int64(5) AS a, Utf8(\"bb\") AS b [a:Int64, b:Utf8]",
+        "                  EmptyRelation []",
+        "                Projection: Int64(7) AS a, Utf8(\"bb\") AS b [a:Int64, b:Utf8]",
+        "                  EmptyRelation []",
+    ];
+    let formatted = plan.display_indent_schema().to_string();
+    let actual: Vec<&str> = formatted.trim().lines().collect();
+    assert_eq!(
+        expected, actual,
+        "\n\nexpected:\n\n{:#?}\nactual:\n\n{:#?}\n\n",
+        expected, actual
+    );
+
+    let results = execute_to_batches(&ctx, sql).await;
+    let expected = vec![
+        "+----+-------+",
+        "| b  | max_a |",
+        "+----+-------+",
+        "| aa | 3     |",
+        "| bb | 7     |",
+        "+----+-------+",
+    ];
+
+    assert_batches_eq!(expected, &results);
+
+    // window expr is referenced by the output, keep it
+    let sql = "WITH _sample_data AS (
+                SELECT 1 as a, 'aa' AS b
+                UNION ALL
+                SELECT 3 as a, 'aa' AS b
+                UNION ALL
+                SELECT 5 as a, 'bb' AS b
+                UNION ALL
+                SELECT 7 as a, 'bb' AS b
+            ), _data2 AS (
+                SELECT
+                row_number() OVER (PARTITION BY s.b ORDER BY s.a) AS seq,
+                s.a,
+                s.b
+                FROM _sample_data s
+            )
+            SELECT d.b, MAX(d.a) AS max_a, max(d.seq)
+            FROM _data2 d
+            GROUP BY d.b
+            ORDER BY d.b;";
+
+    let plan = ctx
+        .create_logical_plan(&("explain ".to_owned() + sql))
+        .expect(&msg);
+    let plan = state.optimize(&plan)?;
+    let expected = vec![
+        "Explain [plan_type:Utf8, plan:Utf8]",
+        "  Sort: #d.b ASC NULLS LAST [b:Utf8, max_a:Int64;N, MAX(d.seq):UInt64;N]",
+        "    Projection: #d.b, #MAX(d.a) AS max_a, #MAX(d.seq) [b:Utf8, max_a:Int64;N, MAX(d.seq):UInt64;N]",
+        "      Aggregate: groupBy=[[#d.b]], aggr=[[MAX(#d.a), MAX(#d.seq)]] [b:Utf8, MAX(d.a):Int64;N, MAX(d.seq):UInt64;N]",
+        "        Projection: #_data2.seq, #_data2.a, #_data2.b, alias=d [seq:UInt64;N, a:Int64, b:Utf8]",
+        "          Projection: #ROW_NUMBER() PARTITION BY [#s.b] ORDER BY [#s.a ASC NULLS LAST] AS seq, #s.a, #s.b, alias=_data2 [seq:UInt64;N, a:Int64, b:Utf8]",
+        "            WindowAggr: windowExpr=[[ROW_NUMBER() PARTITION BY [#s.b] ORDER BY [#s.a ASC NULLS LAST]]] [ROW_NUMBER() PARTITION BY [#s.b] ORDER BY [#s.a ASC NULLS LAST]:UInt64;N, a:Int64, b:Utf8]",
+        "              Projection: #a, #b, alias=s [a:Int64, b:Utf8]",
+        "                Union [a:Int64, b:Utf8]",
+        "                  Projection: Int64(1) AS a, Utf8(\"aa\") AS b [a:Int64, b:Utf8]",
+        "                    EmptyRelation []",
+        "                  Projection: Int64(3) AS a, Utf8(\"aa\") AS b [a:Int64, b:Utf8]",
+        "                    EmptyRelation []",
+        "                  Projection: Int64(5) AS a, Utf8(\"bb\") AS b [a:Int64, b:Utf8]",
+        "                    EmptyRelation []",
+        "                  Projection: Int64(7) AS a, Utf8(\"bb\") AS b [a:Int64, b:Utf8]",
+        "                    EmptyRelation []",
+    ];
+    let formatted = plan.display_indent_schema().to_string();
+    let actual: Vec<&str> = formatted.trim().lines().collect();
+    assert_eq!(
+        expected, actual,
+        "\n\nexpected:\n\n{:#?}\nactual:\n\n{:#?}\n\n",
+        expected, actual
+    );
+
+    let results = execute_to_batches(&ctx, sql).await;
+    let expected = vec![
+        "+----+-------+------------+",
+        "| b  | max_a | MAX(d.seq) |",
+        "+----+-------+------------+",
+        "| aa | 3     | 2          |",
+        "| bb | 7     | 2          |",
+        "+----+-------+------------+",
+    ];
+
+    assert_batches_eq!(expected, &results);
+    Ok(())
+}