Recompute schemas after union schema change (#1635)

* Recompute schemas after union schema change

* Recompute schemas after union schema change

Author: osipovartem

crates/core-executor/src/datafusion/logical_analyzer/union_schema_analyzer.rs

@@ -55,13 +55,21 @@ impl AnalyzerRule for UnionSchemaAnalyzer {
 }
 
 fn analyze_internal(plan: &LogicalPlan) -> DFResult<LogicalPlan> {
-    Ok(plan
-        .clone()
-        .transform_down(|node| match node {
-            LogicalPlan::Union(union) => Ok(Transformed::yes(rewrite_union(&union)?)),
-            _ => Ok(Transformed::no(node.clone())),
-        })?
-        .data)
+    let mut new_plan = plan.clone().transform_up(|node| match node {
+        LogicalPlan::Union(union) => Ok(Transformed::yes(rewrite_union(&union)?)),
+        _ => Ok(Transformed::no(node.clone())),
+    })?;
+
+    // Recompute schemas for Projection and SubqueryAlias nodes above the rewritten Union
+    if new_plan.transformed {
+        new_plan = new_plan.data.transform_up(|node| match node {
+            LogicalPlan::Projection(_) | LogicalPlan::SubqueryAlias(_) => {
+                Ok(Transformed::yes(node.recompute_schema()?))
+            }
+            _ => Ok(Transformed::no(node)),
+        })?;
+    }
+    Ok(new_plan.data)
 }
 
 fn rewrite_union(union: &Union) -> datafusion_common::Result<LogicalPlan> {

crates/core-executor/src/tests/sql/logical_analyzer/snapshots/union_schema_analyzer/query_union_with_timestamp_microseconds.snap

@@ -1,14 +1,14 @@
 ---
 source: crates/core-executor/src/tests/sql/logical_analyzer/union_schema_analyzer.rs
-description: "\"SELECT '2024-12-31 10:00:00.000'::TIMESTAMP as t\n        UNION ALL\n    SELECT '9999-12-31 00:00:00.000 +0000' AS t;\""
+description: "\"SELECT '2024-12-31 10:00:00.000'::TIMESTAMP as t\n        UNION ALL\n    SELECT '9999-12-31 00:00:00.000 +0000' AS t\n    ORDER BY 1\""
 ---
 Ok(
     [
         "+---------------------+",
         "| t                   |",
         "+---------------------+",
-        "| 9999-12-31T00:00:00 |",
         "| 2024-12-31T10:00:00 |",
+        "| 9999-12-31T00:00:00 |",
         "+---------------------+",
     ],
 )

crates/core-executor/src/tests/sql/logical_analyzer/snapshots/union_schema_analyzer/query_union_with_timestamp_microseconds_cte.snap

@@ -0,0 +1,14 @@
+---
+source: crates/core-executor/src/tests/sql/logical_analyzer/union_schema_analyzer.rs
+description: "\"WITH res AS (\n        SELECT '2024-12-31 10:00:00.000'::TIMESTAMP as t\n        UNION ALL\n        SELECT '9999-12-31 00:00:00.000 +0000'\n    )\n    SELECT t FROM res\n    ORDER BY t\""
+---
+Ok(
+    [
+        "+---------------------+",
+        "| t                   |",
+        "+---------------------+",
+        "| 2024-12-31T10:00:00 |",
+        "| 9999-12-31T00:00:00 |",
+        "+---------------------+",
+    ],
+)

crates/core-executor/src/tests/sql/logical_analyzer/union_schema_analyzer.rs

@@ -9,7 +9,19 @@ test_query!(
     union_with_timestamp_microseconds,
     "SELECT '2024-12-31 10:00:00.000'::TIMESTAMP as t
         UNION ALL
-    SELECT '9999-12-31 00:00:00.000 +0000' AS t",
-    sort_all = true,
+    SELECT '9999-12-31 00:00:00.000 +0000' AS t
+    ORDER BY 1",
+    snapshot_path = "union_schema_analyzer"
+);
+
+test_query!(
+    union_with_timestamp_microseconds_cte,
+    "WITH res AS (
+        SELECT '2024-12-31 10:00:00.000'::TIMESTAMP as t
+        UNION ALL
+        SELECT '9999-12-31 00:00:00.000 +0000'
+    )
+    SELECT t FROM res
+    ORDER BY t",
     snapshot_path = "union_schema_analyzer"
 );