linkrs_query/optimizer/factorization/
remove_factorization_rewriter.rs

1use crate::planning::plan::factorization::FactorizationError;
2use crate::planning::plan::factorization::FactorizedSchema;
3use crate::planning::plan::factorization::FactorizedSchemaCompute;
4use crate::planning::plan::logical::logical_node_enum::LogicalNodeEnum;
5use crate::planning::plan::logical::logical_node_traits::{
6    LogicalBinaryInputNode, LogicalMultipleInputNode, LogicalSingleInputNode,
7};
8
9/// RemoveFactorizationRewriter: flatten all groups and eliminate LogicalFlatten nodes.
10///
11/// Mirrors `lbug::optimizer::RemoveFactorizationRewriter` in
12/// `ref/ladybug/src/optimizer/remove_factorization_rewriter.cpp`.
13/// Run at the very beginning of the optimizer pipeline so that heuristic
14/// and CBO passes operate on a fully flat view.
15pub struct RemoveFactorizationRewriter;
16
17impl RemoveFactorizationRewriter {
18    pub fn new() -> Self {
19        Self
20    }
21
22    /// Rewrite the plan in place: bottom-up traversal, replace every
23    /// `LogicalFlatten` with its child, and recompute flat schemas on
24    /// every operator so downstream passes see a fully flat view.
25    pub fn rewrite(&self, plan: &mut LogicalNodeEnum) -> Result<(), FactorizationError> {
26        let old = std::mem::take(plan);
27        let (new_root, _schema) = Self::visit_operator(old)?;
28        *plan = new_root;
29        debug_assert!(
30            !Self::has_flatten(plan),
31            "RemoveFactorizationRewriter: residual LogicalFlatten after rewrite"
32        );
33        Ok(())
34    }
35
36    /// Bottom-up traversal returning a new tree without Flatten nodes and
37    /// the flat schema for the rewritten subtree.
38    fn visit_operator(
39        node: LogicalNodeEnum,
40    ) -> Result<(LogicalNodeEnum, FactorizedSchema), FactorizationError> {
41        Self::visit_operator_replace(node)
42    }
43
44    fn visit_operator_replace(
45        node: LogicalNodeEnum,
46    ) -> Result<(LogicalNodeEnum, FactorizedSchema), FactorizationError> {
47        match node {
48            LogicalNodeEnum::Flatten(mut flatten) => {
49                let child = flatten
50                    .input
51                    .take()
52                    .map(|b| *b)
53                    .expect("flatten missing input");
54                Self::visit_operator(child)
55            }
56            LogicalNodeEnum::Project(n) => Self::visit_single_input(n, LogicalNodeEnum::Project),
57            LogicalNodeEnum::Filter(n) => Self::visit_single_input(n, LogicalNodeEnum::Filter),
58            LogicalNodeEnum::Sort(n) => Self::visit_single_input(n, LogicalNodeEnum::Sort),
59            LogicalNodeEnum::Limit(n) => Self::visit_single_input(n, LogicalNodeEnum::Limit),
60            LogicalNodeEnum::Skip(n) => Self::visit_single_input(n, LogicalNodeEnum::Skip),
61            LogicalNodeEnum::TopN(n) => Self::visit_single_input(n, LogicalNodeEnum::TopN),
62            LogicalNodeEnum::Sample(n) => Self::visit_single_input(n, LogicalNodeEnum::Sample),
63            LogicalNodeEnum::Dedup(n) => Self::visit_single_input(n, LogicalNodeEnum::Dedup),
64            LogicalNodeEnum::Aggregate(n) => {
65                Self::visit_single_input(n, LogicalNodeEnum::Aggregate)
66            }
67            LogicalNodeEnum::Window(n) => Self::visit_single_input(n, LogicalNodeEnum::Window),
68            LogicalNodeEnum::GetVertices(n) => Self::visit_deps(n, LogicalNodeEnum::GetVertices),
69            LogicalNodeEnum::GetNeighbors(n) => Self::visit_deps(n, LogicalNodeEnum::GetNeighbors),
70            LogicalNodeEnum::Assign(n) => Self::visit_single_input(n, LogicalNodeEnum::Assign),
71            LogicalNodeEnum::DataCollect(n) => {
72                Self::visit_single_input(n, LogicalNodeEnum::DataCollect)
73            }
74            LogicalNodeEnum::Materialize(n) => {
75                Self::visit_single_input(n, LogicalNodeEnum::Materialize)
76            }
77            LogicalNodeEnum::RollUpApply(n) => {
78                Self::visit_single_input(n, LogicalNodeEnum::RollUpApply)
79            }
80            LogicalNodeEnum::Union(n) => Self::visit_deps(n, LogicalNodeEnum::Union),
81            LogicalNodeEnum::Minus(n) => Self::visit_deps(n, LogicalNodeEnum::Minus),
82            LogicalNodeEnum::Intersect(n) => Self::visit_deps(n, LogicalNodeEnum::Intersect),
83            LogicalNodeEnum::WcoIntersect(n) => Self::visit_deps(n, LogicalNodeEnum::WcoIntersect),
84            LogicalNodeEnum::InnerJoin(n) => Self::visit_binary(n, LogicalNodeEnum::InnerJoin),
85            LogicalNodeEnum::LeftJoin(n) => Self::visit_binary(n, LogicalNodeEnum::LeftJoin),
86            LogicalNodeEnum::RightJoin(n) => Self::visit_binary(n, LogicalNodeEnum::RightJoin),
87            LogicalNodeEnum::CrossJoin(n) => Self::visit_binary(n, LogicalNodeEnum::CrossJoin),
88            LogicalNodeEnum::FullOuterJoin(n) => {
89                Self::visit_binary(n, LogicalNodeEnum::FullOuterJoin)
90            }
91            LogicalNodeEnum::SemiJoin(n) => Self::visit_binary(n, LogicalNodeEnum::SemiJoin),
92            LogicalNodeEnum::PatternApply(n) => {
93                Self::visit_binary(n, LogicalNodeEnum::PatternApply)
94            }
95            LogicalNodeEnum::CorrelatedApply(n) => {
96                Self::visit_binary(n, LogicalNodeEnum::CorrelatedApply)
97            }
98            LogicalNodeEnum::Apply(mut n) => {
99                let (left, left_schema) = Self::visit_operator(n.left_input().clone())?;
100                let (right, right_schema) = Self::visit_operator(n.right_input().clone())?;
101                n.set_left_input(left);
102                n.set_right_input(right);
103                let mut node = LogicalNodeEnum::Apply(n);
104                let schema = node.compute_flat_schema(&[left_schema, right_schema])?;
105                Ok((node, schema))
106            }
107            LogicalNodeEnum::Traverse(n) => Self::visit_single_input(n, LogicalNodeEnum::Traverse),
108            LogicalNodeEnum::Expand(n) => Self::visit_deps(n, LogicalNodeEnum::Expand),
109            LogicalNodeEnum::ExpandAll(n) => Self::visit_deps(n, LogicalNodeEnum::ExpandAll),
110            LogicalNodeEnum::AppendVertices(n) => {
111                Self::visit_deps(n, LogicalNodeEnum::AppendVertices)
112            }
113            LogicalNodeEnum::BiExpand(n) => Self::visit_binary(n, LogicalNodeEnum::BiExpand),
114            LogicalNodeEnum::BiTraverse(n) => Self::visit_binary(n, LogicalNodeEnum::BiTraverse),
115            LogicalNodeEnum::MultiShortestPath(n) => {
116                Self::visit_binary(n, LogicalNodeEnum::MultiShortestPath)
117            }
118            LogicalNodeEnum::BFSShortest(n) => Self::visit_binary(n, LogicalNodeEnum::BFSShortest),
119            LogicalNodeEnum::AllPaths(n) => Self::visit_binary(n, LogicalNodeEnum::AllPaths),
120            LogicalNodeEnum::ShortestPath(n) => {
121                Self::visit_binary(n, LogicalNodeEnum::ShortestPath)
122            }
123            LogicalNodeEnum::Unwind(n) => Self::visit_single_input(n, LogicalNodeEnum::Unwind),
124            LogicalNodeEnum::Select(mut n) => {
125                let mut child_schemas = Vec::new();
126                if let Some(branch) = n.take_if_branch() {
127                    let (new_branch, schema) = Self::visit_operator(*branch)?;
128                    n.set_if_branch(new_branch);
129                    child_schemas.push(schema);
130                }
131                if let Some(branch) = n.take_else_branch() {
132                    let (new_branch, schema) = Self::visit_operator(*branch)?;
133                    n.set_else_branch(new_branch);
134                    child_schemas.push(schema);
135                }
136                let mut node = LogicalNodeEnum::Select(n);
137                let schema = node.compute_flat_schema(&child_schemas)?;
138                Ok((node, schema))
139            }
140            LogicalNodeEnum::Loop(mut n) => {
141                let mut child_schemas = Vec::new();
142                if let Some(body) = n.take_body() {
143                    let (new_body, schema) = Self::visit_operator(*body)?;
144                    n.set_body(new_body);
145                    child_schemas.push(schema);
146                }
147                let mut node = LogicalNodeEnum::Loop(n);
148                let schema = node.compute_flat_schema(&child_schemas)?;
149                Ok((node, schema))
150            }
151            LogicalNodeEnum::PassThrough(n) => {
152                let mut node = LogicalNodeEnum::PassThrough(n);
153                let schema = node.compute_flat_schema(&[])?;
154                Ok((node, schema))
155            }
156            LogicalNodeEnum::Argument(n) => {
157                let mut node = LogicalNodeEnum::Argument(n);
158                let schema = node.compute_flat_schema(&[])?;
159                Ok((node, schema))
160            }
161            LogicalNodeEnum::Start(n) => {
162                let mut node = LogicalNodeEnum::Start(n);
163                let schema = node.compute_flat_schema(&[])?;
164                Ok((node, schema))
165            }
166            LogicalNodeEnum::GetEdges(n) => {
167                let mut node = LogicalNodeEnum::GetEdges(n);
168                let schema = node.compute_flat_schema(&[])?;
169                Ok((node, schema))
170            }
171            LogicalNodeEnum::ScanVertices(n) => {
172                let mut node = LogicalNodeEnum::ScanVertices(n);
173                let schema = node.compute_flat_schema(&[])?;
174                Ok((node, schema))
175            }
176            LogicalNodeEnum::ScanEdges(n) => {
177                let mut node = LogicalNodeEnum::ScanEdges(n);
178                let schema = node.compute_flat_schema(&[])?;
179                Ok((node, schema))
180            }
181            LogicalNodeEnum::BeginTransaction(n) => {
182                let mut node = LogicalNodeEnum::BeginTransaction(n);
183                let schema = node.compute_flat_schema(&[])?;
184                Ok((node, schema))
185            }
186            LogicalNodeEnum::Commit(n) => {
187                let mut node = LogicalNodeEnum::Commit(n);
188                let schema = node.compute_flat_schema(&[])?;
189                Ok((node, schema))
190            }
191            LogicalNodeEnum::Rollback(n) => {
192                let mut node = LogicalNodeEnum::Rollback(n);
193                let schema = node.compute_flat_schema(&[])?;
194                Ok((node, schema))
195            }
196            LogicalNodeEnum::InsertVertices(n) => {
197                let mut node = LogicalNodeEnum::InsertVertices(n);
198                let schema = node.compute_flat_schema(&[])?;
199                Ok((node, schema))
200            }
201            LogicalNodeEnum::InsertEdges(n) => {
202                let mut node = LogicalNodeEnum::InsertEdges(n);
203                let schema = node.compute_flat_schema(&[])?;
204                Ok((node, schema))
205            }
206            LogicalNodeEnum::Update(n) => {
207                let mut node = LogicalNodeEnum::Update(n);
208                let schema = node.compute_flat_schema(&[])?;
209                Ok((node, schema))
210            }
211            LogicalNodeEnum::DeleteVertices(n) => {
212                let mut node = LogicalNodeEnum::DeleteVertices(n);
213                let schema = node.compute_flat_schema(&[])?;
214                Ok((node, schema))
215            }
216            LogicalNodeEnum::DeleteEdges(n) => {
217                let mut node = LogicalNodeEnum::DeleteEdges(n);
218                let schema = node.compute_flat_schema(&[])?;
219                Ok((node, schema))
220            }
221            LogicalNodeEnum::DeleteIndex(n) => {
222                let mut node = LogicalNodeEnum::DeleteIndex(n);
223                let schema = node.compute_flat_schema(&[])?;
224                Ok((node, schema))
225            }
226            LogicalNodeEnum::PipeDeleteVertices(n) => {
227                let mut node = LogicalNodeEnum::PipeDeleteVertices(n);
228                let schema = node.compute_flat_schema(&[])?;
229                Ok((node, schema))
230            }
231            LogicalNodeEnum::PipeDeleteEdges(n) => {
232                let mut node = LogicalNodeEnum::PipeDeleteEdges(n);
233                let schema = node.compute_flat_schema(&[])?;
234                Ok((node, schema))
235            }
236            LogicalNodeEnum::CopyFrom(n) => {
237                let mut node = LogicalNodeEnum::CopyFrom(n);
238                let schema = node.compute_flat_schema(&[])?;
239                Ok((node, schema))
240            }
241            LogicalNodeEnum::CopyTo(n) => {
242                let mut node = LogicalNodeEnum::CopyTo(n);
243                let schema = node.compute_flat_schema(&[])?;
244                Ok((node, schema))
245            }
246            LogicalNodeEnum::FulltextSearch(n) => {
247                let mut node = LogicalNodeEnum::FulltextSearch(n);
248                let schema = node.compute_flat_schema(&[])?;
249                Ok((node, schema))
250            }
251            LogicalNodeEnum::FulltextLookup(n) => {
252                let mut node = LogicalNodeEnum::FulltextLookup(n);
253                let schema = node.compute_flat_schema(&[])?;
254                Ok((node, schema))
255            }
256            LogicalNodeEnum::MatchFulltext(n) => {
257                let mut node = LogicalNodeEnum::MatchFulltext(n);
258                let schema = node.compute_flat_schema(&[])?;
259                Ok((node, schema))
260            }
261            #[cfg(feature = "vector")]
262            LogicalNodeEnum::VectorSearch(n) => {
263                let mut node = LogicalNodeEnum::VectorSearch(n);
264                let schema = node.compute_flat_schema(&[])?;
265                Ok((node, schema))
266            }
267            #[cfg(feature = "vector")]
268            LogicalNodeEnum::VectorLookup(n) => {
269                let mut node = LogicalNodeEnum::VectorLookup(n);
270                let schema = node.compute_flat_schema(&[])?;
271                Ok((node, schema))
272            }
273            #[cfg(feature = "vector")]
274            LogicalNodeEnum::VectorMatch(n) => {
275                let mut node = LogicalNodeEnum::VectorMatch(n);
276                let schema = node.compute_flat_schema(&[])?;
277                Ok((node, schema))
278            }
279        }
280    }
281
282    pub fn has_flatten_public(node: &LogicalNodeEnum) -> bool {
283        Self::has_flatten(node)
284    }
285
286    fn has_flatten(node: &LogicalNodeEnum) -> bool {
287        if matches!(node, LogicalNodeEnum::Flatten(_)) {
288            return true;
289        }
290        match node {
291            LogicalNodeEnum::Project(n) => n.input.as_deref().is_some_and(Self::has_flatten),
292            LogicalNodeEnum::Filter(n) => n.input.as_deref().is_some_and(Self::has_flatten),
293            LogicalNodeEnum::Sort(n) => n.input.as_deref().is_some_and(Self::has_flatten),
294            LogicalNodeEnum::Limit(n) => n.input.as_deref().is_some_and(Self::has_flatten),
295            LogicalNodeEnum::Skip(n) => n.input.as_deref().is_some_and(Self::has_flatten),
296            LogicalNodeEnum::TopN(n) => n.input.as_deref().is_some_and(Self::has_flatten),
297            LogicalNodeEnum::Sample(n) => n.input.as_deref().is_some_and(Self::has_flatten),
298            LogicalNodeEnum::Dedup(n) => n.input.as_deref().is_some_and(Self::has_flatten),
299            LogicalNodeEnum::Aggregate(n) => n.input.as_deref().is_some_and(Self::has_flatten),
300            LogicalNodeEnum::Window(n) => n.input.as_deref().is_some_and(Self::has_flatten),
301            LogicalNodeEnum::Traverse(n) => n.input.as_deref().is_some_and(Self::has_flatten),
302            LogicalNodeEnum::Expand(n) => n.deps.iter().any(Self::has_flatten),
303            LogicalNodeEnum::ExpandAll(n) => n.deps.iter().any(Self::has_flatten),
304            LogicalNodeEnum::AppendVertices(n) => n.deps.iter().any(Self::has_flatten),
305            LogicalNodeEnum::BiExpand(n) => {
306                Self::has_flatten(&n.left) || Self::has_flatten(&n.right)
307            }
308            LogicalNodeEnum::BiTraverse(n) => {
309                Self::has_flatten(&n.left) || Self::has_flatten(&n.right)
310            }
311            LogicalNodeEnum::GetVertices(n) => n.deps.iter().any(Self::has_flatten),
312            LogicalNodeEnum::GetNeighbors(n) => n.deps.iter().any(Self::has_flatten),
313            LogicalNodeEnum::Assign(n) => n.input.as_deref().is_some_and(Self::has_flatten),
314            LogicalNodeEnum::DataCollect(n) => n.input.as_deref().is_some_and(Self::has_flatten),
315            LogicalNodeEnum::Materialize(n) => n.input.as_deref().is_some_and(Self::has_flatten),
316            LogicalNodeEnum::RollUpApply(n) => n.input.as_deref().is_some_and(Self::has_flatten),
317            LogicalNodeEnum::Union(n) => n.deps.iter().any(Self::has_flatten),
318            LogicalNodeEnum::Minus(n) => n.deps.iter().any(Self::has_flatten),
319            LogicalNodeEnum::Intersect(n) => n.deps.iter().any(Self::has_flatten),
320            LogicalNodeEnum::InnerJoin(n) => {
321                Self::has_flatten(&n.left) || Self::has_flatten(&n.right)
322            }
323            LogicalNodeEnum::LeftJoin(n) => {
324                Self::has_flatten(&n.left) || Self::has_flatten(&n.right)
325            }
326            LogicalNodeEnum::RightJoin(n) => {
327                Self::has_flatten(&n.left) || Self::has_flatten(&n.right)
328            }
329            LogicalNodeEnum::CrossJoin(n) => {
330                Self::has_flatten(&n.left) || Self::has_flatten(&n.right)
331            }
332            LogicalNodeEnum::FullOuterJoin(n) => {
333                Self::has_flatten(&n.left) || Self::has_flatten(&n.right)
334            }
335            LogicalNodeEnum::SemiJoin(n) => {
336                Self::has_flatten(&n.left) || Self::has_flatten(&n.right)
337            }
338            LogicalNodeEnum::PatternApply(n) => {
339                Self::has_flatten(&n.left) || Self::has_flatten(&n.right)
340            }
341            LogicalNodeEnum::CorrelatedApply(n) => {
342                Self::has_flatten(&n.left) || Self::has_flatten(&n.right)
343            }
344            LogicalNodeEnum::Apply(n) => {
345                Self::has_flatten(n.left_input()) || Self::has_flatten(n.right_input())
346            }
347            LogicalNodeEnum::MultiShortestPath(n) => {
348                Self::has_flatten(&n.left) || Self::has_flatten(&n.right)
349            }
350            LogicalNodeEnum::BFSShortest(n) => {
351                Self::has_flatten(&n.left) || Self::has_flatten(&n.right)
352            }
353            LogicalNodeEnum::AllPaths(n) => {
354                Self::has_flatten(&n.left) || Self::has_flatten(&n.right)
355            }
356            LogicalNodeEnum::ShortestPath(n) => {
357                Self::has_flatten(&n.left) || Self::has_flatten(&n.right)
358            }
359            LogicalNodeEnum::Unwind(n) => n.input.as_deref().is_some_and(Self::has_flatten),
360            LogicalNodeEnum::Select(n) => {
361                n.if_branch().is_some_and(Self::has_flatten)
362                    || n.else_branch().is_some_and(Self::has_flatten)
363            }
364            LogicalNodeEnum::Loop(n) => n.body().is_some_and(Self::has_flatten),
365            LogicalNodeEnum::Flatten(_) => true,
366            _ => false,
367        }
368    }
369
370    /// Visit a node with exactly one optional `input` child.
371    ///
372    /// `#[inline(never)]` on purpose: [`Self::visit_operator_replace`] is a
373    /// giant `match` whose arms each keep a large concrete node temporary
374    /// alive across the recursive call, so the compiler reserves stack for
375    /// the union of every arm in a single frame. Routing the uniform arms
376    /// through thin helpers keeps the recursion frames small enough for the
377    /// default 2 MiB test-thread stack.
378    #[inline(never)]
379    fn visit_single_input<N, F>(
380        mut n: N,
381        wrap: F,
382    ) -> Result<(LogicalNodeEnum, FactorizedSchema), FactorizationError>
383    where
384        N: LogicalSingleInputNode,
385        F: FnOnce(N) -> LogicalNodeEnum,
386    {
387        let mut child_schemas = Vec::new();
388        if let Some(input) = n.take_input() {
389            let (new_input, schema) = Self::visit_operator(input)?;
390            n.set_input(new_input);
391            child_schemas.push(schema);
392        }
393        let mut node = wrap(n);
394        let schema = node.compute_flat_schema(&child_schemas)?;
395        Ok((node, schema))
396    }
397
398    /// Visit a node whose children live in the macro-generated `deps` vector.
399    #[inline(never)]
400    fn visit_deps<N, F>(
401        mut n: N,
402        wrap: F,
403    ) -> Result<(LogicalNodeEnum, FactorizedSchema), FactorizationError>
404    where
405        N: LogicalMultipleInputNode,
406        F: FnOnce(N) -> LogicalNodeEnum,
407    {
408        let child_schemas: Vec<FactorizedSchema> = n
409            .inputs_mut()
410            .iter_mut()
411            .map(|dep| {
412                let old = std::mem::take(dep);
413                let (new_dep, schema) = Self::visit_operator(old)?;
414                *dep = new_dep;
415                Ok(schema)
416            })
417            .collect::<Result<Vec<FactorizedSchema>, FactorizationError>>()?;
418        let mut node = wrap(n);
419        let schema = node.compute_flat_schema(&child_schemas)?;
420        Ok((node, schema))
421    }
422
423    /// Visit a node with a boxed `left`/`right` child pair.
424    #[inline(never)]
425    fn visit_binary<N, F>(
426        mut n: N,
427        wrap: F,
428    ) -> Result<(LogicalNodeEnum, FactorizedSchema), FactorizationError>
429    where
430        N: LogicalBinaryInputNode,
431        F: FnOnce(N) -> LogicalNodeEnum,
432    {
433        let left_old = std::mem::take(n.left_input_mut());
434        let right_old = std::mem::take(n.right_input_mut());
435        let (left, left_schema) = Self::visit_operator(left_old)?;
436        let (right, right_schema) = Self::visit_operator(right_old)?;
437        n.set_left_input(left);
438        n.set_right_input(right);
439        let mut node = wrap(n);
440        let schema = node.compute_flat_schema(&[left_schema, right_schema])?;
441        Ok((node, schema))
442    }
443}
444
445impl Default for RemoveFactorizationRewriter {
446    fn default() -> Self {
447        Self::new()
448    }
449}
450
451#[cfg(test)]
452mod tests {
453    use super::*;
454    use crate::planning::plan::core::node_id_generator::next_node_id;
455    use crate::planning::plan::logical::logical_node_enum::LogicalNodeEnum;
456    use crate::planning::plan::logical::logical_nodes::access::LogicalScanVerticesNode;
457    use crate::planning::plan::logical::logical_nodes::flatten::LogicalFlattenNode;
458
459    fn scan() -> LogicalNodeEnum {
460        LogicalNodeEnum::ScanVertices(LogicalScanVerticesNode {
461            id: next_node_id(),
462            space_id: 1,
463            space_name: "test".to_string(),
464            tag: Some("person".to_string()),
465            expression: None,
466            limit: None,
467            projected_properties: vec![],
468            index_hint: None,
469            estimated_cardinality: None,
470            output_var: None,
471            col_names: vec!["a.name".to_string()],
472            column_types: vec![],
473        })
474    }
475
476    #[test]
477    fn remove_single_flatten() {
478        let scan = scan();
479        let flatten = LogicalNodeEnum::Flatten(LogicalFlattenNode::new(0, scan));
480        let mut root = flatten;
481        let rewriter = RemoveFactorizationRewriter::new();
482        rewriter.rewrite(&mut root).unwrap();
483        assert!(!matches!(root, LogicalNodeEnum::Flatten(_)));
484        assert_eq!(root.type_name(), "ScanVertices");
485    }
486
487    #[test]
488    fn remove_nested_flatten() {
489        let scan = scan();
490        let f1 = LogicalNodeEnum::Flatten(LogicalFlattenNode::new(0, scan));
491        let f2 = LogicalNodeEnum::Flatten(LogicalFlattenNode::new(1, f1));
492        let mut root = f2;
493        RemoveFactorizationRewriter::new()
494            .rewrite(&mut root)
495            .unwrap();
496        assert_eq!(root.type_name(), "ScanVertices");
497    }
498
499    #[test]
500    fn remove_nested_under_assign() {
501        let scan = scan();
502        let flatten = LogicalNodeEnum::Flatten(LogicalFlattenNode::new(0, scan.clone()));
503        let mut assign = LogicalNodeEnum::Assign(
504            crate::planning::plan::logical::logical_nodes::graph_ops::LogicalAssignNode {
505                id: next_node_id(),
506                input: Some(Box::new(flatten)),
507                assignments: vec![],
508                output_var: None,
509                col_names: vec![],
510                column_types: vec![],
511            },
512        );
513        RemoveFactorizationRewriter::new()
514            .rewrite(&mut assign)
515            .unwrap();
516        assert!(!RemoveFactorizationRewriter::has_flatten(&assign));
517    }
518
519    #[test]
520    fn remove_under_bi_traverse() {
521        let left = scan();
522        let right = scan();
523        let flatten_left = LogicalNodeEnum::Flatten(LogicalFlattenNode::new(0, left));
524        let mut bi = LogicalNodeEnum::BiTraverse(
525            crate::planning::plan::logical::logical_nodes::traversal::LogicalBiTraverseNode {
526                id: next_node_id(),
527                left: Box::new(flatten_left),
528                right: Box::new(right),
529                space_id: 1,
530                left_src_var: "a".to_string(),
531                right_src_var: "b".to_string(),
532                edge_types: vec![],
533                left_direction: linkrs_core::types::EdgeDirection::Out,
534                right_direction: linkrs_core::types::EdgeDirection::Out,
535                min_hops: 1,
536                max_hops: 3,
537                path_var: "p".to_string(),
538                edge_alias: None,
539                vertex_alias: None,
540                dst_tag: None,
541                output_var: None,
542                col_names: vec![],
543                column_types: vec![],
544            },
545        );
546        RemoveFactorizationRewriter::new().rewrite(&mut bi).unwrap();
547        assert!(!RemoveFactorizationRewriter::has_flatten(&bi));
548    }
549
550    #[test]
551    fn has_flatten_deep_loop() {
552        let scan = scan();
553        let flatten = LogicalNodeEnum::Flatten(LogicalFlattenNode::new(0, scan));
554        let ctx = std::sync::Arc::new(
555            linkrs_core::types::expr::expression_context::ExpressionAnalysisContext::new(),
556        );
557        let expr = linkrs_core::Expression::Variable("x".to_string());
558        let meta = linkrs_core::types::expr::ExpressionMeta::new(expr);
559        let id = ctx.register_expression(meta);
560        let cond = linkrs_core::types::expr::contextual::ContextualExpression::new(id, ctx);
561        let loop_node = LogicalNodeEnum::Loop(
562            crate::planning::plan::logical::logical_nodes::control_flow::LogicalLoopNode::new_with_body(cond, flatten),
563        );
564        assert!(RemoveFactorizationRewriter::has_flatten(&loop_node));
565    }
566}
linkrs_query/optimizer/factorization/remove_factorization_rewriter.rs

linkrs_query/optimizer/factorization/
remove_factorization_rewriter.rs
