weavatrix-search-vector 0.2.0

use crate::{MappedVectorIndex, SearchError, SearchHit, VectorIndex};
use std::collections::{BTreeMap, BTreeSet};

/// Deterministically ordered metadata value attached to a vector key.
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[non_exhaustive]
pub enum MetadataValue {
    Bool(bool),
    I64(i64),
    U64(u64),
    Text(String),
    Bytes(Vec<u8>),
}

/// Metadata fields for one vector.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct Metadata {
    fields: BTreeMap<String, MetadataValue>,
}

impl Metadata {
    #[must_use]
    pub const fn new() -> Self {
        Self {
            fields: BTreeMap::new(),
        }
    }

    pub fn insert(
        &mut self,
        field: impl Into<String>,
        value: MetadataValue,
    ) -> Option<MetadataValue> {
        self.fields.insert(field.into(), value)
    }

    #[must_use]
    pub fn get(&self, field: &str) -> Option<&MetadataValue> {
        self.fields.get(field)
    }

    pub fn remove(&mut self, field: &str) -> Option<MetadataValue> {
        self.fields.remove(field)
    }

    #[must_use]
    pub fn len(&self) -> usize {
        self.fields.len()
    }

    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.fields.is_empty()
    }

    #[must_use]
    pub fn iter(&self) -> impl ExactSizeIterator<Item = (&str, &MetadataValue)> {
        self.fields
            .iter()
            .map(|(field, value)| (field.as_str(), value))
    }
}

/// Composable deterministic predicate over vector metadata.
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub enum MetadataFilter {
    All,
    Exists(String),
    Equal {
        field: String,
        value: MetadataValue,
    },
    I64Range {
        field: String,
        minimum: Option<i64>,
        maximum: Option<i64>,
    },
    U64Range {
        field: String,
        minimum: Option<u64>,
        maximum: Option<u64>,
    },
    TextPrefix {
        field: String,
        prefix: String,
    },
    And(Vec<Self>),
    Or(Vec<Self>),
    Not(Box<Self>),
}

impl MetadataFilter {
    #[must_use]
    pub fn matches(&self, metadata: &Metadata) -> bool {
        match self {
            Self::All => true,
            Self::Exists(field) => metadata.get(field).is_some(),
            Self::Equal { field, value } => metadata.get(field) == Some(value),
            Self::I64Range {
                field,
                minimum,
                maximum,
            } => matches!(
                metadata.get(field),
                Some(MetadataValue::I64(value))
                    if minimum.is_none_or(|minimum| *value >= minimum)
                        && maximum.is_none_or(|maximum| *value <= maximum)
            ),
            Self::U64Range {
                field,
                minimum,
                maximum,
            } => matches!(
                metadata.get(field),
                Some(MetadataValue::U64(value))
                    if minimum.is_none_or(|minimum| *value >= minimum)
                        && maximum.is_none_or(|maximum| *value <= maximum)
            ),
            Self::TextPrefix { field, prefix } => matches!(
                metadata.get(field),
                Some(MetadataValue::Text(value)) if value.starts_with(prefix)
            ),
            Self::And(filters) => filters.iter().all(|filter| filter.matches(metadata)),
            Self::Or(filters) => filters.iter().any(|filter| filter.matches(metadata)),
            Self::Not(filter) => !filter.matches(metadata),
        }
    }
}

/// Key-addressed metadata used to filter owned or memory-mapped indexes.
#[derive(Debug, Clone, Default)]
pub struct MetadataIndex {
    records: BTreeMap<u64, Metadata>,
}

impl MetadataIndex {
    #[must_use]
    pub const fn new() -> Self {
        Self {
            records: BTreeMap::new(),
        }
    }

    pub fn insert(&mut self, key: u64, metadata: Metadata) -> Option<Metadata> {
        self.records.insert(key, metadata)
    }

    pub fn remove(&mut self, key: u64) -> Option<Metadata> {
        self.records.remove(&key)
    }

    #[must_use]
    pub fn get(&self, key: u64) -> Option<&Metadata> {
        self.records.get(&key)
    }

    #[must_use]
    pub fn len(&self) -> usize {
        self.records.len()
    }

    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.records.is_empty()
    }

    #[must_use]
    pub fn matches(&self, key: u64, filter: &MetadataFilter) -> bool {
        self.records
            .get(&key)
            .is_some_and(|metadata| filter.matches(metadata))
    }

    #[must_use]
    pub fn matching_keys(&self, filter: &MetadataFilter) -> BTreeSet<u64> {
        self.records
            .iter()
            .filter_map(|(key, metadata)| filter.matches(metadata).then_some(*key))
            .collect()
    }

    /// Searches an owned index and guarantees enough filtered hits through the
    /// index's exact fallback.
    ///
    /// # Errors
    ///
    /// Returns a typed query or allocation error.
    pub fn search(
        &self,
        index: &VectorIndex,
        query: &[f32],
        count: usize,
        filter: &MetadataFilter,
    ) -> Result<Vec<SearchHit>, SearchError> {
        index.search_filtered(query, count, |key| self.matches(key, filter))
    }

    /// Searches a memory-mapped index with the same filter semantics.
    ///
    /// # Errors
    ///
    /// Returns a typed query or allocation error.
    pub fn search_mapped(
        &self,
        index: &MappedVectorIndex,
        query: &[f32],
        count: usize,
        filter: &MetadataFilter,
    ) -> Result<Vec<SearchHit>, SearchError> {
        index.search_filtered(query, count, |key| self.matches(key, filter))
    }
}