linkrs_query/optimizer/cost/node_estimators/
join.rs

1//! Connection Operation Estimator
2//!
3//! Provide a cost estimate for connecting the nodes:
4//! - InnerJoin
5//! - LeftJoin
6//! - InnerJoin
7//! - LeftJoin
8//! - CrossJoin
9//! - FullOuterJoin
10
11use super::{get_input_rows, NodeEstimator};
12use crate::optimizer::cost::estimate::NodeCostEstimate;
13use crate::optimizer::cost::CostCalculator;
14use crate::optimizer::cost_based::ndv::{
15    join_output_rows, join_selectivity, DEFAULT_JOIN_SELECTIVITY,
16};
17use crate::optimizer::error::CostError;
18use crate::optimizer::stats::StatsView;
19use crate::planning::plan::PlanNodeEnum;
20
21/// Connection Operation Estimator
22pub struct JoinEstimator<'a> {
23    cost_calculator: &'a CostCalculator,
24}
25
26impl<'a> JoinEstimator<'a> {
27    /// Create a new connection estimator.
28    pub fn new(cost_calculator: &'a CostCalculator) -> Self {
29        Self { cost_calculator }
30    }
31}
32
33impl<'a> NodeEstimator for JoinEstimator<'a> {
34    fn estimate(
35        &self,
36        stats: &StatsView,
37        node: &PlanNodeEnum,
38        child_estimates: &[NodeCostEstimate],
39    ) -> Result<(f64, u64), CostError> {
40        let left_rows = get_input_rows(child_estimates, 0);
41        let right_rows = get_input_rows(child_estimates, 1);
42
43        match node {
44            PlanNodeEnum::InnerJoin(n) => {
45                // Containment estimate when NDV is known, else fixed fallback.
46                let selectivity = join_selectivity(stats, n.hash_keys(), n.probe_keys())
47                    .unwrap_or(DEFAULT_JOIN_SELECTIVITY);
48                let output_rows = join_output_rows(left_rows, right_rows, selectivity).max(1);
49                let cost = self
50                    .cost_calculator
51                    .calculate_hash_join_cost(left_rows, right_rows);
52                Ok((cost, output_rows))
53            }
54            PlanNodeEnum::LeftJoin(_) => {
55                // The left join retains all rows from the left table.
56                let output_rows = left_rows;
57                let cost = self
58                    .cost_calculator
59                    .calculate_hash_left_join_cost(left_rows, right_rows);
60                Ok((cost, output_rows))
61            }
62            PlanNodeEnum::CrossJoin(_) => {
63                let output_rows = left_rows.saturating_mul(right_rows);
64                let cost = self
65                    .cost_calculator
66                    .calculate_cross_join_cost(left_rows, right_rows);
67                Ok((cost, output_rows.max(1)))
68            }
69            PlanNodeEnum::FullOuterJoin(_) => {
70                let output_rows = left_rows.saturating_add(right_rows);
71                let cost = self
72                    .cost_calculator
73                    .calculate_full_outer_join_cost(left_rows, right_rows);
74                Ok((cost, output_rows.max(1)))
75            }
76            _ => Err(CostError::UnsupportedNodeType(format!(
77                "The connection estimator does not support node type: {:?}",
78                std::mem::discriminant(node)
79            ))),
80        }
81    }
82}
83
84#[cfg(test)]
85mod tests {
86    use super::*;
87    use crate::optimizer::cost::config::CostModelConfig;
88    use crate::planning::plan::core::nodes::control_flow::start_node::StartNode;
89    use crate::planning::plan::core::nodes::join::join_node::*;
90    use std::sync::Arc;
91
92    fn create_test_calculator() -> CostCalculator {
93        let stats_manager = Arc::new(crate::optimizer::stats::StatisticsManager::new());
94        let config = CostModelConfig::default();
95        CostCalculator::with_config(stats_manager, config)
96    }
97
98    fn create_test_start_node() -> PlanNodeEnum {
99        PlanNodeEnum::Start(StartNode::new())
100    }
101
102    #[test]
103    fn test_hash_inner_join_estimation() {
104        let calculator = create_test_calculator();
105        let estimator = JoinEstimator::new(&calculator);
106
107        let left = create_test_start_node();
108        let right = create_test_start_node();
109        let node =
110            InnerJoinNode::new(left, right, vec![], vec![]).expect("Node creation should succeed");
111        let plan_node = PlanNodeEnum::InnerJoin(node);
112
113        let child_estimates = vec![
114            NodeCostEstimate::new(10.0, 10.0, 100),
115            NodeCostEstimate::new(20.0, 20.0, 200),
116        ];
117        let result = estimator.estimate(
118            &calculator.stats_view(Some("test")),
119            &plan_node,
120            &child_estimates,
121        );
122
123        assert!(result.is_ok());
124        let (cost, output_rows) = result.expect("Estimation should succeed");
125        assert!(cost > 0.0);
126        assert!(output_rows >= 1);
127        assert_eq!(output_rows, 6000);
128    }
129
130    #[test]
131    fn test_hash_left_join_estimation() {
132        let calculator = create_test_calculator();
133        let estimator = JoinEstimator::new(&calculator);
134
135        let left = create_test_start_node();
136        let right = create_test_start_node();
137        let node =
138            LeftJoinNode::new(left, right, vec![], vec![]).expect("Node creation should succeed");
139        let plan_node = PlanNodeEnum::LeftJoin(node);
140
141        let child_estimates = vec![
142            NodeCostEstimate::new(10.0, 10.0, 100),
143            NodeCostEstimate::new(20.0, 20.0, 200),
144        ];
145        let result = estimator.estimate(
146            &calculator.stats_view(Some("test")),
147            &plan_node,
148            &child_estimates,
149        );
150
151        assert!(result.is_ok());
152        let (cost, output_rows) = result.expect("Estimation should succeed");
153        assert!(cost > 0.0);
154        assert_eq!(output_rows, 100);
155    }
156
157    #[test]
158    fn test_inner_join_estimation() {
159        let calculator = create_test_calculator();
160        let estimator = JoinEstimator::new(&calculator);
161
162        let left = create_test_start_node();
163        let right = create_test_start_node();
164        let node =
165            InnerJoinNode::new(left, right, vec![], vec![]).expect("Node creation should succeed");
166        let plan_node = PlanNodeEnum::InnerJoin(node);
167
168        let child_estimates = vec![
169            NodeCostEstimate::new(10.0, 10.0, 100),
170            NodeCostEstimate::new(20.0, 20.0, 200),
171        ];
172        let result = estimator.estimate(
173            &calculator.stats_view(Some("test")),
174            &plan_node,
175            &child_estimates,
176        );
177
178        assert!(result.is_ok());
179        let (cost, output_rows) = result.expect("Estimation should succeed");
180        assert!(cost > 0.0);
181        assert!(output_rows >= 1);
182        assert_eq!(output_rows, 6000);
183    }
184
185    #[test]
186    fn test_left_join_estimation() {
187        let calculator = create_test_calculator();
188        let estimator = JoinEstimator::new(&calculator);
189
190        let left = create_test_start_node();
191        let right = create_test_start_node();
192        let node =
193            LeftJoinNode::new(left, right, vec![], vec![]).expect("Node creation should succeed");
194        let plan_node = PlanNodeEnum::LeftJoin(node);
195
196        let child_estimates = vec![
197            NodeCostEstimate::new(10.0, 10.0, 100),
198            NodeCostEstimate::new(20.0, 20.0, 200),
199        ];
200        let result = estimator.estimate(
201            &calculator.stats_view(Some("test")),
202            &plan_node,
203            &child_estimates,
204        );
205
206        assert!(result.is_ok());
207        let (cost, output_rows) = result.expect("Estimation should succeed");
208        assert!(cost > 0.0);
209        assert_eq!(output_rows, 100);
210    }
211
212    #[test]
213    fn test_cross_join_estimation() {
214        let calculator = create_test_calculator();
215        let estimator = JoinEstimator::new(&calculator);
216
217        let left = create_test_start_node();
218        let right = create_test_start_node();
219        let node = CrossJoinNode::new(left, right).expect("Node creation should succeed");
220        let plan_node = PlanNodeEnum::CrossJoin(node);
221
222        let child_estimates = vec![
223            NodeCostEstimate::new(10.0, 10.0, 100),
224            NodeCostEstimate::new(20.0, 20.0, 200),
225        ];
226        let result = estimator.estimate(
227            &calculator.stats_view(Some("test")),
228            &plan_node,
229            &child_estimates,
230        );
231
232        assert!(result.is_ok());
233        let (cost, output_rows) = result.expect("Estimation should succeed");
234        assert!(cost > 0.0);
235        assert_eq!(output_rows, 20000);
236    }
237
238    #[test]
239    fn test_full_outer_join_estimation() {
240        let calculator = create_test_calculator();
241        let estimator = JoinEstimator::new(&calculator);
242
243        let left = create_test_start_node();
244        let right = create_test_start_node();
245        let node = FullOuterJoinNode::new(left, right, vec![], vec![])
246            .expect("Node creation should succeed");
247        let plan_node = PlanNodeEnum::FullOuterJoin(node);
248
249        let child_estimates = vec![
250            NodeCostEstimate::new(10.0, 10.0, 100),
251            NodeCostEstimate::new(20.0, 20.0, 200),
252        ];
253        let result = estimator.estimate(
254            &calculator.stats_view(Some("test")),
255            &plan_node,
256            &child_estimates,
257        );
258
259        assert!(result.is_ok());
260        let (cost, output_rows) = result.expect("Estimation should succeed");
261        assert!(cost > 0.0);
262        assert_eq!(output_rows, 300);
263    }
264
265    #[test]
266    fn test_unsupported_node_type() {
267        let calculator = create_test_calculator();
268        let estimator = JoinEstimator::new(&calculator);
269
270        let node = PlanNodeEnum::Start(StartNode::new());
271        let child_estimates = vec![];
272        let result = estimator.estimate(
273            &calculator.stats_view(Some("test")),
274            &node,
275            &child_estimates,
276        );
277
278        assert!(result.is_err());
279    }
280
281    #[test]
282    fn test_join_with_different_input_sizes() {
283        let calculator = create_test_calculator();
284        let estimator = JoinEstimator::new(&calculator);
285
286        for (left_rows, right_rows) in [(10, 20), (100, 200), (1000, 500)] {
287            let left = create_test_start_node();
288            let right = create_test_start_node();
289            let node = InnerJoinNode::new(left, right, vec![], vec![])
290                .expect("Node creation should succeed");
291            let plan_node = PlanNodeEnum::InnerJoin(node);
292
293            let child_estimates = vec![
294                NodeCostEstimate::new(10.0, 10.0, left_rows),
295                NodeCostEstimate::new(20.0, 20.0, right_rows),
296            ];
297            let result = estimator.estimate(
298                &calculator.stats_view(Some("test")),
299                &plan_node,
300                &child_estimates,
301            );
302
303            assert!(result.is_ok());
304            let (cost, output_rows) = result.expect("Estimation should succeed");
305            assert!(cost > 0.0);
306            assert!(output_rows >= 1);
307        }
308    }
309
310    #[test]
311    fn test_cross_join_large_inputs() {
312        let calculator = create_test_calculator();
313        let estimator = JoinEstimator::new(&calculator);
314
315        let left = create_test_start_node();
316        let right = create_test_start_node();
317        let node = CrossJoinNode::new(left, right).expect("Node creation should succeed");
318        let plan_node = PlanNodeEnum::CrossJoin(node);
319
320        let child_estimates = vec![
321            NodeCostEstimate::new(10.0, 10.0, 1000),
322            NodeCostEstimate::new(20.0, 20.0, 1000),
323        ];
324        let result = estimator.estimate(
325            &calculator.stats_view(Some("test")),
326            &plan_node,
327            &child_estimates,
328        );
329
330        assert!(result.is_ok());
331        let (cost, output_rows) = result.expect("Estimation should succeed");
332        assert!(cost > 0.0);
333        assert_eq!(output_rows, 1_000_000);
334    }
335
336    #[test]
337    fn test_join_with_zero_input_rows() {
338        let calculator = create_test_calculator();
339        let estimator = JoinEstimator::new(&calculator);
340
341        let left = create_test_start_node();
342        let right = create_test_start_node();
343        let node =
344            InnerJoinNode::new(left, right, vec![], vec![]).expect("Node creation should succeed");
345        let plan_node = PlanNodeEnum::InnerJoin(node);
346
347        let child_estimates = vec![
348            NodeCostEstimate::new(0.0, 0.0, 0),
349            NodeCostEstimate::new(0.0, 0.0, 0),
350        ];
351        let result = estimator.estimate(
352            &calculator.stats_view(Some("test")),
353            &plan_node,
354            &child_estimates,
355        );
356
357        assert!(result.is_ok());
358        let (cost, output_rows) = result.expect("Estimation should succeed");
359        assert!(cost >= 0.0);
360        assert_eq!(output_rows, 1);
361    }
362
363    #[test]
364    fn test_left_join_preserves_left_rows() {
365        let calculator = create_test_calculator();
366        let estimator = JoinEstimator::new(&calculator);
367
368        let left = create_test_start_node();
369        let right = create_test_start_node();
370        let node =
371            LeftJoinNode::new(left, right, vec![], vec![]).expect("Node creation should succeed");
372        let plan_node = PlanNodeEnum::LeftJoin(node);
373
374        let child_estimates = vec![
375            NodeCostEstimate::new(10.0, 10.0, 500),
376            NodeCostEstimate::new(20.0, 20.0, 1000),
377        ];
378        let result = estimator.estimate(
379            &calculator.stats_view(Some("test")),
380            &plan_node,
381            &child_estimates,
382        );
383
384        assert!(result.is_ok());
385        let (_, output_rows) = result.expect("Estimation should succeed");
386        assert_eq!(output_rows, 500);
387    }
388
389    #[test]
390    fn test_hash_left_join_preserves_left_rows() {
391        let calculator = create_test_calculator();
392        let estimator = JoinEstimator::new(&calculator);
393
394        let left = create_test_start_node();
395        let right = create_test_start_node();
396        let node =
397            LeftJoinNode::new(left, right, vec![], vec![]).expect("Node creation should succeed");
398        let plan_node = PlanNodeEnum::LeftJoin(node);
399
400        let child_estimates = vec![
401            NodeCostEstimate::new(10.0, 10.0, 750),
402            NodeCostEstimate::new(20.0, 20.0, 1500),
403        ];
404        let result = estimator.estimate(
405            &calculator.stats_view(Some("test")),
406            &plan_node,
407            &child_estimates,
408        );
409
410        assert!(result.is_ok());
411        let (_, output_rows) = result.expect("Estimation should succeed");
412        assert_eq!(output_rows, 750);
413    }
414}
linkrs_query/optimizer/cost/node_estimators/join.rs

linkrs_query/optimizer/cost/node_estimators/
join.rs
