support pivot

Author: simonvandel

datafusion/sql/src/query.rs

@@ -20,17 +20,19 @@ use std::sync::Arc;
 use crate::planner::{ContextProvider, PlannerContext, SqlToRel};
 
 use crate::stack::StackGuard;
+use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion};
 use datafusion_common::{not_impl_err, Constraints, DFSchema, Result};
-use datafusion_expr::expr::{Sort, WildcardOptions};
+use datafusion_expr::expr::{AggregateFunction, Sort, WildcardOptions};
 
 use datafusion_expr::select_expr::SelectExpr;
 use datafusion_expr::{
-    CreateMemoryTable, DdlStatement, Distinct, Expr, LogicalPlan, LogicalPlanBuilder,
+    case, col, lit, CreateMemoryTable, DdlStatement, Distinct, Expr, LogicalPlan,
+    LogicalPlanBuilder,
 };
 use sqlparser::ast::{
     Expr as SQLExpr, ExprWithAliasAndOrderBy, Ident, LimitClause, Offset, OffsetRows,
-    OrderBy, OrderByExpr, OrderByKind, PipeOperator, Query, SelectInto, SetExpr,
-    SetOperator, SetQuantifier, TableAlias,
+    OrderBy, OrderByExpr, OrderByKind, PipeOperator, PivotValueSource, Query, SelectInto,
+    SetExpr, SetOperator, SetQuantifier, TableAlias,
 };
 use sqlparser::tokenizer::Span;
 
@@ -200,6 +202,19 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
             PipeOperator::Join(join) => {
                 self.parse_relation_join(plan, join, planner_context)
             }
+            PipeOperator::Pivot {
+                aggregate_functions,
+                value_column,
+                value_source,
+                alias,
+            } => self.pipe_operator_pivot(
+                plan,
+                aggregate_functions,
+                value_column,
+                value_source,
+                alias,
+                planner_context,
+            ),
 
             x => not_impl_err!("`{x}` pipe operator is not supported yet"),
         }
@@ -376,6 +391,143 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
             .build()
     }
 
+    /// Handle PIVOT pipe operator
+    fn pipe_operator_pivot(
+        &self,
+        plan: LogicalPlan,
+        aggregate_functions: Vec<sqlparser::ast::ExprWithAlias>,
+        value_column: Vec<Ident>,
+        value_source: PivotValueSource,
+        alias: Option<Ident>,
+        planner_context: &mut PlannerContext,
+    ) -> Result<LogicalPlan> {
+        // Extract pivot values from the value source
+        let pivot_values = if let PivotValueSource::List(values) = value_source {
+            values
+        } else {
+            return not_impl_err!(
+                "Only static pivot value lists are supported currently"
+            );
+        };
+
+        // Convert pivot column to DataFusion expression
+        if value_column.len() != 1 {
+            return not_impl_err!("Multi-column pivot is not supported yet");
+        }
+        let pivot_col_name = &value_column[0].value;
+        let pivot_col_expr = col(pivot_col_name);
+
+        let input_schema = plan.schema();
+
+        // Convert sql to DF exprs
+        let aggregate_functions = aggregate_functions
+            .into_iter()
+            .map(|f| self.sql_to_expr_with_alias(f, input_schema, planner_context))
+            .collect::<Result<Vec<_>, _>>()?;
+
+        // Convert aggregate functions to logical expressions to extract measure columns
+        let mut measure_columns = std::collections::HashSet::new();
+        for agg_func_with_alias in &aggregate_functions {
+            agg_func_with_alias.apply(|e| {
+                if let Expr::Column(col) = e {
+                    measure_columns.insert(col.name.clone());
+                };
+                Ok(TreeNodeRecursion::Continue)
+            })?;
+        }
+
+        // Get all column names from the input plan to determine group-by columns.
+        // Add all columns except the pivot column and measure columns to group by
+        let mut group_by_cols = Vec::new();
+        for field in input_schema.fields() {
+            let col_name = field.name();
+            if col_name != pivot_col_name && !measure_columns.contains(col_name) {
+                group_by_cols.push(col(col_name));
+            }
+        }
+
+        // Create aggregate expressions for each pivot value
+        let mut aggr_exprs = Vec::new();
+
+        // For each pivot value and aggregate function combination, create a conditional aggregate
+        // Process pivot values first to get the desired column order
+        for pivot_value in pivot_values {
+            let pivot_value_expr = self.sql_to_expr(
+                pivot_value.expr.clone(),
+                input_schema,
+                planner_context,
+            )?;
+            for agg_func_with_alias in &aggregate_functions {
+                let (alias_name, mut agg_fn) = match agg_func_with_alias {
+                    Expr::Alias(alias) => match *alias.expr.clone() {
+                        Expr::Alias(inner_alias) => {
+                            let Expr::AggregateFunction(
+                                agg_func @ AggregateFunction { .. },
+                            ) = *inner_alias.expr.clone()
+                            else {
+                                return not_impl_err!("Only function expressions are supported in PIVOT aggregate functions");
+                            };
+                            (Some(alias.name.clone()), agg_func)
+                        }
+                        Expr::AggregateFunction(agg_func @ AggregateFunction { .. }) => {
+                            (Some(alias.name.clone()), agg_func)
+                        }
+                        _ => {
+                            return not_impl_err!("Only function expressions are supported in PIVOT aggregate functions");
+                        }
+                    },
+                    Expr::AggregateFunction(agg_func) => (None, agg_func.clone()),
+                    _ => {
+                        return not_impl_err!("Expected aggregate function");
+                    }
+                };
+
+                if agg_fn.params.args.len() != 1 {
+                    return not_impl_err!(
+                        "Only exactly one aggregate function argument is supported"
+                    );
+                }
+                let arg = &mut agg_fn.params.args[0];
+                *arg = case(pivot_col_expr.clone())
+                    .when(pivot_value_expr.clone(), arg.clone())
+                    .otherwise(lit(datafusion_common::ScalarValue::Null))?;
+
+                let agg_expr = Expr::AggregateFunction(agg_fn);
+                let aggr_func_alias = alias_name.unwrap_or(agg_expr.name_for_alias()?);
+
+                let pivot_value_name = if let Some(alias) = &pivot_value.alias {
+                    alias.value.clone()
+                } else {
+                    // Use the pivot value as column name, stripping quotes
+                    pivot_value.expr.to_string().trim_matches('\'').to_string()
+                };
+
+                aggr_exprs.push(
+                    // Give unique name based on pivot column name
+                    agg_expr.alias(format!("{aggr_func_alias}_{pivot_value_name}")),
+                );
+            }
+        }
+
+        // Create the aggregate logical plan
+        let result_plan = LogicalPlanBuilder::from(plan)
+            .aggregate(group_by_cols, aggr_exprs)?
+            .build()?;
+
+        // Apply table alias if provided
+        if let Some(table_alias) = alias {
+            self.apply_table_alias(
+                result_plan,
+                TableAlias {
+                    name: table_alias,
+                    columns: vec![],
+                },
+            )
+        } else {
+            Ok(result_plan)
+        }
+    }
+
     /// Wrap the logical plan in a `SelectInto`
     fn select_into(
         &self,

datafusion/sqllogictest/test_files/pipe_operator.slt

@@ -195,3 +195,95 @@ query TII
 ----
 apples 2 123
 bananas 5 NULL
+
+# PIVOT pipe
+
+statement ok
+CREATE TABLE pipe_test(
+    product VARCHAR,
+    sales INT,
+    quarter VARCHAR,
+    year INT
+) AS VALUES
+   ('Kale', 51, 'Q1', 2020),
+   ('Kale', 23, 'Q2', 2020),
+   ('Kale', 45, 'Q3', 2020),
+   ('Kale', 3, 'Q4', 2020),
+   ('Kale', 70, 'Q1', 2021),
+   ('Kale', 85, 'Q2', 2021),
+   ('Apple', 77, 'Q1', 2020),
+   ('Apple', 0, 'Q2', 2020),
+   ('Apple', 1, 'Q1', 2021)
+;
+
+query TIIIII rowsort
+SELECT * FROM pipe_test
+|> PIVOT(SUM(sales) FOR quarter IN ('Q1', 'Q2', 'Q3', 'Q4'));
+----
+Apple 2020 77 0 NULL NULL
+Apple 2021 1 NULL NULL NULL
+Kale 2020 51 23 45 3
+Kale 2021 70 85 NULL NULL
+
+query TIIII rowsort
+SELECT * FROM pipe_test
+|> select product, sales, quarter
+|> PIVOT(SUM(sales) FOR quarter IN ('Q1', 'Q2', 'Q3', 'Q4'));
+----
+Apple 78 0 NULL NULL
+Kale 121 108 45 3
+
+query TIII rowsort
+SELECT * FROM pipe_test
+|> select product, sales, quarter
+|> PIVOT(SUM(sales) FOR quarter IN ('Q1', 'Q2', 'Q3'));
+----
+Apple 78 0 NULL
+Kale 121 108 45
+
+query TIIII rowsort
+SELECT * FROM pipe_test
+|> select product, sales, quarter
+|> PIVOT(SUM(sales) as total_sales, count(*) as num_records FOR quarter IN ('Q1', 'Q2'));
+----
+Apple 78 2 0 1
+Kale 121 2 108 2
+
+
+query TT
+EXPLAIN SELECT * FROM pipe_test
+|> select product, sales, quarter
+|> PIVOT(SUM(sales) as total_sales, count(*) as num_records FOR quarter IN ('Q1', 'Q2'));
+----
+logical_plan
+01)Aggregate: groupBy=[[pipe_test.product]], aggr=[[sum(CAST(CASE pipe_test.quarter WHEN Utf8View("Q1") THEN pipe_test.sales ELSE Int32(NULL) END AS Int64)) AS total_sales_Q1, count(CASE pipe_test.quarter WHEN Utf8View("Q1") THEN Int64(1) ELSE Int64(NULL) END) AS num_records_Q1, sum(CAST(CASE pipe_test.quarter WHEN Utf8View("Q2") THEN pipe_test.sales ELSE Int32(NULL) END AS Int64)) AS total_sales_Q2, count(CASE pipe_test.quarter WHEN Utf8View("Q2") THEN Int64(1) ELSE Int64(NULL) END) AS num_records_Q2]]
+02)--TableScan: pipe_test projection=[product, sales, quarter]
+physical_plan
+01)AggregateExec: mode=FinalPartitioned, gby=[product@0 as product], aggr=[total_sales_Q1, num_records_Q1, total_sales_Q2, num_records_Q2]
+02)--CoalesceBatchesExec: target_batch_size=8192
+03)----RepartitionExec: partitioning=Hash([product@0], 4), input_partitions=4
+04)------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
+05)--------AggregateExec: mode=Partial, gby=[product@0 as product], aggr=[total_sales_Q1, num_records_Q1, total_sales_Q2, num_records_Q2]
+06)----------DataSourceExec: partitions=1, partition_sizes=[1]
+
+# With explicit pivot value alias
+query TT
+EXPLAIN SELECT * FROM pipe_test
+|> select product, sales, quarter
+|> PIVOT(SUM(sales) as total_sales, count(*) as num_records FOR quarter IN ('Q1' as q1, 'Q2'));
+----
+logical_plan
+01)Aggregate: groupBy=[[pipe_test.product]], aggr=[[sum(CAST(CASE pipe_test.quarter WHEN Utf8View("Q1") THEN pipe_test.sales ELSE Int32(NULL) END AS Int64)) AS total_sales_q1, count(CASE pipe_test.quarter WHEN Utf8View("Q1") THEN Int64(1) ELSE Int64(NULL) END) AS num_records_q1, sum(CAST(CASE pipe_test.quarter WHEN Utf8View("Q2") THEN pipe_test.sales ELSE Int32(NULL) END AS Int64)) AS total_sales_Q2, count(CASE pipe_test.quarter WHEN Utf8View("Q2") THEN Int64(1) ELSE Int64(NULL) END) AS num_records_Q2]]
+02)--TableScan: pipe_test projection=[product, sales, quarter]
+physical_plan
+01)AggregateExec: mode=FinalPartitioned, gby=[product@0 as product], aggr=[total_sales_q1, num_records_q1, total_sales_Q2, num_records_Q2]
+02)--CoalesceBatchesExec: target_batch_size=8192
+03)----RepartitionExec: partitioning=Hash([product@0], 4), input_partitions=4
+04)------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
+05)--------AggregateExec: mode=Partial, gby=[product@0 as product], aggr=[total_sales_q1, num_records_q1, total_sales_Q2, num_records_Q2]
+06)----------DataSourceExec: partitions=1, partition_sizes=[1]
+
+# Aggregation functions with multiple parameters
+query error DataFusion error: This feature is not implemented: Only exactly one aggregate function argument is supported
+SELECT product, sales, quarter FROM pipe_test
+|> PIVOT(string_agg(sales, '_' order by sales) as agg FOR quarter IN ('Q1', 'Q2'));

docs/source/user-guide/sql/operators.md

@@ -636,6 +636,7 @@ DataFusion currently supports the following pipe operators:
 - [INTERSECT](#pipe_intersect)
 - [EXCEPT](#pipe_except)
 - [AGGREGATE](#pipe_aggregate)
+- [PIVOT](#pipe_pivot)
 
 (pipe_where)=
 
@@ -825,3 +826,30 @@ DataFusion currently supports the following pipe operators:
 | 3     |
 +-------+
 ```
+
+(pipe_pivot)=
+
+### PIVOT
+
+Rotates rows into columns.
+
+```sql
+> (
+  SELECT 'kale' AS product, 51 AS sales, 'Q1' AS quarter
+  UNION ALL
+  SELECT 'kale' AS product, 4 AS sales, 'Q1' AS quarter
+  UNION ALL
+  SELECT 'kale' AS product, 45 AS sales, 'Q2' AS quarter
+  UNION ALL
+  SELECT 'apple' AS product, 8 AS sales, 'Q1' AS quarter
+  UNION ALL
+  SELECT 'apple' AS product, 10 AS sales, 'Q2' AS quarter
+)
+|> PIVOT(SUM(sales) FOR quarter IN ('Q1', 'Q2'));
++---------+-----+-----+
+| product | Q1  | Q2  |
++---------+-----+-----+
+| apple   | 8   | 10  |
+| kale    | 55  | 45  |
++---------+-----+-----+
+```