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]
117 IvfPq,
118 Tq,
122 IvfTq,
126}
127
128#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
134#[serde(rename_all = "snake_case")]
135pub enum IvfRoutingMode {
136 #[default]
139 Auto,
140 Flat,
142 TwoLevel,
144 Hnsw,
146}
147
148#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
154#[serde(rename_all = "snake_case")]
155pub enum DenseVectorQuantization {
156 #[default]
158 F32,
159 F16,
161 UInt8,
163 Binary,
166}
167
168impl DenseVectorQuantization {
169 pub fn element_size(self) -> usize {
172 match self {
173 Self::F32 => 4,
174 Self::F16 => 2,
175 Self::UInt8 => 1,
176 Self::Binary => panic!("element_size() not valid for Binary; use dim.div_ceil(8)"),
177 }
178 }
179
180 pub fn tag(self) -> u8 {
182 match self {
183 Self::F32 => 0,
184 Self::F16 => 1,
185 Self::UInt8 => 2,
186 Self::Binary => 3,
187 }
188 }
189
190 pub fn from_tag(tag: u8) -> Option<Self> {
192 match tag {
193 0 => Some(Self::F32),
194 1 => Some(Self::F16),
195 2 => Some(Self::UInt8),
196 3 => Some(Self::Binary),
197 _ => None,
198 }
199 }
200}
201
202#[derive(Debug, Clone, Serialize, Deserialize)]
212pub struct DenseVectorConfig {
213 pub dim: usize,
215 #[serde(default)]
218 pub index_type: VectorIndexType,
219 #[serde(default)]
221 pub quantization: DenseVectorQuantization,
222 #[serde(default, skip_serializing_if = "Option::is_none")]
226 pub num_clusters: Option<usize>,
227 #[serde(default)]
230 pub ivf_routing: IvfRoutingMode,
231 #[serde(default = "default_nprobe")]
233 pub nprobe: usize,
234 #[serde(default = "default_unit_norm")]
239 pub unit_norm: bool,
240 #[serde(default, skip_serializing_if = "Option::is_none")]
245 pub soar: Option<crate::structures::SoarConfig>,
246}
247
248fn default_nprobe() -> usize {
249 64
250}
251
252fn default_unit_norm() -> bool {
253 true
254}
255
256impl DenseVectorConfig {
257 pub fn new(dim: usize) -> Self {
258 Self {
259 dim,
260 index_type: VectorIndexType::IvfPq,
261 quantization: DenseVectorQuantization::F32,
262 num_clusters: None,
263 ivf_routing: IvfRoutingMode::Auto,
264 nprobe: 64,
265 unit_norm: true,
266 soar: None,
267 }
268 }
269
270 pub fn with_ivf_pq(dim: usize, num_clusters: Option<usize>, nprobe: usize) -> Self {
272 Self {
273 dim,
274 index_type: VectorIndexType::IvfPq,
275 quantization: DenseVectorQuantization::F32,
276 num_clusters,
277 ivf_routing: IvfRoutingMode::Auto,
278 nprobe,
279 unit_norm: true,
280 soar: None,
281 }
282 }
283
284 pub fn flat(dim: usize) -> Self {
286 Self {
287 dim,
288 index_type: VectorIndexType::Flat,
289 quantization: DenseVectorQuantization::F32,
290 num_clusters: None,
291 ivf_routing: IvfRoutingMode::Auto,
292 nprobe: 0,
293 unit_norm: true,
294 soar: None,
295 }
296 }
297
298 pub fn tq(dim: usize) -> Self {
300 Self {
301 dim,
302 index_type: VectorIndexType::Tq,
303 quantization: DenseVectorQuantization::F32,
304 num_clusters: None,
305 ivf_routing: IvfRoutingMode::Flat,
306 nprobe: 0,
307 unit_norm: true,
308 soar: None,
309 }
310 }
311
312 pub fn ivf_tq(dim: usize, num_clusters: Option<usize>, nprobe: usize) -> Self {
314 Self {
315 dim,
316 index_type: VectorIndexType::IvfTq,
317 quantization: DenseVectorQuantization::F32,
318 num_clusters,
319 ivf_routing: IvfRoutingMode::Auto,
320 nprobe,
321 unit_norm: true,
322 soar: None,
323 }
324 }
325
326 pub fn with_quantization(mut self, quantization: DenseVectorQuantization) -> Self {
328 self.quantization = quantization;
329 self
330 }
331
332 pub fn with_unit_norm(mut self) -> Self {
334 self.unit_norm = true;
335 self
336 }
337
338 pub fn with_num_clusters(mut self, num_clusters: usize) -> Self {
340 self.num_clusters = Some(num_clusters);
341 self
342 }
343
344 pub fn with_ivf_routing(mut self, routing: IvfRoutingMode) -> Self {
346 self.ivf_routing = routing;
347 self
348 }
349 pub fn with_soar(mut self, soar: crate::structures::SoarConfig) -> Self {
351 self.soar = Some(soar);
352 self
353 }
354
355 pub fn uses_ivf(&self) -> bool {
357 matches!(
358 self.index_type,
359 VectorIndexType::IvfPq | VectorIndexType::IvfTq
360 )
361 }
362
363 pub fn is_flat(&self) -> bool {
365 self.index_type == VectorIndexType::Flat
366 }
367
368 pub fn optimal_num_clusters(&self, num_vectors: usize) -> usize {
370 self.num_clusters.unwrap_or_else(|| {
371 let optimal = 8.0 * (num_vectors as f64).sqrt();
375 (optimal as usize).clamp(16, 1_048_576)
376 })
377 }
378}
379
380#[derive(Debug, Clone, Serialize, Deserialize)]
386pub struct BinaryDenseVectorConfig {
387 pub dim: usize,
389 #[serde(default)]
393 pub index_type: BinaryIndexType,
394 #[serde(default, skip_serializing_if = "Option::is_none")]
396 pub num_clusters: Option<usize>,
397 #[serde(default)]
400 pub ivf_routing: IvfRoutingMode,
401 #[serde(default = "default_nprobe")]
403 pub nprobe: usize,
404}
405
406#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
408#[serde(rename_all = "snake_case")]
409pub enum BinaryIndexType {
410 Flat,
412 #[default]
414 Ivf,
415}
416
417impl BinaryDenseVectorConfig {
418 pub fn new(dim: usize) -> Self {
419 assert!(
420 dim.is_multiple_of(8),
421 "BinaryDenseVector dimension must be a multiple of 8, got {dim}"
422 );
423 Self {
424 dim,
425 index_type: BinaryIndexType::Ivf,
426 num_clusters: None,
427 ivf_routing: IvfRoutingMode::Auto,
428 nprobe: 64,
429 }
430 }
431
432 pub fn with_ivf(mut self, num_clusters: Option<usize>, nprobe: usize) -> Self {
434 self.index_type = BinaryIndexType::Ivf;
435 self.num_clusters = num_clusters;
436 self.nprobe = nprobe;
437 self
438 }
439
440 pub fn with_ivf_routing(mut self, routing: IvfRoutingMode) -> Self {
442 self.ivf_routing = routing;
443 self
444 }
445
446 pub fn optimal_num_clusters(&self, num_vectors: usize) -> usize {
448 self.num_clusters.unwrap_or_else(|| {
449 let optimal = 8.0 * (num_vectors as f64).sqrt();
450 (optimal as usize).clamp(16, 1_048_576)
451 })
452 }
453
454 pub fn byte_len(&self) -> usize {
456 self.dim.div_ceil(8)
457 }
458}
459
460use super::query_field_router::QueryRouterRule;
461
462#[derive(Debug, Clone, Default, Serialize, Deserialize)]
464pub struct Schema {
465 fields: Vec<FieldEntry>,
466 name_to_field: HashMap<String, Field>,
467 #[serde(default)]
469 default_fields: Vec<Field>,
470 #[serde(default)]
472 query_routers: Vec<QueryRouterRule>,
473 #[serde(default)]
478 reorder_on_merge: bool,
479 #[serde(default)]
483 index_name: String,
484}
485
486impl Schema {
487 pub fn builder() -> SchemaBuilder {
488 SchemaBuilder::default()
489 }
490
491 pub fn get_field(&self, name: &str) -> Option<Field> {
492 self.name_to_field.get(name).copied()
493 }
494
495 pub fn get_field_entry(&self, field: Field) -> Option<&FieldEntry> {
496 self.fields.get(field.0 as usize)
497 }
498
499 pub fn get_field_name(&self, field: Field) -> Option<&str> {
500 self.fields.get(field.0 as usize).map(|e| e.name.as_str())
501 }
502
503 pub fn fields(&self) -> impl Iterator<Item = (Field, &FieldEntry)> {
504 self.fields
505 .iter()
506 .enumerate()
507 .map(|(i, e)| (Field(i as u32), e))
508 }
509
510 pub fn num_fields(&self) -> usize {
511 self.fields.len()
512 }
513
514 pub fn has_reorder_fields(&self) -> bool {
517 self.fields.iter().any(|e| e.reorder)
518 }
519
520 pub fn reorder_on_merge(&self) -> bool {
523 self.reorder_on_merge
524 }
525
526 pub fn index_label(&self) -> &str {
529 if self.index_name.is_empty() {
530 "unknown"
531 } else {
532 &self.index_name
533 }
534 }
535
536 pub fn set_index_name(&mut self, name: impl Into<String>) {
538 self.index_name = name.into();
539 }
540
541 pub fn default_fields(&self) -> &[Field] {
543 &self.default_fields
544 }
545
546 pub fn set_default_fields(&mut self, fields: Vec<Field>) {
548 self.default_fields = fields;
549 }
550
551 pub fn query_routers(&self) -> &[QueryRouterRule] {
553 &self.query_routers
554 }
555
556 pub fn set_query_routers(&mut self, rules: Vec<QueryRouterRule>) {
558 self.query_routers = rules;
559 }
560
561 pub fn primary_field(&self) -> Option<Field> {
563 self.fields
564 .iter()
565 .enumerate()
566 .find(|(_, e)| e.primary_key)
567 .map(|(i, _)| Field(i as u32))
568 }
569}
570
571#[derive(Debug, Default)]
573pub struct SchemaBuilder {
574 fields: Vec<FieldEntry>,
575 default_fields: Vec<String>,
576 query_routers: Vec<QueryRouterRule>,
577 reorder_on_merge: bool,
578 index_name: String,
579}
580
581impl SchemaBuilder {
582 pub fn add_text_field(&mut self, name: &str, indexed: bool, stored: bool) -> Field {
583 self.add_field_with_tokenizer(
584 name,
585 FieldType::Text,
586 indexed,
587 stored,
588 Some("simple".to_string()),
589 )
590 }
591
592 pub fn add_text_field_with_tokenizer(
593 &mut self,
594 name: &str,
595 indexed: bool,
596 stored: bool,
597 tokenizer: &str,
598 ) -> Field {
599 self.add_field_with_tokenizer(
600 name,
601 FieldType::Text,
602 indexed,
603 stored,
604 Some(tokenizer.to_string()),
605 )
606 }
607
608 pub fn add_u64_field(&mut self, name: &str, indexed: bool, stored: bool) -> Field {
609 self.add_field(name, FieldType::U64, indexed, stored)
610 }
611
612 pub fn add_i64_field(&mut self, name: &str, indexed: bool, stored: bool) -> Field {
613 self.add_field(name, FieldType::I64, indexed, stored)
614 }
615
616 pub fn add_f64_field(&mut self, name: &str, indexed: bool, stored: bool) -> Field {
617 self.add_field(name, FieldType::F64, indexed, stored)
618 }
619
620 pub fn add_bytes_field(&mut self, name: &str, stored: bool) -> Field {
621 self.add_field(name, FieldType::Bytes, false, stored)
622 }
623
624 pub fn add_json_field(&mut self, name: &str, stored: bool) -> Field {
629 self.add_field(name, FieldType::Json, false, stored)
630 }
631
632 pub fn add_sparse_vector_field(&mut self, name: &str, indexed: bool, stored: bool) -> Field {
637 self.add_sparse_vector_field_with_config(
638 name,
639 indexed,
640 stored,
641 crate::structures::SparseVectorConfig::default(),
642 )
643 }
644
645 pub fn add_sparse_vector_field_with_config(
650 &mut self,
651 name: &str,
652 indexed: bool,
653 stored: bool,
654 config: crate::structures::SparseVectorConfig,
655 ) -> Field {
656 let field = Field(self.fields.len() as u32);
657 self.fields.push(FieldEntry {
658 name: name.to_string(),
659 field_type: FieldType::SparseVector,
660 indexed,
661 stored,
662 tokenizer: None,
663 multi: false,
664 positions: None,
665 sparse_vector_config: Some(config),
666 dense_vector_config: None,
667 binary_dense_vector_config: None,
668 fast: false,
669 primary_key: false,
670 reorder: false,
671 });
672 field
673 }
674
675 pub fn set_sparse_vector_config(
677 &mut self,
678 field: Field,
679 config: crate::structures::SparseVectorConfig,
680 ) {
681 if let Some(entry) = self.fields.get_mut(field.0 as usize) {
682 entry.sparse_vector_config = Some(config);
683 }
684 }
685
686 pub fn add_dense_vector_field(
691 &mut self,
692 name: &str,
693 dim: usize,
694 indexed: bool,
695 stored: bool,
696 ) -> Field {
697 self.add_dense_vector_field_with_config(name, indexed, stored, DenseVectorConfig::new(dim))
698 }
699
700 pub fn add_dense_vector_field_with_config(
702 &mut self,
703 name: &str,
704 indexed: bool,
705 stored: bool,
706 config: DenseVectorConfig,
707 ) -> Field {
708 let field = Field(self.fields.len() as u32);
709 self.fields.push(FieldEntry {
710 name: name.to_string(),
711 field_type: FieldType::DenseVector,
712 indexed,
713 stored,
714 tokenizer: None,
715 multi: false,
716 positions: None,
717 sparse_vector_config: None,
718 dense_vector_config: Some(config),
719 binary_dense_vector_config: None,
720 fast: false,
721 primary_key: false,
722 reorder: false,
723 });
724 field
725 }
726
727 pub fn add_binary_dense_vector_field(
733 &mut self,
734 name: &str,
735 dim: usize,
736 indexed: bool,
737 stored: bool,
738 ) -> Field {
739 self.add_binary_dense_vector_field_with_config(
740 name,
741 indexed,
742 stored,
743 BinaryDenseVectorConfig::new(dim),
744 )
745 }
746
747 pub fn add_binary_dense_vector_field_with_config(
749 &mut self,
750 name: &str,
751 indexed: bool,
752 stored: bool,
753 config: BinaryDenseVectorConfig,
754 ) -> Field {
755 let field = Field(self.fields.len() as u32);
756 self.fields.push(FieldEntry {
757 name: name.to_string(),
758 field_type: FieldType::BinaryDenseVector,
759 indexed,
760 stored,
761 tokenizer: None,
762 multi: false,
763 positions: None,
764 sparse_vector_config: None,
765 dense_vector_config: None,
766 binary_dense_vector_config: Some(config),
767 fast: false,
768 primary_key: false,
769 reorder: false,
770 });
771 field
772 }
773
774 fn add_field(
775 &mut self,
776 name: &str,
777 field_type: FieldType,
778 indexed: bool,
779 stored: bool,
780 ) -> Field {
781 self.add_field_with_tokenizer(name, field_type, indexed, stored, None)
782 }
783
784 fn add_field_with_tokenizer(
785 &mut self,
786 name: &str,
787 field_type: FieldType,
788 indexed: bool,
789 stored: bool,
790 tokenizer: Option<String>,
791 ) -> Field {
792 self.add_field_full(name, field_type, indexed, stored, tokenizer, false)
793 }
794
795 fn add_field_full(
796 &mut self,
797 name: &str,
798 field_type: FieldType,
799 indexed: bool,
800 stored: bool,
801 tokenizer: Option<String>,
802 multi: bool,
803 ) -> Field {
804 let field = Field(self.fields.len() as u32);
805 self.fields.push(FieldEntry {
806 name: name.to_string(),
807 field_type,
808 indexed,
809 stored,
810 tokenizer,
811 multi,
812 positions: None,
813 sparse_vector_config: None,
814 dense_vector_config: None,
815 binary_dense_vector_config: None,
816 fast: false,
817 primary_key: false,
818 reorder: false,
819 });
820 field
821 }
822
823 pub fn set_multi(&mut self, field: Field, multi: bool) {
825 if let Some(entry) = self.fields.get_mut(field.0 as usize) {
826 entry.multi = multi;
827 }
828 }
829
830 pub fn set_fast(&mut self, field: Field, fast: bool) {
833 if let Some(entry) = self.fields.get_mut(field.0 as usize) {
834 entry.fast = fast;
835 }
836 }
837
838 pub fn set_primary_key(&mut self, field: Field) {
844 if let Some(entry) = self.fields.get_mut(field.0 as usize) {
845 entry.primary_key = true;
846 entry.fast = true;
847 entry.indexed = true;
848 }
849 }
850
851 pub fn set_reorder(&mut self, field: Field, reorder: bool) {
853 if let Some(entry) = self.fields.get_mut(field.0 as usize) {
854 entry.reorder = reorder;
855 }
856 }
857
858 pub fn set_reorder_on_merge(&mut self, on: bool) {
861 self.reorder_on_merge = on;
862 }
863
864 pub fn set_index_name(&mut self, name: impl Into<String>) {
866 self.index_name = name.into();
867 }
868
869 pub fn set_positions(&mut self, field: Field, mode: PositionMode) {
871 if let Some(entry) = self.fields.get_mut(field.0 as usize) {
872 entry.positions = Some(mode);
873 }
874 }
875
876 pub fn set_default_fields(&mut self, field_names: Vec<String>) {
878 self.default_fields = field_names;
879 }
880
881 pub fn set_query_routers(&mut self, rules: Vec<QueryRouterRule>) {
883 self.query_routers = rules;
884 }
885
886 pub fn build(self) -> Schema {
887 let mut name_to_field = HashMap::new();
888 for (i, entry) in self.fields.iter().enumerate() {
889 name_to_field.insert(entry.name.clone(), Field(i as u32));
890 }
891
892 let default_fields: Vec<Field> = self
894 .default_fields
895 .iter()
896 .filter_map(|name| name_to_field.get(name).copied())
897 .collect();
898
899 Schema {
900 fields: self.fields,
901 name_to_field,
902 default_fields,
903 query_routers: self.query_routers,
904 reorder_on_merge: self.reorder_on_merge,
905 index_name: self.index_name,
906 }
907 }
908}
909
910#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
912pub enum FieldValue {
913 #[serde(rename = "text")]
914 Text(String),
915 #[serde(rename = "u64")]
916 U64(u64),
917 #[serde(rename = "i64")]
918 I64(i64),
919 #[serde(rename = "f64")]
920 F64(f64),
921 #[serde(rename = "bytes")]
922 Bytes(Vec<u8>),
923 #[serde(rename = "sparse_vector")]
925 SparseVector(Vec<(u32, f32)>),
926 #[serde(rename = "dense_vector")]
928 DenseVector(Vec<f32>),
929 #[serde(rename = "json")]
931 Json(serde_json::Value),
932 #[serde(rename = "binary_dense_vector")]
934 BinaryDenseVector(Vec<u8>),
935}
936
937impl FieldValue {
938 pub fn as_text(&self) -> Option<&str> {
939 match self {
940 FieldValue::Text(s) => Some(s),
941 _ => None,
942 }
943 }
944
945 pub fn as_u64(&self) -> Option<u64> {
946 match self {
947 FieldValue::U64(v) => Some(*v),
948 _ => None,
949 }
950 }
951
952 pub fn as_i64(&self) -> Option<i64> {
953 match self {
954 FieldValue::I64(v) => Some(*v),
955 _ => None,
956 }
957 }
958
959 pub fn as_f64(&self) -> Option<f64> {
960 match self {
961 FieldValue::F64(v) => Some(*v),
962 _ => None,
963 }
964 }
965
966 pub fn as_bytes(&self) -> Option<&[u8]> {
967 match self {
968 FieldValue::Bytes(b) => Some(b),
969 _ => None,
970 }
971 }
972
973 pub fn as_sparse_vector(&self) -> Option<&[(u32, f32)]> {
974 match self {
975 FieldValue::SparseVector(entries) => Some(entries),
976 _ => None,
977 }
978 }
979
980 pub fn as_dense_vector(&self) -> Option<&[f32]> {
981 match self {
982 FieldValue::DenseVector(v) => Some(v),
983 _ => None,
984 }
985 }
986
987 pub fn as_json(&self) -> Option<&serde_json::Value> {
988 match self {
989 FieldValue::Json(v) => Some(v),
990 _ => None,
991 }
992 }
993
994 pub fn as_binary_dense_vector(&self) -> Option<&[u8]> {
995 match self {
996 FieldValue::BinaryDenseVector(v) => Some(v),
997 _ => None,
998 }
999 }
1000}
1001
1002#[derive(Debug, Clone, Default, Serialize, Deserialize)]
1004pub struct Document {
1005 field_values: Vec<(Field, FieldValue)>,
1006}
1007
1008impl Document {
1009 pub fn new() -> Self {
1010 Self::default()
1011 }
1012
1013 pub fn add_text(&mut self, field: Field, value: impl Into<String>) {
1014 self.field_values
1015 .push((field, FieldValue::Text(value.into())));
1016 }
1017
1018 pub fn add_u64(&mut self, field: Field, value: u64) {
1019 self.field_values.push((field, FieldValue::U64(value)));
1020 }
1021
1022 pub fn add_i64(&mut self, field: Field, value: i64) {
1023 self.field_values.push((field, FieldValue::I64(value)));
1024 }
1025
1026 pub fn add_f64(&mut self, field: Field, value: f64) {
1027 self.field_values.push((field, FieldValue::F64(value)));
1028 }
1029
1030 pub fn add_bytes(&mut self, field: Field, value: Vec<u8>) {
1031 self.field_values.push((field, FieldValue::Bytes(value)));
1032 }
1033
1034 pub fn add_sparse_vector(&mut self, field: Field, entries: Vec<(u32, f32)>) {
1035 self.field_values
1036 .push((field, FieldValue::SparseVector(entries)));
1037 }
1038
1039 pub fn add_dense_vector(&mut self, field: Field, values: Vec<f32>) {
1040 self.field_values
1041 .push((field, FieldValue::DenseVector(values)));
1042 }
1043
1044 pub fn add_json(&mut self, field: Field, value: serde_json::Value) {
1045 self.field_values.push((field, FieldValue::Json(value)));
1046 }
1047
1048 pub fn add_binary_dense_vector(&mut self, field: Field, values: Vec<u8>) {
1049 self.field_values
1050 .push((field, FieldValue::BinaryDenseVector(values)));
1051 }
1052
1053 pub fn get_first(&self, field: Field) -> Option<&FieldValue> {
1054 self.field_values
1055 .iter()
1056 .find(|(f, _)| *f == field)
1057 .map(|(_, v)| v)
1058 }
1059
1060 pub fn get_all(&self, field: Field) -> impl Iterator<Item = &FieldValue> {
1061 self.field_values
1062 .iter()
1063 .filter(move |(f, _)| *f == field)
1064 .map(|(_, v)| v)
1065 }
1066
1067 pub fn field_values(&self) -> &[(Field, FieldValue)] {
1068 &self.field_values
1069 }
1070
1071 pub fn filter_stored(&self, schema: &Schema) -> Document {
1073 Document {
1074 field_values: self
1075 .field_values
1076 .iter()
1077 .filter(|(field, _)| {
1078 schema
1079 .get_field_entry(*field)
1080 .is_some_and(|entry| entry.stored)
1081 })
1082 .cloned()
1083 .collect(),
1084 }
1085 }
1086
1087 pub fn to_json(&self, schema: &Schema) -> serde_json::Value {
1093 use std::collections::HashMap;
1094
1095 let mut field_values_map: HashMap<Field, (String, bool, Vec<serde_json::Value>)> =
1097 HashMap::new();
1098
1099 for (field, value) in &self.field_values {
1100 if let Some(entry) = schema.get_field_entry(*field) {
1101 let json_value = match value {
1102 FieldValue::Text(s) => serde_json::Value::String(s.clone()),
1103 FieldValue::U64(n) => serde_json::Value::Number((*n).into()),
1104 FieldValue::I64(n) => serde_json::Value::Number((*n).into()),
1105 FieldValue::F64(n) => serde_json::json!(n),
1106 FieldValue::Bytes(b) => {
1107 use base64::Engine;
1108 serde_json::Value::String(
1109 base64::engine::general_purpose::STANDARD.encode(b),
1110 )
1111 }
1112 FieldValue::SparseVector(entries) => {
1113 let indices: Vec<u32> = entries.iter().map(|(i, _)| *i).collect();
1114 let values: Vec<f32> = entries.iter().map(|(_, v)| *v).collect();
1115 serde_json::json!({
1116 "indices": indices,
1117 "values": values
1118 })
1119 }
1120 FieldValue::DenseVector(values) => {
1121 serde_json::json!(values)
1122 }
1123 FieldValue::Json(v) => v.clone(),
1124 FieldValue::BinaryDenseVector(b) => {
1125 use base64::Engine;
1126 serde_json::Value::String(
1127 base64::engine::general_purpose::STANDARD.encode(b),
1128 )
1129 }
1130 };
1131 field_values_map
1132 .entry(*field)
1133 .or_insert_with(|| (entry.name.clone(), entry.multi, Vec::new()))
1134 .2
1135 .push(json_value);
1136 }
1137 }
1138
1139 let mut map = serde_json::Map::new();
1141 for (_field, (name, is_multi, values)) in field_values_map {
1142 let json_value = if is_multi || values.len() > 1 {
1143 serde_json::Value::Array(values)
1144 } else {
1145 values.into_iter().next().unwrap()
1146 };
1147 map.insert(name, json_value);
1148 }
1149
1150 serde_json::Value::Object(map)
1151 }
1152
1153 pub fn from_json(json: &serde_json::Value, schema: &Schema) -> Option<Self> {
1162 let obj = json.as_object()?;
1163 let mut doc = Document::new();
1164
1165 for (key, value) in obj {
1166 if let Some(field) = schema.get_field(key) {
1167 let field_entry = schema.get_field_entry(field)?;
1168 Self::add_json_value(&mut doc, field, &field_entry.field_type, value);
1169 }
1170 }
1171
1172 Some(doc)
1173 }
1174
1175 fn add_json_value(
1177 doc: &mut Document,
1178 field: Field,
1179 field_type: &FieldType,
1180 value: &serde_json::Value,
1181 ) {
1182 match value {
1183 serde_json::Value::String(s) => {
1184 if matches!(field_type, FieldType::Text) {
1185 doc.add_text(field, s.clone());
1186 }
1187 }
1188 serde_json::Value::Number(n) => {
1189 match field_type {
1190 FieldType::I64 => {
1191 if let Some(i) = n.as_i64() {
1192 doc.add_i64(field, i);
1193 }
1194 }
1195 FieldType::U64 => {
1196 if let Some(u) = n.as_u64() {
1197 doc.add_u64(field, u);
1198 } else if let Some(i) = n.as_i64() {
1199 if i >= 0 {
1201 doc.add_u64(field, i as u64);
1202 }
1203 }
1204 }
1205 FieldType::F64 => {
1206 if let Some(f) = n.as_f64() {
1207 doc.add_f64(field, f);
1208 }
1209 }
1210 _ => {}
1211 }
1212 }
1213 serde_json::Value::Array(arr) => {
1215 for item in arr {
1216 Self::add_json_value(doc, field, field_type, item);
1217 }
1218 }
1219 serde_json::Value::Object(obj) if matches!(field_type, FieldType::SparseVector) => {
1221 if let (Some(indices_val), Some(values_val)) =
1222 (obj.get("indices"), obj.get("values"))
1223 {
1224 let indices: Vec<u32> = indices_val
1225 .as_array()
1226 .map(|arr| {
1227 arr.iter()
1228 .filter_map(|v| v.as_u64().map(|n| n as u32))
1229 .collect()
1230 })
1231 .unwrap_or_default();
1232 let values: Vec<f32> = values_val
1233 .as_array()
1234 .map(|arr| {
1235 arr.iter()
1236 .filter_map(|v| v.as_f64().map(|n| n as f32))
1237 .collect()
1238 })
1239 .unwrap_or_default();
1240 if indices.len() == values.len() {
1241 let entries: Vec<(u32, f32)> = indices.into_iter().zip(values).collect();
1242 doc.add_sparse_vector(field, entries);
1243 }
1244 }
1245 }
1246 _ if matches!(field_type, FieldType::Json) => {
1248 doc.add_json(field, value.clone());
1249 }
1250 serde_json::Value::Object(_) => {}
1251 _ => {}
1252 }
1253 }
1254}
1255
1256#[cfg(test)]
1257mod tests {
1258 use super::*;
1259
1260 #[test]
1261 fn test_schema_builder() {
1262 let mut builder = Schema::builder();
1263 let title = builder.add_text_field("title", true, true);
1264 let body = builder.add_text_field("body", true, false);
1265 let count = builder.add_u64_field("count", true, true);
1266 let schema = builder.build();
1267
1268 assert_eq!(schema.get_field("title"), Some(title));
1269 assert_eq!(schema.get_field("body"), Some(body));
1270 assert_eq!(schema.get_field("count"), Some(count));
1271 assert_eq!(schema.get_field("nonexistent"), None);
1272 }
1273
1274 #[test]
1275 fn test_set_primary_key_forces_fast_and_indexed() {
1276 let mut builder = Schema::builder();
1281 let id = builder.add_text_field("id", false, true);
1282 builder.set_primary_key(id);
1283 let schema = builder.build();
1284
1285 let entry = schema.get_field_entry(id).unwrap();
1286 assert!(entry.primary_key);
1287 assert!(
1288 entry.fast,
1289 "primary key must imply fast (dedup reads the fast-field text dict)"
1290 );
1291 assert!(entry.indexed, "primary key must imply indexed");
1292 }
1293
1294 #[test]
1295 fn test_document() {
1296 let mut builder = Schema::builder();
1297 let title = builder.add_text_field("title", true, true);
1298 let count = builder.add_u64_field("count", true, true);
1299 let _schema = builder.build();
1300
1301 let mut doc = Document::new();
1302 doc.add_text(title, "Hello World");
1303 doc.add_u64(count, 42);
1304
1305 assert_eq!(doc.get_first(title).unwrap().as_text(), Some("Hello World"));
1306 assert_eq!(doc.get_first(count).unwrap().as_u64(), Some(42));
1307 }
1308
1309 #[test]
1310 fn test_document_serialization() {
1311 let mut builder = Schema::builder();
1312 let title = builder.add_text_field("title", true, true);
1313 let count = builder.add_u64_field("count", true, true);
1314 let _schema = builder.build();
1315
1316 let mut doc = Document::new();
1317 doc.add_text(title, "Hello World");
1318 doc.add_u64(count, 42);
1319
1320 let json = serde_json::to_string(&doc).unwrap();
1322 println!("Serialized doc: {}", json);
1323
1324 let doc2: Document = serde_json::from_str(&json).unwrap();
1326 assert_eq!(
1327 doc2.field_values().len(),
1328 2,
1329 "Should have 2 field values after deserialization"
1330 );
1331 assert_eq!(
1332 doc2.get_first(title).unwrap().as_text(),
1333 Some("Hello World")
1334 );
1335 assert_eq!(doc2.get_first(count).unwrap().as_u64(), Some(42));
1336 }
1337
1338 #[test]
1339 fn test_multivalue_field() {
1340 let mut builder = Schema::builder();
1341 let uris = builder.add_text_field("uris", true, true);
1342 let title = builder.add_text_field("title", true, true);
1343 let schema = builder.build();
1344
1345 let mut doc = Document::new();
1347 doc.add_text(uris, "one");
1348 doc.add_text(uris, "two");
1349 doc.add_text(title, "Test Document");
1350
1351 assert_eq!(doc.get_first(uris).unwrap().as_text(), Some("one"));
1353
1354 let all_uris: Vec<_> = doc.get_all(uris).collect();
1356 assert_eq!(all_uris.len(), 2);
1357 assert_eq!(all_uris[0].as_text(), Some("one"));
1358 assert_eq!(all_uris[1].as_text(), Some("two"));
1359
1360 let json = doc.to_json(&schema);
1362 let uris_json = json.get("uris").unwrap();
1363 assert!(uris_json.is_array(), "Multi-value field should be an array");
1364 let uris_arr = uris_json.as_array().unwrap();
1365 assert_eq!(uris_arr.len(), 2);
1366 assert_eq!(uris_arr[0].as_str(), Some("one"));
1367 assert_eq!(uris_arr[1].as_str(), Some("two"));
1368
1369 let title_json = json.get("title").unwrap();
1371 assert!(
1372 title_json.is_string(),
1373 "Single-value field should be a string"
1374 );
1375 assert_eq!(title_json.as_str(), Some("Test Document"));
1376 }
1377
1378 #[test]
1379 fn test_multivalue_from_json() {
1380 let mut builder = Schema::builder();
1381 let uris = builder.add_text_field("uris", true, true);
1382 let title = builder.add_text_field("title", true, true);
1383 let schema = builder.build();
1384
1385 let json = serde_json::json!({
1387 "uris": ["one", "two"],
1388 "title": "Test Document"
1389 });
1390
1391 let doc = Document::from_json(&json, &schema).unwrap();
1393
1394 let all_uris: Vec<_> = doc.get_all(uris).collect();
1396 assert_eq!(all_uris.len(), 2);
1397 assert_eq!(all_uris[0].as_text(), Some("one"));
1398 assert_eq!(all_uris[1].as_text(), Some("two"));
1399
1400 assert_eq!(
1402 doc.get_first(title).unwrap().as_text(),
1403 Some("Test Document")
1404 );
1405
1406 let json_out = doc.to_json(&schema);
1408 let uris_out = json_out.get("uris").unwrap().as_array().unwrap();
1409 assert_eq!(uris_out.len(), 2);
1410 assert_eq!(uris_out[0].as_str(), Some("one"));
1411 assert_eq!(uris_out[1].as_str(), Some("two"));
1412 }
1413
1414 #[test]
1415 fn test_multi_attribute_forces_array() {
1416 let mut builder = Schema::builder();
1419 let uris = builder.add_text_field("uris", true, true);
1420 builder.set_multi(uris, true); let title = builder.add_text_field("title", true, true);
1422 let schema = builder.build();
1423
1424 assert!(schema.get_field_entry(uris).unwrap().multi);
1426 assert!(!schema.get_field_entry(title).unwrap().multi);
1427
1428 let mut doc = Document::new();
1430 doc.add_text(uris, "only_one");
1431 doc.add_text(title, "Test Document");
1432
1433 let json = doc.to_json(&schema);
1435
1436 let uris_json = json.get("uris").unwrap();
1437 assert!(
1438 uris_json.is_array(),
1439 "Multi field should be array even with single value"
1440 );
1441 let uris_arr = uris_json.as_array().unwrap();
1442 assert_eq!(uris_arr.len(), 1);
1443 assert_eq!(uris_arr[0].as_str(), Some("only_one"));
1444
1445 let title_json = json.get("title").unwrap();
1447 assert!(
1448 title_json.is_string(),
1449 "Non-multi single-value field should be a string"
1450 );
1451 assert_eq!(title_json.as_str(), Some("Test Document"));
1452 }
1453
1454 #[test]
1455 fn test_sparse_vector_field() {
1456 let mut builder = Schema::builder();
1457 let embedding = builder.add_sparse_vector_field("embedding", true, true);
1458 let title = builder.add_text_field("title", true, true);
1459 let schema = builder.build();
1460
1461 assert_eq!(schema.get_field("embedding"), Some(embedding));
1462 assert_eq!(
1463 schema.get_field_entry(embedding).unwrap().field_type,
1464 FieldType::SparseVector
1465 );
1466
1467 let mut doc = Document::new();
1469 doc.add_sparse_vector(embedding, vec![(0, 1.0), (5, 2.5), (10, 0.5)]);
1470 doc.add_text(title, "Test Document");
1471
1472 let entries = doc
1474 .get_first(embedding)
1475 .unwrap()
1476 .as_sparse_vector()
1477 .unwrap();
1478 assert_eq!(entries, &[(0, 1.0), (5, 2.5), (10, 0.5)]);
1479
1480 let json = doc.to_json(&schema);
1482 let embedding_json = json.get("embedding").unwrap();
1483 assert!(embedding_json.is_object());
1484 assert_eq!(
1485 embedding_json
1486 .get("indices")
1487 .unwrap()
1488 .as_array()
1489 .unwrap()
1490 .len(),
1491 3
1492 );
1493
1494 let doc2 = Document::from_json(&json, &schema).unwrap();
1496 let entries2 = doc2
1497 .get_first(embedding)
1498 .unwrap()
1499 .as_sparse_vector()
1500 .unwrap();
1501 assert_eq!(entries2[0].0, 0);
1502 assert!((entries2[0].1 - 1.0).abs() < 1e-6);
1503 assert_eq!(entries2[1].0, 5);
1504 assert!((entries2[1].1 - 2.5).abs() < 1e-6);
1505 assert_eq!(entries2[2].0, 10);
1506 assert!((entries2[2].1 - 0.5).abs() < 1e-6);
1507 }
1508
1509 #[test]
1510 fn test_json_field() {
1511 let mut builder = Schema::builder();
1512 let metadata = builder.add_json_field("metadata", true);
1513 let title = builder.add_text_field("title", true, true);
1514 let schema = builder.build();
1515
1516 assert_eq!(schema.get_field("metadata"), Some(metadata));
1517 assert_eq!(
1518 schema.get_field_entry(metadata).unwrap().field_type,
1519 FieldType::Json
1520 );
1521 assert!(!schema.get_field_entry(metadata).unwrap().indexed);
1523 assert!(schema.get_field_entry(metadata).unwrap().stored);
1524
1525 let json_value = serde_json::json!({
1527 "author": "John Doe",
1528 "tags": ["rust", "search"],
1529 "nested": {"key": "value"}
1530 });
1531 let mut doc = Document::new();
1532 doc.add_json(metadata, json_value.clone());
1533 doc.add_text(title, "Test Document");
1534
1535 let stored_json = doc.get_first(metadata).unwrap().as_json().unwrap();
1537 assert_eq!(stored_json, &json_value);
1538 assert_eq!(
1539 stored_json.get("author").unwrap().as_str(),
1540 Some("John Doe")
1541 );
1542
1543 let doc_json = doc.to_json(&schema);
1545 let metadata_out = doc_json.get("metadata").unwrap();
1546 assert_eq!(metadata_out, &json_value);
1547
1548 let doc2 = Document::from_json(&doc_json, &schema).unwrap();
1550 let stored_json2 = doc2.get_first(metadata).unwrap().as_json().unwrap();
1551 assert_eq!(stored_json2, &json_value);
1552 }
1553
1554 #[test]
1555 fn test_json_field_various_types() {
1556 let mut builder = Schema::builder();
1557 let data = builder.add_json_field("data", true);
1558 let _schema = builder.build();
1559
1560 let arr_value = serde_json::json!([1, 2, 3, "four", null]);
1562 let mut doc = Document::new();
1563 doc.add_json(data, arr_value.clone());
1564 assert_eq!(doc.get_first(data).unwrap().as_json().unwrap(), &arr_value);
1565
1566 let str_value = serde_json::json!("just a string");
1568 let mut doc2 = Document::new();
1569 doc2.add_json(data, str_value.clone());
1570 assert_eq!(doc2.get_first(data).unwrap().as_json().unwrap(), &str_value);
1571
1572 let num_value = serde_json::json!(42.5);
1574 let mut doc3 = Document::new();
1575 doc3.add_json(data, num_value.clone());
1576 assert_eq!(doc3.get_first(data).unwrap().as_json().unwrap(), &num_value);
1577
1578 let null_value = serde_json::Value::Null;
1580 let mut doc4 = Document::new();
1581 doc4.add_json(data, null_value.clone());
1582 assert_eq!(
1583 doc4.get_first(data).unwrap().as_json().unwrap(),
1584 &null_value
1585 );
1586
1587 let bool_value = serde_json::json!(true);
1589 let mut doc5 = Document::new();
1590 doc5.add_json(data, bool_value.clone());
1591 assert_eq!(
1592 doc5.get_first(data).unwrap().as_json().unwrap(),
1593 &bool_value
1594 );
1595 }
1596}