Skip to main content

linkrs_query/optimizer/cost_based/join_order/
optimizer.rs

1//! Join order optimizer: dynamic programming and greedy strategies.
2
3use 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/// Connection Order Optimizer
15#[derive(Debug)]
16pub struct JoinOrderOptimizer {
17    cost_calculator: Arc<CostCalculator>,
18    /// Threshold for the size of the dynamic programming table (if this value is exceeded, the greedy algorithm is used)
19    dp_threshold: usize,
20}
21
22/// Subproblem solution (used in dynamic programming)
23#[derive(Debug, Clone)]
24struct SubproblemSolution {
25    /// The set of included tables (bitmask)
26    pub table_set: u32,
27    /// The last table that was connected
28    pub last_table: String,
29    /// Total cost
30    pub total_cost: f64,
31    /// Estimated output rows of the subplan
32    pub output_rows: u64,
33    /// Connected tree (represented as a string)
34    pub join_tree: String,
35}
36
37impl JoinOrderOptimizer {
38    /// Create a new connection order optimizer.
39    pub fn new(cost_calculator: Arc<CostCalculator>) -> Self {
40        Self {
41            cost_calculator,
42            dp_threshold: 8, // By default, dynamic programming is used for datasets with fewer than 8 tables.
43        }
44    }
45
46    /// Setting the DP threshold
47    pub fn with_dp_threshold(mut self, threshold: usize) -> Self {
48        self.dp_threshold = threshold;
49        self
50    }
51
52    /// Optimize the order of the connections.
53    ///
54    /// # Parameters
55    /// “tables”: The list of tables involved in the connection.
56    /// “conditions”: A list of connection conditions.
57    ///
58    /// # Return
59    /// Results of the optimization of the connection sequence
60    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    /// Price a given table order with the same left-deep accumulation the
83    /// DP minimizes, so the reviewer can compare the current chain against
84    /// the proposed optimum apples-to-apples.
85    ///
86    /// Returns `None` when `order` is not a permutation of `tables`
87    /// (unpriceable — the caller must keep the current shape).
88    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    /// Optimize the connection order using dynamic programming.
135    fn optimize_with_dp(
136        &self,
137        tables: &[TableInfo],
138        conditions: &[JoinCondition],
139    ) -> JoinOrderResult {
140        let n = tables.len();
141
142        // Construct a lookup table for connection conditions
143        let condition_map = self.build_condition_map(conditions);
144
145        // DP table: key = bit mask representing the set of tables; value = the optimal solution
146        let mut dp: HashMap<u32, SubproblemSolution> = HashMap::new();
147
148        // Initialization: The case of a single table
149        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        // Map bit -> table for quick lookup of table id / rows.
161        let bit_to_table: std::collections::HashMap<u32, &TableInfo> =
162            tables.iter().map(|t| (1u32 << t.bit_id, t)).collect();
163
164        // Dynamic Programming: build subsets from smallest to largest. The DP is
165        // left-deep only — each step extends an optimal subplan by exactly one
166        // table — because the physical planner materializes left-deep join
167        // chains from the reconstructed `order`. Costing bushy trees here would
168        // report a cost the executed plan can never realize.
169        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                    // Selectivity between the accumulated set and the new table:
187                    // product of all predicates crossing the cut; 1.0 (cross
188                    // product) when no predicate connects them.
189                    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        // Obtaining the optimal solution
226        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        // Reorganize the order of the connections.
243        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    /// Optimize the connection order using a greedy algorithm.
256    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        // Select the starting table (the one with the fewest rows).
272        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        // Greedyly choosing the next table…
279        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                // Retrieve the index information for the current table and the next table.
302                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                // Choosing a connection algorithm
318                let algorithm = self.select_algorithm(
319                    current_rows,
320                    output_rows,
321                    current_has_index,
322                    next_has_index,
323                    &current_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    /// Construct a lookup table for connection conditions
344    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    /// Generate a subset of the specified size.
356    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    /// Estimate join selectivity between two table sets.
383    ///
384    /// Looks for the most selective predicate that connects the two sets
385    /// (Ladybug `SubplansTable` style). When multiple predicates cross the
386    /// cut, their selectivities are multiplied (correlation-agnostic). When no
387    /// predicate connects the sets, it is a cross product (`selectivity = 1.0`
388    /// — no row reduction, maximal cost).
389    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        // Map bit -> table id for quick lookup.
397        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            // No predicate between the two sets → cross product (no reduction).
433            1.0
434        } else {
435            // Product of crossing selectivities, clamped to avoid catastrophic under-estimation.
436            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    /// Calculating the connection cost for a single right table joining an existing left set
445    ///
446    /// Used by the greedy path. The selectivity is the most selective
447    /// condition that touches `right_table` (conservative lower bound on the
448    /// join output); fallback is the shared join selectivity.
449    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    /// Choosing a connection algorithm
476    ///
477    /// Selecting the optimal connection algorithm based on a cost model:
478    /// 1. If an index is available on one side and the amount of data is moderate, prefer an index-based join.
479    /// 2. Compare the costs of hash join and nested loop join, and choose the one with the lower cost.
480    /// 3. When using hash joins, it is advisable to choose the smaller table as the building side.
481    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        // Threshold definition
491        const NESTED_LOOP_MAX_ROWS: u64 = 100; // The maximum number of rows that can be connected using nested loops
492        const INDEX_JOIN_MAX_ROWS: u64 = 10000; // The maximum number of rows applicable for index joins
493
494        // Strategy 1: If an index is available on one side and the amount of data on the other side is moderate, use the index to perform the join operation.
495        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        // Strategy 2: If the amount of data is small in each case, use nested loops for joining the data (to avoid the overhead associated with building hash tables).
507        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        // Strategy 3: Use hash joins by default, and select smaller tables for the construction side.
515        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    /// Reorganize the order of the connections.
529    ///
530    /// Walks the left-deep DP chain backwards via `last_table` and emits the
531    /// tables in join order. The result is the exact plan shape whose cost the
532    /// DP minimized, so `JoinOrderResult::total_cost` always describes the
533    /// returned order.
534    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(&current_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        // Defensive completion: if the walk missed any table (should not
556        // happen for a well-formed DP table), append it to keep the order
557        // a permutation of all input tables.
558        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    /// Select an algorithm for the determination of the connection order.
569    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            // Obtain index information for the table
584            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            // Select using a cost-based algorithm
593            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    /// Generate a `JoinOrderDecision`
609    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        // Use DP (number of tables <= 8).
685        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    /// The DP minimizes the cost of a left-deep chain and the reconstruction
694    /// emits exactly that chain, so re-costing the returned order with the
695    /// same per-step model must reproduce the reported total cost.
696    #[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        // Re-cost the emitted order as a left-deep chain, mirroring the DP's
716        // per-step model (set-based crossing selectivity + hash join cost).
717        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[&current_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        // Reversed input order prices independently of the optimum.
775        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        // Unknown ids and duplicates are unpriceable, never silently priced.
782        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}