Only check projections in case of subqueries

Author: Vedin

datafusion/core/src/physical_planner.rs

@@ -2160,27 +2160,19 @@ 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
+        // When we detect a PIVOT-derived plan with a value_subquery, ensure all generated columns are preserved
         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() {
+            LogicalPlan::Pivot(pivot) => {
+                if pivot.value_subquery.is_some() && 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(Column::new(field.name(), i)) as Arc<dyn PhysicalExpr>,
                                     field.name().clone(),