Apply PIVOT aliases before schema validation

datafusion/core/tests/sql/pivot_unpivot.rs

@@ -43,6 +43,42 @@ async fn pivot_list_is_lowered_to_filtered_aggregates() -> Result<()> {
     Ok(())
 }
 
+#[tokio::test]
+async fn chained_pivots_apply_explicit_column_aliases() -> Result<()> {
+    let batches = SessionContext::new()
+        .sql(
+            "SELECT SUM(q1_sales) AS q1_sales,
+                    SUM(q2_sales) AS q2_sales,
+                    MAX(q1_discount) AS q1_discount,
+                    MAX(q2_discount) AS q2_discount
+             FROM (
+                 SELECT amount,
+                        quarter AS sales_quarter,
+                        quarter AS discount_quarter,
+                        discount
+                 FROM (VALUES
+                     (100, 'Q1', 10),
+                     (200, 'Q2', 20)
+                 ) sales(amount, quarter, discount)
+             )
+             PIVOT(SUM(amount) FOR sales_quarter IN ('Q1', 'Q2'))
+             PIVOT(MAX(discount) FOR discount_quarter IN ('Q1', 'Q2'))
+             AS p(q1_sales, q2_sales, q1_discount, q2_discount)",
+        )
+        .await?
+        .collect()
+        .await?;
+
+    insta::assert_snapshot!(batches_to_string(&batches), @r"
+    +----------+----------+-------------+-------------+
+    | q1_sales | q2_sales | q1_discount | q2_discount |
+    +----------+----------+-------------+-------------+
+    | 100      | 200      | 10          | 20          |
+    +----------+----------+-------------+-------------+
+    ");
+    Ok(())
+}
+
 #[tokio::test]
 async fn unpivot_excludes_nulls_by_default() -> Result<()> {
     let batches = SessionContext::new()

datafusion/sql/src/relation/mod.rs

@@ -394,12 +394,24 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
                         }
                     })
                     .collect::<Result<Vec<_>>>()?;
+                let output_aliases = alias.as_ref().and_then(|alias| {
+                    (!alias.columns.is_empty()).then(|| {
+                        alias
+                            .columns
+                            .iter()
+                            .map(|column| {
+                                self.ident_normalizer.normalize(column.name.clone())
+                            })
+                            .collect::<Vec<_>>()
+                    })
+                });
                 let plan = transform_pivot_to_aggregate(
                     input_plan,
                     &aggregate_expr,
                     &pivot_column,
                     &pivot_values,
                     default_on_null_expr.as_ref(),
+                    output_aliases.as_deref(),
                 )?;
                 (plan, alias)
             }
@@ -562,9 +574,10 @@ fn transform_pivot_to_aggregate(
     pivot_column: &Column,
     pivot_values: &[ResolvedPivotValue],
     default_on_null_expr: Option<&Expr>,
+    output_aliases: Option<&[String]>,
 ) -> Result<LogicalPlan> {
     let input_schema = input.schema();
-    let group_by = input_schema
+    let group_by_columns = input_schema
         .columns()
         .into_iter()
         .filter(|column| {
@@ -574,7 +587,27 @@ fn transform_pivot_to_aggregate(
                     .iter()
                     .any(|aggregate_column| aggregate_column.name == column.name)
         })
-        .map(Expr::Column)
+        .collect::<Vec<_>>();
+    let group_by_len = group_by_columns.len();
+    if let Some(output_aliases) = output_aliases
+        && output_aliases.len() != group_by_len + pivot_values.len()
+    {
+        return plan_err!(
+            "Source table contains {} columns but only {} names given as column alias",
+            group_by_len + pivot_values.len(),
+            output_aliases.len()
+        );
+    }
+    let group_by = group_by_columns
+        .into_iter()
+        .enumerate()
+        .map(|(index, column)| {
+            if let Some(aliases) = output_aliases {
+                Expr::Column(column).alias(aliases[index].clone())
+            } else {
+                Expr::Column(column)
+            }
+        })
         .collect::<Vec<_>>();
     let pivot_index = input_schema.index_of_column(pivot_column).map_err(|_| {
         datafusion_common::plan_datafusion_err!(
@@ -584,7 +617,8 @@ fn transform_pivot_to_aggregate(
     let pivot_type = input_schema.field(pivot_index).data_type().clone();
     let aggregates = pivot_values
         .iter()
-        .map(|pivot_value| {
+        .enumerate()
+        .map(|(index, pivot_value)| {
             let filter = Expr::BinaryExpr(BinaryExpr::new(
                 Box::new(Expr::Column(pivot_column.clone())),
                 Operator::IsNotDistinctFrom,
@@ -598,13 +632,17 @@ fn transform_pivot_to_aggregate(
             };
             let mut params = aggregate.params.clone();
             params.filter = Some(Box::new(filter));
+            let name = output_aliases.map_or_else(
+                || pivot_value.name.clone(),
+                |aliases| aliases[group_by_len + index].clone(),
+            );
             Ok(Expr::Alias(Alias {
                 expr: Box::new(Expr::AggregateFunction(AggregateFunction {
                     func: Arc::clone(&aggregate.func),
                     params,
                 })),
                 relation: None,
-                name: pivot_value.name.clone(),
+                name,
                 metadata: None,
             }))
         })
@@ -618,7 +656,13 @@ fn transform_pivot_to_aggregate(
     };
     let pivot_names = pivot_values
         .iter()
-        .map(|value| value.name.clone())
+        .enumerate()
+        .map(|(index, value)| {
+            output_aliases.map_or_else(
+                || value.name.clone(),
+                |aliases| aliases[group_by_len + index].clone(),
+            )
+        })
         .collect::<Vec<_>>();
     let projection = aggregate_plan
         .schema()