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