spiral-sdk 0.1.11

//! Fixed search results and bounded pages over their exact source.

use std::collections::HashMap;
use std::ops::Range;
use std::sync::Arc;

use arrow_array::{Array as _, RecordBatch, UInt32Array, UInt64Array};
use arrow_schema::{DataType, Schema, SchemaRef};
use arrow_select::concat::concat_batches;
use arrow_select::take::take_record_batch;
use futures::{Stream, TryStreamExt as _};
use uuid::Uuid;

use super::VectorIndex;
use crate::api::tables::StorageBinding as IssuedBinding;
use crate::api::vector_indexes::{
    DistanceMetric, ReadVersionResponse, StorageAccessRequest, StorageMode,
};
use crate::scandal::api::{VectorIndexArtifact, VectorIndexOrigin, VectorMetric, VectorSearchInfo};
use crate::scandal::vector::{DISTANCE_COLUMN, INFO_KEY, RANK_COLUMN, ROW_ID_COLUMN};
use crate::scandal::{Evaluation, StorageBinding};
use crate::{Error, Expression, ScanOptions, Table};

const MAX_PAGE_BYTES: usize = 64 << 20;

/// Recall controls and the eligible population for one fixed search.
#[derive(Clone, Debug)]
pub struct VectorSearchOptions {
    /// The maximum number of neighbors. The default is 20.
    pub limit: u32,
    /// The number of partitions to probe, or the index default.
    pub probes: Option<u32>,
    /// The candidate limit for exact refinement, or the index default.
    pub rerank: Option<u32>,
    /// The source predicate, applied before final neighbor selection.
    pub filter: Option<Expression>,
    /// Stable source RowIDs to exclude, for example the query sample.
    pub exclude_row_ids: Vec<u64>,
}

impl Default for VectorSearchOptions {
    fn default() -> Self {
        Self {
            limit: 20,
            probes: None,
            rerank: None,
            filter: None,
            exclude_row_ids: Vec::new(),
        }
    }
}

/// Bounded neighbors from one query against exact index and source versions.
/// Pages reuse these neighbors and do not repeat the index search.
#[derive(Clone)]
pub struct VectorSearch {
    source: Table,
    query: Arc<[f64]>,
    options: VectorSearchOptions,
    info: VectorSearchInfo,
    neighbors: RecordBatch,
}

impl VectorIndex {
    /// Search this index against `source` with a fixed candidate budget.
    /// A source version mismatch requires a refresh or the matching source snapshot.
    pub async fn search(
        &self,
        source: &Table,
        query: Vec<f64>,
        options: VectorSearchOptions,
    ) -> Result<VectorSearch, Error> {
        if !Arc::ptr_eq(&source.connection, &self.project.client.connection) {
            return Err(Error::invalid(
                "The index and source must use the same client",
            ));
        }
        let client = self.project.client.vector_indexes();
        let read = client
            .read_version(self.resource_id, self.version_id)
            .await?;
        let artifact = self.artifact(read, source)?;
        let binding = client
            .storage_binding(
                self.resource_id,
                StorageAccessRequest {
                    version_id: self.version_id,
                    mode: StorageMode::ReadOnly,
                },
            )
            .await?;
        if binding.mode != StorageMode::ReadOnly {
            return Err(protocol(
                "Vector Index storage grant has the wrong access mode",
            ));
        }
        let issued: IssuedBinding = serde_json::to_value(binding)
            .and_then(serde_json::from_value)
            .map_err(|error| protocol(format!("Invalid Vector Index storage binding: {error}")))?;
        self.project
            .client
            .connection
            .tables()?
            .search_vectors(
                source.clone(),
                artifact,
                StorageBinding::for_resource(issued)?,
                query,
                options,
            )
            .await
    }

    fn artifact(
        &self,
        read: ReadVersionResponse,
        source: &Table,
    ) -> Result<VectorIndexArtifact, Error> {
        let definition = read.definition;
        let version = read.version;
        if version.vector_index_id != self.resource_id || version.version_id != self.version_id {
            return Err(protocol("Vector Index response changed its exact identity"));
        }
        let source_id = definition
            .source_resource_id
            .parse::<Uuid>()
            .map_err(|_| protocol("Vector Index source is not a resource ID"))?;
        let (Some(source_version), Some(physical)) = (version.source_version_id, version.physical)
        else {
            return Err(Error::invalid("Refresh the Vector Index before search"));
        };
        let source_version = source_version
            .parse::<Uuid>()
            .map_err(|_| protocol("Vector Index source version is not a UUID"))?;
        if source_id != source.identity().resource_id
            || definition.source_table_id != source.identity().resource_id
            || source_version != source.identity().version_id
        {
            return Err(Error::invalid(
                "The source does not match the exact Vector Index version",
            ));
        }
        let metric = match definition.metric {
            DistanceMetric::Cosine => VectorMetric::Cosine,
            DistanceMetric::Euclidean => VectorMetric::Euclidean,
            DistanceMetric::Dot => VectorMetric::Dot,
            DistanceMetric::Unknown(_) => {
                return Err(Error::invalid(
                    "The Vector Index uses an unsupported metric",
                ));
            }
        };
        Ok(VectorIndexArtifact {
            origin: VectorIndexOrigin {
                vector_index_id: self.resource_id,
                version_id: self.version_id,
                source_resource_id: source_id,
                source_version_id: source_version,
                source_field_id: definition.source_field_id,
            },
            dimensions: definition.dimensions,
            metric,
            codec: physical.codec,
            codec_version: physical.codec_version,
            manifest_path: physical.manifest_path,
            indexed_rows: physical.indexed_rows,
            partitions: physical.partitions,
            metadata: physical.metadata.0,
        })
    }
}

impl VectorSearch {
    pub(crate) async fn read(
        source: Table,
        query: Vec<f64>,
        options: VectorSearchOptions,
        artifact: &VectorIndexArtifact,
        evaluation: Evaluation,
    ) -> Result<Self, Error> {
        let schema = evaluation.schema;
        let expected = [
            (ROW_ID_COLUMN, DataType::UInt64),
            (DISTANCE_COLUMN, DataType::Float64),
            (RANK_COLUMN, DataType::UInt32),
        ];
        if schema.fields().len() != expected.len()
            || schema
                .fields()
                .iter()
                .zip(&expected)
                .any(|(field, (name, dtype))| {
                    field.name() != name || field.data_type() != dtype || field.is_nullable()
                })
        {
            return Err(protocol(
                "Vector search returned an invalid neighbor schema",
            ));
        }
        let info: VectorSearchInfo = serde_json::from_str(
            schema
                .metadata()
                .get(INFO_KEY)
                .ok_or_else(|| protocol("Vector search omitted its provenance"))?,
        )
        .map_err(|error| protocol(format!("Invalid vector search provenance: {error}")))?;
        if info.origin != artifact.origin
            || info.limit != options.limit
            || info.metric != artifact.metric
        {
            return Err(protocol("Vector search changed its exact query provenance"));
        }
        let neighbors = collect(schema, evaluation.batches, options.limit as usize).await?;
        Ok(Self {
            source,
            query: query.into(),
            options,
            info,
            neighbors,
        })
    }

    /// Return the exact source handle for these results.
    pub fn source(&self) -> &Table {
        &self.source
    }

    /// Return the fixed query vector.
    pub fn query(&self) -> &[f64] {
        &self.query
    }

    /// Return the fixed predicate, exclusions, and recall controls.
    pub fn options(&self) -> &VectorSearchOptions {
        &self.options
    }

    /// Return exact provenance and the effective candidate budget.
    pub fn info(&self) -> &VectorSearchInfo {
        &self.info
    }

    /// Return the fixed RowIDs, float64 distances, and one-based ranks.
    pub fn neighbors(&self) -> &RecordBatch {
        &self.neighbors
    }

    /// Return the number of neighbors within this fixed budget.
    pub fn len(&self) -> usize {
        self.neighbors.num_rows()
    }

    /// Return whether this fixed budget produced no neighbors.
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Read requested columns only for `rows`, in neighbor order.
    /// Each page accepts at most 64 MiB of Arrow buffers.
    pub async fn page(
        &self,
        rows: Range<usize>,
        columns: Vec<String>,
    ) -> Result<RecordBatch, Error> {
        if rows.start > rows.end || rows.end > self.len() {
            return Err(Error::invalid(
                "The page range exceeds the fixed vector results",
            ));
        }
        let page = self.neighbors.slice(rows.start, rows.end - rows.start);
        if columns.is_empty() {
            return Ok(page);
        }
        if columns
            .iter()
            .any(|name| page.schema().column_with_name(name).is_some())
        {
            return Err(Error::invalid(
                "Page columns use a reserved vector result name",
            ));
        }
        let ids = row_ids(&page)?;
        let scan = self
            .source
            .scan(ScanOptions {
                columns: Some(columns),
                row_id_column: Some(ROW_ID_COLUMN.into()),
                row_ids: Some(row_ranges(ids.values())?),
                limit: Some(page.num_rows() as u64),
                batch_rows: Some(256),
                ..Default::default()
            })
            .await?;
        let metadata = collect(scan.schema(), scan.into_stream(), page.num_rows()).await?;
        let positions = row_ids(&metadata)?
            .values()
            .iter()
            .enumerate()
            .map(|(position, id)| (*id, position as u32))
            .collect::<HashMap<_, _>>();
        let indices = ids
            .values()
            .iter()
            .map(|id| {
                positions
                    .get(id)
                    .copied()
                    .ok_or_else(|| protocol("The source page omitted a neighbor RowID"))
            })
            .collect::<Result<Vec<_>, _>>()?;
        let metadata =
            take_record_batch(&metadata, &UInt32Array::from(indices)).map_err(Error::arrow)?;
        let mut fields = page.schema().fields().iter().cloned().collect::<Vec<_>>();
        let mut arrays = page.columns().to_vec();
        for (field, array) in metadata.schema().fields().iter().zip(metadata.columns()) {
            if field.name() != ROW_ID_COLUMN {
                fields.push(Arc::clone(field));
                arrays.push(Arc::clone(array));
            }
        }
        RecordBatch::try_new(
            Arc::new(Schema::new_with_metadata(
                fields,
                page.schema().metadata().clone(),
            )),
            arrays,
        )
        .map_err(Error::arrow)
    }
}

fn row_ids(batch: &RecordBatch) -> Result<&UInt64Array, Error> {
    batch
        .column_by_name(ROW_ID_COLUMN)
        .and_then(|array| array.as_any().downcast_ref())
        .ok_or_else(|| protocol("Vector results omitted their stable RowIDs"))
}

fn row_ranges(ids: &[u64]) -> Result<Vec<Range<u64>>, Error> {
    let mut ids = ids.to_vec();
    ids.sort_unstable();
    let mut ranges: Vec<Range<u64>> = Vec::new();
    for id in ids {
        let end = id
            .checked_add(1)
            .ok_or_else(|| protocol("A neighbor RowID exceeds the supported range"))?;
        match ranges.last_mut() {
            Some(range) if range.end == id => range.end = end,
            _ => ranges.push(id..end),
        }
    }
    Ok(ranges)
}

async fn collect<S, E>(
    schema: SchemaRef,
    mut batches: S,
    limit: usize,
) -> Result<RecordBatch, Error>
where
    S: Stream<Item = Result<RecordBatch, E>> + Unpin,
    E: Into<Error>,
{
    let mut rows = 0usize;
    let mut bytes = 0usize;
    let mut collected = Vec::new();
    while let Some(batch) = batches.try_next().await.map_err(Into::into)? {
        rows = rows.saturating_add(batch.num_rows());
        bytes = bytes.saturating_add(batch.get_array_memory_size());
        if rows > limit {
            return Err(protocol("Vector results exceeded their row limit"));
        }
        if bytes > MAX_PAGE_BYTES {
            return Err(Error::invalid("The selected page exceeds its byte limit"));
        }
        collected.push(batch);
    }
    concat_batches(&schema, &collected).map_err(Error::arrow)
}

fn protocol(message: impl Into<String>) -> Error {
    Error::Protocol(message.into())
}

#[cfg(test)]
mod tests;