spiral-sdk 0.1.11

//! Vector Indexes resolved on a branch, and the Tasks that refresh them.

use std::collections::BTreeMap;
#[cfg(feature = "tables")]
use std::sync::Arc;

#[cfg(feature = "tables")]
use arrow_schema::DataType;
use uuid::Uuid;

use crate::api::tasks::{
    TaskDefinition, TaskReadTarget, TaskReadTargetLatestOnBranch, VectorIndexBuildMode,
    VectorIndexRefreshDefinition,
};
pub use crate::api::vector_indexes::DistanceMetric as VectorIndexMetric;
#[cfg(feature = "tables")]
use crate::api::vector_indexes::{CreateVectorIndexRequest, Definition};
use crate::catalog_types::ProjectBranchId;
use crate::{Error, Project, ResourceAddress, Run, Snapshot, TaskHandle};
#[cfg(feature = "tables")]
use crate::{ScanOptions, Table};

#[cfg(feature = "tables")]
mod search;
#[cfg(feature = "tables")]
pub use search::{VectorSearch, VectorSearchOptions};

#[cfg(feature = "tables")]
pub(crate) const VECTOR_INDEX_SERVICE: &str = "vector-index";
pub(crate) const VECTOR_INDEX_KIND: &str = "vector-index";

/// This handle selects one exact Vector Index version.
#[derive(Clone)]
pub struct VectorIndex {
    pub(crate) project: Project,
    binding: Binding,
    resource_id: Uuid,
    version_id: Uuid,
}

/// The binding identifies the origin of an exact index handle.
#[derive(Clone, Copy)]
enum Binding {
    /// Creation returns an unbound handle before publication.
    #[cfg(feature = "tables")]
    Unbound,
    /// An exact commit supplies this handle.
    Commit,
    /// A branch snapshot supplies this handle and its refresh target.
    Branch(ProjectBranchId),
}

/// These options select the metric and build controls for a new index.
#[derive(Clone, Debug)]
pub struct VectorIndexOptions {
    /// The metric determines neighbor order.
    pub metric: VectorIndexMetric,
    /// Engine options contain string values for the selected build controls.
    pub engine_options: BTreeMap<String, String>,
}

impl Default for VectorIndexOptions {
    fn default() -> Self {
        Self {
            metric: VectorIndexMetric::Cosine,
            engine_options: BTreeMap::new(),
        }
    }
}

impl Project {
    /// Create an unnamed Vector Index over the `field` column of `source`.
    ///
    /// Creation resolves the column through a zero-row plan for the exact source schema.
    /// A commit binds the index name before a branch snapshot can refresh it.
    #[cfg(feature = "tables")]
    pub async fn create_vector_index(
        &self,
        source: &Table,
        field: &str,
        options: VectorIndexOptions,
    ) -> Result<VectorIndex, Error> {
        if !Arc::ptr_eq(&source.connection, &self.client.connection) {
            return Err(Error::invalid(
                "The index and source must use the same client",
            ));
        }
        if matches!(&options.metric, VectorIndexMetric::Unknown(_)) {
            return Err(Error::invalid("The Vector Index metric is unsupported"));
        }
        let (source_field_id, dimensions) = resolve_vector_field(source, field).await?;
        let created = self
            .client
            .vector_indexes()
            .create_vector_index(CreateVectorIndexRequest {
                project_id: self.id.to_string(),
                definition: Definition {
                    source_resource_id: source.identity().resource_id.to_string(),
                    source_table_id: source.identity().resource_id,
                    source_field_id,
                    dimensions,
                    metric: options.metric,
                    engine_id: "vortex.ivf-rabitq".to_owned(),
                    engine_options: options.engine_options,
                },
            })
            .await?;
        Ok(VectorIndex {
            project: self.clone(),
            binding: Binding::Unbound,
            resource_id: created.vector_index_id,
            version_id: created.version_id,
        })
    }

    /// Pin the head of `branch` and resolve the Vector Index at `path`.
    pub async fn vector_index(&self, branch: &str, path: &str) -> Result<VectorIndex, Error> {
        self.branch(branch)
            .await?
            .snapshot()
            .vector_index(path)
            .await
    }
}

/// Resolve the stable field identity and dimensions of `field` from a zero-row scan plan.
#[cfg(feature = "tables")]
async fn resolve_vector_field(source: &Table, field: &str) -> Result<(Uuid, u32), Error> {
    let scan = source
        .scan(ScanOptions {
            columns: Some(vec![field.to_owned()]),
            limit: Some(0),
            ..Default::default()
        })
        .await?;
    let fields = scan
        .source_fields()
        .ok_or_else(|| Error::invalid("ScanDAL does not expose stable source field identities"))?;
    let source_field_id = fields
        .get(field)
        .copied()
        .ok_or_else(|| Error::invalid("The source column has no stable field identity"))?;
    let schema = scan.schema();
    let dimensions = match schema
        .field_with_name(field)
        .map_err(Error::arrow)?
        .data_type()
    {
        DataType::FixedSizeList(element, size) if element.data_type().is_floating() => {
            u32::try_from(*size).ok().filter(|&size| size > 0)
        }
        _ => None,
    }
    .ok_or_else(|| {
        Error::invalid(
            "A Vector Index requires a nonempty fixed-size list of floating-point coordinates",
        )
    })?;
    Ok((source_field_id, dimensions))
}

impl Snapshot {
    /// Resolve the Vector Index at `path` in this snapshot.
    pub async fn vector_index(&self, path: &str) -> Result<VectorIndex, Error> {
        let address = path.parse::<ResourceAddress>()?;
        if address.reference().is_some() {
            return Err(Error::invalid(
                "select the Vector Index's branch with the snapshot, not an `@` reference",
            ));
        }
        let resource = self.resource_at_path(address.path()).await?;
        if resource.descriptor.kind.0 != VECTOR_INDEX_KIND {
            return Err(Error::invalid(format!("`{path}` is not a Vector Index")));
        }
        Ok(VectorIndex {
            project: self.project.clone(),
            binding: match self.branch_id {
                Some(branch_id) => Binding::Branch(branch_id),
                None => Binding::Commit,
            },
            resource_id: resource.descriptor.resource_id.as_uuid(),
            version_id: resource.version_id.as_uuid(),
        })
    }
}

impl VectorIndex {
    /// The Vector Index's stable resource identifier.
    pub fn resource_id(&self) -> Uuid {
        self.resource_id
    }

    /// This method returns the exact version of the index.
    pub fn version_id(&self) -> Uuid {
        self.version_id
    }

    /// This check distinguishes a new index from a snapshot handle.
    #[cfg(feature = "tables")]
    pub(crate) fn is_unbound(&self) -> bool {
        matches!(self.binding, Binding::Unbound)
    }

    /// Submit a build of the index's successor from the latest source Table on this
    /// snapshot's branch.
    ///
    /// Its result lists the new Vector Index version. Returns `None` when a refresh of this
    /// index and branch is already waiting to run.
    pub async fn refresh(
        &self,
        mode: VectorIndexBuildMode,
        run: Run,
    ) -> Result<Option<TaskHandle>, Error> {
        let Binding::Branch(branch_id) = self.binding else {
            return Err(Error::invalid(
                "refreshing a Vector Index needs a branch snapshot, not an exact commit",
            ));
        };
        let read = self
            .project
            .client
            .vector_indexes()
            .read_version(self.resource_id, self.version_id)
            .await?;
        let source_resource_id =
            read.definition
                .source_resource_id
                .parse::<Uuid>()
                .map_err(|_| {
                    Error::Protocol("the Vector Index's source is not a resource ID".into())
                })?;
        let definition = TaskDefinition::VectorIndexRefresh(VectorIndexRefreshDefinition {
            read_target: TaskReadTarget::LatestOnBranch(TaskReadTargetLatestOnBranch {
                branch_id: branch_id.as_uuid(),
            }),
            vector_index_id: self.resource_id,
            vector_index_resource_id: self.resource_id,
            source_resource_id,
            mode,
            latest: None,
        });
        self.project.submit_task(definition, None, run).await
    }
}

#[cfg(test)]
#[cfg(feature = "tables")]
mod tests;