peter-toth commented on code in PR #9946: URL: https://github.com/apache/arrow-datafusion/pull/9946#discussion_r1555045657
########## datafusion/expr/src/logical_plan/rewrite.rs: ########## @@ -0,0 +1,228 @@ +// 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. + +//! Methods for rewriting logical plans + +use crate::{ + Aggregate, CrossJoin, Distinct, DistinctOn, EmptyRelation, Filter, Join, Limit, + LogicalPlan, Prepare, Projection, RecursiveQuery, Repartition, Sort, Subquery, + SubqueryAlias, Union, Unnest, UserDefinedLogicalNode, Window, +}; +use datafusion_common::tree_node::{Transformed, TreeNodeIterator}; +use datafusion_common::{DFSchema, DFSchemaRef, Result}; +use std::sync::{Arc, OnceLock}; + +/// A temporary node that is left in place while rewriting the children of a +/// [`LogicalPlan`]. This is necessary to ensure that the `LogicalPlan` is +/// always in a valid state (from the Rust perspective) +static PLACEHOLDER: OnceLock<Arc<LogicalPlan>> = OnceLock::new(); + +/// its inputs, so this code would not be needed. However, for now we try and +/// unwrap the `Arc` which avoids `clone`ing in most cases. +/// +/// On error, node be left with a placeholder logical plan +fn rewrite_arc<F>( + node: &mut Arc<LogicalPlan>, + mut f: F, +) -> datafusion_common::Result<Transformed<&mut Arc<LogicalPlan>>> +where + F: FnMut(LogicalPlan) -> Result<Transformed<LogicalPlan>>, +{ + // We need to leave a valid node in the Arc, while we rewrite the existing + // one, so use a single global static placeholder node + let mut new_node = PLACEHOLDER + .get_or_init(|| { + Arc::new(LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: DFSchemaRef::new(DFSchema::empty()), + })) + }) + .clone(); + + // take the old value out of the Arc + std::mem::swap(node, &mut new_node); + + // try to update existing node, if it isn't shared with others + let new_node = Arc::try_unwrap(new_node) + // if None is returned, there is another reference to this + // LogicalPlan, so we must clone instead + .unwrap_or_else(|node| node.as_ref().clone()); + + // apply the actual transform + let result = f(new_node)?; + + // put the new value back into the Arc + let mut new_node = Arc::new(result.data); + std::mem::swap(node, &mut new_node); + + // return the `node` back + Ok(Transformed::new(node, result.transformed, result.tnr)) +} + +/// Rewrite the arc and discard the contents of Transformed +fn rewrite_arc_no_data<F>( + node: &mut Arc<LogicalPlan>, + f: F, +) -> datafusion_common::Result<Transformed<()>> +where + F: FnMut(LogicalPlan) -> Result<Transformed<LogicalPlan>>, +{ + rewrite_arc(node, f).map(|res| res.discard_data()) +} + +/// Rewrites all inputs for an Extension node "in place" +/// (it currently has to copy values because there are no APIs for in place modification) +/// +/// Should be removed when we have an API for in place modifications of the +/// extension to avoid these copies +fn rewrite_extension_inputs<F>( + node: &mut Arc<dyn UserDefinedLogicalNode>, + f: F, +) -> datafusion_common::Result<Transformed<()>> +where + F: FnMut(LogicalPlan) -> Result<Transformed<LogicalPlan>>, +{ + let Transformed { + data: new_inputs, + transformed, + tnr, + } = node + .inputs() + .into_iter() + .cloned() + .map_until_stop_and_collect(f)?; + + let exprs = node.expressions(); + let mut new_node = node.from_template(&exprs, &new_inputs); + std::mem::swap(node, &mut new_node); + Ok(Transformed { + data: (), + transformed, + tnr, + }) +} + +impl LogicalPlan { + /// Applies `f` to each child (input) of this plan node, rewriting them *in place.* + /// + /// Note that this function returns `Transformed<()>` because it does not + /// consume `self`, but instead modifies it in place. However, `F` transforms + /// the children by ownership + /// + /// # Notes + /// + /// Inputs include ONLY direct children, not embedded subquery + /// `LogicalPlan`s, for example such as are in [`Expr::Exists`]. + /// + /// [`Expr::Exists`]: crate::expr::Expr::Exists + pub(crate) fn rewrite_children<F>(&mut self, mut f: F) -> Result<Transformed<()>> Review Comment: Actually, just one thing came into my mind: as far as I see in `rewrite_arc()` you need an owned `Arc<LogicalPlan>` to call `Arc::try_unwrap`. But, you only have `&mut Arc<LogicalPlan>` and that's why you need `std::mem::swap` 2 times with that `PLACEHOLDER`. If your `rewrite_children()`/`map_children()` worked with owned `Arc<LogicalPlan>` and not with `&mut Arc<LogicalPlan>` the swap wouldn't be needed. But in that case the implementation would be a bit complex like: ``` fn map_children<F: FnMut(Self) -> Result<Transformed<Self>>>( self, f: F, ) -> Result<Transformed<Self>> { Ok(match self { LogicalPlan::Projection(Projection { expr, input, schema }) => { rewrite_arc(input, f)?.update_data(|input| LogicalPlan::Projection(Projection { expr, input, schema })) } LogicalPlan::Filter(Filter { predicate, input }) => { rewrite_arc(input, f)?.update_data(|input| LogicalPlan::Filter(Filter { predicate, input })) } ... }) } ``` Also `discard_data()` won't be required. BTW, this is how `Expr::map_children()` is implemented but there are `Box`es so the `transform_box()` implementation is simpler than this `rewrite_arc()` is. -- 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]
