spiral-sdk 0.1.12

//! Fixed-budget vector search over exact resource versions.

use super::api::TableVersion;
pub use super::api::{
    SearchVectorsRequest, VectorIndexArtifact, VectorIndexOrigin, VectorMetric, VectorSearchInfo,
};
use super::error::ScandalError;

/// This column contains the stable source RowID.
pub const ROW_ID_COLUMN: &str = "scandal.row_id";
/// This column contains the exact float64 distance.
pub const DISTANCE_COLUMN: &str = "scandal.distance";
/// This column contains the one-based rank.
pub const RANK_COLUMN: &str = "scandal.rank";
/// The Arrow schema uses this key for a JSON `VectorSearchInfo` value.
pub const INFO_KEY: &str = "scandal.vector_search";
/// This limit caps the final rows in one request.
pub const MAX_RESULTS: u32 = 1_000;
/// This limit caps the index candidates in one request.
pub const MAX_CANDIDATES: u32 = 100_000;

const MAX_PARTITIONS: u32 = 4096;
const MAX_METADATA_BYTES: usize = 512;
const MAX_COLUMNS: usize = 64;
const MAX_COLUMN_NAME_BYTES: usize = 1024;

impl VectorIndexOrigin {
    /// Return the exact source Table version for this artifact.
    pub fn source(&self) -> TableVersion {
        TableVersion {
            resource_id: self.source_resource_id,
            version_id: self.source_version_id,
        }
    }
}

impl SearchVectorsRequest {
    /// Check query shape, recall controls, and projected field names.
    pub fn validate(&self) -> Result<(), ScandalError> {
        if self.limit == 0 || self.limit > MAX_RESULTS {
            return Err(invalid(format!(
                "Vector search limit must be between 1 and {MAX_RESULTS}"
            )));
        }
        if self.probes == Some(0)
            || self
                .rerank
                .is_some_and(|count| count < self.limit || count > MAX_CANDIDATES)
        {
            return Err(invalid(format!(
                "Vector search probes must be positive and rerank must be between limit and {MAX_CANDIDATES}"
            )));
        }
        if (self.index.partitions == 0 && self.index.indexed_rows.0 != 0)
            || self.index.partitions > MAX_PARTITIONS
        {
            return Err(invalid(format!(
                "Vector Index partitions must be between 1 and {MAX_PARTITIONS}"
            )));
        }
        if self.index.dimensions == 0
            || self.query.len() != self.index.dimensions as usize
            || self.query.iter().any(|value| !value.is_finite())
        {
            return Err(invalid(
                "Vector search requires a finite query with the index dimensions",
            ));
        }
        if !matches!(self.index.codec_version, 3 | 4) {
            return Err(invalid(
                "Vector search requires a resource refresh with exact artifact identities",
            ));
        }
        if self.index.metadata.len() > MAX_METADATA_BYTES {
            return Err(invalid("Vector Index checksum metadata exceeds its limit"));
        }
        let columns = self.columns.as_deref().unwrap_or_default();
        if columns.len() > MAX_COLUMNS
            || columns.iter().any(|name| {
                name.len() > MAX_COLUMN_NAME_BYTES
                    || [ROW_ID_COLUMN, DISTANCE_COLUMN, RANK_COLUMN].contains(&name.as_str())
            })
        {
            return Err(invalid(
                "Vector search columns exceed their bounds or use a reserved result name",
            ));
        }
        if self
            .exclude_row_ids
            .as_ref()
            .is_some_and(|ids| ids.len() > MAX_RESULTS as usize)
        {
            return Err(invalid(format!(
                "Vector search permits at most {MAX_RESULTS} excluded RowIDs"
            )));
        }
        Ok(())
    }
}

fn invalid(message: impl Into<String>) -> ScandalError {
    ScandalError::InvalidArgument {
        message: message.into(),
    }
}