polars-plan 0.54.2

Lazy query engine for the Polars DataFrame library
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131

use std::marker::PhantomData;
use std::ops::ControlFlow;

use crate::traversal::edge_provider::NodeEdgesProvider;

pub enum SubtreeVisit {
    Visit,
    Skip,
}

pub trait NodeVisitor {
    type Key;
    type Storage;
    type Edge;
    type BreakValue;

    fn default_edge(
        &mut self,
        key: Self::Key,
        parent_key_and_port: Option<(Self::Key, usize)>,
    ) -> Self::Edge;

    fn pre_visit(
        &mut self,
        key: Self::Key,
        storage: &mut Self::Storage,
        edges: &mut dyn NodeEdgesProvider<Self::Edge>,
    ) -> ControlFlow<Self::BreakValue, SubtreeVisit>;

    /// Indicates this edge is deleted. This edge will not appear when visiting
    /// the node pointed to by this edge. If this edge is the sole edge pointing
    /// to that node, that node will not be visited.
    ///
    /// Called before pre_visit of each node.
    fn is_deleted_edge(&mut self, _edge: &Self::Edge) -> Option<bool> {
        None
    }

    fn post_visit(
        &mut self,
        key: Self::Key,
        storage: &mut Self::Storage,
        edges: &mut dyn NodeEdgesProvider<Self::Edge>,
    ) -> ControlFlow<Self::BreakValue>;
}

pub struct FnVisitors<Key, Storage, Edge, BreakValue, DefaultEdgeFn, PreVisitFn, PostVisitFn>
where
    PreVisitFn: FnMut(
        Key,
        &mut Storage,
        &mut dyn NodeEdgesProvider<Edge>,
    ) -> ControlFlow<BreakValue, SubtreeVisit>,
    PostVisitFn:
        FnMut(Key, &mut Storage, &mut dyn NodeEdgesProvider<Edge>) -> ControlFlow<BreakValue>,
{
    default_edge_fn: DefaultEdgeFn,
    pre_visit_fn: PreVisitFn,
    post_visit_fn: PostVisitFn,
    phantom: PhantomData<(Key, Storage, Edge, BreakValue)>,
}

impl<Key, Storage, Edge, BreakValue, DefaultEdgeFn, PreVisitFn, PostVisitFn>
    FnVisitors<Key, Storage, Edge, BreakValue, DefaultEdgeFn, PreVisitFn, PostVisitFn>
where
    DefaultEdgeFn: FnMut() -> Edge,
    PreVisitFn: FnMut(
        Key,
        &mut Storage,
        &mut dyn NodeEdgesProvider<Edge>,
    ) -> ControlFlow<BreakValue, SubtreeVisit>,
    PostVisitFn:
        FnMut(Key, &mut Storage, &mut dyn NodeEdgesProvider<Edge>) -> ControlFlow<BreakValue>,
{
    pub fn new(
        default_edge_fn: DefaultEdgeFn,
        pre_visit_fn: PreVisitFn,
        post_visit_fn: PostVisitFn,
    ) -> Self {
        Self {
            default_edge_fn,
            pre_visit_fn,
            post_visit_fn,
            phantom: PhantomData,
        }
    }
}

impl<Key, Storage, Edge, BreakValue, DefaultEdgeFn, PreVisitFn, PostVisitFn> NodeVisitor
    for FnVisitors<Key, Storage, Edge, BreakValue, DefaultEdgeFn, PreVisitFn, PostVisitFn>
where
    DefaultEdgeFn: FnMut() -> Edge,
    PreVisitFn: FnMut(
        Key,
        &mut Storage,
        &mut dyn NodeEdgesProvider<Edge>,
    ) -> ControlFlow<BreakValue, SubtreeVisit>,
    PostVisitFn:
        FnMut(Key, &mut Storage, &mut dyn NodeEdgesProvider<Edge>) -> ControlFlow<BreakValue>,
{
    type Key = Key;
    type Storage = Storage;
    type Edge = Edge;
    type BreakValue = BreakValue;

    fn default_edge(
        &mut self,
        _key: Self::Key,
        _parent_key_and_port: Option<(Self::Key, usize)>,
    ) -> Self::Edge {
        (self.default_edge_fn)()
    }

    fn pre_visit(
        &mut self,
        key: Self::Key,
        storage: &mut Self::Storage,
        edges: &mut dyn NodeEdgesProvider<Self::Edge>,
    ) -> ControlFlow<Self::BreakValue, SubtreeVisit> {
        (self.pre_visit_fn)(key, storage, edges)
    }

    fn post_visit(
        &mut self,
        key: Self::Key,
        storage: &mut Self::Storage,
        edges: &mut dyn NodeEdgesProvider<Self::Edge>,
    ) -> ControlFlow<Self::BreakValue> {
        (self.post_visit_fn)(key, storage, edges)
    }
}