uni_plugin/traits/
index.rs

1//! Index kind plugins — custom vector, FTS, geo indexes.
2
3use arrow_array::BooleanArray;
4use arrow_schema::SchemaRef;
5use datafusion::arrow::record_batch::RecordBatch;
6use smol_str::SmolStr;
7
8use crate::errors::FnError;
9
10/// Identifier for an index kind (`"vector"`, `"fts"`, `"hnsw"`, …).
11#[derive(Clone, Debug, PartialEq, Eq, Hash)]
12pub struct IndexKind(pub SmolStr);
13
14impl IndexKind {
15    /// Construct an `IndexKind` from a string.
16    #[must_use]
17    pub fn new(s: impl Into<SmolStr>) -> Self {
18        Self(s.into())
19    }
20}
21
22/// An index-kind provider that knows how to build / open / probe / persist
23/// a custom index.
24pub trait IndexKindProvider: Send + Sync {
25    /// The index kind this provider implements.
26    fn kind(&self) -> IndexKind;
27
28    /// Build a new index from a source record batch.
29    ///
30    /// `options` is free-form JSON configuration.
31    ///
32    /// # Errors
33    ///
34    /// Returns [`FnError`] on build failure (out of memory, bad
35    /// configuration).
36    fn build(&self, source: &RecordBatch, options: &str) -> Result<Box<dyn IndexBuild>, FnError>;
37
38    /// Open an index from previously-persisted bytes.
39    ///
40    /// # Errors
41    ///
42    /// Returns [`FnError`] if the bytes are malformed or incompatible.
43    fn open(&self, persisted: &[u8]) -> Result<Box<dyn IndexHandle>, FnError>;
44}
45
46/// In-flight index build (write side).
47pub trait IndexBuild: Send + Sync {
48    /// Finalize the build and produce a queryable handle.
49    ///
50    /// # Errors
51    ///
52    /// Returns [`FnError`] if finalization fails.
53    fn finalize(self: Box<Self>) -> Result<Box<dyn IndexHandle>, FnError>;
54}
55
56/// Queryable, persistable index handle (read side).
57pub trait IndexHandle: Send + Sync {
58    /// Probe the index with `query` and return up to `k` matches.
59    ///
60    /// # Errors
61    ///
62    /// Returns [`FnError`] on probe failure.
63    fn probe(&self, query: &RecordBatch, k: usize) -> Result<RecordBatch, FnError>;
64
65    /// Whether this index supports per-probe filter pushdown.
66    fn supports_filter(&self) -> bool {
67        false
68    }
69
70    /// Probe with a row-level filter applied.
71    ///
72    /// # Errors
73    ///
74    /// Returns [`FnError`] if filtered probe is unsupported or fails.
75    fn probe_filtered(
76        &self,
77        _query: &RecordBatch,
78        _k: usize,
79        _filter: &BooleanArray,
80    ) -> Result<RecordBatch, FnError> {
81        Err(FnError::new(
82            0x20,
83            "index does not support filter-pushdown probe",
84        ))
85    }
86
87    /// Serialize this index for persistence.
88    ///
89    /// # Errors
90    ///
91    /// Returns [`FnError`] if serialization fails.
92    fn persist(&self) -> Result<Vec<u8>, FnError>;
93
94    /// Output schema of `probe` / `probe_filtered`.
95    fn schema(&self) -> SchemaRef;
96}
uni_plugin/traits/index.rs

uni_plugin/traits/
index.rs
