1use std::collections::HashMap;
22
23#[derive(Debug, Clone, Copy, PartialEq, Eq)]
27pub enum VectorAlgorithm {
28 Hnsw,
31 Flat,
34}
35
36impl std::fmt::Display for VectorAlgorithm {
37 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
38 match self {
39 VectorAlgorithm::Hnsw => write!(f, "HNSW"),
40 VectorAlgorithm::Flat => write!(f, "FLAT"),
41 }
42 }
43}
44
45#[derive(Debug, Clone, Copy, PartialEq, Eq)]
47pub enum DistanceMetric {
48 L2,
50 InnerProduct,
52 Cosine,
54}
55
56impl std::fmt::Display for DistanceMetric {
57 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
58 match self {
59 DistanceMetric::L2 => write!(f, "L2"),
60 DistanceMetric::InnerProduct => write!(f, "IP"),
61 DistanceMetric::Cosine => write!(f, "COSINE"),
62 }
63 }
64}
65
66#[derive(Debug, Clone, PartialEq)]
68pub struct VectorParams {
69 pub algorithm: VectorAlgorithm,
71 pub dim: usize,
73 pub distance_metric: DistanceMetric,
75 pub hnsw_m: Option<usize>,
77 pub hnsw_ef_construction: Option<usize>,
79}
80
81impl VectorParams {
82 pub fn hnsw(dim: usize, distance_metric: DistanceMetric) -> Self {
85 Self {
86 algorithm: VectorAlgorithm::Hnsw,
87 dim,
88 distance_metric,
89 hnsw_m: None,
90 hnsw_ef_construction: None,
91 }
92 }
93
94 pub fn flat(dim: usize, distance_metric: DistanceMetric) -> Self {
96 Self {
97 algorithm: VectorAlgorithm::Flat,
98 dim,
99 distance_metric,
100 hnsw_m: None,
101 hnsw_ef_construction: None,
102 }
103 }
104
105 pub fn with_m(mut self, m: usize) -> Self {
108 self.hnsw_m = Some(m);
109 self
110 }
111
112 pub fn with_ef_construction(mut self, ef: usize) -> Self {
115 self.hnsw_ef_construction = Some(ef);
116 self
117 }
118
119 fn to_schema_args(&self) -> Vec<String> {
122 let mut args = vec![
123 "TYPE".to_string(),
124 "FLOAT32".to_string(),
125 "DIM".to_string(),
126 self.dim.to_string(),
127 "DISTANCE_METRIC".to_string(),
128 self.distance_metric.to_string(),
129 ];
130
131 if self.algorithm == VectorAlgorithm::Hnsw {
133 if let Some(m) = self.hnsw_m {
134 args.push("M".to_string());
135 args.push(m.to_string());
136 }
137 if let Some(ef) = self.hnsw_ef_construction {
138 args.push("EF_CONSTRUCTION".to_string());
139 args.push(ef.to_string());
140 }
141 }
142
143 let mut result = vec![args.len().to_string()];
145 result.extend(args);
146 result
147 }
148}
149
150#[derive(Debug, Clone)]
152pub struct SearchIndex {
153 pub name: String,
155 pub prefix: String,
157 pub fields: Vec<SearchField>,
159}
160
161impl SearchIndex {
162 pub fn new(name: impl Into<String>, prefix: impl Into<String>) -> Self {
164 Self {
165 name: name.into(),
166 prefix: prefix.into(),
167 fields: Vec::new(),
168 }
169 }
170
171 pub fn text(mut self, name: impl Into<String>) -> Self {
173 self.fields.push(SearchField {
174 name: name.into(),
175 json_path: None,
176 field_type: SearchFieldType::Text,
177 sortable: false,
178 no_index: false,
179 });
180 self
181 }
182
183 pub fn text_at(mut self, name: impl Into<String>, json_path: impl Into<String>) -> Self {
185 self.fields.push(SearchField {
186 name: name.into(),
187 json_path: Some(json_path.into()),
188 field_type: SearchFieldType::Text,
189 sortable: false,
190 no_index: false,
191 });
192 self
193 }
194
195 pub fn text_sortable(mut self, name: impl Into<String>) -> Self {
197 self.fields.push(SearchField {
198 name: name.into(),
199 json_path: None,
200 field_type: SearchFieldType::Text,
201 sortable: true,
202 no_index: false,
203 });
204 self
205 }
206
207 pub fn text_sortable_at(mut self, name: impl Into<String>, json_path: impl Into<String>) -> Self {
209 self.fields.push(SearchField {
210 name: name.into(),
211 json_path: Some(json_path.into()),
212 field_type: SearchFieldType::Text,
213 sortable: true,
214 no_index: false,
215 });
216 self
217 }
218
219 pub fn numeric(mut self, name: impl Into<String>) -> Self {
221 self.fields.push(SearchField {
222 name: name.into(),
223 json_path: None,
224 field_type: SearchFieldType::Numeric,
225 sortable: false,
226 no_index: false,
227 });
228 self
229 }
230
231 pub fn numeric_at(mut self, name: impl Into<String>, json_path: impl Into<String>) -> Self {
233 self.fields.push(SearchField {
234 name: name.into(),
235 json_path: Some(json_path.into()),
236 field_type: SearchFieldType::Numeric,
237 sortable: false,
238 no_index: false,
239 });
240 self
241 }
242
243 pub fn numeric_sortable(mut self, name: impl Into<String>) -> Self {
245 self.fields.push(SearchField {
246 name: name.into(),
247 json_path: None,
248 field_type: SearchFieldType::Numeric,
249 sortable: true,
250 no_index: false,
251 });
252 self
253 }
254
255 pub fn numeric_sortable_at(mut self, name: impl Into<String>, json_path: impl Into<String>) -> Self {
257 self.fields.push(SearchField {
258 name: name.into(),
259 json_path: Some(json_path.into()),
260 field_type: SearchFieldType::Numeric,
261 sortable: true,
262 no_index: false,
263 });
264 self
265 }
266
267 pub fn tag(mut self, name: impl Into<String>) -> Self {
269 self.fields.push(SearchField {
270 name: name.into(),
271 json_path: None,
272 field_type: SearchFieldType::Tag,
273 sortable: false,
274 no_index: false,
275 });
276 self
277 }
278
279 pub fn tag_at(mut self, name: impl Into<String>, json_path: impl Into<String>) -> Self {
281 self.fields.push(SearchField {
282 name: name.into(),
283 json_path: Some(json_path.into()),
284 field_type: SearchFieldType::Tag,
285 sortable: false,
286 no_index: false,
287 });
288 self
289 }
290
291 pub fn geo(mut self, name: impl Into<String>) -> Self {
293 self.fields.push(SearchField {
294 name: name.into(),
295 json_path: None,
296 field_type: SearchFieldType::Geo,
297 sortable: false,
298 no_index: false,
299 });
300 self
301 }
302
303 pub fn geo_at(mut self, name: impl Into<String>, json_path: impl Into<String>) -> Self {
305 self.fields.push(SearchField {
306 name: name.into(),
307 json_path: Some(json_path.into()),
308 field_type: SearchFieldType::Geo,
309 sortable: false,
310 no_index: false,
311 });
312 self
313 }
314
315 pub fn vector_hnsw(
331 mut self,
332 name: impl Into<String>,
333 dim: usize,
334 metric: DistanceMetric,
335 ) -> Self {
336 self.fields.push(SearchField {
337 name: name.into(),
338 json_path: None,
339 field_type: SearchFieldType::Vector(VectorParams::hnsw(dim, metric)),
340 sortable: false,
341 no_index: false,
342 });
343 self
344 }
345
346 pub fn vector_hnsw_at(
348 mut self,
349 name: impl Into<String>,
350 json_path: impl Into<String>,
351 dim: usize,
352 metric: DistanceMetric,
353 ) -> Self {
354 self.fields.push(SearchField {
355 name: name.into(),
356 json_path: Some(json_path.into()),
357 field_type: SearchFieldType::Vector(VectorParams::hnsw(dim, metric)),
358 sortable: false,
359 no_index: false,
360 });
361 self
362 }
363
364 pub fn vector_with_params(mut self, name: impl Into<String>, params: VectorParams) -> Self {
380 self.fields.push(SearchField {
381 name: name.into(),
382 json_path: None,
383 field_type: SearchFieldType::Vector(params),
384 sortable: false,
385 no_index: false,
386 });
387 self
388 }
389
390 pub fn vector_flat(
400 mut self,
401 name: impl Into<String>,
402 dim: usize,
403 metric: DistanceMetric,
404 ) -> Self {
405 self.fields.push(SearchField {
406 name: name.into(),
407 json_path: None,
408 field_type: SearchFieldType::Vector(VectorParams::flat(dim, metric)),
409 sortable: false,
410 no_index: false,
411 });
412 self
413 }
414
415 pub fn vector_flat_at(
417 mut self,
418 name: impl Into<String>,
419 json_path: impl Into<String>,
420 dim: usize,
421 metric: DistanceMetric,
422 ) -> Self {
423 self.fields.push(SearchField {
424 name: name.into(),
425 json_path: Some(json_path.into()),
426 field_type: SearchFieldType::Vector(VectorParams::flat(dim, metric)),
427 sortable: false,
428 no_index: false,
429 });
430 self
431 }
432
433 pub fn to_ft_create_args(&self) -> Vec<String> {
435 self.to_ft_create_args_with_prefix(None)
436 }
437
438 pub fn to_ft_create_args_with_prefix(&self, redis_prefix: Option<&str>) -> Vec<String> {
443 let prefix = redis_prefix.unwrap_or("");
444
445 let mut args = vec![
446 format!("{}idx:{}", prefix, self.name),
447 "ON".to_string(),
448 "JSON".to_string(),
449 "PREFIX".to_string(),
450 "1".to_string(),
451 format!("{}{}", prefix, self.prefix),
452 "SCHEMA".to_string(),
453 ];
454
455 for field in &self.fields {
456 args.extend(field.to_schema_args());
457 }
458
459 args
460 }
461}
462
463#[derive(Debug, Clone)]
465pub struct SearchField {
466 pub name: String,
468 pub json_path: Option<String>,
470 pub field_type: SearchFieldType,
472 pub sortable: bool,
474 pub no_index: bool,
476}
477
478impl SearchField {
479 fn to_schema_args(&self) -> Vec<String> {
480 let json_path = self
482 .json_path
483 .clone()
484 .unwrap_or_else(|| format!("$.payload.{}", self.name));
485
486 let mut args = vec![
487 json_path,
488 "AS".to_string(),
489 self.name.clone(),
490 ];
491
492 match &self.field_type {
494 SearchFieldType::Vector(params) => {
495 args.push("VECTOR".to_string());
496 args.push(params.algorithm.to_string());
497 args.extend(params.to_schema_args());
498 }
499 _ => {
500 args.push(self.field_type.to_string());
501 }
502 }
503
504 if self.sortable {
505 args.push("SORTABLE".to_string());
506 }
507
508 if self.no_index {
509 args.push("NOINDEX".to_string());
510 }
511
512 args
513 }
514}
515
516#[derive(Debug, Clone, PartialEq)]
518pub enum SearchFieldType {
519 Text,
521 Numeric,
523 Tag,
525 Geo,
527 Vector(VectorParams),
529}
530
531impl std::fmt::Display for SearchFieldType {
532 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
533 match self {
534 SearchFieldType::Text => write!(f, "TEXT"),
535 SearchFieldType::Numeric => write!(f, "NUMERIC"),
536 SearchFieldType::Tag => write!(f, "TAG"),
537 SearchFieldType::Geo => write!(f, "GEO"),
538 SearchFieldType::Vector(params) => write!(f, "VECTOR {}", params.algorithm),
539 }
540 }
541}
542
543pub struct IndexManager {
545 indexes: HashMap<String, SearchIndex>,
547}
548
549impl IndexManager {
550 pub fn new() -> Self {
552 Self {
553 indexes: HashMap::new(),
554 }
555 }
556
557 pub fn register(&mut self, index: SearchIndex) {
559 self.indexes.insert(index.name.clone(), index);
560 }
561
562 pub fn get(&self, name: &str) -> Option<&SearchIndex> {
564 self.indexes.get(name)
565 }
566
567 pub fn all(&self) -> impl Iterator<Item = &SearchIndex> {
569 self.indexes.values()
570 }
571
572 pub fn ft_create_args(&self, name: &str) -> Option<Vec<String>> {
574 self.indexes.get(name).map(|idx| idx.to_ft_create_args())
575 }
576
577 pub fn find_by_prefix(&self, key: &str) -> Option<&SearchIndex> {
579 self.indexes.values().find(|idx| key.starts_with(&idx.prefix))
580 }
581}
582
583impl Default for IndexManager {
584 fn default() -> Self {
585 Self::new()
586 }
587}
588
589#[cfg(test)]
590mod tests {
591 use super::*;
592
593 #[test]
594 fn test_simple_index() {
595 let index = SearchIndex::new("users", "crdt:users:")
596 .text("name")
597 .text("email")
598 .numeric("age");
599
600 let args = index.to_ft_create_args();
601 assert_eq!(args[0], "idx:users");
602 assert_eq!(args[1], "ON");
603 assert_eq!(args[2], "JSON");
604 assert_eq!(args[3], "PREFIX");
605 assert_eq!(args[4], "1");
606 assert_eq!(args[5], "crdt:users:");
607 assert_eq!(args[6], "SCHEMA");
608 assert!(args.contains(&"$.payload.name".to_string()));
610 assert!(args.contains(&"name".to_string()));
611 assert!(args.contains(&"TEXT".to_string()));
612 }
613
614 #[test]
615 fn test_sortable_fields() {
616 let index = SearchIndex::new("users", "crdt:users:")
617 .text_sortable("name")
618 .numeric_sortable("age");
619
620 let args = index.to_ft_create_args();
621 let sortable_count = args.iter().filter(|a| *a == "SORTABLE").count();
623 assert_eq!(sortable_count, 2);
624 }
625
626 #[test]
627 fn test_tag_field() {
628 let index = SearchIndex::new("items", "crdt:items:").tag("tags");
629
630 let args = index.to_ft_create_args();
631 assert!(args.contains(&"TAG".to_string()));
632 }
633
634 #[test]
635 fn test_custom_json_path() {
636 let index =
637 SearchIndex::new("users", "crdt:users:").text_at("username", "$.profile.name");
638
639 let args = index.to_ft_create_args();
640 assert!(args.contains(&"$.profile.name".to_string()));
641 assert!(args.contains(&"username".to_string()));
642 }
643
644 #[test]
645 fn test_index_manager_register() {
646 let mut manager = IndexManager::new();
647
648 let index = SearchIndex::new("users", "crdt:users:")
649 .text("name")
650 .numeric("age");
651
652 manager.register(index);
653
654 assert!(manager.get("users").is_some());
655 assert!(manager.get("unknown").is_none());
656 }
657
658 #[test]
659 fn test_index_manager_find_by_prefix() {
660 let mut manager = IndexManager::new();
661
662 manager.register(SearchIndex::new("users", "crdt:users:"));
663 manager.register(SearchIndex::new("posts", "crdt:posts:"));
664
665 let found = manager.find_by_prefix("crdt:users:abc123");
666 assert!(found.is_some());
667 assert_eq!(found.unwrap().name, "users");
668
669 let found = manager.find_by_prefix("crdt:posts:xyz");
670 assert!(found.is_some());
671 assert_eq!(found.unwrap().name, "posts");
672
673 let not_found = manager.find_by_prefix("crdt:comments:1");
674 assert!(not_found.is_none());
675 }
676
677 #[test]
678 fn test_ft_create_full_command() {
679 let index = SearchIndex::new("users", "crdt:users:")
680 .text_sortable("name")
681 .text("email")
682 .numeric_sortable("age")
683 .tag("roles");
684
685 let args = index.to_ft_create_args();
686
687 assert_eq!(args[0], "idx:users");
696 assert_eq!(args[6], "SCHEMA");
697
698 let cmd = format!("FT.CREATE {}", args.join(" "));
700 assert!(cmd.contains("idx:users"));
701 assert!(cmd.contains("ON JSON"));
702 assert!(cmd.contains("PREFIX 1 crdt:users:"));
703 assert!(cmd.contains("$.payload.name AS name TEXT SORTABLE"));
704 assert!(cmd.contains("$.payload.email AS email TEXT"));
705 assert!(cmd.contains("$.payload.age AS age NUMERIC SORTABLE"));
706 assert!(cmd.contains("$.payload.roles AS roles TAG"));
707 }
708
709 #[test]
710 fn test_vector_hnsw_basic() {
711 let index = SearchIndex::new("docs", "crdt:docs:")
712 .text("title")
713 .vector_hnsw("embedding", 1536, DistanceMetric::Cosine);
714
715 let args = index.to_ft_create_args();
716 let cmd = format!("FT.CREATE {}", args.join(" "));
717
718 assert!(cmd.contains("$.payload.embedding AS embedding VECTOR HNSW"));
720 assert!(cmd.contains("TYPE FLOAT32"));
721 assert!(cmd.contains("DIM 1536"));
722 assert!(cmd.contains("DISTANCE_METRIC COSINE"));
723 }
724
725 #[test]
726 fn test_vector_hnsw_with_params() {
727 let params = VectorParams::hnsw(768, DistanceMetric::L2)
728 .with_m(32)
729 .with_ef_construction(400);
730
731 let index =
732 SearchIndex::new("embeddings", "crdt:embeddings:").vector_with_params("vec", params);
733
734 let args = index.to_ft_create_args();
735 let cmd = format!("FT.CREATE {}", args.join(" "));
736
737 assert!(cmd.contains("VECTOR HNSW"));
738 assert!(cmd.contains("DIM 768"));
739 assert!(cmd.contains("DISTANCE_METRIC L2"));
740 assert!(cmd.contains("M 32"));
741 assert!(cmd.contains("EF_CONSTRUCTION 400"));
742 }
743
744 #[test]
745 fn test_vector_flat() {
746 let index = SearchIndex::new("small", "crdt:small:")
747 .vector_flat("embedding", 384, DistanceMetric::InnerProduct);
748
749 let args = index.to_ft_create_args();
750 let cmd = format!("FT.CREATE {}", args.join(" "));
751
752 assert!(cmd.contains("VECTOR FLAT"));
753 assert!(cmd.contains("DIM 384"));
754 assert!(cmd.contains("DISTANCE_METRIC IP"));
755 }
756
757 #[test]
758 fn test_vector_custom_path() {
759 let index = SearchIndex::new("docs", "crdt:docs:").vector_hnsw_at(
760 "embedding",
761 "$.metadata.vector",
762 512,
763 DistanceMetric::Cosine,
764 );
765
766 let args = index.to_ft_create_args();
767 let cmd = format!("FT.CREATE {}", args.join(" "));
768
769 assert!(cmd.contains("$.metadata.vector AS embedding VECTOR HNSW"));
770 }
771
772 #[test]
773 fn test_vector_params_nargs() {
774 let params = VectorParams::hnsw(1536, DistanceMetric::Cosine);
776 let args = params.to_schema_args();
777
778 assert_eq!(args[0], "6");
780 assert_eq!(args[1], "TYPE");
781 assert_eq!(args[2], "FLOAT32");
782 assert_eq!(args[3], "DIM");
783 assert_eq!(args[4], "1536");
784 assert_eq!(args[5], "DISTANCE_METRIC");
785 assert_eq!(args[6], "COSINE");
786 }
787
788 #[test]
789 fn test_vector_params_nargs_with_hnsw_options() {
790 let params = VectorParams::hnsw(1536, DistanceMetric::Cosine)
791 .with_m(24)
792 .with_ef_construction(300);
793 let args = params.to_schema_args();
794
795 assert_eq!(args[0], "10");
797 assert!(args.contains(&"M".to_string()));
798 assert!(args.contains(&"24".to_string()));
799 assert!(args.contains(&"EF_CONSTRUCTION".to_string()));
800 assert!(args.contains(&"300".to_string()));
801 }
802
803 #[test]
804 fn test_mixed_index_with_vectors() {
805 let index = SearchIndex::new("documents", "crdt:documents:")
806 .text_sortable("title")
807 .text("content")
808 .tag("category")
809 .numeric("created_at")
810 .vector_hnsw("embedding", 1536, DistanceMetric::Cosine);
811
812 let args = index.to_ft_create_args();
813 let cmd = format!("FT.CREATE {}", args.join(" "));
814
815 assert!(cmd.contains("AS title TEXT SORTABLE"));
817 assert!(cmd.contains("AS content TEXT"));
818 assert!(cmd.contains("AS category TAG"));
819 assert!(cmd.contains("AS created_at NUMERIC"));
820 assert!(cmd.contains("AS embedding VECTOR HNSW"));
821 }
822
823 #[test]
824 fn test_distance_metrics_display() {
825 assert_eq!(format!("{}", DistanceMetric::L2), "L2");
826 assert_eq!(format!("{}", DistanceMetric::InnerProduct), "IP");
827 assert_eq!(format!("{}", DistanceMetric::Cosine), "COSINE");
828 }
829
830 #[test]
831 fn test_vector_algorithm_display() {
832 assert_eq!(format!("{}", VectorAlgorithm::Hnsw), "HNSW");
833 assert_eq!(format!("{}", VectorAlgorithm::Flat), "FLAT");
834 }
835}