1use serde::{Deserialize, Serialize};
4use std::collections::HashMap;
5
6#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
8pub struct Field(pub u32);
9
10#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
12pub enum FieldType {
13 #[serde(rename = "text")]
15 Text,
16 #[serde(rename = "u64")]
18 U64,
19 #[serde(rename = "i64")]
21 I64,
22 #[serde(rename = "f64")]
24 F64,
25 #[serde(rename = "bytes")]
27 Bytes,
28 #[serde(rename = "sparse_vector")]
30 SparseVector,
31 #[serde(rename = "dense_vector")]
33 DenseVector,
34 #[serde(rename = "json")]
36 Json,
37 #[serde(rename = "binary_dense_vector")]
39 BinaryDenseVector,
40}
41
42#[derive(Debug, Clone, Serialize, Deserialize)]
44pub struct FieldEntry {
45 pub name: String,
46 pub field_type: FieldType,
47 pub indexed: bool,
48 pub stored: bool,
49 pub tokenizer: Option<String>,
51 #[serde(default)]
53 pub multi: bool,
54 #[serde(default, skip_serializing_if = "Option::is_none")]
56 pub positions: Option<PositionMode>,
57 #[serde(default, skip_serializing_if = "Option::is_none")]
59 pub sparse_vector_config: Option<crate::structures::SparseVectorConfig>,
60 #[serde(default, skip_serializing_if = "Option::is_none")]
62 pub dense_vector_config: Option<DenseVectorConfig>,
63 #[serde(default, skip_serializing_if = "Option::is_none")]
65 pub binary_dense_vector_config: Option<BinaryDenseVectorConfig>,
66 #[serde(default)]
69 pub fast: bool,
70 #[serde(default)]
72 pub primary_key: bool,
73 #[serde(default)]
77 pub reorder: bool,
78}
79
80#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
82#[serde(rename_all = "snake_case")]
83pub enum PositionMode {
84 Ordinal,
87 TokenPosition,
90 Full,
93}
94
95impl PositionMode {
96 pub fn tracks_ordinal(&self) -> bool {
98 matches!(self, PositionMode::Ordinal | PositionMode::Full)
99 }
100
101 pub fn tracks_token_position(&self) -> bool {
103 matches!(self, PositionMode::TokenPosition | PositionMode::Full)
104 }
105}
106
107#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
109#[serde(rename_all = "snake_case")]
110pub enum VectorIndexType {
111 Flat,
113 #[default]
115 RaBitQ,
116 IvfRaBitQ,
118 ScaNN,
120}
121
122#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
128#[serde(rename_all = "snake_case")]
129pub enum DenseVectorQuantization {
130 #[default]
132 F32,
133 F16,
135 UInt8,
137 Binary,
140}
141
142impl DenseVectorQuantization {
143 pub fn element_size(self) -> usize {
146 match self {
147 Self::F32 => 4,
148 Self::F16 => 2,
149 Self::UInt8 => 1,
150 Self::Binary => panic!("element_size() not valid for Binary; use dim.div_ceil(8)"),
151 }
152 }
153
154 pub fn tag(self) -> u8 {
156 match self {
157 Self::F32 => 0,
158 Self::F16 => 1,
159 Self::UInt8 => 2,
160 Self::Binary => 3,
161 }
162 }
163
164 pub fn from_tag(tag: u8) -> Option<Self> {
166 match tag {
167 0 => Some(Self::F32),
168 1 => Some(Self::F16),
169 2 => Some(Self::UInt8),
170 3 => Some(Self::Binary),
171 _ => None,
172 }
173 }
174}
175
176#[derive(Debug, Clone, Serialize, Deserialize)]
184pub struct DenseVectorConfig {
185 pub dim: usize,
187 #[serde(default)]
190 pub index_type: VectorIndexType,
191 #[serde(default)]
193 pub quantization: DenseVectorQuantization,
194 #[serde(default, skip_serializing_if = "Option::is_none")]
197 pub num_clusters: Option<usize>,
198 #[serde(default = "default_nprobe")]
200 pub nprobe: usize,
201 #[serde(default, skip_serializing_if = "Option::is_none")]
205 pub build_threshold: Option<usize>,
206 #[serde(default = "default_unit_norm")]
211 pub unit_norm: bool,
212 #[serde(default, skip_serializing_if = "Option::is_none")]
218 pub rabitq_bits: Option<u8>,
219 #[serde(default, skip_serializing_if = "Option::is_none")]
224 pub soar: Option<crate::structures::SoarConfig>,
225}
226
227fn default_nprobe() -> usize {
228 32
229}
230
231fn default_unit_norm() -> bool {
232 true
233}
234
235impl DenseVectorConfig {
236 pub fn new(dim: usize) -> Self {
237 Self {
238 dim,
239 index_type: VectorIndexType::RaBitQ,
240 quantization: DenseVectorQuantization::F32,
241 num_clusters: None,
242 nprobe: 32,
243 build_threshold: None,
244 unit_norm: true,
245 soar: None,
246 rabitq_bits: None,
247 }
248 }
249
250 pub fn with_ivf(dim: usize, num_clusters: Option<usize>, nprobe: usize) -> Self {
252 Self {
253 dim,
254 index_type: VectorIndexType::IvfRaBitQ,
255 quantization: DenseVectorQuantization::F32,
256 num_clusters,
257 nprobe,
258 build_threshold: None,
259 unit_norm: true,
260 soar: None,
261 rabitq_bits: None,
262 }
263 }
264
265 pub fn with_scann(dim: usize, num_clusters: Option<usize>, nprobe: usize) -> Self {
267 Self {
268 dim,
269 index_type: VectorIndexType::ScaNN,
270 quantization: DenseVectorQuantization::F32,
271 num_clusters,
272 nprobe,
273 build_threshold: None,
274 unit_norm: true,
275 soar: None,
276 rabitq_bits: None,
277 }
278 }
279
280 pub fn flat(dim: usize) -> Self {
282 Self {
283 dim,
284 index_type: VectorIndexType::Flat,
285 quantization: DenseVectorQuantization::F32,
286 num_clusters: None,
287 nprobe: 0,
288 build_threshold: None,
289 unit_norm: true,
290 soar: None,
291 rabitq_bits: None,
292 }
293 }
294
295 pub fn with_quantization(mut self, quantization: DenseVectorQuantization) -> Self {
297 self.quantization = quantization;
298 self
299 }
300
301 pub fn with_build_threshold(mut self, threshold: usize) -> Self {
303 self.build_threshold = Some(threshold);
304 self
305 }
306
307 pub fn with_unit_norm(mut self) -> Self {
309 self.unit_norm = true;
310 self
311 }
312
313 pub fn with_num_clusters(mut self, num_clusters: usize) -> Self {
315 self.num_clusters = Some(num_clusters);
316 self
317 }
318
319 pub fn with_rabitq_bits(mut self, bits: u8) -> Self {
321 self.rabitq_bits = Some(bits.clamp(1, 8));
322 self
323 }
324
325 pub fn with_soar(mut self, soar: crate::structures::SoarConfig) -> Self {
327 self.soar = Some(soar);
328 self
329 }
330
331 pub fn uses_ivf(&self) -> bool {
333 matches!(
334 self.index_type,
335 VectorIndexType::IvfRaBitQ | VectorIndexType::ScaNN
336 )
337 }
338
339 pub fn uses_scann(&self) -> bool {
341 self.index_type == VectorIndexType::ScaNN
342 }
343
344 pub fn is_flat(&self) -> bool {
346 self.index_type == VectorIndexType::Flat
347 }
348
349 pub fn default_build_threshold(&self) -> usize {
351 self.build_threshold.unwrap_or(match self.index_type {
352 VectorIndexType::Flat => usize::MAX, VectorIndexType::RaBitQ => 1000,
354 VectorIndexType::IvfRaBitQ | VectorIndexType::ScaNN => 10000,
355 })
356 }
357
358 pub fn optimal_num_clusters(&self, num_vectors: usize) -> usize {
360 self.num_clusters.unwrap_or_else(|| {
361 let optimal = (num_vectors as f64).sqrt() as usize;
363 optimal.clamp(16, 4096)
364 })
365 }
366}
367
368#[derive(Debug, Clone, Serialize, Deserialize)]
374pub struct BinaryDenseVectorConfig {
375 pub dim: usize,
377 #[serde(default)]
381 pub index_type: BinaryIndexType,
382 #[serde(default, skip_serializing_if = "Option::is_none")]
384 pub num_clusters: Option<usize>,
385 #[serde(default = "default_nprobe")]
387 pub nprobe: usize,
388 #[serde(default, skip_serializing_if = "Option::is_none")]
391 pub build_threshold: Option<usize>,
392}
393
394#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
396#[serde(rename_all = "snake_case")]
397pub enum BinaryIndexType {
398 #[default]
400 Flat,
401 Ivf,
403}
404
405impl BinaryDenseVectorConfig {
406 pub fn new(dim: usize) -> Self {
407 assert!(
408 dim.is_multiple_of(8),
409 "BinaryDenseVector dimension must be a multiple of 8, got {dim}"
410 );
411 Self {
412 dim,
413 index_type: BinaryIndexType::Flat,
414 num_clusters: None,
415 nprobe: 32,
416 build_threshold: None,
417 }
418 }
419
420 pub fn with_ivf(mut self, num_clusters: Option<usize>, nprobe: usize) -> Self {
422 self.index_type = BinaryIndexType::Ivf;
423 self.num_clusters = num_clusters;
424 self.nprobe = nprobe;
425 self
426 }
427
428 pub fn with_build_threshold(mut self, threshold: usize) -> Self {
430 self.build_threshold = Some(threshold);
431 self
432 }
433
434 pub fn default_build_threshold(&self) -> usize {
436 self.build_threshold.unwrap_or(100_000)
437 }
438
439 pub fn optimal_num_clusters(&self, num_vectors: usize) -> usize {
441 self.num_clusters.unwrap_or_else(|| {
442 let optimal = (num_vectors as f64).sqrt() as usize;
443 optimal.clamp(16, 4096)
444 })
445 }
446
447 pub fn byte_len(&self) -> usize {
449 self.dim.div_ceil(8)
450 }
451}
452
453use super::query_field_router::QueryRouterRule;
454
455#[derive(Debug, Clone, Default, Serialize, Deserialize)]
457pub struct Schema {
458 fields: Vec<FieldEntry>,
459 name_to_field: HashMap<String, Field>,
460 #[serde(default)]
462 default_fields: Vec<Field>,
463 #[serde(default)]
465 query_routers: Vec<QueryRouterRule>,
466 #[serde(default)]
471 reorder_on_merge: bool,
472}
473
474impl Schema {
475 pub fn builder() -> SchemaBuilder {
476 SchemaBuilder::default()
477 }
478
479 pub fn get_field(&self, name: &str) -> Option<Field> {
480 self.name_to_field.get(name).copied()
481 }
482
483 pub fn get_field_entry(&self, field: Field) -> Option<&FieldEntry> {
484 self.fields.get(field.0 as usize)
485 }
486
487 pub fn get_field_name(&self, field: Field) -> Option<&str> {
488 self.fields.get(field.0 as usize).map(|e| e.name.as_str())
489 }
490
491 pub fn fields(&self) -> impl Iterator<Item = (Field, &FieldEntry)> {
492 self.fields
493 .iter()
494 .enumerate()
495 .map(|(i, e)| (Field(i as u32), e))
496 }
497
498 pub fn num_fields(&self) -> usize {
499 self.fields.len()
500 }
501
502 pub fn has_reorder_fields(&self) -> bool {
505 self.fields.iter().any(|e| e.reorder)
506 }
507
508 pub fn reorder_on_merge(&self) -> bool {
511 self.reorder_on_merge
512 }
513
514 pub fn default_fields(&self) -> &[Field] {
516 &self.default_fields
517 }
518
519 pub fn set_default_fields(&mut self, fields: Vec<Field>) {
521 self.default_fields = fields;
522 }
523
524 pub fn query_routers(&self) -> &[QueryRouterRule] {
526 &self.query_routers
527 }
528
529 pub fn set_query_routers(&mut self, rules: Vec<QueryRouterRule>) {
531 self.query_routers = rules;
532 }
533
534 pub fn primary_field(&self) -> Option<Field> {
536 self.fields
537 .iter()
538 .enumerate()
539 .find(|(_, e)| e.primary_key)
540 .map(|(i, _)| Field(i as u32))
541 }
542}
543
544#[derive(Debug, Default)]
546pub struct SchemaBuilder {
547 fields: Vec<FieldEntry>,
548 default_fields: Vec<String>,
549 query_routers: Vec<QueryRouterRule>,
550 reorder_on_merge: bool,
551}
552
553impl SchemaBuilder {
554 pub fn add_text_field(&mut self, name: &str, indexed: bool, stored: bool) -> Field {
555 self.add_field_with_tokenizer(
556 name,
557 FieldType::Text,
558 indexed,
559 stored,
560 Some("simple".to_string()),
561 )
562 }
563
564 pub fn add_text_field_with_tokenizer(
565 &mut self,
566 name: &str,
567 indexed: bool,
568 stored: bool,
569 tokenizer: &str,
570 ) -> Field {
571 self.add_field_with_tokenizer(
572 name,
573 FieldType::Text,
574 indexed,
575 stored,
576 Some(tokenizer.to_string()),
577 )
578 }
579
580 pub fn add_u64_field(&mut self, name: &str, indexed: bool, stored: bool) -> Field {
581 self.add_field(name, FieldType::U64, indexed, stored)
582 }
583
584 pub fn add_i64_field(&mut self, name: &str, indexed: bool, stored: bool) -> Field {
585 self.add_field(name, FieldType::I64, indexed, stored)
586 }
587
588 pub fn add_f64_field(&mut self, name: &str, indexed: bool, stored: bool) -> Field {
589 self.add_field(name, FieldType::F64, indexed, stored)
590 }
591
592 pub fn add_bytes_field(&mut self, name: &str, stored: bool) -> Field {
593 self.add_field(name, FieldType::Bytes, false, stored)
594 }
595
596 pub fn add_json_field(&mut self, name: &str, stored: bool) -> Field {
601 self.add_field(name, FieldType::Json, false, stored)
602 }
603
604 pub fn add_sparse_vector_field(&mut self, name: &str, indexed: bool, stored: bool) -> Field {
609 self.add_sparse_vector_field_with_config(
610 name,
611 indexed,
612 stored,
613 crate::structures::SparseVectorConfig::default(),
614 )
615 }
616
617 pub fn add_sparse_vector_field_with_config(
622 &mut self,
623 name: &str,
624 indexed: bool,
625 stored: bool,
626 config: crate::structures::SparseVectorConfig,
627 ) -> Field {
628 let field = Field(self.fields.len() as u32);
629 self.fields.push(FieldEntry {
630 name: name.to_string(),
631 field_type: FieldType::SparseVector,
632 indexed,
633 stored,
634 tokenizer: None,
635 multi: false,
636 positions: None,
637 sparse_vector_config: Some(config),
638 dense_vector_config: None,
639 binary_dense_vector_config: None,
640 fast: false,
641 primary_key: false,
642 reorder: false,
643 });
644 field
645 }
646
647 pub fn set_sparse_vector_config(
649 &mut self,
650 field: Field,
651 config: crate::structures::SparseVectorConfig,
652 ) {
653 if let Some(entry) = self.fields.get_mut(field.0 as usize) {
654 entry.sparse_vector_config = Some(config);
655 }
656 }
657
658 pub fn add_dense_vector_field(
663 &mut self,
664 name: &str,
665 dim: usize,
666 indexed: bool,
667 stored: bool,
668 ) -> Field {
669 self.add_dense_vector_field_with_config(name, indexed, stored, DenseVectorConfig::new(dim))
670 }
671
672 pub fn add_dense_vector_field_with_config(
674 &mut self,
675 name: &str,
676 indexed: bool,
677 stored: bool,
678 config: DenseVectorConfig,
679 ) -> Field {
680 let field = Field(self.fields.len() as u32);
681 self.fields.push(FieldEntry {
682 name: name.to_string(),
683 field_type: FieldType::DenseVector,
684 indexed,
685 stored,
686 tokenizer: None,
687 multi: false,
688 positions: None,
689 sparse_vector_config: None,
690 dense_vector_config: Some(config),
691 binary_dense_vector_config: None,
692 fast: false,
693 primary_key: false,
694 reorder: false,
695 });
696 field
697 }
698
699 pub fn add_binary_dense_vector_field(
704 &mut self,
705 name: &str,
706 dim: usize,
707 indexed: bool,
708 stored: bool,
709 ) -> Field {
710 self.add_binary_dense_vector_field_with_config(
711 name,
712 indexed,
713 stored,
714 BinaryDenseVectorConfig::new(dim),
715 )
716 }
717
718 pub fn add_binary_dense_vector_field_with_config(
720 &mut self,
721 name: &str,
722 indexed: bool,
723 stored: bool,
724 config: BinaryDenseVectorConfig,
725 ) -> Field {
726 let field = Field(self.fields.len() as u32);
727 self.fields.push(FieldEntry {
728 name: name.to_string(),
729 field_type: FieldType::BinaryDenseVector,
730 indexed,
731 stored,
732 tokenizer: None,
733 multi: false,
734 positions: None,
735 sparse_vector_config: None,
736 dense_vector_config: None,
737 binary_dense_vector_config: Some(config),
738 fast: false,
739 primary_key: false,
740 reorder: false,
741 });
742 field
743 }
744
745 fn add_field(
746 &mut self,
747 name: &str,
748 field_type: FieldType,
749 indexed: bool,
750 stored: bool,
751 ) -> Field {
752 self.add_field_with_tokenizer(name, field_type, indexed, stored, None)
753 }
754
755 fn add_field_with_tokenizer(
756 &mut self,
757 name: &str,
758 field_type: FieldType,
759 indexed: bool,
760 stored: bool,
761 tokenizer: Option<String>,
762 ) -> Field {
763 self.add_field_full(name, field_type, indexed, stored, tokenizer, false)
764 }
765
766 fn add_field_full(
767 &mut self,
768 name: &str,
769 field_type: FieldType,
770 indexed: bool,
771 stored: bool,
772 tokenizer: Option<String>,
773 multi: bool,
774 ) -> Field {
775 let field = Field(self.fields.len() as u32);
776 self.fields.push(FieldEntry {
777 name: name.to_string(),
778 field_type,
779 indexed,
780 stored,
781 tokenizer,
782 multi,
783 positions: None,
784 sparse_vector_config: None,
785 dense_vector_config: None,
786 binary_dense_vector_config: None,
787 fast: false,
788 primary_key: false,
789 reorder: false,
790 });
791 field
792 }
793
794 pub fn set_multi(&mut self, field: Field, multi: bool) {
796 if let Some(entry) = self.fields.get_mut(field.0 as usize) {
797 entry.multi = multi;
798 }
799 }
800
801 pub fn set_fast(&mut self, field: Field, fast: bool) {
804 if let Some(entry) = self.fields.get_mut(field.0 as usize) {
805 entry.fast = fast;
806 }
807 }
808
809 pub fn set_primary_key(&mut self, field: Field) {
811 if let Some(entry) = self.fields.get_mut(field.0 as usize) {
812 entry.primary_key = true;
813 }
814 }
815
816 pub fn set_reorder(&mut self, field: Field, reorder: bool) {
818 if let Some(entry) = self.fields.get_mut(field.0 as usize) {
819 entry.reorder = reorder;
820 }
821 }
822
823 pub fn set_reorder_on_merge(&mut self, on: bool) {
826 self.reorder_on_merge = on;
827 }
828
829 pub fn set_positions(&mut self, field: Field, mode: PositionMode) {
831 if let Some(entry) = self.fields.get_mut(field.0 as usize) {
832 entry.positions = Some(mode);
833 }
834 }
835
836 pub fn set_default_fields(&mut self, field_names: Vec<String>) {
838 self.default_fields = field_names;
839 }
840
841 pub fn set_query_routers(&mut self, rules: Vec<QueryRouterRule>) {
843 self.query_routers = rules;
844 }
845
846 pub fn build(self) -> Schema {
847 let mut name_to_field = HashMap::new();
848 for (i, entry) in self.fields.iter().enumerate() {
849 name_to_field.insert(entry.name.clone(), Field(i as u32));
850 }
851
852 let default_fields: Vec<Field> = self
854 .default_fields
855 .iter()
856 .filter_map(|name| name_to_field.get(name).copied())
857 .collect();
858
859 Schema {
860 fields: self.fields,
861 name_to_field,
862 default_fields,
863 query_routers: self.query_routers,
864 reorder_on_merge: self.reorder_on_merge,
865 }
866 }
867}
868
869#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
871pub enum FieldValue {
872 #[serde(rename = "text")]
873 Text(String),
874 #[serde(rename = "u64")]
875 U64(u64),
876 #[serde(rename = "i64")]
877 I64(i64),
878 #[serde(rename = "f64")]
879 F64(f64),
880 #[serde(rename = "bytes")]
881 Bytes(Vec<u8>),
882 #[serde(rename = "sparse_vector")]
884 SparseVector(Vec<(u32, f32)>),
885 #[serde(rename = "dense_vector")]
887 DenseVector(Vec<f32>),
888 #[serde(rename = "json")]
890 Json(serde_json::Value),
891 #[serde(rename = "binary_dense_vector")]
893 BinaryDenseVector(Vec<u8>),
894}
895
896impl FieldValue {
897 pub fn as_text(&self) -> Option<&str> {
898 match self {
899 FieldValue::Text(s) => Some(s),
900 _ => None,
901 }
902 }
903
904 pub fn as_u64(&self) -> Option<u64> {
905 match self {
906 FieldValue::U64(v) => Some(*v),
907 _ => None,
908 }
909 }
910
911 pub fn as_i64(&self) -> Option<i64> {
912 match self {
913 FieldValue::I64(v) => Some(*v),
914 _ => None,
915 }
916 }
917
918 pub fn as_f64(&self) -> Option<f64> {
919 match self {
920 FieldValue::F64(v) => Some(*v),
921 _ => None,
922 }
923 }
924
925 pub fn as_bytes(&self) -> Option<&[u8]> {
926 match self {
927 FieldValue::Bytes(b) => Some(b),
928 _ => None,
929 }
930 }
931
932 pub fn as_sparse_vector(&self) -> Option<&[(u32, f32)]> {
933 match self {
934 FieldValue::SparseVector(entries) => Some(entries),
935 _ => None,
936 }
937 }
938
939 pub fn as_dense_vector(&self) -> Option<&[f32]> {
940 match self {
941 FieldValue::DenseVector(v) => Some(v),
942 _ => None,
943 }
944 }
945
946 pub fn as_json(&self) -> Option<&serde_json::Value> {
947 match self {
948 FieldValue::Json(v) => Some(v),
949 _ => None,
950 }
951 }
952
953 pub fn as_binary_dense_vector(&self) -> Option<&[u8]> {
954 match self {
955 FieldValue::BinaryDenseVector(v) => Some(v),
956 _ => None,
957 }
958 }
959}
960
961#[derive(Debug, Clone, Default, Serialize, Deserialize)]
963pub struct Document {
964 field_values: Vec<(Field, FieldValue)>,
965}
966
967impl Document {
968 pub fn new() -> Self {
969 Self::default()
970 }
971
972 pub fn add_text(&mut self, field: Field, value: impl Into<String>) {
973 self.field_values
974 .push((field, FieldValue::Text(value.into())));
975 }
976
977 pub fn add_u64(&mut self, field: Field, value: u64) {
978 self.field_values.push((field, FieldValue::U64(value)));
979 }
980
981 pub fn add_i64(&mut self, field: Field, value: i64) {
982 self.field_values.push((field, FieldValue::I64(value)));
983 }
984
985 pub fn add_f64(&mut self, field: Field, value: f64) {
986 self.field_values.push((field, FieldValue::F64(value)));
987 }
988
989 pub fn add_bytes(&mut self, field: Field, value: Vec<u8>) {
990 self.field_values.push((field, FieldValue::Bytes(value)));
991 }
992
993 pub fn add_sparse_vector(&mut self, field: Field, entries: Vec<(u32, f32)>) {
994 self.field_values
995 .push((field, FieldValue::SparseVector(entries)));
996 }
997
998 pub fn add_dense_vector(&mut self, field: Field, values: Vec<f32>) {
999 self.field_values
1000 .push((field, FieldValue::DenseVector(values)));
1001 }
1002
1003 pub fn add_json(&mut self, field: Field, value: serde_json::Value) {
1004 self.field_values.push((field, FieldValue::Json(value)));
1005 }
1006
1007 pub fn add_binary_dense_vector(&mut self, field: Field, values: Vec<u8>) {
1008 self.field_values
1009 .push((field, FieldValue::BinaryDenseVector(values)));
1010 }
1011
1012 pub fn get_first(&self, field: Field) -> Option<&FieldValue> {
1013 self.field_values
1014 .iter()
1015 .find(|(f, _)| *f == field)
1016 .map(|(_, v)| v)
1017 }
1018
1019 pub fn get_all(&self, field: Field) -> impl Iterator<Item = &FieldValue> {
1020 self.field_values
1021 .iter()
1022 .filter(move |(f, _)| *f == field)
1023 .map(|(_, v)| v)
1024 }
1025
1026 pub fn field_values(&self) -> &[(Field, FieldValue)] {
1027 &self.field_values
1028 }
1029
1030 pub fn filter_stored(&self, schema: &Schema) -> Document {
1032 Document {
1033 field_values: self
1034 .field_values
1035 .iter()
1036 .filter(|(field, _)| {
1037 schema
1038 .get_field_entry(*field)
1039 .is_some_and(|entry| entry.stored)
1040 })
1041 .cloned()
1042 .collect(),
1043 }
1044 }
1045
1046 pub fn to_json(&self, schema: &Schema) -> serde_json::Value {
1052 use std::collections::HashMap;
1053
1054 let mut field_values_map: HashMap<Field, (String, bool, Vec<serde_json::Value>)> =
1056 HashMap::new();
1057
1058 for (field, value) in &self.field_values {
1059 if let Some(entry) = schema.get_field_entry(*field) {
1060 let json_value = match value {
1061 FieldValue::Text(s) => serde_json::Value::String(s.clone()),
1062 FieldValue::U64(n) => serde_json::Value::Number((*n).into()),
1063 FieldValue::I64(n) => serde_json::Value::Number((*n).into()),
1064 FieldValue::F64(n) => serde_json::json!(n),
1065 FieldValue::Bytes(b) => {
1066 use base64::Engine;
1067 serde_json::Value::String(
1068 base64::engine::general_purpose::STANDARD.encode(b),
1069 )
1070 }
1071 FieldValue::SparseVector(entries) => {
1072 let indices: Vec<u32> = entries.iter().map(|(i, _)| *i).collect();
1073 let values: Vec<f32> = entries.iter().map(|(_, v)| *v).collect();
1074 serde_json::json!({
1075 "indices": indices,
1076 "values": values
1077 })
1078 }
1079 FieldValue::DenseVector(values) => {
1080 serde_json::json!(values)
1081 }
1082 FieldValue::Json(v) => v.clone(),
1083 FieldValue::BinaryDenseVector(b) => {
1084 use base64::Engine;
1085 serde_json::Value::String(
1086 base64::engine::general_purpose::STANDARD.encode(b),
1087 )
1088 }
1089 };
1090 field_values_map
1091 .entry(*field)
1092 .or_insert_with(|| (entry.name.clone(), entry.multi, Vec::new()))
1093 .2
1094 .push(json_value);
1095 }
1096 }
1097
1098 let mut map = serde_json::Map::new();
1100 for (_field, (name, is_multi, values)) in field_values_map {
1101 let json_value = if is_multi || values.len() > 1 {
1102 serde_json::Value::Array(values)
1103 } else {
1104 values.into_iter().next().unwrap()
1105 };
1106 map.insert(name, json_value);
1107 }
1108
1109 serde_json::Value::Object(map)
1110 }
1111
1112 pub fn from_json(json: &serde_json::Value, schema: &Schema) -> Option<Self> {
1121 let obj = json.as_object()?;
1122 let mut doc = Document::new();
1123
1124 for (key, value) in obj {
1125 if let Some(field) = schema.get_field(key) {
1126 let field_entry = schema.get_field_entry(field)?;
1127 Self::add_json_value(&mut doc, field, &field_entry.field_type, value);
1128 }
1129 }
1130
1131 Some(doc)
1132 }
1133
1134 fn add_json_value(
1136 doc: &mut Document,
1137 field: Field,
1138 field_type: &FieldType,
1139 value: &serde_json::Value,
1140 ) {
1141 match value {
1142 serde_json::Value::String(s) => {
1143 if matches!(field_type, FieldType::Text) {
1144 doc.add_text(field, s.clone());
1145 }
1146 }
1147 serde_json::Value::Number(n) => {
1148 match field_type {
1149 FieldType::I64 => {
1150 if let Some(i) = n.as_i64() {
1151 doc.add_i64(field, i);
1152 }
1153 }
1154 FieldType::U64 => {
1155 if let Some(u) = n.as_u64() {
1156 doc.add_u64(field, u);
1157 } else if let Some(i) = n.as_i64() {
1158 if i >= 0 {
1160 doc.add_u64(field, i as u64);
1161 }
1162 }
1163 }
1164 FieldType::F64 => {
1165 if let Some(f) = n.as_f64() {
1166 doc.add_f64(field, f);
1167 }
1168 }
1169 _ => {}
1170 }
1171 }
1172 serde_json::Value::Array(arr) => {
1174 for item in arr {
1175 Self::add_json_value(doc, field, field_type, item);
1176 }
1177 }
1178 serde_json::Value::Object(obj) if matches!(field_type, FieldType::SparseVector) => {
1180 if let (Some(indices_val), Some(values_val)) =
1181 (obj.get("indices"), obj.get("values"))
1182 {
1183 let indices: Vec<u32> = indices_val
1184 .as_array()
1185 .map(|arr| {
1186 arr.iter()
1187 .filter_map(|v| v.as_u64().map(|n| n as u32))
1188 .collect()
1189 })
1190 .unwrap_or_default();
1191 let values: Vec<f32> = values_val
1192 .as_array()
1193 .map(|arr| {
1194 arr.iter()
1195 .filter_map(|v| v.as_f64().map(|n| n as f32))
1196 .collect()
1197 })
1198 .unwrap_or_default();
1199 if indices.len() == values.len() {
1200 let entries: Vec<(u32, f32)> = indices.into_iter().zip(values).collect();
1201 doc.add_sparse_vector(field, entries);
1202 }
1203 }
1204 }
1205 _ if matches!(field_type, FieldType::Json) => {
1207 doc.add_json(field, value.clone());
1208 }
1209 serde_json::Value::Object(_) => {}
1210 _ => {}
1211 }
1212 }
1213}
1214
1215#[cfg(test)]
1216mod tests {
1217 use super::*;
1218
1219 #[test]
1220 fn test_schema_builder() {
1221 let mut builder = Schema::builder();
1222 let title = builder.add_text_field("title", true, true);
1223 let body = builder.add_text_field("body", true, false);
1224 let count = builder.add_u64_field("count", true, true);
1225 let schema = builder.build();
1226
1227 assert_eq!(schema.get_field("title"), Some(title));
1228 assert_eq!(schema.get_field("body"), Some(body));
1229 assert_eq!(schema.get_field("count"), Some(count));
1230 assert_eq!(schema.get_field("nonexistent"), None);
1231 }
1232
1233 #[test]
1234 fn test_document() {
1235 let mut builder = Schema::builder();
1236 let title = builder.add_text_field("title", true, true);
1237 let count = builder.add_u64_field("count", true, true);
1238 let _schema = builder.build();
1239
1240 let mut doc = Document::new();
1241 doc.add_text(title, "Hello World");
1242 doc.add_u64(count, 42);
1243
1244 assert_eq!(doc.get_first(title).unwrap().as_text(), Some("Hello World"));
1245 assert_eq!(doc.get_first(count).unwrap().as_u64(), Some(42));
1246 }
1247
1248 #[test]
1249 fn test_document_serialization() {
1250 let mut builder = Schema::builder();
1251 let title = builder.add_text_field("title", true, true);
1252 let count = builder.add_u64_field("count", true, true);
1253 let _schema = builder.build();
1254
1255 let mut doc = Document::new();
1256 doc.add_text(title, "Hello World");
1257 doc.add_u64(count, 42);
1258
1259 let json = serde_json::to_string(&doc).unwrap();
1261 println!("Serialized doc: {}", json);
1262
1263 let doc2: Document = serde_json::from_str(&json).unwrap();
1265 assert_eq!(
1266 doc2.field_values().len(),
1267 2,
1268 "Should have 2 field values after deserialization"
1269 );
1270 assert_eq!(
1271 doc2.get_first(title).unwrap().as_text(),
1272 Some("Hello World")
1273 );
1274 assert_eq!(doc2.get_first(count).unwrap().as_u64(), Some(42));
1275 }
1276
1277 #[test]
1278 fn test_multivalue_field() {
1279 let mut builder = Schema::builder();
1280 let uris = builder.add_text_field("uris", true, true);
1281 let title = builder.add_text_field("title", true, true);
1282 let schema = builder.build();
1283
1284 let mut doc = Document::new();
1286 doc.add_text(uris, "one");
1287 doc.add_text(uris, "two");
1288 doc.add_text(title, "Test Document");
1289
1290 assert_eq!(doc.get_first(uris).unwrap().as_text(), Some("one"));
1292
1293 let all_uris: Vec<_> = doc.get_all(uris).collect();
1295 assert_eq!(all_uris.len(), 2);
1296 assert_eq!(all_uris[0].as_text(), Some("one"));
1297 assert_eq!(all_uris[1].as_text(), Some("two"));
1298
1299 let json = doc.to_json(&schema);
1301 let uris_json = json.get("uris").unwrap();
1302 assert!(uris_json.is_array(), "Multi-value field should be an array");
1303 let uris_arr = uris_json.as_array().unwrap();
1304 assert_eq!(uris_arr.len(), 2);
1305 assert_eq!(uris_arr[0].as_str(), Some("one"));
1306 assert_eq!(uris_arr[1].as_str(), Some("two"));
1307
1308 let title_json = json.get("title").unwrap();
1310 assert!(
1311 title_json.is_string(),
1312 "Single-value field should be a string"
1313 );
1314 assert_eq!(title_json.as_str(), Some("Test Document"));
1315 }
1316
1317 #[test]
1318 fn test_multivalue_from_json() {
1319 let mut builder = Schema::builder();
1320 let uris = builder.add_text_field("uris", true, true);
1321 let title = builder.add_text_field("title", true, true);
1322 let schema = builder.build();
1323
1324 let json = serde_json::json!({
1326 "uris": ["one", "two"],
1327 "title": "Test Document"
1328 });
1329
1330 let doc = Document::from_json(&json, &schema).unwrap();
1332
1333 let all_uris: Vec<_> = doc.get_all(uris).collect();
1335 assert_eq!(all_uris.len(), 2);
1336 assert_eq!(all_uris[0].as_text(), Some("one"));
1337 assert_eq!(all_uris[1].as_text(), Some("two"));
1338
1339 assert_eq!(
1341 doc.get_first(title).unwrap().as_text(),
1342 Some("Test Document")
1343 );
1344
1345 let json_out = doc.to_json(&schema);
1347 let uris_out = json_out.get("uris").unwrap().as_array().unwrap();
1348 assert_eq!(uris_out.len(), 2);
1349 assert_eq!(uris_out[0].as_str(), Some("one"));
1350 assert_eq!(uris_out[1].as_str(), Some("two"));
1351 }
1352
1353 #[test]
1354 fn test_multi_attribute_forces_array() {
1355 let mut builder = Schema::builder();
1358 let uris = builder.add_text_field("uris", true, true);
1359 builder.set_multi(uris, true); let title = builder.add_text_field("title", true, true);
1361 let schema = builder.build();
1362
1363 assert!(schema.get_field_entry(uris).unwrap().multi);
1365 assert!(!schema.get_field_entry(title).unwrap().multi);
1366
1367 let mut doc = Document::new();
1369 doc.add_text(uris, "only_one");
1370 doc.add_text(title, "Test Document");
1371
1372 let json = doc.to_json(&schema);
1374
1375 let uris_json = json.get("uris").unwrap();
1376 assert!(
1377 uris_json.is_array(),
1378 "Multi field should be array even with single value"
1379 );
1380 let uris_arr = uris_json.as_array().unwrap();
1381 assert_eq!(uris_arr.len(), 1);
1382 assert_eq!(uris_arr[0].as_str(), Some("only_one"));
1383
1384 let title_json = json.get("title").unwrap();
1386 assert!(
1387 title_json.is_string(),
1388 "Non-multi single-value field should be a string"
1389 );
1390 assert_eq!(title_json.as_str(), Some("Test Document"));
1391 }
1392
1393 #[test]
1394 fn test_sparse_vector_field() {
1395 let mut builder = Schema::builder();
1396 let embedding = builder.add_sparse_vector_field("embedding", true, true);
1397 let title = builder.add_text_field("title", true, true);
1398 let schema = builder.build();
1399
1400 assert_eq!(schema.get_field("embedding"), Some(embedding));
1401 assert_eq!(
1402 schema.get_field_entry(embedding).unwrap().field_type,
1403 FieldType::SparseVector
1404 );
1405
1406 let mut doc = Document::new();
1408 doc.add_sparse_vector(embedding, vec![(0, 1.0), (5, 2.5), (10, 0.5)]);
1409 doc.add_text(title, "Test Document");
1410
1411 let entries = doc
1413 .get_first(embedding)
1414 .unwrap()
1415 .as_sparse_vector()
1416 .unwrap();
1417 assert_eq!(entries, &[(0, 1.0), (5, 2.5), (10, 0.5)]);
1418
1419 let json = doc.to_json(&schema);
1421 let embedding_json = json.get("embedding").unwrap();
1422 assert!(embedding_json.is_object());
1423 assert_eq!(
1424 embedding_json
1425 .get("indices")
1426 .unwrap()
1427 .as_array()
1428 .unwrap()
1429 .len(),
1430 3
1431 );
1432
1433 let doc2 = Document::from_json(&json, &schema).unwrap();
1435 let entries2 = doc2
1436 .get_first(embedding)
1437 .unwrap()
1438 .as_sparse_vector()
1439 .unwrap();
1440 assert_eq!(entries2[0].0, 0);
1441 assert!((entries2[0].1 - 1.0).abs() < 1e-6);
1442 assert_eq!(entries2[1].0, 5);
1443 assert!((entries2[1].1 - 2.5).abs() < 1e-6);
1444 assert_eq!(entries2[2].0, 10);
1445 assert!((entries2[2].1 - 0.5).abs() < 1e-6);
1446 }
1447
1448 #[test]
1449 fn test_json_field() {
1450 let mut builder = Schema::builder();
1451 let metadata = builder.add_json_field("metadata", true);
1452 let title = builder.add_text_field("title", true, true);
1453 let schema = builder.build();
1454
1455 assert_eq!(schema.get_field("metadata"), Some(metadata));
1456 assert_eq!(
1457 schema.get_field_entry(metadata).unwrap().field_type,
1458 FieldType::Json
1459 );
1460 assert!(!schema.get_field_entry(metadata).unwrap().indexed);
1462 assert!(schema.get_field_entry(metadata).unwrap().stored);
1463
1464 let json_value = serde_json::json!({
1466 "author": "John Doe",
1467 "tags": ["rust", "search"],
1468 "nested": {"key": "value"}
1469 });
1470 let mut doc = Document::new();
1471 doc.add_json(metadata, json_value.clone());
1472 doc.add_text(title, "Test Document");
1473
1474 let stored_json = doc.get_first(metadata).unwrap().as_json().unwrap();
1476 assert_eq!(stored_json, &json_value);
1477 assert_eq!(
1478 stored_json.get("author").unwrap().as_str(),
1479 Some("John Doe")
1480 );
1481
1482 let doc_json = doc.to_json(&schema);
1484 let metadata_out = doc_json.get("metadata").unwrap();
1485 assert_eq!(metadata_out, &json_value);
1486
1487 let doc2 = Document::from_json(&doc_json, &schema).unwrap();
1489 let stored_json2 = doc2.get_first(metadata).unwrap().as_json().unwrap();
1490 assert_eq!(stored_json2, &json_value);
1491 }
1492
1493 #[test]
1494 fn test_json_field_various_types() {
1495 let mut builder = Schema::builder();
1496 let data = builder.add_json_field("data", true);
1497 let _schema = builder.build();
1498
1499 let arr_value = serde_json::json!([1, 2, 3, "four", null]);
1501 let mut doc = Document::new();
1502 doc.add_json(data, arr_value.clone());
1503 assert_eq!(doc.get_first(data).unwrap().as_json().unwrap(), &arr_value);
1504
1505 let str_value = serde_json::json!("just a string");
1507 let mut doc2 = Document::new();
1508 doc2.add_json(data, str_value.clone());
1509 assert_eq!(doc2.get_first(data).unwrap().as_json().unwrap(), &str_value);
1510
1511 let num_value = serde_json::json!(42.5);
1513 let mut doc3 = Document::new();
1514 doc3.add_json(data, num_value.clone());
1515 assert_eq!(doc3.get_first(data).unwrap().as_json().unwrap(), &num_value);
1516
1517 let null_value = serde_json::Value::Null;
1519 let mut doc4 = Document::new();
1520 doc4.add_json(data, null_value.clone());
1521 assert_eq!(
1522 doc4.get_first(data).unwrap().as_json().unwrap(),
1523 &null_value
1524 );
1525
1526 let bool_value = serde_json::json!(true);
1528 let mut doc5 = Document::new();
1529 doc5.add_json(data, bool_value.clone());
1530 assert_eq!(
1531 doc5.get_first(data).unwrap().as_json().unwrap(),
1532 &bool_value
1533 );
1534 }
1535}