engine/

auto_index.rs

//! Auto-Index Selection
//!
//! Automatically selects the optimal index type based on data characteristics:
//! - Dataset size (number of vectors)
//! - Dimensionality
//! - Memory constraints
//! - Query latency requirements
//! - Accuracy requirements
//!
//! # Index Selection Matrix
//!
//! | Vectors | Dimensions | Memory | Latency | Recommended Index |
//! |---------|------------|--------|---------|-------------------|
//! | < 10K   | Any        | Any    | Any     | Flat (brute force) |
//! | 10K-100K| < 256      | Low    | Medium  | IVF + SQ8         |
//! | 10K-100K| < 256      | Normal | Low     | HNSW              |
//! | 100K-1M | Any        | Low    | Medium  | IVF-PQ            |
//! | 100K-1M | Any        | Normal | Low     | HNSW              |
//! | > 1M    | Any        | Low    | Any     | IVF-PQ            |
//! | > 1M    | Any        | Normal | Low     | HNSW + SQ8        |

use serde::{Deserialize, Serialize};

/// Memory constraint level
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
pub enum MemoryConstraint {
    /// Minimize memory usage at all costs
    Minimal,
    /// Balance memory and performance
    #[default]
    Balanced,
    /// Prioritize performance over memory
    Unlimited,
}

/// Query latency requirement
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
pub enum LatencyRequirement {
    /// Sub-millisecond queries required
    UltraLow,
    /// < 10ms queries required
    #[default]
    Low,
    /// < 100ms queries acceptable
    Medium,
    /// Latency not critical
    Relaxed,
}

/// Accuracy requirement for search results
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
pub enum AccuracyRequirement {
    /// Exact nearest neighbors (recall = 1.0)
    Exact,
    /// Very high recall (> 0.99)
    VeryHigh,
    /// High recall (> 0.95)
    #[default]
    High,
    /// Moderate recall (> 0.90)
    Moderate,
    /// Lower recall acceptable (> 0.80)
    Relaxed,
}

/// Data characteristics for index selection
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DataCharacteristics {
    /// Expected number of vectors
    pub num_vectors: usize,
    /// Vector dimensionality
    pub dimensions: usize,
    /// Whether vectors will be frequently updated
    pub frequent_updates: bool,
    /// Whether vectors will be frequently deleted
    pub frequent_deletes: bool,
}

/// Requirements for the index
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct IndexRequirements {
    /// Memory constraint level
    pub memory: MemoryConstraint,
    /// Query latency requirement
    pub latency: LatencyRequirement,
    /// Accuracy requirement
    pub accuracy: AccuracyRequirement,
}

/// Recommended index type
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum RecommendedIndex {
    /// Flat index (brute force) - exact but slow for large datasets
    Flat,
    /// HNSW - fast queries, moderate memory, good for < 1M vectors
    Hnsw,
    /// IVF - balanced performance, works well with clustering
    Ivf,
    /// IVF with Product Quantization - memory efficient for large datasets
    IvfPq,
    /// IVF with Scalar Quantization - good balance of memory and accuracy
    IvfSq,
    /// HNSW with Scalar Quantization - fast queries with reduced memory
    HnswSq,
    /// SPFresh - optimized for cold storage with streaming updates
    SpFresh,
}

impl std::fmt::Display for RecommendedIndex {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            RecommendedIndex::Flat => write!(f, "Flat (Brute Force)"),
            RecommendedIndex::Hnsw => write!(f, "HNSW"),
            RecommendedIndex::Ivf => write!(f, "IVF"),
            RecommendedIndex::IvfPq => write!(f, "IVF-PQ"),
            RecommendedIndex::IvfSq => write!(f, "IVF-SQ"),
            RecommendedIndex::HnswSq => write!(f, "HNSW-SQ"),
            RecommendedIndex::SpFresh => write!(f, "SPFresh"),
        }
    }
}

/// Index recommendation with reasoning
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct IndexRecommendation {
    /// Primary recommended index type
    pub primary: RecommendedIndex,
    /// Alternative index types that could work
    pub alternatives: Vec<RecommendedIndex>,
    /// Explanation of why this index was chosen
    pub reasoning: String,
    /// Estimated memory usage in bytes
    pub estimated_memory_bytes: usize,
    /// Estimated query latency in milliseconds
    pub estimated_latency_ms: f32,
    /// Estimated recall/accuracy
    pub estimated_recall: f32,
    /// Suggested configuration parameters
    pub suggested_params: IndexParams,
}

/// Suggested index parameters
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct IndexParams {
    // HNSW params
    pub hnsw_m: Option<usize>,
    pub hnsw_ef_construction: Option<usize>,
    pub hnsw_ef_search: Option<usize>,

    // IVF params
    pub ivf_num_clusters: Option<usize>,
    pub ivf_num_probes: Option<usize>,

    // PQ params
    pub pq_num_subquantizers: Option<usize>,
    pub pq_bits_per_subquantizer: Option<usize>,

    // SQ params
    pub sq_bits: Option<usize>,
}

/// Auto-index selector
pub struct AutoIndexSelector;

impl AutoIndexSelector {
    /// Select the optimal index based on data characteristics and requirements
    pub fn select(
        data: &DataCharacteristics,
        requirements: &IndexRequirements,
    ) -> IndexRecommendation {
        let num_vectors = data.num_vectors;
        let _dimensions = data.dimensions;

        // Decision logic based on dataset size
        let (primary, alternatives, base_reasoning) = if num_vectors < 10_000 {
            // Small dataset: flat index is fine
            if requirements.accuracy == AccuracyRequirement::Exact {
                (
                    RecommendedIndex::Flat,
                    vec![RecommendedIndex::Hnsw],
                    "Small dataset (<10K vectors) - flat index provides exact results efficiently",
                )
            } else {
                (
                    RecommendedIndex::Hnsw,
                    vec![RecommendedIndex::Flat, RecommendedIndex::Ivf],
                    "Small dataset (<10K vectors) - HNSW provides fast approximate search",
                )
            }
        } else if num_vectors < 100_000 {
            // Medium dataset
            match requirements.memory {
                MemoryConstraint::Minimal => {
                    (
                        RecommendedIndex::IvfSq,
                        vec![RecommendedIndex::IvfPq, RecommendedIndex::Ivf],
                        "Medium dataset (10K-100K) with memory constraints - IVF-SQ balances memory and accuracy",
                    )
                }
                MemoryConstraint::Balanced => {
                    match requirements.latency {
                        LatencyRequirement::UltraLow | LatencyRequirement::Low => {
                            (
                                RecommendedIndex::Hnsw,
                                vec![RecommendedIndex::HnswSq, RecommendedIndex::Ivf],
                                "Medium dataset (10K-100K) with low latency requirement - HNSW provides fast queries",
                            )
                        }
                        _ => {
                            (
                                RecommendedIndex::Ivf,
                                vec![RecommendedIndex::Hnsw, RecommendedIndex::IvfPq],
                                "Medium dataset (10K-100K) - IVF provides good balance of speed and memory",
                            )
                        }
                    }
                }
                MemoryConstraint::Unlimited => {
                    (
                        RecommendedIndex::Hnsw,
                        vec![RecommendedIndex::Ivf],
                        "Medium dataset (10K-100K) with no memory constraints - HNSW provides best query performance",
                    )
                }
            }
        } else if num_vectors < 1_000_000 {
            // Large dataset
            match requirements.memory {
                MemoryConstraint::Minimal => {
                    (
                        RecommendedIndex::IvfPq,
                        vec![RecommendedIndex::IvfSq, RecommendedIndex::SpFresh],
                        "Large dataset (100K-1M) with memory constraints - IVF-PQ provides best compression",
                    )
                }
                MemoryConstraint::Balanced => {
                    if data.frequent_updates || data.frequent_deletes {
                        (
                            RecommendedIndex::SpFresh,
                            vec![RecommendedIndex::Hnsw, RecommendedIndex::Ivf],
                            "Large dataset (100K-1M) with frequent updates - SPFresh handles streaming updates well",
                        )
                    } else {
                        match requirements.latency {
                            LatencyRequirement::UltraLow | LatencyRequirement::Low => {
                                (
                                    RecommendedIndex::HnswSq,
                                    vec![RecommendedIndex::Hnsw, RecommendedIndex::IvfPq],
                                    "Large dataset (100K-1M) with low latency requirement - HNSW-SQ balances speed and memory",
                                )
                            }
                            _ => {
                                (
                                    RecommendedIndex::IvfPq,
                                    vec![RecommendedIndex::IvfSq, RecommendedIndex::Hnsw],
                                    "Large dataset (100K-1M) - IVF-PQ provides good compression with acceptable latency",
                                )
                            }
                        }
                    }
                }
                MemoryConstraint::Unlimited => {
                    (
                        RecommendedIndex::Hnsw,
                        vec![RecommendedIndex::HnswSq],
                        "Large dataset (100K-1M) with no memory constraints - HNSW provides best query performance",
                    )
                }
            }
        } else {
            // Very large dataset (> 1M)
            match requirements.memory {
                MemoryConstraint::Minimal => {
                    (
                        RecommendedIndex::IvfPq,
                        vec![RecommendedIndex::SpFresh],
                        "Very large dataset (>1M) with memory constraints - IVF-PQ is essential for memory efficiency",
                    )
                }
                MemoryConstraint::Balanced => {
                    if data.frequent_updates || data.frequent_deletes {
                        (
                            RecommendedIndex::SpFresh,
                            vec![RecommendedIndex::IvfPq],
                            "Very large dataset (>1M) with frequent updates - SPFresh handles streaming workloads",
                        )
                    } else {
                        (
                            RecommendedIndex::IvfPq,
                            vec![RecommendedIndex::HnswSq, RecommendedIndex::SpFresh],
                            "Very large dataset (>1M) - IVF-PQ provides best memory/performance tradeoff",
                        )
                    }
                }
                MemoryConstraint::Unlimited => {
                    (
                        RecommendedIndex::HnswSq,
                        vec![RecommendedIndex::Hnsw, RecommendedIndex::IvfPq],
                        "Very large dataset (>1M) with no memory constraints - HNSW-SQ provides fast queries with reduced memory",
                    )
                }
            }
        };

        // Calculate estimates
        let (estimated_memory, estimated_latency, estimated_recall) =
            Self::estimate_performance(&primary, data, requirements);

        // Generate suggested parameters
        let suggested_params = Self::suggest_params(&primary, data, requirements);

        IndexRecommendation {
            primary,
            alternatives,
            reasoning: base_reasoning.to_string(),
            estimated_memory_bytes: estimated_memory,
            estimated_latency_ms: estimated_latency,
            estimated_recall,
            suggested_params,
        }
    }

    /// Estimate performance characteristics for an index type
    fn estimate_performance(
        index_type: &RecommendedIndex,
        data: &DataCharacteristics,
        requirements: &IndexRequirements,
    ) -> (usize, f32, f32) {
        let n = data.num_vectors;
        let d = data.dimensions;
        let base_vector_size = n * d * 4; // 4 bytes per f32

        match index_type {
            RecommendedIndex::Flat => {
                // Flat: stores all vectors, O(n) search
                let memory = base_vector_size;
                let latency = (n as f32 / 100_000.0) * 10.0; // ~10ms per 100K vectors
                (memory, latency, 1.0) // Exact recall
            }
            RecommendedIndex::Hnsw => {
                // HNSW: vectors + graph structure (~2x overhead)
                let memory = base_vector_size * 2;
                let latency = ((n as f32).log2() * 0.1).max(0.5); // O(log n)
                let recall = match requirements.accuracy {
                    AccuracyRequirement::Exact => 0.999,
                    AccuracyRequirement::VeryHigh => 0.995,
                    _ => 0.98,
                };
                (memory, latency, recall)
            }
            RecommendedIndex::Ivf => {
                // IVF: vectors + centroids
                let num_clusters = (n as f32).sqrt() as usize;
                let memory = base_vector_size + num_clusters * d * 4;
                let latency = (n as f32 / num_clusters as f32 / 10_000.0) * 5.0;
                (memory, latency.max(1.0), 0.95)
            }
            RecommendedIndex::IvfPq => {
                // IVF-PQ: heavily compressed
                let compression_ratio = 32; // Typically 32x compression
                let memory = base_vector_size / compression_ratio;
                let latency = 5.0 + (n as f32 / 1_000_000.0) * 2.0;
                (memory, latency, 0.90)
            }
            RecommendedIndex::IvfSq => {
                // IVF-SQ: 4x compression
                let memory = base_vector_size / 4;
                let latency = 3.0 + (n as f32 / 500_000.0) * 2.0;
                (memory, latency, 0.95)
            }
            RecommendedIndex::HnswSq => {
                // HNSW-SQ: HNSW with quantized vectors
                let memory = (base_vector_size / 4) * 2; // Quantized + graph
                let latency = ((n as f32).log2() * 0.15).max(0.5);
                (memory, latency, 0.96)
            }
            RecommendedIndex::SpFresh => {
                // SPFresh: optimized for cold storage
                let memory = base_vector_size / 2; // Partial in-memory
                let latency = 10.0 + (n as f32 / 100_000.0) * 1.0;
                (memory, latency, 0.92)
            }
        }
    }

    /// Suggest optimal parameters for the index type
    fn suggest_params(
        index_type: &RecommendedIndex,
        data: &DataCharacteristics,
        requirements: &IndexRequirements,
    ) -> IndexParams {
        let mut params = IndexParams::default();
        let n = data.num_vectors;
        let d = data.dimensions;

        match index_type {
            RecommendedIndex::Hnsw | RecommendedIndex::HnswSq => {
                // HNSW parameters
                params.hnsw_m = Some(match requirements.accuracy {
                    AccuracyRequirement::Exact | AccuracyRequirement::VeryHigh => 32,
                    AccuracyRequirement::High => 16,
                    _ => 12,
                });
                params.hnsw_ef_construction = Some(params.hnsw_m.unwrap_or(16) * 10);
                params.hnsw_ef_search = Some(match requirements.latency {
                    LatencyRequirement::UltraLow => 50,
                    LatencyRequirement::Low => 100,
                    LatencyRequirement::Medium => 200,
                    LatencyRequirement::Relaxed => 400,
                });

                if matches!(index_type, RecommendedIndex::HnswSq) {
                    params.sq_bits = Some(8);
                }
            }
            RecommendedIndex::Ivf | RecommendedIndex::IvfPq | RecommendedIndex::IvfSq => {
                // IVF parameters
                let num_clusters = ((n as f32).sqrt() as usize).clamp(16, 65536);
                params.ivf_num_clusters = Some(num_clusters);
                params.ivf_num_probes = Some(
                    match requirements.accuracy {
                        AccuracyRequirement::Exact | AccuracyRequirement::VeryHigh => {
                            num_clusters / 4
                        }
                        AccuracyRequirement::High => num_clusters / 8,
                        AccuracyRequirement::Moderate => num_clusters / 16,
                        AccuracyRequirement::Relaxed => num_clusters / 32,
                    }
                    .max(1),
                );

                if matches!(index_type, RecommendedIndex::IvfPq) {
                    // PQ parameters
                    let num_subquantizers = (d / 4).clamp(1, 64);
                    params.pq_num_subquantizers = Some(num_subquantizers);
                    params.pq_bits_per_subquantizer = Some(8);
                }

                if matches!(index_type, RecommendedIndex::IvfSq) {
                    params.sq_bits = Some(8);
                }
            }
            RecommendedIndex::SpFresh => {
                // SPFresh uses IVF-like clustering
                let num_clusters = ((n as f32).sqrt() as usize).clamp(16, 4096);
                params.ivf_num_clusters = Some(num_clusters);
                params.ivf_num_probes = Some((num_clusters / 10).max(1));
            }
            RecommendedIndex::Flat => {
                // No parameters needed for flat index
            }
        }

        params
    }

    /// Quick recommendation based on just vector count and dimensions
    pub fn quick_select(num_vectors: usize, dimensions: usize) -> RecommendedIndex {
        let data = DataCharacteristics {
            num_vectors,
            dimensions,
            frequent_updates: false,
            frequent_deletes: false,
        };
        let requirements = IndexRequirements::default();
        Self::select(&data, &requirements).primary
    }
}

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

    #[test]
    fn test_small_dataset() {
        let data = DataCharacteristics {
            num_vectors: 1000,
            dimensions: 128,
            frequent_updates: false,
            frequent_deletes: false,
        };
        let requirements = IndexRequirements::default();

        let rec = AutoIndexSelector::select(&data, &requirements);
        assert!(matches!(
            rec.primary,
            RecommendedIndex::Hnsw | RecommendedIndex::Flat
        ));
    }

    #[test]
    fn test_large_dataset_memory_constrained() {
        let data = DataCharacteristics {
            num_vectors: 5_000_000,
            dimensions: 768,
            frequent_updates: false,
            frequent_deletes: false,
        };
        let requirements = IndexRequirements {
            memory: MemoryConstraint::Minimal,
            latency: LatencyRequirement::Medium,
            accuracy: AccuracyRequirement::Moderate,
        };

        let rec = AutoIndexSelector::select(&data, &requirements);
        assert_eq!(rec.primary, RecommendedIndex::IvfPq);
        assert!(rec.estimated_memory_bytes < data.num_vectors * data.dimensions * 4 / 10);
    }

    #[test]
    fn test_streaming_workload() {
        let data = DataCharacteristics {
            num_vectors: 500_000,
            dimensions: 256,
            frequent_updates: true,
            frequent_deletes: true,
        };
        let requirements = IndexRequirements {
            memory: MemoryConstraint::Balanced,
            latency: LatencyRequirement::Medium,
            accuracy: AccuracyRequirement::High,
        };

        let rec = AutoIndexSelector::select(&data, &requirements);
        assert_eq!(rec.primary, RecommendedIndex::SpFresh);
    }

    #[test]
    fn test_exact_search() {
        let data = DataCharacteristics {
            num_vectors: 5000,
            dimensions: 64,
            frequent_updates: false,
            frequent_deletes: false,
        };
        let requirements = IndexRequirements {
            memory: MemoryConstraint::Unlimited,
            latency: LatencyRequirement::Relaxed,
            accuracy: AccuracyRequirement::Exact,
        };

        let rec = AutoIndexSelector::select(&data, &requirements);
        assert_eq!(rec.primary, RecommendedIndex::Flat);
        assert_eq!(rec.estimated_recall, 1.0);
    }

    #[test]
    fn test_quick_select() {
        // Small: HNSW or Flat
        let small = AutoIndexSelector::quick_select(5000, 128);
        assert!(matches!(
            small,
            RecommendedIndex::Hnsw | RecommendedIndex::Flat
        ));

        // Medium: HNSW
        let medium = AutoIndexSelector::quick_select(50_000, 256);
        assert!(matches!(
            medium,
            RecommendedIndex::Hnsw | RecommendedIndex::Ivf
        ));

        // Large: usually IVF-PQ or HNSW-SQ
        let large = AutoIndexSelector::quick_select(2_000_000, 512);
        assert!(matches!(
            large,
            RecommendedIndex::IvfPq | RecommendedIndex::HnswSq
        ));
    }

    #[test]
    fn test_suggested_params() {
        let data = DataCharacteristics {
            num_vectors: 100_000,
            dimensions: 128,
            frequent_updates: false,
            frequent_deletes: false,
        };
        let requirements = IndexRequirements::default();

        let rec = AutoIndexSelector::select(&data, &requirements);

        // Should have suggested parameters
        if matches!(rec.primary, RecommendedIndex::Hnsw) {
            assert!(rec.suggested_params.hnsw_m.is_some());
            assert!(rec.suggested_params.hnsw_ef_construction.is_some());
            assert!(rec.suggested_params.hnsw_ef_search.is_some());
        }
    }
}