linkrs_query/optimizer/factorization/
flatten_resolver.rs

1use std::collections::{HashMap, HashSet};
2
3use linkrs_core::types::expr::ExpressionId;
4
5use crate::planning::plan::factorization::{FGroupPos, FactorizedSchema, INVALID_F_GROUP_POS};
6
7use super::group_dependency_analyzer::GroupDependencyAnalyzer;
8
9/// FlattenAllButOne: flatten all dependent unflat groups except the leading one.
10///
11/// Mirrors `lbug::planner::FlattenAllButOne` in
12/// `ref/ladybug/src/planner/operator/factorization/flatten_resolver.cpp`.
13pub struct FlattenAllButOne;
14
15impl FlattenAllButOne {
16    /// For a set of expressions: collect dependent groups, keep first unflat as leading, flatten rest.
17    pub fn get_groups_pos_to_flatten_for_exprs(
18        exprs: &[ExpressionId],
19        schema: &FactorizedSchema,
20        expr_store: &HashMap<ExpressionId, linkrs_core::Expression>,
21    ) -> HashSet<FGroupPos> {
22        let mut result = HashSet::new();
23        let mut dependent_groups = HashSet::new();
24
25        for expr_id in exprs {
26            let mut analyzer = GroupDependencyAnalyzer::with_expr_store(schema, false, expr_store);
27            analyzer.visit(expr_id);
28            result.extend(analyzer.required_flat_groups().iter().copied());
29            dependent_groups.extend(analyzer.dependent_groups().iter().copied());
30        }
31
32        let mut candidates: Vec<FGroupPos> = dependent_groups
33            .into_iter()
34            .filter(|pos| {
35                schema
36                    .get_group(*pos)
37                    .map(|g| !g.is_flat() && !result.contains(pos))
38                    .unwrap_or(false)
39            })
40            .collect();
41        // Deterministic order
42        candidates.sort_unstable();
43        for pos in candidates.iter().skip(1) {
44            result.insert(*pos);
45        }
46        result
47    }
48
49    pub fn get_groups_pos_to_flatten_for_exprs_with_leading(
50        exprs: &[ExpressionId],
51        schema: &FactorizedSchema,
52        expr_store: &HashMap<ExpressionId, linkrs_core::Expression>,
53    ) -> (FGroupPos, HashSet<FGroupPos>) {
54        let mut result = HashSet::new();
55        let mut dependent_groups = HashSet::new();
56
57        for expr_id in exprs {
58            let mut analyzer = GroupDependencyAnalyzer::with_expr_store(schema, false, expr_store);
59            analyzer.visit(expr_id);
60            result.extend(analyzer.required_flat_groups().iter().copied());
61            dependent_groups.extend(analyzer.dependent_groups().iter().copied());
62        }
63
64        let mut candidates: Vec<FGroupPos> = dependent_groups
65            .iter()
66            .filter(|pos| {
67                schema
68                    .get_group(**pos)
69                    .map(|g| !g.is_flat() && !result.contains(*pos))
70                    .unwrap_or(false)
71            })
72            .copied()
73            .collect();
74        candidates.sort_unstable();
75
76        for pos in candidates.iter().skip(1) {
77            result.insert(*pos);
78        }
79
80        if candidates.is_empty() {
81            (INVALID_F_GROUP_POS, result)
82        } else {
83            (candidates[0], result)
84        }
85    }
86
87    /// Single expression version.
88    pub fn get_groups_pos_to_flatten_for_expr(
89        expr_id: &ExpressionId,
90        schema: &FactorizedSchema,
91        expr_store: &HashMap<ExpressionId, linkrs_core::Expression>,
92    ) -> HashSet<FGroupPos> {
93        Self::get_groups_pos_to_flatten_for_exprs(std::slice::from_ref(expr_id), schema, expr_store)
94    }
95
96    /// Group set version: flatten all but one.
97    pub fn get_groups_pos_to_flatten_for_groups(
98        dependent_groups: &HashSet<FGroupPos>,
99        schema: &FactorizedSchema,
100    ) -> HashSet<FGroupPos> {
101        let mut candidates: Vec<FGroupPos> = dependent_groups
102            .iter()
103            .filter(|pos| {
104                schema
105                    .get_group(**pos)
106                    .map(|g| !g.is_flat())
107                    .unwrap_or(false)
108            })
109            .copied()
110            .collect();
111        candidates.sort_unstable();
112        let mut result = HashSet::new();
113        for pos in candidates.iter().skip(1) {
114            result.insert(*pos);
115        }
116        result
117    }
118
119    /// Group set version with leading return.
120    pub fn get_groups_pos_to_flatten_for_groups_with_leading(
121        dependent_groups: &HashSet<FGroupPos>,
122        schema: &FactorizedSchema,
123    ) -> (FGroupPos, HashSet<FGroupPos>) {
124        let mut candidates: Vec<FGroupPos> = dependent_groups
125            .iter()
126            .filter(|pos| {
127                schema
128                    .get_group(**pos)
129                    .map(|g| !g.is_flat())
130                    .unwrap_or(false)
131            })
132            .copied()
133            .collect();
134        candidates.sort_unstable();
135        let mut result = HashSet::new();
136        for pos in candidates.iter().skip(1) {
137            result.insert(*pos);
138        }
139        if candidates.is_empty() {
140            (INVALID_F_GROUP_POS, result)
141        } else {
142            (candidates[0], result)
143        }
144    }
145}
146
147/// FlattenAll: flatten every dependent unflat group.
148pub struct FlattenAll;
149
150impl FlattenAll {
151    pub fn get_groups_pos_to_flatten_for_exprs(
152        exprs: &[ExpressionId],
153        schema: &FactorizedSchema,
154        expr_store: &HashMap<ExpressionId, linkrs_core::Expression>,
155    ) -> HashSet<FGroupPos> {
156        let mut result = HashSet::new();
157        for expr_id in exprs {
158            result.extend(Self::get_groups_pos_to_flatten_for_expr(
159                expr_id, schema, expr_store,
160            ));
161        }
162        result
163    }
164
165    pub fn get_groups_pos_to_flatten_for_expr(
166        expr_id: &ExpressionId,
167        schema: &FactorizedSchema,
168        expr_store: &HashMap<ExpressionId, linkrs_core::Expression>,
169    ) -> HashSet<FGroupPos> {
170        // Baseline `FlattenAll(expr)` flattens only dependent groups and
171        // ignores `required_flat` (see ladybug `flatten_resolver.cpp:92-97`).
172        // Callers that need lambda-body flatness use `FlattenAllButOne`,
173        // which merges `required_flat` explicitly.
174        let mut analyzer = GroupDependencyAnalyzer::with_expr_store(schema, false, expr_store);
175        analyzer.visit(expr_id);
176        Self::get_groups_pos_to_flatten_for_groups(analyzer.dependent_groups(), schema)
177    }
178
179    pub fn get_groups_pos_to_flatten_for_groups(
180        dependent_groups: &HashSet<FGroupPos>,
181        schema: &FactorizedSchema,
182    ) -> HashSet<FGroupPos> {
183        dependent_groups
184            .iter()
185            .filter(|pos| {
186                schema
187                    .get_group(**pos)
188                    .map(|g| !g.is_flat())
189                    .unwrap_or(false)
190            })
191            .copied()
192            .collect()
193    }
194}
195
196/// Shared two-stage aggregate flatten rule mirroring
197/// `LogicalAggregate::getGroupsPosToFlatten` (`logical_aggregate.cpp:32-57`).
198///
199/// Stage 1 runs `FlattenAllButOne` over the group keys and returns the
200/// surviving `leading` group. Stage 2 walks every aggregate payload: lambda
201/// bodies (`required_flat`) always flatten, while `distinct` payloads
202/// additionally flatten any dependent group other than `leading`. Plain
203/// payloads on the leading group stay factorized.
204pub fn aggregate_groups_to_flatten(
205    key_ids: &[ExpressionId],
206    key_store: &HashMap<ExpressionId, linkrs_core::Expression>,
207    aggregate_args: &[Vec<linkrs_core::Expression>],
208    aggregate_distinct: &[bool],
209    child_schema: &FactorizedSchema,
210) -> (FGroupPos, HashSet<FGroupPos>) {
211    let (leading, mut to_flatten) =
212        FlattenAllButOne::get_groups_pos_to_flatten_for_exprs_with_leading(
213            key_ids,
214            child_schema,
215            key_store,
216        );
217    if leading == INVALID_F_GROUP_POS {
218        return (leading, to_flatten);
219    }
220    for (idx, args) in aggregate_args.iter().enumerate() {
221        let is_distinct = aggregate_distinct.get(idx).copied().unwrap_or(false);
222        for payload in args {
223            let mut analyzer =
224                GroupDependencyAnalyzer::with_expr_store(child_schema, false, key_store);
225            analyzer.visit_expression(payload);
226            for pos in analyzer.required_flat_groups().iter().copied() {
227                if child_schema.get_group(pos).is_some_and(|g| !g.is_flat()) {
228                    to_flatten.insert(pos);
229                }
230            }
231            if is_distinct {
232                for pos in analyzer.dependent_groups().iter().copied() {
233                    if pos != leading && child_schema.get_group(pos).is_some_and(|g| !g.is_flat()) {
234                        to_flatten.insert(pos);
235                    }
236                }
237            }
238        }
239    }
240    (leading, to_flatten)
241}
242
243/// Legacy wrapper providing the Ladybug-style static API.
244pub struct FlattenResolver;
245
246impl FlattenResolver {
247    pub fn flatten_all_but_one(
248        dependent_groups: &[FGroupPos],
249        schema: &FactorizedSchema,
250    ) -> (FGroupPos, Vec<FGroupPos>) {
251        let set: HashSet<FGroupPos> = dependent_groups.iter().copied().collect();
252        let (leading, to_flatten) =
253            FlattenAllButOne::get_groups_pos_to_flatten_for_groups_with_leading(&set, schema);
254        let mut vec: Vec<FGroupPos> = to_flatten.into_iter().collect();
255        vec.sort_unstable();
256        (leading, vec)
257    }
258
259    pub fn flatten_all(
260        dependent_groups: &[FGroupPos],
261        schema: &FactorizedSchema,
262    ) -> Vec<FGroupPos> {
263        let set: HashSet<FGroupPos> = dependent_groups.iter().copied().collect();
264        let mut vec: Vec<FGroupPos> =
265            FlattenAll::get_groups_pos_to_flatten_for_groups(&set, schema)
266                .into_iter()
267                .collect();
268        vec.sort_unstable();
269        vec
270    }
271}
272
273#[cfg(test)]
274mod tests {
275    use super::*;
276    use crate::planning::plan::factorization::FactorizedSchema;
277    use linkrs_core::types::expr::ExpressionId;
278    use linkrs_core::types::expr::FunctionArg;
279    use std::collections::HashSet;
280
281    fn expr(id: u64) -> ExpressionId {
282        ExpressionId::new(id)
283    }
284
285    #[test]
286    fn flatten_all_single_unflat() {
287        let mut schema = FactorizedSchema::new();
288        let g0 = schema.create_flat_group(false);
289        let g1 = schema.create_group();
290        schema.insert_to_group_and_scope(expr(1), g0).unwrap();
291        schema.insert_to_group_and_scope(expr(2), g1).unwrap();
292
293        let mut set = HashSet::new();
294        set.insert(g1);
295        let res = FlattenAll::get_groups_pos_to_flatten_for_groups(&set, &schema);
296        assert_eq!(res, HashSet::from([g1]));
297
298        let res2 = FlattenAllButOne::get_groups_pos_to_flatten_for_groups(&set, &schema);
299        assert!(
300            res2.is_empty(),
301            "single unflat should not be flattened in AllButOne"
302        );
303    }
304
305    #[test]
306    fn flatten_all_but_one_two_unflats() {
307        let mut schema = FactorizedSchema::new();
308        let g0 = schema.create_group();
309        let g1 = schema.create_group();
310        schema.insert_to_group_and_scope(expr(1), g0).unwrap();
311        schema.insert_to_group_and_scope(expr(2), g1).unwrap();
312        let mut set = HashSet::new();
313        set.insert(g0);
314        set.insert(g1);
315        let res = FlattenAllButOne::get_groups_pos_to_flatten_for_groups(&set, &schema);
316        assert_eq!(res.len(), 1);
317        let (_leading, to_flatten) =
318            FlattenAllButOne::get_groups_pos_to_flatten_for_groups_with_leading(&set, &schema);
319        assert_eq!(to_flatten.len(), 1);
320    }
321
322    #[test]
323    fn flatten_resolver_api() {
324        let mut schema = FactorizedSchema::new();
325        let g0 = schema.create_group();
326        let g1 = schema.create_group();
327        schema.insert_to_group_and_scope(expr(10), g0).unwrap();
328        schema.insert_to_group_and_scope(expr(20), g1).unwrap();
329        let (leading, to_flatten) = FlattenResolver::flatten_all_but_one(&[g0, g1], &schema);
330        assert_ne!(leading, INVALID_F_GROUP_POS);
331        assert_eq!(to_flatten.len(), 1);
332        let all = FlattenResolver::flatten_all(&[g0, g1], &schema);
333        assert_eq!(all.len(), 2);
334    }
335
336    #[test]
337    fn flatten_all_expr_matches_dependent_only() {
338        use crate::optimizer::factorization::GroupDependencyAnalyzer;
339        let mut schema = FactorizedSchema::new();
340        let g0 = schema.create_flat_group(false);
341        let g1 = schema.create_group();
342        schema
343            .insert_to_group_and_scope_with_name(expr(10), Some("a".to_string()), g0)
344            .unwrap();
345        schema
346            .insert_to_group_and_scope_with_name(expr(20), Some("x".to_string()), g1)
347            .unwrap();
348        // Simulate list_extract(lambda) where lambda body depends on x (unflat).
349        // Baseline `FlattenAll(expr)` flattens only dependent groups and
350        // ignores `required_flat`; `AllButOne` merges `required_flat`.
351        let the_expr = linkrs_core::Expression::Function {
352            name: "list_extract".to_string(),
353            args: vec![
354                FunctionArg::Positional(linkrs_core::Expression::List(vec![
355                    linkrs_core::Expression::Literal(linkrs_core::Value::BigInt(1)),
356                ])),
357                FunctionArg::Positional(linkrs_core::Expression::Variable("x".to_string())),
358            ],
359        };
360        let mut store = HashMap::new();
361        let fake_id = expr(999);
362        store.insert(fake_id.clone(), the_expr);
363        let mut analyzer = GroupDependencyAnalyzer::with_expr_store(&schema, false, &store);
364        analyzer.visit(&fake_id);
365        let expected =
366            FlattenAll::get_groups_pos_to_flatten_for_groups(analyzer.dependent_groups(), &schema);
367        let all = FlattenAll::get_groups_pos_to_flatten_for_expr(&fake_id, &schema, &store);
368        assert_eq!(all, expected, "FlattenAll(expr) must equal dependent-only");
369        assert!(all.contains(&g1));
370        let but_one =
371            FlattenAllButOne::get_groups_pos_to_flatten_for_expr(&fake_id, &schema, &store);
372        assert!(
373            but_one.contains(&g1),
374            "required_flat group g1 should be flattened via AllButOne, got {:?}",
375            but_one
376        );
377    }
378}
linkrs_query/optimizer/factorization/flatten_resolver.rs

linkrs_query/optimizer/factorization/
flatten_resolver.rs
