Coverage Report

Created: 2024-10-13 08:39

/Users/andrewlamb/Software/datafusion/datafusion/physical-expr-common/src/tree_node.rs
Line
Count
Source (jump to first uncovered line)
1
// Licensed to the Apache Software Foundation (ASF) under one
2
// or more contributor license agreements.  See the NOTICE file
3
// distributed with this work for additional information
4
// regarding copyright ownership.  The ASF licenses this file
5
// to you under the Apache License, Version 2.0 (the
6
// "License"); you may not use this file except in compliance
7
// with the License.  You may obtain a copy of the License at
8
//
9
//   http://www.apache.org/licenses/LICENSE-2.0
10
//
11
// Unless required by applicable law or agreed to in writing,
12
// software distributed under the License is distributed on an
13
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14
// KIND, either express or implied.  See the License for the
15
// specific language governing permissions and limitations
16
// under the License.
17
18
//! This module provides common traits for visiting or rewriting tree nodes easily.
19
20
use std::fmt::{self, Display, Formatter};
21
use std::sync::Arc;
22
23
use crate::physical_expr::{with_new_children_if_necessary, PhysicalExpr};
24
25
use datafusion_common::tree_node::{ConcreteTreeNode, DynTreeNode};
26
use datafusion_common::Result;
27
28
impl DynTreeNode for dyn PhysicalExpr {
29
14.2k
    fn arc_children(&self) -> Vec<&Arc<Self>> {
30
14.2k
        self.children()
31
14.2k
    }
32
33
66
    fn with_new_arc_children(
34
66
        &self,
35
66
        arc_self: Arc<Self>,
36
66
        new_children: Vec<Arc<Self>>,
37
66
    ) -> Result<Arc<Self>> {
38
66
        with_new_children_if_necessary(arc_self, new_children)
39
66
    }
40
}
41
42
/// A node object encapsulating a [`PhysicalExpr`] node with a payload. Since there are
43
/// two ways to access child plans—directly from the plan  and through child nodes—it's
44
/// recommended to perform mutable operations via [`Self::update_expr_from_children`].
45
#[derive(Debug)]
46
pub struct ExprContext<T: Sized> {
47
    /// The physical expression associated with this context.
48
    pub expr: Arc<dyn PhysicalExpr>,
49
    /// Custom data payload of the node.
50
    pub data: T,
51
    /// Child contexts of this node.
52
    pub children: Vec<Self>,
53
}
54
55
impl<T> ExprContext<T> {
56
17.0k
    pub fn new(expr: Arc<dyn PhysicalExpr>, data: T, children: Vec<Self>) -> Self {
57
17.0k
        Self {
58
17.0k
            expr,
59
17.0k
            data,
60
17.0k
            children,
61
17.0k
        }
62
17.0k
    }
63
64
8.10k
    pub fn update_expr_from_children(mut self) -> Result<Self> {
65
15.9k
        let children_expr = self.children.iter().map(|c| c.expr.clone()
).collect()8.10k
;
66
8.10k
        self.expr = with_new_children_if_necessary(self.expr, children_expr)
?0
;
67
8.10k
        Ok(self)
68
8.10k
    }
69
}
70
71
impl<T: Default> ExprContext<T> {
72
17.0k
    pub fn new_default(plan: Arc<dyn PhysicalExpr>) -> Self {
73
17.0k
        let children = plan
74
17.0k
            .children()
75
17.0k
            .into_iter()
76
17.0k
            .cloned()
77
17.0k
            .map(Self::new_default)
78
17.0k
            .collect();
79
17.0k
        Self::new(plan, Default::default(), children)
80
17.0k
    }
81
}
82
83
impl<T: Display> Display for ExprContext<T> {
84
0
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
85
0
        write!(f, "expr: {:?}", self.expr)?;
86
0
        write!(f, "data:{}", self.data)?;
87
0
        write!(f, "")
88
0
    }
89
}
90
91
impl<T> ConcreteTreeNode for ExprContext<T> {
92
0
    fn children(&self) -> &[Self] {
93
0
        &self.children
94
0
    }
95
96
17.5k
    fn take_children(mut self) -> (Self, Vec<Self>) {
97
17.5k
        let children = std::mem::take(&mut self.children);
98
17.5k
        (self, children)
99
17.5k
    }
100
101
8.10k
    fn with_new_children(mut self, children: Vec<Self>) -> Result<Self> {
102
8.10k
        self.children = children;
103
8.10k
        self.update_expr_from_children()
104
8.10k
    }
105
}