[GitHub] [arrow-datafusion] alamb commented on a diff in pull request #3717: Use column aliases specified by `WITH` statements

[GitBox <gi...@apache.org>]

alamb commented on code in PR #3717:
URL: https://github.com/apache/arrow-datafusion/pull/3717#discussion_r988368205


##########
datafusion/sql/src/planner.rs:
##########
@@ -5046,6 +5058,67 @@ mod tests {
         quick_test(sql, expected)
     }
 
+    #[test]
+    fn cte_with_no_column_names() {
+        let sql = "WITH \
+        numbers AS ( \
+            SELECT 1 as a, 2 as b, 3 as c \
+        ) \
+        SELECT * FROM numbers;";
+
+        let expected = "Projection: numbers.a, numbers.b, numbers.c\
+        \n  Projection: Int64(1) AS a, Int64(2) AS b, Int64(3) AS c, alias=numbers\
+        \n    EmptyRelation";
+
+        quick_test(sql, expected)
+    }
+
+    #[test]
+    fn cte_with_column_names() {
+        let sql = "WITH \
+        numbers(a, b, c) AS ( \
+            SELECT 1, 2, 3 \
+        ) \
+        SELECT * FROM numbers;";
+
+        let expected = "Projection: numbers.a, numbers.b, numbers.c\
+        \n  Projection: numbers.Int64(1) AS a, numbers.Int64(2) AS b, numbers.Int64(3) AS c, alias=numbers\
+        \n    Projection: Int64(1), Int64(2), Int64(3), alias=numbers\
+        \n      EmptyRelation";
+
+        quick_test(sql, expected)
+    }
+
+    #[test]
+    fn cte_with_column_aliases_precedence() {
+        // The end result should always be what CTE specification says
+        let sql = "WITH \
+        numbers(a, b, c) AS ( \
+            SELECT 1 as x, 2 as y, 3 as z \
+        ) \
+        SELECT * FROM numbers;";
+
+        let expected = "Projection: numbers.a, numbers.b, numbers.c\

Review Comment:
   👍 



##########
datafusion/sql/src/planner.rs:
##########
@@ -785,33 +790,40 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
             }
         };
         if let Some(alias) = alias {
-            let columns_alias = alias.clone().columns;
-            if columns_alias.is_empty() {
-                // sqlparser-rs encodes AS t as an empty list of column alias
-                Ok(plan)
-            } else if columns_alias.len() != plan.schema().fields().len() {
-                Err(DataFusionError::Plan(format!(
-                    "Source table contains {} columns but only {} names given as column alias",
-                    plan.schema().fields().len(),
-                    columns_alias.len(),
-                )))
-            } else {
-                Ok(LogicalPlanBuilder::from(plan.clone())
-                    .project_with_alias(
-                        plan.schema().fields().iter().zip(columns_alias.iter()).map(
-                            |(field, ident)| {
-                                col(field.name()).alias(&normalize_ident(ident))
-                            },
-                        ),
-                        Some(normalize_ident(&alias.name)),
-                    )?
-                    .build()?)
-            }
+            self.apply_table_alias(plan, alias)
         } else {
             Ok(plan)
         }
     }
 
+    /// Apply the given TableAlias to the top-level projection.
+    fn apply_table_alias(

Review Comment:
   👌 



-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org