Skip to main content

kevy_client/
index.rs

1//! Declarative secondary indexes: `IDX.CREATE` / `IDX.QUERY` /
2//! `IDX.DROP` / `IDX.LIST`.
3//!
4//! **Remote-only.** The embedded backend answers `Unsupported`: the
5//! wire face coerces query bounds through the index's declared value
6//! type server-side, which the client cannot replicate without the
7//! catalog. Embedded users match the public
8//! [`Connection::Embedded`](crate::Connection) variant and call
9//! [`kevy_embedded::Store`]'s typed `idx_*` API directly.
10//!
11//! Shape strategy: `IDX.CREATE` / `IDX.QUERY` have a large argument
12//! face (see docs/verb-reference.md), so the wrap is two-layered —
13//! typed shortcuts for the common forms (`RANGE` / `EQ` / `MATCH` /
14//! `KNN`, plain range-index create) plus `*_raw` argv passthroughs
15//! that keep every server capability reachable (COMPOSE, HYBRID,
16//! GROUPS, ANN create options, …) without this crate chasing the verb
17//! grammar release-by-release.
18
19use crate::{KevyError, KevyResult};
20
21use kevy_resp::Reply;
22
23use crate::{Connection, num_f64, num_u64, string, unexpected};
24
25/// Declared scalar type for [`Connection::idx_create_range`]
26/// (`TYPE i64|f64|str`). Vector/ANN indexes have extra required
27/// options — declare those via [`Connection::idx_create_raw`].
28#[derive(Debug, Clone, Copy, PartialEq, Eq)]
29pub enum IdxType {
30    /// `TYPE i64` — signed 64-bit integer field.
31    I64,
32    /// `TYPE f64` — finite 64-bit float field.
33    F64,
34    /// `TYPE str` — raw bytes, memcmp order.
35    Str,
36}
37
38impl IdxType {
39    fn tag(self) -> &'static [u8] {
40        match self {
41            Self::I64 => b"i64",
42            Self::F64 => b"f64",
43            Self::Str => b"str",
44        }
45    }
46}
47
48/// One `IDX.QUERY` hit: the row's key plus the indexed value's string
49/// form (the same repr the wire carries).
50#[derive(Debug, Clone, PartialEq, Eq)]
51pub struct IdxRow {
52    /// The matching key.
53    pub key: Vec<u8>,
54    /// The indexed field value, in its wire string form.
55    pub value: Vec<u8>,
56}
57
58/// One page of `IDX.QUERY RANGE`/`EQ` results.
59#[derive(Debug, Clone, PartialEq, Eq)]
60pub struct IdxPage {
61    /// Cursor for the next page — `None` when the scan is complete;
62    /// otherwise pass it back via `idx_query_range`'s `cursor`.
63    pub cursor: Option<Vec<u8>>,
64    /// The page's hits in `(value, key)` order.
65    pub rows: Vec<IdxRow>,
66}
67
68/// One declared index, as reported by `IDX.LIST`.
69#[derive(Debug, Clone, PartialEq, Eq)]
70pub struct IdxInfo {
71    /// Index name.
72    pub name: Vec<u8>,
73    /// Key prefix the index covers.
74    pub prefix: Vec<u8>,
75    /// Kind tag (`range` / `unique` / `text` / `ann` / `agg`).
76    pub kind: String,
77    /// Build state (`ready` / `building`).
78    pub state: String,
79    /// Total indexed entries across shards.
80    pub entries: u64,
81    /// Total index bytes across shards.
82    pub bytes: u64,
83}
84
85impl Connection {
86    /// `IDX.CREATE name ON PREFIX prefix FIELD field TYPE ty KIND range`
87    /// — declare a plain range index (the common case: range + EQ
88    /// queries over one hash field under a key prefix).
89    pub fn idx_create_range(
90        &mut self,
91        name: &[u8],
92        prefix: &[u8],
93        field: &[u8],
94        ty: IdxType,
95    ) -> KevyResult<()> {
96        let args: &[&[u8]] = &[
97            name,
98            b"ON",
99            b"PREFIX",
100            prefix,
101            b"FIELD",
102            field,
103            b"TYPE",
104            ty.tag(),
105            b"KIND",
106            b"range",
107        ];
108        self.idx_create_raw(args)
109    }
110
111    /// `IDX.CREATE <args…>` — raw passthrough; `args` is everything
112    /// after the verb (see docs/verb-reference.md for the full
113    /// grammar: MAXMEM, DIM/DISTANCE/M/EF for ANN, GROUPBY for agg…).
114    pub fn idx_create_raw(&mut self, args: &[&[u8]]) -> KevyResult<()> {
115        match self.remote("IDX.CREATE")?.request(&raw_argv(b"IDX.CREATE", args))? {
116            Reply::Simple(s) if s == b"OK" => Ok(()),
117            Reply::Error(e) => Err(KevyError::Protocol(string(e))),
118            other => Err(unexpected(other)),
119        }
120    }
121
122    /// `IDX.DROP name` — returns whether the index existed.
123    pub fn idx_drop(&mut self, name: &[u8]) -> KevyResult<bool> {
124        match self.remote("IDX.DROP")?.request_borrowed(&[b"IDX.DROP", name])? {
125            Reply::Int(1) => Ok(true),
126            Reply::Int(0) => Ok(false),
127            Reply::Error(e) => Err(KevyError::Protocol(string(e))),
128            other => Err(unexpected(other)),
129        }
130    }
131
132    /// `IDX.LIST` — declared indexes with build state and stats.
133    pub fn idx_list(&mut self) -> KevyResult<Vec<IdxInfo>> {
134        match self.remote("IDX.LIST")?.request_borrowed(&[b"IDX.LIST"])? {
135            Reply::Array(items) => items.into_iter().map(parse_info).collect(),
136            Reply::Error(e) => Err(KevyError::Protocol(string(e))),
137            other => Err(unexpected(other)),
138        }
139    }
140
141    /// `IDX.QUERY name RANGE min max [LIMIT n] [CURSOR c]` — one page
142    /// of `(key, value)` hits in `(value, key)` order. `min`/`max` are
143    /// the wire string forms (e.g. `b"18"`), coerced server-side per
144    /// the index's declared type. Resume with [`IdxPage::cursor`].
145    pub fn idx_query_range(
146        &mut self,
147        name: &[u8],
148        min: &[u8],
149        max: &[u8],
150        limit: usize,
151        cursor: Option<&[u8]>,
152    ) -> KevyResult<IdxPage> {
153        let lim = limit.to_string();
154        let mut args: Vec<&[u8]> = vec![name, b"RANGE", min, max, b"LIMIT", lim.as_bytes()];
155        if let Some(c) = cursor {
156            args.push(b"CURSOR");
157            args.push(c);
158        }
159        parse_page(self.idx_query_raw(&args)?)
160    }
161
162    /// `IDX.QUERY name EQ value [LIMIT n]` — point-lookup page.
163    pub fn idx_query_eq(&mut self, name: &[u8], value: &[u8], limit: usize) -> KevyResult<IdxPage> {
164        let lim = limit.to_string();
165        parse_page(self.idx_query_raw(&[name, b"EQ", value, b"LIMIT", lim.as_bytes()])?)
166    }
167
168    /// `IDX.QUERY name MATCH text [LIMIT n]` — BM25-ranked full-text
169    /// hits as `(key, score)`, best first (needs a `KIND text` index).
170    pub fn idx_query_match(
171        &mut self,
172        name: &[u8],
173        text: &[u8],
174        limit: usize,
175    ) -> KevyResult<Vec<(Vec<u8>, f64)>> {
176        let lim = limit.to_string();
177        parse_ranked(self.idx_query_raw(&[name, b"MATCH", text, b"LIMIT", lim.as_bytes()])?)
178    }
179
180    /// `IDX.QUERY name KNN vector [LIMIT k]` — nearest neighbours as
181    /// `(key, distance)`, closest first (needs a `KIND ann` index).
182    /// `vector` is encoded as the f32 little-endian blob the index
183    /// stores.
184    pub fn idx_query_knn(
185        &mut self,
186        name: &[u8],
187        vector: &[f32],
188        k: usize,
189    ) -> KevyResult<Vec<(Vec<u8>, f64)>> {
190        let blob: Vec<u8> = vector.iter().flat_map(|f| f.to_le_bytes()).collect();
191        let lim = k.to_string();
192        parse_ranked(self.idx_query_raw(&[name, b"KNN", &blob, b"LIMIT", lim.as_bytes()])?)
193    }
194
195    /// `IDX.QUERY <args…>` — raw passthrough returning the raw
196    /// [`Reply`]; `args` is everything after the verb. The escape
197    /// hatch for COMPOSE / HYBRID / GROUPS / FIELDS hydration and any
198    /// future query shape.
199    pub fn idx_query_raw(&mut self, args: &[&[u8]]) -> KevyResult<Reply> {
200        match self.remote("IDX.QUERY")?.request(&raw_argv(b"IDX.QUERY", args))? {
201            Reply::Error(e) => Err(KevyError::Protocol(string(e))),
202            other => Ok(other),
203        }
204    }
205}
206
207fn raw_argv(verb: &[u8], args: &[&[u8]]) -> Vec<Vec<u8>> {
208    let mut argv = Vec::with_capacity(args.len() + 1);
209    argv.push(verb.to_vec());
210    argv.extend(args.iter().map(|a| a.to_vec()));
211    argv
212}
213
214/// `*2 [cursor, *2N [key, value]…]` → [`IdxPage`]. Cursor `0` = done.
215fn parse_page(reply: Reply) -> KevyResult<IdxPage> {
216    let Reply::Array(items) = reply else {
217        return Err(unexpected(reply));
218    };
219    if items.len() != 2 {
220        return Err(KevyError::Protocol("IDX.QUERY page: expected [cursor, rows]".into()));
221    }
222    let mut it = items.into_iter();
223    let cursor = match it.next().unwrap() {
224        Reply::Bulk(c) if c == b"0" => None,
225        Reply::Bulk(c) => Some(c),
226        other => return Err(unexpected(other)),
227    };
228    let Reply::Array(flat) = it.next().unwrap() else {
229        return Err(KevyError::Protocol("IDX.QUERY page: rows not an array".into()));
230    };
231    let mut rows = Vec::with_capacity(flat.len() / 2);
232    let mut flat = flat.into_iter();
233    while let Some(k) = flat.next() {
234        let (Reply::Bulk(key), Some(Reply::Bulk(value))) = (k, flat.next()) else {
235            return Err(KevyError::Protocol("IDX.QUERY page: odd or non-bulk row pair".into()));
236        };
237        rows.push(IdxRow { key, value });
238    }
239    Ok(IdxPage { cursor, rows })
240}
241
242/// MATCH/KNN shape: `*N` rows of `*(2+2F) [key, score, fields…]` →
243/// `(key, score)` (shortcuts request no FIELDS, so F = 0).
244fn parse_ranked(reply: Reply) -> KevyResult<Vec<(Vec<u8>, f64)>> {
245    let Reply::Array(items) = reply else {
246        return Err(unexpected(reply));
247    };
248    items
249        .into_iter()
250        .map(|row| {
251            let Reply::Array(cells) = row else {
252                return Err(unexpected(row));
253            };
254            let mut it = cells.into_iter();
255            match (it.next(), it.next()) {
256                (Some(Reply::Bulk(key)), Some(Reply::Bulk(score))) => Ok((key, num_f64(&score)?)),
257                _ => Err(KevyError::Protocol("IDX.QUERY ranked row: expected [key, score]".into())),
258            }
259        })
260        .collect()
261}
262
263/// One `IDX.LIST` entry: a flat label/value bulk array
264/// (`name … prefix … kind … state … entries … bytes …`).
265fn parse_info(entry: Reply) -> KevyResult<IdxInfo> {
266    let Reply::Array(cells) = entry else {
267        return Err(unexpected(entry));
268    };
269    let mut info = IdxInfo {
270        name: Vec::new(),
271        prefix: Vec::new(),
272        kind: String::new(),
273        state: String::new(),
274        entries: 0,
275        bytes: 0,
276    };
277    let mut it = cells.into_iter();
278    while let (Some(Reply::Bulk(label)), Some(Reply::Bulk(value))) = (it.next(), it.next()) {
279        match label.as_slice() {
280            b"name" => info.name = value,
281            b"prefix" => info.prefix = value,
282            b"kind" => info.kind = string(value),
283            b"state" => info.state = string(value),
284            b"entries" => info.entries = num_u64(&value)?,
285            b"bytes" => info.bytes = num_u64(&value)?,
286            _ => {} // forward-compatible: skip labels this version doesn't know
287        }
288    }
289    Ok(info)
290}
291
292#[cfg(test)]
293mod tests {
294    use super::*;
295
296    #[test]
297    fn embedded_idx_is_unsupported() {
298        let mut c = Connection::connect("mem://").unwrap();
299        let err = c.idx_list().unwrap_err();
300        assert!(matches!(err, KevyError::Unsupported(_)));
301        let err = c.idx_create_range(b"i", b"user:", b"age", IdxType::I64).unwrap_err();
302        assert!(matches!(err, KevyError::Unsupported(_)));
303    }
304
305    #[test]
306    fn page_parser_maps_cursor_and_rows() {
307        let reply = Reply::Array(vec![
308            Reply::Bulk(b"abc1".to_vec()),
309            Reply::Array(vec![
310                Reply::Bulk(b"user:1".to_vec()),
311                Reply::Bulk(b"21".to_vec()),
312                Reply::Bulk(b"user:2".to_vec()),
313                Reply::Bulk(b"22".to_vec()),
314            ]),
315        ]);
316        let page = parse_page(reply).unwrap();
317        assert_eq!(page.cursor, Some(b"abc1".to_vec()));
318        assert_eq!(page.rows.len(), 2);
319        assert_eq!(page.rows[0].key, b"user:1");
320        assert_eq!(page.rows[0].value, b"21");
321
322        let done = parse_page(Reply::Array(vec![Reply::Bulk(b"0".to_vec()), Reply::Array(vec![])]))
323            .unwrap();
324        assert_eq!(done.cursor, None);
325        assert!(done.rows.is_empty());
326    }
327
328    #[test]
329    fn ranked_parser_maps_key_score() {
330        let reply = Reply::Array(vec![Reply::Array(vec![
331            Reply::Bulk(b"doc:1".to_vec()),
332            Reply::Bulk(b"1.5".to_vec()),
333        ])]);
334        assert_eq!(parse_ranked(reply).unwrap(), vec![(b"doc:1".to_vec(), 1.5)]);
335    }
336}