peter-toth commented on code in PR #13467:
URL: https://github.com/apache/datafusion/pull/13467#discussion_r1848073920
##########
datafusion/common/src/tree_node.rs:
##########
@@ -769,6 +770,263 @@ impl<T> Transformed<T> {
}
}
+/// [`Container`] contains elements that a function can be applied on or
mapped. The
+/// elements of the container are siblings so the continuation rules are
similar to
+/// [`TreeNodeRecursion::visit_sibling`] / [`Transformed::transform_sibling`].
+pub trait Container<'a, T: 'a>: Sized {
+ fn apply_elements<F: FnMut(&'a T) -> Result<TreeNodeRecursion>>(
+ &'a self,
+ f: F,
+ ) -> Result<TreeNodeRecursion>;
+
+ fn map_elements<F: FnMut(T) -> Result<Transformed<T>>>(
+ self,
+ f: F,
+ ) -> Result<Transformed<Self>>;
+}
+
+impl<'a, T: 'a, C: Container<'a, T>> Container<'a, T> for Box<C> {
+ fn apply_elements<F: FnMut(&'a T) -> Result<TreeNodeRecursion>>(
+ &'a self,
+ f: F,
+ ) -> Result<TreeNodeRecursion> {
+ self.as_ref().apply_elements(f)
+ }
+
+ fn map_elements<F: FnMut(T) -> Result<Transformed<T>>>(
+ self,
+ f: F,
+ ) -> Result<Transformed<Self>> {
+ (*self).map_elements(f)?.map_data(|c| Ok(Self::new(c)))
+ }
+}
+
+impl<'a, T: 'a, C: Container<'a, T> + Clone> Container<'a, T> for Arc<C> {
+ fn apply_elements<F: FnMut(&'a T) -> Result<TreeNodeRecursion>>(
+ &'a self,
+ f: F,
+ ) -> Result<TreeNodeRecursion> {
+ self.as_ref().apply_elements(f)
+ }
+
+ fn map_elements<F: FnMut(T) -> Result<Transformed<T>>>(
+ self,
+ f: F,
+ ) -> Result<Transformed<Self>> {
+ Arc::unwrap_or_clone(self)
+ .map_elements(f)?
+ .map_data(|c| Ok(Arc::new(c)))
+ }
+}
+
+impl<'a, T: 'a, C: Container<'a, T>> Container<'a, T> for Option<C> {
+ fn apply_elements<F: FnMut(&'a T) -> Result<TreeNodeRecursion>>(
+ &'a self,
+ f: F,
+ ) -> Result<TreeNodeRecursion> {
+ match self {
+ Some(t) => t.apply_elements(f),
+ None => Ok(TreeNodeRecursion::Continue),
+ }
+ }
+
+ fn map_elements<F: FnMut(T) -> Result<Transformed<T>>>(
+ self,
+ f: F,
+ ) -> Result<Transformed<Self>> {
+ self.map_or(Ok(Transformed::no(None)), |c| {
+ c.map_elements(f)?.map_data(|c| Ok(Some(c)))
+ })
+ }
+}
+
+impl<'a, T: 'a, C: Container<'a, T>> Container<'a, T> for Vec<C> {
+ fn apply_elements<F: FnMut(&'a T) -> Result<TreeNodeRecursion>>(
+ &'a self,
+ mut f: F,
+ ) -> Result<TreeNodeRecursion> {
+ let mut tnr = TreeNodeRecursion::Continue;
+ for c in self {
+ tnr = c.apply_elements(&mut f)?;
+ match tnr {
+ TreeNodeRecursion::Continue | TreeNodeRecursion::Jump => {}
+ TreeNodeRecursion::Stop => return Ok(TreeNodeRecursion::Stop),
+ }
+ }
+ Ok(tnr)
+ }
+
+ fn map_elements<F: FnMut(T) -> Result<Transformed<T>>>(
+ self,
+ mut f: F,
+ ) -> Result<Transformed<Self>> {
+ let mut tnr = TreeNodeRecursion::Continue;
+ let mut transformed = false;
+ self.into_iter()
+ .map(|c| match tnr {
+ TreeNodeRecursion::Continue | TreeNodeRecursion::Jump => {
+ c.map_elements(&mut f).map(|result| {
+ tnr = result.tnr;
+ transformed |= result.transformed;
+ result.data
+ })
+ }
+ TreeNodeRecursion::Stop => Ok(c),
+ })
+ .collect::<Result<Vec<_>>>()
+ .map(|data| Transformed::new(data, transformed, tnr))
+ }
+}
+
+impl<'a, T: 'a, K: Eq + Hash, C: Container<'a, T>> Container<'a, T> for
HashMap<K, C> {
+ fn apply_elements<F: FnMut(&'a T) -> Result<TreeNodeRecursion>>(
+ &'a self,
+ mut f: F,
+ ) -> Result<TreeNodeRecursion> {
+ let mut tnr = TreeNodeRecursion::Continue;
+ for c in self.values() {
+ tnr = c.apply_elements(&mut f)?;
+ match tnr {
+ TreeNodeRecursion::Continue | TreeNodeRecursion::Jump => {}
+ TreeNodeRecursion::Stop => return Ok(TreeNodeRecursion::Stop),
+ }
+ }
+ Ok(tnr)
+ }
+
+ fn map_elements<F: FnMut(T) -> Result<Transformed<T>>>(
+ self,
+ mut f: F,
+ ) -> Result<Transformed<Self>> {
+ let mut tnr = TreeNodeRecursion::Continue;
+ let mut transformed = false;
+ self.into_iter()
+ .map(|(k, c)| match tnr {
+ TreeNodeRecursion::Continue | TreeNodeRecursion::Jump => {
+ c.map_elements(&mut f).map(|result| {
+ tnr = result.tnr;
+ transformed |= result.transformed;
+ (k, result.data)
+ })
+ }
+ TreeNodeRecursion::Stop => Ok((k, c)),
+ })
+ .collect::<Result<HashMap<_, _>>>()
+ .map(|data| Transformed::new(data, transformed, tnr))
+ }
+}
+
+impl<'a, T: 'a, C0: Container<'a, T>, C1: Container<'a, T>> Container<'a, T>
Review Comment:
I updated the PR description with an example to show how the new container
can be used. I tried to use the complex example of `Expr::Case`. Its fields
contain `Box`, `Option`, `Vec` and 2-tuple as well.
The 3-tuple blanket implementation is not used in any `TreeNodes`, but it
makes possible to get rid of the ugly `map_until_stop_and_collect` macro as we
can put different typed fields into a 3-tupple and just call `apply_elements()`
/ `map_elements()` on it.
##########
datafusion/expr/src/logical_plan/tree_node.rs:
##########
@@ -81,7 +79,7 @@ impl TreeNode for LogicalPlan {
expr,
input,
schema,
- }) => rewrite_arc(input, f)?.update_data(|input| {
+ }) => input.map_elements(f)?.update_data(|input| {
Review Comment:
> IIUC map_children defines how LogicalPlan's construct a tree topology,
while map_elements is resolving the link between nodes.
Yes, the container trait supposes that the elements in it are sibling nodes.
Basically this is why we can use it when we implement `apply_children()` /
`map_children()`. The blanket implementations for the trait makes it very
simple to use with `Expr` and `LogicalPlan` as those logical plan treenodes
define their children as a composition of `Vec`s, `Option`s, tuples, `Box`es
and `Arc`s.
Here in `LogicalPlan` a treenode has mostly only 1 child (except for `Join`
and `Union` maybe) so input (`Arc<LogicalPlan>`) is just a single element
container so this PR doesn't have much benefit, but I think it is worth to do
it to allign `Expr` and `LogicalPlan`.
`TreeNode::map_children` and `TreeNodeContainer:: map_elements` are similar
in a way that both apply `f` to the children/elements of a `TreeNode` /
`TreeNodeContainer`. But we use `TreeNode` to define different kids of logical
and physical plan trees, while `TreeNodeContainer` can be used to simplify
implementation of `apply_children` / `map_children()`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org
For additional commands, e-mail: github-h...@datafusion.apache.org
