linkrs_query/optimizer/
builder.rs1use std::sync::Arc;
7
8use crate::optimizer::OptimizerEngine;
9use crate::optimizer::{CostModelConfig, CteCacheManager, StatisticsManager};
10use linkrs_core::types::expr::expression_context::ExpressionAnalysisContext;
11
12pub struct OptimizerEngineBuilder {
14 cost_config: Option<CostModelConfig>,
15 expression_context: Option<Arc<ExpressionAnalysisContext>>,
16 stats_manager: Option<Arc<StatisticsManager>>,
17 cte_cache_manager: Option<Arc<CteCacheManager>>,
18 metrics_stats: Option<Arc<linkrs_metrics::StatsManager>>,
19}
20
21impl Default for OptimizerEngineBuilder {
22 fn default() -> Self {
23 Self::new()
24 }
25}
26
27impl OptimizerEngineBuilder {
28 pub fn new() -> Self {
30 Self {
31 cost_config: None,
32 expression_context: None,
33 stats_manager: None,
34 cte_cache_manager: None,
35 metrics_stats: None,
36 }
37 }
38
39 pub fn with_cost_config(mut self, config: CostModelConfig) -> Self {
41 self.cost_config = Some(config);
42 self
43 }
44
45 pub fn with_expression_context(mut self, ctx: Arc<ExpressionAnalysisContext>) -> Self {
47 self.expression_context = Some(ctx);
48 self
49 }
50
51 pub fn with_stats_manager(mut self, manager: Arc<StatisticsManager>) -> Self {
53 self.stats_manager = Some(manager);
54 self
55 }
56
57 pub fn with_cte_cache_manager(mut self, manager: Arc<CteCacheManager>) -> Self {
59 self.cte_cache_manager = Some(manager);
60 self
61 }
62
63 pub fn with_metrics_stats(mut self, stats: Arc<linkrs_metrics::StatsManager>) -> Self {
65 self.metrics_stats = Some(stats);
66 self
67 }
68
69 pub fn build(self) -> OptimizerEngine {
71 let cost_config = self.cost_config.unwrap_or_default();
72 let expression_context = self
73 .expression_context
74 .unwrap_or_else(|| Arc::new(ExpressionAnalysisContext::new()));
75 let stats_manager = self
76 .stats_manager
77 .unwrap_or_else(|| Arc::new(StatisticsManager::new()));
78 let cte_cache_manager = self
79 .cte_cache_manager
80 .unwrap_or_else(|| Arc::new(CteCacheManager::new()));
81
82 OptimizerEngine::with_components(
83 expression_context,
84 stats_manager,
85 cte_cache_manager,
86 cost_config,
87 self.metrics_stats,
88 )
89 }
90}
91
92#[cfg(test)]
93mod tests {
94 use super::*;
95
96 #[test]
97 fn test_builder_default() {
98 let builder = OptimizerEngineBuilder::default();
99 let engine = builder.build();
100 assert_eq!(engine.cost_config(), &CostModelConfig::default());
101 }
102
103 #[test]
104 fn test_builder_with_cost_config() {
105 let config = CostModelConfig::for_ssd();
106 let builder = OptimizerEngineBuilder::new().with_cost_config(config);
107 let engine = builder.build();
108 assert_eq!(engine.cost_config(), &config);
109 }
110
111 #[test]
112 fn test_builder_with_expression_context() {
113 let ctx = Arc::new(ExpressionAnalysisContext::new());
114 let builder = OptimizerEngineBuilder::new().with_expression_context(ctx.clone());
115 let engine = builder.build();
116 assert!(Arc::ptr_eq(engine.expression_context(), &ctx));
117 }
118
119 #[test]
120 fn test_builder_with_stats_manager() {
121 let stats = Arc::new(StatisticsManager::new());
122 let builder = OptimizerEngineBuilder::new().with_stats_manager(stats.clone());
123 let engine = builder.build();
124 assert_eq!(
125 engine.stats_manager().as_ref() as *const _,
126 stats.as_ref() as *const _
127 );
128 }
129
130 #[test]
131 fn test_builder_chain() {
132 let config = CostModelConfig::for_ssd();
133 let ctx = Arc::new(ExpressionAnalysisContext::new());
134 let stats = Arc::new(StatisticsManager::new());
135
136 let engine = OptimizerEngineBuilder::new()
137 .with_cost_config(config)
138 .with_expression_context(ctx)
139 .with_stats_manager(stats)
140 .build();
141
142 assert_eq!(engine.cost_config(), &config);
143 }
144}