oxirs_gql/
cost_based_optimizer.rs

1//! Cost-Based Query Optimization
2//!
3//! This module provides sophisticated cost-based optimization for GraphQL queries,
4//! using statistical models and historical data to generate optimal execution plans.
5//!
6//! # Features
7//!
8//! - **Cost Estimation**: Accurate cost estimation for different execution strategies
9//! - **Statistics Tracking**: Track query statistics for better cost models
10//! - **Plan Comparison**: Compare multiple execution plans and choose the best
11//! - **Adaptive Optimization**: Learn from execution patterns over time
12//! - **Join Optimization**: Optimize joins and data fetching operations
13//! - **Index Recommendations**: Suggest indexes for better performance
14//!
15//! # Example
16//!
17//! ```rust,ignore
18//! use oxirs_gql::cost_based_optimizer::{CostBasedOptimizer, OptimizationConfig};
19//!
20//! let config = OptimizationConfig::new()
21//!     .with_statistics_collection(true)
22//!     .with_adaptive_learning(true);
23//!
24//! let optimizer = CostBasedOptimizer::new(config);
25//! let query = /* ... GraphQL query ... */;
26//!
27//! // Generate optimized plan
28//! let plan = optimizer.optimize(query).await?;
29//! ```
30
31use serde::{Deserialize, Serialize};
32use std::collections::HashMap;
33use std::sync::Arc;
34use tokio::sync::RwLock;
35
36/// Configuration for cost-based optimization
37#[derive(Debug, Clone)]
38pub struct OptimizationConfig {
39    /// Enable statistics collection
40    pub enable_statistics: bool,
41    /// Enable adaptive learning
42    pub enable_adaptive_learning: bool,
43    /// Maximum number of alternative plans to consider
44    pub max_alternative_plans: usize,
45    /// Cost threshold for plan selection
46    pub cost_threshold: f64,
47    /// Enable join optimization
48    pub enable_join_optimization: bool,
49    /// Enable index recommendations
50    pub enable_index_recommendations: bool,
51    /// Statistics sample rate (0.0-1.0)
52    pub statistics_sample_rate: f64,
53}
54
55impl Default for OptimizationConfig {
56    fn default() -> Self {
57        Self {
58            enable_statistics: true,
59            enable_adaptive_learning: true,
60            max_alternative_plans: 5,
61            cost_threshold: 1000.0,
62            enable_join_optimization: true,
63            enable_index_recommendations: true,
64            statistics_sample_rate: 1.0,
65        }
66    }
67}
68
69impl OptimizationConfig {
70    /// Create new configuration
71    pub fn new() -> Self {
72        Self::default()
73    }
74
75    /// Enable statistics collection
76    pub fn with_statistics_collection(mut self, enabled: bool) -> Self {
77        self.enable_statistics = enabled;
78        self
79    }
80
81    /// Enable adaptive learning
82    pub fn with_adaptive_learning(mut self, enabled: bool) -> Self {
83        self.enable_adaptive_learning = enabled;
84        self
85    }
86
87    /// Set maximum alternative plans
88    pub fn with_max_alternative_plans(mut self, max: usize) -> Self {
89        self.max_alternative_plans = max;
90        self
91    }
92
93    /// Set cost threshold
94    pub fn with_cost_threshold(mut self, threshold: f64) -> Self {
95        self.cost_threshold = threshold;
96        self
97    }
98
99    /// Enable join optimization
100    pub fn with_join_optimization(mut self, enabled: bool) -> Self {
101        self.enable_join_optimization = enabled;
102        self
103    }
104
105    /// Enable index recommendations
106    pub fn with_index_recommendations(mut self, enabled: bool) -> Self {
107        self.enable_index_recommendations = enabled;
108        self
109    }
110
111    /// Set statistics sample rate
112    pub fn with_statistics_sample_rate(mut self, rate: f64) -> Self {
113        self.statistics_sample_rate = rate.clamp(0.0, 1.0);
114        self
115    }
116}
117
118/// Query operation type
119#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
120pub enum OperationType {
121    /// Field selection
122    FieldSelect,
123    /// Object fetch
124    ObjectFetch,
125    /// List fetch
126    ListFetch,
127    /// Join operation
128    Join,
129    /// Filter operation
130    Filter,
131    /// Aggregation
132    Aggregation,
133    /// Sort operation
134    Sort,
135}
136
137/// Cost factors for an operation
138#[derive(Debug, Clone, Serialize, Deserialize)]
139pub struct CostFactors {
140    /// CPU cost
141    pub cpu_cost: f64,
142    /// I/O cost
143    pub io_cost: f64,
144    /// Network cost
145    pub network_cost: f64,
146    /// Memory cost
147    pub memory_cost: f64,
148}
149
150impl CostFactors {
151    /// Create new cost factors
152    pub fn new() -> Self {
153        Self {
154            cpu_cost: 0.0,
155            io_cost: 0.0,
156            network_cost: 0.0,
157            memory_cost: 0.0,
158        }
159    }
160
161    /// Total cost
162    pub fn total(&self) -> f64 {
163        self.cpu_cost + self.io_cost + self.network_cost + self.memory_cost
164    }
165
166    /// Add another cost
167    pub fn add(&mut self, other: &CostFactors) {
168        self.cpu_cost += other.cpu_cost;
169        self.io_cost += other.io_cost;
170        self.network_cost += other.network_cost;
171        self.memory_cost += other.memory_cost;
172    }
173}
174
175impl Default for CostFactors {
176    fn default() -> Self {
177        Self::new()
178    }
179}
180
181/// Operation cost estimate
182#[derive(Debug, Clone, Serialize, Deserialize)]
183pub struct OperationCost {
184    /// Operation type
185    pub operation: OperationType,
186    /// Operation name/identifier
187    pub name: String,
188    /// Cost factors
189    pub factors: CostFactors,
190    /// Estimated row count
191    pub estimated_rows: usize,
192    /// Estimated execution time
193    pub estimated_time_ms: u64,
194    /// Confidence score (0.0-1.0)
195    pub confidence: f64,
196}
197
198impl OperationCost {
199    /// Create new operation cost
200    pub fn new(operation: OperationType, name: String) -> Self {
201        Self {
202            operation,
203            name,
204            factors: CostFactors::new(),
205            estimated_rows: 0,
206            estimated_time_ms: 0,
207            confidence: 0.5,
208        }
209    }
210
211    /// Set CPU cost
212    pub fn with_cpu_cost(mut self, cost: f64) -> Self {
213        self.factors.cpu_cost = cost;
214        self
215    }
216
217    /// Set I/O cost
218    pub fn with_io_cost(mut self, cost: f64) -> Self {
219        self.factors.io_cost = cost;
220        self
221    }
222
223    /// Set network cost
224    pub fn with_network_cost(mut self, cost: f64) -> Self {
225        self.factors.network_cost = cost;
226        self
227    }
228
229    /// Set memory cost
230    pub fn with_memory_cost(mut self, cost: f64) -> Self {
231        self.factors.memory_cost = cost;
232        self
233    }
234
235    /// Set estimated rows
236    pub fn with_estimated_rows(mut self, rows: usize) -> Self {
237        self.estimated_rows = rows;
238        self
239    }
240
241    /// Set estimated time
242    pub fn with_estimated_time(mut self, time_ms: u64) -> Self {
243        self.estimated_time_ms = time_ms;
244        self
245    }
246
247    /// Set confidence
248    pub fn with_confidence(mut self, confidence: f64) -> Self {
249        self.confidence = confidence.clamp(0.0, 1.0);
250        self
251    }
252
253    /// Total cost
254    pub fn total_cost(&self) -> f64 {
255        self.factors.total()
256    }
257}
258
259/// Execution plan with costs
260#[derive(Debug, Clone, Serialize, Deserialize)]
261pub struct CostBasedPlan {
262    /// Plan identifier
263    pub plan_id: String,
264    /// Operations in the plan
265    pub operations: Vec<OperationCost>,
266    /// Total cost
267    pub total_cost: f64,
268    /// Total estimated time
269    pub total_estimated_time_ms: u64,
270    /// Plan strategy
271    pub strategy: PlanStrategy,
272    /// Index recommendations
273    pub index_recommendations: Vec<IndexRecommendation>,
274}
275
276impl CostBasedPlan {
277    /// Create a new plan
278    pub fn new(plan_id: String, strategy: PlanStrategy) -> Self {
279        Self {
280            plan_id,
281            operations: Vec::new(),
282            total_cost: 0.0,
283            total_estimated_time_ms: 0,
284            strategy,
285            index_recommendations: Vec::new(),
286        }
287    }
288
289    /// Add operation
290    pub fn add_operation(&mut self, op: OperationCost) {
291        self.total_cost += op.total_cost();
292        self.total_estimated_time_ms += op.estimated_time_ms;
293        self.operations.push(op);
294    }
295
296    /// Add index recommendation
297    pub fn add_index_recommendation(&mut self, rec: IndexRecommendation) {
298        self.index_recommendations.push(rec);
299    }
300
301    /// Compare with another plan
302    pub fn is_better_than(&self, other: &CostBasedPlan) -> bool {
303        self.total_cost < other.total_cost
304    }
305}
306
307/// Plan strategy
308#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
309pub enum PlanStrategy {
310    /// Sequential execution
311    Sequential,
312    /// Parallel execution
313    Parallel,
314    /// Batch execution
315    Batch,
316    /// Streaming execution
317    Streaming,
318    /// Hybrid strategy
319    Hybrid,
320}
321
322/// Index recommendation
323#[derive(Debug, Clone, Serialize, Deserialize)]
324pub struct IndexRecommendation {
325    /// Table/type name
326    pub table: String,
327    /// Fields to index
328    pub fields: Vec<String>,
329    /// Index type
330    pub index_type: IndexType,
331    /// Expected performance improvement (percentage)
332    pub expected_improvement: f64,
333    /// Rationale for recommendation
334    pub rationale: String,
335}
336
337/// Index type
338#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
339pub enum IndexType {
340    /// B-tree index
341    BTree,
342    /// Hash index
343    Hash,
344    /// Full-text index
345    FullText,
346    /// Composite index
347    Composite,
348}
349
350/// Query statistics
351#[derive(Debug, Clone, Serialize, Deserialize)]
352pub struct QueryStatistics {
353    /// Query fingerprint
354    pub fingerprint: String,
355    /// Number of executions
356    pub execution_count: usize,
357    /// Average execution time
358    pub avg_execution_time_ms: f64,
359    /// Average row count
360    pub avg_row_count: f64,
361    /// Minimum execution time
362    pub min_execution_time_ms: u64,
363    /// Maximum execution time
364    pub max_execution_time_ms: u64,
365    /// Last execution timestamp
366    pub last_execution: u64,
367}
368
369impl QueryStatistics {
370    /// Create new statistics
371    pub fn new(fingerprint: String) -> Self {
372        Self {
373            fingerprint,
374            execution_count: 0,
375            avg_execution_time_ms: 0.0,
376            avg_row_count: 0.0,
377            min_execution_time_ms: u64::MAX,
378            max_execution_time_ms: 0,
379            last_execution: 0,
380        }
381    }
382
383    /// Update with new execution
384    pub fn update(&mut self, execution_time_ms: u64, row_count: usize) {
385        self.execution_count += 1;
386
387        // Update average execution time
388        let old_avg = self.avg_execution_time_ms;
389        self.avg_execution_time_ms = (old_avg * (self.execution_count - 1) as f64
390            + execution_time_ms as f64)
391            / self.execution_count as f64;
392
393        // Update average row count
394        let old_row_avg = self.avg_row_count;
395        self.avg_row_count = (old_row_avg * (self.execution_count - 1) as f64 + row_count as f64)
396            / self.execution_count as f64;
397
398        // Update min/max
399        self.min_execution_time_ms = self.min_execution_time_ms.min(execution_time_ms);
400        self.max_execution_time_ms = self.max_execution_time_ms.max(execution_time_ms);
401
402        self.last_execution = chrono::Utc::now().timestamp() as u64;
403    }
404}
405
406/// Cost model for estimation
407pub struct CostModel {
408    /// Cost weights
409    weights: HashMap<OperationType, CostFactors>,
410}
411
412impl CostModel {
413    /// Create a new cost model
414    pub fn new() -> Self {
415        let mut weights = HashMap::new();
416
417        // Initialize default weights
418        let mut field_cost = CostFactors::new();
419        field_cost.cpu_cost = 1.0;
420        field_cost.io_cost = 0.5;
421        weights.insert(OperationType::FieldSelect, field_cost);
422
423        let mut object_cost = CostFactors::new();
424        object_cost.cpu_cost = 2.0;
425        object_cost.io_cost = 5.0;
426        weights.insert(OperationType::ObjectFetch, object_cost);
427
428        let mut list_cost = CostFactors::new();
429        list_cost.cpu_cost = 5.0;
430        list_cost.io_cost = 10.0;
431        weights.insert(OperationType::ListFetch, list_cost);
432
433        let mut join_cost = CostFactors::new();
434        join_cost.cpu_cost = 10.0;
435        join_cost.io_cost = 20.0;
436        weights.insert(OperationType::Join, join_cost);
437
438        let mut filter_cost = CostFactors::new();
439        filter_cost.cpu_cost = 3.0;
440        weights.insert(OperationType::Filter, filter_cost);
441
442        let mut agg_cost = CostFactors::new();
443        agg_cost.cpu_cost = 8.0;
444        agg_cost.memory_cost = 5.0;
445        weights.insert(OperationType::Aggregation, agg_cost);
446
447        let mut sort_cost = CostFactors::new();
448        sort_cost.cpu_cost = 6.0;
449        sort_cost.memory_cost = 4.0;
450        weights.insert(OperationType::Sort, sort_cost);
451
452        Self { weights }
453    }
454
455    /// Estimate cost for an operation
456    pub fn estimate_cost(
457        &mut self,
458        op_type: OperationType,
459        row_count: usize,
460        statistics: Option<&QueryStatistics>,
461    ) -> OperationCost {
462        let base_cost = self.weights.get(&op_type).cloned().unwrap_or_default();
463
464        let mut factors = base_cost.clone();
465
466        // Scale by row count
467        let scale_factor = (row_count as f64).ln().max(1.0);
468        factors.cpu_cost *= scale_factor;
469        factors.io_cost *= scale_factor;
470
471        // Use statistics if available
472        let estimated_time_ms = if let Some(stats) = statistics {
473            (stats.avg_execution_time_ms * (row_count as f64 / stats.avg_row_count)) as u64
474        } else {
475            // Rough estimation: 1ms per 100 rows
476            (row_count as f64 / 100.0).max(1.0) as u64
477        };
478
479        let confidence = if statistics.is_some() { 0.9 } else { 0.5 };
480
481        OperationCost::new(op_type, "operation".to_string())
482            .with_cpu_cost(factors.cpu_cost)
483            .with_io_cost(factors.io_cost)
484            .with_network_cost(factors.network_cost)
485            .with_memory_cost(factors.memory_cost)
486            .with_estimated_rows(row_count)
487            .with_estimated_time(estimated_time_ms)
488            .with_confidence(confidence)
489    }
490
491    /// Monte Carlo cost simulation
492    pub fn simulate_cost(&mut self, plan: &CostBasedPlan, iterations: usize) -> f64 {
493        let mut total = 0.0;
494
495        for _ in 0..iterations {
496            let mut cost = 0.0;
497            for op in &plan.operations {
498                // Add random variation using fastrand (already in dependencies)
499                let variation = 0.9 + fastrand::f64() * 0.2; // Random between 0.9 and 1.1
500                cost += op.total_cost() * variation;
501            }
502            total += cost;
503        }
504
505        total / iterations as f64
506    }
507}
508
509impl Default for CostModel {
510    fn default() -> Self {
511        Self::new()
512    }
513}
514
515/// Cost-based query optimizer
516pub struct CostBasedOptimizer {
517    config: OptimizationConfig,
518    cost_model: Arc<RwLock<CostModel>>,
519    statistics: Arc<RwLock<HashMap<String, QueryStatistics>>>,
520}
521
522impl CostBasedOptimizer {
523    /// Create a new cost-based optimizer
524    pub fn new(config: OptimizationConfig) -> Self {
525        Self {
526            config,
527            cost_model: Arc::new(RwLock::new(CostModel::new())),
528            statistics: Arc::new(RwLock::new(HashMap::new())),
529        }
530    }
531
532    /// Optimize a query
533    pub async fn optimize(&self, query: &str) -> Result<CostBasedPlan, OptimizationError> {
534        // Generate alternative plans
535        let plans = self.generate_alternative_plans(query).await?;
536
537        // Select best plan
538        let best_plan = self.select_best_plan(plans).await?;
539
540        Ok(best_plan)
541    }
542
543    /// Generate alternative execution plans
544    async fn generate_alternative_plans(
545        &self,
546        query: &str,
547    ) -> Result<Vec<CostBasedPlan>, OptimizationError> {
548        let mut plans = Vec::new();
549
550        // Generate sequential plan
551        let sequential = self.generate_plan(query, PlanStrategy::Sequential).await?;
552        plans.push(sequential);
553
554        // Generate parallel plan
555        if self.config.max_alternative_plans > 1 {
556            let parallel = self.generate_plan(query, PlanStrategy::Parallel).await?;
557            plans.push(parallel);
558        }
559
560        // Generate batch plan
561        if self.config.max_alternative_plans > 2 {
562            let batch = self.generate_plan(query, PlanStrategy::Batch).await?;
563            plans.push(batch);
564        }
565
566        // Generate streaming plan
567        if self.config.max_alternative_plans > 3 {
568            let streaming = self.generate_plan(query, PlanStrategy::Streaming).await?;
569            plans.push(streaming);
570        }
571
572        // Generate hybrid plan
573        if self.config.max_alternative_plans > 4 {
574            let hybrid = self.generate_plan(query, PlanStrategy::Hybrid).await?;
575            plans.push(hybrid);
576        }
577
578        Ok(plans)
579    }
580
581    /// Generate a specific plan
582    async fn generate_plan(
583        &self,
584        query: &str,
585        strategy: PlanStrategy,
586    ) -> Result<CostBasedPlan, OptimizationError> {
587        let fingerprint = self.fingerprint_query(query);
588        let mut plan = CostBasedPlan::new(format!("plan_{:?}", strategy), strategy);
589
590        // Get statistics
591        let stats = {
592            let stats_map = self.statistics.read().await;
593            stats_map.get(&fingerprint).cloned()
594        };
595
596        // Estimate operations (simplified for mock)
597        let row_count = 100; // Mock value
598
599        let mut cost_model = self.cost_model.write().await;
600
601        // Add mock operations
602        let field_op = cost_model.estimate_cost(OperationType::FieldSelect, 10, stats.as_ref());
603        plan.add_operation(field_op);
604
605        let object_op = cost_model.estimate_cost(OperationType::ObjectFetch, 50, stats.as_ref());
606        plan.add_operation(object_op);
607
608        let list_op = cost_model.estimate_cost(OperationType::ListFetch, row_count, stats.as_ref());
609        plan.add_operation(list_op);
610
611        // Add index recommendations if enabled
612        if self.config.enable_index_recommendations {
613            let rec = IndexRecommendation {
614                table: "users".to_string(),
615                fields: vec!["id".to_string()],
616                index_type: IndexType::BTree,
617                expected_improvement: 25.0,
618                rationale: "Frequent lookups on id field".to_string(),
619            };
620            plan.add_index_recommendation(rec);
621        }
622
623        Ok(plan)
624    }
625
626    /// Select the best plan from alternatives
627    async fn select_best_plan(
628        &self,
629        plans: Vec<CostBasedPlan>,
630    ) -> Result<CostBasedPlan, OptimizationError> {
631        if plans.is_empty() {
632            return Err(OptimizationError::NoPlanGenerated);
633        }
634
635        // Find plan with minimum cost
636        let best = plans
637            .into_iter()
638            .min_by(|a, b| {
639                a.total_cost
640                    .partial_cmp(&b.total_cost)
641                    .expect("cost values should not be NaN")
642            })
643            .expect("plans is not empty (checked above)");
644
645        Ok(best)
646    }
647
648    /// Record execution statistics
649    pub async fn record_execution(
650        &self,
651        query: &str,
652        execution_time_ms: u64,
653        row_count: usize,
654    ) -> Result<(), OptimizationError> {
655        if !self.config.enable_statistics {
656            return Ok(());
657        }
658
659        // Sample based on configured rate
660        if fastrand::f64() > self.config.statistics_sample_rate {
661            return Ok(());
662        }
663
664        let fingerprint = self.fingerprint_query(query);
665
666        let mut stats_map = self.statistics.write().await;
667        let stats = stats_map
668            .entry(fingerprint.clone())
669            .or_insert_with(|| QueryStatistics::new(fingerprint));
670
671        stats.update(execution_time_ms, row_count);
672
673        Ok(())
674    }
675
676    /// Get statistics for a query
677    pub async fn get_statistics(&self, query: &str) -> Option<QueryStatistics> {
678        let fingerprint = self.fingerprint_query(query);
679        let stats_map = self.statistics.read().await;
680        stats_map.get(&fingerprint).cloned()
681    }
682
683    /// Generate query fingerprint
684    fn fingerprint_query(&self, query: &str) -> String {
685        // Simple fingerprint: normalize whitespace and use simple hash
686        let normalized = query.split_whitespace().collect::<Vec<_>>().join(" ");
687
688        // Simple hash function (FNV-1a)
689        let mut hash: u64 = 0xcbf29ce484222325;
690        for byte in normalized.bytes() {
691            hash ^= byte as u64;
692            hash = hash.wrapping_mul(0x100000001b3);
693        }
694
695        format!("{:016x}", hash)
696    }
697
698    /// Clear all statistics
699    pub async fn clear_statistics(&self) {
700        let mut stats_map = self.statistics.write().await;
701        stats_map.clear();
702    }
703
704    /// Get number of tracked queries
705    pub async fn statistics_count(&self) -> usize {
706        let stats_map = self.statistics.read().await;
707        stats_map.len()
708    }
709}
710
711/// Errors that can occur during optimization
712#[derive(Debug, thiserror::Error)]
713pub enum OptimizationError {
714    /// No plan could be generated
715    #[error("No execution plan could be generated")]
716    NoPlanGenerated,
717
718    /// Cost estimation failed
719    #[error("Cost estimation failed: {0}")]
720    CostEstimationFailed(String),
721
722    /// Statistics error
723    #[error("Statistics error: {0}")]
724    StatisticsError(String),
725
726    /// Invalid query
727    #[error("Invalid query: {0}")]
728    InvalidQuery(String),
729}
730
731#[cfg(test)]
732mod tests {
733    use super::*;
734
735    #[test]
736    fn test_optimization_config_builder() {
737        let config = OptimizationConfig::new()
738            .with_statistics_collection(true)
739            .with_adaptive_learning(true)
740            .with_max_alternative_plans(10)
741            .with_cost_threshold(500.0)
742            .with_join_optimization(true)
743            .with_index_recommendations(true)
744            .with_statistics_sample_rate(0.5);
745
746        assert!(config.enable_statistics);
747        assert!(config.enable_adaptive_learning);
748        assert_eq!(config.max_alternative_plans, 10);
749        assert_eq!(config.cost_threshold, 500.0);
750        assert!(config.enable_join_optimization);
751        assert!(config.enable_index_recommendations);
752        assert_eq!(config.statistics_sample_rate, 0.5);
753    }
754
755    #[test]
756    fn test_cost_factors_total() {
757        let mut factors = CostFactors::new();
758        factors.cpu_cost = 10.0;
759        factors.io_cost = 5.0;
760        factors.network_cost = 3.0;
761        factors.memory_cost = 2.0;
762
763        assert_eq!(factors.total(), 20.0);
764    }
765
766    #[test]
767    fn test_cost_factors_add() {
768        let mut factors1 = CostFactors::new();
769        factors1.cpu_cost = 10.0;
770        factors1.io_cost = 5.0;
771
772        let mut factors2 = CostFactors::new();
773        factors2.cpu_cost = 3.0;
774        factors2.network_cost = 2.0;
775
776        factors1.add(&factors2);
777
778        assert_eq!(factors1.cpu_cost, 13.0);
779        assert_eq!(factors1.io_cost, 5.0);
780        assert_eq!(factors1.network_cost, 2.0);
781    }
782
783    #[test]
784    fn test_operation_cost_creation() {
785        let cost = OperationCost::new(OperationType::FieldSelect, "test".to_string())
786            .with_cpu_cost(5.0)
787            .with_io_cost(2.0)
788            .with_estimated_rows(100)
789            .with_estimated_time(50)
790            .with_confidence(0.8);
791
792        assert_eq!(cost.operation, OperationType::FieldSelect);
793        assert_eq!(cost.factors.cpu_cost, 5.0);
794        assert_eq!(cost.factors.io_cost, 2.0);
795        assert_eq!(cost.estimated_rows, 100);
796        assert_eq!(cost.estimated_time_ms, 50);
797        assert_eq!(cost.confidence, 0.8);
798    }
799
800    #[test]
801    fn test_operation_cost_total() {
802        let cost = OperationCost::new(OperationType::Join, "join".to_string())
803            .with_cpu_cost(10.0)
804            .with_io_cost(20.0)
805            .with_network_cost(5.0);
806
807        assert_eq!(cost.total_cost(), 35.0);
808    }
809
810    #[test]
811    fn test_plan_creation() {
812        let plan = CostBasedPlan::new("plan1".to_string(), PlanStrategy::Sequential);
813
814        assert_eq!(plan.plan_id, "plan1");
815        assert_eq!(plan.strategy, PlanStrategy::Sequential);
816        assert_eq!(plan.total_cost, 0.0);
817        assert_eq!(plan.total_estimated_time_ms, 0);
818    }
819
820    #[test]
821    fn test_plan_add_operation() {
822        let mut plan = CostBasedPlan::new("plan1".to_string(), PlanStrategy::Parallel);
823
824        let op1 = OperationCost::new(OperationType::FieldSelect, "field1".to_string())
825            .with_cpu_cost(5.0)
826            .with_estimated_time(10);
827
828        let op2 = OperationCost::new(OperationType::ObjectFetch, "obj1".to_string())
829            .with_cpu_cost(10.0)
830            .with_estimated_time(20);
831
832        plan.add_operation(op1);
833        plan.add_operation(op2);
834
835        assert_eq!(plan.operations.len(), 2);
836        assert_eq!(plan.total_cost, 15.0);
837        assert_eq!(plan.total_estimated_time_ms, 30);
838    }
839
840    #[test]
841    fn test_plan_comparison() {
842        let mut plan1 = CostBasedPlan::new("plan1".to_string(), PlanStrategy::Sequential);
843        plan1.total_cost = 100.0;
844
845        let mut plan2 = CostBasedPlan::new("plan2".to_string(), PlanStrategy::Parallel);
846        plan2.total_cost = 50.0;
847
848        assert!(plan2.is_better_than(&plan1));
849        assert!(!plan1.is_better_than(&plan2));
850    }
851
852    #[test]
853    fn test_query_statistics_creation() {
854        let stats = QueryStatistics::new("fingerprint123".to_string());
855
856        assert_eq!(stats.fingerprint, "fingerprint123");
857        assert_eq!(stats.execution_count, 0);
858        assert_eq!(stats.avg_execution_time_ms, 0.0);
859    }
860
861    #[test]
862    fn test_query_statistics_update() {
863        let mut stats = QueryStatistics::new("test".to_string());
864
865        stats.update(100, 50);
866        assert_eq!(stats.execution_count, 1);
867        assert_eq!(stats.avg_execution_time_ms, 100.0);
868        assert_eq!(stats.avg_row_count, 50.0);
869        assert_eq!(stats.min_execution_time_ms, 100);
870        assert_eq!(stats.max_execution_time_ms, 100);
871
872        stats.update(200, 100);
873        assert_eq!(stats.execution_count, 2);
874        assert_eq!(stats.avg_execution_time_ms, 150.0);
875        assert_eq!(stats.avg_row_count, 75.0);
876        assert_eq!(stats.min_execution_time_ms, 100);
877        assert_eq!(stats.max_execution_time_ms, 200);
878    }
879
880    #[test]
881    fn test_cost_model_creation() {
882        let model = CostModel::new();
883        assert!(!model.weights.is_empty());
884    }
885
886    #[test]
887    fn test_cost_model_estimate() {
888        let mut model = CostModel::new();
889
890        let cost = model.estimate_cost(OperationType::FieldSelect, 100, None);
891
892        assert_eq!(cost.operation, OperationType::FieldSelect);
893        assert!(cost.total_cost() > 0.0);
894        assert_eq!(cost.estimated_rows, 100);
895        assert_eq!(cost.confidence, 0.5);
896    }
897
898    #[test]
899    fn test_cost_model_with_statistics() {
900        let mut model = CostModel::new();
901        let mut stats = QueryStatistics::new("test".to_string());
902        stats.update(100, 50);
903        stats.update(200, 100);
904
905        let cost = model.estimate_cost(OperationType::ObjectFetch, 100, Some(&stats));
906
907        assert!(cost.confidence > 0.5);
908        assert!(cost.estimated_time_ms > 0);
909    }
910
911    #[test]
912    fn test_monte_carlo_simulation() {
913        let mut model = CostModel::new();
914        let mut plan = CostBasedPlan::new("test".to_string(), PlanStrategy::Sequential);
915
916        let op = OperationCost::new(OperationType::Join, "join".to_string()).with_cpu_cost(50.0);
917
918        plan.add_operation(op);
919
920        let simulated_cost = model.simulate_cost(&plan, 100);
921        assert!(simulated_cost > 0.0);
922    }
923
924    #[tokio::test]
925    async fn test_optimizer_creation() {
926        let config = OptimizationConfig::new();
927        let optimizer = CostBasedOptimizer::new(config);
928
929        assert_eq!(optimizer.statistics_count().await, 0);
930    }
931
932    #[tokio::test]
933    async fn test_query_optimization() {
934        let config = OptimizationConfig::new();
935        let optimizer = CostBasedOptimizer::new(config);
936
937        let query = "{ user { id name posts { title } } }";
938        let plan = optimizer.optimize(query).await;
939
940        assert!(plan.is_ok());
941        let plan = plan.expect("should succeed");
942        assert!(!plan.operations.is_empty());
943        assert!(plan.total_cost > 0.0);
944    }
945
946    #[tokio::test]
947    async fn test_alternative_plans_generation() {
948        let config = OptimizationConfig::new().with_max_alternative_plans(5);
949        let optimizer = CostBasedOptimizer::new(config);
950
951        let query = "{ test }";
952        let plans = optimizer
953            .generate_alternative_plans(query)
954            .await
955            .expect("should succeed");
956
957        assert_eq!(plans.len(), 5);
958        assert!(plans.iter().any(|p| p.strategy == PlanStrategy::Sequential));
959        assert!(plans.iter().any(|p| p.strategy == PlanStrategy::Parallel));
960        assert!(plans.iter().any(|p| p.strategy == PlanStrategy::Batch));
961        assert!(plans.iter().any(|p| p.strategy == PlanStrategy::Streaming));
962        assert!(plans.iter().any(|p| p.strategy == PlanStrategy::Hybrid));
963    }
964
965    #[tokio::test]
966    async fn test_best_plan_selection() {
967        let config = OptimizationConfig::new();
968        let optimizer = CostBasedOptimizer::new(config);
969
970        let mut plan1 = CostBasedPlan::new("plan1".to_string(), PlanStrategy::Sequential);
971        plan1.total_cost = 100.0;
972
973        let mut plan2 = CostBasedPlan::new("plan2".to_string(), PlanStrategy::Parallel);
974        plan2.total_cost = 50.0;
975
976        let mut plan3 = CostBasedPlan::new("plan3".to_string(), PlanStrategy::Batch);
977        plan3.total_cost = 75.0;
978
979        let plans = vec![plan1, plan2.clone(), plan3];
980        let best = optimizer
981            .select_best_plan(plans)
982            .await
983            .expect("should succeed");
984
985        assert_eq!(best.plan_id, plan2.plan_id);
986        assert_eq!(best.total_cost, 50.0);
987    }
988
989    #[tokio::test]
990    async fn test_record_execution() {
991        let config = OptimizationConfig::new().with_statistics_collection(true);
992        let optimizer = CostBasedOptimizer::new(config);
993
994        let query = "{ test }";
995        optimizer
996            .record_execution(query, 100, 50)
997            .await
998            .expect("should succeed");
999
1000        let stats = optimizer.get_statistics(query).await;
1001        assert!(stats.is_some());
1002
1003        let stats = stats.expect("should succeed");
1004        assert_eq!(stats.execution_count, 1);
1005        assert_eq!(stats.avg_execution_time_ms, 100.0);
1006    }
1007
1008    #[tokio::test]
1009    async fn test_statistics_disabled() {
1010        let config = OptimizationConfig::new().with_statistics_collection(false);
1011        let optimizer = CostBasedOptimizer::new(config);
1012
1013        let query = "{ test }";
1014        optimizer
1015            .record_execution(query, 100, 50)
1016            .await
1017            .expect("should succeed");
1018
1019        assert_eq!(optimizer.statistics_count().await, 0);
1020    }
1021
1022    #[tokio::test]
1023    async fn test_clear_statistics() {
1024        let config = OptimizationConfig::new();
1025        let optimizer = CostBasedOptimizer::new(config);
1026
1027        optimizer
1028            .record_execution("{ test1 }", 100, 50)
1029            .await
1030            .expect("should succeed");
1031        optimizer
1032            .record_execution("{ test2 }", 200, 100)
1033            .await
1034            .expect("should succeed");
1035
1036        assert_eq!(optimizer.statistics_count().await, 2);
1037
1038        optimizer.clear_statistics().await;
1039        assert_eq!(optimizer.statistics_count().await, 0);
1040    }
1041
1042    #[tokio::test]
1043    async fn test_query_fingerprinting() {
1044        let config = OptimizationConfig::new();
1045        let optimizer = CostBasedOptimizer::new(config);
1046
1047        let query1 = "{ user { id name } }";
1048        let query2 = "{  user  {  id   name  }  }"; // Different whitespace
1049
1050        let fp1 = optimizer.fingerprint_query(query1);
1051        let fp2 = optimizer.fingerprint_query(query2);
1052
1053        assert_eq!(fp1, fp2);
1054    }
1055
1056    #[tokio::test]
1057    async fn test_index_recommendations() {
1058        let config = OptimizationConfig::new().with_index_recommendations(true);
1059        let optimizer = CostBasedOptimizer::new(config);
1060
1061        let query = "{ users { id } }";
1062        let plan = optimizer.optimize(query).await.expect("should succeed");
1063
1064        assert!(!plan.index_recommendations.is_empty());
1065        assert!(plan.index_recommendations[0].expected_improvement > 0.0);
1066    }
1067
1068    #[test]
1069    fn test_index_recommendation_creation() {
1070        let rec = IndexRecommendation {
1071            table: "users".to_string(),
1072            fields: vec!["email".to_string()],
1073            index_type: IndexType::Hash,
1074            expected_improvement: 30.0,
1075            rationale: "Frequent equality lookups".to_string(),
1076        };
1077
1078        assert_eq!(rec.table, "users");
1079        assert_eq!(rec.fields.len(), 1);
1080        assert_eq!(rec.index_type, IndexType::Hash);
1081        assert_eq!(rec.expected_improvement, 30.0);
1082    }
1083
1084    #[test]
1085    fn test_confidence_clamping() {
1086        let cost =
1087            OperationCost::new(OperationType::FieldSelect, "test".to_string()).with_confidence(1.5);
1088
1089        assert_eq!(cost.confidence, 1.0);
1090
1091        let cost2 = OperationCost::new(OperationType::FieldSelect, "test".to_string())
1092            .with_confidence(-0.5);
1093
1094        assert_eq!(cost2.confidence, 0.0);
1095    }
1096}
oxirs_gql/cost_based_optimizer.rs

oxirs_gql/
cost_based_optimizer.rs
