linkrs_query/planning/plan/factorization/
compute.rs

1use std::collections::{HashMap, HashSet};
2
3use super::{FGroupPos, FactorizationError, FactorizedSchema};
4
5/// Utilities for factorization invariants.
6pub struct SchemaUtils;
7
8impl SchemaUtils {
9    pub fn get_leading_group_pos(
10        group_positions: &HashSet<FGroupPos>,
11        schema: &FactorizedSchema,
12    ) -> Result<FGroupPos, FactorizationError> {
13        if group_positions.is_empty() {
14            return Err(FactorizationError::EmptyGroupPositions);
15        }
16        Self::validate_at_most_one_unflat(group_positions, schema)?;
17        for &pos in group_positions {
18            if let Some(g) = schema.get_group(pos) {
19                if !g.is_flat() {
20                    return Ok(pos);
21                }
22            }
23        }
24        group_positions
25            .iter()
26            .next()
27            .copied()
28            .ok_or(FactorizationError::NonEmptyGroupPositions)
29    }
30
31    pub fn validate_at_most_one_unflat(
32        group_positions: &HashSet<FGroupPos>,
33        schema: &FactorizedSchema,
34    ) -> Result<(), FactorizationError> {
35        let mut unflat = 0;
36        for &pos in group_positions {
37            if let Some(g) = schema.get_group(pos) {
38                if !g.is_flat() {
39                    unflat += 1;
40                }
41            }
42        }
43        if unflat > 1 {
44            return Err(FactorizationError::TooManyUnflatGroups(unflat));
45        }
46        Ok(())
47    }
48
49    pub fn validate_no_unflat(
50        group_positions: &HashSet<FGroupPos>,
51        schema: &FactorizedSchema,
52    ) -> Result<(), FactorizationError> {
53        for &pos in group_positions {
54            if let Some(g) = schema.get_group(pos) {
55                if !g.is_flat() {
56                    return Err(FactorizationError::GroupExpectedFlat(pos));
57                }
58            }
59        }
60        Ok(())
61    }
62}
63
64/// Rebuilds factorized schemas across sink boundaries.
65///
66/// A sink operator collects its whole input before producing output, so the
67/// output groups no longer share the input's nesting structure. Flat
68/// payloads are gathered into one new group (marked single-state when unflat
69/// payloads also exist); each contributing unflat input group is copied into
70/// its own new group with the cardinality multiplier preserved.
71///
72/// Mirrors `lbug::planner::SinkOperatorUtil` in
73/// `ref/ladybug/src/planner/operator/factorization/sink_util.cpp`.
74pub struct SinkOperatorUtil;
75
76impl SinkOperatorUtil {
77    pub fn merge_schema(
78        input_schema: &FactorizedSchema,
79        expressions_to_merge: &[linkrs_core::types::expr::ExpressionId],
80        result_schema: &mut FactorizedSchema,
81    ) -> Result<(), FactorizationError> {
82        let mut flat_payloads = Vec::new();
83        let mut unflat_per_group: HashMap<FGroupPos, Vec<linkrs_core::types::expr::ExpressionId>> =
84            HashMap::new();
85        for expr_id in expressions_to_merge {
86            let Some(pos) = input_schema.get_group_pos(expr_id) else {
87                continue;
88            };
89            let Some(group) = input_schema.get_group(pos) else {
90                continue;
91            };
92            if group.is_flat() {
93                flat_payloads.push(expr_id.clone());
94            } else {
95                unflat_per_group
96                    .entry(pos)
97                    .or_default()
98                    .push(expr_id.clone());
99            }
100        }
101        let mut unflat_groups: Vec<FGroupPos> = unflat_per_group.keys().copied().collect();
102        unflat_groups.sort_unstable();
103        let mut old_to_new: HashMap<FGroupPos, FGroupPos> = HashMap::new();
104        let mut flat_new_pos: Option<FGroupPos> = None;
105        if unflat_groups.is_empty() {
106            if !flat_payloads.is_empty() {
107                let new_pos = result_schema.create_group();
108                for expr_id in &flat_payloads {
109                    result_schema.insert_to_group_and_scope(expr_id.clone(), new_pos)?;
110                }
111                flat_new_pos = Some(new_pos);
112            }
113        } else {
114            if !flat_payloads.is_empty() {
115                let new_pos = result_schema.create_group();
116                for expr_id in &flat_payloads {
117                    result_schema.insert_to_group_and_scope(expr_id.clone(), new_pos)?;
118                }
119                result_schema.set_group_as_single_state(new_pos)?;
120                flat_new_pos = Some(new_pos);
121            }
122            for old_pos in unflat_groups {
123                let new_pos = result_schema.create_group();
124                for expr_id in &unflat_per_group[&old_pos] {
125                    result_schema.insert_to_group_and_scope(expr_id.clone(), new_pos)?;
126                }
127                if let Some(input_group) = input_schema.get_group(old_pos) {
128                    let multiplier = input_group.cardinality_multiplier();
129                    if let Some(new_group) = result_schema.get_group_mut(new_pos) {
130                        new_group.set_multiplier(multiplier);
131                    }
132                }
133                old_to_new.insert(old_pos, new_pos);
134            }
135        }
136        if let Some(flat_new_pos) = flat_new_pos {
137            let mut flat_old: Vec<FGroupPos> = flat_payloads
138                .iter()
139                .filter_map(|eid| input_schema.get_group_pos(eid))
140                .collect();
141            flat_old.sort_unstable();
142            flat_old.dedup();
143            for old_pos in flat_old {
144                old_to_new.entry(old_pos).or_insert(flat_new_pos);
145            }
146        }
147        Self::remap_names(
148            input_schema,
149            expressions_to_merge,
150            &old_to_new,
151            result_schema,
152        )?;
153        result_schema.validate_at_most_one_unflat()?;
154        Ok(())
155    }
156
157    pub fn recompute_schema(
158        input_schema: &FactorizedSchema,
159        expressions_to_merge: &[linkrs_core::types::expr::ExpressionId],
160        result_schema: &mut FactorizedSchema,
161    ) -> Result<(), FactorizationError> {
162        result_schema.clear();
163        Self::merge_schema(input_schema, expressions_to_merge, result_schema)
164    }
165
166    fn remap_names(
167        input_schema: &FactorizedSchema,
168        expressions_to_merge: &[linkrs_core::types::expr::ExpressionId],
169        old_to_new: &HashMap<FGroupPos, FGroupPos>,
170        result_schema: &mut FactorizedSchema,
171    ) -> Result<(), FactorizationError> {
172        let mut owned_names: HashSet<&str> = HashSet::new();
173        for expr_id in expressions_to_merge {
174            if let Some(name) = input_schema.expression_name(expr_id) {
175                owned_names.insert(name);
176            }
177        }
178        let linked_names: HashSet<&str> = input_schema
179            .expression_id_to_name
180            .values()
181            .map(String::as_str)
182            .collect();
183        let mut names: Vec<(&String, &FGroupPos)> =
184            input_schema.expression_name_to_group_iter().collect();
185        names.sort_by(|a, b| a.0.cmp(b.0));
186        for (name, old_pos) in names {
187            let Some(new_pos) = old_to_new.get(old_pos).copied() else {
188                continue;
189            };
190            if !owned_names.contains(name.as_str()) && linked_names.contains(name.as_str()) {
191                continue;
192            }
193            if result_schema.get_group_pos_by_name_opt(name).is_none() {
194                result_schema.insert_name_for_group(name.clone(), new_pos)?;
195            }
196        }
197        Ok(())
198    }
199}
linkrs_query/planning/plan/factorization/compute.rs

linkrs_query/planning/plan/factorization/
compute.rs
