xudong963 commented on a change in pull request #1415:
URL: https://github.com/apache/arrow-datafusion/pull/1415#discussion_r764504778
##########
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()
+ }
+ }
+ _ => {
+ 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()
+ }
+ }
+ }
+
/// Apply a sort
- pub fn sort(&self, exprs: impl IntoIterator<Item = impl Into<Expr>>) ->
Result<Self> {
- Ok(Self::from(LogicalPlan::Sort(Sort {
- expr: normalize_cols(exprs, &self.plan)?,
- input: Arc::new(self.plan.clone()),
- })))
+ pub fn sort(
+ &self,
+ exprs: impl IntoIterator<Item = impl Into<Expr>> + Clone,
+ ) -> Result<Self> {
+ let schema = self.plan.schema();
+
+ // Collect sort columns that are missing in the input plan's schema
+ let mut missing_cols: Vec<Column> = vec![];
+ exprs
+ .clone()
+ .into_iter()
+ .try_for_each::<_, Result<()>>(|expr| {
+ let mut columns: HashSet<Column> = HashSet::new();
+ utils::expr_to_columns(&expr.into(), &mut columns)?;
+
+ columns.into_iter().for_each(|c| {
+ if schema.field_from_column(&c).is_err() {
+ missing_cols.push(c);
+ }
+ });
+
+ Ok(())
+ })?;
+
+ if !missing_cols.is_empty() {
Review comment:
You can use early return and put short code in `if`
```rust
if missing_cols.is_emptry() {
return Ok();
}
...
```
##########
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()
+ }
+ }
+ _ => {
+ 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()
+ }
+ }
+ }
+
/// Apply a sort
- pub fn sort(&self, exprs: impl IntoIterator<Item = impl Into<Expr>>) ->
Result<Self> {
- Ok(Self::from(LogicalPlan::Sort(Sort {
- expr: normalize_cols(exprs, &self.plan)?,
- input: Arc::new(self.plan.clone()),
- })))
+ pub fn sort(
+ &self,
+ exprs: impl IntoIterator<Item = impl Into<Expr>> + Clone,
+ ) -> Result<Self> {
+ let schema = self.plan.schema();
+
+ // Collect sort columns that are missing in the input plan's schema
+ let mut missing_cols: Vec<Column> = vec![];
+ exprs
+ .clone()
+ .into_iter()
+ .try_for_each::<_, Result<()>>(|expr| {
+ let mut columns: HashSet<Column> = HashSet::new();
+ utils::expr_to_columns(&expr.into(), &mut columns)?;
+
+ columns.into_iter().for_each(|c| {
+ if schema.field_from_column(&c).is_err() {
+ missing_cols.push(c);
+ }
+ });
+
+ Ok(())
+ })?;
+
+ if !missing_cols.is_empty() {
+ let plan = self.add_missing_columns(self.plan.clone(),
&missing_cols);
+ let sort_plan = LogicalPlan::Sort(Sort {
+ expr: normalize_cols(exprs, &plan)?,
+ input: Arc::new(plan.clone()),
+ });
+ // remove pushed down sort columns
Review comment:
I think we don't need to remove all pushed-down sort columns, removing
unselected columns may be enough.
##########
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:
duplicate with 524~533, we can simplify it?
##########
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()
+ }
+ }
+ _ => {
+ 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()
+ }
+ }
+ }
+
/// Apply a sort
- pub fn sort(&self, exprs: impl IntoIterator<Item = impl Into<Expr>>) ->
Result<Self> {
- Ok(Self::from(LogicalPlan::Sort(Sort {
- expr: normalize_cols(exprs, &self.plan)?,
- input: Arc::new(self.plan.clone()),
- })))
+ pub fn sort(
+ &self,
+ exprs: impl IntoIterator<Item = impl Into<Expr>> + Clone,
+ ) -> Result<Self> {
+ let schema = self.plan.schema();
+
+ // Collect sort columns that are missing in the input plan's schema
+ let mut missing_cols: Vec<Column> = vec![];
+ exprs
+ .clone()
+ .into_iter()
+ .try_for_each::<_, Result<()>>(|expr| {
+ let mut columns: HashSet<Column> = HashSet::new();
+ utils::expr_to_columns(&expr.into(), &mut columns)?;
+
+ columns.into_iter().for_each(|c| {
+ if schema.field_from_column(&c).is_err() {
+ missing_cols.push(c);
+ }
+ });
+
+ Ok(())
+ })?;
+
+ if !missing_cols.is_empty() {
+ let plan = self.add_missing_columns(self.plan.clone(),
&missing_cols);
+ let sort_plan = LogicalPlan::Sort(Sort {
+ expr: normalize_cols(exprs, &plan)?,
+ input: Arc::new(plan.clone()),
+ });
+ // remove pushed down sort columns
+ let mut new_expr = Vec::new();
Review comment:
```suggestion
let mut new_expr = Vec::with_capacity(schema.fields().len());
```
--
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
