linkrs_query/planning/plan/
execution_plan.rs

1//! Structure definition of the execution plan
2//! Contains the ExecutionPlan and SubPlan structures.
3
4use std::collections::HashMap;
5
6use crate::planning::plan::core::nodes::base::plan_node_traits::{
7    MultipleInputNode, SingleInputNode,
8};
9use crate::planning::plan::logical_plan::LogicalPlan;
10use crate::planning::plan::PlanNodeEnum;
11use crate::planning::planner::PlannerError;
12
13// The `PartitionSpec` / `PartitionSource` / `PartitionSpecError` types now
14// live in `partition_spec.rs`.  Re-export them here so the public
15// `ExecutionPlan` API (which carries an optional `PartitionSpec`) keeps
16// compiling without touching the new module path at every call site.
17pub use crate::planning::plan::partition_spec::{
18    PartitionSource, PartitionSpec, PartitionSpecError,
19};
20
21/// Execution plan structure
22/// Represents the complete executable plan, including the root node and the plan ID.
23#[derive(Debug, Clone)]
24pub struct ExecutionPlan {
25    /// The root node of the planning tree
26    pub root: Option<PlanNodeEnum>,
27
28    /// The unique ID of the plan
29    pub id: i64,
30
31    /// Optimized time (in microseconds)
32    pub optimize_time_in_us: u64,
33
34    /// Of course! Please provide the text you would like to have translated.
35    pub format: String,
36
37    /// Optional physical partition layout. A layout is consumed only by a
38    /// partition-safe streaming plan; unsupported plans fail explicitly.
39    pub partition_spec: Option<PartitionSpec>,
40
41    /// Maximum worker threads for intra-query parallelism.
42    /// 1 means fully serial. Propagated from PartitioningConfig.
43    pub max_workers: usize,
44
45    /// Per-partition output channel capacity for backpressure.
46    pub max_buffered_chunks: usize,
47
48    /// Why parallel partitioning was not selected (empty = partitioning
49    /// active or not requested). Surfaced in EXPLAIN / PROFILE diagnostics.
50    pub parallel_fallback_reason: String,
51
52    /// Cost-based decision notes produced during optimization (e.g. subquery
53    /// unnesting and join-order rewrites). Surfaced in EXPLAIN diagnostics.
54    pub cbo_notes: Vec<String>,
55
56    /// Per-logical-node output row estimates produced by the cost-based
57    /// phase, keyed by logical node id. Written into physical operator
58    /// specs (`estimated_cardinality`) by the `estimated_rows` metadata pass.
59    pub row_estimates: HashMap<i64, u64>,
60
61    /// Cost-based join algorithm decisions keyed by planner node id.
62    ///
63    /// Produced by the join-order rewriter when it reorders a join chain;
64    /// consumed by the arena builder to pick `HashJoin` vs nested-loop
65    /// execution for each `InnerJoin`/`LeftJoin` node. Absent keys fall
66    /// back to the default heuristic (hash join for valid equi keys).
67    pub join_algorithms: HashMap<i64, crate::optimizer::JoinAlgorithm>,
68
69    /// The pure logical plan (if conversion succeeded).
70    /// Used by cost-based optimization to make physical decisions.
71    pub logical_plan: Option<LogicalPlan>,
72}
73
74impl ExecutionPlan {
75    /// Create a new execution plan.
76    pub fn new(root: Option<PlanNodeEnum>) -> Self {
77        Self {
78            root,
79            id: -1,
80            optimize_time_in_us: 0,
81            format: "default".to_string(),
82            partition_spec: None,
83            max_workers: 1,
84            max_buffered_chunks: 10,
85            parallel_fallback_reason: String::new(),
86            cbo_notes: Vec::new(),
87            row_estimates: HashMap::new(),
88            join_algorithms: HashMap::new(),
89            logical_plan: None,
90        }
91    }
92
93    /// Set the root node of the plan.
94    pub fn set_root(&mut self, root: PlanNodeEnum) {
95        self.root = Some(root);
96    }
97
98    /// Obtain the reference to the root node of the plan.
99    pub fn root(&self) -> &Option<PlanNodeEnum> {
100        &self.root
101    }
102
103    /// Obtain a reference to the variable root node.
104    pub fn root_mut(&mut self) -> &mut Option<PlanNodeEnum> {
105        &mut self.root
106    }
107
108    /// Set the ID for the plan.
109    pub fn set_id(&mut self, id: i64) {
110        self.id = id;
111    }
112
113    /// Set the optimization time
114    pub fn set_optimize_time(&mut self, time_us: u64) {
115        self.optimize_time_in_us = time_us;
116    }
117
118    /// Set the output format
119    pub fn set_format(&mut self, format: String) {
120        self.format = format;
121    }
122
123    /// Attach the physical partition layout chosen by the optimizer.
124    pub fn set_partition_spec(&mut self, partition_spec: PartitionSpec) {
125        self.partition_spec = Some(partition_spec);
126    }
127
128    /// Remove partitioning and execute the plan as a single tree.
129    pub fn clear_partition_spec(&mut self) {
130        self.partition_spec = None;
131    }
132
133    pub fn partition_spec(&self) -> Option<&PartitionSpec> {
134        self.partition_spec.as_ref()
135    }
136
137    pub fn set_max_workers(&mut self, max_workers: usize) {
138        self.max_workers = max_workers;
139    }
140
141    pub fn set_max_buffered_chunks(&mut self, max_buffered_chunks: usize) {
142        self.max_buffered_chunks = max_buffered_chunks;
143    }
144
145    /// Attach a pure logical plan for cost-based optimization.
146    pub fn set_logical_plan(&mut self, logical_plan: LogicalPlan) {
147        self.logical_plan = Some(logical_plan);
148    }
149
150    /// Access the logical plan, if available.
151    pub fn logical_plan(&self) -> Option<&LogicalPlan> {
152        self.logical_plan.as_ref()
153    }
154
155    /// Calculate the number of nodes in the plan.
156    /// Traverses the execution plan tree with an explicit stack instead of
157    /// recursion so deeply nested plans cannot overflow the call stack.
158    pub fn node_count(&self) -> usize {
159        let Some(root) = self.root.as_ref() else {
160            return 0;
161        };
162        let mut count = 0;
163        let mut stack = vec![root];
164        while let Some(node) = stack.pop() {
165            count += 1;
166            stack.extend(node.children());
167        }
168        count
169    }
170}
171
172/// SubPlan structure
173/// Represents a sub-part of the execution plan, which contains the root node and the tail node.
174/// Segmented planning for complex queries
175#[derive(Debug, Clone)]
176pub struct SubPlan {
177    /// The root node of the sub-plan
178    pub root: Option<PlanNodeEnum>,
179
180    /// The end node of the sub-plan
181    /// Used to connect multiple sub-plans
182    pub tail: Option<PlanNodeEnum>,
183
184    /// Native logical tree, when the planner produced one directly.
185    ///
186    /// Migrated planners build the pure [`LogicalNodeEnum`] tree first and
187    /// convert it to the physical root exactly once at the plan exit; the
188    /// logical tree is attached here so the compiler can build the
189    /// [`crate::planning::plan::logical_plan::LogicalPlan`] natively
190    /// instead of stripping it back out of the physical tree.
191    pub logical_root: Option<crate::planning::plan::logical::LogicalNodeEnum>,
192}
193
194impl SubPlan {
195    /// Create a new SubPlan without a logical mirror.
196    ///
197    /// Reserved for planners whose operators stay physical-only by design
198    /// (data modifications, administrative statements, search operators)
199    /// where no `LogicalNodeEnum` counterpart exists. DQL planners build a
200    /// native logical tree instead and attach it via the `SubPlan` struct
201    /// literal or the `plan_combiner` mirror helpers.
202    pub fn new(root: Option<PlanNodeEnum>, tail: Option<PlanNodeEnum>) -> Self {
203        Self {
204            root,
205            tail,
206            logical_root: None,
207        }
208    }
209
210    /// Create a SubPlan from a native logical tree.
211    ///
212    /// Converts the logical root to the physical root exactly once
213    /// (via `convert_logical_to_physical`) and retains the logical tree for
214    /// the compiler.
215    pub fn from_logical_root(
216        logical_root: crate::planning::plan::logical::LogicalNodeEnum,
217    ) -> Self {
218        let physical_root =
219            crate::planning::physical_planner::convert_logical_to_physical(logical_root.clone());
220        Self {
221            root: Some(physical_root),
222            tail: None,
223            logical_root: Some(logical_root),
224        }
225    }
226
227    /// Obtain the native logical tree, if the planner produced one.
228    pub fn logical_root(&self) -> Option<&crate::planning::plan::logical::LogicalNodeEnum> {
229        self.logical_root.as_ref()
230    }
231
232    /// Create a SubPlan that contains only the root node.
233    ///
234    /// Like [`SubPlan::new`], the plan carries no logical mirror and is
235    /// reserved for physical-only operators and transient physical wiring
236    /// seeds.
237    pub fn from_root(root: PlanNodeEnum) -> Self {
238        Self {
239            root: Some(root.clone()),
240            tail: Some(root),
241            logical_root: None,
242        }
243    }
244
245    /// Create a SubPlan that contains only a single node.
246    ///
247    /// Like [`SubPlan::new`], the plan carries no logical mirror and is
248    /// reserved for physical-only operators and transient physical wiring
249    /// seeds.
250    pub fn from_single_node(node: PlanNodeEnum) -> Self {
251        Self {
252            root: Some(node.clone()),
253            tail: Some(node),
254            logical_root: None,
255        }
256    }
257
258    /// Obtain a reference to the root node.
259    pub fn root(&self) -> &Option<PlanNodeEnum> {
260        &self.root
261    }
262
263    /// Obtain the reference to the tail node.
264    pub fn tail(&self) -> &Option<PlanNodeEnum> {
265        &self.tail
266    }
267
268    /// Setting the root node
269    pub fn set_root(&mut self, root: PlanNodeEnum) {
270        self.root = Some(root);
271    }
272
273    /// Setting the tail node
274    pub fn set_tail(&mut self, tail: PlanNodeEnum) {
275        self.tail = Some(tail);
276    }
277
278    /// Check whether SubPlan is empty.
279    pub fn is_empty(&self) -> bool {
280        self.root.is_none()
281    }
282
283    /// Retrieve all nodes from the SubPlan.
284    pub fn collect_nodes(&self) -> Vec<PlanNodeEnum> {
285        let mut nodes = Vec::new();
286
287        if let Some(root) = &self.root {
288            nodes.push(root.clone());
289        }
290
291        if let Some(tail) = &self.tail {
292            nodes.push(tail.clone());
293        }
294
295        nodes
296    }
297
298    /// Connect `upstream` as the input of `downstream`, returning a
299    /// structurally closed plan.
300    ///
301    /// The upstream root is attached as the downstream root's input, so the
302    /// resulting SubPlan is a complete tree that planners and executors
303    /// observe identically.  This replaces the previous pattern of creating
304    /// the downstream node first and manually injecting its input, which
305    /// left the connection outside the SubPlan structure.
306    ///
307    /// The returned plan's tail is the upstream tail (the data flows from
308    /// upstream into downstream).
309    pub fn connect_upstream(
310        mut downstream: SubPlan,
311        upstream: SubPlan,
312    ) -> Result<SubPlan, PlannerError> {
313        let down_root = downstream.root.take().ok_or_else(|| {
314            PlannerError::PlanGenerationFailed("downstream has no root".to_string())
315        })?;
316        let up_root = upstream.root.ok_or_else(|| {
317            PlannerError::PlanGenerationFailed("upstream has no root".to_string())
318        })?;
319
320        let mut connected = down_root;
321        connect_node_input(&mut connected, up_root)?;
322        downstream.root = Some(connected);
323        downstream.tail = upstream.tail;
324        Ok(downstream)
325    }
326}
327
328/// Attach `input` as the structural input of `node`.
329///
330/// Single-input nodes get their input slot and dependency list filled via
331/// [`SingleInputNode::set_input`]; multi-input expansion nodes (Expand /
332/// ExpandAll / AppendVertices) append to their dependency list via
333/// [`MultipleInputNode::add_input`], which is the storage their tree
334/// traversals read.  Binary and leaf nodes cannot host a unary upstream.
335fn connect_node_input(node: &mut PlanNodeEnum, input: PlanNodeEnum) -> Result<(), PlannerError> {
336    match node {
337        PlanNodeEnum::Project(n) => n.set_input(input),
338        PlanNodeEnum::Filter(n) => n.set_input(input),
339        PlanNodeEnum::Sort(n) => n.set_input(input),
340        PlanNodeEnum::Limit(n) => n.set_input(input),
341        PlanNodeEnum::TopN(n) => n.set_input(input),
342        PlanNodeEnum::Sample(n) => n.set_input(input),
343        PlanNodeEnum::Dedup(n) => n.set_input(input),
344        PlanNodeEnum::DataCollect(n) => n.set_input(input),
345        PlanNodeEnum::Aggregate(n) => n.set_input(input),
346        PlanNodeEnum::Window(n) => n.set_input(input),
347        PlanNodeEnum::Unwind(n) => n.set_input(input),
348        PlanNodeEnum::Assign(n) => n.set_input(input),
349        PlanNodeEnum::PatternApply(n) => n.set_input(input),
350        PlanNodeEnum::CorrelatedApply(n) => n.set_input(input),
351        PlanNodeEnum::RollUpApply(n) => n.set_input(input),
352        PlanNodeEnum::Materialize(n) => n.set_input(input),
353        PlanNodeEnum::Traverse(n) => n.set_input(input),
354        PlanNodeEnum::PipeDeleteVertices(n) => n.set_input(input),
355        PlanNodeEnum::PipeDeleteEdges(n) => n.set_input(input),
356        PlanNodeEnum::Union(n) => n.set_input(input),
357        PlanNodeEnum::Minus(n) => n.set_input(input),
358        PlanNodeEnum::Intersect(n) => n.set_input(input),
359        PlanNodeEnum::Expand(n) => n.add_input(input),
360        PlanNodeEnum::ExpandAll(n) => n.add_input(input),
361        PlanNodeEnum::AppendVertices(n) => n.add_input(input),
362        _ => {
363            return Err(PlannerError::PlanGenerationFailed(format!(
364                "Cannot connect upstream to node of type {}",
365                node.name()
366            )));
367        }
368    }
369    Ok(())
370}
371
372#[cfg(test)]
373mod tests {
374    use super::*;
375
376    #[test]
377    fn execution_plan_max_workers_defaults() {
378        let plan = ExecutionPlan::new(None);
379        assert_eq!(plan.max_workers, 1);
380        assert_eq!(plan.max_buffered_chunks, 10);
381    }
382
383    #[test]
384    fn execution_plan_set_max_workers() {
385        let mut plan = ExecutionPlan::new(None);
386        plan.set_max_workers(4);
387        plan.set_max_buffered_chunks(20);
388        assert_eq!(plan.max_workers, 4);
389        assert_eq!(plan.max_buffered_chunks, 20);
390    }
391
392    #[test]
393    fn connect_upstream_attaches_input_structurally() {
394        use crate::planning::plan::core::nodes::{ExpandAllNode, StartNode};
395
396        let downstream = SubPlan::from_single_node(PlanNodeEnum::ExpandAll(ExpandAllNode::new(
397            1,
398            vec![],
399            "out",
400        )));
401        let upstream = SubPlan::from_single_node(PlanNodeEnum::Start(StartNode::new()));
402
403        let connected = SubPlan::connect_upstream(downstream, upstream)
404            .expect("connect_upstream should succeed");
405        let root = connected.root.expect("connected plan should have a root");
406        let children = root.children();
407        assert_eq!(children.len(), 1, "expand node must own exactly one input");
408        assert!(
409            matches!(children[0], PlanNodeEnum::Start(_)),
410            "expand node input should be the upstream root"
411        );
412        assert!(
413            connected.tail.is_some(),
414            "connected plan tail should carry the upstream tail"
415        );
416    }
417
418    #[test]
419    fn connect_upstream_requires_both_roots() {
420        use crate::planning::plan::core::nodes::StartNode;
421
422        let empty = SubPlan::new(None, None);
423        let populated = SubPlan::from_single_node(PlanNodeEnum::Start(StartNode::new()));
424
425        assert!(SubPlan::connect_upstream(empty.clone(), populated.clone()).is_err());
426        assert!(SubPlan::connect_upstream(populated, empty).is_err());
427    }
428
429    #[test]
430    fn connect_upstream_rejects_leaf_downstream() {
431        use crate::planning::plan::core::nodes::StartNode;
432
433        let downstream = SubPlan::from_single_node(PlanNodeEnum::Start(StartNode::new()));
434        let upstream = SubPlan::from_single_node(PlanNodeEnum::Start(StartNode::new()));
435
436        assert!(
437            SubPlan::connect_upstream(downstream, upstream).is_err(),
438            "a leaf node cannot host a unary upstream"
439        );
440    }
441}
linkrs_query/planning/plan/execution_plan.rs

linkrs_query/planning/plan/
execution_plan.rs
