diff --git a/datafusion/core/tests/sql/pivot_unpivot.rs b/datafusion/core/tests/sql/pivot_unpivot.rs index 22c5b7b523280..7d68187165453 100644 --- a/datafusion/core/tests/sql/pivot_unpivot.rs +++ b/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(()) } diff --git a/datafusion/sql/src/relation/mod.rs b/datafusion/sql/src/relation/mod.rs index 0f1a1137c23d8..c0d298d494c76 100644 --- a/datafusion/sql/src/relation/mod.rs +++ b/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 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 { 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::>(); let projection = aggregate_plan .schema() diff --git a/datafusion/sql/tests/common/mod.rs b/datafusion/sql/tests/common/mod.rs index dfe393c1952ac..f8b88498ff51b 100644 --- a/datafusion/sql/tests/common/mod.rs +++ b/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), diff --git a/datafusion/sql/tests/sql_integration.rs b/datafusion/sql/tests/sql_integration.rs index a83d070355e0b..7ca7eccdd42b9 100644 --- a/datafusion/sql/tests/sql_integration.rs +++ b/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::>(); + assert_eq!(field_names, vec!["empid", "'2023_Q1'", "q2", "Mixed Case"]); +} + #[test] fn parse_decimals_1() { let sql = "SELECT 1";