1use std::collections::{HashMap, HashSet};
2
3use linkrs_core::types::expr::ExpressionId;
4
5use super::flatten_resolver::{FlattenAll, FlattenAllButOne};
6use crate::planning::plan::factorization::{
7 FGroupPos, FactorizationError, FactorizedSchema, FactorizedSchemaCompute,
8};
9use crate::planning::plan::logical::LogicalNodeEnum;
10
11mod flatten;
12mod join;
13
14pub struct FactorizationRewriter {
15 pub enabled: bool,
16 skipped_flat_groups: Vec<FGroupPos>,
17}
18
19impl FactorizationRewriter {
20 pub fn new() -> Self {
21 Self {
22 enabled: true,
23 skipped_flat_groups: Vec::new(),
24 }
25 }
26
27 pub fn disabled() -> Self {
28 Self {
29 enabled: false,
30 skipped_flat_groups: Vec::new(),
31 }
32 }
33
34 pub fn rewrite(&mut self, plan: &mut LogicalNodeEnum) -> Result<(), FactorizationError> {
35 if !self.enabled {
36 return Ok(());
37 }
38 self.visit_operator(plan)?;
39 Ok(())
40 }
41
42 pub fn take_skipped_flat_groups(&mut self) -> Vec<FGroupPos> {
43 std::mem::take(&mut self.skipped_flat_groups)
44 }
45
46 fn visit_operator(
47 &mut self,
48 node: &mut LogicalNodeEnum,
49 ) -> Result<FactorizedSchema, FactorizationError> {
50 match node {
51 LogicalNodeEnum::Project(n) => {
52 let mut child_schema = if let Some(child) = n.input.as_mut() {
53 self.visit_operator(child)?
54 } else {
55 FactorizedSchema::new()
56 };
57 let store = Self::build_store_for_project(n);
58 let has_random = store
59 .values()
60 .any(crate::optimizer::analysis::expression::NondeterministicChecker::contains_random);
61 if has_random {
62 let groups = child_schema.groups_pos_in_scope();
63 let to_flatten =
64 FlattenAll::get_groups_pos_to_flatten_for_groups(&groups, &child_schema);
65 if !to_flatten.is_empty() {
66 if let Some(child) = n.input.as_mut() {
67 self.replace_child_and_flatten(child, &to_flatten, &mut child_schema)?;
68 }
69 for pos in &to_flatten {
70 child_schema.flatten_group(*pos)?;
71 }
72 }
73 } else {
74 for expr_id in Self::expr_ids_for_project(n) {
75 let to_flatten = FlattenAllButOne::get_groups_pos_to_flatten_for_expr(
76 &expr_id,
77 &child_schema,
78 &store,
79 );
80 if !to_flatten.is_empty() {
81 if let Some(child) = n.input.as_mut() {
82 self.replace_child_and_flatten(
83 child,
84 &to_flatten,
85 &mut child_schema,
86 )?;
87 }
88 for pos in &to_flatten {
89 child_schema.flatten_group(*pos)?;
90 }
91 }
92 }
93 }
94 node.compute_factorized_schema(&[child_schema])
95 }
96 LogicalNodeEnum::Filter(n) => {
97 let mut child_schema = if let Some(child) = n.input.as_mut() {
98 self.visit_operator(child)?
99 } else {
100 FactorizedSchema::new()
101 };
102 let expr_id = n.condition.id().clone();
103 let mut store = HashMap::new();
104 if let Some(expr) = n.condition.get_expression() {
105 store.insert(expr_id.clone(), expr);
106 }
107 let to_flatten = FlattenAllButOne::get_groups_pos_to_flatten_for_expr(
108 &expr_id,
109 &child_schema,
110 &store,
111 );
112 if !to_flatten.is_empty() {
113 if let Some(child) = n.input.as_mut() {
114 self.replace_child_and_flatten(child, &to_flatten, &mut child_schema)?;
115 }
116 for pos in &to_flatten {
117 child_schema.flatten_group(*pos)?;
118 }
119 return node.compute_factorized_schema(&[child_schema]);
120 }
121 node.compute_factorized_schema(&[child_schema])
122 }
123 LogicalNodeEnum::Aggregate(n) => {
124 let mut child_schema = if let Some(child) = n.input.as_mut() {
125 self.visit_operator(child)?
126 } else {
127 FactorizedSchema::new()
128 };
129 let key_ids: Vec<ExpressionId> =
130 n.group_key_exprs.iter().map(|e| e.id().clone()).collect();
131 let mut store = HashMap::new();
132 for expr in &n.group_key_exprs {
133 if let Some(inner) = expr.get_expression() {
134 store.insert(expr.id().clone(), inner);
135 }
136 }
137 let (_leading, to_flatten) = super::flatten_resolver::aggregate_groups_to_flatten(
138 &key_ids,
139 &store,
140 &n.aggregation_args,
141 &n.aggregation_distinct,
142 &child_schema,
143 );
144 if !to_flatten.is_empty() {
145 if let Some(child) = n.input.as_mut() {
146 self.replace_child_and_flatten(child, &to_flatten, &mut child_schema)?;
147 }
148 for pos in &to_flatten {
149 child_schema.flatten_group(*pos)?;
150 }
151 return node.compute_factorized_schema(&[child_schema]);
152 }
153 node.compute_factorized_schema(&[child_schema])
154 }
155 LogicalNodeEnum::Sort(n) => {
156 let mut child_schema = if let Some(child) = n.input.as_mut() {
157 self.visit_operator(child)?
158 } else {
159 FactorizedSchema::new()
160 };
161 let to_flatten = if child_schema.num_groups() > 1 {
162 let (expr_ids, store) = Self::sort_key_store(n, &child_schema);
163 let mut dependent = HashSet::new();
164 let mut unresolved = false;
165 for expr_id in &expr_ids {
166 let mut analyzer =
167 super::group_dependency_analyzer::GroupDependencyAnalyzer::with_expr_store(
168 &child_schema,
169 false,
170 &store,
171 );
172 analyzer.visit(expr_id);
173 dependent.extend(analyzer.dependent_groups().iter().copied());
174 unresolved |= analyzer.has_unresolved();
175 }
176 if unresolved {
177 FlattenAll::get_groups_pos_to_flatten_for_groups(
178 &child_schema.groups_pos_in_scope(),
179 &child_schema,
180 )
181 } else {
182 FlattenAll::get_groups_pos_to_flatten_for_groups(&dependent, &child_schema)
183 }
184 } else {
185 HashSet::new()
186 };
187 if !to_flatten.is_empty() {
188 if let Some(child) = n.input.as_mut() {
189 self.replace_child_and_flatten(child, &to_flatten, &mut child_schema)?;
190 }
191 for pos in &to_flatten {
192 child_schema.flatten_group(*pos)?;
193 }
194 return node.compute_factorized_schema(&[child_schema]);
195 }
196 node.compute_factorized_schema(&[child_schema])
197 }
198 LogicalNodeEnum::Window(n) => {
199 let mut child_schema = if let Some(child) = n.input.as_mut() {
200 self.visit_operator(child)?
201 } else {
202 FactorizedSchema::new()
203 };
204 let groups = child_schema.groups_pos_in_scope();
205 let to_flatten =
206 FlattenAllButOne::get_groups_pos_to_flatten_for_groups(&groups, &child_schema);
207 if !to_flatten.is_empty() {
208 if let Some(child) = n.input.as_mut() {
209 self.replace_child_and_flatten(child, &to_flatten, &mut child_schema)?;
210 }
211 for pos in &to_flatten {
212 child_schema.flatten_group(*pos)?;
213 }
214 return node.compute_factorized_schema(&[child_schema]);
215 }
216 node.compute_factorized_schema(&[child_schema])
217 }
218 LogicalNodeEnum::Limit(n) => {
219 let mut child_schema = if let Some(child) = n.input.as_mut() {
220 self.visit_operator(child)?
221 } else {
222 FactorizedSchema::new()
223 };
224 let groups = child_schema.groups_pos_in_scope();
225 let to_flatten =
226 FlattenAllButOne::get_groups_pos_to_flatten_for_groups(&groups, &child_schema);
227 if !to_flatten.is_empty() {
228 if let Some(child) = n.input.as_mut() {
229 self.replace_child_and_flatten(child, &to_flatten, &mut child_schema)?;
230 }
231 for pos in &to_flatten {
232 child_schema.flatten_group(*pos)?;
233 }
234 return node.compute_factorized_schema(&[child_schema]);
235 }
236 node.compute_factorized_schema(&[child_schema])
237 }
238 LogicalNodeEnum::Skip(n) => {
239 let child_schema = if let Some(child) = n.input.as_mut() {
240 self.visit_operator(child)?
241 } else {
242 FactorizedSchema::new()
243 };
244 node.compute_factorized_schema(&[child_schema])
245 }
246 LogicalNodeEnum::TopN(n) => {
247 let mut child_schema = if let Some(child) = n.input.as_mut() {
248 self.visit_operator(child)?
249 } else {
250 FactorizedSchema::new()
251 };
252 let groups = child_schema.groups_pos_in_scope();
253 let to_flatten =
254 FlattenAllButOne::get_groups_pos_to_flatten_for_groups(&groups, &child_schema);
255 if !to_flatten.is_empty() {
256 if let Some(child) = n.input.as_mut() {
257 self.replace_child_and_flatten(child, &to_flatten, &mut child_schema)?;
258 }
259 for pos in &to_flatten {
260 child_schema.flatten_group(*pos)?;
261 }
262 return node.compute_factorized_schema(&[child_schema]);
263 }
264 node.compute_factorized_schema(&[child_schema])
265 }
266 LogicalNodeEnum::Dedup(n) => {
267 let mut child_schema = if let Some(child) = n.input.as_mut() {
268 self.visit_operator(child)?
269 } else {
270 FactorizedSchema::new()
271 };
272 let groups = child_schema.groups_pos_in_scope();
273 let to_flatten =
274 FlattenAll::get_groups_pos_to_flatten_for_groups(&groups, &child_schema);
275 if !to_flatten.is_empty() {
276 if let Some(child) = n.input.as_mut() {
277 self.replace_child_and_flatten(child, &to_flatten, &mut child_schema)?;
278 }
279 for pos in &to_flatten {
280 child_schema.flatten_group(*pos)?;
281 }
282 return node.compute_factorized_schema(&[child_schema]);
283 }
284 node.compute_factorized_schema(&[child_schema])
285 }
286 LogicalNodeEnum::InnerJoin(n) => {
287 let mut left_schema = self.visit_operator(&mut n.left)?;
288 let mut right_schema = self.visit_operator(&mut n.right)?;
289 self.visit_hash_join_inner(n, &mut left_schema, &mut right_schema)?;
290 node.compute_factorized_schema(&[left_schema, right_schema])
291 }
292 LogicalNodeEnum::LeftJoin(n) => {
293 let mut left_schema = self.visit_operator(&mut n.left)?;
294 let mut right_schema = self.visit_operator(&mut n.right)?;
295 self.visit_hash_join_left(n, &mut left_schema, &mut right_schema)?;
296 node.compute_factorized_schema(&[left_schema, right_schema])
297 }
298 LogicalNodeEnum::RightJoin(n) => {
299 let mut left_schema = self.visit_operator(&mut n.left)?;
300 let mut right_schema = self.visit_operator(&mut n.right)?;
301 self.visit_hash_join_right(
302 &mut left_schema,
303 &mut right_schema,
304 &n.hash_keys,
305 &n.probe_keys,
306 &mut n.left,
307 &mut n.right,
308 )?;
309 node.compute_factorized_schema(&[left_schema, right_schema])
310 }
311 LogicalNodeEnum::CrossJoin(n) => {
312 let mut left_schema = self.visit_operator(&mut n.left)?;
313 let mut right_schema = self.visit_operator(&mut n.right)?;
314 if n.hash_keys.is_empty() && n.probe_keys.is_empty() {
315 } else {
316 self.visit_hash_join_generic_inner(
317 &mut left_schema,
318 &mut right_schema,
319 &n.hash_keys,
320 &n.probe_keys,
321 &mut n.left,
322 &mut n.right,
323 false,
324 )?;
325 }
326 node.compute_factorized_schema(&[left_schema, right_schema])
327 }
328 LogicalNodeEnum::FullOuterJoin(n) => {
329 let mut left_schema = self.visit_operator(&mut n.left)?;
330 let mut right_schema = self.visit_operator(&mut n.right)?;
331 self.visit_hash_join_full_outer(
332 &mut left_schema,
333 &mut right_schema,
334 &n.hash_keys,
335 &n.probe_keys,
336 &mut n.left,
337 &mut n.right,
338 )?;
339 node.compute_factorized_schema(&[left_schema, right_schema])
340 }
341 LogicalNodeEnum::SemiJoin(n) => {
342 let mut left_schema = self.visit_operator(&mut n.left)?;
343 let mut right_schema = self.visit_operator(&mut n.right)?;
344 self.visit_hash_join_generic_inner(
345 &mut left_schema,
346 &mut right_schema,
347 &n.hash_keys,
348 &n.probe_keys,
349 &mut n.left,
350 &mut n.right,
351 false,
352 )?;
353 node.compute_factorized_schema(&[left_schema, right_schema])
354 }
355 LogicalNodeEnum::Traverse(n) => {
356 let effective_schema = if let Some(child) = n.input.as_mut() {
357 let mut schema = self.visit_operator(child)?;
358 if let Some(pos) = schema.unflat_group_pos() {
359 let mut to_flatten = HashSet::new();
360 to_flatten.insert(pos);
361 self.replace_child_and_flatten(child, &to_flatten, &mut schema)?;
362 schema.flatten_group(pos)?;
363 schema
364 } else {
365 schema
366 }
367 } else {
368 FactorizedSchema::new()
369 };
370 node.compute_factorized_schema(&[effective_schema])
371 }
372 LogicalNodeEnum::Expand(n) => {
373 let mut child_schema = FactorizedSchema::new();
374 for dep in &mut n.deps {
375 let mut schema = self.visit_operator(dep)?;
376 if let Some(pos) = schema.unflat_group_pos() {
377 let mut to_flatten = HashSet::new();
378 to_flatten.insert(pos);
379 self.replace_node_and_flatten(dep, &to_flatten, &schema)?;
380 schema.flatten_group(pos)?;
381 }
382 if child_schema.num_groups() == 0 {
383 child_schema = schema;
384 }
385 }
386 node.compute_factorized_schema(&[child_schema])
387 }
388 LogicalNodeEnum::ExpandAll(n) => {
389 let mut child_schema = FactorizedSchema::new();
390 for dep in &mut n.deps {
391 let mut schema = self.visit_operator(dep)?;
392 if let Some(pos) = schema.unflat_group_pos() {
393 let mut to_flatten = HashSet::new();
394 to_flatten.insert(pos);
395 self.replace_node_and_flatten(dep, &to_flatten, &schema)?;
396 schema.flatten_group(pos)?;
397 }
398 if child_schema.num_groups() == 0 {
399 child_schema = schema;
400 }
401 }
402 node.compute_factorized_schema(&[child_schema])
403 }
404 LogicalNodeEnum::BiExpand(n) => {
405 let mut left_schema = self.visit_operator(&mut n.left)?;
406 let mut right_schema = self.visit_operator(&mut n.right)?;
407 if let Some(pos) = left_schema.unflat_group_pos() {
408 let mut to_flatten = HashSet::new();
409 to_flatten.insert(pos);
410 self.replace_child_and_flatten(&mut n.left, &to_flatten, &mut left_schema)?;
411 left_schema.flatten_group(pos)?;
412 }
413 if let Some(pos) = right_schema.unflat_group_pos() {
414 let mut to_flatten = HashSet::new();
415 to_flatten.insert(pos);
416 self.replace_child_and_flatten(&mut n.right, &to_flatten, &mut right_schema)?;
417 right_schema.flatten_group(pos)?;
418 }
419 node.compute_factorized_schema(&[left_schema, right_schema])
420 }
421 LogicalNodeEnum::BiTraverse(n) => {
422 let mut left_schema = self.visit_operator(&mut n.left)?;
423 let mut right_schema = self.visit_operator(&mut n.right)?;
424 if let Some(pos) = left_schema.unflat_group_pos() {
425 let mut to_flatten = HashSet::new();
426 to_flatten.insert(pos);
427 self.replace_child_and_flatten(&mut n.left, &to_flatten, &mut left_schema)?;
428 left_schema.flatten_group(pos)?;
429 }
430 if let Some(pos) = right_schema.unflat_group_pos() {
431 let mut to_flatten = HashSet::new();
432 to_flatten.insert(pos);
433 self.replace_child_and_flatten(&mut n.right, &to_flatten, &mut right_schema)?;
434 right_schema.flatten_group(pos)?;
435 }
436 node.compute_factorized_schema(&[left_schema, right_schema])
437 }
438 LogicalNodeEnum::AppendVertices(n) => {
439 let mut child_schema = FactorizedSchema::new();
440 for dep in &mut n.deps {
441 let mut schema = self.visit_operator(dep)?;
442 if let Some(pos) = schema.unflat_group_pos() {
443 let mut to_flatten = HashSet::new();
444 to_flatten.insert(pos);
445 self.replace_node_and_flatten(dep, &to_flatten, &schema)?;
446 schema.flatten_group(pos)?;
447 }
448 if child_schema.num_groups() == 0 {
449 child_schema = schema;
450 }
451 }
452 node.compute_factorized_schema(&[child_schema])
453 }
454 LogicalNodeEnum::Unwind(n) => {
455 let mut child_schema = if let Some(child) = n.input.as_mut() {
456 self.visit_operator(child)?
457 } else {
458 FactorizedSchema::new()
459 };
460 let list_id = n.list_expression.id().clone();
461 let mut store = HashMap::new();
462 if let Some(expr) = n.list_expression.get_expression() {
463 store.insert(list_id.clone(), expr);
464 }
465 let to_flatten =
466 FlattenAll::get_groups_pos_to_flatten_for_expr(&list_id, &child_schema, &store);
467 let to_flatten = if to_flatten.is_empty() && !child_schema.is_flat_schema() {
468 let mut analyzer =
469 crate::optimizer::factorization::GroupDependencyAnalyzer::with_expr_store(
470 &child_schema,
471 false,
472 &store,
473 );
474 analyzer.visit(&list_id);
475 if analyzer.has_unresolved() {
476 FlattenAll::get_groups_pos_to_flatten_for_groups(
477 &child_schema.groups_pos_in_scope(),
478 &child_schema,
479 )
480 } else {
481 to_flatten
482 }
483 } else {
484 to_flatten
485 };
486 if !to_flatten.is_empty() {
487 if let Some(child) = n.input.as_mut() {
488 self.replace_child_and_flatten(child, &to_flatten, &mut child_schema)?;
489 }
490 for pos in &to_flatten {
491 child_schema.flatten_group(*pos)?;
492 }
493 return node.compute_factorized_schema(&[child_schema]);
494 }
495 node.compute_factorized_schema(&[child_schema])
496 }
497 LogicalNodeEnum::Union(n) => {
498 let mut child_schemas = Vec::new();
499 for dep in &mut n.deps {
500 let mut schema = self.visit_operator(dep)?;
501 self.flatten_barrier_child(dep, &schema)?;
502 schema.flatten_all()?;
503 child_schemas.push(schema);
504 }
505 node.compute_factorized_schema(&child_schemas)
506 }
507 LogicalNodeEnum::Minus(n) => {
508 let mut child_schemas = Vec::new();
509 for dep in &mut n.deps {
510 let mut schema = self.visit_operator(dep)?;
511 self.flatten_barrier_child(dep, &schema)?;
512 schema.flatten_all()?;
513 child_schemas.push(schema);
514 }
515 node.compute_factorized_schema(&child_schemas)
516 }
517 LogicalNodeEnum::Intersect(n) => {
518 let mut child_schemas = Vec::new();
519 for dep in &mut n.deps {
520 let mut schema = self.visit_operator(dep)?;
521 self.flatten_barrier_child(dep, &schema)?;
522 schema.flatten_all()?;
523 child_schemas.push(schema);
524 }
525 node.compute_factorized_schema(&child_schemas)
526 }
527 LogicalNodeEnum::WcoIntersect(n) => {
528 let mut child_schemas = Vec::new();
529 for dep in &mut n.deps {
530 child_schemas.push(self.visit_operator(dep)?);
531 }
532 let probe_to_flatten = child_schemas
533 .first()
534 .map(|probe_schema| n.get_groups_to_flatten_on_probe_side(probe_schema))
535 .unwrap_or_default();
536 if !probe_to_flatten.is_empty() {
537 self.replace_node_and_flatten(
538 &mut n.deps[0],
539 &probe_to_flatten,
540 &child_schemas[0],
541 )?;
542 for pos in &probe_to_flatten {
543 child_schemas[0].flatten_group(*pos)?;
544 }
545 }
546 for build_idx in 0..n.num_builds() {
547 let child_idx = build_idx + 1;
548 if child_idx >= child_schemas.len() {
549 break;
550 }
551 let to_flatten =
552 n.get_groups_to_flatten_on_build_side(build_idx, &child_schemas[child_idx]);
553 if !to_flatten.is_empty() {
554 self.replace_node_and_flatten(
555 &mut n.deps[child_idx],
556 &to_flatten,
557 &child_schemas[child_idx],
558 )?;
559 for pos in &to_flatten {
560 child_schemas[child_idx].flatten_group(*pos)?;
561 }
562 }
563 }
564 node.compute_factorized_schema(&child_schemas)
565 }
566 LogicalNodeEnum::Flatten(n) => {
567 let child_schema = if let Some(child) = n.input.as_mut() {
568 self.visit_operator(child)?
569 } else {
570 FactorizedSchema::new()
571 };
572 node.compute_factorized_schema(&[child_schema])
573 }
574 LogicalNodeEnum::GetVertices(n) => {
575 let mut child_schema = FactorizedSchema::new();
576 for dep in &mut n.deps {
577 let schema = self.visit_operator(dep)?;
578 if child_schema.num_groups() == 0 {
579 child_schema = schema;
580 }
581 }
582 node.compute_factorized_schema(&[child_schema])
583 }
584 LogicalNodeEnum::GetNeighbors(n) => {
585 let mut child_schema = FactorizedSchema::new();
586 for dep in &mut n.deps {
587 let mut schema = self.visit_operator(dep)?;
588 if let Some(pos) = schema.unflat_group_pos() {
589 let mut to_flatten = HashSet::new();
590 to_flatten.insert(pos);
591 self.replace_node_and_flatten(dep, &to_flatten, &schema)?;
592 schema.flatten_group(pos)?;
593 }
594 if child_schema.num_groups() == 0 {
595 child_schema = schema;
596 }
597 }
598 node.compute_factorized_schema(&[child_schema])
599 }
600 LogicalNodeEnum::GetEdges(_)
601 | LogicalNodeEnum::ScanVertices(_)
602 | LogicalNodeEnum::ScanEdges(_)
603 | LogicalNodeEnum::Start(_)
604 | LogicalNodeEnum::Argument(_)
605 | LogicalNodeEnum::PassThrough(_)
606 | LogicalNodeEnum::BeginTransaction(_)
607 | LogicalNodeEnum::Commit(_)
608 | LogicalNodeEnum::Rollback(_)
609 | LogicalNodeEnum::InsertVertices(_)
610 | LogicalNodeEnum::InsertEdges(_)
611 | LogicalNodeEnum::Update(_)
612 | LogicalNodeEnum::DeleteVertices(_)
613 | LogicalNodeEnum::DeleteEdges(_)
614 | LogicalNodeEnum::DeleteIndex(_)
615 | LogicalNodeEnum::PipeDeleteVertices(_)
616 | LogicalNodeEnum::PipeDeleteEdges(_)
617 | LogicalNodeEnum::CopyFrom(_)
618 | LogicalNodeEnum::CopyTo(_)
619 | LogicalNodeEnum::FulltextSearch(_)
620 | LogicalNodeEnum::FulltextLookup(_)
621 | LogicalNodeEnum::MatchFulltext(_) => node.compute_factorized_schema(&[]),
622 #[cfg(feature = "vector")]
623 LogicalNodeEnum::VectorSearch(_)
624 | LogicalNodeEnum::VectorLookup(_)
625 | LogicalNodeEnum::VectorMatch(_) => node.compute_factorized_schema(&[]),
626 LogicalNodeEnum::Assign(n) => {
627 let mut child_schema = if let Some(child) = n.input.as_mut() {
628 self.visit_operator(child)?
629 } else {
630 FactorizedSchema::new()
631 };
632 for (_, rhs) in &n.assignments {
633 let rhs_id = rhs.id().clone();
634 let mut store = HashMap::new();
635 if let Some(inner) = rhs.get_expression() {
636 store.insert(rhs_id.clone(), inner);
637 }
638 let to_flatten = FlattenAllButOne::get_groups_pos_to_flatten_for_expr(
639 &rhs_id,
640 &child_schema,
641 &store,
642 );
643 if !to_flatten.is_empty() {
644 if let Some(child) = n.input.as_mut() {
645 self.replace_child_and_flatten(child, &to_flatten, &mut child_schema)?;
646 }
647 for pos in &to_flatten {
648 child_schema.flatten_group(*pos)?;
649 }
650 }
651 }
652 node.compute_factorized_schema(&[child_schema])
653 }
654 LogicalNodeEnum::DataCollect(n) => {
655 let mut child_schema = if let Some(child) = n.input.as_mut() {
656 self.visit_operator(child)?
657 } else {
658 FactorizedSchema::new()
659 };
660 self.flatten_barrier_single(n.input.as_mut(), &child_schema)?;
661 child_schema.flatten_all()?;
662 node.compute_factorized_schema(&[child_schema])
663 }
664 LogicalNodeEnum::Materialize(n) => {
665 let mut child_schema = if let Some(child) = n.input.as_mut() {
666 self.visit_operator(child)?
667 } else {
668 FactorizedSchema::new()
669 };
670 self.flatten_barrier_single(n.input.as_mut(), &child_schema)?;
671 child_schema.flatten_all()?;
672 node.compute_factorized_schema(&[child_schema])
673 }
674 LogicalNodeEnum::RollUpApply(n) => {
675 let mut child_schema = if let Some(child) = n.input.as_mut() {
676 self.visit_operator(child)?
677 } else {
678 FactorizedSchema::new()
679 };
680 self.flatten_barrier_single(n.input.as_mut(), &child_schema)?;
681 child_schema.flatten_all()?;
682 node.compute_factorized_schema(&[child_schema])
683 }
684 LogicalNodeEnum::Sample(n) => {
685 let child_schema = if let Some(child) = n.input.as_mut() {
686 self.visit_operator(child)?
687 } else {
688 FactorizedSchema::new()
689 };
690 node.compute_factorized_schema(&[child_schema])
691 }
692 LogicalNodeEnum::Select(n) => {
693 let cond_id = n.condition.id().clone();
694 let mut store = HashMap::new();
695 if let Some(expr) = n.condition.get_expression() {
696 store.insert(cond_id.clone(), expr);
697 }
698 let mut branch_schemas = Vec::new();
699 if let Some(branch) = n.if_branch.as_mut() {
700 let mut schema = self.visit_operator(branch)?;
701 let to_flatten = FlattenAllButOne::get_groups_pos_to_flatten_for_expr(
702 &cond_id, &schema, &store,
703 );
704 if !to_flatten.is_empty() {
705 self.replace_child_and_flatten(branch, &to_flatten, &mut schema)?;
706 for pos in &to_flatten {
707 schema.flatten_group(*pos)?;
708 }
709 branch_schemas.push(schema);
710 } else {
711 branch_schemas.push(schema);
712 }
713 }
714 if let Some(branch) = n.else_branch.as_mut() {
715 let mut schema = self.visit_operator(branch)?;
716 let to_flatten = FlattenAllButOne::get_groups_pos_to_flatten_for_expr(
717 &cond_id, &schema, &store,
718 );
719 if !to_flatten.is_empty() {
720 self.replace_child_and_flatten(branch, &to_flatten, &mut schema)?;
721 for pos in &to_flatten {
722 schema.flatten_group(*pos)?;
723 }
724 branch_schemas.push(schema);
725 } else {
726 branch_schemas.push(schema);
727 }
728 }
729 if let Some(effective) = branch_schemas.into_iter().next() {
730 node.compute_factorized_schema(&[effective])
731 } else {
732 node.compute_factorized_schema(&[])
733 }
734 }
735 LogicalNodeEnum::Loop(n) => {
736 let cond_id = n.condition.id().clone();
737 let mut store = HashMap::new();
738 if let Some(expr) = n.condition.get_expression() {
739 store.insert(cond_id.clone(), expr);
740 }
741 let child_schema = if let Some(body) = n.body.as_mut() {
742 let mut schema = self.visit_operator(body)?;
743 let to_flatten = FlattenAllButOne::get_groups_pos_to_flatten_for_expr(
744 &cond_id, &schema, &store,
745 );
746 if !to_flatten.is_empty() {
747 self.replace_child_and_flatten(body, &to_flatten, &mut schema)?;
748 for pos in &to_flatten {
749 schema.flatten_group(*pos)?;
750 }
751 schema
752 } else {
753 schema
754 }
755 } else {
756 FactorizedSchema::new()
757 };
758 node.compute_factorized_schema(&[child_schema])
759 }
760 LogicalNodeEnum::PatternApply(n) => {
761 let mut left_schema = self.visit_operator(&mut n.left)?;
762 let mut right_schema = self.visit_operator(&mut n.right)?;
763 self.flatten_barrier_binary(
764 &mut n.left,
765 &mut n.right,
766 &left_schema,
767 &right_schema,
768 )?;
769 left_schema.flatten_all()?;
770 right_schema.flatten_all()?;
771 node.compute_factorized_schema(&[left_schema, right_schema])
772 }
773 LogicalNodeEnum::CorrelatedApply(n) => {
774 let mut left_schema = self.visit_operator(&mut n.left)?;
775 let mut right_schema = self.visit_operator(&mut n.right)?;
776 self.flatten_barrier_binary(
777 &mut n.left,
778 &mut n.right,
779 &left_schema,
780 &right_schema,
781 )?;
782 left_schema.flatten_all()?;
783 right_schema.flatten_all()?;
784 node.compute_factorized_schema(&[left_schema, right_schema])
785 }
786 LogicalNodeEnum::Apply(n) => {
787 let mut left_schema = self.visit_operator(n.left_input_mut())?;
788 let mut right_schema = self.visit_operator(n.right_input_mut())?;
789 {
790 let left = n.left_input_mut();
791 self.flatten_barrier_child(left, &left_schema)?;
792 }
793 {
794 let right = n.right_input_mut();
795 self.flatten_barrier_child(right, &right_schema)?;
796 }
797 left_schema.flatten_all()?;
798 right_schema.flatten_all()?;
799 node.compute_factorized_schema(&[left_schema, right_schema])
800 }
801 LogicalNodeEnum::MultiShortestPath(n) => {
802 let mut left_schema = self.visit_operator(&mut n.left)?;
803 let mut right_schema = self.visit_operator(&mut n.right)?;
804 self.flatten_barrier_binary(
805 &mut n.left,
806 &mut n.right,
807 &left_schema,
808 &right_schema,
809 )?;
810 left_schema.flatten_all()?;
811 right_schema.flatten_all()?;
812 node.compute_factorized_schema(&[left_schema, right_schema])
813 }
814 LogicalNodeEnum::BFSShortest(n) => {
815 let mut left_schema = self.visit_operator(&mut n.left)?;
816 let mut right_schema = self.visit_operator(&mut n.right)?;
817 self.flatten_barrier_binary(
818 &mut n.left,
819 &mut n.right,
820 &left_schema,
821 &right_schema,
822 )?;
823 left_schema.flatten_all()?;
824 right_schema.flatten_all()?;
825 node.compute_factorized_schema(&[left_schema, right_schema])
826 }
827 LogicalNodeEnum::AllPaths(n) => {
828 let mut left_schema = self.visit_operator(&mut n.left)?;
829 let mut right_schema = self.visit_operator(&mut n.right)?;
830 self.flatten_barrier_binary(
831 &mut n.left,
832 &mut n.right,
833 &left_schema,
834 &right_schema,
835 )?;
836 left_schema.flatten_all()?;
837 right_schema.flatten_all()?;
838 node.compute_factorized_schema(&[left_schema, right_schema])
839 }
840 LogicalNodeEnum::ShortestPath(n) => {
841 let mut left_schema = self.visit_operator(&mut n.left)?;
842 let mut right_schema = self.visit_operator(&mut n.right)?;
843 self.flatten_barrier_binary(
844 &mut n.left,
845 &mut n.right,
846 &left_schema,
847 &right_schema,
848 )?;
849 left_schema.flatten_all()?;
850 right_schema.flatten_all()?;
851 node.compute_factorized_schema(&[left_schema, right_schema])
852 }
853 }
854 }
855
856 fn build_store_for_project(
857 node: &crate::planning::plan::logical::logical_nodes::operation::LogicalProjectNode,
858 ) -> HashMap<ExpressionId, linkrs_core::Expression> {
859 let mut store = HashMap::new();
860 for col in &node.columns {
861 if let Some(expr) = col.expression.get_expression() {
862 store.insert(col.expression.id().clone(), expr);
863 }
864 }
865 store
866 }
867
868 fn expr_ids_for_project(
869 node: &crate::planning::plan::logical::logical_nodes::operation::LogicalProjectNode,
870 ) -> Vec<linkrs_core::types::expr::ExpressionId> {
871 node.columns
872 .iter()
873 .map(|c| c.expression.id().clone())
874 .collect()
875 }
876
877 fn sort_key_store(
878 node: &crate::planning::plan::logical::logical_nodes::operation::LogicalSortNode,
879 schema: &FactorizedSchema,
880 ) -> (
881 Vec<linkrs_core::types::expr::ExpressionId>,
882 HashMap<ExpressionId, linkrs_core::Expression>,
883 ) {
884 let mut ids = Vec::with_capacity(node.sort_items.len());
885 let mut store = HashMap::new();
886 let mut candidate: u64 = 1 << 63;
887 for item in &node.sort_items {
888 while schema.is_expression_in_scope(&ExpressionId::new(candidate))
889 || store.contains_key(&ExpressionId::new(candidate))
890 {
891 candidate = candidate.wrapping_add(1);
892 }
893 let id = ExpressionId::new(candidate);
894 candidate = candidate.wrapping_add(1);
895 store.insert(id.clone(), item.expression.clone());
896 ids.push(id);
897 }
898 (ids, store)
899 }
900}
901
902impl Default for FactorizationRewriter {
903 fn default() -> Self {
904 Self::new()
905 }
906}