1use crate::{KevyError, KevyResult};
20
21use kevy_resp::Reply;
22
23use crate::{Connection, num_f64, num_u64, string, unexpected};
24
25#[derive(Debug, Clone, Copy, PartialEq, Eq)]
29pub enum IdxType {
30 I64,
32 F64,
34 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#[derive(Debug, Clone, PartialEq, Eq)]
51pub struct IdxRow {
52 pub key: Vec<u8>,
54 pub value: Vec<u8>,
56}
57
58#[derive(Debug, Clone, PartialEq, Eq)]
60pub struct IdxPage {
61 pub cursor: Option<Vec<u8>>,
64 pub rows: Vec<IdxRow>,
66}
67
68#[derive(Debug, Clone, PartialEq, Eq)]
70pub struct IdxInfo {
71 pub name: Vec<u8>,
73 pub prefix: Vec<u8>,
75 pub kind: String,
77 pub state: String,
79 pub entries: u64,
81 pub bytes: u64,
83}
84
85impl Connection {
86 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 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 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 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 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 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 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 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 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
214fn 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
242fn 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
263fn 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 _ => {} }
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}