Skip to main content

linkrs_query/planning/plan/factorization/
group.rs

1use std::collections::HashMap;
2
3use super::FactorizationError;
4use linkrs_core::types::expr::ExpressionId;
5
6/// A group of expressions sharing the same nesting level.
7///
8/// Mirrors `lbug::planner::FactorizationGroup` in
9/// `ref/ladybug/src/include/planner/operator/schema.h`.
10#[derive(Debug, Clone)]
11pub struct FactorizationGroup {
12    flat: bool,
13    single_state: bool,
14    cardinality_multiplier: f64,
15    expressions: Vec<ExpressionId>,
16    expression_id_to_pos: HashMap<ExpressionId, usize>,
17    expression_name_to_pos: HashMap<String, usize>,
18}
19
20impl FactorizationGroup {
21    pub fn new() -> Self {
22        Self {
23            flat: false,
24            single_state: false,
25            cardinality_multiplier: 1.0,
26            expressions: Vec::new(),
27            expression_id_to_pos: HashMap::new(),
28            expression_name_to_pos: HashMap::new(),
29        }
30    }
31
32    pub fn new_flat(single_state: bool) -> Self {
33        Self {
34            flat: true,
35            single_state,
36            cardinality_multiplier: 1.0,
37            expressions: Vec::new(),
38            expression_id_to_pos: HashMap::new(),
39            expression_name_to_pos: HashMap::new(),
40        }
41    }
42
43    pub fn is_flat(&self) -> bool {
44        self.flat
45    }
46
47    pub fn is_single_state(&self) -> bool {
48        self.single_state
49    }
50
51    pub fn set_flat(&mut self) -> Result<(), FactorizationError> {
52        if self.flat {
53            return Err(FactorizationError::GroupAlreadyFlat);
54        }
55        self.flat = true;
56        Ok(())
57    }
58
59    pub fn set_single_state(&mut self) -> Result<(), FactorizationError> {
60        if self.single_state {
61            return Err(FactorizationError::GroupAlreadySingleState);
62        }
63        self.single_state = true;
64        if !self.flat {
65            self.flat = true;
66        }
67        Ok(())
68    }
69
70    pub fn cardinality_multiplier(&self) -> f64 {
71        self.cardinality_multiplier
72    }
73
74    pub fn set_multiplier(&mut self, multiplier: f64) {
75        self.cardinality_multiplier = multiplier;
76    }
77
78    pub fn expressions(&self) -> &[ExpressionId] {
79        &self.expressions
80    }
81
82    pub fn len(&self) -> usize {
83        self.expressions.len()
84    }
85
86    pub fn is_empty(&self) -> bool {
87        self.expressions.is_empty()
88    }
89
90    pub fn insert_expression(&mut self, expr_id: ExpressionId) -> Result<(), FactorizationError> {
91        self.insert_expression_with_name(expr_id, None)
92    }
93
94    pub fn insert_expression_with_name(
95        &mut self,
96        expr_id: ExpressionId,
97        name: Option<String>,
98    ) -> Result<(), FactorizationError> {
99        if self.expression_id_to_pos.contains_key(&expr_id) {
100            return Err(FactorizationError::DuplicateExpressionId(expr_id));
101        }
102        if let Some(n) = name {
103            if self.expression_name_to_pos.contains_key(&n) {
104                return Err(FactorizationError::DuplicateExpressionName(n));
105            }
106            self.expression_name_to_pos
107                .insert(n, self.expressions.len());
108        }
109        self.expression_id_to_pos
110            .insert(expr_id.clone(), self.expressions.len());
111        self.expressions.push(expr_id);
112        Ok(())
113    }
114
115    pub fn get_expression_pos(&self, expr_id: &ExpressionId) -> Option<usize> {
116        self.expression_id_to_pos.get(expr_id).copied()
117    }
118
119    pub fn contains(&self, expr_id: &ExpressionId) -> bool {
120        self.expression_id_to_pos.contains_key(expr_id)
121    }
122
123    pub fn contains_name(&self, name: &str) -> bool {
124        self.expression_name_to_pos.contains_key(name)
125    }
126}
127
128impl Default for FactorizationGroup {
129    fn default() -> Self {
130        Self::new()
131    }
132}