1use std::collections::{HashMap, HashSet};
4use std::sync::Arc;
5
6use crate::optimizer::cost::CostCalculator;
7use crate::optimizer::cost_based::ndv::DEFAULT_JOIN_SELECTIVITY;
8use crate::optimizer::decision::{JoinAlgorithm, JoinOrderDecision};
9
10use super::condition::JoinCondition;
11use super::result::{JoinOrderResult, OptimizationMethod};
12use super::table_info::TableInfo;
13
14#[derive(Debug)]
16pub struct JoinOrderOptimizer {
17 cost_calculator: Arc<CostCalculator>,
18 dp_threshold: usize,
20}
21
22#[derive(Debug, Clone)]
24struct SubproblemSolution {
25 pub table_set: u32,
27 pub last_table: String,
29 pub total_cost: f64,
31 pub output_rows: u64,
33 pub join_tree: String,
35}
36
37impl JoinOrderOptimizer {
38 pub fn new(cost_calculator: Arc<CostCalculator>) -> Self {
40 Self {
41 cost_calculator,
42 dp_threshold: 8, }
44 }
45
46 pub fn with_dp_threshold(mut self, threshold: usize) -> Self {
48 self.dp_threshold = threshold;
49 self
50 }
51
52 pub fn optimize_join_order(
61 &self,
62 tables: &[TableInfo],
63 conditions: &[JoinCondition],
64 ) -> JoinOrderResult {
65 if tables.len() <= 1 {
66 return JoinOrderResult {
67 order: tables.iter().map(|t| t.id.clone()).collect(),
68 algorithms: Vec::new(),
69 total_cost: 0.0,
70 final_output_rows: tables.first().map(|t| t.estimated_rows).unwrap_or(0),
71 optimization_method: OptimizationMethod::Heuristic,
72 };
73 }
74
75 if tables.len() <= self.dp_threshold {
76 self.optimize_with_dp(tables, conditions)
77 } else {
78 self.optimize_with_greedy(tables, conditions)
79 }
80 }
81
82 pub fn cost_of_order(
89 &self,
90 tables: &[TableInfo],
91 conditions: &[JoinCondition],
92 order: &[String],
93 ) -> Option<f64> {
94 if order.len() != tables.len() {
95 return None;
96 }
97 let bit_to_table: HashMap<u32, &TableInfo> =
98 tables.iter().map(|t| (1u32 << t.bit_id, t)).collect();
99 let id_to_table: HashMap<&str, &TableInfo> =
100 tables.iter().map(|t| (t.id.as_str(), t)).collect();
101 let condition_map = self.build_condition_map(conditions);
102
103 let mut accumulated_mask = 0u32;
104 let mut accumulated_rows = 0u64;
105 let mut total_cost = 0.0;
106 for (position, id) in order.iter().enumerate() {
107 let table = id_to_table.get(id.as_str())?;
108 if position == 0 {
109 accumulated_mask |= 1 << table.bit_id;
110 accumulated_rows = table.estimated_rows;
111 continue;
112 }
113 let table_bit = 1 << table.bit_id;
114 if accumulated_mask & table_bit != 0 {
115 return None;
116 }
117 let selectivity = self.estimate_selectivity_for_sets(
118 accumulated_mask,
119 table_bit,
120 &bit_to_table,
121 &condition_map,
122 );
123 total_cost += self
124 .cost_calculator
125 .calculate_hash_join_cost(accumulated_rows, table.estimated_rows);
126 accumulated_rows =
127 ((accumulated_rows as f64 * table.estimated_rows as f64 * selectivity) as u64)
128 .max(1);
129 accumulated_mask |= table_bit;
130 }
131 Some(total_cost)
132 }
133
134 fn optimize_with_dp(
136 &self,
137 tables: &[TableInfo],
138 conditions: &[JoinCondition],
139 ) -> JoinOrderResult {
140 let n = tables.len();
141
142 let condition_map = self.build_condition_map(conditions);
144
145 let mut dp: HashMap<u32, SubproblemSolution> = HashMap::new();
147
148 for table in tables {
150 let solution = SubproblemSolution {
151 table_set: 1 << table.bit_id,
152 last_table: table.id.clone(),
153 total_cost: 0.0,
154 output_rows: table.estimated_rows,
155 join_tree: table.id.clone(),
156 };
157 dp.insert(solution.table_set, solution);
158 }
159
160 let bit_to_table: std::collections::HashMap<u32, &TableInfo> =
162 tables.iter().map(|t| (1u32 << t.bit_id, t)).collect();
163
164 for subset_size in 2..=n {
170 for subset in self.generate_subsets(n, subset_size) {
171 let mut best_solution: Option<SubproblemSolution> = None;
172
173 for table in tables {
174 let table_bit = 1 << table.bit_id;
175 if subset & table_bit == 0 {
176 continue;
177 }
178 let remaining = subset ^ table_bit;
179 if remaining == 0 {
180 continue;
181 }
182 let Some(left_solution) = dp.get(&remaining) else {
183 continue;
184 };
185
186 let selectivity = self.estimate_selectivity_for_sets(
190 remaining,
191 table_bit,
192 &bit_to_table,
193 &condition_map,
194 );
195 let join_cost = self
196 .cost_calculator
197 .calculate_hash_join_cost(left_solution.output_rows, table.estimated_rows);
198 let output_rows = ((left_solution.output_rows as f64
199 * table.estimated_rows as f64
200 * selectivity) as u64)
201 .max(1);
202 let total_cost = left_solution.total_cost + join_cost;
203
204 let solution = SubproblemSolution {
205 table_set: subset,
206 last_table: table.id.clone(),
207 total_cost,
208 output_rows,
209 join_tree: format!("Join({}, {})", left_solution.join_tree, table.id),
210 };
211 if best_solution
212 .as_ref()
213 .is_none_or(|best| solution.total_cost < best.total_cost)
214 {
215 best_solution = Some(solution);
216 }
217 }
218
219 if let Some(solution) = best_solution {
220 dp.insert(subset, solution);
221 }
222 }
223 }
224
225 let full_set = (1 << n) - 1;
227 let best_solution = dp
228 .get(&full_set)
229 .cloned()
230 .unwrap_or_else(|| SubproblemSolution {
231 table_set: full_set,
232 last_table: tables
233 .last()
234 .expect("tables collection is not empty")
235 .id
236 .clone(),
237 total_cost: f64::MAX,
238 output_rows: 0,
239 join_tree: "fallback".to_string(),
240 });
241
242 let order = self.reconstruct_order(&best_solution, &dp, tables);
244 let algorithms = self.select_algorithms(&order, conditions, tables);
245
246 JoinOrderResult {
247 order,
248 algorithms,
249 total_cost: best_solution.total_cost,
250 final_output_rows: best_solution.output_rows,
251 optimization_method: OptimizationMethod::DynamicProgramming,
252 }
253 }
254
255 fn optimize_with_greedy(
257 &self,
258 tables: &[TableInfo],
259 conditions: &[JoinCondition],
260 ) -> JoinOrderResult {
261 let condition_map = self.build_condition_map(conditions);
262 let mut remaining: HashSet<String> = tables.iter().map(|t| t.id.clone()).collect();
263 let table_map: HashMap<String, &TableInfo> =
264 tables.iter().map(|t| (t.id.clone(), t)).collect();
265
266 let mut order = Vec::new();
267 let mut algorithms = Vec::new();
268 let mut total_cost = 0.0;
269 let mut current_rows = 0u64;
270
271 if let Some(start_table) = tables.iter().min_by_key(|t| t.estimated_rows) {
273 order.push(start_table.id.clone());
274 remaining.remove(&start_table.id);
275 current_rows = start_table.estimated_rows;
276 }
277
278 while !remaining.is_empty() {
280 let mut best_next: Option<(String, f64, u64)> = None;
281
282 for table_id in &remaining {
283 if let Some(table) = table_map.get(table_id) {
284 let (cost, output_rows) = self.calculate_join_cost(
285 current_rows,
286 table.estimated_rows,
287 &table.id,
288 &condition_map,
289 );
290
291 if best_next
292 .as_ref()
293 .is_none_or(|(_, best_cost, _)| cost < *best_cost)
294 {
295 best_next = Some((table_id.clone(), cost, output_rows));
296 }
297 }
298 }
299
300 if let Some((next_id, cost, output_rows)) = best_next {
301 let current_has_index = order
303 .last()
304 .and_then(|id| table_map.get(id))
305 .map(|t| t.has_index)
306 .unwrap_or(false);
307 let next_has_index = table_map
308 .get(&next_id)
309 .map(|t| t.has_index)
310 .unwrap_or(false);
311 let current_id = order.last().cloned().unwrap_or_default();
312
313 order.push(next_id.clone());
314 remaining.remove(&next_id);
315 total_cost += cost;
316
317 let algorithm = self.select_algorithm(
319 current_rows,
320 output_rows,
321 current_has_index,
322 next_has_index,
323 ¤t_id,
324 &next_id,
325 );
326 algorithms.push(algorithm);
327
328 current_rows = output_rows;
329 } else {
330 break;
331 }
332 }
333
334 JoinOrderResult {
335 order,
336 algorithms,
337 total_cost,
338 final_output_rows: current_rows,
339 optimization_method: OptimizationMethod::Greedy,
340 }
341 }
342
343 fn build_condition_map(&self, conditions: &[JoinCondition]) -> HashMap<(String, String), f64> {
345 let mut map = HashMap::new();
346 for cond in conditions {
347 let key = (cond.left_table.clone(), cond.right_table.clone());
348 let reverse_key = (cond.right_table.clone(), cond.left_table.clone());
349 map.insert(key, cond.selectivity);
350 map.insert(reverse_key, cond.selectivity);
351 }
352 map
353 }
354
355 fn generate_subsets(&self, n: usize, k: usize) -> Vec<u32> {
357 let mut result = Vec::new();
358 self.generate_subsets_recursive(0, n, k, 0, &mut result);
359 result
360 }
361
362 fn generate_subsets_recursive(
363 &self,
364 start: usize,
365 n: usize,
366 k: usize,
367 current: u32,
368 result: &mut Vec<u32>,
369 ) {
370 if k == 0 {
371 result.push(current);
372 return;
373 }
374 if start >= n {
375 return;
376 }
377 for i in start..n {
378 self.generate_subsets_recursive(i + 1, n, k - 1, current | (1 << i), result);
379 }
380 }
381
382 fn estimate_selectivity_for_sets(
390 &self,
391 left_set: u32,
392 right_set: u32,
393 bit_to_table: &HashMap<u32, &TableInfo>,
394 condition_map: &HashMap<(String, String), f64>,
395 ) -> f64 {
396 let table_id_for_bit =
398 |bit: u32| -> Option<String> { bit_to_table.get(&bit).map(|t| t.id.clone()) };
399 let left_ids: Vec<String> = {
400 let mut ids = Vec::new();
401 let mut mask = left_set;
402 while mask != 0 {
403 let bit = mask & mask.wrapping_neg();
404 if let Some(id) = table_id_for_bit(bit) {
405 ids.push(id);
406 }
407 mask ^= bit;
408 }
409 ids
410 };
411 let right_ids: Vec<String> = {
412 let mut ids = Vec::new();
413 let mut mask = right_set;
414 while mask != 0 {
415 let bit = mask & mask.wrapping_neg();
416 if let Some(id) = table_id_for_bit(bit) {
417 ids.push(id);
418 }
419 mask ^= bit;
420 }
421 ids
422 };
423 let mut crossing: Vec<f64> = Vec::new();
424 for l in &left_ids {
425 for r in &right_ids {
426 if let Some(sel) = condition_map.get(&(l.clone(), r.clone())) {
427 crossing.push(*sel);
428 }
429 }
430 }
431 if crossing.is_empty() {
432 1.0
434 } else {
435 let mut prod = 1.0;
437 for s in crossing {
438 prod *= s;
439 }
440 prod.clamp(0.001, 1.0)
441 }
442 }
443
444 fn calculate_join_cost(
450 &self,
451 left_rows: u64,
452 right_rows: u64,
453 right_table: &str,
454 condition_map: &HashMap<(String, String), f64>,
455 ) -> (f64, u64) {
456 let selectivity = condition_map
457 .iter()
458 .filter(|((l, r), _)| l == right_table || r == right_table)
459 .map(|(_, s)| *s)
460 .fold(f64::INFINITY, f64::min)
461 .clamp(0.001, 1.0);
462 let selectivity = if selectivity.is_finite() {
463 selectivity
464 } else {
465 DEFAULT_JOIN_SELECTIVITY
466 };
467
468 let output_rows = ((left_rows as f64 * right_rows as f64 * selectivity) as u64).max(1);
469 let cost = self
470 .cost_calculator
471 .calculate_hash_join_cost(left_rows, right_rows);
472 (cost, output_rows)
473 }
474
475 fn select_algorithm(
482 &self,
483 left_rows: u64,
484 right_rows: u64,
485 left_has_index: bool,
486 right_has_index: bool,
487 left_id: &str,
488 right_id: &str,
489 ) -> JoinAlgorithm {
490 const NESTED_LOOP_MAX_ROWS: u64 = 100; const INDEX_JOIN_MAX_ROWS: u64 = 10000; if left_has_index && right_rows <= INDEX_JOIN_MAX_ROWS {
496 return JoinAlgorithm::IndexJoin {
497 indexed_side: left_id.to_string(),
498 };
499 }
500 if right_has_index && left_rows <= INDEX_JOIN_MAX_ROWS {
501 return JoinAlgorithm::IndexJoin {
502 indexed_side: right_id.to_string(),
503 };
504 }
505
506 if left_rows <= NESTED_LOOP_MAX_ROWS && right_rows <= NESTED_LOOP_MAX_ROWS {
508 return JoinAlgorithm::NestedLoopJoin {
509 outer: left_id.to_string(),
510 inner: right_id.to_string(),
511 };
512 }
513
514 if left_rows <= right_rows {
516 JoinAlgorithm::HashJoin {
517 build_side: left_id.to_string(),
518 probe_side: right_id.to_string(),
519 }
520 } else {
521 JoinAlgorithm::HashJoin {
522 build_side: right_id.to_string(),
523 probe_side: left_id.to_string(),
524 }
525 }
526 }
527
528 fn reconstruct_order(
535 &self,
536 solution: &SubproblemSolution,
537 dp: &HashMap<u32, SubproblemSolution>,
538 tables: &[TableInfo],
539 ) -> Vec<String> {
540 let mut order = Vec::new();
541 let mut current_set = solution.table_set;
542 while current_set != 0 {
543 if let Some(sol) = dp.get(¤t_set) {
544 order.push(sol.last_table.clone());
545 if let Some(table) = tables.iter().find(|t| t.id == sol.last_table) {
546 current_set &= !(1 << table.bit_id);
547 } else {
548 break;
549 }
550 } else {
551 break;
552 }
553 }
554 order.reverse();
555 if order.len() != tables.len() {
559 for t in tables {
560 if !order.contains(&t.id) {
561 order.push(t.id.clone());
562 }
563 }
564 }
565 order
566 }
567
568 fn select_algorithms(
570 &self,
571 order: &[String],
572 _conditions: &[JoinCondition],
573 tables: &[TableInfo],
574 ) -> Vec<JoinAlgorithm> {
575 let mut algorithms = Vec::new();
576 let table_map: HashMap<String, &TableInfo> =
577 tables.iter().map(|t| (t.id.clone(), t)).collect();
578
579 for i in 1..order.len() {
580 let left = &order[i - 1];
581 let right = &order[i];
582
583 let left_info = table_map.get(left);
585 let right_info = table_map.get(right);
586
587 let left_rows = left_info.map(|t| t.estimated_rows).unwrap_or(0);
588 let right_rows = right_info.map(|t| t.estimated_rows).unwrap_or(0);
589 let left_has_index = left_info.map(|t| t.has_index).unwrap_or(false);
590 let right_has_index = right_info.map(|t| t.has_index).unwrap_or(false);
591
592 let algorithm = self.select_algorithm(
594 left_rows,
595 right_rows,
596 left_has_index,
597 right_has_index,
598 left,
599 right,
600 );
601
602 algorithms.push(algorithm);
603 }
604
605 algorithms
606 }
607
608 pub fn to_decision(&self, result: &JoinOrderResult) -> JoinOrderDecision {
610 let mut decision = JoinOrderDecision::empty();
611
612 for (i, table) in result.order.iter().enumerate() {
613 if i < result.algorithms.len() {
614 decision.add_join_step(table.clone(), result.algorithms[i].clone());
615 } else {
616 decision.add_join_step(
617 table.clone(),
618 JoinAlgorithm::HashJoin {
619 build_side: "default".to_string(),
620 probe_side: "default".to_string(),
621 },
622 );
623 }
624 }
625
626 decision
627 }
628}
629
630#[cfg(test)]
631mod tests {
632 use super::*;
633 use crate::optimizer::stats::StatisticsManager;
634
635 fn create_test_optimizer() -> JoinOrderOptimizer {
636 let stats_manager = Arc::new(StatisticsManager::new());
637 let cost_calculator = Arc::new(CostCalculator::new(stats_manager));
638 JoinOrderOptimizer::new(cost_calculator)
639 }
640
641 fn create_test_tables() -> Vec<TableInfo> {
642 vec![
643 TableInfo::new("A".into(), 1000).with_bit_id(0),
644 TableInfo::new("B".into(), 500).with_bit_id(1),
645 TableInfo::new("C".into(), 2000).with_bit_id(2),
646 ]
647 }
648
649 #[test]
650 fn test_single_table() {
651 let optimizer = create_test_optimizer();
652 let tables = vec![TableInfo::new("A".into(), 1000)];
653 let result = optimizer.optimize_join_order(&tables, &[]);
654
655 assert_eq!(result.order.len(), 1);
656 assert_eq!(result.order[0], "A");
657 assert_eq!(result.total_cost, 0.0);
658 }
659
660 #[test]
661 fn test_two_tables() {
662 let optimizer = create_test_optimizer();
663 let tables = vec![
664 TableInfo::new("A".into(), 1000).with_bit_id(0),
665 TableInfo::new("B".into(), 500).with_bit_id(1),
666 ];
667 let conditions = vec![JoinCondition::new("A".into(), "B".to_string())];
668
669 let result = optimizer.optimize_join_order(&tables, &conditions);
670
671 assert_eq!(result.order.len(), 2);
672 assert!(!result.algorithms.is_empty());
673 }
674
675 #[test]
676 fn test_dp_vs_greedy() {
677 let optimizer = create_test_optimizer();
678 let tables = create_test_tables();
679 let conditions = vec![
680 JoinCondition::new("A".into(), "B".to_string()).with_selectivity(0.1),
681 JoinCondition::new("B".into(), "C".to_string()).with_selectivity(0.2),
682 ];
683
684 let result = optimizer.optimize_join_order(&tables, &conditions);
686 assert_eq!(
687 result.optimization_method,
688 OptimizationMethod::DynamicProgramming
689 );
690 assert_eq!(result.order.len(), 3);
691 }
692
693 #[test]
697 fn test_dp_cost_matches_reconstructed_order() {
698 let optimizer = create_test_optimizer();
699 let tables = vec![
700 TableInfo::new("A".into(), 10000).with_bit_id(0),
701 TableInfo::new("B".into(), 800).with_bit_id(1),
702 TableInfo::new("C".into(), 5000).with_bit_id(2),
703 TableInfo::new("D".into(), 300).with_bit_id(3),
704 ];
705 let conditions = vec![
706 JoinCondition::new("A".into(), "B".to_string()).with_selectivity(0.1),
707 JoinCondition::new("B".into(), "C".to_string()).with_selectivity(0.05),
708 JoinCondition::new("C".into(), "D".to_string()).with_selectivity(0.2),
709 JoinCondition::new("A".into(), "D".to_string()).with_selectivity(0.5),
710 ];
711
712 let result = optimizer.optimize_join_order(&tables, &conditions);
713 assert_eq!(result.order.len(), 4);
714
715 let condition_map = optimizer.build_condition_map(&conditions);
718 let bit_to_table: std::collections::HashMap<u32, &TableInfo> =
719 tables.iter().map(|t| (1u32 << t.bit_id, t)).collect();
720 let id_to_bit: HashMap<&String, u32> =
721 tables.iter().map(|t| (&t.id, 1u32 << t.bit_id)).collect();
722
723 let mut current_mask = id_to_bit[&result.order[0]];
724 let mut current_rows = bit_to_table[¤t_mask].estimated_rows;
725 let mut total_cost = 0.0f64;
726 for id in &result.order[1..] {
727 let table_bit = id_to_bit[id];
728 let table = bit_to_table[&table_bit];
729 let selectivity = optimizer.estimate_selectivity_for_sets(
730 current_mask,
731 table_bit,
732 &bit_to_table,
733 &condition_map,
734 );
735 total_cost += optimizer
736 .cost_calculator
737 .calculate_hash_join_cost(current_rows, table.estimated_rows);
738 current_rows =
739 ((current_rows as f64 * table.estimated_rows as f64 * selectivity) as u64).max(1);
740 current_mask |= table_bit;
741 }
742 assert!((total_cost - result.total_cost).abs() < 1e-6 * total_cost.max(1.0));
743 }
744
745 #[test]
746 fn test_table_with_selectivity() {
747 let table = TableInfo::new("A".into(), 1000)
748 .with_selectivity(0.5)
749 .with_index(true);
750
751 assert_eq!(table.selectivity, 0.5);
752 assert!(table.has_index);
753 }
754
755 #[test]
756 fn test_cost_of_order_matches_dp_total_and_rejects_non_permutation() {
757 let optimizer = create_test_optimizer();
758 let tables = vec![
759 TableInfo::new("A".into(), 10000).with_bit_id(0),
760 TableInfo::new("B".into(), 800).with_bit_id(1),
761 TableInfo::new("C".into(), 5000).with_bit_id(2),
762 ];
763 let conditions = vec![
764 JoinCondition::new("A".into(), "B".to_string()).with_selectivity(0.1),
765 JoinCondition::new("B".into(), "C".to_string()).with_selectivity(0.05),
766 ];
767
768 let result = optimizer.optimize_join_order(&tables, &conditions);
769 let priced = optimizer
770 .cost_of_order(&tables, &conditions, &result.order)
771 .expect("emitted order is priceable");
772 assert!((priced - result.total_cost).abs() < 1e-6 * result.total_cost.max(1.0));
773
774 let reversed = vec!["C".to_string(), "B".to_string(), "A".to_string()];
776 let reversed_cost = optimizer
777 .cost_of_order(&tables, &conditions, &reversed)
778 .expect("permutation is priceable");
779 assert!(reversed_cost >= result.total_cost);
780
781 assert!(optimizer
783 .cost_of_order(&tables, &conditions, &["A".to_string(), "B".to_string()])
784 .is_none());
785 assert!(optimizer
786 .cost_of_order(
787 &tables,
788 &conditions,
789 &["A".to_string(), "B".to_string(), "Z".to_string()]
790 )
791 .is_none());
792 assert!(optimizer
793 .cost_of_order(
794 &tables,
795 &conditions,
796 &["A".to_string(), "B".to_string(), "A".to_string()]
797 )
798 .is_none());
799 }
800
801 #[test]
802 fn test_condition_with_selectivity() {
803 let condition = JoinCondition::new("A".into(), "B".to_string()).with_selectivity(0.25);
804
805 assert_eq!(condition.selectivity, 0.25);
806 }
807
808 #[test]
809 fn test_to_decision() {
810 let optimizer = create_test_optimizer();
811 let tables = create_test_tables();
812 let conditions = vec![JoinCondition::new("A".into(), "B".to_string())];
813
814 let result = optimizer.optimize_join_order(&tables, &conditions);
815 let decision = optimizer.to_decision(&result);
816
817 assert_eq!(decision.join_order.len(), result.order.len());
818 }
819}