linkrs_query/planning/plan/
factorization.rs

1//! Factorized-plan schema description for the query optimizer.
2//!
3//! **Naming note (R1).** The types here (`FactorizedSchema`,
4//! `FactorizationGroup`, `FactorizationRewriter`) reuse Kuzu/Ladybug
5//! terminology. In linkrs they describe a *planning-side* grouping of
6//! expressions into nested levels and drive a *minimal-flatten* rewrite —
7//! they are **not** a "factorized-table execution engine" (no compressed
8//! factor tables, `SEMI_MASKER`, or multiplicity reduction). The benefit is
9//! fewer materialization points and smaller fan-out, not skipping the
10//! cross-product. Keep this distinction in mind when reading the code or the
11//! EXPLAIN output: a `LogicalFlatten` is a cross-product materialization
12//! point, not a factorization step.
13
14use linkrs_core::types::expr::ExpressionId;
15use thiserror::Error;
16
17mod compute;
18mod expression;
19mod group;
20
21pub use compute::{SchemaUtils, SinkOperatorUtil};
22pub use expression::FactorizedSchema;
23pub use group::FactorizationGroup;
24
25/// Factorization group position identifier.
26pub type FGroupPos = u32;
27
28/// Invalid group position sentinel.
29pub const INVALID_F_GROUP_POS: FGroupPos = u32::MAX;
30
31/// Errors raised while building or validating a [`FactorizedSchema`].
32///
33/// These replace the previous `assert!`-based invariant checks so a malformed
34/// plan surfaces as a query error instead of aborting the server process.
35/// See `AGENTS.md`: "Never use unwrap".
36#[derive(Error, Debug, Clone, PartialEq)]
37pub enum FactorizationError {
38    /// More than one group is left unflat at the same nesting level.
39    #[error("at most one unflat group allowed, found {0}")]
40    TooManyUnflatGroups(usize),
41    /// A group that was expected to be flat is still unflat.
42    #[error("group {0} expected flat but is unflat")]
43    GroupExpectedFlat(FGroupPos),
44    /// A group position passed to a flatten/set operation is out of range.
45    #[error("group_pos {0} out of range")]
46    GroupPosOutOfRange(FGroupPos),
47    /// An expression was inserted into scope more than once.
48    #[error("expression {0:?} already in scope")]
49    ExpressionAlreadyInScope(ExpressionId),
50    /// An expression was mapped to a group more than once.
51    #[error("expression {0:?} already mapped to group")]
52    ExpressionAlreadyMapped(ExpressionId),
53    /// A duplicate expression id was registered inside a group.
54    #[error("duplicate expression id {0:?} in group")]
55    DuplicateExpressionId(ExpressionId),
56    /// A duplicate expression name was registered inside a group.
57    #[error("duplicate expression name {0} in group")]
58    DuplicateExpressionName(String),
59    /// `flatten_group` was given an invalid group position.
60    #[error("flatten_group: invalid pos {0}")]
61    InvalidFlattenPos(FGroupPos),
62    /// `set_group_as_single_state` was given an invalid group position.
63    #[error("set_group_as_single_state: invalid pos {0}")]
64    InvalidSingleStatePos(FGroupPos),
65    /// `set_flat` was called on a group that is already flat.
66    #[error("group already flat")]
67    GroupAlreadyFlat,
68    /// `set_single_state` was called on a group that is already single-state.
69    #[error("group already single state")]
70    GroupAlreadySingleState,
71    /// A group position set passed to a leader/validator was empty.
72    #[error("groupPositions empty")]
73    EmptyGroupPositions,
74    /// A non-empty group position set was required but found empty.
75    #[error("expected non-empty group positions")]
76    NonEmptyGroupPositions,
77}
78
79/// Trait for operators that can compute factorized schemas.
80///
81/// `child_schemas` must be the bottom-up computed results for the direct children;
82/// passing an empty slice forces recomputation and violates the factorization invariant.
83pub trait FactorizedSchemaCompute {
84    fn compute_factorized_schema(
85        &mut self,
86        child_schemas: &[FactorizedSchema],
87    ) -> Result<FactorizedSchema, FactorizationError>;
88    fn compute_flat_schema(
89        &mut self,
90        child_schemas: &[FactorizedSchema],
91    ) -> Result<FactorizedSchema, FactorizationError>;
92}
93
94#[cfg(test)]
95mod tests {
96    use super::*;
97    use std::collections::HashSet;
98
99    fn expr(id: u64) -> ExpressionId {
100        ExpressionId::new(id)
101    }
102
103    #[test]
104    fn group_basic() {
105        let mut g = FactorizationGroup::new();
106        assert!(!g.is_flat());
107        g.set_flat().unwrap();
108        assert!(g.is_flat());
109    }
110
111    #[test]
112    fn group_single_state_forces_flat() {
113        let mut g = FactorizationGroup::new();
114        g.set_single_state().unwrap();
115        assert!(g.is_flat());
116        assert!(g.is_single_state());
117    }
118
119    #[test]
120    fn schema_single_flat_group() {
121        let mut schema = FactorizedSchema::new();
122        let pos = schema.create_flat_group(false);
123        assert_eq!(pos, 0);
124        schema.insert_to_group_and_scope(expr(1), pos).unwrap();
125        schema.insert_to_group_and_scope(expr(2), pos).unwrap();
126        assert_eq!(schema.num_groups(), 1);
127        assert!(schema.get_group(pos).expect("group").is_flat());
128        assert_eq!(schema.get_group_pos(&expr(1)), Some(0));
129    }
130
131    #[test]
132    fn schema_unflat_group_and_flatten() {
133        let mut schema = FactorizedSchema::new();
134        let flat_pos = schema.create_flat_group(false);
135        let unflat_pos = schema.create_group();
136        schema
137            .insert_to_group_and_scope(expr(10), flat_pos)
138            .unwrap();
139        schema
140            .insert_to_group_and_scope(expr(20), unflat_pos)
141            .unwrap();
142        assert!(!schema.get_group(unflat_pos).expect("unflat").is_flat());
143        assert!(schema.has_unflat_group());
144        assert_eq!(schema.unflat_group_pos(), Some(unflat_pos));
145        schema.validate_at_most_one_unflat().unwrap();
146        schema.flatten_group(unflat_pos).unwrap();
147        assert!(schema.is_flat_schema());
148        assert!(!schema.has_unflat_group());
149    }
150
151    #[test]
152    fn schema_at_most_one_unflat_invariant() {
153        let mut schema = FactorizedSchema::new();
154        let g0 = schema.create_group();
155        let g1 = schema.create_group();
156        schema.insert_to_group_and_scope(expr(1), g0).unwrap();
157        schema.insert_to_group_and_scope(expr(2), g1).unwrap();
158        assert!(schema.validate_at_most_one_unflat().is_err());
159    }
160
161    #[test]
162    fn flatten_group_validates_invariant_at_runtime() {
163        let mut schema = FactorizedSchema::new();
164        let flat_pos = schema.create_flat_group(false);
165        let g0 = schema.create_group();
166        let g1 = schema.create_group();
167        schema.insert_to_group_and_scope(expr(1), flat_pos).unwrap();
168        schema.insert_to_group_and_scope(expr(2), g0).unwrap();
169        schema.insert_to_group_and_scope(expr(3), g1).unwrap();
170        assert!(schema.flatten_group(flat_pos).is_err());
171    }
172
173    #[test]
174    fn schema_copy_and_flat_copy() {
175        let mut schema = FactorizedSchema::new();
176        let g0 = schema.create_flat_group(false);
177        let g1 = schema.create_group();
178        schema.insert_to_group_and_scope(expr(1), g0).unwrap();
179        schema.insert_to_group_and_scope(expr(2), g1).unwrap();
180        let flat = schema.flat_copy().unwrap();
181        assert!(flat.is_flat_schema());
182        assert!(!schema.is_flat_schema());
183        let copied = schema.copy();
184        assert_eq!(copied.num_groups(), 2);
185    }
186
187    #[test]
188    fn schema_utils_leading_group() {
189        let mut schema = FactorizedSchema::new();
190        let flat = schema.create_flat_group(false);
191        let unflat = schema.create_group();
192        schema.insert_to_group_and_scope(expr(1), flat).unwrap();
193        schema.insert_to_group_and_scope(expr(2), unflat).unwrap();
194        let mut set = HashSet::new();
195        set.insert(flat);
196        set.insert(unflat);
197        let leading = SchemaUtils::get_leading_group_pos(&set, &schema).unwrap();
198        assert_eq!(leading, unflat);
199        let mut flat_only = HashSet::new();
200        flat_only.insert(flat);
201        let leading2 = SchemaUtils::get_leading_group_pos(&flat_only, &schema).unwrap();
202        assert_eq!(leading2, flat);
203    }
204
205    #[test]
206    fn extend_schema_simulation() {
207        let mut scan_schema = FactorizedSchema::new();
208        let g0 = scan_schema.create_flat_group(false);
209        scan_schema
210            .insert_to_group_and_scope(expr(100), g0)
211            .unwrap();
212        scan_schema
213            .insert_to_group_and_scope(expr(101), g0)
214            .unwrap();
215
216        let mut extend_schema = scan_schema.copy();
217        let g1 = extend_schema.create_group();
218        extend_schema
219            .insert_to_group_and_scope(expr(200), g1)
220            .unwrap();
221        extend_schema
222            .insert_to_group_and_scope(expr(201), g1)
223            .unwrap();
224        assert_eq!(extend_schema.num_groups(), 2);
225        assert!(extend_schema.get_group(g0).expect("g0").is_flat());
226        assert!(!extend_schema.get_group(g1).expect("g1").is_flat());
227        extend_schema.validate_at_most_one_unflat().unwrap();
228    }
229
230    #[test]
231    fn member_names_reports_sorted_aliases_per_group() {
232        let mut schema = FactorizedSchema::new();
233        let g0 = schema.create_flat_group(false);
234        let g1 = schema.create_group();
235        schema.insert_to_group_and_scope(expr(1), g0).unwrap();
236        schema
237            .insert_to_group_and_scope_with_name(expr(2), Some("b".to_string()), g1)
238            .unwrap();
239        schema
240            .insert_to_group_and_scope_with_name(expr(3), Some("a".to_string()), g1)
241            .unwrap();
242        assert_eq!(
243            schema.member_names(g1),
244            vec!["a".to_string(), "b".to_string()]
245        );
246        assert!(schema.member_names(g0).is_empty());
247        assert!(schema.member_names(99).is_empty());
248    }
249
250    #[test]
251    fn hash_join_merge_schema() {
252        let mut left = FactorizedSchema::new();
253        let lg = left.create_flat_group(false);
254        left.insert_to_group_and_scope(expr(1), lg).unwrap();
255
256        let mut right = FactorizedSchema::new();
257        let rg = right.create_group();
258        right.insert_to_group_and_scope(expr(2), rg).unwrap();
259
260        let mut merged = left.copy();
261        let mapping = merged.merge_groups_from(&right);
262        assert_eq!(mapping.get(&0), Some(&1));
263        assert_eq!(merged.num_groups(), 2);
264    }
265
266    #[test]
267    fn merge_groups_from_keeps_expression_id_lookup() {
268        let mut left = FactorizedSchema::new();
269        let lg = left.create_flat_group(false);
270        left.insert_to_group_and_scope(expr(1), lg).unwrap();
271
272        let mut right = FactorizedSchema::new();
273        let rg = right.create_group();
274        right
275            .insert_to_group_and_scope_with_name(expr(2), Some("b".to_string()), rg)
276            .unwrap();
277
278        let mut merged = left.copy();
279        merged.merge_groups_from(&right);
280        assert_eq!(merged.get_group_pos(&expr(1)), Some(0));
281        assert_eq!(merged.get_group_pos(&expr(2)), Some(1));
282        assert!(merged.is_expression_in_scope(&expr(2)));
283        assert!(merged.expressions_in_scope().contains(&expr(2)));
284        assert_eq!(merged.get_group_pos_by_name_opt("b"), Some(1));
285        assert_eq!(merged.expression_name(&expr(2)), Some("b"));
286    }
287
288    #[test]
289    fn aggregate_flattens_all_but_one() {
290        let mut schema = FactorizedSchema::new();
291        let g0 = schema.create_flat_group(false);
292        let g1 = schema.create_group();
293        schema.insert_to_group_and_scope(expr(1), g0).unwrap();
294        schema.insert_to_group_and_scope(expr(2), g1).unwrap();
295        let mut agg_schema = FactorizedSchema::new();
296        let out = agg_schema.create_flat_group(false);
297        agg_schema.insert_to_group_and_scope(expr(10), out).unwrap();
298        assert!(agg_schema.is_flat_schema());
299    }
300
301    #[test]
302    fn expression_id_to_name_roundtrip() {
303        let mut schema = FactorizedSchema::new();
304        let g = schema.create_flat_group(false);
305        schema
306            .insert_to_group_and_scope_with_name(expr(1), Some("a".to_string()), g)
307            .unwrap();
308        assert_eq!(schema.expression_name(&expr(1)), Some("a"));
309        assert_eq!(schema.expression_name(&expr(2)), None);
310        let mut other = FactorizedSchema::new();
311        other.merge_groups_from(&schema);
312        other.insert_to_scope_may_repeat(expr(1), 0).unwrap();
313        assert_eq!(other.expression_name(&expr(1)), Some("a"));
314        other.clear_expressions_in_scope();
315        assert_eq!(other.expression_name(&expr(1)), None);
316    }
317
318    #[test]
319    fn sink_merge_preserves_multiplier_and_single_state() {
320        let mut input = FactorizedSchema::new();
321        let flat_pos = input.create_flat_group(false);
322        input
323            .insert_to_group_and_scope_with_name(expr(1), Some("a".to_string()), flat_pos)
324            .unwrap();
325        let unflat_pos = input.create_group();
326        input
327            .insert_to_group_and_scope_with_name(expr(2), Some("b".to_string()), unflat_pos)
328            .unwrap();
329        input
330            .get_group_mut(unflat_pos)
331            .expect("unflat group")
332            .set_multiplier(2.5);
333        let mut out = FactorizedSchema::new();
334        SinkOperatorUtil::recompute_schema(&input, &[expr(1), expr(2)], &mut out).unwrap();
335        out.validate_at_most_one_unflat().unwrap();
336        assert_eq!(out.num_groups(), 2);
337        let a_pos = out.get_group_pos(&expr(1)).expect("a placed");
338        let a_group = out.get_group(a_pos).expect("a group");
339        assert!(a_group.is_flat());
340        assert!(a_group.is_single_state());
341        let b_pos = out.get_group_pos(&expr(2)).expect("b placed");
342        assert_ne!(a_pos, b_pos);
343        let b_group = out.get_group(b_pos).expect("b group");
344        assert!(!b_group.is_flat());
345        assert_eq!(b_group.cardinality_multiplier(), 2.5);
346        assert_eq!(out.get_group_pos_by_name_opt("a"), Some(a_pos));
347        assert_eq!(out.get_group_pos_by_name_opt("b"), Some(b_pos));
348    }
349
350    #[test]
351    fn sink_recompute_all_flat_input() {
352        let mut input = FactorizedSchema::new();
353        let g0 = input.create_flat_group(false);
354        input
355            .insert_to_group_and_scope_with_name(expr(1), Some("a".to_string()), g0)
356            .unwrap();
357        let g1 = input.create_flat_group(false);
358        input
359            .insert_to_group_and_scope_with_name(expr(2), Some("b".to_string()), g1)
360            .unwrap();
361        let mut out = FactorizedSchema::new();
362        SinkOperatorUtil::recompute_schema(&input, &[expr(1), expr(2)], &mut out).unwrap();
363        assert_eq!(out.num_groups(), 1);
364        assert!(out.is_expression_in_scope(&expr(1)));
365        assert!(out.is_expression_in_scope(&expr(2)));
366        assert!(out.get_group_pos_by_name_opt("a").is_some());
367        assert!(out.get_group_pos_by_name_opt("b").is_some());
368        out.validate_at_most_one_unflat().unwrap();
369    }
370
371    #[test]
372    fn sink_merge_partial_scope_leaves_unmerged_names() {
373        let mut input = FactorizedSchema::new();
374        let flat_pos = input.create_flat_group(false);
375        input
376            .insert_to_group_and_scope_with_name(expr(1), Some("a".to_string()), flat_pos)
377            .unwrap();
378        let unflat_pos = input.create_group();
379        input
380            .insert_to_group_and_scope_with_name(expr(2), Some("b".to_string()), unflat_pos)
381            .unwrap();
382        input
383            .insert_name_for_group("bare".into(), unflat_pos)
384            .unwrap();
385        let mut out = FactorizedSchema::new();
386        SinkOperatorUtil::recompute_schema(&input, &[expr(1)], &mut out).unwrap();
387        assert!(out.is_expression_in_scope(&expr(1)));
388        assert!(!out.is_expression_in_scope(&expr(2)));
389        assert!(out.get_group_pos_by_name_opt("a").is_some());
390        assert!(out.get_group_pos_by_name_opt("b").is_none());
391        assert!(out.get_group_pos_by_name_opt("bare").is_none());
392    }
393
394    #[test]
395    fn scope_and_bare_name_registration_reject_out_of_range() {
396        let mut schema = FactorizedSchema::new();
397        let err = schema
398            .insert_to_scope(expr(1), 99)
399            .expect_err("scope OOR must fail");
400        assert_eq!(err, FactorizationError::GroupPosOutOfRange(99));
401        let err = schema
402            .insert_to_scope_with_name(expr(1), "a".to_string(), 99)
403            .expect_err("named scope OOR must fail");
404        assert_eq!(err, FactorizationError::GroupPosOutOfRange(99));
405        let err = schema
406            .insert_name_for_group("bare".into(), 99)
407            .expect_err("bare name OOR must fail");
408        assert_eq!(err, FactorizationError::GroupPosOutOfRange(99));
409    }
410
411    #[test]
412    fn duplicate_scope_registration_distinguishes_mapping_from_scope() {
413        let mut schema = FactorizedSchema::new();
414        let g = schema.create_flat_group(false);
415        schema.insert_to_scope(expr(1), g).unwrap();
416        let err = schema
417            .insert_to_scope(expr(1), g)
418            .expect_err("repeat must fail");
419        assert_eq!(err, FactorizationError::ExpressionAlreadyMapped(expr(1)));
420        let mut other = FactorizedSchema::new();
421        let h = other.create_flat_group(false);
422        other.expressions_in_scope.insert(expr(2));
423        let err = other
424            .insert_to_scope(expr(2), h)
425            .expect_err("scope repeat must fail");
426        assert_eq!(err, FactorizationError::ExpressionAlreadyInScope(expr(2)));
427    }
428
429    #[test]
430    fn resolve_group_pos_prefers_id_over_bare_name() {
431        let mut schema = FactorizedSchema::new();
432        let g0 = schema.create_flat_group(false);
433        let g1 = schema.create_group();
434        schema
435            .insert_to_group_and_scope_with_name(expr(1), Some("shared".to_string()), g0)
436            .unwrap();
437        schema.insert_name_for_group("shared".into(), g1).unwrap();
438        assert_eq!(
439            schema.resolve_group_pos(Some(&expr(1)), Some("shared")),
440            Some(g0)
441        );
442        assert_eq!(schema.resolve_group_pos(None, Some("shared")), Some(g1));
443        assert_eq!(schema.resolve_group_pos(None, Some("ghost")), None);
444    }
445}
linkrs_query/planning/plan/factorization.rs

linkrs_query/planning/plan/
factorization.rs
