spiral/vector_index/
search.rs

1//! Fixed search results and bounded pages over their exact source.
2
3use std::collections::HashMap;
4use std::ops::Range;
5use std::sync::Arc;
6
7use arrow_array::{Array as _, RecordBatch, UInt32Array, UInt64Array};
8use arrow_schema::{DataType, Schema, SchemaRef};
9use arrow_select::concat::concat_batches;
10use arrow_select::take::take_record_batch;
11use futures::{Stream, TryStreamExt as _};
12use uuid::Uuid;
13
14use super::VectorIndex;
15use crate::api::tables::StorageBinding as IssuedBinding;
16use crate::api::vector_indexes::{
17    DistanceMetric, ReadVersionResponse, StorageAccessRequest, StorageMode,
18};
19use crate::scandal::api::{VectorIndexArtifact, VectorIndexOrigin, VectorMetric, VectorSearchInfo};
20use crate::scandal::vector::{DISTANCE_COLUMN, INFO_KEY, RANK_COLUMN, ROW_ID_COLUMN};
21use crate::scandal::{Evaluation, StorageBinding};
22use crate::{Error, Expression, ScanOptions, Table};
23
24const MAX_PAGE_BYTES: usize = 64 << 20;
25
26/// Recall controls and the eligible population for one fixed search.
27#[derive(Clone, Debug)]
28pub struct VectorSearchOptions {
29    /// The maximum number of neighbors. The default is 20.
30    pub limit: u32,
31    /// The number of partitions to probe, or the index default.
32    pub probes: Option<u32>,
33    /// The candidate limit for exact refinement, or the index default.
34    pub rerank: Option<u32>,
35    /// The source predicate, applied before final neighbor selection.
36    pub filter: Option<Expression>,
37    /// Stable source RowIDs to exclude, for example the query sample.
38    pub exclude_row_ids: Vec<u64>,
39}
40
41impl Default for VectorSearchOptions {
42    fn default() -> Self {
43        Self {
44            limit: 20,
45            probes: None,
46            rerank: None,
47            filter: None,
48            exclude_row_ids: Vec::new(),
49        }
50    }
51}
52
53/// Bounded neighbors from one query against exact index and source versions.
54/// Pages reuse these neighbors and do not repeat the index search.
55#[derive(Clone)]
56pub struct VectorSearch {
57    source: Table,
58    query: Arc<[f64]>,
59    options: VectorSearchOptions,
60    info: VectorSearchInfo,
61    neighbors: RecordBatch,
62}
63
64impl VectorIndex {
65    /// Search this index against `source` with a fixed candidate budget.
66    /// A source version mismatch requires a refresh or the matching source snapshot.
67    pub async fn search(
68        &self,
69        source: &Table,
70        query: Vec<f64>,
71        options: VectorSearchOptions,
72    ) -> Result<VectorSearch, Error> {
73        if !Arc::ptr_eq(&source.connection, &self.project.client.connection) {
74            return Err(Error::invalid(
75                "The index and source must use the same client",
76            ));
77        }
78        let client = self.project.client.vector_indexes();
79        let read = client
80            .read_version(self.resource_id, self.version_id)
81            .await?;
82        let artifact = self.artifact(read, source)?;
83        let binding = client
84            .storage_binding(
85                self.resource_id,
86                StorageAccessRequest {
87                    version_id: self.version_id,
88                    mode: StorageMode::ReadOnly,
89                },
90            )
91            .await?;
92        if binding.mode != StorageMode::ReadOnly {
93            return Err(protocol(
94                "Vector Index storage grant has the wrong access mode",
95            ));
96        }
97        let issued: IssuedBinding = serde_json::to_value(binding)
98            .and_then(serde_json::from_value)
99            .map_err(|error| protocol(format!("Invalid Vector Index storage binding: {error}")))?;
100        self.project
101            .client
102            .connection
103            .tables()?
104            .search_vectors(
105                source.clone(),
106                artifact,
107                StorageBinding::for_resource(issued)?,
108                query,
109                options,
110            )
111            .await
112    }
113
114    fn artifact(
115        &self,
116        read: ReadVersionResponse,
117        source: &Table,
118    ) -> Result<VectorIndexArtifact, Error> {
119        let definition = read.definition;
120        let version = read.version;
121        if version.vector_index_id != self.resource_id || version.version_id != self.version_id {
122            return Err(protocol("Vector Index response changed its exact identity"));
123        }
124        let source_id = definition
125            .source_resource_id
126            .parse::<Uuid>()
127            .map_err(|_| protocol("Vector Index source is not a resource ID"))?;
128        let (Some(source_version), Some(physical)) = (version.source_version_id, version.physical)
129        else {
130            return Err(Error::invalid("Refresh the Vector Index before search"));
131        };
132        let source_version = source_version
133            .parse::<Uuid>()
134            .map_err(|_| protocol("Vector Index source version is not a UUID"))?;
135        if source_id != source.identity().resource_id
136            || definition.source_table_id != source.identity().resource_id
137            || source_version != source.identity().version_id
138        {
139            return Err(Error::invalid(
140                "The source does not match the exact Vector Index version",
141            ));
142        }
143        let metric = match definition.metric {
144            DistanceMetric::Cosine => VectorMetric::Cosine,
145            DistanceMetric::Euclidean => VectorMetric::Euclidean,
146            DistanceMetric::Dot => VectorMetric::Dot,
147            DistanceMetric::Unknown(_) => {
148                return Err(Error::invalid(
149                    "The Vector Index uses an unsupported metric",
150                ));
151            }
152        };
153        Ok(VectorIndexArtifact {
154            origin: VectorIndexOrigin {
155                vector_index_id: self.resource_id,
156                version_id: self.version_id,
157                source_resource_id: source_id,
158                source_version_id: source_version,
159                source_field_id: definition.source_field_id,
160            },
161            dimensions: definition.dimensions,
162            metric,
163            codec: physical.codec,
164            codec_version: physical.codec_version,
165            manifest_path: physical.manifest_path,
166            indexed_rows: physical.indexed_rows,
167            partitions: physical.partitions,
168            metadata: physical.metadata.0,
169        })
170    }
171}
172
173impl VectorSearch {
174    pub(crate) async fn read(
175        source: Table,
176        query: Vec<f64>,
177        options: VectorSearchOptions,
178        artifact: &VectorIndexArtifact,
179        evaluation: Evaluation,
180    ) -> Result<Self, Error> {
181        let schema = evaluation.schema;
182        let expected = [
183            (ROW_ID_COLUMN, DataType::UInt64),
184            (DISTANCE_COLUMN, DataType::Float64),
185            (RANK_COLUMN, DataType::UInt32),
186        ];
187        if schema.fields().len() != expected.len()
188            || schema
189                .fields()
190                .iter()
191                .zip(&expected)
192                .any(|(field, (name, dtype))| {
193                    field.name() != name || field.data_type() != dtype || field.is_nullable()
194                })
195        {
196            return Err(protocol(
197                "Vector search returned an invalid neighbor schema",
198            ));
199        }
200        let info: VectorSearchInfo = serde_json::from_str(
201            schema
202                .metadata()
203                .get(INFO_KEY)
204                .ok_or_else(|| protocol("Vector search omitted its provenance"))?,
205        )
206        .map_err(|error| protocol(format!("Invalid vector search provenance: {error}")))?;
207        if info.origin != artifact.origin
208            || info.limit != options.limit
209            || info.metric != artifact.metric
210        {
211            return Err(protocol("Vector search changed its exact query provenance"));
212        }
213        let neighbors = collect(schema, evaluation.batches, options.limit as usize).await?;
214        Ok(Self {
215            source,
216            query: query.into(),
217            options,
218            info,
219            neighbors,
220        })
221    }
222
223    /// Return the exact source handle for these results.
224    pub fn source(&self) -> &Table {
225        &self.source
226    }
227
228    /// Return the fixed query vector.
229    pub fn query(&self) -> &[f64] {
230        &self.query
231    }
232
233    /// Return the fixed predicate, exclusions, and recall controls.
234    pub fn options(&self) -> &VectorSearchOptions {
235        &self.options
236    }
237
238    /// Return exact provenance and the effective candidate budget.
239    pub fn info(&self) -> &VectorSearchInfo {
240        &self.info
241    }
242
243    /// Return the fixed RowIDs, float64 distances, and one-based ranks.
244    pub fn neighbors(&self) -> &RecordBatch {
245        &self.neighbors
246    }
247
248    /// Return the number of neighbors within this fixed budget.
249    pub fn len(&self) -> usize {
250        self.neighbors.num_rows()
251    }
252
253    /// Return whether this fixed budget produced no neighbors.
254    pub fn is_empty(&self) -> bool {
255        self.len() == 0
256    }
257
258    /// Read requested columns only for `rows`, in neighbor order.
259    /// Each page accepts at most 64 MiB of Arrow buffers.
260    pub async fn page(
261        &self,
262        rows: Range<usize>,
263        columns: Vec<String>,
264    ) -> Result<RecordBatch, Error> {
265        if rows.start > rows.end || rows.end > self.len() {
266            return Err(Error::invalid(
267                "The page range exceeds the fixed vector results",
268            ));
269        }
270        let page = self.neighbors.slice(rows.start, rows.end - rows.start);
271        if columns.is_empty() {
272            return Ok(page);
273        }
274        if columns
275            .iter()
276            .any(|name| page.schema().column_with_name(name).is_some())
277        {
278            return Err(Error::invalid(
279                "Page columns use a reserved vector result name",
280            ));
281        }
282        let ids = row_ids(&page)?;
283        let scan = self
284            .source
285            .scan(ScanOptions {
286                columns: Some(columns),
287                row_id_column: Some(ROW_ID_COLUMN.into()),
288                row_ids: Some(row_ranges(ids.values())?),
289                limit: Some(page.num_rows() as u64),
290                batch_rows: Some(256),
291                ..Default::default()
292            })
293            .await?;
294        let metadata = collect(scan.schema(), scan.into_stream(), page.num_rows()).await?;
295        let positions = row_ids(&metadata)?
296            .values()
297            .iter()
298            .enumerate()
299            .map(|(position, id)| (*id, position as u32))
300            .collect::<HashMap<_, _>>();
301        let indices = ids
302            .values()
303            .iter()
304            .map(|id| {
305                positions
306                    .get(id)
307                    .copied()
308                    .ok_or_else(|| protocol("The source page omitted a neighbor RowID"))
309            })
310            .collect::<Result<Vec<_>, _>>()?;
311        let metadata =
312            take_record_batch(&metadata, &UInt32Array::from(indices)).map_err(Error::arrow)?;
313        let mut fields = page.schema().fields().iter().cloned().collect::<Vec<_>>();
314        let mut arrays = page.columns().to_vec();
315        for (field, array) in metadata.schema().fields().iter().zip(metadata.columns()) {
316            if field.name() != ROW_ID_COLUMN {
317                fields.push(Arc::clone(field));
318                arrays.push(Arc::clone(array));
319            }
320        }
321        RecordBatch::try_new(
322            Arc::new(Schema::new_with_metadata(
323                fields,
324                page.schema().metadata().clone(),
325            )),
326            arrays,
327        )
328        .map_err(Error::arrow)
329    }
330}
331
332fn row_ids(batch: &RecordBatch) -> Result<&UInt64Array, Error> {
333    batch
334        .column_by_name(ROW_ID_COLUMN)
335        .and_then(|array| array.as_any().downcast_ref())
336        .ok_or_else(|| protocol("Vector results omitted their stable RowIDs"))
337}
338
339fn row_ranges(ids: &[u64]) -> Result<Vec<Range<u64>>, Error> {
340    let mut ids = ids.to_vec();
341    ids.sort_unstable();
342    let mut ranges: Vec<Range<u64>> = Vec::new();
343    for id in ids {
344        let end = id
345            .checked_add(1)
346            .ok_or_else(|| protocol("A neighbor RowID exceeds the supported range"))?;
347        match ranges.last_mut() {
348            Some(range) if range.end == id => range.end = end,
349            _ => ranges.push(id..end),
350        }
351    }
352    Ok(ranges)
353}
354
355async fn collect<S, E>(
356    schema: SchemaRef,
357    mut batches: S,
358    limit: usize,
359) -> Result<RecordBatch, Error>
360where
361    S: Stream<Item = Result<RecordBatch, E>> + Unpin,
362    E: Into<Error>,
363{
364    let mut rows = 0usize;
365    let mut bytes = 0usize;
366    let mut collected = Vec::new();
367    while let Some(batch) = batches.try_next().await.map_err(Into::into)? {
368        rows = rows.saturating_add(batch.num_rows());
369        bytes = bytes.saturating_add(batch.get_array_memory_size());
370        if rows > limit {
371            return Err(protocol("Vector results exceeded their row limit"));
372        }
373        if bytes > MAX_PAGE_BYTES {
374            return Err(Error::invalid("The selected page exceeds its byte limit"));
375        }
376        collected.push(batch);
377    }
378    concat_batches(&schema, &collected).map_err(Error::arrow)
379}
380
381fn protocol(message: impl Into<String>) -> Error {
382    Error::Protocol(message.into())
383}
384
385#[cfg(test)]
386mod tests;
spiral/vector_index/search.rs

spiral/vector_index/
search.rs
