linkrs_query/planning/plan/factorization/
expression.rs

1use std::collections::{HashMap, HashSet};
2
3use super::{FGroupPos, FactorizationError, FactorizationGroup};
4use linkrs_core::types::expr::ExpressionId;
5
6/// Output schema with factorization structure.
7///
8/// Tracks flat/unflat groups and which expression belongs to which group.
9/// Enforces the invariant that at most one group is unflat at any time.
10#[derive(Debug, Clone, Default)]
11pub struct FactorizedSchema {
12    pub(super) groups: Vec<FactorizationGroup>,
13    pub(super) expression_to_group: HashMap<ExpressionId, FGroupPos>,
14    pub(super) expression_name_to_group: HashMap<String, FGroupPos>,
15    pub(super) expression_id_to_name: HashMap<ExpressionId, String>,
16    pub(super) expressions_in_scope: HashSet<ExpressionId>,
17}
18
19impl FactorizedSchema {
20    pub fn new() -> Self {
21        Self::default()
22    }
23
24    pub fn num_groups(&self) -> usize {
25        self.groups.len()
26    }
27
28    pub fn groups(&self) -> &[FactorizationGroup] {
29        &self.groups
30    }
31
32    pub fn groups_mut(&mut self) -> &mut Vec<FactorizationGroup> {
33        &mut self.groups
34    }
35
36    pub fn get_group(&self, pos: FGroupPos) -> Option<&FactorizationGroup> {
37        self.groups.get(pos as usize)
38    }
39
40    pub fn get_group_mut(&mut self, pos: FGroupPos) -> Option<&mut FactorizationGroup> {
41        self.groups.get_mut(pos as usize)
42    }
43
44    pub fn get_group_by_expression(&self, expr_id: &ExpressionId) -> Option<&FactorizationGroup> {
45        let pos = self.get_group_pos(expr_id)?;
46        self.get_group(pos)
47    }
48
49    #[cfg(test)]
50    pub fn get_group_by_name(&self, name: &str) -> Option<&FactorizationGroup> {
51        let pos = self.get_group_pos_by_name(name)?;
52        self.get_group(pos)
53    }
54
55    pub fn create_group(&mut self) -> FGroupPos {
56        let pos = self.groups.len() as FGroupPos;
57        self.groups.push(FactorizationGroup::new());
58        pos
59    }
60
61    pub fn create_flat_group(&mut self, single_state: bool) -> FGroupPos {
62        let pos = self.groups.len() as FGroupPos;
63        self.groups.push(FactorizationGroup::new_flat(single_state));
64        pos
65    }
66
67    pub fn insert_to_scope(
68        &mut self,
69        expr_id: ExpressionId,
70        group_pos: FGroupPos,
71    ) -> Result<(), FactorizationError> {
72        if (group_pos as usize) >= self.groups.len() {
73            return Err(FactorizationError::GroupPosOutOfRange(group_pos));
74        }
75        if self.expression_to_group.contains_key(&expr_id) {
76            return Err(FactorizationError::ExpressionAlreadyMapped(expr_id));
77        }
78        if self.expressions_in_scope.contains(&expr_id) {
79            return Err(FactorizationError::ExpressionAlreadyInScope(expr_id));
80        }
81        self.expression_to_group.insert(expr_id.clone(), group_pos);
82        self.expressions_in_scope.insert(expr_id);
83        Ok(())
84    }
85
86    pub fn insert_to_scope_with_name(
87        &mut self,
88        expr_id: ExpressionId,
89        name: String,
90        group_pos: FGroupPos,
91    ) -> Result<(), FactorizationError> {
92        if (group_pos as usize) >= self.groups.len() {
93            return Err(FactorizationError::GroupPosOutOfRange(group_pos));
94        }
95        self.expression_name_to_group
96            .insert(name.clone(), group_pos);
97        self.expression_id_to_name.insert(expr_id.clone(), name);
98        self.insert_to_scope(expr_id, group_pos)?;
99        Ok(())
100    }
101
102    pub fn insert_to_group_and_scope(
103        &mut self,
104        expr_id: ExpressionId,
105        group_pos: FGroupPos,
106    ) -> Result<(), FactorizationError> {
107        self.insert_to_group_and_scope_with_name(expr_id, None, group_pos)
108    }
109
110    pub fn insert_to_group_and_scope_with_name(
111        &mut self,
112        expr_id: ExpressionId,
113        name: Option<String>,
114        group_pos: FGroupPos,
115    ) -> Result<(), FactorizationError> {
116        if (group_pos as usize) >= self.groups.len() {
117            return Err(FactorizationError::GroupPosOutOfRange(group_pos));
118        }
119        if self.expression_to_group.contains_key(&expr_id) {
120            return Err(FactorizationError::ExpressionAlreadyMapped(expr_id));
121        }
122        if self.expressions_in_scope.contains(&expr_id) {
123            return Err(FactorizationError::ExpressionAlreadyInScope(expr_id));
124        }
125        let group = &mut self.groups[group_pos as usize];
126        group.insert_expression_with_name(expr_id.clone(), name.clone())?;
127        if let Some(n) = name {
128            self.expression_name_to_group.insert(n.clone(), group_pos);
129            self.expression_id_to_name.insert(expr_id.clone(), n);
130        }
131        self.expression_to_group.insert(expr_id.clone(), group_pos);
132        self.expressions_in_scope.insert(expr_id);
133        Ok(())
134    }
135
136    pub fn insert_to_group_and_scope_batch(
137        &mut self,
138        exprs: Vec<ExpressionId>,
139        group_pos: FGroupPos,
140    ) -> Result<(), FactorizationError> {
141        for e in exprs {
142            self.insert_to_group_and_scope(e, group_pos)?;
143        }
144        Ok(())
145    }
146
147    pub fn insert_to_scope_may_repeat(
148        &mut self,
149        expr_id: ExpressionId,
150        group_pos: FGroupPos,
151    ) -> Result<(), FactorizationError> {
152        if (group_pos as usize) >= self.groups.len() {
153            return Err(FactorizationError::GroupPosOutOfRange(group_pos));
154        }
155        self.expression_to_group.insert(expr_id.clone(), group_pos);
156        self.expressions_in_scope.insert(expr_id);
157        Ok(())
158    }
159
160    pub fn insert_to_group_and_scope_may_repeat(
161        &mut self,
162        expr_id: ExpressionId,
163        group_pos: FGroupPos,
164    ) -> Result<(), FactorizationError> {
165        if (group_pos as usize) >= self.groups.len() {
166            return Err(FactorizationError::GroupPosOutOfRange(group_pos));
167        }
168        let group = &mut self.groups[group_pos as usize];
169        if !group.contains(&expr_id) {
170            group.insert_expression(expr_id.clone())?;
171        }
172        self.expression_to_group.insert(expr_id.clone(), group_pos);
173        self.expressions_in_scope.insert(expr_id);
174        Ok(())
175    }
176
177    pub fn get_group_pos(&self, expr_id: &ExpressionId) -> Option<FGroupPos> {
178        self.expression_to_group.get(expr_id).copied()
179    }
180
181    #[cfg(test)]
182    pub fn get_group_pos_by_name(&self, name: &str) -> Option<FGroupPos> {
183        self.get_group_pos_by_name_opt(name)
184    }
185
186    pub fn get_group_pos_by_name_opt(&self, name: &str) -> Option<FGroupPos> {
187        self.expression_name_to_group.get(name).copied()
188    }
189
190    pub fn expression_name(&self, expr_id: &ExpressionId) -> Option<&str> {
191        self.expression_id_to_name.get(expr_id).map(String::as_str)
192    }
193
194    pub fn insert_name_for_group(
195        &mut self,
196        name: String,
197        group_pos: FGroupPos,
198    ) -> Result<(), FactorizationError> {
199        if (group_pos as usize) >= self.groups.len() {
200            return Err(FactorizationError::GroupPosOutOfRange(group_pos));
201        }
202        self.expression_name_to_group.insert(name, group_pos);
203        Ok(())
204    }
205
206    pub fn resolve_group_pos(
207        &self,
208        expr_id: Option<&ExpressionId>,
209        name: Option<&str>,
210    ) -> Option<FGroupPos> {
211        if let Some(id) = expr_id {
212            if let Some(pos) = self.get_group_pos(id) {
213                return Some(pos);
214            }
215            if let Some(owned) = self.expression_name(id) {
216                let owned = owned.to_string();
217                if let Some(pos) = self.get_group_pos_by_name_opt(&owned) {
218                    return Some(pos);
219                }
220            }
221        }
222        if let Some(n) = name {
223            if let Some(pos) = self.get_group_pos_by_name_opt(n) {
224                return Some(pos);
225            }
226        }
227        None
228    }
229
230    pub fn member_names(&self, group_pos: FGroupPos) -> Vec<String> {
231        let mut names: Vec<String> = self
232            .expression_name_to_group
233            .iter()
234            .filter(|(_, pos)| **pos == group_pos)
235            .map(|(name, _)| name.clone())
236            .collect();
237        names.sort();
238        names
239    }
240
241    pub fn get_expression_pos(&self, expr_id: &ExpressionId) -> Option<(FGroupPos, usize)> {
242        let gpos = self.get_group_pos(expr_id)?;
243        let group = self.get_group(gpos)?;
244        let pos = group.get_expression_pos(expr_id)?;
245        Some((gpos, pos))
246    }
247
248    pub fn flatten_group(&mut self, pos: FGroupPos) -> Result<(), FactorizationError> {
249        let group = self
250            .get_group_mut(pos)
251            .ok_or(FactorizationError::InvalidFlattenPos(pos))?;
252        if !group.is_flat() {
253            group.set_flat()?;
254        }
255        self.validate_at_most_one_unflat()
256    }
257
258    pub fn flatten_all(&mut self) -> Result<(), FactorizationError> {
259        for i in 0..self.groups.len() {
260            let pos = i as FGroupPos;
261            if let Some(g) = self.get_group(pos) {
262                if !g.is_flat() {
263                    self.flatten_group(pos)?;
264                }
265            }
266        }
267        Ok(())
268    }
269
270    pub fn set_group_as_single_state(&mut self, pos: FGroupPos) -> Result<(), FactorizationError> {
271        let group = self
272            .get_group_mut(pos)
273            .ok_or(FactorizationError::InvalidSingleStatePos(pos))?;
274        if !group.is_single_state() {
275            group.set_single_state()?;
276        }
277        Ok(())
278    }
279
280    pub fn is_expression_in_scope(&self, expr_id: &ExpressionId) -> bool {
281        self.expression_to_group.contains_key(expr_id)
282    }
283
284    #[cfg(test)]
285    pub fn is_name_in_scope(&self, name: &str) -> bool {
286        self.expression_name_to_group.contains_key(name)
287    }
288
289    pub fn expressions_in_scope(&self) -> &HashSet<ExpressionId> {
290        &self.expressions_in_scope
291    }
292
293    pub fn expressions_in_scope_for_group(&self, pos: FGroupPos) -> Vec<ExpressionId> {
294        let group = match self.get_group(pos) {
295            Some(g) => g,
296            None => return Vec::new(),
297        };
298        group
299            .expressions()
300            .iter()
301            .filter(|e| self.expressions_in_scope.contains(e))
302            .cloned()
303            .collect()
304    }
305
306    pub fn evaluable(&self, expr_id: &ExpressionId) -> bool {
307        self.is_expression_in_scope(expr_id)
308    }
309
310    pub fn clear_expressions_in_scope(&mut self) {
311        self.expression_to_group.clear();
312        self.expression_name_to_group.clear();
313        self.expression_id_to_name.clear();
314        self.expressions_in_scope.clear();
315    }
316
317    pub fn groups_pos_in_scope(&self) -> HashSet<FGroupPos> {
318        self.expression_to_group.values().copied().collect()
319    }
320
321    pub fn copy(&self) -> Self {
322        self.clone()
323    }
324
325    pub fn clear(&mut self) {
326        self.groups.clear();
327        self.clear_expressions_in_scope();
328    }
329
330    pub fn has_unflat_group(&self) -> bool {
331        self.groups.iter().any(|g| !g.is_flat())
332    }
333
334    pub fn unflat_group_pos(&self) -> Option<FGroupPos> {
335        self.groups
336            .iter()
337            .enumerate()
338            .find(|(_, g)| !g.is_flat())
339            .map(|(i, _)| i as FGroupPos)
340    }
341
342    pub fn validate_at_most_one_unflat(&self) -> Result<(), FactorizationError> {
343        let unflat = self.groups.iter().filter(|g| !g.is_flat()).count();
344        if unflat > 1 {
345            return Err(FactorizationError::TooManyUnflatGroups(unflat));
346        }
347        Ok(())
348    }
349
350    pub fn has_at_most_one_unflat(&self) -> bool {
351        self.groups.iter().filter(|g| !g.is_flat()).count() <= 1
352    }
353
354    pub fn is_flat_schema(&self) -> bool {
355        self.groups.iter().all(|g| g.is_flat())
356    }
357
358    pub fn flat_copy(&self) -> Result<Self, FactorizationError> {
359        let mut copy = self.clone();
360        copy.flatten_all()?;
361        Ok(copy)
362    }
363
364    pub fn merge_groups_from(&mut self, other: &FactorizedSchema) -> HashMap<FGroupPos, FGroupPos> {
365        let mut mapping = HashMap::new();
366        for (idx, group) in other.groups.iter().enumerate() {
367            let old_pos = idx as FGroupPos;
368            let new_pos = self.groups.len() as FGroupPos;
369            self.groups.push(group.clone());
370            mapping.insert(old_pos, new_pos);
371        }
372        for (name, pos) in &other.expression_name_to_group {
373            if let Some(new_pos) = mapping.get(pos) {
374                self.expression_name_to_group.insert(name.clone(), *new_pos);
375            }
376        }
377        for (expr_id, name) in &other.expression_id_to_name {
378            self.expression_id_to_name
379                .insert(expr_id.clone(), name.clone());
380        }
381        for (expr_id, pos) in &other.expression_to_group {
382            if let Some(new_pos) = mapping.get(pos) {
383                self.expression_to_group.insert(expr_id.clone(), *new_pos);
384            }
385        }
386        for expr_id in &other.expressions_in_scope {
387            if self.expression_to_group.contains_key(expr_id) {
388                self.expressions_in_scope.insert(expr_id.clone());
389            }
390        }
391        mapping
392    }
393
394    pub fn expression_to_group_iter(&self) -> impl Iterator<Item = (&ExpressionId, &FGroupPos)> {
395        self.expression_to_group.iter()
396    }
397
398    pub fn expression_name_to_group_iter(&self) -> impl Iterator<Item = (&String, &FGroupPos)> {
399        self.expression_name_to_group.iter()
400    }
401}
linkrs_query/planning/plan/factorization/expression.rs

linkrs_query/planning/plan/factorization/
expression.rs
