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` as 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.
--
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
