Skip to main content

kevy_client/
index.rs

1//! Declarative secondary indexes: `IDX.CREATE` / `IDX.QUERY` /
2//! `IDX.DROP` / `IDX.LIST` (v1.14.0).
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 std::io;
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    ) -> io::Result<()> {
96        let args: &[&[u8]] = &[
97            name, b"ON", b"PREFIX", prefix, b"FIELD", field, b"TYPE", ty.tag(), b"KIND",
98            b"range",
99        ];
100        self.idx_create_raw(args)
101    }
102
103    /// `IDX.CREATE <args…>` — raw passthrough; `args` is everything
104    /// after the verb (see docs/verb-reference.md for the full
105    /// grammar: MAXMEM, DIM/DISTANCE/M/EF for ANN, GROUPBY for agg…).
106    pub fn idx_create_raw(&mut self, args: &[&[u8]]) -> io::Result<()> {
107        match self.remote("IDX.CREATE")?.request(&raw_argv(b"IDX.CREATE", args))? {
108            Reply::Simple(s) if s == b"OK" => Ok(()),
109            Reply::Error(e) => Err(io::Error::other(string(e))),
110            other => Err(unexpected(other)),
111        }
112    }
113
114    /// `IDX.DROP name` — returns whether the index existed.
115    pub fn idx_drop(&mut self, name: &[u8]) -> io::Result<bool> {
116        match self.remote("IDX.DROP")?.request_borrowed(&[b"IDX.DROP", name])? {
117            Reply::Int(1) => Ok(true),
118            Reply::Int(0) => Ok(false),
119            Reply::Error(e) => Err(io::Error::other(string(e))),
120            other => Err(unexpected(other)),
121        }
122    }
123
124    /// `IDX.LIST` — declared indexes with build state and stats.
125    pub fn idx_list(&mut self) -> io::Result<Vec<IdxInfo>> {
126        match self.remote("IDX.LIST")?.request_borrowed(&[b"IDX.LIST"])? {
127            Reply::Array(items) => items.into_iter().map(parse_info).collect(),
128            Reply::Error(e) => Err(io::Error::other(string(e))),
129            other => Err(unexpected(other)),
130        }
131    }
132
133    /// `IDX.QUERY name RANGE min max [LIMIT n] [CURSOR c]` — one page
134    /// of `(key, value)` hits in `(value, key)` order. `min`/`max` are
135    /// the wire string forms (e.g. `b"18"`), coerced server-side per
136    /// the index's declared type. Resume with [`IdxPage::cursor`].
137    pub fn idx_query_range(
138        &mut self,
139        name: &[u8],
140        min: &[u8],
141        max: &[u8],
142        limit: usize,
143        cursor: Option<&[u8]>,
144    ) -> io::Result<IdxPage> {
145        let lim = limit.to_string();
146        let mut args: Vec<&[u8]> = vec![name, b"RANGE", min, max, b"LIMIT", lim.as_bytes()];
147        if let Some(c) = cursor {
148            args.push(b"CURSOR");
149            args.push(c);
150        }
151        parse_page(self.idx_query_raw(&args)?)
152    }
153
154    /// `IDX.QUERY name EQ value [LIMIT n]` — point-lookup page.
155    pub fn idx_query_eq(&mut self, name: &[u8], value: &[u8], limit: usize) -> io::Result<IdxPage> {
156        let lim = limit.to_string();
157        parse_page(self.idx_query_raw(&[name, b"EQ", value, b"LIMIT", lim.as_bytes()])?)
158    }
159
160    /// `IDX.QUERY name MATCH text [LIMIT n]` — BM25-ranked full-text
161    /// hits as `(key, score)`, best first (needs a `KIND text` index).
162    pub fn idx_query_match(
163        &mut self,
164        name: &[u8],
165        text: &[u8],
166        limit: usize,
167    ) -> io::Result<Vec<(Vec<u8>, f64)>> {
168        let lim = limit.to_string();
169        parse_ranked(self.idx_query_raw(&[name, b"MATCH", text, b"LIMIT", lim.as_bytes()])?)
170    }
171
172    /// `IDX.QUERY name KNN vector [LIMIT k]` — nearest neighbours as
173    /// `(key, distance)`, closest first (needs a `KIND ann` index).
174    /// `vector` is encoded as the f32 little-endian blob the index
175    /// stores.
176    pub fn idx_query_knn(
177        &mut self,
178        name: &[u8],
179        vector: &[f32],
180        k: usize,
181    ) -> io::Result<Vec<(Vec<u8>, f64)>> {
182        let blob: Vec<u8> = vector.iter().flat_map(|f| f.to_le_bytes()).collect();
183        let lim = k.to_string();
184        parse_ranked(self.idx_query_raw(&[name, b"KNN", &blob, b"LIMIT", lim.as_bytes()])?)
185    }
186
187    /// `IDX.QUERY <args…>` — raw passthrough returning the raw
188    /// [`Reply`]; `args` is everything after the verb. The escape
189    /// hatch for COMPOSE / HYBRID / GROUPS / FIELDS hydration and any
190    /// future query shape.
191    pub fn idx_query_raw(&mut self, args: &[&[u8]]) -> io::Result<Reply> {
192        match self.remote("IDX.QUERY")?.request(&raw_argv(b"IDX.QUERY", args))? {
193            Reply::Error(e) => Err(io::Error::other(string(e))),
194            other => Ok(other),
195        }
196    }
197}
198
199fn raw_argv(verb: &[u8], args: &[&[u8]]) -> Vec<Vec<u8>> {
200    let mut argv = Vec::with_capacity(args.len() + 1);
201    argv.push(verb.to_vec());
202    argv.extend(args.iter().map(|a| a.to_vec()));
203    argv
204}
205
206/// `*2 [cursor, *2N [key, value]…]` → [`IdxPage`]. Cursor `0` = done.
207fn parse_page(reply: Reply) -> io::Result<IdxPage> {
208    let Reply::Array(items) = reply else {
209        return Err(unexpected(reply));
210    };
211    if items.len() != 2 {
212        return Err(io::Error::other("IDX.QUERY page: expected [cursor, rows]"));
213    }
214    let mut it = items.into_iter();
215    let cursor = match it.next().unwrap() {
216        Reply::Bulk(c) if c == b"0" => None,
217        Reply::Bulk(c) => Some(c),
218        other => return Err(unexpected(other)),
219    };
220    let Reply::Array(flat) = it.next().unwrap() else {
221        return Err(io::Error::other("IDX.QUERY page: rows not an array"));
222    };
223    let mut rows = Vec::with_capacity(flat.len() / 2);
224    let mut flat = flat.into_iter();
225    while let Some(k) = flat.next() {
226        let (Reply::Bulk(key), Some(Reply::Bulk(value))) = (k, flat.next()) else {
227            return Err(io::Error::other("IDX.QUERY page: odd or non-bulk row pair"));
228        };
229        rows.push(IdxRow { key, value });
230    }
231    Ok(IdxPage { cursor, rows })
232}
233
234/// MATCH/KNN shape: `*N` rows of `*(2+2F) [key, score, fields…]` →
235/// `(key, score)` (shortcuts request no FIELDS, so F = 0).
236fn parse_ranked(reply: Reply) -> io::Result<Vec<(Vec<u8>, f64)>> {
237    let Reply::Array(items) = reply else {
238        return Err(unexpected(reply));
239    };
240    items
241        .into_iter()
242        .map(|row| {
243            let Reply::Array(cells) = row else {
244                return Err(unexpected(row));
245            };
246            let mut it = cells.into_iter();
247            match (it.next(), it.next()) {
248                (Some(Reply::Bulk(key)), Some(Reply::Bulk(score))) => {
249                    Ok((key, num_f64(&score)?))
250                }
251                _ => Err(io::Error::other("IDX.QUERY ranked row: expected [key, score]")),
252            }
253        })
254        .collect()
255}
256
257/// One `IDX.LIST` entry: a flat label/value bulk array
258/// (`name … prefix … kind … state … entries … bytes …`).
259fn parse_info(entry: Reply) -> io::Result<IdxInfo> {
260    let Reply::Array(cells) = entry else {
261        return Err(unexpected(entry));
262    };
263    let mut info = IdxInfo {
264        name: Vec::new(),
265        prefix: Vec::new(),
266        kind: String::new(),
267        state: String::new(),
268        entries: 0,
269        bytes: 0,
270    };
271    let mut it = cells.into_iter();
272    while let (Some(Reply::Bulk(label)), Some(Reply::Bulk(value))) = (it.next(), it.next()) {
273        match label.as_slice() {
274            b"name" => info.name = value,
275            b"prefix" => info.prefix = value,
276            b"kind" => info.kind = string(value),
277            b"state" => info.state = string(value),
278            b"entries" => info.entries = num_u64(&value)?,
279            b"bytes" => info.bytes = num_u64(&value)?,
280            _ => {} // forward-compatible: skip labels this version doesn't know
281        }
282    }
283    Ok(info)
284}
285
286#[cfg(test)]
287mod tests {
288    use super::*;
289
290    #[test]
291    fn embedded_idx_is_unsupported() {
292        let mut c = Connection::open("mem://").unwrap();
293        let err = c.idx_list().unwrap_err();
294        assert_eq!(err.kind(), io::ErrorKind::Unsupported);
295        let err = c
296            .idx_create_range(b"i", b"user:", b"age", IdxType::I64)
297            .unwrap_err();
298        assert_eq!(err.kind(), io::ErrorKind::Unsupported);
299    }
300
301    #[test]
302    fn page_parser_maps_cursor_and_rows() {
303        let reply = Reply::Array(vec![
304            Reply::Bulk(b"abc1".to_vec()),
305            Reply::Array(vec![
306                Reply::Bulk(b"user:1".to_vec()),
307                Reply::Bulk(b"21".to_vec()),
308                Reply::Bulk(b"user:2".to_vec()),
309                Reply::Bulk(b"22".to_vec()),
310            ]),
311        ]);
312        let page = parse_page(reply).unwrap();
313        assert_eq!(page.cursor, Some(b"abc1".to_vec()));
314        assert_eq!(page.rows.len(), 2);
315        assert_eq!(page.rows[0].key, b"user:1");
316        assert_eq!(page.rows[0].value, b"21");
317
318        let done = parse_page(Reply::Array(vec![
319            Reply::Bulk(b"0".to_vec()),
320            Reply::Array(vec![]),
321        ]))
322        .unwrap();
323        assert_eq!(done.cursor, None);
324        assert!(done.rows.is_empty());
325    }
326
327    #[test]
328    fn ranked_parser_maps_key_score() {
329        let reply = Reply::Array(vec![Reply::Array(vec![
330            Reply::Bulk(b"doc:1".to_vec()),
331            Reply::Bulk(b"1.5".to_vec()),
332        ])]);
333        assert_eq!(parse_ranked(reply).unwrap(), vec![(b"doc:1".to_vec(), 1.5)]);
334    }
335}