1use crate::optimizer::analysis::BatchPlanAnalysis;
4use crate::optimizer::cost::SelectivityEstimator;
5use crate::optimizer::cost_based::row_estimates::estimate_node_output_rows_corrected;
6use crate::optimizer::stats::feedback::cardinality::CardinalityFeedbackManager;
7use crate::optimizer::stats::feedback::decision::DecorrelationAdvice;
8use crate::optimizer::stats::StatsView;
9use crate::planning::plan::core::nodes::base::plan_node_traits::SingleInputNode;
10use crate::planning::plan::core::nodes::PlanNodeEnum;
11use crate::planning::plan::core::nodes::{PatternApplyNode, SemiJoinNode};
12use linkrs_core::types::operators::BinaryOperator;
13use linkrs_core::Expression;
14
15use super::types::{KeepReason, UnnestDecision, UnnestReason, MAX_BOUNDED_SUBQUERY_ROWS};
16
17#[derive(Debug, Clone)]
22pub struct SubqueryUnnestingOptimizer {
23 max_subquery_rows: u64,
25 max_complexity: u32,
27}
28
29impl Default for SubqueryUnnestingOptimizer {
30 fn default() -> Self {
31 Self::new()
32 }
33}
34
35impl SubqueryUnnestingOptimizer {
36 pub fn new() -> Self {
38 Self {
39 max_subquery_rows: 1000,
40 max_complexity: 50,
41 }
42 }
43
44 pub fn with_max_rows(mut self, max_rows: u64) -> Self {
46 self.max_subquery_rows = max_rows;
47 self
48 }
49
50 pub fn with_max_complexity(mut self, max_complexity: u32) -> Self {
52 self.max_complexity = max_complexity;
53 self
54 }
55
56 pub fn should_unnest(
72 &self,
73 pattern_apply: &PatternApplyNode,
74 analysis: &BatchPlanAnalysis,
75 stats: &StatsView,
76 selectivity: &SelectivityEstimator,
77 cardinality: &CardinalityFeedbackManager,
78 advice: &DecorrelationAdvice,
79 ) -> UnnestDecision {
80 if !analysis.expression_summary.is_fully_deterministic {
82 return UnnestDecision::KeepPatternApply {
83 reason: KeepReason::NonDeterministic,
84 };
85 }
86
87 if analysis.expression_summary.total_complexity > self.max_complexity {
89 return UnnestDecision::KeepPatternApply {
90 reason: KeepReason::ComplexCondition,
91 };
92 }
93
94 if !Self::is_simple_subquery_shape_with_row_bound(pattern_apply.right_input(), i64::MAX) {
98 return UnnestDecision::KeepPatternApply {
99 reason: KeepReason::TooComplex,
100 };
101 }
102
103 if Self::contains_aggregation(pattern_apply.right_input()) {
105 return UnnestDecision::KeepPatternApply {
106 reason: KeepReason::AggregateSubquery,
107 };
108 }
109
110 if advice.prefer_unnest && advice.confidence >= 0.7 {
113 return UnnestDecision::ShouldUnnest {
114 reason: UnnestReason::Empirical,
115 original_cost: 0.0,
116 unnested_cost: 0.0,
117 };
118 }
119 if advice.prefer_keep && advice.confidence >= 0.7 {
120 return UnnestDecision::KeepPatternApply {
121 reason: KeepReason::EmpiricalKeep,
122 };
123 }
124
125 let estimated_rows = self.estimate_subquery_rows(
127 pattern_apply.right_input(),
128 stats,
129 selectivity,
130 cardinality,
131 );
132 if estimated_rows > self.max_subquery_rows {
133 return UnnestDecision::KeepPatternApply {
134 reason: KeepReason::TooManyRows,
135 };
136 }
137
138 let left_rows = estimate_node_output_rows_corrected(
140 pattern_apply.left_input(),
141 stats,
142 selectivity,
143 cardinality,
144 )
145 .max(1) as f64;
146 let original_cost =
147 self.estimate_pattern_apply_cost(left_rows, estimated_rows, advice.apply_cost_per_row);
148 let unnested_cost = self.estimate_hash_join_cost(left_rows, estimated_rows);
149
150 if unnested_cost < original_cost {
151 UnnestDecision::ShouldUnnest {
152 reason: UnnestReason::CostBased,
153 original_cost,
154 unnested_cost,
155 }
156 } else {
157 UnnestDecision::ShouldUnnest {
158 reason: UnnestReason::SimpleSubquery,
159 original_cost,
160 unnested_cost,
161 }
162 }
163 }
164
165 pub fn is_simple_subquery_shape(node: &PlanNodeEnum) -> bool {
173 Self::is_simple_subquery_shape_with_row_bound(node, MAX_BOUNDED_SUBQUERY_ROWS)
174 }
175
176 pub fn is_simple_subquery_shape_with_row_bound(node: &PlanNodeEnum, max_rows: i64) -> bool {
183 match node {
184 PlanNodeEnum::ScanVertices(_)
186 | PlanNodeEnum::ScanEdges(_)
187 | PlanNodeEnum::IndexScan(_) => true,
188
189 PlanNodeEnum::Limit(n) => {
193 let count = n.count();
194 (0..=max_rows).contains(&count)
195 && Self::is_simple_subquery_shape_with_row_bound(
196 SingleInputNode::input(n),
197 max_rows,
198 )
199 }
200
201 PlanNodeEnum::Filter(n) => {
203 let condition = n.condition();
204 if let Some(expr_meta) = condition.expression() {
205 if !Self::is_simple_equality_condition(expr_meta.inner()) {
206 return false;
207 }
208 }
209 Self::is_simple_subquery_shape_with_row_bound(SingleInputNode::input(n), max_rows)
210 }
211
212 PlanNodeEnum::Project(n) => {
214 Self::is_simple_subquery_shape_with_row_bound(SingleInputNode::input(n), max_rows)
215 }
216
217 PlanNodeEnum::InnerJoin(n) => {
222 Self::is_simple_subquery_shape_with_row_bound(n.left_input(), max_rows)
223 && Self::is_simple_subquery_shape_with_row_bound(n.right_input(), max_rows)
224 }
225
226 _ => false,
228 }
229 }
230
231 pub fn is_mark_join_shape(node: &PlanNodeEnum) -> bool {
241 match node {
242 PlanNodeEnum::ScanVertices(_)
244 | PlanNodeEnum::ScanEdges(_)
245 | PlanNodeEnum::IndexScan(_) => true,
246
247 PlanNodeEnum::Limit(n) => {
249 let count = n.count();
250 (0..=MAX_BOUNDED_SUBQUERY_ROWS).contains(&count)
251 && Self::is_mark_join_shape(SingleInputNode::input(n))
252 }
253
254 PlanNodeEnum::Filter(n) => Self::is_mark_join_shape(SingleInputNode::input(n)),
256
257 PlanNodeEnum::Project(n) => Self::is_mark_join_shape(SingleInputNode::input(n)),
259
260 _ => false,
262 }
263 }
264
265 pub fn has_bounded_input(node: &PlanNodeEnum) -> bool {
269 match node {
270 PlanNodeEnum::ScanVertices(_)
271 | PlanNodeEnum::ScanEdges(_)
272 | PlanNodeEnum::IndexScan(_) => true,
273 PlanNodeEnum::Limit(n) => {
274 let count = n.count();
275 (0..=MAX_BOUNDED_SUBQUERY_ROWS).contains(&count)
276 }
277 _ => false,
278 }
279 }
280
281 pub fn contains_aggregation(node: &PlanNodeEnum) -> bool {
285 match node {
286 PlanNodeEnum::Aggregate(_) | PlanNodeEnum::Window(_) => true,
287 _ => node.children().into_iter().any(Self::contains_aggregation),
288 }
289 }
290
291 fn is_simple_equality_condition(expr: &Expression) -> bool {
295 match expr {
296 Expression::Binary { op, left, right } => match op {
297 BinaryOperator::Equal => {
298 Self::is_simple_expression(left.as_ref())
299 && Self::is_simple_expression(right.as_ref())
300 }
301 BinaryOperator::And | BinaryOperator::Or => {
302 Self::is_simple_equality_condition(left.as_ref())
303 && Self::is_simple_equality_condition(right.as_ref())
304 }
305 _ => false,
306 },
307 _ => false,
308 }
309 }
310
311 fn is_simple_expression(expr: &Expression) -> bool {
314 matches!(
315 expr,
316 Expression::Literal(_) | Expression::Variable(_) | Expression::Property { .. }
317 )
318 }
319
320 fn estimate_subquery_rows(
327 &self,
328 node: &PlanNodeEnum,
329 stats: &StatsView,
330 selectivity: &SelectivityEstimator,
331 cardinality: &CardinalityFeedbackManager,
332 ) -> u64 {
333 estimate_node_output_rows_corrected(node, stats, selectivity, cardinality)
334 }
335
336 fn estimate_pattern_apply_cost(
341 &self,
342 left_rows: f64,
343 subquery_rows: u64,
344 measured_per_row_cost: Option<f64>,
345 ) -> f64 {
346 match measured_per_row_cost {
350 Some(cost) => left_rows * cost * 0.001,
351 None => left_rows * (subquery_rows as f64 * 0.1),
352 }
353 }
354
355 fn estimate_hash_join_cost(&self, left_rows: f64, subquery_rows: u64) -> f64 {
357 let right_rows = subquery_rows as f64;
359
360 let build_cost = right_rows;
362 let probe_cost = left_rows * 0.5; build_cost + probe_cost
365 }
366
367 pub fn unnest(
375 &self,
376 pattern_apply: PatternApplyNode,
377 ) -> Result<PlanNodeEnum, crate::planning::planner::PlannerError> {
378 Self::build_semi_join_from_pattern_apply(pattern_apply)
379 }
380
381 pub fn build_semi_join_from_pattern_apply(
393 pattern_apply: PatternApplyNode,
394 ) -> Result<PlanNodeEnum, crate::planning::planner::PlannerError> {
395 let left_input = pattern_apply.left_input().clone();
396 let right_input = pattern_apply.right_input().clone();
397
398 let join_node = if pattern_apply.is_anti_predicate() {
399 SemiJoinNode::new_anti(
400 left_input,
401 right_input,
402 pattern_apply.hash_keys().to_vec(),
403 pattern_apply.probe_keys().to_vec(),
404 )?
405 } else {
406 SemiJoinNode::new_semi(
407 left_input,
408 right_input,
409 pattern_apply.hash_keys().to_vec(),
410 pattern_apply.probe_keys().to_vec(),
411 )?
412 };
413
414 Ok(PlanNodeEnum::SemiJoin(join_node))
415 }
416}
417
418#[cfg(test)]
419mod tests {
420 use super::*;
421 use crate::optimizer::analysis::BatchPlanAnalyzer;
422 use crate::optimizer::stats::StatisticsManager;
423 use crate::planning::plan::core::nodes::access::graph_scan_node::ScanVerticesNode;
424 use crate::planning::plan::core::nodes::operation::filter_node::FilterNode;
425 use crate::planning::plan::core::nodes::operation::sort_node::LimitNode;
426 use crate::planning::plan::core::nodes::PlanNodeEnum;
427 use linkrs_core::types::expr::expression_context::ExpressionAnalysisContext;
428 use linkrs_core::types::expr::ExpressionMeta;
429 use linkrs_core::types::ContextualExpression;
430 use std::sync::Arc;
431
432 fn test_scan() -> PlanNodeEnum {
433 let mut scan = ScanVerticesNode::new(1, "test");
434 scan.set_tag("person");
435 PlanNodeEnum::ScanVertices(scan)
436 }
437
438 fn test_selectivity() -> SelectivityEstimator {
439 SelectivityEstimator::new(std::sync::Arc::new(StatisticsManager::new()))
440 }
441
442 #[test]
443 fn test_optimizer_creation() {
444 let optimizer = SubqueryUnnestingOptimizer::new();
445 assert_eq!(optimizer.max_subquery_rows, 1000);
446 assert_eq!(optimizer.max_complexity, 50);
447 }
448
449 #[test]
450 fn test_optimizer_with_config() {
451 let _optimizer = SubqueryUnnestingOptimizer::new()
452 .with_max_rows(500)
453 .with_max_complexity(30);
454 }
455
456 #[test]
457 fn test_simple_expression_check() {
458 let literal = Expression::Literal(linkrs_core::Value::Int(42));
459 assert!(SubqueryUnnestingOptimizer::is_simple_expression(&literal));
460
461 let variable = Expression::Variable("n".to_string());
462 assert!(SubqueryUnnestingOptimizer::is_simple_expression(&variable));
463
464 let property = Expression::Property {
465 object: Box::new(Expression::Variable("n".to_string())),
466 property: "name".to_string(),
467 };
468 assert!(SubqueryUnnestingOptimizer::is_simple_expression(&property));
469
470 let binary = Expression::Binary {
471 left: Box::new(Expression::Literal(linkrs_core::Value::Int(1))),
472 op: linkrs_core::types::operators::BinaryOperator::Add,
473 right: Box::new(Expression::Literal(linkrs_core::Value::Int(2))),
474 };
475 assert!(!SubqueryUnnestingOptimizer::is_simple_expression(&binary));
476 }
477
478 #[test]
479 fn test_shape_accepts_scans_and_limits() {
480 assert!(SubqueryUnnestingOptimizer::is_simple_subquery_shape(
481 &test_scan()
482 ));
483
484 let limit = LimitNode::new(test_scan(), 0, 10).expect("limit should build");
485 assert!(SubqueryUnnestingOptimizer::is_simple_subquery_shape(
486 &PlanNodeEnum::Limit(limit)
487 ));
488 assert!(SubqueryUnnestingOptimizer::has_bounded_input(
489 &PlanNodeEnum::Limit(LimitNode::new(test_scan(), 0, 10).expect("limit should build"))
490 ));
491
492 let large = LimitNode::new(test_scan(), 0, 5000).expect("limit should build");
493 assert!(!SubqueryUnnestingOptimizer::has_bounded_input(
494 &PlanNodeEnum::Limit(large)
495 ));
496 }
497
498 #[test]
499 fn test_shape_rejects_aggregation() {
500 use crate::planning::plan::core::nodes::AggregateNode;
501 let aggregate =
502 AggregateNode::new(test_scan(), vec![], vec![]).expect("aggregate should build");
503 assert!(SubqueryUnnestingOptimizer::contains_aggregation(
504 &PlanNodeEnum::Aggregate(aggregate)
505 ));
506 }
507
508 #[test]
509 fn test_shape_accepts_disjunctive_equality_filter() {
510 use linkrs_core::types::expr::expression_context::ExpressionAnalysisContext;
511 use linkrs_core::types::ContextualExpression;
512 let condition = {
513 let ctx = Arc::new(ExpressionAnalysisContext::new());
514 let prop = |name: &str| Expression::Property {
515 object: Box::new(Expression::Variable("n".to_string())),
516 property: name.to_string(),
517 };
518 let expr = Expression::Binary {
519 left: Box::new(Expression::Binary {
520 left: Box::new(prop("city")),
521 op: BinaryOperator::Equal,
522 right: Box::new(Expression::Literal(linkrs_core::Value::Int(1))),
523 }),
524 op: BinaryOperator::Or,
525 right: Box::new(Expression::Binary {
526 left: Box::new(prop("age")),
527 op: BinaryOperator::Equal,
528 right: Box::new(Expression::Literal(linkrs_core::Value::Int(2))),
529 }),
530 };
531 let id = ctx.register_expression(linkrs_core::types::expr::ExpressionMeta::new(expr));
532 ContextualExpression::new(id, ctx)
533 };
534 let or_filter = PlanNodeEnum::Filter(
536 FilterNode::new(test_scan(), condition.clone()).expect("filter should build"),
537 );
538 assert!(SubqueryUnnestingOptimizer::is_simple_subquery_shape(
539 &or_filter
540 ));
541
542 let range_condition = {
544 let ctx = Arc::new(ExpressionAnalysisContext::new());
545 let expr = Expression::Binary {
546 left: Box::new(Expression::Property {
547 object: Box::new(Expression::Variable("n".to_string())),
548 property: "age".to_string(),
549 }),
550 op: BinaryOperator::GreaterThan,
551 right: Box::new(Expression::Literal(linkrs_core::Value::Int(30))),
552 };
553 let id = ctx.register_expression(linkrs_core::types::expr::ExpressionMeta::new(expr));
554 ContextualExpression::new(id, ctx)
555 };
556 let range_filter = PlanNodeEnum::Filter(
557 FilterNode::new(test_scan(), range_condition).expect("filter should build"),
558 );
559 assert!(!SubqueryUnnestingOptimizer::is_simple_subquery_shape(
560 &range_filter
561 ));
562 }
563
564 #[test]
565 fn test_shape_accepts_two_table_inner_join() {
566 use crate::planning::plan::core::nodes::join::join_node::InnerJoinNode;
567 use linkrs_core::types::expr::expression_context::ExpressionAnalysisContext;
568 use linkrs_core::types::ContextualExpression;
569 let key = || {
570 let ctx = Arc::new(ExpressionAnalysisContext::new());
571 let expr = Expression::Binary {
572 left: Box::new(Expression::Property {
573 object: Box::new(Expression::Variable("a".to_string())),
574 property: "id".to_string(),
575 }),
576 op: BinaryOperator::Equal,
577 right: Box::new(Expression::Property {
578 object: Box::new(Expression::Variable("b".to_string())),
579 property: "id".to_string(),
580 }),
581 };
582 let id = ctx.register_expression(linkrs_core::types::expr::ExpressionMeta::new(expr));
583 ContextualExpression::new(id, ctx)
584 };
585 let join = PlanNodeEnum::InnerJoin(
586 InnerJoinNode::new(test_scan(), test_scan(), vec![key()], vec![key()])
587 .expect("inner join should build"),
588 );
589 assert!(SubqueryUnnestingOptimizer::is_simple_subquery_shape(&join));
590
591 use crate::planning::plan::core::nodes::AggregateNode;
593 let agg_side = PlanNodeEnum::Aggregate(
594 AggregateNode::new(test_scan(), vec![], vec![]).expect("aggregate should build"),
595 );
596 let bad_join = PlanNodeEnum::InnerJoin(
597 InnerJoinNode::new(agg_side, test_scan(), vec![key()], vec![key()])
598 .expect("inner join should build"),
599 );
600 assert!(!SubqueryUnnestingOptimizer::is_simple_subquery_shape(
601 &bad_join
602 ));
603 }
604
605 #[test]
606 fn test_row_bound_is_configurable_for_cbo() {
607 let large = PlanNodeEnum::Limit(
610 crate::planning::plan::core::nodes::operation::sort_node::LimitNode::new(
611 test_scan(),
612 0,
613 5000,
614 )
615 .expect("limit should build"),
616 );
617 assert!(!SubqueryUnnestingOptimizer::is_simple_subquery_shape(
618 &large
619 ));
620 assert!(
621 SubqueryUnnestingOptimizer::is_simple_subquery_shape_with_row_bound(&large, i64::MAX)
622 );
623 }
624
625 #[test]
626 fn test_simple_subquery_unnests() {
627 let scan = test_scan();
628 let condition = {
629 let ctx = Arc::new(ExpressionAnalysisContext::new());
630 let expr = Expression::Binary {
631 left: Box::new(Expression::Property {
632 object: Box::new(Expression::Variable("n".to_string())),
633 property: "age".to_string(),
634 }),
635 op: BinaryOperator::Equal,
636 right: Box::new(Expression::Literal(linkrs_core::Value::Int(18))),
637 };
638 let id = ctx.register_expression(ExpressionMeta::new(expr));
639 ContextualExpression::new(id, ctx)
640 };
641 let filter = PlanNodeEnum::Filter(FilterNode::new(scan, condition).expect("filter"));
642
643 let pattern_apply = PatternApplyNode::new(test_scan(), filter, vec![], vec![], false)
644 .expect("pattern apply should build");
645
646 let analysis =
647 BatchPlanAnalyzer::new().analyze(&PlanNodeEnum::PatternApply(pattern_apply.clone()));
648 let selectivity = test_selectivity();
649 let stats_manager = StatisticsManager::new();
650 let stats = StatsView::new(&stats_manager, Some("test"));
651 let decision = SubqueryUnnestingOptimizer::new().should_unnest(
652 &pattern_apply,
653 &analysis,
654 &stats,
655 &selectivity,
656 &CardinalityFeedbackManager::new(),
657 &DecorrelationAdvice::default(),
658 );
659 assert!(
660 matches!(decision, UnnestDecision::ShouldUnnest { .. }),
661 "a deterministic single-scan equality-filtered subquery must unnest"
662 );
663 }
664
665 #[test]
666 fn test_unnest_produces_semi_join_with_split_keys() {
667 let ctx = Arc::new(ExpressionAnalysisContext::new());
668 let hash_key_expr = Expression::Property {
669 object: Box::new(Expression::Variable("t".to_string())),
670 property: "city".to_string(),
671 };
672 let probe_key_expr = Expression::Property {
673 object: Box::new(Expression::Variable("p".to_string())),
674 property: "city".to_string(),
675 };
676 let hash_key_id = ctx.register_expression(ExpressionMeta::new(hash_key_expr));
677 let hash_key = ContextualExpression::new(hash_key_id, ctx.clone());
678 let probe_key_id = ctx.register_expression(ExpressionMeta::new(probe_key_expr));
679 let probe_key = ContextualExpression::new(probe_key_id, ctx);
680
681 let pattern_apply = PatternApplyNode::new(
682 test_scan(),
683 test_scan(),
684 vec![hash_key],
685 vec![probe_key],
686 false,
687 )
688 .expect("pattern apply should build");
689
690 let transformed = SubqueryUnnestingOptimizer::new()
691 .unnest(pattern_apply)
692 .expect("unnest should succeed");
693 match &transformed {
694 PlanNodeEnum::SemiJoin(join) => {
695 assert!(!join.is_anti());
696 assert_eq!(join.hash_keys().len(), 1);
697 assert_eq!(join.probe_keys().len(), 1);
698 let hash_meta = join.hash_keys()[0].expression().expect("hash key meta");
699 let probe_meta = join.probe_keys()[0].expression().expect("probe key meta");
700 assert_eq!(
701 hash_meta.inner(),
702 &Expression::Property {
703 object: Box::new(Expression::Variable("t".to_string())),
704 property: "city".to_string(),
705 }
706 );
707 assert_eq!(
708 probe_meta.inner(),
709 &Expression::Property {
710 object: Box::new(Expression::Variable("p".to_string())),
711 property: "city".to_string(),
712 }
713 );
714 }
715 _ => panic!("expected SemiJoin, got {:?}", transformed.name()),
716 }
717 }
718
719 #[test]
720 fn test_unnest_anti_produces_anti_join() {
721 let pattern_apply = PatternApplyNode::new(test_scan(), test_scan(), vec![], vec![], true)
722 .expect("pattern apply should build");
723
724 let transformed = SubqueryUnnestingOptimizer::new()
725 .unnest(pattern_apply)
726 .expect("unnest should succeed");
727 match &transformed {
728 PlanNodeEnum::SemiJoin(join) => assert!(join.is_anti()),
729 _ => panic!("expected AntiJoin, got {:?}", transformed.name()),
730 }
731 }
732}