oxirs_tdb/

query_join_optimizer.rs

//! Join Order Optimization for SPARQL Query Planning
//!
//! This module implements advanced join order optimization for SPARQL queries,
//! determining the optimal execution order for triple patterns to minimize
//! intermediate result sizes and overall query execution time.
//!
//! # Algorithms
//!
//! - **Greedy Algorithm**: Fast heuristic-based join ordering
//! - **Dynamic Programming**: Optimal join order with O(n * 2^n) complexity
//! - **Genetic Algorithm**: For large queries with many joins
//! - **Left-Deep Trees**: Generate left-deep join trees for pipeline execution
//!
//! # Cost Model
//!
//! Join cost = size(left) × size(right) × selectivity
//!
//! Selectivity estimation:
//! - Uses statistics from index (cardinality, distinct values)
//! - Histogram-based estimation for range predicates
//! - Join selectivity based on correlation statistics
//!
//! # Example
//!
//! ```rust,no_run
//! use oxirs_tdb::query_join_optimizer::{JoinOptimizer, JoinOptimizerConfig, TriplePattern};
//!
//! # fn example() -> anyhow::Result<()> {
//! // Create optimizer
//! let config = JoinOptimizerConfig::default();
//! let mut optimizer = JoinOptimizer::new(config);
//!
//! // Add triple patterns
//! let patterns = vec![
//!     TriplePattern::new("?s", "type", "Person"),
//!     TriplePattern::new("?s", "name", "?name"),
//!     TriplePattern::new("?s", "age", "?age"),
//! ];
//!
//! // Optimize join order
//! let optimized = optimizer.optimize(patterns)?;
//! # Ok(())
//! # }
//! ```

use crate::error::{Result, TdbError};
use crate::statistics::StatisticsSnapshot;
use anyhow::Context;
use parking_lot::Mutex;
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet};
use std::sync::Arc;

/// Triple pattern in SPARQL query
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct TriplePattern {
    /// Subject (variable or constant)
    pub subject: String,
    /// Predicate (variable or constant)
    pub predicate: String,
    /// Object (variable or constant)
    pub object: String,
    /// Estimated cardinality
    pub estimated_cardinality: Option<u64>,
}

impl TriplePattern {
    /// Create a new triple pattern
    pub fn new(subject: &str, predicate: &str, object: &str) -> Self {
        Self {
            subject: subject.to_string(),
            predicate: predicate.to_string(),
            object: object.to_string(),
            estimated_cardinality: None,
        }
    }

    /// Check if a term is a variable (starts with '?')
    pub fn is_variable(term: &str) -> bool {
        term.starts_with('?')
    }

    /// Get all variables in the pattern
    pub fn variables(&self) -> HashSet<String> {
        let mut vars = HashSet::new();
        if Self::is_variable(&self.subject) {
            vars.insert(self.subject.clone());
        }
        if Self::is_variable(&self.predicate) {
            vars.insert(self.predicate.clone());
        }
        if Self::is_variable(&self.object) {
            vars.insert(self.object.clone());
        }
        vars
    }

    /// Get join variables with another pattern
    pub fn join_variables(&self, other: &TriplePattern) -> HashSet<String> {
        self.variables()
            .intersection(&other.variables())
            .cloned()
            .collect()
    }

    /// Count bound terms (non-variables)
    pub fn bound_count(&self) -> usize {
        let mut count = 0;
        if !Self::is_variable(&self.subject) {
            count += 1;
        }
        if !Self::is_variable(&self.predicate) {
            count += 1;
        }
        if !Self::is_variable(&self.object) {
            count += 1;
        }
        count
    }
}

/// Join node in execution plan
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JoinNode {
    /// Left input pattern or join
    pub left: Box<JoinPlan>,
    /// Right input pattern or join
    pub right: Box<JoinPlan>,
    /// Join variables
    pub join_vars: HashSet<String>,
    /// Estimated cost
    pub cost: f64,
    /// Estimated output cardinality
    pub cardinality: u64,
}

/// Join execution plan (tree structure)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum JoinPlan {
    /// Leaf: single triple pattern
    Pattern(TriplePattern),
    /// Inner node: join of two sub-plans
    Join(JoinNode),
}

impl JoinPlan {
    /// Get all variables in the plan
    pub fn variables(&self) -> HashSet<String> {
        match self {
            JoinPlan::Pattern(pattern) => pattern.variables(),
            JoinPlan::Join(node) => {
                let mut vars = node.left.variables();
                vars.extend(node.right.variables());
                vars
            }
        }
    }

    /// Get estimated cardinality
    pub fn cardinality(&self) -> u64 {
        match self {
            JoinPlan::Pattern(pattern) => pattern.estimated_cardinality.unwrap_or(1000),
            JoinPlan::Join(node) => node.cardinality,
        }
    }

    /// Get total cost
    pub fn cost(&self) -> f64 {
        match self {
            JoinPlan::Pattern(_) => 0.0,
            JoinPlan::Join(node) => node.cost,
        }
    }
}

/// Join optimization algorithm
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum JoinAlgorithm {
    /// Greedy algorithm (fast, O(n^2))
    Greedy,
    /// Dynamic programming (optimal, O(n * 2^n))
    DynamicProgramming,
    /// Genetic algorithm (for large queries)
    Genetic,
}

/// Join optimizer configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JoinOptimizerConfig {
    /// Optimization algorithm
    pub algorithm: JoinAlgorithm,
    /// Maximum patterns for dynamic programming (switch to greedy if exceeded)
    pub dp_max_patterns: usize,
    /// Default selectivity for unknown joins
    pub default_join_selectivity: f64,
    /// Enable cost-based optimization
    pub cost_based: bool,
}

impl Default for JoinOptimizerConfig {
    fn default() -> Self {
        Self {
            algorithm: JoinAlgorithm::Greedy,
            dp_max_patterns: 12,
            default_join_selectivity: 0.1,
            cost_based: true,
        }
    }
}

/// Join Optimizer
///
/// Optimizes join order for SPARQL query execution.
pub struct JoinOptimizer {
    /// Configuration
    config: JoinOptimizerConfig,
    /// Statistics for cost estimation
    stats: Option<Arc<StatisticsSnapshot>>,
    /// Optimization statistics
    opt_stats: Arc<Mutex<OptimizationStats>>,
}

/// Optimization statistics
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct OptimizationStats {
    /// Total optimizations performed
    pub total_optimizations: u64,
    /// Patterns optimized with greedy
    pub greedy_optimizations: u64,
    /// Patterns optimized with DP
    pub dp_optimizations: u64,
    /// Average optimization time (milliseconds)
    pub avg_optimization_time_ms: f64,
    /// Total optimization time
    total_time_ms: f64,
}

impl JoinOptimizer {
    /// Create a new Join Optimizer
    pub fn new(config: JoinOptimizerConfig) -> Self {
        Self {
            config,
            stats: None,
            opt_stats: Arc::new(Mutex::new(OptimizationStats::default())),
        }
    }

    /// Set statistics for cost estimation
    pub fn set_statistics(&mut self, stats: Arc<StatisticsSnapshot>) {
        self.stats = Some(stats);
    }

    /// Optimize join order for a set of triple patterns
    pub fn optimize(&mut self, mut patterns: Vec<TriplePattern>) -> Result<JoinPlan> {
        if patterns.is_empty() {
            return Err(TdbError::Other("No patterns to optimize".to_string()));
        }

        if patterns.len() == 1 {
            return Ok(JoinPlan::Pattern(
                patterns
                    .pop()
                    .expect("collection validated to be non-empty"),
            ));
        }

        // Estimate cardinalities
        self.estimate_cardinalities(&mut patterns)?;

        // Choose algorithm based on pattern count
        let algorithm = if patterns.len() > self.config.dp_max_patterns {
            JoinAlgorithm::Greedy
        } else {
            self.config.algorithm
        };

        let start = std::time::Instant::now();

        let plan = match algorithm {
            JoinAlgorithm::Greedy => self.greedy_optimize(patterns)?,
            JoinAlgorithm::DynamicProgramming => self.dp_optimize(patterns)?,
            JoinAlgorithm::Genetic => {
                // Fallback to greedy for now
                // TODO: Implement genetic algorithm
                self.greedy_optimize(patterns)?
            }
        };

        let duration = start.elapsed().as_millis() as f64;

        // Update statistics
        let mut opt_stats = self.opt_stats.lock();
        opt_stats.total_optimizations += 1;
        opt_stats.total_time_ms += duration;
        opt_stats.avg_optimization_time_ms =
            opt_stats.total_time_ms / opt_stats.total_optimizations as f64;

        match algorithm {
            JoinAlgorithm::Greedy => opt_stats.greedy_optimizations += 1,
            JoinAlgorithm::DynamicProgramming => opt_stats.dp_optimizations += 1,
            JoinAlgorithm::Genetic => opt_stats.greedy_optimizations += 1,
        }

        Ok(plan)
    }

    /// Estimate cardinalities for all patterns
    fn estimate_cardinalities(&self, patterns: &mut [TriplePattern]) -> Result<()> {
        for pattern in patterns.iter_mut() {
            let cardinality = self.estimate_pattern_cardinality(pattern);
            pattern.estimated_cardinality = Some(cardinality);
        }
        Ok(())
    }

    /// Estimate cardinality for a single pattern
    fn estimate_pattern_cardinality(&self, pattern: &TriplePattern) -> u64 {
        if let Some(ref stats) = self.stats {
            // Use statistics if available
            let bound_count = pattern.bound_count();

            match bound_count {
                3 => 1,                         // All bound: exact match
                2 => stats.total_triples / 100, // Two bound: highly selective
                1 => stats.total_triples / 10,  // One bound: moderately selective
                0 => stats.total_triples,       // No bound: full scan
                _ => stats.total_triples / 10,
            }
        } else {
            // Default estimates
            match pattern.bound_count() {
                3 => 1,
                2 => 1000,
                1 => 10000,
                0 => 100000,
                _ => 10000,
            }
        }
    }

    /// Greedy join optimization (O(n^2))
    ///
    /// Algorithm:
    /// 1. Start with pattern with smallest cardinality
    /// 2. Iteratively add pattern that minimizes join cost
    /// 3. Prefer patterns that join with existing variables
    fn greedy_optimize(&self, patterns: Vec<TriplePattern>) -> Result<JoinPlan> {
        if patterns.len() == 1 {
            return Ok(JoinPlan::Pattern(patterns[0].clone()));
        }

        // Find pattern with smallest cardinality
        let mut remaining: Vec<_> = patterns.clone();
        let start_idx = remaining
            .iter()
            .enumerate()
            .min_by_key(|(_, p)| p.estimated_cardinality.unwrap_or(u64::MAX))
            .map(|(i, _)| i)
            .expect("collection validated to be non-empty");

        let mut current_plan = JoinPlan::Pattern(remaining.remove(start_idx));
        let mut current_vars = current_plan.variables();

        // Greedily add remaining patterns
        while !remaining.is_empty() {
            // Find best pattern to add
            let (best_idx, _best_cost) = remaining
                .iter()
                .enumerate()
                .map(|(i, pattern)| {
                    let join_vars = current_vars
                        .intersection(&pattern.variables())
                        .cloned()
                        .collect();
                    let cost = self.estimate_join_cost(&current_plan, pattern, &join_vars);
                    (i, cost)
                })
                .min_by(|(_, cost1), (_, cost2)| {
                    cost1
                        .partial_cmp(cost2)
                        .unwrap_or(std::cmp::Ordering::Equal)
                })
                .expect("collection validated to be non-empty");

            let next_pattern = remaining.remove(best_idx);
            let join_vars = current_vars
                .intersection(&next_pattern.variables())
                .cloned()
                .collect();

            // Create join node
            current_plan = self.create_join(current_plan, next_pattern, join_vars);
            current_vars = current_plan.variables();
        }

        Ok(current_plan)
    }

    /// Dynamic programming join optimization (O(n * 2^n))
    ///
    /// Finds optimal join order by considering all possible join trees.
    /// Uses memoization to avoid recomputing subproblems.
    fn dp_optimize(&self, patterns: Vec<TriplePattern>) -> Result<JoinPlan> {
        let n = patterns.len();
        if n == 1 {
            return Ok(JoinPlan::Pattern(patterns[0].clone()));
        }

        // DP table: maps subset (as bitmask) to best plan
        let mut dp: HashMap<u64, JoinPlan> = HashMap::new();

        // Base case: single patterns
        for (i, pattern) in patterns.iter().enumerate() {
            let mask = 1u64 << i;
            dp.insert(mask, JoinPlan::Pattern(pattern.clone()));
        }

        // Build up larger subsets
        for size in 2..=n {
            // Generate all subsets of given size
            let subsets = self.generate_subsets(n, size);

            for subset in subsets {
                // Try all possible splits of this subset
                let mut best_plan: Option<JoinPlan> = None;
                let mut best_cost = f64::MAX;

                // Try splitting subset into left and right
                for left_mask in 1..subset {
                    if (left_mask & subset) != left_mask {
                        continue;
                    }

                    let right_mask = subset & !left_mask;
                    if right_mask == 0 {
                        continue;
                    }

                    // Get optimal plans for left and right
                    let left_plan = dp.get(&left_mask).cloned();
                    let right_plan = dp.get(&right_mask).cloned();

                    if let (Some(left), Some(right)) = (left_plan, right_plan) {
                        let join_vars = left
                            .variables()
                            .intersection(&right.variables())
                            .cloned()
                            .collect();

                        let cost = self.estimate_join_cost_plans(&left, &right, &join_vars);

                        if cost < best_cost {
                            best_cost = cost;
                            best_plan = Some(self.create_join_plans(left, right, join_vars));
                        }
                    }
                }

                if let Some(plan) = best_plan {
                    dp.insert(subset, plan);
                }
            }
        }

        // Return plan for full set
        let full_mask = (1u64 << n) - 1;
        dp.remove(&full_mask)
            .ok_or_else(|| TdbError::Other("DP optimization failed".to_string()))
    }

    /// Generate all subsets of size k from n elements
    fn generate_subsets(&self, n: usize, k: usize) -> Vec<u64> {
        let mut subsets = Vec::new();
        self.generate_subsets_recursive(n, k, 0, 0, &mut subsets);
        subsets
    }

    /// Recursive helper for subset generation
    #[allow(clippy::only_used_in_recursion)]
    fn generate_subsets_recursive(
        &self,
        n: usize,
        k: usize,
        start: usize,
        current: u64,
        subsets: &mut Vec<u64>,
    ) {
        if k == 0 {
            subsets.push(current);
            return;
        }

        for i in start..n {
            let next = current | (1u64 << i);
            self.generate_subsets_recursive(n, k - 1, i + 1, next, subsets);
        }
    }

    /// Estimate join cost between current plan and a pattern
    fn estimate_join_cost(
        &self,
        plan: &JoinPlan,
        pattern: &TriplePattern,
        join_vars: &HashSet<String>,
    ) -> f64 {
        let left_card = plan.cardinality() as f64;
        let right_card = pattern.estimated_cardinality.unwrap_or(1000) as f64;

        let selectivity = if join_vars.is_empty() {
            // Cross product
            1.0
        } else {
            // Join with shared variables
            self.config.default_join_selectivity
        };

        // Cost = |left| × |right| × selectivity
        left_card * right_card * selectivity
    }

    /// Estimate join cost between two plans
    fn estimate_join_cost_plans(
        &self,
        left: &JoinPlan,
        right: &JoinPlan,
        join_vars: &HashSet<String>,
    ) -> f64 {
        let left_card = left.cardinality() as f64;
        let right_card = right.cardinality() as f64;

        let selectivity = if join_vars.is_empty() {
            1.0
        } else {
            self.config.default_join_selectivity
        };

        left_card * right_card * selectivity + left.cost() + right.cost()
    }

    /// Create join node from plan and pattern
    fn create_join(
        &self,
        left: JoinPlan,
        right: TriplePattern,
        join_vars: HashSet<String>,
    ) -> JoinPlan {
        let right_plan = JoinPlan::Pattern(right);
        self.create_join_plans(left, right_plan, join_vars)
    }

    /// Create join node from two plans
    fn create_join_plans(
        &self,
        left: JoinPlan,
        right: JoinPlan,
        join_vars: HashSet<String>,
    ) -> JoinPlan {
        let cost = self.estimate_join_cost_plans(&left, &right, &join_vars);
        let left_card = left.cardinality();
        let right_card = right.cardinality();

        let cardinality = if join_vars.is_empty() {
            // Use saturating_mul to prevent overflow for cross products
            left_card.saturating_mul(right_card)
        } else {
            ((left_card as f64 * right_card as f64 * self.config.default_join_selectivity) as u64)
                .max(1)
        };

        JoinPlan::Join(JoinNode {
            left: Box::new(left),
            right: Box::new(right),
            join_vars,
            cost,
            cardinality,
        })
    }

    /// Get optimization statistics
    pub fn stats(&self) -> OptimizationStats {
        self.opt_stats.lock().clone()
    }
}

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

    #[test]
    fn test_triple_pattern_variables() {
        let pattern = TriplePattern::new("?s", "type", "Person");
        let vars = pattern.variables();

        assert_eq!(vars.len(), 1);
        assert!(vars.contains("?s"));
    }

    #[test]
    fn test_triple_pattern_bound_count() {
        let pattern1 = TriplePattern::new("?s", "type", "Person");
        assert_eq!(pattern1.bound_count(), 2);

        let pattern2 = TriplePattern::new("?s", "?p", "?o");
        assert_eq!(pattern2.bound_count(), 0);

        let pattern3 = TriplePattern::new("Alice", "name", "?name");
        assert_eq!(pattern3.bound_count(), 2);
    }

    #[test]
    fn test_join_variables() {
        let pattern1 = TriplePattern::new("?s", "type", "Person");
        let pattern2 = TriplePattern::new("?s", "name", "?name");

        let join_vars = pattern1.join_variables(&pattern2);
        assert_eq!(join_vars.len(), 1);
        assert!(join_vars.contains("?s"));
    }

    #[test]
    fn test_optimizer_creation() {
        let config = JoinOptimizerConfig::default();
        let optimizer = JoinOptimizer::new(config);

        assert_eq!(optimizer.config.algorithm, JoinAlgorithm::Greedy);
    }

    #[test]
    fn test_single_pattern_optimization() {
        let config = JoinOptimizerConfig::default();
        let mut optimizer = JoinOptimizer::new(config);

        let patterns = vec![TriplePattern::new("?s", "type", "Person")];

        let plan = optimizer.optimize(patterns).unwrap();

        match plan {
            JoinPlan::Pattern(pattern) => {
                assert_eq!(pattern.subject, "?s");
                assert_eq!(pattern.predicate, "type");
                assert_eq!(pattern.object, "Person");
            }
            _ => panic!("Expected pattern plan"),
        }
    }

    #[test]
    fn test_greedy_optimization() {
        let config = JoinOptimizerConfig {
            algorithm: JoinAlgorithm::Greedy,
            ..Default::default()
        };
        let mut optimizer = JoinOptimizer::new(config);

        let patterns = vec![
            TriplePattern::new("?s", "type", "Person"),
            TriplePattern::new("?s", "name", "?name"),
            TriplePattern::new("?s", "age", "?age"),
        ];

        let plan = optimizer.optimize(patterns).unwrap();

        // Should produce a join plan
        match plan {
            JoinPlan::Join(_) => {
                // Success
            }
            _ => panic!("Expected join plan"),
        }

        let stats = optimizer.stats();
        assert_eq!(stats.total_optimizations, 1);
        assert_eq!(stats.greedy_optimizations, 1);
    }

    #[test]
    fn test_dp_optimization() {
        let config = JoinOptimizerConfig {
            algorithm: JoinAlgorithm::DynamicProgramming,
            ..Default::default()
        };
        let mut optimizer = JoinOptimizer::new(config);

        let patterns = vec![
            TriplePattern::new("?s", "type", "Person"),
            TriplePattern::new("?s", "name", "?name"),
            TriplePattern::new("?s", "age", "?age"),
        ];

        let plan = optimizer.optimize(patterns).unwrap();

        match plan {
            JoinPlan::Join(_) => {
                // Success
            }
            _ => panic!("Expected join plan"),
        }

        let stats = optimizer.stats();
        assert_eq!(stats.dp_optimizations, 1);
    }

    #[test]
    fn test_cardinality_estimation() {
        let config = JoinOptimizerConfig::default();
        let optimizer = JoinOptimizer::new(config);

        // All bound: should estimate 1
        let pattern1 = TriplePattern::new("Alice", "name", "\"Alice\"");
        let card1 = optimizer.estimate_pattern_cardinality(&pattern1);
        assert_eq!(card1, 1);

        // Two bound: should estimate moderate
        let pattern2 = TriplePattern::new("?s", "name", "\"Alice\"");
        let card2 = optimizer.estimate_pattern_cardinality(&pattern2);
        assert!(card2 > 1 && card2 < 10000);

        // All unbound: should estimate large
        let pattern3 = TriplePattern::new("?s", "?p", "?o");
        let card3 = optimizer.estimate_pattern_cardinality(&pattern3);
        assert!(card3 > 10000);
    }

    #[test]
    fn test_join_plan_variables() {
        let pattern1 = TriplePattern::new("?s", "type", "Person");
        let pattern2 = TriplePattern::new("?s", "name", "?name");

        let plan1 = JoinPlan::Pattern(pattern1.clone());
        let plan2 = JoinPlan::Pattern(pattern2.clone());

        let join_vars: HashSet<String> = vec!["?s".to_string()].into_iter().collect();

        let join_plan = JoinPlan::Join(JoinNode {
            left: Box::new(plan1),
            right: Box::new(plan2),
            join_vars,
            cost: 100.0,
            cardinality: 1000,
        });

        let vars = join_plan.variables();
        assert_eq!(vars.len(), 2);
        assert!(vars.contains("?s"));
        assert!(vars.contains("?name"));
    }

    #[test]
    fn test_multiple_optimizations() {
        let config = JoinOptimizerConfig::default();
        let mut optimizer = JoinOptimizer::new(config);

        // First optimization
        let patterns1 = vec![
            TriplePattern::new("?s", "type", "Person"),
            TriplePattern::new("?s", "name", "?name"),
        ];
        optimizer.optimize(patterns1).unwrap();

        // Second optimization
        let patterns2 = vec![
            TriplePattern::new("?x", "knows", "?y"),
            TriplePattern::new("?y", "age", "?age"),
        ];
        optimizer.optimize(patterns2).unwrap();

        let stats = optimizer.stats();
        assert_eq!(stats.total_optimizations, 2);
        // On fast systems, optimization may complete so quickly that avg time is 0.0
        assert!(stats.avg_optimization_time_ms >= 0.0);
    }

    #[test]
    fn test_empty_patterns() {
        let config = JoinOptimizerConfig::default();
        let mut optimizer = JoinOptimizer::new(config);

        let patterns = vec![];
        let result = optimizer.optimize(patterns);

        assert!(result.is_err());
    }

    #[test]
    fn test_large_pattern_set() {
        let config = JoinOptimizerConfig {
            algorithm: JoinAlgorithm::DynamicProgramming,
            dp_max_patterns: 5,
            ..Default::default()
        };
        let mut optimizer = JoinOptimizer::new(config);

        // Create 10 patterns (exceeds dp_max_patterns)
        let patterns: Vec<_> = (0..10)
            .map(|i| TriplePattern::new(&format!("?s{}", i), "type", "Person"))
            .collect();

        let plan = optimizer.optimize(patterns).unwrap();

        // Should fallback to greedy
        let stats = optimizer.stats();
        assert_eq!(stats.greedy_optimizations, 1);
        assert_eq!(stats.dp_optimizations, 0);
    }
}