uqa_storage/

index_manager.rs

//
// Unified Query Algebra
//
// Copyright (c) 2023-2026 Cognica, Inc.
//

//! Index manager: registry that creates / drops / looks up indexes.
//!
//! Owns the in-memory map of
//! `Box<dyn Index>` and resolves `find_covering_index` lookups for
//! the planner. The registry stays in memory and delegates persistence to the
//! catalog when wired by the engine.

#![allow(clippy::needless_pass_by_value, clippy::map_unwrap_or)]

use parking_lot::Mutex;
use std::collections::BTreeMap;

use uqa_core::{Predicate, Value};

use crate::btree_index::BTreeIndex;
use crate::index_abc::Index;
use crate::index_types::{IndexDef, IndexType};
use crate::StorageBackendError;

/// Thin [`Index`] adapter over an in-memory [`BTreeIndex`]. Each
/// adapter owns its [`IndexDef`] so the registry can route lookups
/// by table/column.
pub struct BTreeIndexHandle {
    def: IndexDef,
    inner: BTreeIndex,
}

impl BTreeIndexHandle {
    pub fn new(def: IndexDef) -> Result<Self, StorageBackendError> {
        if def.columns.len() != 1 {
            return Err(StorageBackendError::Other(format!(
                "B-tree index `{}` requires exactly one column; got {}",
                def.name,
                def.columns.len()
            )));
        }
        let field = def.columns[0].clone();
        Ok(Self {
            def,
            inner: BTreeIndex::new(field),
        })
    }

    pub fn insert(&mut self, doc_id: u64, value: Value) {
        self.inner.insert(doc_id, value);
    }

    pub fn remove(&mut self, doc_id: u64, value: &Value) {
        self.inner.remove(doc_id, value);
    }

    pub fn clear(&mut self) {
        self.inner.clear();
    }

    pub fn inner(&self) -> &BTreeIndex {
        &self.inner
    }

    pub fn inner_mut(&mut self) -> &mut BTreeIndex {
        &mut self.inner
    }
}

impl Index for BTreeIndexHandle {
    fn index_def(&self) -> &IndexDef {
        &self.def
    }
    fn scan(&self, predicate: &Predicate) -> uqa_core::PostingList {
        self.inner.scan(predicate)
    }
    fn estimate_cardinality(&self, predicate: &Predicate) -> usize {
        self.inner.scan(predicate).len()
    }
    fn scan_cost(&self, predicate: &Predicate) -> f64 {
        // Cost proxy: equality predicates are cheap (one bucket lookup);
        // ranges scale with the predicted matching cardinality. The
        // planner reads relative numbers so the absolute scale is
        // unimportant.
        let card = self.estimate_cardinality(predicate) as f64;
        match predicate {
            Predicate::Equals(_) => 1.0 + card * 0.1,
            _ => card.max(1.0),
        }
    }
    fn build(&mut self) -> Result<(), StorageBackendError> {
        Ok(())
    }
    fn drop_index(&mut self) -> Result<(), StorageBackendError> {
        self.inner.clear();
        Ok(())
    }
}

/// In-memory physical index registry, independent of catalog persistence.
#[derive(Default)]
pub struct IndexManager {
    indexes: Mutex<BTreeMap<String, Box<dyn Index>>>,
}

impl IndexManager {
    pub fn new() -> Self {
        Self {
            indexes: Mutex::new(BTreeMap::new()),
        }
    }

    /// Build a physical index and register the definition under
    /// `index_def.name`. Returns an error if an index with the same
    /// name is already registered.
    pub fn create_index(&self, index_def: IndexDef) -> Result<(), StorageBackendError> {
        if index_def.index_type != IndexType::BTree {
            return Err(StorageBackendError::Other(format!(
                "IndexManager has no physical `{}` implementation for index `{}`; engine-specific index backends must be registered through their owning engine",
                index_def.index_type.as_str(),
                index_def.name
            )));
        }
        let mut index: Box<dyn Index> = Box::new(BTreeIndexHandle::new(index_def.clone())?);
        let mut guard = self.indexes.lock();
        if guard.contains_key(&index_def.name) {
            return Err(StorageBackendError::Other(format!(
                "index `{}` is already registered",
                index_def.name
            )));
        }
        index.build()?;
        guard.insert(index_def.name.clone(), index);
        Ok(())
    }

    pub fn drop_index(&self, name: &str) -> Result<bool, StorageBackendError> {
        let mut guard = self.indexes.lock();
        if let Some(mut idx) = guard.remove(name) {
            idx.drop_index()?;
            Ok(true)
        } else {
            Ok(false)
        }
    }

    pub fn drop_indexes_for_table(&self, table_name: &str) -> Result<(), StorageBackendError> {
        let mut guard = self.indexes.lock();
        let names: Vec<String> = guard
            .iter()
            .filter(|(_, idx)| idx.index_def().table_name == table_name)
            .map(|(n, _)| n.clone())
            .collect();
        for name in names {
            if let Some(mut idx) = guard.remove(&name) {
                idx.drop_index()?;
            }
        }
        Ok(())
    }

    pub fn find_covering_index_name(
        &self,
        table_name: &str,
        column: &str,
        predicate: &Predicate,
    ) -> Option<String> {
        self.find_covering_index_with_cost(table_name, column, predicate)
            .map(|(name, _)| name)
    }

    /// Like [`Self::find_covering_index_name`] but returns the chosen
    /// index's name together with its `scan_cost(predicate)` so the caller can
    /// require `scan_cost < full_scan_cost` before committing to the rewrite.
    pub fn find_covering_index_with_cost(
        &self,
        table_name: &str,
        column: &str,
        predicate: &Predicate,
    ) -> Option<(String, f64)> {
        let guard = self.indexes.lock();
        let mut best: Option<(String, f64)> = None;
        for (name, idx) in guard.iter() {
            let def = idx.index_def();
            if def.table_name != table_name {
                continue;
            }
            if def.columns.first().map(String::as_str) != Some(column) {
                continue;
            }
            let cost = idx.scan_cost(predicate);
            if best.as_ref().map(|(_, c)| cost < *c).unwrap_or(true) {
                best = Some((name.clone(), cost));
            }
        }
        best
    }

    pub fn has_index(&self, name: &str) -> bool {
        self.indexes.lock().contains_key(name)
    }

    pub fn index_count(&self) -> usize {
        self.indexes.lock().len()
    }
}

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

    fn fresh() -> IndexManager {
        IndexManager::new()
    }

    #[test]
    fn fresh_manager_has_no_indexes() {
        let mgr = fresh();
        assert_eq!(mgr.index_count(), 0);
        assert!(!mgr.has_index("missing"));
    }

    #[test]
    fn drop_unknown_index_is_noop() {
        let mgr = fresh();
        assert!(!mgr.drop_index("missing").unwrap());
    }

    #[test]
    fn create_then_drop_btree_index_reflects_in_count() {
        let mgr = fresh();
        let def = IndexDef::new(
            "users_age_idx",
            IndexType::BTree,
            "users",
            vec!["age".into()],
        );
        mgr.create_index(def).unwrap();
        assert_eq!(mgr.index_count(), 1);
        assert!(mgr.has_index("users_age_idx"));
        assert!(mgr.drop_index("users_age_idx").unwrap());
        assert_eq!(mgr.index_count(), 0);
    }

    #[test]
    fn find_covering_index_picks_matching_btree() {
        let mgr = fresh();
        mgr.create_index(IndexDef::new(
            "users_age_idx",
            IndexType::BTree,
            "users",
            vec!["age".into()],
        ))
        .unwrap();
        let pred = Predicate::Equals(Value::Int(42));
        let pick = mgr.find_covering_index_name("users", "age", &pred);
        assert_eq!(pick, Some("users_age_idx".into()));
    }

    #[test]
    fn invalid_or_unimplemented_physical_indexes_are_explicit_errors() {
        let mgr = fresh();
        for columns in [Vec::new(), vec!["a".into(), "b".into()]] {
            let error = mgr
                .create_index(IndexDef::new(
                    "bad_btree",
                    IndexType::BTree,
                    "users",
                    columns,
                ))
                .unwrap_err();
            assert!(error.to_string().contains("requires exactly one column"));
        }

        let error = mgr
            .create_index(IndexDef::new(
                "not_a_btree",
                IndexType::Gin,
                "users",
                vec!["body".into()],
            ))
            .unwrap_err();
        assert!(error
            .to_string()
            .contains("no physical `gin` implementation"));
        assert_eq!(mgr.index_count(), 0);
    }
}