linkrs_query/planning/plan/
factorization_compute.rs

1mod access;
2mod assign;
3mod control_flow;
4mod flat_leaf;
5mod join;
6mod operation;
7mod set_ops;
8mod traversal;
9mod unwind;
10
11use linkrs_core::types::expr::contextual::ContextualExpression;
12use linkrs_core::types::expr::ExpressionId;
13
14use crate::planning::plan::factorization::{
15    FGroupPos, FactorizationError, FactorizedSchema, FactorizedSchemaCompute,
16};
17
18use crate::planning::plan::logical::logical_node_enum::LogicalNodeEnum;
19
20pub(super) fn resolve_id(expr: &ContextualExpression) -> ExpressionId {
21    expr.id().clone()
22}
23
24/// Register bare output names on an expansion output group.
25///
26/// Traversal nodes fan out into a new unflat group but carry no tracked
27/// `ExpressionId` for the output. Downstream `Variable` references resolve
28/// by name through `expression_name_to_group`, so the aliases must be
29/// recorded explicitly. Without this the output group stays empty and every
30/// downstream use falls back to flatten-all via
31/// `GroupDependencyAnalyzer::mark_unresolved`.
32pub(super) fn register_output_names(
33    schema: &mut FactorizedSchema,
34    output_var: Option<&str>,
35    col_names: &[String],
36    group: FGroupPos,
37) -> Result<(), FactorizationError> {
38    if let Some(var) = output_var {
39        schema.insert_name_for_group(var.to_string(), group)?;
40    }
41    for name in col_names {
42        schema.insert_name_for_group(name.clone(), group)?;
43    }
44    Ok(())
45}
46
47/// Split binary child schemas into `(left, right)`, defaulting missing sides
48/// to an empty schema so key-aware join helpers never index out of bounds.
49fn pair(child_schemas: &[FactorizedSchema]) -> (&FactorizedSchema, &FactorizedSchema) {
50    static EMPTY: std::sync::OnceLock<FactorizedSchema> = std::sync::OnceLock::new();
51    let empty = EMPTY.get_or_init(FactorizedSchema::new);
52    let left = child_schemas.first().unwrap_or(empty);
53    let right = child_schemas.get(1).unwrap_or(empty);
54    (left, right)
55}
56
57/// Handling module for each logical operator in `compute_factorized_schema`.
58///
59/// This match is intentionally exhaustive with no wildcard arm: adding a new
60/// `LogicalNodeEnum` variant fails compilation here, which forces the author
61/// to update both this table and the `FactorizationRewriter::visit_operator`
62/// dispatch (R5 drift guard). The returned string names the compute submodule
63/// that owns the operator.
64pub(crate) fn dispatch_module(node: &LogicalNodeEnum) -> &'static str {
65    match node {
66        LogicalNodeEnum::ScanVertices(_) => "access",
67        LogicalNodeEnum::ScanEdges(_) => "access",
68        LogicalNodeEnum::GetVertices(_) => "access",
69        LogicalNodeEnum::GetEdges(_) => "access",
70        LogicalNodeEnum::GetNeighbors(_) => "access",
71        LogicalNodeEnum::Start(_) => "access",
72        LogicalNodeEnum::Project(_) => "operation",
73        LogicalNodeEnum::Filter(_) => "operation",
74        LogicalNodeEnum::Aggregate(_) => "operation",
75        LogicalNodeEnum::Flatten(_) => "operation",
76        LogicalNodeEnum::Sort(_) => "operation",
77        LogicalNodeEnum::TopN(_) => "operation",
78        LogicalNodeEnum::Window(_) => "operation",
79        LogicalNodeEnum::Dedup(_) => "operation",
80        LogicalNodeEnum::Limit(_) => "operation",
81        LogicalNodeEnum::Skip(_) => "operation",
82        LogicalNodeEnum::Sample(_) => "operation",
83        LogicalNodeEnum::InnerJoin(_) => "join",
84        LogicalNodeEnum::LeftJoin(_) => "join",
85        LogicalNodeEnum::RightJoin(_) => "join",
86        LogicalNodeEnum::CrossJoin(_) => "join",
87        LogicalNodeEnum::FullOuterJoin(_) => "join",
88        LogicalNodeEnum::SemiJoin(_) => "join",
89        LogicalNodeEnum::Traverse(_) => "traversal",
90        LogicalNodeEnum::Expand(_) => "traversal",
91        LogicalNodeEnum::ExpandAll(_) => "traversal",
92        LogicalNodeEnum::AppendVertices(_) => "traversal",
93        LogicalNodeEnum::BiExpand(_) => "traversal",
94        LogicalNodeEnum::BiTraverse(_) => "traversal",
95        LogicalNodeEnum::Union(_) => "set_ops",
96        LogicalNodeEnum::Minus(_) => "set_ops",
97        LogicalNodeEnum::Intersect(_) => "set_ops",
98        LogicalNodeEnum::WcoIntersect(_) => "set_ops",
99        LogicalNodeEnum::Assign(_) => "assign",
100        LogicalNodeEnum::Select(_) => "control_flow",
101        LogicalNodeEnum::Loop(_) => "control_flow",
102        LogicalNodeEnum::BeginTransaction(_) => "control_flow",
103        LogicalNodeEnum::Commit(_) => "control_flow",
104        LogicalNodeEnum::Rollback(_) => "control_flow",
105        LogicalNodeEnum::PassThrough(_) => "control_flow",
106        LogicalNodeEnum::Argument(_) => "control_flow",
107        LogicalNodeEnum::Unwind(_) => "unwind",
108        LogicalNodeEnum::FulltextSearch(_) => "flat_leaf",
109        LogicalNodeEnum::FulltextLookup(_) => "flat_leaf",
110        LogicalNodeEnum::MatchFulltext(_) => "flat_leaf",
111        #[cfg(feature = "vector")]
112        LogicalNodeEnum::VectorSearch(_) => "flat_leaf",
113        #[cfg(feature = "vector")]
114        LogicalNodeEnum::VectorLookup(_) => "flat_leaf",
115        #[cfg(feature = "vector")]
116        LogicalNodeEnum::VectorMatch(_) => "flat_leaf",
117        LogicalNodeEnum::InsertVertices(_) => "flat_leaf",
118        LogicalNodeEnum::InsertEdges(_) => "flat_leaf",
119        LogicalNodeEnum::Update(_) => "flat_leaf",
120        LogicalNodeEnum::DeleteVertices(_) => "flat_leaf",
121        LogicalNodeEnum::DeleteEdges(_) => "flat_leaf",
122        LogicalNodeEnum::DeleteIndex(_) => "flat_leaf",
123        LogicalNodeEnum::PipeDeleteVertices(_) => "flat_leaf",
124        LogicalNodeEnum::PipeDeleteEdges(_) => "flat_leaf",
125        LogicalNodeEnum::CopyFrom(_) => "flat_leaf",
126        LogicalNodeEnum::CopyTo(_) => "flat_leaf",
127        LogicalNodeEnum::DataCollect(_) => "flat_leaf",
128        LogicalNodeEnum::Materialize(_) => "flat_leaf",
129        LogicalNodeEnum::RollUpApply(_) => "flat_leaf",
130        LogicalNodeEnum::Apply(_) => "flat_leaf",
131        LogicalNodeEnum::PatternApply(_) => "flat_leaf",
132        LogicalNodeEnum::CorrelatedApply(_) => "flat_leaf",
133        LogicalNodeEnum::MultiShortestPath(_) => "flat_leaf",
134        LogicalNodeEnum::BFSShortest(_) => "flat_leaf",
135        LogicalNodeEnum::AllPaths(_) => "flat_leaf",
136        LogicalNodeEnum::ShortestPath(_) => "flat_leaf",
137    }
138}
139
140/// Schema for bidirectional expansion over two child schemas.
141///
142/// Child order follows the node inputs: the first child is the probe side
143/// and is flattened before fan-out, the second child is the build side
144/// whose bindings are merged into scope. Any remaining nested group is
145/// flattened so the expansion output group is the single unflat group.
146/// `output_aliases` are registered on the new group via
147/// `register_output_names` semantics so downstream variables resolve
148/// precisely instead of via flatten-all fallback.
149pub(super) fn bi_expand_schema(
150    child_schemas: &[FactorizedSchema],
151    output_aliases: &[String],
152) -> Result<FactorizedSchema, FactorizationError> {
153    let mut schema = child_schemas.first().cloned().unwrap_or_default();
154    if schema.has_unflat_group() {
155        if let Some(pos) = schema.unflat_group_pos() {
156            schema.flatten_group(pos)?;
157        }
158    }
159    if let Some(build) = child_schemas.get(1) {
160        let mapping = schema.merge_groups_from(build);
161        for (expr_id, gpos) in build.expression_to_group_iter() {
162            let new_pos = mapping.get(gpos).copied().unwrap_or(*gpos);
163            schema.insert_to_scope_may_repeat(expr_id.clone(), new_pos)?;
164        }
165        if schema.has_unflat_group() {
166            if let Some(pos) = schema.unflat_group_pos() {
167                schema.flatten_group(pos)?;
168            }
169        }
170    }
171    let out = schema.create_group();
172    for alias in output_aliases {
173        schema.insert_name_for_group(alias.clone(), out)?;
174    }
175    schema.validate_at_most_one_unflat()?;
176    Ok(schema)
177}
178
179impl FactorizedSchemaCompute for LogicalNodeEnum {
180    fn compute_factorized_schema(
181        &mut self,
182        child_schemas: &[FactorizedSchema],
183    ) -> Result<FactorizedSchema, FactorizationError> {
184        // Keep the R5 drift table linked: every operator must have a dispatch
185        // entry, otherwise the tables have drifted.
186        debug_assert!(!dispatch_module(self).is_empty());
187        let schema = match self {
188            LogicalNodeEnum::ScanVertices(n) => access::scan_vertices(n)?,
189            LogicalNodeEnum::ScanEdges(n) => access::scan_edges(n)?,
190            LogicalNodeEnum::GetVertices(n) => access::get_vertices(n, child_schemas)?,
191            LogicalNodeEnum::GetEdges(n) => access::get_edges(n)?,
192            LogicalNodeEnum::GetNeighbors(n) => access::get_neighbors(n, child_schemas)?,
193            LogicalNodeEnum::Start(_) => access::start()?,
194
195            LogicalNodeEnum::Project(n) => operation::project(n, child_schemas)?,
196            LogicalNodeEnum::Filter(n) => operation::filter(n, child_schemas)?,
197            LogicalNodeEnum::Aggregate(n) => operation::aggregate(n, child_schemas)?,
198            LogicalNodeEnum::Flatten(n) => operation::flatten(n, child_schemas)?,
199            LogicalNodeEnum::Sort(n) => operation::sort(n, child_schemas)?,
200            LogicalNodeEnum::TopN(n) => operation::top_n(n, child_schemas)?,
201            LogicalNodeEnum::Window(n) => operation::window(n, child_schemas)?,
202            LogicalNodeEnum::Dedup(n) => operation::dedup(n, child_schemas)?,
203            LogicalNodeEnum::Limit(n) => operation::limit(n, child_schemas)?,
204            LogicalNodeEnum::Skip(n) => operation::skip(n, child_schemas)?,
205            LogicalNodeEnum::Sample(n) => operation::sample(n, child_schemas)?,
206
207            LogicalNodeEnum::InnerJoin(n) => {
208                let (left, right) = pair(child_schemas);
209                join::binary_join_inner(left, right, &n.hash_keys, &n.probe_keys)?
210            }
211            LogicalNodeEnum::LeftJoin(n) => {
212                let (left, right) = pair(child_schemas);
213                join::binary_join_inner(left, right, &n.hash_keys, &n.probe_keys)?
214            }
215            LogicalNodeEnum::RightJoin(n) => {
216                let (left, right) = pair(child_schemas);
217                join::binary_join_right(left, right, &n.hash_keys, &n.probe_keys)?
218            }
219            LogicalNodeEnum::CrossJoin(n) => {
220                let (left, right) = pair(child_schemas);
221                if n.hash_keys.is_empty() && n.probe_keys.is_empty() {
222                    join::cross_join_no_keys(left, right)?
223                } else {
224                    join::binary_join_inner(left, right, &n.hash_keys, &n.probe_keys)?
225                }
226            }
227            LogicalNodeEnum::FullOuterJoin(n) => {
228                let (left, right) = pair(child_schemas);
229                join::binary_join_full_outer(left, right, &n.hash_keys, &n.probe_keys)?
230            }
231            LogicalNodeEnum::SemiJoin(n) => {
232                let (left, right) = pair(child_schemas);
233                join::binary_join_inner(left, right, &n.hash_keys, &n.probe_keys)?
234            }
235
236            LogicalNodeEnum::Traverse(n) => traversal::traverse(n, child_schemas)?,
237            LogicalNodeEnum::Expand(n) => traversal::expand(n, child_schemas)?,
238            LogicalNodeEnum::ExpandAll(n) => traversal::expand_all(n, child_schemas)?,
239            LogicalNodeEnum::AppendVertices(n) => traversal::append_vertices(n, child_schemas)?,
240            LogicalNodeEnum::BiExpand(n) => traversal::bi_expand(n, child_schemas)?,
241            LogicalNodeEnum::BiTraverse(n) => traversal::bi_traverse(n, child_schemas)?,
242
243            LogicalNodeEnum::Union(_) | LogicalNodeEnum::Minus(_) => {
244                set_ops::union_minus(child_schemas)?
245            }
246            LogicalNodeEnum::Intersect(_) => set_ops::intersect(child_schemas)?,
247            LogicalNodeEnum::WcoIntersect(n) => set_ops::wco_intersect(n, child_schemas)?,
248
249            LogicalNodeEnum::Assign(n) => assign::assign(n, child_schemas)?,
250
251            LogicalNodeEnum::Select(n) => control_flow::select(n, child_schemas)?,
252            LogicalNodeEnum::Loop(n) => control_flow::loop_node(n, child_schemas)?,
253            LogicalNodeEnum::BeginTransaction(_)
254            | LogicalNodeEnum::Commit(_)
255            | LogicalNodeEnum::Rollback(_)
256            | LogicalNodeEnum::PassThrough(_)
257            | LogicalNodeEnum::Argument(_) => control_flow::passthrough(child_schemas)?,
258
259            LogicalNodeEnum::Unwind(n) => unwind::unwind(n, child_schemas)?,
260
261            LogicalNodeEnum::FulltextSearch(_)
262            | LogicalNodeEnum::FulltextLookup(_)
263            | LogicalNodeEnum::MatchFulltext(_) => flat_leaf::flat_leaf()?,
264            #[cfg(feature = "vector")]
265            LogicalNodeEnum::VectorSearch(_)
266            | LogicalNodeEnum::VectorLookup(_)
267            | LogicalNodeEnum::VectorMatch(_) => flat_leaf::flat_leaf()?,
268
269            // Single-input barrier: output is fully flat. The rewriter
270            // inserts the matching `FlattenAll` nodes explicitly, so compute
271            // and rewrite stay in sync (barrier semantics).
272            LogicalNodeEnum::DataCollect(_)
273            | LogicalNodeEnum::Materialize(_)
274            | LogicalNodeEnum::RollUpApply(_) => flat_leaf::flatten_all_from_child(child_schemas)?,
275
276            LogicalNodeEnum::InsertVertices(_)
277            | LogicalNodeEnum::InsertEdges(_)
278            | LogicalNodeEnum::Update(_) => flat_leaf::flat_leaf()?,
279
280            // Binary barrier (Apply family, shortest-path family): no
281            // per-operator flatten rule, so merge both children and flatten.
282            // The rewriter inserts `FlattenAll` on each side explicitly.
283            LogicalNodeEnum::Apply(_)
284            | LogicalNodeEnum::PatternApply(_)
285            | LogicalNodeEnum::CorrelatedApply(_)
286            | LogicalNodeEnum::MultiShortestPath(_)
287            | LogicalNodeEnum::BFSShortest(_)
288            | LogicalNodeEnum::AllPaths(_)
289            | LogicalNodeEnum::ShortestPath(_) => flat_leaf::barrier_binary(child_schemas)?,
290
291            _ => flat_leaf::flatten_all_from_child(child_schemas)?,
292        };
293        debug_assert!(
294            schema.has_at_most_one_unflat(),
295            "compute_factorized_schema: at most one unflat group invariant violated"
296        );
297        Ok(schema)
298    }
299
300    fn compute_flat_schema(
301        &mut self,
302        child_schemas: &[FactorizedSchema],
303    ) -> Result<FactorizedSchema, FactorizationError> {
304        let flat_children: Vec<FactorizedSchema> = child_schemas
305            .iter()
306            .map(|cs| cs.flat_copy())
307            .collect::<Result<Vec<_>, _>>()?;
308        let mut result = self.compute_factorized_schema(&flat_children)?;
309        result.flatten_all()?;
310        result.validate_at_most_one_unflat()?;
311        Ok(result)
312    }
313}
314
315#[cfg(test)]
316mod tests;
linkrs_query/planning/plan/factorization_compute.rs

linkrs_query/planning/plan/
factorization_compute.rs
