You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by GitBox <gi...@apache.org> on 2021/12/08 06:27:35 UTC
[GitHub] [arrow-datafusion] houqp commented on a change in pull request #1415: Support sort on unprojected columns
houqp commented on a change in pull request #1415:
URL: https://github.com/apache/arrow-datafusion/pull/1415#discussion_r764579879
##########
File path: datafusion/src/logical_plan/builder.rs
##########
@@ -466,12 +467,112 @@ 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,
+ }) if missing_cols
+ .iter()
+ .all(|c| input.schema().field_from_column(c).is_ok()) =>
+ {
+ let input_schema = input.schema();
+
+ 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,
+ })
+ }
+ _ => {
+ 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::with_capacity(schema.fields().len());
+ schema.fields().iter().for_each(|f| {
+ let col = f.qualified_column();
+ new_expr.push(Expr::Column(col));
+ });
Review comment:
nitpick, you can collect the new_expr directly to leave more room for the compiler to do optimization. I didn't check the machine code here to see if it will make a difference here, but in general a good practice to keep.
```suggestion
let new_expr = schema.fields().iter().map(|f| Expr::Column(f.qualified_column())).collect();
```
##########
File path: datafusion/src/logical_plan/builder.rs
##########
@@ -466,12 +467,112 @@ 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,
+ }) if missing_cols
+ .iter()
+ .all(|c| input.schema().field_from_column(c).is_ok()) =>
+ {
+ let input_schema = input.schema();
+
+ 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,
+ })
+ }
+ _ => {
+ 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::with_capacity(schema.fields().len());
+ schema.fields().iter().for_each(|f| {
+ let col = f.qualified_column();
+ new_expr.push(Expr::Column(col));
+ });
+ let new_schema =
+ DFSchema::new(exprlist_to_fields(&new_expr, schema).unwrap()).unwrap();
Review comment:
any reason why we want to call unwrap here instead of propagating the error back to the caller?
--
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org