spiral/
vector_index.rs

1//! Vector Indexes resolved on a branch, and the Tasks that refresh them.
2
3use std::collections::BTreeMap;
4#[cfg(feature = "tables")]
5use std::sync::Arc;
6
7#[cfg(feature = "tables")]
8use arrow_schema::DataType;
9use uuid::Uuid;
10
11use crate::api::tasks::{
12    TaskDefinition, TaskReadTarget, TaskReadTargetLatestOnBranch, VectorIndexBuildMode,
13    VectorIndexRefreshDefinition,
14};
15pub use crate::api::vector_indexes::DistanceMetric as VectorIndexMetric;
16#[cfg(feature = "tables")]
17use crate::api::vector_indexes::{CreateVectorIndexRequest, Definition};
18use crate::catalog_types::ProjectBranchId;
19use crate::{Error, Project, ResourceAddress, Run, Snapshot, TaskHandle};
20#[cfg(feature = "tables")]
21use crate::{ScanOptions, Table};
22
23#[cfg(feature = "tables")]
24mod search;
25#[cfg(feature = "tables")]
26pub use search::{VectorSearch, VectorSearchOptions};
27
28#[cfg(feature = "tables")]
29pub(crate) const VECTOR_INDEX_SERVICE: &str = "vector-index";
30pub(crate) const VECTOR_INDEX_KIND: &str = "vector-index";
31
32/// This handle selects one exact Vector Index version.
33#[derive(Clone)]
34pub struct VectorIndex {
35    pub(crate) project: Project,
36    binding: Binding,
37    resource_id: Uuid,
38    version_id: Uuid,
39}
40
41/// The binding identifies the origin of an exact index handle.
42#[derive(Clone, Copy)]
43enum Binding {
44    /// Creation returns an unbound handle before publication.
45    #[cfg(feature = "tables")]
46    Unbound,
47    /// An exact commit supplies this handle.
48    Commit,
49    /// A branch snapshot supplies this handle and its refresh target.
50    Branch(ProjectBranchId),
51}
52
53/// These options select the metric and build controls for a new index.
54#[derive(Clone, Debug)]
55pub struct VectorIndexOptions {
56    /// The metric determines neighbor order.
57    pub metric: VectorIndexMetric,
58    /// Engine options contain string values for the selected build controls.
59    pub engine_options: BTreeMap<String, String>,
60}
61
62impl Default for VectorIndexOptions {
63    fn default() -> Self {
64        Self {
65            metric: VectorIndexMetric::Cosine,
66            engine_options: BTreeMap::new(),
67        }
68    }
69}
70
71impl Project {
72    /// Create an unnamed Vector Index over the `field` column of `source`.
73    ///
74    /// Creation resolves the column through a zero-row plan for the exact source schema.
75    /// A commit binds the index name before a branch snapshot can refresh it.
76    #[cfg(feature = "tables")]
77    pub async fn create_vector_index(
78        &self,
79        source: &Table,
80        field: &str,
81        options: VectorIndexOptions,
82    ) -> Result<VectorIndex, Error> {
83        if !Arc::ptr_eq(&source.connection, &self.client.connection) {
84            return Err(Error::invalid(
85                "The index and source must use the same client",
86            ));
87        }
88        if matches!(&options.metric, VectorIndexMetric::Unknown(_)) {
89            return Err(Error::invalid("The Vector Index metric is unsupported"));
90        }
91        let (source_field_id, dimensions) = resolve_vector_field(source, field).await?;
92        let created = self
93            .client
94            .vector_indexes()
95            .create_vector_index(CreateVectorIndexRequest {
96                project_id: self.id.to_string(),
97                definition: Definition {
98                    source_resource_id: source.identity().resource_id.to_string(),
99                    source_table_id: source.identity().resource_id,
100                    source_field_id,
101                    dimensions,
102                    metric: options.metric,
103                    engine_id: "vortex.ivf-rabitq".to_owned(),
104                    engine_options: options.engine_options,
105                },
106            })
107            .await?;
108        Ok(VectorIndex {
109            project: self.clone(),
110            binding: Binding::Unbound,
111            resource_id: created.vector_index_id,
112            version_id: created.version_id,
113        })
114    }
115
116    /// Pin the head of `branch` and resolve the Vector Index at `path`.
117    pub async fn vector_index(&self, branch: &str, path: &str) -> Result<VectorIndex, Error> {
118        self.branch(branch)
119            .await?
120            .snapshot()
121            .vector_index(path)
122            .await
123    }
124}
125
126/// Resolve the stable field identity and dimensions of `field` from a zero-row scan plan.
127#[cfg(feature = "tables")]
128async fn resolve_vector_field(source: &Table, field: &str) -> Result<(Uuid, u32), Error> {
129    let scan = source
130        .scan(ScanOptions {
131            columns: Some(vec![field.to_owned()]),
132            limit: Some(0),
133            ..Default::default()
134        })
135        .await?;
136    let fields = scan
137        .source_fields()
138        .ok_or_else(|| Error::invalid("ScanDAL does not expose stable source field identities"))?;
139    let source_field_id = fields
140        .get(field)
141        .copied()
142        .ok_or_else(|| Error::invalid("The source column has no stable field identity"))?;
143    let schema = scan.schema();
144    let dimensions = match schema
145        .field_with_name(field)
146        .map_err(Error::arrow)?
147        .data_type()
148    {
149        DataType::FixedSizeList(element, size) if element.data_type().is_floating() => {
150            u32::try_from(*size).ok().filter(|&size| size > 0)
151        }
152        _ => None,
153    }
154    .ok_or_else(|| {
155        Error::invalid(
156            "A Vector Index requires a nonempty fixed-size list of floating-point coordinates",
157        )
158    })?;
159    Ok((source_field_id, dimensions))
160}
161
162impl Snapshot {
163    /// Resolve the Vector Index at `path` in this snapshot.
164    pub async fn vector_index(&self, path: &str) -> Result<VectorIndex, Error> {
165        let address = path.parse::<ResourceAddress>()?;
166        if address.reference().is_some() {
167            return Err(Error::invalid(
168                "select the Vector Index's branch with the snapshot, not an `@` reference",
169            ));
170        }
171        let resource = self.resource_at_path(address.path()).await?;
172        if resource.descriptor.kind.0 != VECTOR_INDEX_KIND {
173            return Err(Error::invalid(format!("`{path}` is not a Vector Index")));
174        }
175        Ok(VectorIndex {
176            project: self.project.clone(),
177            binding: match self.branch_id {
178                Some(branch_id) => Binding::Branch(branch_id),
179                None => Binding::Commit,
180            },
181            resource_id: resource.descriptor.resource_id.as_uuid(),
182            version_id: resource.version_id.as_uuid(),
183        })
184    }
185}
186
187impl VectorIndex {
188    /// The Vector Index's stable resource identifier.
189    pub fn resource_id(&self) -> Uuid {
190        self.resource_id
191    }
192
193    /// This method returns the exact version of the index.
194    pub fn version_id(&self) -> Uuid {
195        self.version_id
196    }
197
198    /// This check distinguishes a new index from a snapshot handle.
199    #[cfg(feature = "tables")]
200    pub(crate) fn is_unbound(&self) -> bool {
201        matches!(self.binding, Binding::Unbound)
202    }
203
204    /// Submit a build of the index's successor from the latest source Table on this
205    /// snapshot's branch.
206    ///
207    /// Its result lists the new Vector Index version. Returns `None` when a refresh of this
208    /// index and branch is already waiting to run.
209    pub async fn refresh(
210        &self,
211        mode: VectorIndexBuildMode,
212        run: Run,
213    ) -> Result<Option<TaskHandle>, Error> {
214        let Binding::Branch(branch_id) = self.binding else {
215            return Err(Error::invalid(
216                "refreshing a Vector Index needs a branch snapshot, not an exact commit",
217            ));
218        };
219        let read = self
220            .project
221            .client
222            .vector_indexes()
223            .read_version(self.resource_id, self.version_id)
224            .await?;
225        let source_resource_id =
226            read.definition
227                .source_resource_id
228                .parse::<Uuid>()
229                .map_err(|_| {
230                    Error::Protocol("the Vector Index's source is not a resource ID".into())
231                })?;
232        let definition = TaskDefinition::VectorIndexRefresh(VectorIndexRefreshDefinition {
233            read_target: TaskReadTarget::LatestOnBranch(TaskReadTargetLatestOnBranch {
234                branch_id: branch_id.as_uuid(),
235            }),
236            vector_index_id: self.resource_id,
237            vector_index_resource_id: self.resource_id,
238            source_resource_id,
239            mode,
240            latest: None,
241        });
242        self.project.submit_task(definition, None, run).await
243    }
244}
245
246#[cfg(test)]
247#[cfg(feature = "tables")]
248mod tests;
spiral/vector_index.rs

spiral/
vector_index.rs
