1use std::io;
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 ) -> 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 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 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 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 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 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 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 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 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
206fn 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
234fn 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
257fn 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 _ => {} }
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}