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 IvfPq,
118 Tq,
122 #[default]
126 IvfTq,
127}
128
129pub(crate) fn reject_removed_vector_index_types(schema: &Schema) -> Result<(), String> {
133 for (_, entry) in schema.fields() {
134 if let Some(config) = entry.dense_vector_config.as_ref()
135 && config.index_type == VectorIndexType::IvfPq
136 {
137 return Err(format!(
138 "dense field '{}' uses index_type `ivf_pq`, which was removed; \
139 recreate the index with `ivf_tq` (trained router, training-free \
140 TurboQuant leaves) and reindex — see docs/turboquant-quantization.md",
141 entry.name,
142 ));
143 }
144 }
145 Ok(())
146}
147
148#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
154#[serde(rename_all = "snake_case")]
155pub enum IvfRoutingMode {
156 #[default]
159 Auto,
160 Flat,
162 TwoLevel,
164 Hnsw,
166}
167
168#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
174#[serde(rename_all = "snake_case")]
175pub enum DenseVectorQuantization {
176 #[default]
178 F32,
179 F16,
181 UInt8,
183 Binary,
186}
187
188impl DenseVectorQuantization {
189 pub fn element_size(self) -> usize {
192 match self {
193 Self::F32 => 4,
194 Self::F16 => 2,
195 Self::UInt8 => 1,
196 Self::Binary => panic!("element_size() not valid for Binary; use dim.div_ceil(8)"),
197 }
198 }
199
200 pub fn tag(self) -> u8 {
202 match self {
203 Self::F32 => 0,
204 Self::F16 => 1,
205 Self::UInt8 => 2,
206 Self::Binary => 3,
207 }
208 }
209
210 pub fn from_tag(tag: u8) -> Option<Self> {
212 match tag {
213 0 => Some(Self::F32),
214 1 => Some(Self::F16),
215 2 => Some(Self::UInt8),
216 3 => Some(Self::Binary),
217 _ => None,
218 }
219 }
220}
221
222#[derive(Debug, Clone, Serialize, Deserialize)]
232pub struct DenseVectorConfig {
233 pub dim: usize,
235 #[serde(default)]
238 pub index_type: VectorIndexType,
239 #[serde(default)]
241 pub quantization: DenseVectorQuantization,
242 #[serde(default, skip_serializing_if = "Option::is_none")]
246 pub num_clusters: Option<usize>,
247 #[serde(default)]
250 pub ivf_routing: IvfRoutingMode,
251 #[serde(default = "default_nprobe")]
253 pub nprobe: usize,
254 #[serde(default = "default_unit_norm")]
262 pub unit_norm: bool,
263 #[serde(default = "default_soar")]
274 pub soar: Option<crate::structures::SoarConfig>,
275}
276
277fn default_nprobe() -> usize {
278 64
279}
280
281fn default_unit_norm() -> bool {
282 true
283}
284
285fn default_soar() -> Option<crate::structures::SoarConfig> {
286 Some(crate::structures::SoarConfig::default())
287}
288
289impl DenseVectorConfig {
290 pub fn new(dim: usize) -> Self {
291 Self {
292 dim,
293 index_type: VectorIndexType::IvfTq,
294 quantization: DenseVectorQuantization::F32,
295 num_clusters: None,
296 ivf_routing: IvfRoutingMode::Auto,
297 nprobe: 64,
298 unit_norm: true,
299 soar: Some(crate::structures::SoarConfig::default()),
300 }
301 }
302
303 pub fn flat(dim: usize) -> Self {
305 Self {
306 dim,
307 index_type: VectorIndexType::Flat,
308 quantization: DenseVectorQuantization::F32,
309 num_clusters: None,
310 ivf_routing: IvfRoutingMode::Auto,
311 nprobe: 0,
312 unit_norm: true,
313 soar: None,
314 }
315 }
316
317 pub fn tq(dim: usize) -> Self {
319 Self {
320 dim,
321 index_type: VectorIndexType::Tq,
322 quantization: DenseVectorQuantization::F32,
323 num_clusters: None,
324 ivf_routing: IvfRoutingMode::Flat,
325 nprobe: 0,
326 unit_norm: true,
327 soar: None,
328 }
329 }
330
331 pub fn ivf_tq(dim: usize, num_clusters: Option<usize>, nprobe: usize) -> Self {
333 Self {
334 dim,
335 index_type: VectorIndexType::IvfTq,
336 quantization: DenseVectorQuantization::F32,
337 num_clusters,
338 ivf_routing: IvfRoutingMode::Auto,
339 nprobe,
340 unit_norm: true,
341 soar: Some(crate::structures::SoarConfig::default()),
342 }
343 }
344
345 pub fn with_quantization(mut self, quantization: DenseVectorQuantization) -> Self {
347 self.quantization = quantization;
348 self
349 }
350
351 pub fn with_unit_norm(mut self) -> Self {
353 self.unit_norm = true;
354 self
355 }
356
357 pub fn with_num_clusters(mut self, num_clusters: usize) -> Self {
359 self.num_clusters = Some(num_clusters);
360 self
361 }
362
363 pub fn with_ivf_routing(mut self, routing: IvfRoutingMode) -> Self {
365 self.ivf_routing = routing;
366 self
367 }
368 pub fn with_soar(mut self, soar: crate::structures::SoarConfig) -> Self {
370 self.soar = Some(soar);
371 self
372 }
373
374 pub fn without_soar(mut self) -> Self {
376 self.soar = None;
377 self
378 }
379
380 pub fn uses_ivf(&self) -> bool {
382 self.index_type == VectorIndexType::IvfTq
383 }
384
385 pub fn is_flat(&self) -> bool {
387 self.index_type == VectorIndexType::Flat
388 }
389
390 pub fn optimal_num_clusters(&self, num_vectors: usize) -> usize {
392 self.num_clusters.unwrap_or_else(|| {
393 let optimal = 8.0 * (num_vectors as f64).sqrt();
397 (optimal as usize).clamp(16, 1_048_576)
398 })
399 }
400}
401
402#[derive(Debug, Clone, Serialize, Deserialize)]
408pub struct BinaryDenseVectorConfig {
409 pub dim: usize,
411 #[serde(default)]
415 pub index_type: BinaryIndexType,
416 #[serde(default, skip_serializing_if = "Option::is_none")]
418 pub num_clusters: Option<usize>,
419 #[serde(default)]
422 pub ivf_routing: IvfRoutingMode,
423 #[serde(default = "default_nprobe")]
425 pub nprobe: usize,
426}
427
428#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
430#[serde(rename_all = "snake_case")]
431pub enum BinaryIndexType {
432 Flat,
434 #[default]
436 Ivf,
437}
438
439impl BinaryDenseVectorConfig {
440 pub fn new(dim: usize) -> Self {
441 assert!(
442 dim.is_multiple_of(8),
443 "BinaryDenseVector dimension must be a multiple of 8, got {dim}"
444 );
445 Self {
446 dim,
447 index_type: BinaryIndexType::Ivf,
448 num_clusters: None,
449 ivf_routing: IvfRoutingMode::Auto,
450 nprobe: 64,
451 }
452 }
453
454 pub fn with_ivf(mut self, num_clusters: Option<usize>, nprobe: usize) -> Self {
456 self.index_type = BinaryIndexType::Ivf;
457 self.num_clusters = num_clusters;
458 self.nprobe = nprobe;
459 self
460 }
461
462 pub fn with_ivf_routing(mut self, routing: IvfRoutingMode) -> Self {
464 self.ivf_routing = routing;
465 self
466 }
467
468 pub fn optimal_num_clusters(&self, num_vectors: usize) -> usize {
470 self.num_clusters.unwrap_or_else(|| {
471 let optimal = 8.0 * (num_vectors as f64).sqrt();
472 (optimal as usize).clamp(16, 1_048_576)
473 })
474 }
475
476 pub fn byte_len(&self) -> usize {
478 self.dim.div_ceil(8)
479 }
480}
481
482use super::query_field_router::QueryRouterRule;
483
484#[derive(Debug, Clone, Default, Serialize, Deserialize)]
486pub struct Schema {
487 fields: Vec<FieldEntry>,
488 name_to_field: HashMap<String, Field>,
489 #[serde(default)]
491 default_fields: Vec<Field>,
492 #[serde(default)]
494 query_routers: Vec<QueryRouterRule>,
495 #[serde(default)]
500 reorder_on_merge: bool,
501 #[serde(default)]
505 index_name: String,
506}
507
508impl Schema {
509 pub fn builder() -> SchemaBuilder {
510 SchemaBuilder::default()
511 }
512
513 pub fn get_field(&self, name: &str) -> Option<Field> {
514 self.name_to_field.get(name).copied()
515 }
516
517 pub fn get_field_entry(&self, field: Field) -> Option<&FieldEntry> {
518 self.fields.get(field.0 as usize)
519 }
520
521 pub fn get_field_name(&self, field: Field) -> Option<&str> {
522 self.fields.get(field.0 as usize).map(|e| e.name.as_str())
523 }
524
525 pub fn fields(&self) -> impl Iterator<Item = (Field, &FieldEntry)> {
526 self.fields
527 .iter()
528 .enumerate()
529 .map(|(i, e)| (Field(i as u32), e))
530 }
531
532 pub fn num_fields(&self) -> usize {
533 self.fields.len()
534 }
535
536 pub fn has_reorder_fields(&self) -> bool {
539 self.fields.iter().any(|e| e.reorder)
540 }
541
542 pub fn reorder_on_merge(&self) -> bool {
545 self.reorder_on_merge
546 }
547
548 pub fn index_label(&self) -> &str {
551 if self.index_name.is_empty() {
552 "unknown"
553 } else {
554 &self.index_name
555 }
556 }
557
558 pub fn set_index_name(&mut self, name: impl Into<String>) {
560 self.index_name = name.into();
561 }
562
563 pub fn default_fields(&self) -> &[Field] {
565 &self.default_fields
566 }
567
568 pub fn set_default_fields(&mut self, fields: Vec<Field>) {
570 self.default_fields = fields;
571 }
572
573 pub fn query_routers(&self) -> &[QueryRouterRule] {
575 &self.query_routers
576 }
577
578 pub fn set_query_routers(&mut self, rules: Vec<QueryRouterRule>) {
580 self.query_routers = rules;
581 }
582
583 pub fn primary_field(&self) -> Option<Field> {
585 self.fields
586 .iter()
587 .enumerate()
588 .find(|(_, e)| e.primary_key)
589 .map(|(i, _)| Field(i as u32))
590 }
591}
592
593#[derive(Debug, Default)]
595pub struct SchemaBuilder {
596 fields: Vec<FieldEntry>,
597 default_fields: Vec<String>,
598 query_routers: Vec<QueryRouterRule>,
599 reorder_on_merge: bool,
600 index_name: String,
601}
602
603impl SchemaBuilder {
604 pub fn add_text_field(&mut self, name: &str, indexed: bool, stored: bool) -> Field {
605 self.add_field_with_tokenizer(
606 name,
607 FieldType::Text,
608 indexed,
609 stored,
610 Some("simple".to_string()),
611 )
612 }
613
614 pub fn add_text_field_with_tokenizer(
615 &mut self,
616 name: &str,
617 indexed: bool,
618 stored: bool,
619 tokenizer: &str,
620 ) -> Field {
621 self.add_field_with_tokenizer(
622 name,
623 FieldType::Text,
624 indexed,
625 stored,
626 Some(tokenizer.to_string()),
627 )
628 }
629
630 pub fn add_u64_field(&mut self, name: &str, indexed: bool, stored: bool) -> Field {
631 self.add_field(name, FieldType::U64, indexed, stored)
632 }
633
634 pub fn add_i64_field(&mut self, name: &str, indexed: bool, stored: bool) -> Field {
635 self.add_field(name, FieldType::I64, indexed, stored)
636 }
637
638 pub fn add_f64_field(&mut self, name: &str, indexed: bool, stored: bool) -> Field {
639 self.add_field(name, FieldType::F64, indexed, stored)
640 }
641
642 pub fn add_bytes_field(&mut self, name: &str, stored: bool) -> Field {
643 self.add_field(name, FieldType::Bytes, false, stored)
644 }
645
646 pub fn add_json_field(&mut self, name: &str, stored: bool) -> Field {
651 self.add_field(name, FieldType::Json, false, stored)
652 }
653
654 pub fn add_sparse_vector_field(&mut self, name: &str, indexed: bool, stored: bool) -> Field {
659 self.add_sparse_vector_field_with_config(
660 name,
661 indexed,
662 stored,
663 crate::structures::SparseVectorConfig::default(),
664 )
665 }
666
667 pub fn add_sparse_vector_field_with_config(
672 &mut self,
673 name: &str,
674 indexed: bool,
675 stored: bool,
676 config: crate::structures::SparseVectorConfig,
677 ) -> Field {
678 let field = Field(self.fields.len() as u32);
679 self.fields.push(FieldEntry {
680 name: name.to_string(),
681 field_type: FieldType::SparseVector,
682 indexed,
683 stored,
684 tokenizer: None,
685 multi: false,
686 positions: None,
687 sparse_vector_config: Some(config),
688 dense_vector_config: None,
689 binary_dense_vector_config: None,
690 fast: false,
691 primary_key: false,
692 reorder: false,
693 });
694 field
695 }
696
697 pub fn set_sparse_vector_config(
699 &mut self,
700 field: Field,
701 config: crate::structures::SparseVectorConfig,
702 ) {
703 if let Some(entry) = self.fields.get_mut(field.0 as usize) {
704 entry.sparse_vector_config = Some(config);
705 }
706 }
707
708 pub fn add_dense_vector_field(
713 &mut self,
714 name: &str,
715 dim: usize,
716 indexed: bool,
717 stored: bool,
718 ) -> Field {
719 self.add_dense_vector_field_with_config(name, indexed, stored, DenseVectorConfig::new(dim))
720 }
721
722 pub fn add_dense_vector_field_with_config(
724 &mut self,
725 name: &str,
726 indexed: bool,
727 stored: bool,
728 config: DenseVectorConfig,
729 ) -> Field {
730 let field = Field(self.fields.len() as u32);
731 self.fields.push(FieldEntry {
732 name: name.to_string(),
733 field_type: FieldType::DenseVector,
734 indexed,
735 stored,
736 tokenizer: None,
737 multi: false,
738 positions: None,
739 sparse_vector_config: None,
740 dense_vector_config: Some(config),
741 binary_dense_vector_config: None,
742 fast: false,
743 primary_key: false,
744 reorder: false,
745 });
746 field
747 }
748
749 pub fn add_binary_dense_vector_field(
755 &mut self,
756 name: &str,
757 dim: usize,
758 indexed: bool,
759 stored: bool,
760 ) -> Field {
761 self.add_binary_dense_vector_field_with_config(
762 name,
763 indexed,
764 stored,
765 BinaryDenseVectorConfig::new(dim),
766 )
767 }
768
769 pub fn add_binary_dense_vector_field_with_config(
771 &mut self,
772 name: &str,
773 indexed: bool,
774 stored: bool,
775 config: BinaryDenseVectorConfig,
776 ) -> Field {
777 let field = Field(self.fields.len() as u32);
778 self.fields.push(FieldEntry {
779 name: name.to_string(),
780 field_type: FieldType::BinaryDenseVector,
781 indexed,
782 stored,
783 tokenizer: None,
784 multi: false,
785 positions: None,
786 sparse_vector_config: None,
787 dense_vector_config: None,
788 binary_dense_vector_config: Some(config),
789 fast: false,
790 primary_key: false,
791 reorder: false,
792 });
793 field
794 }
795
796 fn add_field(
797 &mut self,
798 name: &str,
799 field_type: FieldType,
800 indexed: bool,
801 stored: bool,
802 ) -> Field {
803 self.add_field_with_tokenizer(name, field_type, indexed, stored, None)
804 }
805
806 fn add_field_with_tokenizer(
807 &mut self,
808 name: &str,
809 field_type: FieldType,
810 indexed: bool,
811 stored: bool,
812 tokenizer: Option<String>,
813 ) -> Field {
814 self.add_field_full(name, field_type, indexed, stored, tokenizer, false)
815 }
816
817 fn add_field_full(
818 &mut self,
819 name: &str,
820 field_type: FieldType,
821 indexed: bool,
822 stored: bool,
823 tokenizer: Option<String>,
824 multi: bool,
825 ) -> Field {
826 let field = Field(self.fields.len() as u32);
827 self.fields.push(FieldEntry {
828 name: name.to_string(),
829 field_type,
830 indexed,
831 stored,
832 tokenizer,
833 multi,
834 positions: None,
835 sparse_vector_config: None,
836 dense_vector_config: None,
837 binary_dense_vector_config: None,
838 fast: false,
839 primary_key: false,
840 reorder: false,
841 });
842 field
843 }
844
845 pub fn set_multi(&mut self, field: Field, multi: bool) {
847 if let Some(entry) = self.fields.get_mut(field.0 as usize) {
848 entry.multi = multi;
849 }
850 }
851
852 pub fn set_fast(&mut self, field: Field, fast: bool) {
855 if let Some(entry) = self.fields.get_mut(field.0 as usize) {
856 entry.fast = fast;
857 }
858 }
859
860 pub fn set_primary_key(&mut self, field: Field) {
866 if let Some(entry) = self.fields.get_mut(field.0 as usize) {
867 entry.primary_key = true;
868 entry.fast = true;
869 entry.indexed = true;
870 }
871 }
872
873 pub fn set_reorder(&mut self, field: Field, reorder: bool) {
875 if let Some(entry) = self.fields.get_mut(field.0 as usize) {
876 entry.reorder = reorder;
877 }
878 }
879
880 pub fn set_reorder_on_merge(&mut self, on: bool) {
883 self.reorder_on_merge = on;
884 }
885
886 pub fn set_index_name(&mut self, name: impl Into<String>) {
888 self.index_name = name.into();
889 }
890
891 pub fn set_positions(&mut self, field: Field, mode: PositionMode) {
893 if let Some(entry) = self.fields.get_mut(field.0 as usize) {
894 entry.positions = Some(mode);
895 }
896 }
897
898 pub fn set_default_fields(&mut self, field_names: Vec<String>) {
900 self.default_fields = field_names;
901 }
902
903 pub fn set_query_routers(&mut self, rules: Vec<QueryRouterRule>) {
905 self.query_routers = rules;
906 }
907
908 pub fn build(self) -> Schema {
909 let mut name_to_field = HashMap::new();
910 for (i, entry) in self.fields.iter().enumerate() {
911 name_to_field.insert(entry.name.clone(), Field(i as u32));
912 }
913
914 let default_fields: Vec<Field> = self
916 .default_fields
917 .iter()
918 .filter_map(|name| name_to_field.get(name).copied())
919 .collect();
920
921 Schema {
922 fields: self.fields,
923 name_to_field,
924 default_fields,
925 query_routers: self.query_routers,
926 reorder_on_merge: self.reorder_on_merge,
927 index_name: self.index_name,
928 }
929 }
930}
931
932#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
934pub enum FieldValue {
935 #[serde(rename = "text")]
936 Text(String),
937 #[serde(rename = "u64")]
938 U64(u64),
939 #[serde(rename = "i64")]
940 I64(i64),
941 #[serde(rename = "f64")]
942 F64(f64),
943 #[serde(rename = "bytes")]
944 Bytes(Vec<u8>),
945 #[serde(rename = "sparse_vector")]
947 SparseVector(Vec<(u32, f32)>),
948 #[serde(rename = "dense_vector")]
950 DenseVector(Vec<f32>),
951 #[serde(rename = "json")]
953 Json(serde_json::Value),
954 #[serde(rename = "binary_dense_vector")]
956 BinaryDenseVector(Vec<u8>),
957}
958
959impl FieldValue {
960 pub fn as_text(&self) -> Option<&str> {
961 match self {
962 FieldValue::Text(s) => Some(s),
963 _ => None,
964 }
965 }
966
967 pub fn as_u64(&self) -> Option<u64> {
968 match self {
969 FieldValue::U64(v) => Some(*v),
970 _ => None,
971 }
972 }
973
974 pub fn as_i64(&self) -> Option<i64> {
975 match self {
976 FieldValue::I64(v) => Some(*v),
977 _ => None,
978 }
979 }
980
981 pub fn as_f64(&self) -> Option<f64> {
982 match self {
983 FieldValue::F64(v) => Some(*v),
984 _ => None,
985 }
986 }
987
988 pub fn as_bytes(&self) -> Option<&[u8]> {
989 match self {
990 FieldValue::Bytes(b) => Some(b),
991 _ => None,
992 }
993 }
994
995 pub fn as_sparse_vector(&self) -> Option<&[(u32, f32)]> {
996 match self {
997 FieldValue::SparseVector(entries) => Some(entries),
998 _ => None,
999 }
1000 }
1001
1002 pub fn as_dense_vector(&self) -> Option<&[f32]> {
1003 match self {
1004 FieldValue::DenseVector(v) => Some(v),
1005 _ => None,
1006 }
1007 }
1008
1009 pub fn as_json(&self) -> Option<&serde_json::Value> {
1010 match self {
1011 FieldValue::Json(v) => Some(v),
1012 _ => None,
1013 }
1014 }
1015
1016 pub fn as_binary_dense_vector(&self) -> Option<&[u8]> {
1017 match self {
1018 FieldValue::BinaryDenseVector(v) => Some(v),
1019 _ => None,
1020 }
1021 }
1022}
1023
1024#[derive(Debug, Clone, Default, Serialize, Deserialize)]
1026pub struct Document {
1027 field_values: Vec<(Field, FieldValue)>,
1028}
1029
1030impl Document {
1031 pub fn new() -> Self {
1032 Self::default()
1033 }
1034
1035 pub fn add_text(&mut self, field: Field, value: impl Into<String>) {
1036 self.field_values
1037 .push((field, FieldValue::Text(value.into())));
1038 }
1039
1040 pub fn add_u64(&mut self, field: Field, value: u64) {
1041 self.field_values.push((field, FieldValue::U64(value)));
1042 }
1043
1044 pub fn add_i64(&mut self, field: Field, value: i64) {
1045 self.field_values.push((field, FieldValue::I64(value)));
1046 }
1047
1048 pub fn add_f64(&mut self, field: Field, value: f64) {
1049 self.field_values.push((field, FieldValue::F64(value)));
1050 }
1051
1052 pub fn add_bytes(&mut self, field: Field, value: Vec<u8>) {
1053 self.field_values.push((field, FieldValue::Bytes(value)));
1054 }
1055
1056 pub fn add_sparse_vector(&mut self, field: Field, entries: Vec<(u32, f32)>) {
1057 self.field_values
1058 .push((field, FieldValue::SparseVector(entries)));
1059 }
1060
1061 pub fn add_dense_vector(&mut self, field: Field, values: Vec<f32>) {
1062 self.field_values
1063 .push((field, FieldValue::DenseVector(values)));
1064 }
1065
1066 pub fn add_json(&mut self, field: Field, value: serde_json::Value) {
1067 self.field_values.push((field, FieldValue::Json(value)));
1068 }
1069
1070 pub fn add_binary_dense_vector(&mut self, field: Field, values: Vec<u8>) {
1071 self.field_values
1072 .push((field, FieldValue::BinaryDenseVector(values)));
1073 }
1074
1075 pub fn get_first(&self, field: Field) -> Option<&FieldValue> {
1076 self.field_values
1077 .iter()
1078 .find(|(f, _)| *f == field)
1079 .map(|(_, v)| v)
1080 }
1081
1082 pub fn get_all(&self, field: Field) -> impl Iterator<Item = &FieldValue> {
1083 self.field_values
1084 .iter()
1085 .filter(move |(f, _)| *f == field)
1086 .map(|(_, v)| v)
1087 }
1088
1089 pub fn field_values(&self) -> &[(Field, FieldValue)] {
1090 &self.field_values
1091 }
1092
1093 pub fn filter_stored(&self, schema: &Schema) -> Document {
1095 Document {
1096 field_values: self
1097 .field_values
1098 .iter()
1099 .filter(|(field, _)| {
1100 schema
1101 .get_field_entry(*field)
1102 .is_some_and(|entry| entry.stored)
1103 })
1104 .cloned()
1105 .collect(),
1106 }
1107 }
1108
1109 pub fn to_json(&self, schema: &Schema) -> serde_json::Value {
1115 use std::collections::HashMap;
1116
1117 let mut field_values_map: HashMap<Field, (String, bool, Vec<serde_json::Value>)> =
1119 HashMap::new();
1120
1121 for (field, value) in &self.field_values {
1122 if let Some(entry) = schema.get_field_entry(*field) {
1123 let json_value = match value {
1124 FieldValue::Text(s) => serde_json::Value::String(s.clone()),
1125 FieldValue::U64(n) => serde_json::Value::Number((*n).into()),
1126 FieldValue::I64(n) => serde_json::Value::Number((*n).into()),
1127 FieldValue::F64(n) => serde_json::json!(n),
1128 FieldValue::Bytes(b) => {
1129 use base64::Engine;
1130 serde_json::Value::String(
1131 base64::engine::general_purpose::STANDARD.encode(b),
1132 )
1133 }
1134 FieldValue::SparseVector(entries) => {
1135 let indices: Vec<u32> = entries.iter().map(|(i, _)| *i).collect();
1136 let values: Vec<f32> = entries.iter().map(|(_, v)| *v).collect();
1137 serde_json::json!({
1138 "indices": indices,
1139 "values": values
1140 })
1141 }
1142 FieldValue::DenseVector(values) => {
1143 serde_json::json!(values)
1144 }
1145 FieldValue::Json(v) => v.clone(),
1146 FieldValue::BinaryDenseVector(b) => {
1147 use base64::Engine;
1148 serde_json::Value::String(
1149 base64::engine::general_purpose::STANDARD.encode(b),
1150 )
1151 }
1152 };
1153 field_values_map
1154 .entry(*field)
1155 .or_insert_with(|| (entry.name.clone(), entry.multi, Vec::new()))
1156 .2
1157 .push(json_value);
1158 }
1159 }
1160
1161 let mut map = serde_json::Map::new();
1163 for (_field, (name, is_multi, values)) in field_values_map {
1164 let json_value = if is_multi || values.len() > 1 {
1165 serde_json::Value::Array(values)
1166 } else {
1167 values.into_iter().next().unwrap()
1168 };
1169 map.insert(name, json_value);
1170 }
1171
1172 serde_json::Value::Object(map)
1173 }
1174
1175 pub fn from_json(json: &serde_json::Value, schema: &Schema) -> Option<Self> {
1184 let obj = json.as_object()?;
1185 let mut doc = Document::new();
1186
1187 for (key, value) in obj {
1188 if let Some(field) = schema.get_field(key) {
1189 let field_entry = schema.get_field_entry(field)?;
1190 Self::add_json_value(&mut doc, field, &field_entry.field_type, value);
1191 }
1192 }
1193
1194 Some(doc)
1195 }
1196
1197 fn add_json_value(
1199 doc: &mut Document,
1200 field: Field,
1201 field_type: &FieldType,
1202 value: &serde_json::Value,
1203 ) {
1204 match value {
1205 serde_json::Value::String(s) => {
1206 if matches!(field_type, FieldType::Text) {
1207 doc.add_text(field, s.clone());
1208 }
1209 }
1210 serde_json::Value::Number(n) => {
1211 match field_type {
1212 FieldType::I64 => {
1213 if let Some(i) = n.as_i64() {
1214 doc.add_i64(field, i);
1215 }
1216 }
1217 FieldType::U64 => {
1218 if let Some(u) = n.as_u64() {
1219 doc.add_u64(field, u);
1220 } else if let Some(i) = n.as_i64() {
1221 if i >= 0 {
1223 doc.add_u64(field, i as u64);
1224 }
1225 }
1226 }
1227 FieldType::F64 => {
1228 if let Some(f) = n.as_f64() {
1229 doc.add_f64(field, f);
1230 }
1231 }
1232 _ => {}
1233 }
1234 }
1235 serde_json::Value::Array(arr) => {
1237 for item in arr {
1238 Self::add_json_value(doc, field, field_type, item);
1239 }
1240 }
1241 serde_json::Value::Object(obj) if matches!(field_type, FieldType::SparseVector) => {
1243 if let (Some(indices_val), Some(values_val)) =
1244 (obj.get("indices"), obj.get("values"))
1245 {
1246 let indices: Vec<u32> = indices_val
1247 .as_array()
1248 .map(|arr| {
1249 arr.iter()
1250 .filter_map(|v| v.as_u64().map(|n| n as u32))
1251 .collect()
1252 })
1253 .unwrap_or_default();
1254 let values: Vec<f32> = values_val
1255 .as_array()
1256 .map(|arr| {
1257 arr.iter()
1258 .filter_map(|v| v.as_f64().map(|n| n as f32))
1259 .collect()
1260 })
1261 .unwrap_or_default();
1262 if indices.len() == values.len() {
1263 let entries: Vec<(u32, f32)> = indices.into_iter().zip(values).collect();
1264 doc.add_sparse_vector(field, entries);
1265 }
1266 }
1267 }
1268 _ if matches!(field_type, FieldType::Json) => {
1270 doc.add_json(field, value.clone());
1271 }
1272 serde_json::Value::Object(_) => {}
1273 _ => {}
1274 }
1275 }
1276}
1277
1278#[cfg(test)]
1279mod tests {
1280 use super::*;
1281
1282 #[test]
1283 fn test_schema_builder() {
1284 let mut builder = Schema::builder();
1285 let title = builder.add_text_field("title", true, true);
1286 let body = builder.add_text_field("body", true, false);
1287 let count = builder.add_u64_field("count", true, true);
1288 let schema = builder.build();
1289
1290 assert_eq!(schema.get_field("title"), Some(title));
1291 assert_eq!(schema.get_field("body"), Some(body));
1292 assert_eq!(schema.get_field("count"), Some(count));
1293 assert_eq!(schema.get_field("nonexistent"), None);
1294 }
1295
1296 #[test]
1297 fn ivf_tq_defaults_to_selective_soar() {
1298 for config in [
1299 DenseVectorConfig::new(8),
1300 DenseVectorConfig::ivf_tq(8, Some(4), 2),
1301 ] {
1302 let soar = config.soar.expect("IVF-TQ should enable SOAR by default");
1303 assert_eq!(soar.num_secondary, 1);
1304 assert!(soar.selective);
1305 assert_eq!(soar.calibration_target(), Some(0.30));
1306 }
1307
1308 assert!(DenseVectorConfig::flat(8).soar.is_none());
1309 assert!(DenseVectorConfig::tq(8).soar.is_none());
1310 }
1311
1312 #[test]
1313 fn omitted_and_explicitly_disabled_soar_are_distinct_in_serde() {
1314 let omitted: DenseVectorConfig = serde_json::from_value(serde_json::json!({
1315 "dim": 8,
1316 "index_type": "ivf_tq"
1317 }))
1318 .unwrap();
1319 let default_soar = omitted
1320 .soar
1321 .as_ref()
1322 .expect("an omitted SOAR setting should enable the selective default");
1323 assert_eq!(default_soar.num_secondary, 1);
1324 assert!(default_soar.selective);
1325 assert_eq!(default_soar.calibration_target(), Some(0.30));
1326
1327 let disabled: DenseVectorConfig = serde_json::from_value(serde_json::json!({
1328 "dim": 8,
1329 "index_type": "ivf_tq",
1330 "soar": null
1331 }))
1332 .unwrap();
1333 assert!(disabled.soar.is_none());
1334
1335 let encoded = serde_json::to_value(&disabled).unwrap();
1336 assert_eq!(encoded.get("soar"), Some(&serde_json::Value::Null));
1337 let round_trip: DenseVectorConfig = serde_json::from_value(encoded).unwrap();
1338 assert!(
1339 round_trip.soar.is_none(),
1340 "explicit off must survive a schema round trip"
1341 );
1342 }
1343
1344 #[test]
1345 fn test_set_primary_key_forces_fast_and_indexed() {
1346 let mut builder = Schema::builder();
1351 let id = builder.add_text_field("id", false, true);
1352 builder.set_primary_key(id);
1353 let schema = builder.build();
1354
1355 let entry = schema.get_field_entry(id).unwrap();
1356 assert!(entry.primary_key);
1357 assert!(
1358 entry.fast,
1359 "primary key must imply fast (dedup reads the fast-field text dict)"
1360 );
1361 assert!(entry.indexed, "primary key must imply indexed");
1362 }
1363
1364 #[test]
1365 fn test_document() {
1366 let mut builder = Schema::builder();
1367 let title = builder.add_text_field("title", true, true);
1368 let count = builder.add_u64_field("count", true, true);
1369 let _schema = builder.build();
1370
1371 let mut doc = Document::new();
1372 doc.add_text(title, "Hello World");
1373 doc.add_u64(count, 42);
1374
1375 assert_eq!(doc.get_first(title).unwrap().as_text(), Some("Hello World"));
1376 assert_eq!(doc.get_first(count).unwrap().as_u64(), Some(42));
1377 }
1378
1379 #[test]
1380 fn test_document_serialization() {
1381 let mut builder = Schema::builder();
1382 let title = builder.add_text_field("title", true, true);
1383 let count = builder.add_u64_field("count", true, true);
1384 let _schema = builder.build();
1385
1386 let mut doc = Document::new();
1387 doc.add_text(title, "Hello World");
1388 doc.add_u64(count, 42);
1389
1390 let json = serde_json::to_string(&doc).unwrap();
1392 println!("Serialized doc: {}", json);
1393
1394 let doc2: Document = serde_json::from_str(&json).unwrap();
1396 assert_eq!(
1397 doc2.field_values().len(),
1398 2,
1399 "Should have 2 field values after deserialization"
1400 );
1401 assert_eq!(
1402 doc2.get_first(title).unwrap().as_text(),
1403 Some("Hello World")
1404 );
1405 assert_eq!(doc2.get_first(count).unwrap().as_u64(), Some(42));
1406 }
1407
1408 #[test]
1409 fn test_multivalue_field() {
1410 let mut builder = Schema::builder();
1411 let uris = builder.add_text_field("uris", true, true);
1412 let title = builder.add_text_field("title", true, true);
1413 let schema = builder.build();
1414
1415 let mut doc = Document::new();
1417 doc.add_text(uris, "one");
1418 doc.add_text(uris, "two");
1419 doc.add_text(title, "Test Document");
1420
1421 assert_eq!(doc.get_first(uris).unwrap().as_text(), Some("one"));
1423
1424 let all_uris: Vec<_> = doc.get_all(uris).collect();
1426 assert_eq!(all_uris.len(), 2);
1427 assert_eq!(all_uris[0].as_text(), Some("one"));
1428 assert_eq!(all_uris[1].as_text(), Some("two"));
1429
1430 let json = doc.to_json(&schema);
1432 let uris_json = json.get("uris").unwrap();
1433 assert!(uris_json.is_array(), "Multi-value field should be an array");
1434 let uris_arr = uris_json.as_array().unwrap();
1435 assert_eq!(uris_arr.len(), 2);
1436 assert_eq!(uris_arr[0].as_str(), Some("one"));
1437 assert_eq!(uris_arr[1].as_str(), Some("two"));
1438
1439 let title_json = json.get("title").unwrap();
1441 assert!(
1442 title_json.is_string(),
1443 "Single-value field should be a string"
1444 );
1445 assert_eq!(title_json.as_str(), Some("Test Document"));
1446 }
1447
1448 #[test]
1449 fn test_multivalue_from_json() {
1450 let mut builder = Schema::builder();
1451 let uris = builder.add_text_field("uris", true, true);
1452 let title = builder.add_text_field("title", true, true);
1453 let schema = builder.build();
1454
1455 let json = serde_json::json!({
1457 "uris": ["one", "two"],
1458 "title": "Test Document"
1459 });
1460
1461 let doc = Document::from_json(&json, &schema).unwrap();
1463
1464 let all_uris: Vec<_> = doc.get_all(uris).collect();
1466 assert_eq!(all_uris.len(), 2);
1467 assert_eq!(all_uris[0].as_text(), Some("one"));
1468 assert_eq!(all_uris[1].as_text(), Some("two"));
1469
1470 assert_eq!(
1472 doc.get_first(title).unwrap().as_text(),
1473 Some("Test Document")
1474 );
1475
1476 let json_out = doc.to_json(&schema);
1478 let uris_out = json_out.get("uris").unwrap().as_array().unwrap();
1479 assert_eq!(uris_out.len(), 2);
1480 assert_eq!(uris_out[0].as_str(), Some("one"));
1481 assert_eq!(uris_out[1].as_str(), Some("two"));
1482 }
1483
1484 #[test]
1485 fn test_multi_attribute_forces_array() {
1486 let mut builder = Schema::builder();
1489 let uris = builder.add_text_field("uris", true, true);
1490 builder.set_multi(uris, true); let title = builder.add_text_field("title", true, true);
1492 let schema = builder.build();
1493
1494 assert!(schema.get_field_entry(uris).unwrap().multi);
1496 assert!(!schema.get_field_entry(title).unwrap().multi);
1497
1498 let mut doc = Document::new();
1500 doc.add_text(uris, "only_one");
1501 doc.add_text(title, "Test Document");
1502
1503 let json = doc.to_json(&schema);
1505
1506 let uris_json = json.get("uris").unwrap();
1507 assert!(
1508 uris_json.is_array(),
1509 "Multi field should be array even with single value"
1510 );
1511 let uris_arr = uris_json.as_array().unwrap();
1512 assert_eq!(uris_arr.len(), 1);
1513 assert_eq!(uris_arr[0].as_str(), Some("only_one"));
1514
1515 let title_json = json.get("title").unwrap();
1517 assert!(
1518 title_json.is_string(),
1519 "Non-multi single-value field should be a string"
1520 );
1521 assert_eq!(title_json.as_str(), Some("Test Document"));
1522 }
1523
1524 #[test]
1525 fn test_sparse_vector_field() {
1526 let mut builder = Schema::builder();
1527 let embedding = builder.add_sparse_vector_field("embedding", true, true);
1528 let title = builder.add_text_field("title", true, true);
1529 let schema = builder.build();
1530
1531 assert_eq!(schema.get_field("embedding"), Some(embedding));
1532 assert_eq!(
1533 schema.get_field_entry(embedding).unwrap().field_type,
1534 FieldType::SparseVector
1535 );
1536
1537 let mut doc = Document::new();
1539 doc.add_sparse_vector(embedding, vec![(0, 1.0), (5, 2.5), (10, 0.5)]);
1540 doc.add_text(title, "Test Document");
1541
1542 let entries = doc
1544 .get_first(embedding)
1545 .unwrap()
1546 .as_sparse_vector()
1547 .unwrap();
1548 assert_eq!(entries, &[(0, 1.0), (5, 2.5), (10, 0.5)]);
1549
1550 let json = doc.to_json(&schema);
1552 let embedding_json = json.get("embedding").unwrap();
1553 assert!(embedding_json.is_object());
1554 assert_eq!(
1555 embedding_json
1556 .get("indices")
1557 .unwrap()
1558 .as_array()
1559 .unwrap()
1560 .len(),
1561 3
1562 );
1563
1564 let doc2 = Document::from_json(&json, &schema).unwrap();
1566 let entries2 = doc2
1567 .get_first(embedding)
1568 .unwrap()
1569 .as_sparse_vector()
1570 .unwrap();
1571 assert_eq!(entries2[0].0, 0);
1572 assert!((entries2[0].1 - 1.0).abs() < 1e-6);
1573 assert_eq!(entries2[1].0, 5);
1574 assert!((entries2[1].1 - 2.5).abs() < 1e-6);
1575 assert_eq!(entries2[2].0, 10);
1576 assert!((entries2[2].1 - 0.5).abs() < 1e-6);
1577 }
1578
1579 #[test]
1580 fn test_json_field() {
1581 let mut builder = Schema::builder();
1582 let metadata = builder.add_json_field("metadata", true);
1583 let title = builder.add_text_field("title", true, true);
1584 let schema = builder.build();
1585
1586 assert_eq!(schema.get_field("metadata"), Some(metadata));
1587 assert_eq!(
1588 schema.get_field_entry(metadata).unwrap().field_type,
1589 FieldType::Json
1590 );
1591 assert!(!schema.get_field_entry(metadata).unwrap().indexed);
1593 assert!(schema.get_field_entry(metadata).unwrap().stored);
1594
1595 let json_value = serde_json::json!({
1597 "author": "John Doe",
1598 "tags": ["rust", "search"],
1599 "nested": {"key": "value"}
1600 });
1601 let mut doc = Document::new();
1602 doc.add_json(metadata, json_value.clone());
1603 doc.add_text(title, "Test Document");
1604
1605 let stored_json = doc.get_first(metadata).unwrap().as_json().unwrap();
1607 assert_eq!(stored_json, &json_value);
1608 assert_eq!(
1609 stored_json.get("author").unwrap().as_str(),
1610 Some("John Doe")
1611 );
1612
1613 let doc_json = doc.to_json(&schema);
1615 let metadata_out = doc_json.get("metadata").unwrap();
1616 assert_eq!(metadata_out, &json_value);
1617
1618 let doc2 = Document::from_json(&doc_json, &schema).unwrap();
1620 let stored_json2 = doc2.get_first(metadata).unwrap().as_json().unwrap();
1621 assert_eq!(stored_json2, &json_value);
1622 }
1623
1624 #[test]
1625 fn test_json_field_various_types() {
1626 let mut builder = Schema::builder();
1627 let data = builder.add_json_field("data", true);
1628 let _schema = builder.build();
1629
1630 let arr_value = serde_json::json!([1, 2, 3, "four", null]);
1632 let mut doc = Document::new();
1633 doc.add_json(data, arr_value.clone());
1634 assert_eq!(doc.get_first(data).unwrap().as_json().unwrap(), &arr_value);
1635
1636 let str_value = serde_json::json!("just a string");
1638 let mut doc2 = Document::new();
1639 doc2.add_json(data, str_value.clone());
1640 assert_eq!(doc2.get_first(data).unwrap().as_json().unwrap(), &str_value);
1641
1642 let num_value = serde_json::json!(42.5);
1644 let mut doc3 = Document::new();
1645 doc3.add_json(data, num_value.clone());
1646 assert_eq!(doc3.get_first(data).unwrap().as_json().unwrap(), &num_value);
1647
1648 let null_value = serde_json::Value::Null;
1650 let mut doc4 = Document::new();
1651 doc4.add_json(data, null_value.clone());
1652 assert_eq!(
1653 doc4.get_first(data).unwrap().as_json().unwrap(),
1654 &null_value
1655 );
1656
1657 let bool_value = serde_json::json!(true);
1659 let mut doc5 = Document::new();
1660 doc5.add_json(data, bool_value.clone());
1661 assert_eq!(
1662 doc5.get_first(data).unwrap().as_json().unwrap(),
1663 &bool_value
1664 );
1665 }
1666}