viirya commented on a change in pull request #1415:
URL: https://github.com/apache/arrow-datafusion/pull/1415#discussion_r764522255
##########
File path: datafusion/src/logical_plan/builder.rs
##########
@@ -466,12 +467,126 @@ impl LogicalPlanBuilder {
})))
}
+ /// Add missing sort columns to downstream projection
+ fn add_missing_columns(
+ &self,
+ curr_plan: LogicalPlan,
+ missing_cols: &[Column],
+ ) -> LogicalPlan {
+ match curr_plan.clone() {
+ LogicalPlan::Projection(Projection {
+ input,
+ expr,
+ schema: _,
+ alias,
+ }) => {
+ let input_schema = input.schema();
+ if missing_cols
+ .iter()
+ .all(|c| input_schema.field_from_column(c).is_ok())
+ {
+ let mut new_expr = Vec::new();
+ new_expr.extend(expr);
+ missing_cols.iter().for_each(|c| {
+ new_expr.push(
+ normalize_col(Expr::Column(c.clone()),
&input).unwrap(),
+ );
+ });
+
+ let new_inputs =
+
curr_plan.inputs().into_iter().cloned().collect::<Vec<_>>();
+
+ let new_schema = DFSchema::new(
+ exprlist_to_fields(&new_expr, input_schema).unwrap(),
+ )
+ .unwrap();
+
+ LogicalPlan::Projection(Projection {
+ expr: new_expr,
+ input: Arc::new(new_inputs.get(0).unwrap().clone()),
+ schema: DFSchemaRef::new(new_schema),
+ alias,
+ })
+ } else {
+ let inputs = curr_plan.inputs();
+ let new_inputs = inputs
+ .iter()
+ .map(|input_plan| {
+ self.add_missing_columns((*input_plan).clone(),
missing_cols)
+ })
+ .collect::<Vec<_>>();
+
+ let expr = curr_plan.expressions();
+ utils::from_plan(&curr_plan, &expr, &new_inputs).unwrap()
Review comment:
oh, `if ...` can be pulled up. I'll update it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscr...@arrow.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
