berkaysynnada commented on code in PR #13467:
URL: https://github.com/apache/datafusion/pull/13467#discussion_r1847754432
##########
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>>(
Review Comment:
Can we somehow simplify these declarations by aliasing like
`type ApplyFn<'a, T> = dyn FnMut(&'a T) -> Result<TreeNodeRecursion>;`
`type MapFn<T> = dyn FnMut(T) -> Result<Transformed<T>>;`
##########
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 agree with @alamb. We can provide them when they are really necessary --
AFAIK there is not any tuple-based trees in DF.
##########
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:
`map_children` also does not apply the f to input. IIUC `map_children`
defines how LogicalPlan's construct a tree topology, while `map_elements` is
resolving the link between nodes.
--
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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
