peter-toth commented on code in PR #9780: URL: https://github.com/apache/arrow-datafusion/pull/9780#discussion_r1547749571
########## datafusion/expr/src/logical_plan/mutate.rs: ########## @@ -0,0 +1,346 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use super::plan::*; +use crate::expr::{Exists, InSubquery}; +use crate::{Expr, UserDefinedLogicalNode}; +use datafusion_common::tree_node::Transformed; +use datafusion_common::{internal_err, Result}; +use datafusion_common::{Column, DFSchema, DFSchemaRef}; +use std::sync::{Arc, OnceLock}; + +impl LogicalPlan { + /// applies `f` to each expression of this node, potentially rewriting it in + /// place + /// + /// If `f` returns an error, the error is returned and the expressions are + /// left in a partially modified state + pub fn rewrite_exprs<F>(&mut self, mut f: F) -> Result<Transformed<()>> + where + F: FnMut(&mut Expr) -> Result<Transformed<()>>, + { + match self { + LogicalPlan::Projection(Projection { expr, .. }) => { + rewrite_expr_iter_mut(expr.iter_mut(), f) + } + LogicalPlan::Values(Values { values, .. }) => { + rewrite_expr_iter_mut(values.iter_mut().flatten(), f) + } + LogicalPlan::Filter(Filter { predicate, .. }) => f(predicate), + LogicalPlan::Repartition(Repartition { + partitioning_scheme, + .. + }) => match partitioning_scheme { + Partitioning::Hash(expr, _) => rewrite_expr_iter_mut(expr.iter_mut(), f), + Partitioning::DistributeBy(expr) => { + rewrite_expr_iter_mut(expr.iter_mut(), f) + } + Partitioning::RoundRobinBatch(_) => Ok(Transformed::no(())), + }, + LogicalPlan::Window(Window { window_expr, .. }) => { + rewrite_expr_iter_mut(window_expr.iter_mut(), f) + } + LogicalPlan::Aggregate(Aggregate { + group_expr, + aggr_expr, + .. + }) => { + let exprs = group_expr.iter_mut().chain(aggr_expr.iter_mut()); + rewrite_expr_iter_mut(exprs, f) + } + // There are two part of expression for join, equijoin(on) and non-equijoin(filter). + // 1. the first part is `on.len()` equijoin expressions, and the struct of each expr is `left-on = right-on`. + // 2. the second part is non-equijoin(filter). + LogicalPlan::Join(Join { on, filter, .. }) => { + let exprs = on + .iter_mut() + .flat_map(|(e1, e2)| std::iter::once(e1).chain(std::iter::once(e2))); + + let result = rewrite_expr_iter_mut(exprs, &mut f)?; + + if let Some(filter) = filter.as_mut() { + result.and_then(|| f(filter)) + } else { + Ok(result) + } + } + LogicalPlan::Sort(Sort { expr, .. }) => { + rewrite_expr_iter_mut(expr.iter_mut(), f) + } + LogicalPlan::Extension(extension) => { + rewrite_extension_exprs(&mut extension.node, f) + } + LogicalPlan::TableScan(TableScan { filters, .. }) => { + rewrite_expr_iter_mut(filters.iter_mut(), f) + } + LogicalPlan::Unnest(Unnest { column, .. }) => rewrite_column(column, f), + LogicalPlan::Distinct(Distinct::On(DistinctOn { + on_expr, + select_expr, + sort_expr, + .. + })) => { + let exprs = on_expr + .iter_mut() + .chain(select_expr.iter_mut()) + .chain(sort_expr.iter_mut().flat_map(|x| x.iter_mut())); + + rewrite_expr_iter_mut(exprs, f) + } + // plans without expressions + LogicalPlan::EmptyRelation(_) + | LogicalPlan::RecursiveQuery(_) + | LogicalPlan::Subquery(_) + | LogicalPlan::SubqueryAlias(_) + | LogicalPlan::Limit(_) + | LogicalPlan::Statement(_) + | LogicalPlan::CrossJoin(_) + | LogicalPlan::Analyze(_) + | LogicalPlan::Explain(_) + | LogicalPlan::Union(_) + | LogicalPlan::Distinct(Distinct::All(_)) + | LogicalPlan::Dml(_) + | LogicalPlan::Ddl(_) + | LogicalPlan::Copy(_) + | LogicalPlan::DescribeTable(_) + | LogicalPlan::Prepare(_) => Ok(Transformed::no(())), + } + } + + /// applies `f` to each input of this node, rewriting them in place. + /// + /// # Notes + /// Inputs include both direct children as well as any embedded subquery + /// `LogicalPlan`s, for example such as are in [`Expr::Exists`]. + /// + /// If `f` returns an `Err`, that Err is returned, and the inputs are left + /// in a partially modified state + pub fn rewrite_inputs<F>(&mut self, mut f: F) -> Result<Transformed<()>> + where + F: FnMut(&mut LogicalPlan) -> Result<Transformed<()>>, + { + let children_result = match self { + LogicalPlan::Projection(Projection { input, .. }) => { + rewrite_arc(input, &mut f) + } + LogicalPlan::Filter(Filter { input, .. }) => rewrite_arc(input, &mut f), + LogicalPlan::Repartition(Repartition { input, .. }) => { + rewrite_arc(input, &mut f) + } + LogicalPlan::Window(Window { input, .. }) => rewrite_arc(input, &mut f), + LogicalPlan::Aggregate(Aggregate { input, .. }) => rewrite_arc(input, &mut f), + LogicalPlan::Sort(Sort { input, .. }) => rewrite_arc(input, &mut f), + LogicalPlan::Join(Join { left, right, .. }) => { + rewrite_arc(left, &mut f)?.and_then(|| rewrite_arc(right, &mut f)) + } + LogicalPlan::CrossJoin(CrossJoin { left, right, .. }) => { + rewrite_arc(left, &mut f)?.and_then(|| rewrite_arc(right, &mut f)) + } + LogicalPlan::Limit(Limit { input, .. }) => rewrite_arc(input, &mut f), + LogicalPlan::Subquery(Subquery { subquery, .. }) => { + rewrite_arc(subquery, &mut f) + } + LogicalPlan::SubqueryAlias(SubqueryAlias { input, .. }) => { + rewrite_arc(input, &mut f) + } + LogicalPlan::Extension(extension) => { + rewrite_extension_inputs(&mut extension.node, &mut f) + } + LogicalPlan::Union(Union { inputs, .. }) => inputs + .iter_mut() + .try_fold(Transformed::no(()), |acc, input| { + acc.and_then(|| rewrite_arc(input, &mut f)) + }), + LogicalPlan::Distinct( + Distinct::All(input) | Distinct::On(DistinctOn { input, .. }), + ) => rewrite_arc(input, &mut f), + LogicalPlan::Explain(explain) => rewrite_arc(&mut explain.plan, &mut f), + LogicalPlan::Analyze(analyze) => rewrite_arc(&mut analyze.input, &mut f), + LogicalPlan::Dml(write) => rewrite_arc(&mut write.input, &mut f), + LogicalPlan::Copy(copy) => rewrite_arc(&mut copy.input, &mut f), + LogicalPlan::Ddl(ddl) => { + if let Some(input) = ddl.input_mut() { + rewrite_arc(input, &mut f) + } else { + Ok(Transformed::no(())) + } + } + LogicalPlan::Unnest(Unnest { input, .. }) => rewrite_arc(input, &mut f), + LogicalPlan::Prepare(Prepare { input, .. }) => rewrite_arc(input, &mut f), + LogicalPlan::RecursiveQuery(RecursiveQuery { + static_term, + recursive_term, + .. + }) => rewrite_arc(static_term, &mut f)? + .and_then(|| rewrite_arc(recursive_term, &mut f)), + // plans without inputs + LogicalPlan::TableScan { .. } + | LogicalPlan::Statement { .. } + | LogicalPlan::EmptyRelation { .. } + | LogicalPlan::Values { .. } + | LogicalPlan::DescribeTable(_) => Ok(Transformed::no(())), + }?; + + // after visiting the actual children we we need to visit any subqueries + // that are inside the expressions + children_result.and_then(|| self.rewrite_subqueries(&mut f)) + } + + /// applies `f` to LogicalPlans in any subquery expressions + /// + /// If Err is returned, the plan may be left in a partially modified state + fn rewrite_subqueries<F>(&mut self, mut f: F) -> Result<Transformed<()>> Review Comment: I'm a bit concerned about this part. I mean the `TreeNode` APIs are currently inconsistent about handling "inner" trees. I'm refferring to `LogicalPlan` for example. It overrides `TreeNode::apply()` and `TreeNode::visit()` to apply `f` on subquery's logical plans too. - But first of all, none of the other trees do that, - secodly, `LogicalPlan` doesn't override `rewrite()` or `transform...()` so transformations are applied only the "current" tree and not on "whole" logical plan. I'm not sure that the current state is right or wrong as I can imagine both usescases in which - a transformations needs to be applied on the current tree only - a transformations need to change the whole tree (subplans trees included) Now, this code seems to apply `f` on subquery logical plans too. If we keep it that way `mutate()` will be inconsistent with other APIs. But I'm also not sure that why it applies `f` on the root node only? Shouldn't it call `subquery.subquery.mutate(f)` to start a traversal on the subquery plan tree? Is it needed because of the `Optimizer`? Does it optimize the whole logical plan at once? If so, I think we should introduce new APIs on `LogicalPlan` that is different to the existing `TreeNode` APIs. In Spark for example we have `TreeNode.transform()` and `QueryPlan.transformWithSubqueries()` for the different purposes. -- 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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
