linkrs_query/executor/streaming/plan/
arena_builder.rs

1//! PhysicalPlanBuilder: builds an arena [`PhysicalPlan`] directly from planner nodes.
2//!
3//! This builder walks the [`PlanNodeEnum`] tree directly, creating arena operators
4//! and a fragment DAG.
5
6use super::context::PhysicalPlanBuildContext;
7use std::collections::HashMap;
8
9use super::types::{
10    CapabilitySet, FragmentGraph, FragmentSpec, LogicalNodeId, OutputContract, PhysicalOperatorId,
11    PhysicalOperatorIdAllocator, PhysicalOperatorSpec, PhysicalPlan, PlanCompatibility,
12    PlanFingerprint,
13};
14use crate::executor::base::ExecutionContext;
15use crate::executor::build_error::PlanBuildError;
16use crate::planning::plan::core::nodes::base::plan_node_enum::PlanNodeEnum;
17
18mod assembler;
19mod metadata;
20mod partition;
21mod specs;
22
23use assembler::{ArenaFragmentAllocator, ArenaPlanAssembler};
24
25/// Builds an arena [`PhysicalPlan`] from a [`PlanNodeEnum`] tree.
26pub struct PhysicalPlanBuilder;
27
28impl PhysicalPlanBuilder {
29    /// Build a complete [`PhysicalPlan`] from a plan node.
30    ///
31    /// Walks the [`PlanNodeEnum`] tree directly, creating arena operators
32    /// and a fragment DAG.  When `ctx.partition_spec` is set and the tree is
33    /// partitionable, a partitioned plan (N local source fragments + one
34    /// exchange fragment + global fragments) is produced instead; otherwise
35    /// the single-tree path is used and the reason is recorded in
36    /// `ctx.parallel_fallback_reason`.
37    pub fn build(
38        node: &PlanNodeEnum,
39        ctx: &mut PhysicalPlanBuildContext,
40        exec_ctx: &ExecutionContext,
41    ) -> Result<PhysicalPlan, PlanBuildError> {
42        let mut partitioned = None;
43        if let Some(spec) = ctx.partition_spec.clone() {
44            match partition::build_partitioned(node, &spec, exec_ctx) {
45                Ok(Some(result)) => {
46                    ctx.parallel_fallback_reason.clear();
47                    partitioned = Some(result);
48                }
49                Ok(None) => {
50                    if ctx.parallel_fallback_reason.is_empty() {
51                        ctx.parallel_fallback_reason =
52                            "plan shape does not support partitioned execution".to_string();
53                    }
54                }
55                Err(error) => {
56                    if ctx.parallel_fallback_reason.is_empty() {
57                        ctx.parallel_fallback_reason =
58                            format!("partitioned execution unavailable ({error})");
59                    }
60                }
61            }
62        }
63
64        let (mut operators, mut fragments, root_fid, root_op_id) = match partitioned {
65            Some((operators, fragments, root_fid, root_op_id)) => {
66                (operators, fragments, root_fid, root_op_id)
67            }
68            None => Self::build_serial(node, exec_ctx)?,
69        };
70
71        Self::finalize(&mut operators, &mut fragments, root_fid, root_op_id, ctx)
72    }
73
74    /// Build the single linear fragment DAG for an unpartitioned plan.
75    fn build_serial(
76        node: &PlanNodeEnum,
77        exec_ctx: &ExecutionContext,
78    ) -> Result<
79        (
80            Vec<PhysicalOperatorSpec>,
81            Vec<FragmentSpec>,
82            super::types::FragmentId,
83            PhysicalOperatorId,
84        ),
85        PlanBuildError,
86    > {
87        let mut operators: Vec<PhysicalOperatorSpec> = Vec::new();
88        let mut fragments: Vec<FragmentSpec> = Vec::new();
89        let mut op_alloc = PhysicalOperatorIdAllocator::new();
90        let mut frag_alloc = ArenaFragmentAllocator::new();
91
92        let (root_fid, root_op_id) = ArenaPlanAssembler::convert_node(
93            node,
94            &mut operators,
95            &mut fragments,
96            &mut op_alloc,
97            &mut frag_alloc,
98            exec_ctx,
99        )?;
100
101        Ok((operators, fragments, root_fid, root_op_id))
102    }
103
104    /// Run the shared metadata passes and assemble the final plan.
105    fn finalize(
106        operators: &mut Vec<PhysicalOperatorSpec>,
107        fragments: &mut Vec<FragmentSpec>,
108        root_fid: super::types::FragmentId,
109        root_op_id: PhysicalOperatorId,
110        ctx: &PhysicalPlanBuildContext,
111    ) -> Result<PhysicalPlan, PlanBuildError> {
112        metadata::propagate_layouts(operators, fragments)?;
113        metadata::populate_input_contracts(operators, fragments)?;
114        metadata::populate_runtime_metadata(operators);
115        metadata::populate_choice_reasons(operators);
116        metadata::populate_estimated_rows(
117            operators,
118            fragments,
119            &ctx.statistics.per_node_row_estimates,
120        );
121
122        let output = operators
123            .iter()
124            .find(|op| op.operator_id == root_op_id)
125            .map(|op| metadata::output_contract(&op.spec, op.output_layout.clone()))
126            .unwrap_or_else(|| OutputContract {
127                output_layout: super::super::slot::SlotLayout::new(vec![]),
128                always_produces_row: false,
129                nullability: Vec::new(),
130                ordering: Vec::new(),
131                delivery_streamable: true,
132                pipeline_mode: super::types::PipelineMode::Pipelined,
133            });
134
135        let mut logical_to_physical: HashMap<LogicalNodeId, Vec<PhysicalOperatorId>> =
136            HashMap::new();
137        for operator in &*operators {
138            if let Some(logical_id) = operator.logical_node_id {
139                logical_to_physical
140                    .entry(logical_id)
141                    .or_default()
142                    .push(operator.operator_id);
143            }
144        }
145        let fingerprint = PlanFingerprint::compute(operators);
146
147        let mut required_capabilities = CapabilitySet::EMPTY;
148        for operator in &*operators {
149            required_capabilities.insert(metadata::capability_for_operator(&operator.spec));
150        }
151
152        let plan = PhysicalPlan {
153            operators: std::mem::take(operators),
154            logical_to_physical,
155            fragments: FragmentGraph::new(std::mem::take(fragments), root_fid),
156            root_fragment: root_fid,
157            output,
158            compatibility: PlanCompatibility {
159                fingerprint,
160                layout_version: ctx.schema.as_ref().map(|s| s.layout_version),
161                required_capabilities: required_capabilities.clone(),
162                planning_config_hash: ctx.config.config_hash,
163                optimizer_version: ctx.config.optimizer_version,
164            },
165            required_capabilities,
166            parameter_schema: ctx.parameter_schema.clone(),
167            parallel_fallback_reason: ctx.parallel_fallback_reason.clone(),
168            cbo_notes: ctx.cbo_notes.clone(),
169            statistics_summary: ctx.statistics_summary.clone(),
170            partition_spec: ctx.partition_spec.clone(),
171        };
172
173        Ok(plan)
174    }
175}
176
177#[cfg(test)]
178mod tests {
179    use super::*;
180    use crate::executor::base::ExecutionContext;
181    use crate::executor::streaming::plan::types::{InputContract, OperatorKindSpec};
182    use crate::executor::streaming::SourceSpec;
183    use crate::planning::plan::core::nodes::control_flow::start_node::StartNode;
184    use linkrs_core::types::expr::expression_context::ExpressionAnalysisContext;
185    use std::sync::Arc;
186
187    fn test_context() -> ExecutionContext {
188        ExecutionContext::new(Arc::new(ExpressionAnalysisContext::new()))
189    }
190
191    #[test]
192    fn test_build_start() {
193        let start = StartNode::new();
194        let node = PlanNodeEnum::Start(start);
195        let mut ctx = PhysicalPlanBuildContext::new();
196        let exec_ctx = test_context();
197
198        let plan = PhysicalPlanBuilder::build(&node, &mut ctx, &exec_ctx).unwrap();
199        assert_eq!(plan.operator_count(), 1);
200        assert!(matches!(
201            plan.operator(PhysicalOperatorId(0)).unwrap().spec,
202            OperatorKindSpec::Source(SourceSpec::Start)
203        ));
204    }
205
206    #[test]
207    fn test_build_start_has_explain_name() {
208        let start = StartNode::new();
209        let node = PlanNodeEnum::Start(start);
210        let mut ctx = PhysicalPlanBuildContext::new();
211        let exec_ctx = test_context();
212
213        let plan = PhysicalPlanBuilder::build(&node, &mut ctx, &exec_ctx).unwrap();
214        let op = plan.operator(PhysicalOperatorId(0)).unwrap();
215        assert_eq!(op.explain_name, "Start");
216    }
217
218    #[test]
219    fn test_build_then_materialize_start() {
220        use crate::executor::base::MemoryBudget;
221        use crate::executor::streaming::instance::{
222            QueryBindings, QueryExecutionInstance, ResultSink,
223        };
224        use crate::executor::streaming::transaction_scope::TransactionScope;
225        use std::sync::Arc;
226
227        let start = StartNode::new();
228        let node = PlanNodeEnum::Start(start);
229        let mut ctx = PhysicalPlanBuildContext::new();
230        let exec_ctx = test_context();
231
232        let plan = PhysicalPlanBuilder::build(&node, &mut ctx, &exec_ctx).unwrap();
233        let plan_arc = Arc::new(plan);
234
235        crate::executor::streaming::plan::validator::PhysicalPlanValidator::validate(&plan_arc)
236            .unwrap();
237
238        let bindings = QueryBindings {
239            parameters: Arc::new(std::collections::HashMap::new()),
240            session_variables: Arc::new(std::collections::HashMap::new()),
241            parameter_frame: None,
242            space_name: None,
243            storage: None,
244            bound_snapshot: None,
245            memory_budget: MemoryBudget::default_budget(),
246            max_workers: 1,
247            chunk_size: 2048,
248            max_buffered_chunks: 4,
249            query_id: 1,
250            cancel_token: None,
251            session_id: None,
252            user_name: None,
253            query_text: None,
254            transaction: TransactionScope::None,
255            shared_scheduler: None,
256            partition_count: 0,
257            arena: None,
258            feedback_history: None,
259            columnar_policy: None,
260            macro_manager: None,
261            type_alias_manager: None,
262            sequence_manager: None,
263            search: crate::executor::base::SearchContext::default(),
264        };
265
266        let result = QueryExecutionInstance::instantiate_plan(
267            plan_arc,
268            bindings,
269            ResultSink::Discard,
270            None,
271            None,
272        );
273        assert!(result.is_ok());
274    }
275
276    #[test]
277    fn propagates_layout_through_blocking_pipeline_and_output_contract() {
278        use crate::planning::plan::core::nodes::control_flow::start_node::StartNode;
279        use crate::planning::plan::core::nodes::operation::sort_node::LimitNode;
280
281        let start = StartNode::new();
282        let limit =
283            LimitNode::new(PlanNodeEnum::Start(start), 0, 10).expect("limit plan should build");
284        let node = PlanNodeEnum::Limit(limit);
285        let mut ctx = PhysicalPlanBuildContext::new();
286        let exec_ctx = test_context();
287
288        let plan = PhysicalPlanBuilder::build(&node, &mut ctx, &exec_ctx).unwrap();
289        let start_op = plan.operator(PhysicalOperatorId(0)).unwrap();
290        let limit_op = plan.operator(PhysicalOperatorId(1)).unwrap();
291
292        assert_eq!(start_op.explain_name, "Start");
293        assert_eq!(limit_op.explain_name, "Limit");
294        assert_eq!(plan.fragment_count(), 1);
295    }
296
297    #[test]
298    fn transaction_command_has_start_input_fragment() {
299        use crate::planning::plan::core::nodes::control_flow::BeginTransactionNode;
300
301        let node = PlanNodeEnum::BeginTransaction(BeginTransactionNode::new(42));
302        let mut ctx = PhysicalPlanBuildContext::new();
303        let exec_ctx = test_context();
304
305        let plan = PhysicalPlanBuilder::build(&node, &mut ctx, &exec_ctx)
306            .expect("transaction plan should build");
307        assert_eq!(plan.operator_count(), 2);
308        assert_eq!(plan.fragment_count(), 2);
309
310        let root = plan
311            .fragments
312            .get(plan.root_fragment)
313            .expect("root fragment should exist");
314        assert_eq!(root.inputs.len(), 1);
315        assert!(matches!(
316            plan.operator(root.root_operator)
317                .expect("transaction root operator should exist")
318                .input_contract,
319            InputContract::UnaryInput(_)
320        ));
321        crate::executor::streaming::plan::validator::PhysicalPlanValidator::validate(&Arc::new(
322            plan,
323        ))
324        .expect("transaction plan should validate");
325    }
326
327    #[test]
328    fn correlated_apply_builds_nested_sub_plan_rooted_at_argument() {
329        use crate::executor::streaming::operators::spec::ApplySpec;
330        use crate::planning::plan::core::nodes::base::plan_node_traits::PlanNode;
331        use crate::planning::plan::core::nodes::control_flow::ArgumentNode;
332        use crate::planning::plan::core::nodes::graph_operations::graph_operations_node::CorrelatedApplyNode;
333        use crate::planning::plan::core::nodes::join::CrossJoinNode;
334        use crate::planning::plan::core::nodes::operation::filter_node::FilterNode;
335        use linkrs_core::types::expr::expression_context::ExpressionAnalysisContext;
336        use linkrs_core::types::{ContextualExpression, ExpressionMeta};
337
338        // Outer (left) plan.
339        let left = PlanNodeEnum::Start(StartNode::new());
340
341        // Right subtree: Filter(CrossJoin(Argument(col_names = outer), Start),
342        // true) — the self-contained plan rebuilt per outer row at runtime.
343        let mut argument = ArgumentNode::new(-2, "_correlated_apply");
344        argument.set_col_names(vec!["t".to_string()]);
345        let cross = CrossJoinNode::new(PlanNodeEnum::Argument(argument), left.clone())
346            .expect("cross join should build");
347        let expr_ctx = Arc::new(ExpressionAnalysisContext::new());
348        let expr_id = expr_ctx.register_expression(ExpressionMeta::new(
349            linkrs_core::Expression::Literal(linkrs_core::Value::Bool(true)),
350        ));
351        let condition = ContextualExpression::new(expr_id, expr_ctx);
352        let filter = FilterNode::new(cross.into_enum(), condition).expect("filter should build");
353        let right = PlanNodeEnum::Filter(filter);
354
355        let apply =
356            CorrelatedApplyNode::new(left, right, false).expect("correlated apply should build");
357        let node = PlanNodeEnum::CorrelatedApply(apply);
358
359        let mut ctx = PhysicalPlanBuildContext::new();
360        let exec_ctx = test_context();
361        let plan =
362            PhysicalPlanBuilder::build(&node, &mut ctx, &exec_ctx).expect("plan should build");
363
364        let apply_op = plan
365            .operators
366            .iter()
367            .find(|op| {
368                matches!(
369                    op.spec,
370                    OperatorKindSpec::Apply(ApplySpec::CorrelatedApply { .. })
371                )
372            })
373            .expect("plan contains a CorrelatedApply operator");
374
375        let sub_plan = match &apply_op.spec {
376            OperatorKindSpec::Apply(ApplySpec::CorrelatedApply { sub_plan, anti }) => {
377                assert!(!anti, "correlated apply defaults to semi");
378                sub_plan
379            }
380            _ => unreachable!("spec match earlier guarantees the variant"),
381        };
382
383        // The nested sub-plan is self-contained and contains an Argument
384        // source whose layout mirrors the outer columns.
385        let argument_source = sub_plan
386            .operators
387            .iter()
388            .find(|op| {
389                matches!(
390                    op.spec,
391                    OperatorKindSpec::Source(SourceSpec::Argument { .. })
392                )
393            })
394            .expect("sub-plan must be rooted at an Argument source");
395        match &argument_source.spec {
396            OperatorKindSpec::Source(SourceSpec::Argument { col_names }) => {
397                assert_eq!(
398                    col_names,
399                    &["t".to_string()],
400                    "Argument col_names mirror the outer layout"
401                );
402            }
403            _ => unreachable!(),
404        }
405    }
406}
linkrs_query/executor/streaming/plan/arena_builder.rs

linkrs_query/executor/streaming/plan/
arena_builder.rs
