fix: fix issue 10326, add missing columns into list directory,instead…

… of traversing the plan looking for projection node
apache · Jan 18, 2025 · 3a8c04c · 3a8c04c
1 parent 0e22172
commit 3a8c04c
Showing 1 changed file with 14 additions and 27 deletions.
diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs
@@ -575,44 +575,31 @@ impl LogicalPlanBuilder {
         missing_cols: &IndexSet<Column>,
         is_distinct: bool,
     ) -> Result<LogicalPlan> {
-        match curr_plan {
-            LogicalPlan::Projection(Projection {
-                input,
-                mut expr,
-                schema: _,
-            }) if missing_cols.iter().all(|c| input.schema().has_column(c)) => {
+        let inputs = curr_plan.inputs();
+        let mut exprs = curr_plan.expressions();
+        for input in inputs {
+            if missing_cols.iter().all(|c| input.schema().has_column(c)) {
                 let mut missing_exprs = missing_cols
                     .iter()
-                    .map(|c| normalize_col(Expr::Column(c.clone()), &input))
+                    .map(|c| normalize_col(Expr::Column(c.clone()), input))
                     .collect::<Result<Vec<_>>>()?;
 
                 // Do not let duplicate columns to be added, some of the
                 // missing_cols may be already present but without the new
                 // projected alias.
-                missing_exprs.retain(|e| !expr.contains(e));
+                missing_exprs.retain(|e| !exprs.contains(e));
                 if is_distinct {
-                    Self::ambiguous_distinct_check(&missing_exprs, missing_cols, &expr)?;
+                    Self::ambiguous_distinct_check(&missing_exprs, missing_cols, &exprs)?;
                 }
-                expr.extend(missing_exprs);
-                project(Arc::unwrap_or_clone(input), expr)
-            }
-            _ => {
-                let is_distinct =
-                    is_distinct || matches!(curr_plan, LogicalPlan::Distinct(_));
-                let new_inputs = curr_plan
-                    .inputs()
-                    .into_iter()
-                    .map(|input_plan| {
-                        Self::add_missing_columns(
-                            (*input_plan).clone(),
-                            missing_cols,
-                            is_distinct,
-                        )
-                    })
-                    .collect::<Result<Vec<_>>>()?;
-                curr_plan.with_new_exprs(curr_plan.expressions(), new_inputs)
+                exprs.extend(missing_exprs);
             }
         }
+        let inputs = curr_plan
+            .inputs()
+            .into_iter()
+            .map(|input| input.to_owned())
+            .collect::<Vec<_>>();
+        curr_plan.with_new_exprs(exprs, inputs)
     }
 
     fn ambiguous_distinct_check(