Polishing

Author: Vedin

datafusion/core/src/execution/session_state.rs

@@ -562,28 +562,6 @@ impl SessionState {
 
     /// Optimizes the logical plan by applying optimizer rules.
     pub fn optimize(&self, plan: &LogicalPlan) -> datafusion_common::Result<LogicalPlan> {
-        // Special case for Pivot nodes with subqueries
-        /*if let LogicalPlan::Pivot(p) = plan {
-            if let Some(subquery) = &p.value_subquery {
-                // Optimize the subquery first
-                let optimized_subquery = self.optimizer.optimize(
-                    subquery.as_ref().clone(),
-                    self,
-                    |_, _| {},
-                )?;
-                
-                // Create a new Pivot with the optimized subquery
-                return Ok(LogicalPlan::Pivot(datafusion_expr::Pivot {
-                    input: p.input.clone(),
-                    aggregate_expr: p.aggregate_expr.clone(),
-                    pivot_column: p.pivot_column.clone(),
-                    pivot_values: p.pivot_values.clone(),
-                    schema: p.schema.clone(),
-                    value_subquery: Some(Arc::new(optimized_subquery)),
-                }));
-            }
-        }*/
-
         if let LogicalPlan::Explain(e) = plan {
             let mut stringified_plans = e.stringified_plans.clone();
 
@@ -673,9 +651,7 @@ impl SessionState {
         &self,
         logical_plan: &LogicalPlan,
     ) -> datafusion_common::Result<Arc<dyn ExecutionPlan>> {
-        println!("logical_plan_beore_optPIVOT328: {:#?}", logical_plan);
         let logical_plan = self.optimize(logical_plan)?;
-        println!("logical_plan_PIVOT329: {:#?}", logical_plan);
         self.query_planner
             .create_physical_plan(&logical_plan, self)
             .await

datafusion/core/src/physical_planner.rs

@@ -64,8 +64,8 @@ use arrow::datatypes::{Schema, SchemaRef};
 use datafusion_common::display::ToStringifiedPlan;
 use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion, TreeNodeVisitor};
 use datafusion_common::{
-    exec_err, internal_datafusion_err, internal_err, not_impl_err, plan_err, DFSchema, DFSchemaRef,
-    ScalarValue, Column, TableReference,
+    exec_err, internal_datafusion_err, internal_err, not_impl_err, plan_err, DFSchema,
+    ScalarValue, Column,
 };
 use datafusion_datasource::memory::MemorySourceConfig;
 use datafusion_expr::dml::{CopyTo, InsertOp, DmlStatement, WriteOp};
@@ -78,7 +78,7 @@ use datafusion_expr::logical_plan::builder::wrap_projection_for_join_if_necessar
 use datafusion_expr::{
     Analyze, DescribeTable, DmlStatement,Explain, ExplainFormat, Extension, FetchType,
     Filter, JoinType, RecursiveQuery, SkipType, SortExpr, StringifiedPlan, WindowFrame,
-    WindowFrameBound, WriteOp, SubqueryAlias,
+    WindowFrameBound, WriteOp, SubqueryAlias, LogicalPlanBuilder, BinaryExpr
 };
 use datafusion_execution::FunctionRegistry;
 use datafusion_physical_expr::aggregate::{AggregateExprBuilder, AggregateFunctionExpr};
@@ -97,8 +97,7 @@ use itertools::{multiunzip, Itertools};
 use log::{debug, trace};
 use sqlparser::ast::NullTreatment;
 use tokio::sync::Mutex;
-
-use datafusion_sql::transform_pivot_to_aggregate;
+use datafusion_expr_common::operator::Operator;
 
 use datafusion_physical_plan::collect;
 
@@ -923,20 +922,17 @@ impl DefaultPhysicalPlanner {
                     pivot.pivot_values.clone()
                 };
 
-                if !pivot_values.is_empty() {
-                    // Transform Pivot into Aggregate plan with the resolved pivot values
+                return if !pivot_values.is_empty() {
                     let agg_plan = transform_pivot_to_aggregate(
                         Arc::new(pivot.input.as_ref().clone()),
                         &pivot.aggregate_expr,
                         &pivot.pivot_column,
-                        Some(pivot_values),
-                        None,
+                        pivot_values,
                     )?;
 
-                    // The schema information is already preserved in the agg_plan
-                    return self.create_physical_plan(&agg_plan, session_state).await;
+                    self.create_physical_plan(&agg_plan, session_state).await
                 } else {
-                    return plan_err!("PIVOT operation requires at least one value to pivot on");
+                    plan_err!("PIVOT operation requires at least one value to pivot on")
                 }
             }
             // 2 Children
@@ -1734,6 +1730,76 @@ pub use datafusion_physical_expr::{
     create_physical_sort_expr, create_physical_sort_exprs,
 };
 
+/// Transform a PIVOT operation into a more standard Aggregate + Projection plan
+/// For known pivot values, we create a projection that includes "IS NOT DISTINCT FROM" conditions
+///
+/// For example, for SUM(amount) PIVOT(quarter FOR quarter in ('2023_Q1', '2023_Q2')), we create:
+/// - SUM(amount) FILTER (WHERE quarter IS NOT DISTINCT FROM '2023_Q1') AS "2023_Q1"
+/// - SUM(amount) FILTER (WHERE quarter IS NOT DISTINCT FROM '2023_Q2') AS "2023_Q2"
+///
+pub fn transform_pivot_to_aggregate(
+    input: Arc<LogicalPlan>,
+    aggregate_expr: &Expr,
+    pivot_column: &Column,
+    pivot_values: Vec<ScalarValue>,
+) -> Result<LogicalPlan> {
+    let df_schema = input.schema();
+
+    let all_columns: Vec<Column> = df_schema.columns();
+
+    // Filter to include only columns we want for GROUP BY
+    // (exclude pivot column and aggregate expression columns)
+    let group_by_columns: Vec<Expr> = all_columns
+        .into_iter()
+        .filter(|col| {
+            col.name != pivot_column.name
+                && !aggregate_expr.column_refs().iter().any(|agg_col| agg_col.name == col.name)
+        })
+        .map(|col| Expr::Column(col))
+        .collect();
+
+    let builder = LogicalPlanBuilder::from(Arc::unwrap_or_clone(input.clone()));
+
+    let mut aggregate_exprs = Vec::new();
+
+    for value in &pivot_values {
+        let filter_condition = Expr::BinaryExpr(BinaryExpr::new(
+            Box::new(Expr::Column(pivot_column.clone())),
+            Operator::IsNotDistinctFrom,
+            Box::new(Expr::Literal(value.clone()))
+        ));
+
+        let filtered_agg = match aggregate_expr {
+            Expr::AggregateFunction(agg) => {
+                let mut new_params = agg.params.clone();
+                new_params.filter = Some(Box::new(filter_condition));
+                Expr::AggregateFunction(AggregateFunction {
+                    func: agg.func.clone(),
+                    params: new_params,
+                })
+            },
+            _ => {
+                return plan_err!("Unsupported aggregate expression should always be AggregateFunction");
+            }
+        };
+
+        // Use the pivot value as the column name
+        let field_name = value.to_string().trim_matches('\'').to_string();
+        let aliased_agg = Expr::Alias(Alias {
+            expr: Box::new(filtered_agg),
+            relation: None,
+            name: field_name,
+            metadata: None,
+        });
+
+        aggregate_exprs.push(aliased_agg);
+    }
+
+    let aggregate_plan = builder.aggregate(group_by_columns, aggregate_exprs)?.build()?;
+
+    Ok(aggregate_plan)
+}
+
 impl DefaultPhysicalPlanner {
     /// Handles capturing the various plans for EXPLAIN queries
     ///

datafusion/expr/src/logical_plan/display.rs

@@ -24,7 +24,7 @@ use crate::{
     expr_vec_fmt, Aggregate, DescribeTable, Distinct, DistinctOn, DmlStatement, Expr,
     Filter, Join, Limit, LogicalPlan, Partitioning, Projection, RecursiveQuery,
     Repartition, Sort, Subquery, SubqueryAlias, TableProviderFilterPushDown, TableScan,
-    Unnest, Values, Window, Pivot,
+    Unnest, Values, Window,
 };
 
 use crate::dml::CopyTo;

datafusion/expr/src/logical_plan/plan.rs

@@ -20,7 +20,6 @@
 use std::cmp::Ordering;
 use std::collections::{HashMap, HashSet};
 use std::fmt::{self, Debug, Display, Formatter};
-use std::fs::metadata;
 use std::hash::{Hash, Hasher};
 use std::str::FromStr;
 use std::sync::{Arc, LazyLock};
@@ -1179,10 +1178,10 @@ impl LogicalPlan {
                 Ok(new_plan)
             }
             LogicalPlan::Pivot(Pivot {
-                aggregate_expr,
+                aggregate_expr: _,
                 pivot_column,
                 pivot_values,
-                schema,
+                schema: _,
                 value_subquery,
                 ..
             }) => {
@@ -2345,7 +2344,6 @@ impl Pivot {
         pivot_column: Column,
         value_subquery: Arc<LogicalPlan>,
     ) -> Result<Self> {
-        // Create an empty schema - will be filled in when the subquery is executed
         let schema = pivot_schema_without_values(
             input.schema(),
             &aggregate_expr,
@@ -2356,15 +2354,13 @@ impl Pivot {
             input,
             aggregate_expr,
             pivot_column,
-            pivot_values: Vec::new(), // Will be populated during physical planning
+            pivot_values: Vec::new(),
             schema: Arc::new(schema),
             value_subquery: Some(value_subquery),
         })
     }
 }
 
-/// Create a pivot schema without knowing the pivot values
-/// This is used when we have a subquery for pivot values
 fn pivot_schema_without_values(
     input_schema: &DFSchemaRef,
     aggregate_expr: &Expr,
@@ -2387,7 +2383,6 @@ fn pivot_schema_without_values(
     DFSchema::new_with_metadata(fields_with_table_ref, input_schema.metadata().clone())
 }
 
-/// Create a pivot schema with known pivot values
 fn pivot_schema(
     input_schema: &DFSchemaRef,
     aggregate_expr: &Expr,
@@ -2396,14 +2391,12 @@ fn pivot_schema(
 ) -> Result<DFSchema> {
     let mut fields = vec![];
 
-    // Include all fields except pivot and value columns
     for field in input_schema.fields() {
         if !aggregate_expr.column_refs().iter().any(|col| col.name() == field.name()) && field.name() != pivot_column.name() {
             fields.push(field.clone());
         }
     }
 
-    // Add new fields for each pivot value
     for pivot_value in pivot_values {
         let field_name = format!("{}", pivot_value);
         let data_type = aggregate_expr.get_type(input_schema)?;
@@ -5558,4 +5551,3 @@ digraph {
         Ok(())
     }
 }
-

datafusion/optimizer/src/optimize_projections/mod.rs

@@ -382,23 +382,8 @@ fn optimize_projections(
                 dependency_indices.clone(),
             )]
         }
-        LogicalPlan::Pivot(pivot) => {
-            return Ok(Transformed::no(plan)); // TODO: Implement this
-            /*
-            // For PIVOT operations, we need columns from the aggregate expression and the pivot column
-            let mut pivot_indices = Vec::new();
-            
-            // Add required indices for the pivot column
-            if let Ok(idx) = pivot.input.schema().index_of_column(&pivot.pivot_column) {
-                pivot_indices.push(idx);
-            }
-            
-            // Create RequiredIndices with these indices and add dependency from aggregate expression
-            let required = RequiredIndices::new_from_indices(pivot_indices)
-                .with_exprs(pivot.input.schema(), std::iter::once(&pivot.aggregate_expr))
-                .with_projection_beneficial();
-                
-            vec![required]*/
+        LogicalPlan::Pivot(_) => {
+            return Ok(Transformed::no(plan));
         },
 
     };

datafusion/proto/src/logical_plan/mod.rs

@@ -1838,19 +1838,6 @@ impl AsLogicalPlan for LogicalPlanNode {
                 })
             }
             LogicalPlan::Pivot(_) => {
-                /*let input =
-                    LogicalPlanNode::try_from_logical_plan(pivot.input.as_ref(), extension_codec)?;
-                Ok(LogicalPlanNode {
-                    logical_plan_type: Some(LogicalPlanType::Pivot(Box::new(
-                        protobuf::PivotNode {
-                            input: Some(Box::new(input)),
-                            aggregate_expr: pivot.aggregate_expr.clone(),
-                            pivot_column: pivot.pivot_column.clone(),
-                            pivot_values: pivot.pivot_values.clone(),
-                            schema:  convert_required!(*pivot.schema)?,
-                        },
-                    ))),
-                })*/
                 Err(proto_error(
                     "LogicalPlan serde is not yet implemented for Statement",
                 ))

datafusion/sql/src/lib.rs

@@ -62,6 +62,3 @@ mod values;
 )]
 pub use datafusion_common::{ResolvedTableReference, TableReference};
 pub use sqlparser;
-
-// Re-export the transform_pivot_to_aggregate function
-pub use relation::transform_pivot_to_aggregate;