Preserve Snowflake PIVOT column names

datafusion/core/tests/sql/pivot_unpivot.rs

@@ -33,12 +33,12 @@ async fn pivot_list_is_lowered_to_filtered_aggregates() -> Result<()> {
         .await?;
 
     insta::assert_snapshot!(batches_to_string(&batches), @r"
-    +----+----+----+
-    | id | a  | b  |
-    +----+----+----+
-    | 1  | 10 | 20 |
-    | 2  | 7  |    |
-    +----+----+----+
+    +----+-----+-----+
+    | id | 'a' | 'b' |
+    +----+-----+-----+
+    | 1  | 10  | 20  |
+    | 2  | 7   |     |
+    +----+-----+-----+
     ");
     Ok(())
 }

datafusion/sql/src/relation/mod.rs

@@ -43,6 +43,11 @@ struct SqlToRelRelationContext<'a, 'b, S: ContextProvider> {
     planner_context: &'a mut PlannerContext,
 }
 
+struct ResolvedPivotValue {
+    value: ScalarValue,
+    name: String,
+}
+
 // Implement RelationPlannerContext
 impl<'a, 'b, S: ContextProvider> RelationPlannerContext
     for SqlToRelRelationContext<'a, 'b, S>
@@ -373,12 +378,18 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
                 let pivot_values = values
                     .into_iter()
                     .map(|value| {
+                        let name = value.alias.map_or_else(
+                            || value.expr.to_string(),
+                            |alias| self.ident_normalizer.normalize(alias),
+                        );
                         match self.sql_expr_to_logical_expr(
                             value.expr,
                             input_plan.schema(),
                             planner_context,
                         )? {
-                            Expr::Literal(value, _) => Ok(value),
+                            Expr::Literal(value, _) => {
+                                Ok(ResolvedPivotValue { value, name })
+                            }
                             _ => plan_err!("PIVOT values must be literals"),
                         }
                     })
@@ -549,7 +560,7 @@ fn transform_pivot_to_aggregate(
     input: LogicalPlan,
     aggregate_expr: &Expr,
     pivot_column: &Column,
-    pivot_values: &[ScalarValue],
+    pivot_values: &[ResolvedPivotValue],
     default_on_null_expr: Option<&Expr>,
 ) -> Result<LogicalPlan> {
     let input_schema = input.schema();
@@ -573,12 +584,12 @@ fn transform_pivot_to_aggregate(
     let pivot_type = input_schema.field(pivot_index).data_type().clone();
     let aggregates = pivot_values
         .iter()
-        .map(|value| {
+        .map(|pivot_value| {
             let filter = Expr::BinaryExpr(BinaryExpr::new(
                 Box::new(Expr::Column(pivot_column.clone())),
                 Operator::IsNotDistinctFrom,
                 Box::new(Expr::Cast(Cast::new(
-                    Box::new(Expr::Literal(value.clone(), None)),
+                    Box::new(Expr::Literal(pivot_value.value.clone(), None)),
                     pivot_type.clone(),
                 ))),
             ));
@@ -587,14 +598,13 @@ fn transform_pivot_to_aggregate(
             };
             let mut params = aggregate.params.clone();
             params.filter = Some(Box::new(filter));
-            let name = value.to_string().trim_matches('\'').to_string();
             Ok(Expr::Alias(Alias {
                 expr: Box::new(Expr::AggregateFunction(AggregateFunction {
                     func: Arc::clone(&aggregate.func),
                     params,
                 })),
                 relation: None,
-                name,
+                name: pivot_value.name.clone(),
                 metadata: None,
             }))
         })
@@ -608,7 +618,7 @@ fn transform_pivot_to_aggregate(
     };
     let pivot_names = pivot_values
         .iter()
-        .map(|value| value.to_string().trim_matches('\'').to_string())
+        .map(|value| value.name.clone())
         .collect::<Vec<_>>();
     let projection = aggregate_plan
         .schema()

datafusion/sql/tests/common/mod.rs

@@ -141,6 +141,11 @@ impl ContextProvider for MockContextProvider {
                 Field::new("id", DataType::Int32, false),
                 Field::new("price", DataType::Decimal128(10, 2), false),
             ])),
+            "quarterly_sales" => Ok(Schema::new(vec![
+                Field::new("empid", DataType::Int32, false),
+                Field::new("amount", DataType::Int32, false),
+                Field::new("quarter", DataType::Utf8, false),
+            ])),
             "person" => Ok(Schema::new(vec![
                 Field::new("id", DataType::UInt32, false),
                 Field::new("first_name", DataType::Utf8, false),

datafusion/sql/tests/sql_integration.rs

@@ -70,6 +70,25 @@ use sqlparser::parser::Parser;
 mod cases;
 mod common;
 
+#[test]
+fn pivot_uses_literal_text_or_explicit_alias_for_column_names() {
+    let plan = logical_plan_with_dialect(
+        "SELECT * FROM quarterly_sales \
+         PIVOT(SUM(amount) FOR quarter IN (\
+           '2023_Q1', '2023_Q2' AS q2, '2023_Q3' AS \"Mixed Case\"))",
+        &SnowflakeDialect {},
+    )
+    .unwrap();
+
+    let field_names = plan
+        .schema()
+        .fields()
+        .iter()
+        .map(|field| field.name().as_str())
+        .collect::<Vec<_>>();
+    assert_eq!(field_names, vec!["empid", "'2023_Q1'", "q2", "Mixed Case"]);
+}
+
 #[test]
 fn parse_decimals_1() {
     let sql = "SELECT 1";