Pivot initial changes

Author: Denys Tsomenko

datafusion/core/src/execution/session_state.rs

@@ -554,6 +554,28 @@ 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();
 
@@ -643,7 +665,9 @@ 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

@@ -65,22 +65,23 @@ 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,
-    ScalarValue,
+    exec_err, internal_datafusion_err, internal_err, not_impl_err, plan_err, DFSchema, DFSchemaRef,
+    ScalarValue, Column, TableReference,
 };
 use datafusion_datasource::memory::MemorySourceConfig;
-use datafusion_expr::dml::{CopyTo, InsertOp};
+use datafusion_expr::dml::{CopyTo, InsertOp, DmlStatement, WriteOp};
 use datafusion_expr::expr::{
     physical_name, AggregateFunction, AggregateFunctionParams, Alias, GroupingSet,
     WindowFunction, WindowFunctionParams,
 };
 use datafusion_expr::expr_rewriter::unnormalize_cols;
 use datafusion_expr::logical_plan::builder::wrap_projection_for_join_if_necessary;
 use datafusion_expr::{
-    Analyze, DescribeTable, DmlStatement, Explain, ExplainFormat, Extension, FetchType,
+    Analyze, DescribeTable, Explain, ExplainFormat, Extension, FetchType,
     Filter, JoinType, RecursiveQuery, SkipType, SortExpr, StringifiedPlan, WindowFrame,
-    WindowFrameBound, WriteOp,
+    WindowFrameBound, SubqueryAlias,
 };
+use datafusion_execution::FunctionRegistry;
 use datafusion_physical_expr::aggregate::{AggregateExprBuilder, AggregateFunctionExpr};
 use datafusion_physical_expr::expressions::Literal;
 use datafusion_physical_expr::LexOrdering;
@@ -97,6 +98,10 @@ use log::{debug, trace};
 use sqlparser::ast::NullTreatment;
 use tokio::sync::Mutex;
 
+use datafusion_sql::transform_pivot_to_aggregate;
+
+use datafusion_physical_plan::collect;
+
 /// Physical query planner that converts a `LogicalPlan` to an
 /// `ExecutionPlan` suitable for execution.
 #[async_trait]
@@ -304,8 +309,10 @@ impl DefaultPhysicalPlanner {
                 }
                 1 => NodeState::ZeroOrOneChild,
                 _ => {
+
                     let ready_children = Vec::with_capacity(node.inputs().len());
                     let ready_children = Mutex::new(ready_children);
+
                     NodeState::TwoOrMoreChildren(ready_children)
                 }
             };
@@ -887,7 +894,53 @@ impl DefaultPhysicalPlanner {
                     options.clone(),
                 ))
             }
-
+            LogicalPlan::Pivot(pivot) => {
+                let pivot_values = if let Some(subquery) = &pivot.value_subquery {
+                    let optimized_subquery = session_state.optimize(subquery.as_ref())?;
+                    
+                    let subquery_physical_plan = self.create_physical_plan(
+                        &optimized_subquery,
+                        session_state
+                    ).await?;
+                    
+                    let subquery_results = collect(subquery_physical_plan.clone(), session_state.task_ctx()).await?;
+                    
+                    let mut pivot_values = Vec::new();
+                    for batch in subquery_results.iter() {
+                        if batch.num_columns() != 1 {
+                            return plan_err!("Pivot subquery must return a single column");
+                        }
+                        
+                        let column = batch.column(0);
+                        for row_idx in 0..batch.num_rows() {
+                            if !column.is_null(row_idx) {
+                                pivot_values.push(
+                                    ScalarValue::try_from_array(column, row_idx)?
+                                );
+                            }
+                        }
+                    }
+                    pivot_values
+                } else {
+                    pivot.pivot_values.clone()
+                };
+                
+                if !pivot_values.is_empty() {
+                    // Transform Pivot into Aggregate plan with the resolved pivot values
+                    let agg_plan = transform_pivot_to_aggregate(
+                        Arc::new(pivot.input.as_ref().clone()),
+                        &pivot.aggregate_expr,
+                        &pivot.pivot_column,
+                        Some(pivot_values),
+                        None, 
+                    )?;
+                    
+                    // The schema information is already preserved in the agg_plan
+                    return self.create_physical_plan(&agg_plan, session_state).await;
+                } else {
+                    return plan_err!("PIVOT operation requires at least one value to pivot on");
+                }
+            }
             // 2 Children
             LogicalPlan::Join(Join {
                 left,
@@ -2067,6 +2120,44 @@ impl DefaultPhysicalPlanner {
             })
             .collect::<Result<Vec<_>>>()?;
 
+        // Check if input plan was transformed from a PIVOT operation
+        // The original logical plan will be a LogicalPlan::Pivot that has been transformed to an Aggregate
+        match input.as_ref() {
+            // Direct PIVOT
+            LogicalPlan::Pivot(_) => {
+                // When we detect a PIVOT-derived plan, ensure all generated columns are preserved
+                if input_exec.as_any().downcast_ref::<AggregateExec>().is_some() {
+                    let agg_exec = input_exec.as_any().downcast_ref::<AggregateExec>().unwrap();
+                    let schema = input_exec.schema();
+                    let group_by_len = agg_exec.group_expr().expr().len();
+
+                    // If we have aggregate expressions that correspond to pivot columns
+                    if group_by_len < schema.fields().len() {
+                        // This is a pivot result - we need to include all columns from the aggregate
+                        let mut all_exprs = physical_exprs.clone();
+
+                        // Add any missing pivot columns (which are all columns after the group_by columns)
+                        for (i, field) in schema.fields().iter().enumerate().skip(group_by_len) {
+                            // Check if this column is already included in the projection
+                            if !physical_exprs.iter().any(|(_, name)| name == field.name()) {
+                                // Add this pivot column to the projection
+                                all_exprs.push((
+                                    Arc::new(crate::physical_plan::expressions::Column::new(field.name(), i)) as Arc<dyn PhysicalExpr>,
+                                    field.name().clone(),
+                                ));
+                            }
+                        }
+
+                        return Ok(Arc::new(ProjectionExec::try_new(
+                            all_exprs,
+                            input_exec,
+                        )?));
+                    }
+                }
+            },
+           _ => {}
+        }
+        
         Ok(Arc::new(ProjectionExec::try_new(
             physical_exprs,
             input_exec,

datafusion/expr/src/logical_plan/builder.rs

@@ -34,7 +34,7 @@ use crate::logical_plan::{
     Aggregate, Analyze, Distinct, DistinctOn, EmptyRelation, Explain, Filter, Join,
     JoinConstraint, JoinType, Limit, LogicalPlan, Partitioning, PlanType, Prepare,
     Projection, Repartition, Sort, SubqueryAlias, TableScan, Union, Unnest, Values,
-    Window,
+    Window, Pivot,
 };
 use crate::select_expr::SelectExpr;
 use crate::utils::{
@@ -1426,6 +1426,21 @@ impl LogicalPlanBuilder {
         unnest_with_options(Arc::unwrap_or_clone(self.plan), columns, options)
             .map(Self::new)
     }
+
+    pub fn pivot(
+        self,
+        aggregate_expr: Expr,
+        pivot_column: Column,
+        pivot_values: Vec<ScalarValue>,
+    ) -> Result<Self> {
+        let pivot_plan = Pivot::try_new(
+            self.plan,
+            aggregate_expr,
+            pivot_column,
+            pivot_values,
+        )?;
+        Ok(Self::new(LogicalPlan::Pivot(pivot_plan)))
+    }
 }
 
 impl From<LogicalPlan> for LogicalPlanBuilder {

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,
+    Unnest, Values, Window, Pivot,
 };
 
 use crate::dml::CopyTo;
@@ -650,6 +650,7 @@ impl<'a, 'b> PgJsonVisitor<'a, 'b> {
                     "StructColumn": expr_vec_fmt!(struct_type_columns),
                 })
             }
+            LogicalPlan::Pivot(_) => todo!(),
         }
     }
 }

datafusion/expr/src/logical_plan/mod.rs

@@ -41,7 +41,7 @@ pub use plan::{
     DistinctOn, EmptyRelation, Explain, ExplainFormat, Extension, FetchType, Filter,
     Join, JoinConstraint, JoinType, Limit, LogicalPlan, Partitioning, PlanType,
     Projection, RecursiveQuery, Repartition, SkipType, Sort, StringifiedPlan, Subquery,
-    SubqueryAlias, TableScan, ToStringifiedPlan, Union, Unnest, Values, Window,
+    SubqueryAlias, TableScan, ToStringifiedPlan, Union, Unnest, Values, Window, Pivot
 };
 pub use statement::{
     Deallocate, Execute, Prepare, SetVariable, Statement, TransactionAccessMode,