oxirs_core/query/

optimizer.rs

//! AI-powered query optimization with learned cost models
//!
//! This module implements advanced query optimization using machine learning
//! techniques to improve query performance based on historical patterns.

#![allow(dead_code)]

use crate::indexing::IndexStats;
use crate::model::Variable;
use crate::query::algebra::{
    AlgebraTriplePattern, GraphPattern, Query as AlgebraQuery, QueryForm, TermPattern,
};
use crate::query::plan::{ExecutionPlan, QueryPlanner};
use crate::OxirsError;
use std::collections::{HashMap, VecDeque};
use std::sync::{Arc, RwLock};
use std::time::{Duration, Instant};

/// Cost model for query optimization
#[derive(Debug, Clone)]
pub struct CostModel {
    /// Historical query execution times
    execution_history: Arc<RwLock<QueryHistory>>,
    /// Learned parameters for cost estimation
    learned_parameters: Arc<RwLock<LearnedParameters>>,
    /// Index statistics
    index_stats: Arc<IndexStats>,
}

/// Historical query execution data
#[derive(Debug, Default)]
struct QueryHistory {
    /// Recent query patterns and their execution times
    patterns: VecDeque<(QueryPattern, ExecutionMetrics)>,
    /// Maximum history size
    max_size: usize,
}

/// Learned parameters from query history
#[derive(Debug, Default)]
struct LearnedParameters {
    /// Cost per triple scan by predicate
    scan_costs: HashMap<String, f64>,
    /// Join selectivity by pattern
    join_selectivities: HashMap<JoinPattern, f64>,
    /// Filter selectivity by expression type
    filter_selectivities: HashMap<String, f64>,
}

/// Query pattern for learning
#[derive(Debug, Clone, Hash, Eq, PartialEq)]
struct QueryPattern {
    /// Number of triple patterns
    num_patterns: usize,
    /// Predicates used
    predicates: Vec<String>,
    /// Join types
    join_types: Vec<JoinType>,
    /// Filter presence
    has_filter: bool,
}

/// Join pattern for selectivity estimation
#[derive(Debug, Clone, Hash, Eq, PartialEq)]
struct JoinPattern {
    /// Variables involved in join
    num_vars: usize,
    /// Types of terms being joined
    term_types: Vec<String>,
}

/// Types of joins
#[derive(Debug, Clone, Hash, Eq, PartialEq)]
enum JoinType {
    SubjectSubject,
    SubjectObject,
    ObjectObject,
    PredicatePredicate,
}

/// Execution metrics for a query
#[derive(Debug, Clone)]
struct ExecutionMetrics {
    /// Total execution time
    execution_time: Duration,
    /// Number of results returned
    result_count: usize,
    /// Memory usage
    memory_used: usize,
    /// CPU utilization
    cpu_percent: f32,
}

/// AI-powered query optimizer
pub struct AIQueryOptimizer {
    /// Base query planner
    base_planner: QueryPlanner,
    /// Cost model for optimization
    cost_model: CostModel,
    /// Query cache for predictive caching
    query_cache: Arc<RwLock<QueryCache>>,
    /// Hardware capabilities
    hardware_info: HardwareInfo,
}

/// Predictive query cache
#[derive(Debug, Default)]
struct QueryCache {
    /// Cached query results
    cache: HashMap<String, CachedResult>,
    /// Access patterns for prediction
    access_patterns: VecDeque<AccessPattern>,
    /// Maximum cache size
    max_size: usize,
}

/// Cached query result
#[derive(Debug, Clone)]
struct CachedResult {
    /// The cached data
    data: Vec<u8>,
    /// When it was cached
    cached_at: Instant,
    /// How many times accessed
    access_count: usize,
    /// Last access time
    last_accessed: Instant,
}

/// Query access pattern
#[derive(Debug, Clone)]
struct AccessPattern {
    /// Query hash
    query_hash: String,
    /// Access time
    accessed_at: Instant,
    /// User or session ID
    session_id: String,
}

/// Hardware information for adaptive optimization
#[derive(Debug, Clone)]
struct HardwareInfo {
    /// Number of CPU cores
    cpu_cores: usize,
    /// Available memory in bytes
    memory_bytes: usize,
    /// CPU architecture features
    cpu_features: CpuFeatures,
    /// GPU availability
    gpu_available: bool,
}

/// CPU features for optimization
#[derive(Debug, Clone)]
struct CpuFeatures {
    /// SIMD support
    has_simd: bool,
    /// AVX2 support
    has_avx2: bool,
    /// Cache line size
    cache_line_size: usize,
}

impl AIQueryOptimizer {
    /// Create a new AI-powered query optimizer
    pub fn new(index_stats: Arc<IndexStats>) -> Self {
        Self {
            base_planner: QueryPlanner::new(),
            cost_model: CostModel::new(index_stats),
            query_cache: Arc::new(RwLock::new(QueryCache::new())),
            hardware_info: HardwareInfo::detect(),
        }
    }

    /// Optimize a query using AI techniques
    pub fn optimize_query(&self, query: &AlgebraQuery) -> Result<OptimizedPlan, OxirsError> {
        // Extract query pattern for learning
        let pattern = self.extract_query_pattern(query)?;

        // Check cache for similar queries
        if let Some(cached) = self.check_predictive_cache(&pattern) {
            return Ok(cached);
        }

        // Generate multiple candidate plans
        let candidates = self.generate_candidate_plans(query)?;

        // Estimate costs using learned model
        let mut best_plan = None;
        let mut best_cost = f64::MAX;

        for candidate in candidates {
            let cost = self.estimate_cost(&candidate, &pattern)?;
            if cost < best_cost {
                best_cost = cost;
                best_plan = Some(candidate);
            }
        }

        let plan = best_plan
            .ok_or_else(|| OxirsError::Query("No valid execution plan found".to_string()))?;

        // Apply hardware-specific optimizations
        let optimized = self.apply_hardware_optimizations(plan)?;

        // Update learning model
        self.update_learning_model(&pattern, &optimized);

        Ok(optimized)
    }

    /// Extract pattern from query for learning
    fn extract_query_pattern(&self, query: &AlgebraQuery) -> Result<QueryPattern, OxirsError> {
        match &query.form {
            QueryForm::Select { where_clause, .. } => {
                let (num_patterns, predicates, join_types) =
                    self.analyze_graph_pattern(where_clause)?;

                Ok(QueryPattern {
                    num_patterns,
                    predicates,
                    join_types,
                    has_filter: self.has_filter(where_clause),
                })
            }
            _ => Err(OxirsError::Query("Unsupported query form".to_string())),
        }
    }

    /// Analyze graph pattern for optimization
    fn analyze_graph_pattern(
        &self,
        pattern: &GraphPattern,
    ) -> Result<(usize, Vec<String>, Vec<JoinType>), OxirsError> {
        match pattern {
            GraphPattern::Bgp(patterns) => {
                let num_patterns = patterns.len();
                let mut predicates = Vec::new();
                let mut join_types = Vec::new();

                // Extract predicates
                for triple in patterns {
                    if let TermPattern::NamedNode(pred) = &triple.predicate {
                        predicates.push(pred.as_str().to_string());
                    }
                }

                // Analyze join types between patterns
                for i in 0..patterns.len() {
                    for j in (i + 1)..patterns.len() {
                        if let Some(join_type) = self.get_join_type(&patterns[i], &patterns[j]) {
                            join_types.push(join_type);
                        }
                    }
                }

                Ok((num_patterns, predicates, join_types))
            }
            _ => Ok((0, Vec::new(), Vec::new())),
        }
    }

    /// Determine join type between triple patterns
    fn get_join_type(
        &self,
        left: &AlgebraTriplePattern,
        right: &AlgebraTriplePattern,
    ) -> Option<JoinType> {
        // Check if subjects match
        if self.patterns_match(&left.subject, &right.subject) {
            return Some(JoinType::SubjectSubject);
        }

        // Check subject-object join
        if self.patterns_match(&left.subject, &right.object) {
            return Some(JoinType::SubjectObject);
        }

        // Check object-object join
        if self.patterns_match(&left.object, &right.object) {
            return Some(JoinType::ObjectObject);
        }

        // Check predicate join (rare)
        if self.patterns_match(&left.predicate, &right.predicate) {
            return Some(JoinType::PredicatePredicate);
        }

        None
    }

    /// Check if two term patterns match (share a variable)
    fn patterns_match(&self, left: &TermPattern, right: &TermPattern) -> bool {
        match (left, right) {
            (TermPattern::Variable(v1), TermPattern::Variable(v2)) => v1 == v2,
            _ => false,
        }
    }

    /// Check if pattern has filters
    #[allow(clippy::only_used_in_recursion)]
    fn has_filter(&self, pattern: &GraphPattern) -> bool {
        match pattern {
            GraphPattern::Filter { .. } => true,
            GraphPattern::Bgp(_) => false,
            GraphPattern::Union(left, right) => self.has_filter(left) || self.has_filter(right),
            _ => false,
        }
    }

    /// Generate candidate execution plans
    fn generate_candidate_plans(
        &self,
        query: &AlgebraQuery,
    ) -> Result<Vec<ExecutionPlan>, OxirsError> {
        let mut candidates = Vec::new();

        // Basic plan from base planner
        let basic_plan = self.base_planner.plan_query(query)?;
        candidates.push(basic_plan.clone());

        // Generate join order variations
        if let QueryForm::Select {
            where_clause: GraphPattern::Bgp(patterns),
            ..
        } = &query.form
        {
            // Try different join orders
            let join_orders = self.generate_join_orders(patterns);
            for order in join_orders {
                if let Ok(plan) = self.create_plan_with_order(patterns, &order) {
                    candidates.push(plan);
                }
            }
        }

        // Add index-based variations
        candidates.extend(self.generate_index_plans(query)?);

        Ok(candidates)
    }

    /// Generate different join orders for optimization
    fn generate_join_orders(&self, patterns: &[AlgebraTriplePattern]) -> Vec<Vec<usize>> {
        let mut orders = Vec::new();

        // Original order
        orders.push((0..patterns.len()).collect());

        // Most selective first (based on statistics)
        let mut selective_order: Vec<usize> = (0..patterns.len()).collect();
        selective_order.sort_by_key(|&i| self.estimate_selectivity(&patterns[i]));
        orders.push(selective_order);

        // Limit to reasonable number of variations
        orders.truncate(5);
        orders
    }

    /// Estimate selectivity of a triple pattern  
    fn estimate_selectivity(&self, pattern: &AlgebraTriplePattern) -> i64 {
        // Lower score = more selective (better to execute first)
        let mut score = 0;

        // Concrete terms are more selective
        if !matches!(pattern.subject, TermPattern::Variable(_)) {
            score -= 1000;
        }
        if !matches!(pattern.predicate, TermPattern::Variable(_)) {
            score -= 100;
        }
        if !matches!(pattern.object, TermPattern::Variable(_)) {
            score -= 1000;
        }

        score
    }

    /// Create execution plan with specific join order
    fn create_plan_with_order(
        &self,
        patterns: &[AlgebraTriplePattern],
        order: &[usize],
    ) -> Result<ExecutionPlan, OxirsError> {
        if order.is_empty() {
            return Err(OxirsError::Query("Empty join order".to_string()));
        }

        let mut plan = ExecutionPlan::TripleScan {
            pattern: crate::query::plan::convert_algebra_triple_pattern(&patterns[order[0]]),
        };

        for &idx in &order[1..] {
            let right_plan = ExecutionPlan::TripleScan {
                pattern: crate::query::plan::convert_algebra_triple_pattern(&patterns[idx]),
            };

            plan = ExecutionPlan::HashJoin {
                left: Box::new(plan),
                right: Box::new(right_plan),
                join_vars: Vec::new(), // Would compute actual join vars
            };
        }

        Ok(plan)
    }

    /// Generate index-based execution plans
    fn generate_index_plans(
        &self,
        _query: &AlgebraQuery,
    ) -> Result<Vec<ExecutionPlan>, OxirsError> {
        // Would generate plans that leverage specific indexes
        Ok(Vec::new())
    }

    /// Estimate cost of execution plan
    fn estimate_cost(
        &self,
        plan: &ExecutionPlan,
        pattern: &QueryPattern,
    ) -> Result<f64, OxirsError> {
        let params = self
            .cost_model
            .learned_parameters
            .read()
            .map_err(|e| OxirsError::Query(format!("Failed to read parameters: {e}")))?;

        let base_cost = self.estimate_plan_cost(plan, &params)?;

        // Adjust based on pattern history
        let history_factor = self.get_history_factor(pattern);

        Ok(base_cost * history_factor)
    }

    /// Estimate base cost of a plan
    #[allow(clippy::only_used_in_recursion)]
    fn estimate_plan_cost(
        &self,
        plan: &ExecutionPlan,
        params: &LearnedParameters,
    ) -> Result<f64, OxirsError> {
        match plan {
            ExecutionPlan::TripleScan { pattern } => {
                // Base scan cost
                let mut cost = 100.0;

                // Adjust based on predicate selectivity
                if let Some(crate::model::pattern::PredicatePattern::NamedNode(pred)) =
                    &pattern.predicate
                {
                    if let Some(&pred_cost) = params.scan_costs.get(pred.as_str()) {
                        cost *= pred_cost;
                    }
                }

                Ok(cost)
            }
            ExecutionPlan::HashJoin { left, right, .. } => {
                let left_cost = self.estimate_plan_cost(left, params)?;
                let right_cost = self.estimate_plan_cost(right, params)?;

                // Join cost depends on input sizes
                Ok(left_cost + right_cost + (left_cost * right_cost * 0.01))
            }
            ExecutionPlan::Filter { input, .. } => {
                let input_cost = self.estimate_plan_cost(input, params)?;
                // Filters typically reduce result size
                Ok(input_cost * 0.5)
            }
            _ => Ok(1000.0), // Default cost
        }
    }

    /// Get historical performance factor
    fn get_history_factor(&self, pattern: &QueryPattern) -> f64 {
        // Check if we've seen similar patterns before
        if let Ok(history) = self.cost_model.execution_history.read() {
            for (hist_pattern, metrics) in history.patterns.iter() {
                if self.patterns_similar(pattern, hist_pattern) {
                    // Adjust based on historical performance
                    return if metrics.execution_time.as_millis() < 100 {
                        0.8 // Performed well historically
                    } else {
                        1.2 // Performed poorly
                    };
                }
            }
        }
        1.0 // No history
    }

    /// Check if patterns are similar
    fn patterns_similar(&self, a: &QueryPattern, b: &QueryPattern) -> bool {
        a.num_patterns == b.num_patterns
            && a.has_filter == b.has_filter
            && a.predicates.len() == b.predicates.len()
    }

    /// Check predictive cache
    fn check_predictive_cache(&self, _pattern: &QueryPattern) -> Option<OptimizedPlan> {
        // Would check cache for similar queries
        None
    }

    /// Apply hardware-specific optimizations
    fn apply_hardware_optimizations(
        &self,
        plan: ExecutionPlan,
    ) -> Result<OptimizedPlan, OxirsError> {
        let mut optimized = OptimizedPlan {
            base_plan: plan,
            parallelism_level: 1,
            use_simd: false,
            use_gpu: false,
            memory_budget: 0,
        };

        // Set parallelism based on CPU cores
        optimized.parallelism_level = self.calculate_optimal_parallelism();

        // Enable SIMD if available
        optimized.use_simd = self.hardware_info.cpu_features.has_simd;

        // Consider GPU for large operations
        optimized.use_gpu =
            self.hardware_info.gpu_available && self.should_use_gpu(&optimized.base_plan);

        // Set memory budget
        optimized.memory_budget = self.calculate_memory_budget();

        Ok(optimized)
    }

    /// Calculate optimal parallelism level
    fn calculate_optimal_parallelism(&self) -> usize {
        // Use 75% of cores to leave room for system
        (self.hardware_info.cpu_cores as f32 * 0.75) as usize
    }

    /// Determine if GPU should be used
    fn should_use_gpu(&self, _plan: &ExecutionPlan) -> bool {
        // Would analyze plan complexity and data size
        false // Placeholder
    }

    /// Calculate memory budget for query
    fn calculate_memory_budget(&self) -> usize {
        // Use 50% of available memory
        self.hardware_info.memory_bytes / 2
    }

    /// Update learning model with execution results
    fn update_learning_model(&self, pattern: &QueryPattern, _plan: &OptimizedPlan) {
        // Record pattern for future learning
        if let Ok(mut history) = self.cost_model.execution_history.write() {
            let metrics = ExecutionMetrics {
                execution_time: Duration::from_millis(50), // Would get actual time
                result_count: 100,                         // Would get actual count
                memory_used: 1024 * 1024,                  // Would measure actual usage
                cpu_percent: 25.0,                         // Would measure actual CPU
            };

            history.add_execution(pattern.clone(), metrics);
        }
    }
}

/// Optimized execution plan with hardware hints
#[derive(Debug)]
pub struct OptimizedPlan {
    /// Base execution plan
    pub base_plan: ExecutionPlan,
    /// Parallelism level to use
    pub parallelism_level: usize,
    /// Whether to use SIMD instructions
    pub use_simd: bool,
    /// Whether to use GPU acceleration
    pub use_gpu: bool,
    /// Memory budget in bytes
    pub memory_budget: usize,
}

impl CostModel {
    fn new(index_stats: Arc<IndexStats>) -> Self {
        Self {
            execution_history: Arc::new(RwLock::new(QueryHistory::new())),
            learned_parameters: Arc::new(RwLock::new(LearnedParameters::default())),
            index_stats,
        }
    }
}

impl QueryHistory {
    fn new() -> Self {
        Self {
            patterns: VecDeque::new(),
            max_size: 10000,
        }
    }

    fn add_execution(&mut self, pattern: QueryPattern, metrics: ExecutionMetrics) {
        self.patterns.push_back((pattern, metrics));

        // Keep history bounded
        while self.patterns.len() > self.max_size {
            self.patterns.pop_front();
        }
    }
}

impl QueryCache {
    fn new() -> Self {
        Self {
            cache: HashMap::new(),
            access_patterns: VecDeque::new(),
            max_size: 1000,
        }
    }
}

impl HardwareInfo {
    fn detect() -> Self {
        Self {
            cpu_cores: std::thread::available_parallelism()
                .map(|p| p.get())
                .unwrap_or(1),
            memory_bytes: 8 * 1024 * 1024 * 1024, // 8GB default
            cpu_features: CpuFeatures {
                has_simd: cfg!(target_feature = "sse2"),
                has_avx2: cfg!(target_feature = "avx2"),
                cache_line_size: 64,
            },
            gpu_available: false, // Would detect actual GPU
        }
    }
}

/// Multi-query optimizer for batch processing
pub struct MultiQueryOptimizer {
    /// Single query optimizer
    single_optimizer: AIQueryOptimizer,
    /// Shared subexpression detection
    subexpression_cache: Arc<RwLock<HashMap<String, ExecutionPlan>>>,
}

impl MultiQueryOptimizer {
    /// Create new multi-query optimizer
    pub fn new(index_stats: Arc<IndexStats>) -> Self {
        Self {
            single_optimizer: AIQueryOptimizer::new(index_stats),
            subexpression_cache: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// Optimize multiple queries together
    pub fn optimize_batch(
        &self,
        queries: &[AlgebraQuery],
    ) -> Result<Vec<OptimizedPlan>, OxirsError> {
        // Detect common subexpressions
        let common_subs = self.detect_common_subexpressions(queries)?;

        // Create shared execution plans
        let mut optimized_plans = Vec::new();

        for query in queries {
            let mut plan = self.single_optimizer.optimize_query(query)?;

            // Replace common subexpressions with shared plans
            plan = self.reuse_common_subexpressions(plan, &common_subs)?;

            optimized_plans.push(plan);
        }

        Ok(optimized_plans)
    }

    /// Detect common subexpressions across queries
    fn detect_common_subexpressions(
        &self,
        queries: &[AlgebraQuery],
    ) -> Result<HashMap<String, ExecutionPlan>, OxirsError> {
        let mut common_subs = HashMap::new();

        // Extract patterns from all queries
        let mut pattern_counts = HashMap::new();

        for query in queries {
            self.count_patterns(query, &mut pattern_counts)?;
        }

        // Find patterns that appear multiple times
        for (pattern_key, count) in pattern_counts {
            if count > 1 {
                // Create shared plan for this pattern
                // (Simplified - would create actual plan)
                common_subs.insert(
                    pattern_key,
                    ExecutionPlan::TripleScan {
                        pattern: crate::model::pattern::TriplePattern::new(
                            Some(crate::model::pattern::SubjectPattern::Variable(
                                Variable::new("?s").unwrap(),
                            )),
                            Some(crate::model::pattern::PredicatePattern::Variable(
                                Variable::new("?p").unwrap(),
                            )),
                            Some(crate::model::pattern::ObjectPattern::Variable(
                                Variable::new("?o").unwrap(),
                            )),
                        ),
                    },
                );
            }
        }

        Ok(common_subs)
    }

    /// Count pattern occurrences
    fn count_patterns(
        &self,
        query: &AlgebraQuery,
        counts: &mut HashMap<String, usize>,
    ) -> Result<(), OxirsError> {
        if let QueryForm::Select { where_clause, .. } = &query.form {
            self.count_graph_patterns(where_clause, counts)?;
        }
        Ok(())
    }

    /// Count patterns in graph pattern
    fn count_graph_patterns(
        &self,
        pattern: &GraphPattern,
        counts: &mut HashMap<String, usize>,
    ) -> Result<(), OxirsError> {
        if let GraphPattern::Bgp(patterns) = pattern {
            for triple in patterns {
                let key = format!("{triple}"); // Simplified
                *counts.entry(key).or_insert(0) += 1;
            }
        }
        Ok(())
    }

    /// Reuse common subexpressions in plan
    fn reuse_common_subexpressions(
        &self,
        plan: OptimizedPlan,
        _common: &HashMap<String, ExecutionPlan>,
    ) -> Result<OptimizedPlan, OxirsError> {
        // Would traverse plan and replace common parts
        Ok(plan)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_ai_optimizer_creation() {
        let stats = Arc::new(IndexStats::new());
        let optimizer = AIQueryOptimizer::new(stats);

        assert!(optimizer.hardware_info.cpu_cores > 0);
    }

    #[test]
    fn test_cost_model() {
        let stats = Arc::new(IndexStats::new());
        let model = CostModel::new(stats);

        let history = model.execution_history.read().unwrap();
        assert_eq!(history.patterns.len(), 0);
    }

    #[test]
    fn test_hardware_detection() {
        let hw = HardwareInfo::detect();

        assert!(hw.cpu_cores > 0);
        assert!(hw.memory_bytes > 0);
        assert_eq!(hw.cpu_features.cache_line_size, 64);
    }
}