linkrs_query/optimizer/engine/
config.rs

1use std::collections::HashMap;
2use std::sync::atomic::{AtomicU64, Ordering};
3use std::sync::Arc;
4
5use crate::optimizer::heuristic::{LogicalBatchOptimizer, PhysicalHeuristicOptimizer};
6use crate::optimizer::partitioning::{PartitioningConfig, PartitioningPlanner};
7use crate::optimizer::stats::feedback::cardinality::CardinalityFeedbackManager;
8use crate::optimizer::stats::feedback::decision::DecisionFeedbackStore;
9use crate::optimizer::stats::feedback::history::QueryFeedbackHistory;
10use crate::optimizer::stats::feedback::selectivity::SelectivityFeedbackManager;
11use crate::optimizer::stats::feedback::trigger::AutoFeedbackTrigger;
12use crate::optimizer::{
13    BatchPlanAnalyzer, CostCalculator, CostModelConfig, CteCacheManager, SelectivityEstimator,
14    StatisticsManager, SubqueryUnnestingOptimizer,
15};
16use linkrs_core::types::expr::expression_context::ExpressionAnalysisContext;
17
18use super::OptimizerEngine;
19
20impl OptimizerEngine {
21    pub fn new(cost_config: CostModelConfig) -> Self {
22        Self::with_expression_context(Arc::new(ExpressionAnalysisContext::new()), cost_config)
23    }
24
25    pub fn with_expression_context(
26        expression_context: Arc<ExpressionAnalysisContext>,
27        cost_config: CostModelConfig,
28    ) -> Self {
29        let stats_manager = Arc::new(StatisticsManager::new());
30        let cte_cache_manager = Arc::new(CteCacheManager::new());
31
32        Self::with_components(
33            expression_context,
34            stats_manager,
35            cte_cache_manager,
36            cost_config,
37            None,
38        )
39    }
40
41    pub(crate) fn with_components(
42        expression_context: Arc<ExpressionAnalysisContext>,
43        stats_manager: Arc<StatisticsManager>,
44        cte_cache_manager: Arc<CteCacheManager>,
45        cost_config: CostModelConfig,
46        metrics_stats: Option<Arc<linkrs_metrics::StatsManager>>,
47    ) -> Self {
48        let cost_calculator = Arc::new(CostCalculator::with_config(
49            stats_manager.clone(),
50            cost_config,
51        ));
52        let selectivity_feedback = Arc::new(SelectivityFeedbackManager::new());
53        let selectivity_estimator = Arc::new(SelectivityEstimator::with_feedback(
54            stats_manager.clone(),
55            selectivity_feedback.clone(),
56        ));
57        let cardinality_feedback = Arc::new(CardinalityFeedbackManager::new());
58        let decision_feedback = Arc::new(DecisionFeedbackStore::new());
59
60        let batch_plan_analyzer = BatchPlanAnalyzer::new();
61        let subquery_unnesting_optimizer = SubqueryUnnestingOptimizer::new();
62
63        let logical_heuristic = LogicalBatchOptimizer::new();
64        let physical_heuristic = PhysicalHeuristicOptimizer::from_registry(
65            crate::optimizer::heuristic::rule_enum::RuleRegistry::default(),
66        );
67
68        Self {
69            expression_context,
70            stats_manager,
71            cte_cache_manager,
72            cost_calculator,
73            selectivity_estimator: Arc::clone(&selectivity_estimator),
74            batch_plan_analyzer,
75            subquery_unnesting_optimizer,
76            cost_config,
77            space_cost_configs: parking_lot::RwLock::new(HashMap::new()),
78            space_calculators: parking_lot::RwLock::new(HashMap::new()),
79            cost_epoch: AtomicU64::new(0),
80            logical_heuristic,
81            physical_heuristic,
82            last_batch_statistics: std::sync::Mutex::new(Vec::new()),
83            partitioning_planner: PartitioningPlanner::new(PartitioningConfig::default())
84                .with_selectivity(Arc::clone(&selectivity_estimator)),
85            enable_heuristic: true,
86            max_heuristic_iterations: 100,
87            feedback_history: Arc::new(QueryFeedbackHistory::default()),
88            selectivity_feedback,
89            cardinality_feedback,
90            decision_feedback,
91            feedback_trigger: AutoFeedbackTrigger::default(),
92            enable_feedback: true,
93            feedback_applied_seq: parking_lot::Mutex::new(HashMap::new()),
94            columnar_policy: Arc::new(crate::executor::streaming::chunk::ColumnarPolicy::default()),
95            metrics_stats,
96        }
97    }
98
99    pub fn for_ssd() -> Self {
100        Self::new(CostModelConfig::for_ssd())
101    }
102
103    pub fn for_in_memory() -> Self {
104        Self::new(CostModelConfig::for_in_memory())
105    }
106
107    pub fn cost_config(&self) -> &CostModelConfig {
108        &self.cost_config
109    }
110
111    pub fn cost_calculator(&self) -> &Arc<CostCalculator> {
112        &self.cost_calculator
113    }
114
115    pub fn cost_calculator_for(&self, space: Option<&str>) -> Arc<CostCalculator> {
116        let Some(space) = space else {
117            return self.cost_calculator.clone();
118        };
119        if let Some(cached) = self.space_calculators.read().get(space) {
120            return cached.clone();
121        }
122        let config = self
123            .space_cost_configs
124            .read()
125            .get(space)
126            .copied()
127            .unwrap_or(self.cost_config);
128        let calculator = Arc::new(CostCalculator::with_config(
129            self.stats_manager.clone(),
130            config,
131        ));
132        self.space_calculators
133            .write()
134            .insert(space.to_string(), calculator.clone());
135        calculator
136    }
137
138    pub fn space_cost_configs(&self) -> HashMap<String, CostModelConfig> {
139        self.space_cost_configs.read().clone()
140    }
141
142    pub fn set_space_cost_configs(&self, configs: HashMap<String, CostModelConfig>) {
143        *self.space_cost_configs.write() = configs;
144        self.space_calculators.write().clear();
145        self.cost_epoch.fetch_add(1, Ordering::Relaxed);
146    }
147
148    pub fn cost_epoch(&self) -> u64 {
149        self.cost_epoch.load(Ordering::Relaxed)
150    }
151
152    pub fn stats_manager(&self) -> &Arc<StatisticsManager> {
153        &self.stats_manager
154    }
155
156    pub fn selectivity_estimator(&self) -> &Arc<SelectivityEstimator> {
157        &self.selectivity_estimator
158    }
159
160    pub fn expression_context(&self) -> &Arc<ExpressionAnalysisContext> {
161        &self.expression_context
162    }
163
164    pub fn feedback_history(&self) -> Arc<QueryFeedbackHistory> {
165        Arc::clone(&self.feedback_history)
166    }
167
168    pub fn selectivity_feedback(&self) -> &Arc<SelectivityFeedbackManager> {
169        &self.selectivity_feedback
170    }
171
172    pub fn cardinality_feedback(&self) -> &Arc<CardinalityFeedbackManager> {
173        &self.cardinality_feedback
174    }
175
176    pub fn decision_feedback(&self) -> &Arc<DecisionFeedbackStore> {
177        &self.decision_feedback
178    }
179
180    pub fn set_enable_feedback(&mut self, enable: bool) {
181        self.enable_feedback = enable;
182        log::info!(
183            "Feedback-driven selectivity correction has been {}",
184            if enable { "enabled" } else { "disabled" }
185        );
186    }
187
188    pub fn feedback_enabled(&self) -> bool {
189        self.enable_feedback
190    }
191
192    pub fn columnar_policy(&self) -> Arc<crate::executor::streaming::chunk::ColumnarPolicy> {
193        Arc::clone(&self.columnar_policy)
194    }
195
196    pub fn set_metrics_stats(&mut self, stats: Arc<linkrs_metrics::StatsManager>) {
197        self.metrics_stats = Some(stats);
198    }
199
200    pub fn metrics_stats(&self) -> Option<Arc<linkrs_metrics::StatsManager>> {
201        self.metrics_stats.clone()
202    }
203
204    pub fn batch_plan_analyzer(&self) -> &BatchPlanAnalyzer {
205        &self.batch_plan_analyzer
206    }
207
208    pub fn subquery_unnesting_optimizer(&self) -> &SubqueryUnnestingOptimizer {
209        &self.subquery_unnesting_optimizer
210    }
211
212    pub fn cte_cache_manager(&self) -> &CteCacheManager {
213        &self.cte_cache_manager
214    }
215
216    pub fn set_cte_cache_stats_manager(&self, stats_manager: Arc<linkrs_metrics::StatsManager>) {
217        self.cte_cache_manager.set_stats_manager(stats_manager);
218    }
219
220    pub fn set_cost_config(&mut self, config: CostModelConfig) {
221        self.cost_config = config;
222        self.cost_calculator = Arc::new(CostCalculator::with_config(
223            self.stats_manager.clone(),
224            self.cost_config,
225        ));
226        self.space_calculators.write().clear();
227        self.cost_epoch.fetch_add(1, Ordering::Relaxed);
228        self.batch_plan_analyzer = BatchPlanAnalyzer::new();
229        self.subquery_unnesting_optimizer = SubqueryUnnestingOptimizer::new();
230        log::info!(
231            "Optimizer cost model configuration has been updated: {:?}",
232            self.cost_config
233        );
234    }
235
236    pub fn set_enable_heuristic(&mut self, enable: bool) {
237        self.enable_heuristic = enable;
238        log::info!(
239            "Heuristic optimization has {}",
240            if enable {
241                "(computing) enable (a feature)"
242            } else {
243                "prohibit the use of sth."
244            }
245        );
246    }
247
248    pub fn set_max_heuristic_iterations(&mut self, max: usize) {
249        self.max_heuristic_iterations = max;
250        log::info!(
251            "The maximum number of heuristic iterations has been set to {}",
252            max
253        );
254    }
255
256    pub fn set_partitioning_config(&mut self, config: PartitioningConfig) {
257        self.partitioning_planner =
258            PartitioningPlanner::new(config).with_selectivity(self.selectivity_estimator.clone());
259    }
260
261    /// Apply the `[optimizer.rules]` selection to the physical heuristic
262    /// phase. `enabled` non-empty keeps only those rules; `disabled` then
263    /// removes rules. Names match rule type names with or without the
264    /// `Rule` suffix.
265    pub fn set_rule_selection(&mut self, disabled: &[String], enabled: &[String]) {
266        if disabled.is_empty() && enabled.is_empty() {
267            return;
268        }
269        self.physical_heuristic = PhysicalHeuristicOptimizer::from_registry(
270            crate::optimizer::heuristic::rule_enum::RuleRegistry::filtered(disabled, enabled),
271        );
272        log::info!(
273            "Optimizer rule selection applied: {} disabled, {} explicitly enabled",
274            disabled.len(),
275            enabled.len()
276        );
277    }
278
279    pub fn partitioning_config(&self) -> &PartitioningConfig {
280        self.partitioning_planner.config()
281    }
282
283    pub fn heuristic_batch(&self) -> &crate::optimizer::heuristic::batch::BatchOptimizer {
284        self.physical_heuristic.batch()
285    }
286
287    pub fn logical_heuristic(&self) -> &LogicalBatchOptimizer {
288        &self.logical_heuristic
289    }
290
291    pub fn physical_heuristic(&self) -> &PhysicalHeuristicOptimizer {
292        &self.physical_heuristic
293    }
294
295    pub fn last_batch_statistics(
296        &self,
297    ) -> Vec<(
298        crate::optimizer::heuristic::batch::OptimizationBatch,
299        crate::optimizer::heuristic::batch::BatchStatistics,
300    )> {
301        match self.last_batch_statistics.lock() {
302            Ok(guard) => guard.clone(),
303            Err(poisoned) => poisoned.into_inner().clone(),
304        }
305    }
306}
linkrs_query/optimizer/engine/config.rs

linkrs_query/optimizer/engine/
config.rs
