linkrs_query/optimizer/cost_based/
topn_optimization.rs

1//! Sorting Elimination Optimizer Module
2//!
3//! A cost-based sorting optimization strategy, focusing on the decision-making process for converting the Sort + Limit operation into a TopN result.
4//!
5//! ## Note
6//!
7//! This module **only contains cost-based optimization logic**.
8//!
9//! Heuristic sorting eliminations (such as those related to the ordering of indexes) have been handled during the rewrite phase using the `EliminateSortRule` mechanism.
10//! This layered design ensures that:
11//! Simple optimization measures (such as index matching) should be carried out as early as possible during the rewrite phase.
12//! 2. Complex cost decision-making processes (TopN conversion) are carried out during the physical optimization phase.
13//!
14//! ## Optimization Strategies
15//!
16//! - Convert “Sort + Limit” to “TopN” (based on cost comparison).
17//!
18//! ## Usage Examples
19//!
20//! ```rust
21//! use linkrs_query::optimizer::cost::CostCalculator;
22//! use linkrs_query::optimizer::cost_based::SortEliminationOptimizer;
23//! use linkrs_query::optimizer::stats::StatisticsManager;
24//! use std::sync::Arc;
25//!
26//! let stats_manager = Arc::new(StatisticsManager::new());
27//! let cost_calculator = Arc::new(CostCalculator::new(stats_manager));
28//! let optimizer = SortEliminationOptimizer::new(cost_calculator);
29//!
30//! // let decision = optimizer.optimize(&sort_context);
31//! ```
32
33use std::sync::Arc;
34
35use crate::optimizer::cost::CostCalculator;
36use crate::planning::plan::core::nodes::{SortItem, SortNode};
37
38/// Sorting-based decision elimination
39#[derive(Debug, Clone, PartialEq)]
40pub enum SortEliminationDecision {
41    /// Maintain the order (this cannot be changed or reversed).
42    KeepSort {
43        /// Reason for retention
44        reason: SortKeepReason,
45        /// Estimated cost
46        estimated_cost: f64,
47    },
48    /// Convert to TopN
49    ConvertToTopN {
50        /// Reason for conversion
51        reason: TopNConversionReason,
52        /// Cost estimate for TopN method
53        topn_cost: f64,
54        /// Original sorting cost
55        original_cost: f64,
56    },
57}
58
59/// The reason for maintaining the order:
60#[derive(Debug, Clone, PartialEq)]
61pub enum SortKeepReason {
62    /// There are no Limit child nodes; therefore, it is not possible to convert the data into a TopN format.
63    NoLimitForTopN,
64    /// The value of “Limit” is too small to be worth converting.
65    LimitTooSmall,
66    /// Retaining the sorted order based on cost analysis is more advantageous.
67    CostBasedDecision,
68}
69
70/// Reasons for converting to TopN
71#[derive(Debug, Clone, PartialEq)]
72pub enum TopNConversionReason {
73    /// The combination of “Sort” and “Limit” results in lower costs when implementing the TopN algorithm.
74    SortWithLimit,
75    /// For small amounts of data, it is better to use the TopN method.
76    SmallLimit,
77    /// Based on cost analysis
78    CostBased,
79    /// Memory constrained - sorting would require too much memory
80    MemoryConstrained,
81}
82
83/// Sorting optimization context
84#[derive(Debug, Clone)]
85pub struct SortContext {
86    /// Sorted nodes
87    pub sort_node: SortNode,
88    /// Estimated number of lines to be translated
89    pub input_rows: u64,
90    /// Is there a Limit child node?
91    pub has_limit_child: bool,
92    /// The Limit value (if any)
93    pub limit_value: Option<i64>,
94}
95
96impl SortContext {
97    /// Create a new sorting context.
98    pub fn new(sort_node: SortNode, input_rows: u64) -> Self {
99        Self {
100            sort_node,
101            input_rows,
102            has_limit_child: false,
103            limit_value: None,
104        }
105    }
106
107    /// Setting the Limit information
108    pub fn with_limit(mut self, limit: i64) -> Self {
109        self.has_limit_child = true;
110        self.limit_value = Some(limit);
111        self
112    }
113}
114
115/// Sorting Elimination Optimizer
116///
117/// A sorting optimizer based on a cost model, which focuses on the decision-making process for converting the “Sort + Limit” operation into the retrieval of the TopN results.
118///
119/// Heuristic optimizations such as the elimination of index orderliness have been completed during the rewrite phase.
120#[derive(Debug)]
121pub struct SortEliminationOptimizer {
122    cost_calculator: Arc<CostCalculator>,
123    /// TopN conversion threshold (conversion is considered when limit < threshold * input_rows)
124    topn_threshold: f64,
125    /// Only the smallest Limit value is considered for the TopN selection.
126    min_limit_for_topn: i64,
127}
128
129impl SortEliminationOptimizer {
130    /// Create a new sorting and elimination optimizer
131    pub fn new(cost_calculator: Arc<CostCalculator>) -> Self {
132        let thresholds = cost_calculator.config().strategy_thresholds;
133        Self {
134            cost_calculator,
135            topn_threshold: thresholds.topn_threshold,
136            min_limit_for_topn: thresholds.topn_default_limit as i64,
137        }
138    }
139
140    /// Set the TopN conversion threshold
141    pub fn with_topn_threshold(mut self, threshold: f64) -> Self {
142        self.topn_threshold = threshold.clamp(0.001, 1.0);
143        self
144    }
145
146    /// Set the minimum limit value
147    pub fn with_min_limit(mut self, min_limit: i64) -> Self {
148        self.min_limit_for_topn = min_limit.max(1);
149        self
150    }
151
152    /// Optimize the sorting operation
153    ///
154    /// The decision to convert from Sort + Limit to TopN is based on a cost analysis.
155    ///
156    /// # Parameters
157    /// **context:** Sort context
158    ///
159    /// # Returns
160    /// Sorting optimization decision (whether to maintain the original sorting order or convert the data into a TopN list)
161    pub fn optimize(&self, context: &SortContext) -> SortEliminationDecision {
162        self.optimize_with_memory(context, None)
163    }
164
165    /// Optimize sorting with memory awareness
166    ///
167    /// This method considers available memory when making the sort vs TopN decision.
168    /// If the sort operation would require more memory than available, it forces TopN conversion.
169    ///
170    /// # Parameters
171    /// - `context`: Sort context
172    /// - `available_memory`: Available memory in bytes (None means unlimited)
173    ///
174    /// # Returns
175    /// Sorting optimization decision
176    pub fn optimize_with_memory(
177        &self,
178        context: &SortContext,
179        available_memory: Option<usize>,
180    ) -> SortEliminationDecision {
181        let sort_items = context.sort_node.sort_items();
182
183        // Check memory constraints first
184        if let Some(memory) = available_memory {
185            let sort_memory = self
186                .cost_calculator
187                .estimate_sort_memory(context.input_rows, sort_items.len());
188
189            // If sorting would exceed available memory, force TopN
190            if sort_memory > memory {
191                if let Some(limit) = context.limit_value {
192                    if limit >= self.min_limit_for_topn {
193                        let original_cost =
194                            self.calculate_sort_cost(context.input_rows, sort_items.len());
195                        let topn_cost = self
196                            .cost_calculator
197                            .calculate_topn_cost(context.input_rows, limit);
198
199                        return SortEliminationDecision::ConvertToTopN {
200                            reason: TopNConversionReason::MemoryConstrained,
201                            topn_cost,
202                            original_cost,
203                        };
204                    }
205                }
206
207                // No limit available but memory constrained - keep sort but warn
208                let sort_cost = self.calculate_sort_cost(context.input_rows, sort_items.len());
209                return SortEliminationDecision::KeepSort {
210                    reason: SortKeepReason::CostBasedDecision,
211                    estimated_cost: sort_cost,
212                };
213            }
214        }
215
216        // Check whether it is possible to convert this into a TopN format.
217        if let Some(decision) = self.check_topn_conversion(context, sort_items) {
218            return decision;
219        }
220
221        // The text cannot be translated; the original order is therefore retained.
222        let sort_cost = self.calculate_sort_cost(context.input_rows, sort_items.len());
223        SortEliminationDecision::KeepSort {
224            reason: SortKeepReason::NoLimitForTopN,
225            estimated_cost: sort_cost,
226        }
227    }
228
229    /// Check whether it is possible to convert this into a TopN format.
230    ///
231    /// The decision to convert from “Sort + Limit” to “TopN” is based on a cost comparison.
232    /// This is a **cost-based decision**, not an unconditional conversion.
233    fn check_topn_conversion(
234        &self,
235        context: &SortContext,
236        sort_items: &[SortItem],
237    ) -> Option<SortEliminationDecision> {
238        let limit = context.limit_value?;
239
240        if limit < self.min_limit_for_topn {
241            return None;
242        }
243
244        // Single TopN rule: limit/input < topn_threshold (== 1/TOPN_RATIO)
245        // or input exceeds the shared sort-memory row threshold.
246        let limit_ratio = limit as f64 / context.input_rows as f64;
247
248        if limit_ratio < self.topn_threshold
249            || context.input_rows > self.cost_calculator.config().memory_sort_threshold
250        {
251            let original_cost = self.calculate_sort_cost(context.input_rows, sort_items.len());
252            let topn_cost = self
253                .cost_calculator
254                .calculate_topn_cost(context.input_rows, limit);
255
256            if topn_cost < original_cost {
257                return Some(SortEliminationDecision::ConvertToTopN {
258                    reason: if context.has_limit_child {
259                        TopNConversionReason::SortWithLimit
260                    } else {
261                        TopNConversionReason::CostBased
262                    },
263                    topn_cost,
264                    original_cost,
265                });
266            }
267        }
268
269        None
270    }
271
272    /// Calculating the cost of sorting
273    fn calculate_sort_cost(&self, input_rows: u64, sort_columns: usize) -> f64 {
274        self.cost_calculator
275            .calculate_sort_cost(input_rows, sort_columns, None)
276    }
277
278    /// Check whether it is possible to convert this into TopN nodes.
279    ///
280    /// # Parameters
281    /// - `sort_items`: Functions to sort items
282    /// - `limit`: The value of the `limit` parameter
283    /// - `input_rows`: The number of input rows
284    ///
285    /// # Back
286    /// - If conversion is possible, return (TopN cost, original sorting cost).
287    pub fn check_topn_conversion_cost(
288        &self,
289        sort_items: &[SortItem],
290        limit: i64,
291        input_rows: u64,
292    ) -> Option<(f64, f64)> {
293        if limit < self.min_limit_for_topn {
294            return None;
295        }
296
297        let limit_ratio = limit as f64 / input_rows as f64;
298
299        if limit_ratio < self.topn_threshold
300            || input_rows > self.cost_calculator.config().memory_sort_threshold
301        {
302            let original_cost = self.calculate_sort_cost(input_rows, sort_items.len());
303            let topn_cost = self.cost_calculator.calculate_topn_cost(input_rows, limit);
304
305            if topn_cost < original_cost {
306                return Some((topn_cost, original_cost));
307            }
308        }
309
310        None
311    }
312
313    /// Obtain suggestions for sorting optimization
314    ///
315    /// Analyze the sorting operation and provide optimization suggestions.
316    pub fn get_optimization_advice(&self, context: &SortContext) -> Vec<String> {
317        let mut advice = Vec::new();
318
319        match self.optimize(context) {
320            SortEliminationDecision::ConvertToTopN {
321                reason,
322                topn_cost,
323                original_cost,
324            } => {
325                let savings = original_cost - topn_cost;
326                advice.push(format!(
327                    "It is recommended to convert Sort + Limit to TopN, reason: {:?} and expected cost savings: {:.2}",
328                    reason, savings
329                ));
330            }
331            SortEliminationDecision::KeepSort { reason, .. } => {
332                advice.push(format!("Preserve the sort operation, cause: {:?}", reason));
333
334                if matches!(reason, SortKeepReason::NoLimitForTopN) {
335                    advice.push(
336                        "If the query contains a LIMIT, consider converting Sort + Limit to TopN"
337                            .to_string(),
338                    );
339                }
340            }
341        }
342
343        advice
344    }
345}
346
347#[cfg(test)]
348mod tests {
349    use super::*;
350    use crate::optimizer::stats::StatisticsManager;
351
352    fn create_test_optimizer() -> SortEliminationOptimizer {
353        let stats_manager = Arc::new(StatisticsManager::new());
354        let cost_calculator = Arc::new(CostCalculator::new(stats_manager));
355        SortEliminationOptimizer::new(cost_calculator)
356    }
357
358    #[test]
359    fn test_sort_elimination_optimizer_creation() {
360        let optimizer = create_test_optimizer();
361        // Values come from StrategyThresholds defaults
362        assert_eq!(optimizer.topn_threshold, 0.1);
363        assert_eq!(optimizer.min_limit_for_topn, 100); // Default from StrategyThresholds
364    }
365
366    #[test]
367    fn test_with_topn_threshold() {
368        let stats_manager = Arc::new(StatisticsManager::new());
369        let cost_calculator = Arc::new(CostCalculator::new(stats_manager));
370        let optimizer = SortEliminationOptimizer::new(cost_calculator).with_topn_threshold(0.2);
371
372        assert_eq!(optimizer.topn_threshold, 0.2);
373    }
374
375    #[test]
376    fn test_with_min_limit() {
377        let stats_manager = Arc::new(StatisticsManager::new());
378        let cost_calculator = Arc::new(CostCalculator::new(stats_manager));
379        let optimizer = SortEliminationOptimizer::new(cost_calculator).with_min_limit(10);
380
381        assert_eq!(optimizer.min_limit_for_topn, 10);
382    }
383
384    #[test]
385    fn test_threshold_clamping() {
386        let stats_manager = Arc::new(StatisticsManager::new());
387        let cost_calculator = Arc::new(CostCalculator::new(stats_manager));
388
389        let optimizer1 =
390            SortEliminationOptimizer::new(cost_calculator.clone()).with_topn_threshold(2.0); // More than 1.0
391        assert_eq!(optimizer1.topn_threshold, 1.0);
392
393        let optimizer2 = SortEliminationOptimizer::new(cost_calculator).with_topn_threshold(0.0001); // Less than 0.001
394        assert_eq!(optimizer2.topn_threshold, 0.001);
395    }
396
397    #[test]
398    fn test_context_creation() {
399        let start_node = crate::planning::plan::core::nodes::StartNode::new();
400        let sort_node = SortNode::new(
401            crate::planning::plan::PlanNodeEnum::Start(start_node),
402            vec![SortItem::column_asc("name".to_string())],
403        )
404        .expect("Failed to create SortNode");
405
406        let context = SortContext::new(sort_node, 1000);
407        assert_eq!(context.input_rows, 1000);
408        assert!(!context.has_limit_child);
409        assert_eq!(context.limit_value, None);
410    }
411
412    #[test]
413    fn test_context_with_limit() {
414        let start_node = crate::planning::plan::core::nodes::StartNode::new();
415        let sort_node = SortNode::new(
416            crate::planning::plan::PlanNodeEnum::Start(start_node),
417            vec![SortItem::column_asc("name".to_string())],
418        )
419        .expect("Failed to create SortNode");
420
421        let context = SortContext::new(sort_node, 1000).with_limit(10);
422
423        assert!(context.has_limit_child);
424        assert_eq!(context.limit_value, Some(10));
425    }
426
427    #[test]
428    fn test_check_topn_conversion_cost_limit_too_small() {
429        let optimizer = create_test_optimizer();
430        let sort_items = vec![SortItem::column_asc("name".to_string())];
431
432        let result = optimizer.check_topn_conversion_cost(&sort_items, 0, 1000);
433        assert_eq!(result, None);
434    }
435
436    #[test]
437    fn test_sort_keep_reason_display() {
438        let reason = SortKeepReason::NoLimitForTopN;
439        assert!(format!("{:?}", reason).contains("NoLimitForTopN"));
440
441        let reason = SortKeepReason::LimitTooSmall;
442        assert!(format!("{:?}", reason).contains("LimitTooSmall"));
443
444        let reason = SortKeepReason::CostBasedDecision;
445        assert!(format!("{:?}", reason).contains("CostBasedDecision"));
446    }
447
448    #[test]
449    fn test_topn_conversion_reason_display() {
450        let reason = TopNConversionReason::SortWithLimit;
451        assert!(format!("{:?}", reason).contains("SortWithLimit"));
452
453        let reason = TopNConversionReason::SmallLimit;
454        assert!(format!("{:?}", reason).contains("SmallLimit"));
455
456        let reason = TopNConversionReason::CostBased;
457        assert!(format!("{:?}", reason).contains("CostBased"));
458
459        let reason = TopNConversionReason::MemoryConstrained;
460        assert!(format!("{:?}", reason).contains("MemoryConstrained"));
461    }
462
463    #[test]
464    fn test_optimize_with_memory() {
465        let optimizer = create_test_optimizer();
466
467        // Create a sort context with limit
468        let start_node = crate::planning::plan::core::nodes::StartNode::new();
469        let sort_node = SortNode::new(
470            crate::planning::plan::PlanNodeEnum::Start(start_node),
471            vec![SortItem::column_asc("name".to_string())],
472        )
473        .expect("Failed to create SortNode");
474
475        let context = SortContext::new(sort_node, 10000).with_limit(100);
476
477        // With plenty of memory, should use normal logic
478        let decision_with_memory = optimizer.optimize_with_memory(&context, Some(10 * 1024 * 1024)); // 10MB
479        assert!(matches!(
480            decision_with_memory,
481            SortEliminationDecision::ConvertToTopN { .. }
482                | SortEliminationDecision::KeepSort { .. }
483        ));
484
485        // With very limited memory, should force TopN if possible
486        let decision_limited = optimizer.optimize_with_memory(&context, Some(1024)); // 1KB
487                                                                                     // Should convert to TopN due to memory constraint
488        match decision_limited {
489            SortEliminationDecision::ConvertToTopN { reason, .. } => {
490                assert_eq!(reason, TopNConversionReason::MemoryConstrained);
491            }
492            SortEliminationDecision::KeepSort { .. } => {
493                // This is also valid if the optimizer decides to keep sort
494            }
495        }
496    }
497
498    #[test]
499    fn test_optimize_with_memory_no_limit() {
500        let optimizer = create_test_optimizer();
501
502        // Create a sort context without limit
503        let start_node = crate::planning::plan::core::nodes::StartNode::new();
504        let sort_node = SortNode::new(
505            crate::planning::plan::PlanNodeEnum::Start(start_node),
506            vec![SortItem::column_asc("name".to_string())],
507        )
508        .expect("Failed to create SortNode");
509
510        let context = SortContext::new(sort_node, 1000); // No limit
511
512        // Without limit, memory constraint should result in KeepSort
513        let decision = optimizer.optimize_with_memory(&context, Some(1024));
514        match decision {
515            SortEliminationDecision::KeepSort { .. } => {
516                // Expected
517            }
518            _ => {
519                // Other outcomes are also acceptable
520            }
521        }
522    }
523}
linkrs_query/optimizer/cost_based/topn_optimization.rs

linkrs_query/optimizer/cost_based/
topn_optimization.rs
