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/07 20:36:33 UTC
[GitHub] [arrow-datafusion] viirya opened a new pull request #1415: Support sort on unprojected columns
viirya opened a new pull request #1415:
URL: https://github.com/apache/arrow-datafusion/pull/1415
# Which issue does this PR close?
<!--
We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. For example `Closes #123` indicates that this PR will close issue #123.
-->
Closes #1372.
# Rationale for this change
<!--
Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes.
-->
# What changes are included in this PR?
<!--
There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR.
-->
# Are there any user-facing changes?
<!--
If there are user-facing changes then we may require documentation to be updated before approving the PR.
-->
<!--
If there are any breaking changes to public APIs, please add the `api change` label.
-->
--
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
[GitHub] [arrow-datafusion] viirya commented on a change in pull request #1415: Support sort on unprojected columns
Posted by GitBox <gi...@apache.org>.
viirya commented on a change in pull request #1415:
URL: https://github.com/apache/arrow-datafusion/pull/1415#discussion_r764534291
##########
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:
Hmm? I think the first plan has only one `Sort` and the second one has two, no? Why only one `Sort` will sort the table twice, but two `Sort` sorts it once?
I also don't get "unselected columns". Currently this pushes missing columns down. These missing columns are unselected at all because they are not in original schema, we need to remove them, no?
--
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
[GitHub] [arrow-datafusion] viirya commented on a change in pull request #1415: Support sort on unprojected columns
Posted by GitBox <gi...@apache.org>.
viirya commented on a change in pull request #1415:
URL: https://github.com/apache/arrow-datafusion/pull/1415#discussion_r764534291
##########
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:
Hmm? I think the first plan has only one `Sort` and the second one has two, no? Why only one `Sort` will sort the table twice, but two `Sort` sorts it once?
I also don't get "unselected columns". Currently this pushes missing columns down. These missing columns are not in original schema, we need to remove them, no?
--
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
[GitHub] [arrow-datafusion] viirya commented on a change in pull request #1415: Support sort on unprojected columns
Posted by GitBox <gi...@apache.org>.
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] houqp commented on a change in pull request #1415: Support sort on unprojected columns
Posted by GitBox <gi...@apache.org>.
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. It will often leverage the size hint to preallocate memory and perform vectorization whenever applicable. 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();
```
--
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
[GitHub] [arrow-datafusion] xudong963 commented on pull request #1415: Support sort on unprojected columns
Posted by GitBox <gi...@apache.org>.
xudong963 commented on pull request #1415:
URL: https://github.com/apache/arrow-datafusion/pull/1415#issuecomment-988433419
It'll be better to also add a test in `datafusion/tests/sql.rs`.
--
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
[GitHub] [arrow-datafusion] xudong963 commented on a change in pull request #1415: Support sort on unprojected columns
Posted by GitBox <gi...@apache.org>.
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-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [arrow-datafusion] alamb merged pull request #1415: Support sort on unprojected columns
Posted by GitBox <gi...@apache.org>.
alamb merged pull request #1415:
URL: https://github.com/apache/arrow-datafusion/pull/1415
--
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
[GitHub] [arrow-datafusion] Dandandan commented on pull request #1415: Support sort on unprojected columns
Posted by GitBox <gi...@apache.org>.
Dandandan commented on pull request #1415:
URL: https://github.com/apache/arrow-datafusion/pull/1415#issuecomment-989986284
Nice 🎉
--
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
[GitHub] [arrow-datafusion] alamb commented on a change in pull request #1415: Support sort on unprojected columns
Posted by GitBox <gi...@apache.org>.
alamb commented on a change in pull request #1415:
URL: https://github.com/apache/arrow-datafusion/pull/1415#discussion_r765914983
##########
File path: datafusion/tests/sql.rs
##########
@@ -6351,6 +6351,21 @@ async fn test_expect_distinct() -> Result<()> {
Ok(())
}
+#[tokio::test]
+async fn test_sort_unprojected_col() -> Result<()> {
+ let mut ctx = ExecutionContext::new();
+ register_alltypes_parquet(&mut ctx).await;
+ // execute the query
+ let sql = "SELECT id FROM alltypes_plain ORDER BY int_col, double_col";
Review comment:
👍
--
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
[GitHub] [arrow-datafusion] viirya commented on a change in pull request #1415: Support sort on unprojected columns
Posted by GitBox <gi...@apache.org>.
viirya commented on a change in pull request #1415:
URL: https://github.com/apache/arrow-datafusion/pull/1415#discussion_r764551149
##########
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:
np, thanks for review.
--
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
[GitHub] [arrow-datafusion] viirya commented on a change in pull request #1415: Support sort on unprojected columns
Posted by GitBox <gi...@apache.org>.
viirya commented on a change in pull request #1415:
URL: https://github.com/apache/arrow-datafusion/pull/1415#discussion_r764531918
##########
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:
Oh, I see. You mean removing `else` block.
--
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
[GitHub] [arrow-datafusion] xudong963 commented on a change in pull request #1415: Support sort on unprojected columns
Posted by GitBox <gi...@apache.org>.
xudong963 commented on a change in pull request #1415:
URL: https://github.com/apache/arrow-datafusion/pull/1415#discussion_r764527294
##########
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:
for example
This is your current implementation, it'll sort the whole table twice.
```
❯ explain select column1 from foo order by column1, column2;
+---------------+-----------------------------------------------------------------------------------------+
| plan_type | plan |
+---------------+-----------------------------------------------------------------------------------------+
| logical_plan | Projection: #foo.column1 |
| | Sort: #foo.column1 ASC NULLS LAST, #foo.column2 ASC NULLS LAST |
| | TableScan: foo projection=Some([0, 1])
```
If we only remove unselected columns, the plan will be the following. This will sort the whole table once and sort column1 once
```
❯ explain select column1 from foo order by column1, column2;
+---------------+-----------------------------------------------------------------------------------------+
| plan_type | plan |
+---------------+-----------------------------------------------------------------------------------------+
| logical_plan | Sort: #foo.column1 ASC NULLS LAST,
| | Projection: #foo.column1 |
| | Sort: #foo.column2 ASC NULLS LAST |
| | TableScan: foo projection=Some([0, 1])
```
--
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
[GitHub] [arrow-datafusion] xudong963 commented on a change in pull request #1415: Support sort on unprojected columns
Posted by GitBox <gi...@apache.org>.
xudong963 commented on a change in pull request #1415:
URL: https://github.com/apache/arrow-datafusion/pull/1415#discussion_r764524380
##########
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:
I just give an example. The full recommendations are below
```suggestion
if missing_cols.is_empty() {
return Ok(Self::from(LogicalPlan::Sort(Sort {
expr: normalize_cols(exprs, &self.plan)?,
input: Arc::new(self.plan.clone()),
})));
}
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();
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();
Ok(Self::from(LogicalPlan::Projection(Projection {
expr: new_expr,
input: Arc::new(sort_plan),
schema: DFSchemaRef::new(new_schema),
alias: None,
})))
```
--
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
[GitHub] [arrow-datafusion] alamb commented on pull request #1415: Support sort on unprojected columns
Posted by GitBox <gi...@apache.org>.
alamb commented on pull request #1415:
URL: https://github.com/apache/arrow-datafusion/pull/1415#issuecomment-989975725
Thanks @viirya !
--
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
[GitHub] [arrow-datafusion] viirya commented on a change in pull request #1415: Support sort on unprojected columns
Posted by GitBox <gi...@apache.org>.
viirya commented on a change in pull request #1415:
URL: https://github.com/apache/arrow-datafusion/pull/1415#discussion_r764520726
##########
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:
Sure, I can put into another function. It is not too many lines, so I thought it may not worth.
--
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
[GitHub] [arrow-datafusion] viirya commented on pull request #1415: Support sort on unprojected columns
Posted by GitBox <gi...@apache.org>.
viirya commented on pull request #1415:
URL: https://github.com/apache/arrow-datafusion/pull/1415#issuecomment-988464940
Thanks @xudong963
--
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
[GitHub] [arrow-datafusion] viirya commented on a change in pull request #1415: Support sort on unprojected columns
Posted by GitBox <gi...@apache.org>.
viirya commented on a change in pull request #1415:
URL: https://github.com/apache/arrow-datafusion/pull/1415#discussion_r764585083
##########
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:
no, actually I think we should propagate the error. let me update it later.
--
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
[GitHub] [arrow-datafusion] xudong963 commented on a change in pull request #1415: Support sort on unprojected columns
Posted by GitBox <gi...@apache.org>.
xudong963 commented on a change in pull request #1415:
URL: https://github.com/apache/arrow-datafusion/pull/1415#discussion_r764548764
##########
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:
Sorry, I misunderstand
--
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
[GitHub] [arrow-datafusion] viirya commented on a change in pull request #1415: Support sort on unprojected columns
Posted by GitBox <gi...@apache.org>.
viirya commented on a change in pull request #1415:
URL: https://github.com/apache/arrow-datafusion/pull/1415#discussion_r764520425
##########
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:
Not sure I follow your comment correctly. We need to keep original schema unchanged, right?
--
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
[GitHub] [arrow-datafusion] viirya commented on a change in pull request #1415: Support sort on unprojected columns
Posted by GitBox <gi...@apache.org>.
viirya commented on a change in pull request #1415:
URL: https://github.com/apache/arrow-datafusion/pull/1415#discussion_r764521230
##########
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:
? `sort` doesn't just return `Ok()`.
--
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
[GitHub] [arrow-datafusion] houqp commented on pull request #1415: Support sort on unprojected columns
Posted by GitBox <gi...@apache.org>.
houqp commented on pull request #1415:
URL: https://github.com/apache/arrow-datafusion/pull/1415#issuecomment-988539934
Thank you @viirya , very cool to see this in action.
--
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
[GitHub] [arrow-datafusion] xudong963 commented on a change in pull request #1415: Support sort on unprojected columns
Posted by GitBox <gi...@apache.org>.
xudong963 commented on a change in pull request #1415:
URL: https://github.com/apache/arrow-datafusion/pull/1415#discussion_r764593940
##########
File path: datafusion/src/logical_plan/builder.rs
##########
@@ -466,11 +467,112 @@ impl LogicalPlanBuilder {
})))
}
+ /// Add missing sort columns to downstream projection
+ fn add_missing_columns(
+ &self,
+ curr_plan: LogicalPlan,
+ missing_cols: &[Column],
+ ) -> Result<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().try_for_each::<_, Result<()>>(|c| {
+ new_expr.push(normalize_col(Expr::Column(c.clone()), &input)?);
+
+ Ok(())
+ });
+
+ let new_inputs =
+ curr_plan.inputs().into_iter().cloned().collect::<Vec<_>>();
+
+ let new_schema =
+ DFSchema::new(exprlist_to_fields(&new_expr, input_schema)?)?;
+
+ Ok(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 mut new_inputs = vec![];
+ inputs.iter().try_for_each::<_, Result<()>>(|input_plan| {
+ let new_input =
+ self.add_missing_columns((*input_plan).clone(), missing_cols)?;
+ new_inputs.push(new_input);
+
+ Ok(())
+ });
+
+ let expr = curr_plan.expressions();
+ utils::from_plan(&curr_plan, &expr, &new_inputs)
+ }
+ }
+ }
+
/// 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:
👍
--
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
[GitHub] [arrow-datafusion] xudong963 commented on a change in pull request #1415: Support sort on unprojected columns
Posted by GitBox <gi...@apache.org>.
xudong963 commented on a change in pull request #1415:
URL: https://github.com/apache/arrow-datafusion/pull/1415#discussion_r764512897
##########
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.
--
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
[GitHub] [arrow-datafusion] houqp commented on a change in pull request #1415: Support sort on unprojected columns
Posted by GitBox <gi...@apache.org>.
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
[GitHub] [arrow-datafusion] viirya commented on pull request #1415: Support sort on unprojected columns
Posted by GitBox <gi...@apache.org>.
viirya commented on pull request #1415:
URL: https://github.com/apache/arrow-datafusion/pull/1415#issuecomment-990066074
Thank you @houqp @xudong963 @alamb @Dandandan !
--
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
[GitHub] [arrow-datafusion] viirya commented on a change in pull request #1415: Support sort on unprojected columns
Posted by GitBox <gi...@apache.org>.
viirya commented on a change in pull request #1415:
URL: https://github.com/apache/arrow-datafusion/pull/1415#discussion_r764534291
##########
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:
Hmm? I think the first plan has only one `Sort` and the second one has two, no? Why only one `Sort` will sort the table twice?
I also don't get "unselected columns". Currently this pushes missing columns down. These missing columns are not in original schema, we need to remove them, no?
--
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
[GitHub] [arrow-datafusion] xudong963 commented on a change in pull request #1415: Support sort on unprojected columns
Posted by GitBox <gi...@apache.org>.
xudong963 commented on a change in pull request #1415:
URL: https://github.com/apache/arrow-datafusion/pull/1415#discussion_r764548764
##########
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:
Sorry, I misunderstand. I check the code again, I find I'm wrong. Thanks @viirya
--
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
[GitHub] [arrow-datafusion] viirya commented on a change in pull request #1415: Support sort on unprojected columns
Posted by GitBox <gi...@apache.org>.
viirya commented on a change in pull request #1415:
URL: https://github.com/apache/arrow-datafusion/pull/1415#discussion_r764585083
##########
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:
no (update: seems it is because I need to recursively call `add_missing_columns` in `map`, so cannot directly propagate the error. I rewrote it to `try_for_each` then.), actually I think we should propagate the error. let me update it later.
--
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
[GitHub] [arrow-datafusion] viirya commented on pull request #1415: Support sort on unprojected columns
Posted by GitBox <gi...@apache.org>.
viirya commented on pull request #1415:
URL: https://github.com/apache/arrow-datafusion/pull/1415#issuecomment-988548399
@houqp thanks for review. Fixed.
--
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
[GitHub] [arrow-datafusion] viirya commented on a change in pull request #1415: Support sort on unprojected columns
Posted by GitBox <gi...@apache.org>.
viirya commented on a change in pull request #1415:
URL: https://github.com/apache/arrow-datafusion/pull/1415#discussion_r764585083
##########
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:
no, actually I think we should propagate the error. let me update it later.
update: seems it is because I need to recursively call `add_missing_columns` in `map`, so cannot directly propagate the error. I rewrote it to `try_for_each` then.
--
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