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 Scann,
131}
132
133pub(crate) fn reject_removed_vector_index_types(schema: &Schema) -> Result<(), String> {
137 for (_, entry) in schema.fields() {
138 if let Some(config) = entry.dense_vector_config.as_ref() {
139 validate_target_vectors(
140 &entry.name,
141 config.target_vectors,
142 !matches!(
143 config.index_type,
144 VectorIndexType::Flat | VectorIndexType::Tq
145 ),
146 )?;
147 if config.index_type == VectorIndexType::IvfPq {
148 return Err(format!(
149 "dense field '{}' uses index_type `ivf_pq`, which was removed; \
150 recreate the index with `ivf_tq` (trained router, training-free \
151 TurboQuant leaves) and reindex — see docs/turboquant-quantization.md",
152 entry.name,
153 ));
154 }
155 if config.index_type == VectorIndexType::Scann && config.soar.is_some() {
156 return Err(format!(
157 "dense field '{}' enables SOAR for ScaNN, but ScaNN SOAR secondary assignments are not implemented; set soar to null/off",
158 entry.name,
159 ));
160 }
161 validate_persisted_scann_options(
162 &entry.name,
163 config.index_type == VectorIndexType::Scann,
164 config.num_clusters,
165 config.tree_levels,
166 config.nprobe,
167 config.ivf_routing,
168 )?;
169 }
170 if let Some(config) = entry.binary_dense_vector_config.as_ref() {
171 validate_target_vectors(
172 &entry.name,
173 config.target_vectors,
174 config.index_type != BinaryIndexType::Flat,
175 )?;
176 if config.soar.is_some() && config.index_type != BinaryIndexType::Scann {
177 return Err(format!(
178 "binary dense field '{}' enables binary SOAR spilling, but it requires the ScaNN index",
179 entry.name,
180 ));
181 }
182 if config.index_type == BinaryIndexType::Scann && !config.dim.is_multiple_of(8) {
183 return Err(format!(
184 "binary dense field '{}' uses ScaNN with dimension {}; binary ScaNN dimensions must be a multiple of 8 bits",
185 entry.name, config.dim,
186 ));
187 }
188 validate_persisted_scann_options(
189 &entry.name,
190 config.index_type == BinaryIndexType::Scann,
191 config.num_clusters,
192 config.tree_levels,
193 config.nprobe,
194 config.ivf_routing,
195 )?;
196 }
197 }
198 Ok(())
199}
200
201fn validate_target_vectors(
202 field_name: &str,
203 target_vectors: Option<u64>,
204 topology_is_automatic: bool,
205) -> Result<(), String> {
206 if target_vectors == Some(0) {
207 return Err(format!(
208 "field '{field_name}' has target_vectors 0; expected a positive steady-state vector count"
209 ));
210 }
211 if target_vectors.is_some() && !topology_is_automatic {
212 return Err(format!(
213 "field '{field_name}' sets target_vectors for a flat/training-free index; the hint is only valid for IVF or ScaNN automatic topology"
214 ));
215 }
216 Ok(())
217}
218
219fn validate_persisted_scann_options(
220 field_name: &str,
221 is_scann: bool,
222 num_clusters: Option<usize>,
223 tree_levels: Option<u8>,
224 nprobe: usize,
225 routing: IvfRoutingMode,
226) -> Result<(), String> {
227 if !is_scann {
228 if tree_levels.is_some() {
229 return Err(format!(
230 "field '{field_name}' sets tree_levels but does not use the ScaNN index"
231 ));
232 }
233 return Ok(());
234 }
235 if routing != IvfRoutingMode::Auto {
236 return Err(format!(
237 "field '{field_name}' sets routing {routing:?} for ScaNN, but ScaNN owns its hierarchical routing; remove the routing option"
238 ));
239 }
240
241 if let Some(levels) = tree_levels
242 && !(1..=3).contains(&levels)
243 {
244 return Err(format!(
245 "field '{field_name}' has ScaNN tree_levels {levels}; expected 1..=3"
246 ));
247 }
248 if let Some(leaves) = num_clusters {
249 if !(2..=30_000_000).contains(&leaves) {
250 return Err(format!(
251 "field '{field_name}' has ScaNN num_clusters {leaves}; expected 2..=30000000"
252 ));
253 }
254 if nprobe > leaves {
255 return Err(format!(
256 "field '{field_name}' has ScaNN nprobe {nprobe} greater than num_clusters {leaves}"
257 ));
258 }
259 }
260 if nprobe == 0 {
261 return Err(format!(
262 "field '{field_name}' has ScaNN nprobe 0; expected a positive probe count"
263 ));
264 }
265 Ok(())
266}
267
268#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
274#[serde(rename_all = "snake_case")]
275pub enum IvfRoutingMode {
276 #[default]
279 Auto,
280 Flat,
282 TwoLevel,
284 Hnsw,
286}
287
288#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
294#[serde(rename_all = "snake_case")]
295pub enum DenseVectorQuantization {
296 #[default]
298 F32,
299 F16,
301 UInt8,
303 Binary,
306}
307
308impl DenseVectorQuantization {
309 pub fn element_size(self) -> usize {
312 match self {
313 Self::F32 => 4,
314 Self::F16 => 2,
315 Self::UInt8 => 1,
316 Self::Binary => panic!("element_size() not valid for Binary; use dim.div_ceil(8)"),
317 }
318 }
319
320 pub fn tag(self) -> u8 {
322 match self {
323 Self::F32 => 0,
324 Self::F16 => 1,
325 Self::UInt8 => 2,
326 Self::Binary => 3,
327 }
328 }
329
330 pub fn from_tag(tag: u8) -> Option<Self> {
332 match tag {
333 0 => Some(Self::F32),
334 1 => Some(Self::F16),
335 2 => Some(Self::UInt8),
336 3 => Some(Self::Binary),
337 _ => None,
338 }
339 }
340}
341
342#[derive(Debug, Clone, Serialize)]
352#[serde(deny_unknown_fields)]
353pub struct DenseVectorConfig {
354 pub dim: usize,
356 #[serde(default)]
359 pub index_type: VectorIndexType,
360 #[serde(default)]
362 pub quantization: DenseVectorQuantization,
363 #[serde(default, skip_serializing_if = "Option::is_none")]
367 pub num_clusters: Option<usize>,
368 #[serde(default, skip_serializing_if = "Option::is_none")]
372 pub target_vectors: Option<u64>,
373 #[serde(default, skip_serializing_if = "Option::is_none")]
376 pub tree_levels: Option<u8>,
377 #[serde(default)]
380 pub ivf_routing: IvfRoutingMode,
381 #[serde(default = "default_nprobe")]
383 pub nprobe: usize,
384 #[serde(default = "default_unit_norm")]
392 pub unit_norm: bool,
393 #[serde(default = "default_soar")]
404 pub soar: Option<crate::structures::SoarConfig>,
405}
406
407#[derive(Default)]
408enum PersistedSoar {
409 #[default]
410 Unspecified,
411 Specified(Option<crate::structures::SoarConfig>),
412}
413
414impl<'de> Deserialize<'de> for PersistedSoar {
415 fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
416 Option::<crate::structures::SoarConfig>::deserialize(deserializer).map(Self::Specified)
417 }
418}
419
420#[derive(Deserialize)]
421#[serde(deny_unknown_fields)]
422struct DenseVectorConfigWire {
423 dim: usize,
424 #[serde(default)]
425 index_type: VectorIndexType,
426 #[serde(default)]
427 quantization: DenseVectorQuantization,
428 #[serde(default)]
429 num_clusters: Option<usize>,
430 #[serde(default)]
431 target_vectors: Option<u64>,
432 #[serde(default)]
433 tree_levels: Option<u8>,
434 #[serde(default)]
435 ivf_routing: IvfRoutingMode,
436 #[serde(default = "default_nprobe")]
437 nprobe: usize,
438 #[serde(default = "default_unit_norm")]
439 unit_norm: bool,
440 #[serde(default)]
441 soar: PersistedSoar,
442}
443
444impl<'de> Deserialize<'de> for DenseVectorConfig {
445 fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
446 let wire = DenseVectorConfigWire::deserialize(deserializer)?;
447 let soar = match wire.soar {
448 PersistedSoar::Specified(soar) => soar,
449 PersistedSoar::Unspecified if wire.index_type == VectorIndexType::IvfTq => {
450 default_soar()
451 }
452 PersistedSoar::Unspecified => None,
453 };
454 Ok(Self {
455 dim: wire.dim,
456 index_type: wire.index_type,
457 quantization: wire.quantization,
458 num_clusters: wire.num_clusters,
459 target_vectors: wire.target_vectors,
460 tree_levels: wire.tree_levels,
461 ivf_routing: wire.ivf_routing,
462 nprobe: wire.nprobe,
463 unit_norm: wire.unit_norm,
464 soar,
465 })
466 }
467}
468
469fn default_nprobe() -> usize {
470 64
471}
472
473fn default_unit_norm() -> bool {
474 true
475}
476
477fn default_soar() -> Option<crate::structures::SoarConfig> {
478 Some(crate::structures::SoarConfig::default())
479}
480
481impl DenseVectorConfig {
482 pub fn new(dim: usize) -> Self {
483 Self {
484 dim,
485 index_type: VectorIndexType::IvfTq,
486 quantization: DenseVectorQuantization::F32,
487 num_clusters: None,
488 target_vectors: None,
489 tree_levels: None,
490 ivf_routing: IvfRoutingMode::Auto,
491 nprobe: 64,
492 unit_norm: true,
493 soar: Some(crate::structures::SoarConfig::default()),
494 }
495 }
496
497 pub fn flat(dim: usize) -> Self {
499 Self {
500 dim,
501 index_type: VectorIndexType::Flat,
502 quantization: DenseVectorQuantization::F32,
503 num_clusters: None,
504 target_vectors: None,
505 tree_levels: None,
506 ivf_routing: IvfRoutingMode::Auto,
507 nprobe: 0,
508 unit_norm: true,
509 soar: None,
510 }
511 }
512
513 pub fn tq(dim: usize) -> Self {
515 Self {
516 dim,
517 index_type: VectorIndexType::Tq,
518 quantization: DenseVectorQuantization::F32,
519 num_clusters: None,
520 target_vectors: None,
521 tree_levels: None,
522 ivf_routing: IvfRoutingMode::Flat,
523 nprobe: 0,
524 unit_norm: true,
525 soar: None,
526 }
527 }
528
529 pub fn ivf_tq(dim: usize, num_clusters: Option<usize>, nprobe: usize) -> Self {
531 Self {
532 dim,
533 index_type: VectorIndexType::IvfTq,
534 quantization: DenseVectorQuantization::F32,
535 num_clusters,
536 target_vectors: None,
537 tree_levels: None,
538 ivf_routing: IvfRoutingMode::Auto,
539 nprobe,
540 unit_norm: true,
541 soar: Some(crate::structures::SoarConfig::default()),
542 }
543 }
544
545 pub fn with_quantization(mut self, quantization: DenseVectorQuantization) -> Self {
547 self.quantization = quantization;
548 self
549 }
550
551 pub fn with_unit_norm(mut self) -> Self {
553 self.unit_norm = true;
554 self
555 }
556
557 pub fn with_num_clusters(mut self, num_clusters: usize) -> Self {
559 self.num_clusters = Some(num_clusters);
560 self
561 }
562
563 pub fn with_target_vectors(mut self, target_vectors: u64) -> Self {
565 self.target_vectors = Some(target_vectors);
566 self
567 }
568
569 pub fn with_ivf_routing(mut self, routing: IvfRoutingMode) -> Self {
571 self.ivf_routing = routing;
572 self
573 }
574 pub fn with_soar(mut self, soar: crate::structures::SoarConfig) -> Self {
576 self.soar = Some(soar);
577 self
578 }
579
580 pub fn without_soar(mut self) -> Self {
582 self.soar = None;
583 self
584 }
585
586 pub fn uses_ivf(&self) -> bool {
588 self.index_type == VectorIndexType::IvfTq
589 }
590
591 pub fn supports_soar(&self) -> bool {
593 self.index_type == VectorIndexType::IvfTq
594 }
595
596 pub fn is_flat(&self) -> bool {
598 self.index_type == VectorIndexType::Flat
599 }
600
601 pub fn optimal_num_clusters(&self, num_vectors: usize) -> usize {
603 self.num_clusters.unwrap_or_else(|| {
604 let num_vectors = self.target_vectors.map_or(num_vectors, |target| {
605 usize::try_from(target)
606 .unwrap_or(usize::MAX)
607 .max(num_vectors)
608 });
609 let optimal = 8.0 * (num_vectors as f64).sqrt();
613 (optimal as usize).clamp(16, 1_048_576)
614 })
615 }
616}
617
618#[derive(Debug, Clone, Serialize, Deserialize)]
624#[serde(deny_unknown_fields)]
625pub struct BinaryDenseVectorConfig {
626 pub dim: usize,
628 #[serde(default)]
632 pub index_type: BinaryIndexType,
633 #[serde(default, skip_serializing_if = "Option::is_none")]
635 pub num_clusters: Option<usize>,
636 #[serde(default, skip_serializing_if = "Option::is_none")]
640 pub target_vectors: Option<u64>,
641 #[serde(default, skip_serializing_if = "Option::is_none")]
644 pub tree_levels: Option<u8>,
645 #[serde(default)]
648 pub ivf_routing: IvfRoutingMode,
649 #[serde(default = "default_nprobe")]
651 pub nprobe: usize,
652 #[serde(default, skip_serializing_if = "Option::is_none")]
656 pub soar: Option<crate::structures::SoarConfig>,
657}
658
659#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
661#[serde(rename_all = "snake_case")]
662pub enum BinaryIndexType {
663 Flat,
665 #[default]
667 Ivf,
668 Scann,
670}
671
672#[derive(Debug, Clone)]
676pub enum VectorIndexAlter {
677 Dense(DenseVectorConfig),
678 Binary(BinaryDenseVectorConfig),
679}
680
681impl BinaryDenseVectorConfig {
682 pub fn new(dim: usize) -> Self {
683 assert!(
684 dim.is_multiple_of(8),
685 "BinaryDenseVector dimension must be a multiple of 8, got {dim}"
686 );
687 Self {
688 dim,
689 index_type: BinaryIndexType::Ivf,
690 num_clusters: None,
691 target_vectors: None,
692 tree_levels: None,
693 ivf_routing: IvfRoutingMode::Auto,
694 nprobe: 64,
695 soar: None,
696 }
697 }
698
699 pub fn with_ivf(mut self, num_clusters: Option<usize>, nprobe: usize) -> Self {
701 self.index_type = BinaryIndexType::Ivf;
702 self.num_clusters = num_clusters;
703 self.nprobe = nprobe;
704 self
705 }
706
707 pub fn with_target_vectors(mut self, target_vectors: u64) -> Self {
709 self.target_vectors = Some(target_vectors);
710 self
711 }
712
713 pub fn with_ivf_routing(mut self, routing: IvfRoutingMode) -> Self {
715 self.ivf_routing = routing;
716 self
717 }
718
719 pub fn with_soar(mut self, soar: crate::structures::SoarConfig) -> Self {
721 self.soar = Some(soar);
722 self
723 }
724
725 pub fn without_soar(mut self) -> Self {
727 self.soar = None;
728 self
729 }
730
731 pub fn optimal_num_clusters(&self, num_vectors: usize) -> usize {
733 self.num_clusters.unwrap_or_else(|| {
734 let num_vectors = self.target_vectors.map_or(num_vectors, |target| {
735 usize::try_from(target)
736 .unwrap_or(usize::MAX)
737 .max(num_vectors)
738 });
739 let balanced = (num_vectors as f64).sqrt().ceil() as usize;
743 balanced.clamp(16, 1_048_576)
744 })
745 }
746
747 pub fn byte_len(&self) -> usize {
749 self.dim.div_ceil(8)
750 }
751}
752
753use super::query_field_router::QueryRouterRule;
754
755#[derive(Debug, Clone, Default, Serialize, Deserialize)]
757pub struct Schema {
758 fields: Vec<FieldEntry>,
759 name_to_field: HashMap<String, Field>,
760 #[serde(default)]
762 default_fields: Vec<Field>,
763 #[serde(default)]
765 query_routers: Vec<QueryRouterRule>,
766 #[serde(default)]
771 reorder_on_merge: bool,
772 #[serde(default)]
776 index_name: String,
777}
778
779impl Schema {
780 pub fn builder() -> SchemaBuilder {
781 SchemaBuilder::default()
782 }
783
784 pub fn get_field(&self, name: &str) -> Option<Field> {
785 self.name_to_field.get(name).copied()
786 }
787
788 pub fn get_field_entry(&self, field: Field) -> Option<&FieldEntry> {
789 self.fields.get(field.0 as usize)
790 }
791
792 pub fn with_vector_index_alter(
795 &self,
796 field: Field,
797 alter: VectorIndexAlter,
798 ) -> Result<Self, String> {
799 let mut next = self.clone();
800 let entry = next
801 .fields
802 .get_mut(field.0 as usize)
803 .ok_or_else(|| format!("vector ALTER references unknown field {}", field.0))?;
804 match alter {
805 VectorIndexAlter::Dense(config) => {
806 let current = entry
807 .dense_vector_config
808 .as_ref()
809 .ok_or_else(|| format!("field '{}' is not a dense vector field", entry.name))?;
810 if config.dim != current.dim || config.quantization != current.quantization {
811 return Err(format!(
812 "field '{}' ALTER cannot change dimension or storage quantization",
813 entry.name
814 ));
815 }
816 if matches!(
817 config.index_type,
818 VectorIndexType::Flat | VectorIndexType::Tq
819 ) {
820 return Err(format!(
821 "field '{}' ALTER target must be `ivf_tq` or `scann`",
822 entry.name
823 ));
824 }
825 entry.dense_vector_config = Some(config);
826 }
827 VectorIndexAlter::Binary(config) => {
828 let current = entry.binary_dense_vector_config.as_ref().ok_or_else(|| {
829 format!("field '{}' is not a binary dense vector field", entry.name)
830 })?;
831 if config.dim != current.dim {
832 return Err(format!(
833 "field '{}' ALTER cannot change binary dimension",
834 entry.name
835 ));
836 }
837 if config.index_type == BinaryIndexType::Flat {
838 return Err(format!(
839 "field '{}' ALTER target must be `ivf` or `scann`",
840 entry.name
841 ));
842 }
843 entry.binary_dense_vector_config = Some(config);
844 }
845 }
846 reject_removed_vector_index_types(&next)?;
847 Ok(next)
848 }
849
850 pub fn get_field_name(&self, field: Field) -> Option<&str> {
851 self.fields.get(field.0 as usize).map(|e| e.name.as_str())
852 }
853
854 pub fn fields(&self) -> impl Iterator<Item = (Field, &FieldEntry)> {
855 self.fields
856 .iter()
857 .enumerate()
858 .map(|(i, e)| (Field(i as u32), e))
859 }
860
861 pub fn num_fields(&self) -> usize {
862 self.fields.len()
863 }
864
865 pub fn has_reorder_fields(&self) -> bool {
868 self.fields.iter().any(|e| e.reorder)
869 }
870
871 pub fn reorder_on_merge(&self) -> bool {
874 self.reorder_on_merge
875 }
876
877 pub fn index_label(&self) -> &str {
880 if self.index_name.is_empty() {
881 "unknown"
882 } else {
883 &self.index_name
884 }
885 }
886
887 pub fn set_index_name(&mut self, name: impl Into<String>) {
889 self.index_name = name.into();
890 }
891
892 pub fn default_fields(&self) -> &[Field] {
894 &self.default_fields
895 }
896
897 pub fn set_default_fields(&mut self, fields: Vec<Field>) {
899 self.default_fields = fields;
900 }
901
902 pub fn query_routers(&self) -> &[QueryRouterRule] {
904 &self.query_routers
905 }
906
907 pub fn set_query_routers(&mut self, rules: Vec<QueryRouterRule>) {
909 self.query_routers = rules;
910 }
911
912 pub fn primary_field(&self) -> Option<Field> {
914 self.fields
915 .iter()
916 .enumerate()
917 .find(|(_, e)| e.primary_key)
918 .map(|(i, _)| Field(i as u32))
919 }
920}
921
922#[derive(Debug, Default)]
924pub struct SchemaBuilder {
925 fields: Vec<FieldEntry>,
926 default_fields: Vec<String>,
927 query_routers: Vec<QueryRouterRule>,
928 reorder_on_merge: bool,
929 index_name: String,
930}
931
932impl SchemaBuilder {
933 pub fn add_text_field(&mut self, name: &str, indexed: bool, stored: bool) -> Field {
934 self.add_field_with_tokenizer(
935 name,
936 FieldType::Text,
937 indexed,
938 stored,
939 Some("simple".to_string()),
940 )
941 }
942
943 pub fn add_text_field_with_tokenizer(
944 &mut self,
945 name: &str,
946 indexed: bool,
947 stored: bool,
948 tokenizer: &str,
949 ) -> Field {
950 self.add_field_with_tokenizer(
951 name,
952 FieldType::Text,
953 indexed,
954 stored,
955 Some(tokenizer.to_string()),
956 )
957 }
958
959 pub fn add_u64_field(&mut self, name: &str, indexed: bool, stored: bool) -> Field {
960 self.add_field(name, FieldType::U64, indexed, stored)
961 }
962
963 pub fn add_i64_field(&mut self, name: &str, indexed: bool, stored: bool) -> Field {
964 self.add_field(name, FieldType::I64, indexed, stored)
965 }
966
967 pub fn add_f64_field(&mut self, name: &str, indexed: bool, stored: bool) -> Field {
968 self.add_field(name, FieldType::F64, indexed, stored)
969 }
970
971 pub fn add_bytes_field(&mut self, name: &str, stored: bool) -> Field {
972 self.add_field(name, FieldType::Bytes, false, stored)
973 }
974
975 pub fn add_json_field(&mut self, name: &str, stored: bool) -> Field {
980 self.add_field(name, FieldType::Json, false, stored)
981 }
982
983 pub fn add_sparse_vector_field(&mut self, name: &str, indexed: bool, stored: bool) -> Field {
988 self.add_sparse_vector_field_with_config(
989 name,
990 indexed,
991 stored,
992 crate::structures::SparseVectorConfig::default(),
993 )
994 }
995
996 pub fn add_sparse_vector_field_with_config(
1001 &mut self,
1002 name: &str,
1003 indexed: bool,
1004 stored: bool,
1005 config: crate::structures::SparseVectorConfig,
1006 ) -> Field {
1007 let field = Field(self.fields.len() as u32);
1008 self.fields.push(FieldEntry {
1009 name: name.to_string(),
1010 field_type: FieldType::SparseVector,
1011 indexed,
1012 stored,
1013 tokenizer: None,
1014 multi: false,
1015 positions: None,
1016 sparse_vector_config: Some(config),
1017 dense_vector_config: None,
1018 binary_dense_vector_config: None,
1019 fast: false,
1020 primary_key: false,
1021 reorder: false,
1022 });
1023 field
1024 }
1025
1026 pub fn set_sparse_vector_config(
1028 &mut self,
1029 field: Field,
1030 config: crate::structures::SparseVectorConfig,
1031 ) {
1032 if let Some(entry) = self.fields.get_mut(field.0 as usize) {
1033 entry.sparse_vector_config = Some(config);
1034 }
1035 }
1036
1037 pub fn add_dense_vector_field(
1042 &mut self,
1043 name: &str,
1044 dim: usize,
1045 indexed: bool,
1046 stored: bool,
1047 ) -> Field {
1048 self.add_dense_vector_field_with_config(name, indexed, stored, DenseVectorConfig::new(dim))
1049 }
1050
1051 pub fn add_dense_vector_field_with_config(
1053 &mut self,
1054 name: &str,
1055 indexed: bool,
1056 stored: bool,
1057 config: DenseVectorConfig,
1058 ) -> Field {
1059 let field = Field(self.fields.len() as u32);
1060 self.fields.push(FieldEntry {
1061 name: name.to_string(),
1062 field_type: FieldType::DenseVector,
1063 indexed,
1064 stored,
1065 tokenizer: None,
1066 multi: false,
1067 positions: None,
1068 sparse_vector_config: None,
1069 dense_vector_config: Some(config),
1070 binary_dense_vector_config: None,
1071 fast: false,
1072 primary_key: false,
1073 reorder: false,
1074 });
1075 field
1076 }
1077
1078 pub fn add_binary_dense_vector_field(
1084 &mut self,
1085 name: &str,
1086 dim: usize,
1087 indexed: bool,
1088 stored: bool,
1089 ) -> Field {
1090 self.add_binary_dense_vector_field_with_config(
1091 name,
1092 indexed,
1093 stored,
1094 BinaryDenseVectorConfig::new(dim),
1095 )
1096 }
1097
1098 pub fn add_binary_dense_vector_field_with_config(
1100 &mut self,
1101 name: &str,
1102 indexed: bool,
1103 stored: bool,
1104 config: BinaryDenseVectorConfig,
1105 ) -> Field {
1106 let field = Field(self.fields.len() as u32);
1107 self.fields.push(FieldEntry {
1108 name: name.to_string(),
1109 field_type: FieldType::BinaryDenseVector,
1110 indexed,
1111 stored,
1112 tokenizer: None,
1113 multi: false,
1114 positions: None,
1115 sparse_vector_config: None,
1116 dense_vector_config: None,
1117 binary_dense_vector_config: Some(config),
1118 fast: false,
1119 primary_key: false,
1120 reorder: false,
1121 });
1122 field
1123 }
1124
1125 fn add_field(
1126 &mut self,
1127 name: &str,
1128 field_type: FieldType,
1129 indexed: bool,
1130 stored: bool,
1131 ) -> Field {
1132 self.add_field_with_tokenizer(name, field_type, indexed, stored, None)
1133 }
1134
1135 fn add_field_with_tokenizer(
1136 &mut self,
1137 name: &str,
1138 field_type: FieldType,
1139 indexed: bool,
1140 stored: bool,
1141 tokenizer: Option<String>,
1142 ) -> Field {
1143 self.add_field_full(name, field_type, indexed, stored, tokenizer, false)
1144 }
1145
1146 fn add_field_full(
1147 &mut self,
1148 name: &str,
1149 field_type: FieldType,
1150 indexed: bool,
1151 stored: bool,
1152 tokenizer: Option<String>,
1153 multi: bool,
1154 ) -> Field {
1155 let field = Field(self.fields.len() as u32);
1156 self.fields.push(FieldEntry {
1157 name: name.to_string(),
1158 field_type,
1159 indexed,
1160 stored,
1161 tokenizer,
1162 multi,
1163 positions: None,
1164 sparse_vector_config: None,
1165 dense_vector_config: None,
1166 binary_dense_vector_config: None,
1167 fast: false,
1168 primary_key: false,
1169 reorder: false,
1170 });
1171 field
1172 }
1173
1174 pub fn set_multi(&mut self, field: Field, multi: bool) {
1176 if let Some(entry) = self.fields.get_mut(field.0 as usize) {
1177 entry.multi = multi;
1178 }
1179 }
1180
1181 pub fn set_fast(&mut self, field: Field, fast: bool) {
1184 if let Some(entry) = self.fields.get_mut(field.0 as usize) {
1185 entry.fast = fast;
1186 }
1187 }
1188
1189 pub fn set_primary_key(&mut self, field: Field) {
1195 if let Some(entry) = self.fields.get_mut(field.0 as usize) {
1196 entry.primary_key = true;
1197 entry.fast = true;
1198 entry.indexed = true;
1199 }
1200 }
1201
1202 pub fn set_reorder(&mut self, field: Field, reorder: bool) {
1204 if let Some(entry) = self.fields.get_mut(field.0 as usize) {
1205 entry.reorder = reorder;
1206 }
1207 }
1208
1209 pub fn set_reorder_on_merge(&mut self, on: bool) {
1212 self.reorder_on_merge = on;
1213 }
1214
1215 pub fn set_index_name(&mut self, name: impl Into<String>) {
1217 self.index_name = name.into();
1218 }
1219
1220 pub fn set_positions(&mut self, field: Field, mode: PositionMode) {
1222 if let Some(entry) = self.fields.get_mut(field.0 as usize) {
1223 entry.positions = Some(mode);
1224 }
1225 }
1226
1227 pub fn set_default_fields(&mut self, field_names: Vec<String>) {
1229 self.default_fields = field_names;
1230 }
1231
1232 pub fn set_query_routers(&mut self, rules: Vec<QueryRouterRule>) {
1234 self.query_routers = rules;
1235 }
1236
1237 pub fn build(self) -> Schema {
1238 let mut name_to_field = HashMap::new();
1239 for (i, entry) in self.fields.iter().enumerate() {
1240 name_to_field.insert(entry.name.clone(), Field(i as u32));
1241 }
1242
1243 let default_fields: Vec<Field> = self
1245 .default_fields
1246 .iter()
1247 .filter_map(|name| name_to_field.get(name).copied())
1248 .collect();
1249
1250 Schema {
1251 fields: self.fields,
1252 name_to_field,
1253 default_fields,
1254 query_routers: self.query_routers,
1255 reorder_on_merge: self.reorder_on_merge,
1256 index_name: self.index_name,
1257 }
1258 }
1259}
1260
1261#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
1263pub enum FieldValue {
1264 #[serde(rename = "text")]
1265 Text(String),
1266 #[serde(rename = "u64")]
1267 U64(u64),
1268 #[serde(rename = "i64")]
1269 I64(i64),
1270 #[serde(rename = "f64")]
1271 F64(f64),
1272 #[serde(rename = "bytes")]
1273 Bytes(Vec<u8>),
1274 #[serde(rename = "sparse_vector")]
1276 SparseVector(Vec<(u32, f32)>),
1277 #[serde(rename = "dense_vector")]
1279 DenseVector(Vec<f32>),
1280 #[serde(rename = "json")]
1282 Json(serde_json::Value),
1283 #[serde(rename = "binary_dense_vector")]
1285 BinaryDenseVector(Vec<u8>),
1286}
1287
1288impl FieldValue {
1289 pub fn as_text(&self) -> Option<&str> {
1290 match self {
1291 FieldValue::Text(s) => Some(s),
1292 _ => None,
1293 }
1294 }
1295
1296 pub fn as_u64(&self) -> Option<u64> {
1297 match self {
1298 FieldValue::U64(v) => Some(*v),
1299 _ => None,
1300 }
1301 }
1302
1303 pub fn as_i64(&self) -> Option<i64> {
1304 match self {
1305 FieldValue::I64(v) => Some(*v),
1306 _ => None,
1307 }
1308 }
1309
1310 pub fn as_f64(&self) -> Option<f64> {
1311 match self {
1312 FieldValue::F64(v) => Some(*v),
1313 _ => None,
1314 }
1315 }
1316
1317 pub fn as_bytes(&self) -> Option<&[u8]> {
1318 match self {
1319 FieldValue::Bytes(b) => Some(b),
1320 _ => None,
1321 }
1322 }
1323
1324 pub fn as_sparse_vector(&self) -> Option<&[(u32, f32)]> {
1325 match self {
1326 FieldValue::SparseVector(entries) => Some(entries),
1327 _ => None,
1328 }
1329 }
1330
1331 pub fn as_dense_vector(&self) -> Option<&[f32]> {
1332 match self {
1333 FieldValue::DenseVector(v) => Some(v),
1334 _ => None,
1335 }
1336 }
1337
1338 pub fn as_json(&self) -> Option<&serde_json::Value> {
1339 match self {
1340 FieldValue::Json(v) => Some(v),
1341 _ => None,
1342 }
1343 }
1344
1345 pub fn as_binary_dense_vector(&self) -> Option<&[u8]> {
1346 match self {
1347 FieldValue::BinaryDenseVector(v) => Some(v),
1348 _ => None,
1349 }
1350 }
1351}
1352
1353#[derive(Debug, Clone, Default, Serialize, Deserialize)]
1355pub struct Document {
1356 field_values: Vec<(Field, FieldValue)>,
1357}
1358
1359impl Document {
1360 pub fn new() -> Self {
1361 Self::default()
1362 }
1363
1364 pub fn add_text(&mut self, field: Field, value: impl Into<String>) {
1365 self.field_values
1366 .push((field, FieldValue::Text(value.into())));
1367 }
1368
1369 pub fn add_u64(&mut self, field: Field, value: u64) {
1370 self.field_values.push((field, FieldValue::U64(value)));
1371 }
1372
1373 pub fn add_i64(&mut self, field: Field, value: i64) {
1374 self.field_values.push((field, FieldValue::I64(value)));
1375 }
1376
1377 pub fn add_f64(&mut self, field: Field, value: f64) {
1378 self.field_values.push((field, FieldValue::F64(value)));
1379 }
1380
1381 pub fn add_bytes(&mut self, field: Field, value: Vec<u8>) {
1382 self.field_values.push((field, FieldValue::Bytes(value)));
1383 }
1384
1385 pub fn add_sparse_vector(&mut self, field: Field, entries: Vec<(u32, f32)>) {
1386 self.field_values
1387 .push((field, FieldValue::SparseVector(entries)));
1388 }
1389
1390 pub fn add_dense_vector(&mut self, field: Field, values: Vec<f32>) {
1391 self.field_values
1392 .push((field, FieldValue::DenseVector(values)));
1393 }
1394
1395 pub fn add_json(&mut self, field: Field, value: serde_json::Value) {
1396 self.field_values.push((field, FieldValue::Json(value)));
1397 }
1398
1399 pub fn add_binary_dense_vector(&mut self, field: Field, values: Vec<u8>) {
1400 self.field_values
1401 .push((field, FieldValue::BinaryDenseVector(values)));
1402 }
1403
1404 pub fn get_first(&self, field: Field) -> Option<&FieldValue> {
1405 self.field_values
1406 .iter()
1407 .find(|(f, _)| *f == field)
1408 .map(|(_, v)| v)
1409 }
1410
1411 pub fn get_all(&self, field: Field) -> impl Iterator<Item = &FieldValue> {
1412 self.field_values
1413 .iter()
1414 .filter(move |(f, _)| *f == field)
1415 .map(|(_, v)| v)
1416 }
1417
1418 pub fn field_values(&self) -> &[(Field, FieldValue)] {
1419 &self.field_values
1420 }
1421
1422 pub fn filter_stored(&self, schema: &Schema) -> Document {
1424 Document {
1425 field_values: self
1426 .field_values
1427 .iter()
1428 .filter(|(field, _)| {
1429 schema
1430 .get_field_entry(*field)
1431 .is_some_and(|entry| entry.stored)
1432 })
1433 .cloned()
1434 .collect(),
1435 }
1436 }
1437
1438 pub fn to_json(&self, schema: &Schema) -> serde_json::Value {
1444 use std::collections::HashMap;
1445
1446 let mut field_values_map: HashMap<Field, (String, bool, Vec<serde_json::Value>)> =
1448 HashMap::new();
1449
1450 for (field, value) in &self.field_values {
1451 if let Some(entry) = schema.get_field_entry(*field) {
1452 let json_value = match value {
1453 FieldValue::Text(s) => serde_json::Value::String(s.clone()),
1454 FieldValue::U64(n) => serde_json::Value::Number((*n).into()),
1455 FieldValue::I64(n) => serde_json::Value::Number((*n).into()),
1456 FieldValue::F64(n) => serde_json::json!(n),
1457 FieldValue::Bytes(b) => {
1458 use base64::Engine;
1459 serde_json::Value::String(
1460 base64::engine::general_purpose::STANDARD.encode(b),
1461 )
1462 }
1463 FieldValue::SparseVector(entries) => {
1464 let indices: Vec<u32> = entries.iter().map(|(i, _)| *i).collect();
1465 let values: Vec<f32> = entries.iter().map(|(_, v)| *v).collect();
1466 serde_json::json!({
1467 "indices": indices,
1468 "values": values
1469 })
1470 }
1471 FieldValue::DenseVector(values) => {
1472 serde_json::json!(values)
1473 }
1474 FieldValue::Json(v) => v.clone(),
1475 FieldValue::BinaryDenseVector(b) => {
1476 use base64::Engine;
1477 serde_json::Value::String(
1478 base64::engine::general_purpose::STANDARD.encode(b),
1479 )
1480 }
1481 };
1482 field_values_map
1483 .entry(*field)
1484 .or_insert_with(|| (entry.name.clone(), entry.multi, Vec::new()))
1485 .2
1486 .push(json_value);
1487 }
1488 }
1489
1490 let mut map = serde_json::Map::new();
1492 for (_field, (name, is_multi, values)) in field_values_map {
1493 let json_value = if is_multi || values.len() > 1 {
1494 serde_json::Value::Array(values)
1495 } else {
1496 values.into_iter().next().unwrap()
1497 };
1498 map.insert(name, json_value);
1499 }
1500
1501 serde_json::Value::Object(map)
1502 }
1503
1504 pub fn from_json(json: &serde_json::Value, schema: &Schema) -> Option<Self> {
1513 let obj = json.as_object()?;
1514 let mut doc = Document::new();
1515
1516 for (key, value) in obj {
1517 if let Some(field) = schema.get_field(key) {
1518 let field_entry = schema.get_field_entry(field)?;
1519 Self::add_json_value(&mut doc, field, &field_entry.field_type, value);
1520 }
1521 }
1522
1523 Some(doc)
1524 }
1525
1526 fn add_json_value(
1528 doc: &mut Document,
1529 field: Field,
1530 field_type: &FieldType,
1531 value: &serde_json::Value,
1532 ) {
1533 match value {
1534 serde_json::Value::String(s) => {
1535 if matches!(field_type, FieldType::Text) {
1536 doc.add_text(field, s.clone());
1537 }
1538 }
1539 serde_json::Value::Number(n) => {
1540 match field_type {
1541 FieldType::I64 => {
1542 if let Some(i) = n.as_i64() {
1543 doc.add_i64(field, i);
1544 }
1545 }
1546 FieldType::U64 => {
1547 if let Some(u) = n.as_u64() {
1548 doc.add_u64(field, u);
1549 } else if let Some(i) = n.as_i64() {
1550 if i >= 0 {
1552 doc.add_u64(field, i as u64);
1553 }
1554 }
1555 }
1556 FieldType::F64 => {
1557 if let Some(f) = n.as_f64() {
1558 doc.add_f64(field, f);
1559 }
1560 }
1561 _ => {}
1562 }
1563 }
1564 serde_json::Value::Array(arr) => {
1566 for item in arr {
1567 Self::add_json_value(doc, field, field_type, item);
1568 }
1569 }
1570 serde_json::Value::Object(obj) if matches!(field_type, FieldType::SparseVector) => {
1572 if let (Some(indices_val), Some(values_val)) =
1573 (obj.get("indices"), obj.get("values"))
1574 {
1575 let indices: Vec<u32> = indices_val
1576 .as_array()
1577 .map(|arr| {
1578 arr.iter()
1579 .filter_map(|v| v.as_u64().map(|n| n as u32))
1580 .collect()
1581 })
1582 .unwrap_or_default();
1583 let values: Vec<f32> = values_val
1584 .as_array()
1585 .map(|arr| {
1586 arr.iter()
1587 .filter_map(|v| v.as_f64().map(|n| n as f32))
1588 .collect()
1589 })
1590 .unwrap_or_default();
1591 if indices.len() == values.len() {
1592 let entries: Vec<(u32, f32)> = indices.into_iter().zip(values).collect();
1593 doc.add_sparse_vector(field, entries);
1594 }
1595 }
1596 }
1597 _ if matches!(field_type, FieldType::Json) => {
1599 doc.add_json(field, value.clone());
1600 }
1601 serde_json::Value::Object(_) => {}
1602 _ => {}
1603 }
1604 }
1605}
1606
1607#[cfg(test)]
1608mod tests {
1609 use super::*;
1610
1611 #[test]
1612 fn test_schema_builder() {
1613 let mut builder = Schema::builder();
1614 let title = builder.add_text_field("title", true, true);
1615 let body = builder.add_text_field("body", true, false);
1616 let count = builder.add_u64_field("count", true, true);
1617 let schema = builder.build();
1618
1619 assert_eq!(schema.get_field("title"), Some(title));
1620 assert_eq!(schema.get_field("body"), Some(body));
1621 assert_eq!(schema.get_field("count"), Some(count));
1622 assert_eq!(schema.get_field("nonexistent"), None);
1623 }
1624
1625 #[test]
1626 fn ivf_tq_defaults_to_selective_soar() {
1627 for config in [
1628 DenseVectorConfig::new(8),
1629 DenseVectorConfig::ivf_tq(8, Some(4), 2),
1630 ] {
1631 let soar = config.soar.expect("IVF-TQ should enable SOAR by default");
1632 assert_eq!(soar.num_secondary, 1);
1633 assert!(soar.selective);
1634 assert_eq!(soar.calibration_target(), Some(0.30));
1635 }
1636
1637 assert!(DenseVectorConfig::flat(8).soar.is_none());
1638 assert!(DenseVectorConfig::tq(8).soar.is_none());
1639 }
1640
1641 #[test]
1642 fn binary_ivf_uses_measured_balanced_fifteen_million_geometry() {
1643 let config = BinaryDenseVectorConfig::new(2_560);
1644 assert_eq!(config.optimal_num_clusters(15_000_000), 3_873);
1645
1646 let explicit = config.with_ivf(Some(8_192), 128);
1647 assert_eq!(explicit.optimal_num_clusters(15_000_000), 8_192);
1648 }
1649
1650 #[test]
1651 fn target_vectors_sizes_automatic_topology_but_explicit_clusters_win() {
1652 let hinted = BinaryDenseVectorConfig::new(2_560).with_target_vectors(1_000_000_000);
1653 assert_eq!(hinted.optimal_num_clusters(1_000_000), 31_623);
1654
1655 let lower_hint = BinaryDenseVectorConfig::new(2_560).with_target_vectors(1_000_000);
1656 assert_eq!(
1657 lower_hint.optimal_num_clusters(15_000_000),
1658 BinaryDenseVectorConfig::new(2_560).optimal_num_clusters(15_000_000),
1659 "a steady-state hint is a lower bound and must not shrink live-corpus geometry"
1660 );
1661
1662 let explicit = hinted.with_ivf(Some(8_192), 128);
1663 assert_eq!(explicit.optimal_num_clusters(1_000_000), 8_192);
1664
1665 let float = DenseVectorConfig::ivf_tq(1_024, None, 64).with_target_vectors(1_000_000_000);
1666 assert_eq!(float.optimal_num_clusters(1_000_000), 252_982);
1667 }
1668
1669 #[test]
1670 fn persisted_target_vectors_must_be_positive_and_topology_bearing() {
1671 let mut zero = BinaryDenseVectorConfig::new(256);
1672 zero.target_vectors = Some(0);
1673 let mut builder = Schema::builder();
1674 builder.add_binary_dense_vector_field_with_config("hash", true, false, zero);
1675 let error = reject_removed_vector_index_types(&builder.build()).unwrap_err();
1676 assert!(error.contains("positive steady-state"), "{error}");
1677
1678 let hinted = DenseVectorConfig::ivf_tq(128, None, 64).with_target_vectors(1_000_000_000);
1679 let encoded = serde_json::to_value(&hinted).unwrap();
1680 let decoded: DenseVectorConfig = serde_json::from_value(encoded).unwrap();
1681 assert_eq!(decoded.target_vectors, Some(1_000_000_000));
1682
1683 let binary = BinaryDenseVectorConfig::new(2_560).with_target_vectors(1_000_000_000);
1684 let encoded = serde_json::to_value(&binary).unwrap();
1685 let decoded: BinaryDenseVectorConfig = serde_json::from_value(encoded).unwrap();
1686 assert_eq!(decoded.target_vectors, Some(1_000_000_000));
1687
1688 let old_dense: DenseVectorConfig = serde_json::from_value(serde_json::json!({
1689 "dim": 128,
1690 "index_type": "ivf_tq"
1691 }))
1692 .unwrap();
1693 assert_eq!(old_dense.target_vectors, None);
1694 let old_binary: BinaryDenseVectorConfig = serde_json::from_value(serde_json::json!({
1695 "dim": 256,
1696 "index_type": "ivf"
1697 }))
1698 .unwrap();
1699 assert_eq!(old_binary.target_vectors, None);
1700
1701 let flat = DenseVectorConfig::flat(128).with_target_vectors(1_000_000);
1702 let mut builder = Schema::builder();
1703 builder.add_dense_vector_field_with_config("embedding", true, false, flat);
1704 let error = reject_removed_vector_index_types(&builder.build()).unwrap_err();
1705 assert!(error.contains("flat/training-free"), "{error}");
1706
1707 let mut binary_flat = BinaryDenseVectorConfig::new(256);
1708 binary_flat.index_type = BinaryIndexType::Flat;
1709 binary_flat.target_vectors = Some(1_000_000);
1710 let mut builder = Schema::builder();
1711 builder.add_binary_dense_vector_field_with_config("hash", true, false, binary_flat);
1712 let error = reject_removed_vector_index_types(&builder.build()).unwrap_err();
1713 assert!(error.contains("flat/training-free"), "{error}");
1714 }
1715
1716 #[test]
1717 fn omitted_and_explicitly_disabled_soar_are_distinct_in_serde() {
1718 let omitted: DenseVectorConfig = serde_json::from_value(serde_json::json!({
1719 "dim": 8,
1720 "index_type": "ivf_tq"
1721 }))
1722 .unwrap();
1723 let default_soar = omitted
1724 .soar
1725 .as_ref()
1726 .expect("an omitted SOAR setting should enable the selective default");
1727 assert_eq!(default_soar.num_secondary, 1);
1728 assert!(default_soar.selective);
1729 assert_eq!(default_soar.calibration_target(), Some(0.30));
1730
1731 let disabled: DenseVectorConfig = serde_json::from_value(serde_json::json!({
1732 "dim": 8,
1733 "index_type": "ivf_tq",
1734 "soar": null
1735 }))
1736 .unwrap();
1737 assert!(disabled.soar.is_none());
1738
1739 let encoded = serde_json::to_value(&disabled).unwrap();
1740 assert_eq!(encoded.get("soar"), Some(&serde_json::Value::Null));
1741 let round_trip: DenseVectorConfig = serde_json::from_value(encoded).unwrap();
1742 assert!(
1743 round_trip.soar.is_none(),
1744 "explicit off must survive a schema round trip"
1745 );
1746 }
1747
1748 #[test]
1749 fn scann_config_serde_preserves_old_json_defaults_and_new_parameters() {
1750 let old_dense: DenseVectorConfig = serde_json::from_value(serde_json::json!({
1751 "dim": 768,
1752 "index_type": "ivf_tq"
1753 }))
1754 .unwrap();
1755 assert_eq!(old_dense.tree_levels, None);
1756 let old_json = serde_json::to_value(&old_dense).unwrap();
1757 assert!(old_json.get("tree_levels").is_none());
1758
1759 let scann: DenseVectorConfig = serde_json::from_value(serde_json::json!({
1760 "dim": 1024,
1761 "index_type": "scann",
1762 "num_clusters": 10_000_000,
1763 "tree_levels": 2,
1764 "nprobe": 1024
1765 }))
1766 .unwrap();
1767 assert_eq!(scann.index_type, VectorIndexType::Scann);
1768 assert_eq!(scann.tree_levels, Some(2));
1769 assert!(scann.soar.is_none());
1770
1771 let binary: BinaryDenseVectorConfig = serde_json::from_value(serde_json::json!({
1772 "dim": 1024,
1773 "index_type": "scann",
1774 "tree_levels": 3
1775 }))
1776 .unwrap();
1777 assert_eq!(binary.index_type, BinaryIndexType::Scann);
1778 assert_eq!(binary.tree_levels, Some(3));
1779 }
1780
1781 #[test]
1782 fn persisted_scann_geometry_is_validated_on_schema_load() {
1783 let mut invalid_levels = DenseVectorConfig::new(128);
1784 invalid_levels.index_type = VectorIndexType::Scann;
1785 invalid_levels.tree_levels = Some(4);
1786 invalid_levels.soar = None;
1787 let mut builder = Schema::builder();
1788 builder.add_dense_vector_field_with_config("embedding", true, false, invalid_levels);
1789 let error = reject_removed_vector_index_types(&builder.build())
1790 .expect_err("invalid persisted ScaNN levels must fail at the schema gate");
1791 assert!(error.contains("1..=3"), "{error}");
1792
1793 let mut wrong_algorithm = BinaryDenseVectorConfig::new(256);
1794 wrong_algorithm.tree_levels = Some(2);
1795 let mut builder = Schema::builder();
1796 builder.add_binary_dense_vector_field_with_config("hash", true, false, wrong_algorithm);
1797 let error = reject_removed_vector_index_types(&builder.build())
1798 .expect_err("ScaNN-only persisted options must fail on IVF");
1799 assert!(error.contains("does not use the ScaNN index"), "{error}");
1800
1801 let mut invalid_soar = DenseVectorConfig::flat(128);
1802 invalid_soar.index_type = VectorIndexType::Scann;
1803 invalid_soar.nprobe = 1;
1804 invalid_soar.soar = Some(crate::structures::SoarConfig::default());
1805 let mut builder = Schema::builder();
1806 builder.add_dense_vector_field_with_config("embedding", true, false, invalid_soar);
1807 let error = reject_removed_vector_index_types(&builder.build())
1808 .expect_err("persisted ScaNN SOAR must fail until assignments exist");
1809 assert!(error.contains("not implemented"), "{error}");
1810
1811 let mut one_leaf = DenseVectorConfig::flat(128);
1812 one_leaf.index_type = VectorIndexType::Scann;
1813 one_leaf.num_clusters = Some(1);
1814 one_leaf.nprobe = 1;
1815 let mut builder = Schema::builder();
1816 builder.add_dense_vector_field_with_config("embedding", true, false, one_leaf);
1817 let error = reject_removed_vector_index_types(&builder.build())
1818 .expect_err("one-leaf ScaNN geometry must fail at schema load");
1819 assert!(error.contains("2..=30000000"), "{error}");
1820
1821 let binary = BinaryDenseVectorConfig {
1822 dim: 255,
1823 index_type: BinaryIndexType::Scann,
1824 num_clusters: Some(2),
1825 target_vectors: None,
1826 tree_levels: Some(1),
1827 ivf_routing: IvfRoutingMode::Auto,
1828 nprobe: 1,
1829 soar: None,
1830 };
1831 let mut builder = Schema::builder();
1832 builder.add_binary_dense_vector_field_with_config("hash", true, false, binary);
1833 let error = reject_removed_vector_index_types(&builder.build())
1834 .expect_err("binary ScaNN dimensions must be byte-aligned");
1835 assert!(error.contains("multiple of 8"), "{error}");
1836 }
1837
1838 #[test]
1839 fn test_set_primary_key_forces_fast_and_indexed() {
1840 let mut builder = Schema::builder();
1845 let id = builder.add_text_field("id", false, true);
1846 builder.set_primary_key(id);
1847 let schema = builder.build();
1848
1849 let entry = schema.get_field_entry(id).unwrap();
1850 assert!(entry.primary_key);
1851 assert!(
1852 entry.fast,
1853 "primary key must imply fast (dedup reads the fast-field text dict)"
1854 );
1855 assert!(entry.indexed, "primary key must imply indexed");
1856 }
1857
1858 #[test]
1859 fn test_document() {
1860 let mut builder = Schema::builder();
1861 let title = builder.add_text_field("title", true, true);
1862 let count = builder.add_u64_field("count", true, true);
1863 let _schema = builder.build();
1864
1865 let mut doc = Document::new();
1866 doc.add_text(title, "Hello World");
1867 doc.add_u64(count, 42);
1868
1869 assert_eq!(doc.get_first(title).unwrap().as_text(), Some("Hello World"));
1870 assert_eq!(doc.get_first(count).unwrap().as_u64(), Some(42));
1871 }
1872
1873 #[test]
1874 fn test_document_serialization() {
1875 let mut builder = Schema::builder();
1876 let title = builder.add_text_field("title", true, true);
1877 let count = builder.add_u64_field("count", true, true);
1878 let _schema = builder.build();
1879
1880 let mut doc = Document::new();
1881 doc.add_text(title, "Hello World");
1882 doc.add_u64(count, 42);
1883
1884 let json = serde_json::to_string(&doc).unwrap();
1886 println!("Serialized doc: {}", json);
1887
1888 let doc2: Document = serde_json::from_str(&json).unwrap();
1890 assert_eq!(
1891 doc2.field_values().len(),
1892 2,
1893 "Should have 2 field values after deserialization"
1894 );
1895 assert_eq!(
1896 doc2.get_first(title).unwrap().as_text(),
1897 Some("Hello World")
1898 );
1899 assert_eq!(doc2.get_first(count).unwrap().as_u64(), Some(42));
1900 }
1901
1902 #[test]
1903 fn test_multivalue_field() {
1904 let mut builder = Schema::builder();
1905 let uris = builder.add_text_field("uris", true, true);
1906 let title = builder.add_text_field("title", true, true);
1907 let schema = builder.build();
1908
1909 let mut doc = Document::new();
1911 doc.add_text(uris, "one");
1912 doc.add_text(uris, "two");
1913 doc.add_text(title, "Test Document");
1914
1915 assert_eq!(doc.get_first(uris).unwrap().as_text(), Some("one"));
1917
1918 let all_uris: Vec<_> = doc.get_all(uris).collect();
1920 assert_eq!(all_uris.len(), 2);
1921 assert_eq!(all_uris[0].as_text(), Some("one"));
1922 assert_eq!(all_uris[1].as_text(), Some("two"));
1923
1924 let json = doc.to_json(&schema);
1926 let uris_json = json.get("uris").unwrap();
1927 assert!(uris_json.is_array(), "Multi-value field should be an array");
1928 let uris_arr = uris_json.as_array().unwrap();
1929 assert_eq!(uris_arr.len(), 2);
1930 assert_eq!(uris_arr[0].as_str(), Some("one"));
1931 assert_eq!(uris_arr[1].as_str(), Some("two"));
1932
1933 let title_json = json.get("title").unwrap();
1935 assert!(
1936 title_json.is_string(),
1937 "Single-value field should be a string"
1938 );
1939 assert_eq!(title_json.as_str(), Some("Test Document"));
1940 }
1941
1942 #[test]
1943 fn test_multivalue_from_json() {
1944 let mut builder = Schema::builder();
1945 let uris = builder.add_text_field("uris", true, true);
1946 let title = builder.add_text_field("title", true, true);
1947 let schema = builder.build();
1948
1949 let json = serde_json::json!({
1951 "uris": ["one", "two"],
1952 "title": "Test Document"
1953 });
1954
1955 let doc = Document::from_json(&json, &schema).unwrap();
1957
1958 let all_uris: Vec<_> = doc.get_all(uris).collect();
1960 assert_eq!(all_uris.len(), 2);
1961 assert_eq!(all_uris[0].as_text(), Some("one"));
1962 assert_eq!(all_uris[1].as_text(), Some("two"));
1963
1964 assert_eq!(
1966 doc.get_first(title).unwrap().as_text(),
1967 Some("Test Document")
1968 );
1969
1970 let json_out = doc.to_json(&schema);
1972 let uris_out = json_out.get("uris").unwrap().as_array().unwrap();
1973 assert_eq!(uris_out.len(), 2);
1974 assert_eq!(uris_out[0].as_str(), Some("one"));
1975 assert_eq!(uris_out[1].as_str(), Some("two"));
1976 }
1977
1978 #[test]
1979 fn test_multi_attribute_forces_array() {
1980 let mut builder = Schema::builder();
1983 let uris = builder.add_text_field("uris", true, true);
1984 builder.set_multi(uris, true); let title = builder.add_text_field("title", true, true);
1986 let schema = builder.build();
1987
1988 assert!(schema.get_field_entry(uris).unwrap().multi);
1990 assert!(!schema.get_field_entry(title).unwrap().multi);
1991
1992 let mut doc = Document::new();
1994 doc.add_text(uris, "only_one");
1995 doc.add_text(title, "Test Document");
1996
1997 let json = doc.to_json(&schema);
1999
2000 let uris_json = json.get("uris").unwrap();
2001 assert!(
2002 uris_json.is_array(),
2003 "Multi field should be array even with single value"
2004 );
2005 let uris_arr = uris_json.as_array().unwrap();
2006 assert_eq!(uris_arr.len(), 1);
2007 assert_eq!(uris_arr[0].as_str(), Some("only_one"));
2008
2009 let title_json = json.get("title").unwrap();
2011 assert!(
2012 title_json.is_string(),
2013 "Non-multi single-value field should be a string"
2014 );
2015 assert_eq!(title_json.as_str(), Some("Test Document"));
2016 }
2017
2018 #[test]
2019 fn test_sparse_vector_field() {
2020 let mut builder = Schema::builder();
2021 let embedding = builder.add_sparse_vector_field("embedding", true, true);
2022 let title = builder.add_text_field("title", true, true);
2023 let schema = builder.build();
2024
2025 assert_eq!(schema.get_field("embedding"), Some(embedding));
2026 assert_eq!(
2027 schema.get_field_entry(embedding).unwrap().field_type,
2028 FieldType::SparseVector
2029 );
2030
2031 let mut doc = Document::new();
2033 doc.add_sparse_vector(embedding, vec![(0, 1.0), (5, 2.5), (10, 0.5)]);
2034 doc.add_text(title, "Test Document");
2035
2036 let entries = doc
2038 .get_first(embedding)
2039 .unwrap()
2040 .as_sparse_vector()
2041 .unwrap();
2042 assert_eq!(entries, &[(0, 1.0), (5, 2.5), (10, 0.5)]);
2043
2044 let json = doc.to_json(&schema);
2046 let embedding_json = json.get("embedding").unwrap();
2047 assert!(embedding_json.is_object());
2048 assert_eq!(
2049 embedding_json
2050 .get("indices")
2051 .unwrap()
2052 .as_array()
2053 .unwrap()
2054 .len(),
2055 3
2056 );
2057
2058 let doc2 = Document::from_json(&json, &schema).unwrap();
2060 let entries2 = doc2
2061 .get_first(embedding)
2062 .unwrap()
2063 .as_sparse_vector()
2064 .unwrap();
2065 assert_eq!(entries2[0].0, 0);
2066 assert!((entries2[0].1 - 1.0).abs() < 1e-6);
2067 assert_eq!(entries2[1].0, 5);
2068 assert!((entries2[1].1 - 2.5).abs() < 1e-6);
2069 assert_eq!(entries2[2].0, 10);
2070 assert!((entries2[2].1 - 0.5).abs() < 1e-6);
2071 }
2072
2073 #[test]
2074 fn test_json_field() {
2075 let mut builder = Schema::builder();
2076 let metadata = builder.add_json_field("metadata", true);
2077 let title = builder.add_text_field("title", true, true);
2078 let schema = builder.build();
2079
2080 assert_eq!(schema.get_field("metadata"), Some(metadata));
2081 assert_eq!(
2082 schema.get_field_entry(metadata).unwrap().field_type,
2083 FieldType::Json
2084 );
2085 assert!(!schema.get_field_entry(metadata).unwrap().indexed);
2087 assert!(schema.get_field_entry(metadata).unwrap().stored);
2088
2089 let json_value = serde_json::json!({
2091 "author": "John Doe",
2092 "tags": ["rust", "search"],
2093 "nested": {"key": "value"}
2094 });
2095 let mut doc = Document::new();
2096 doc.add_json(metadata, json_value.clone());
2097 doc.add_text(title, "Test Document");
2098
2099 let stored_json = doc.get_first(metadata).unwrap().as_json().unwrap();
2101 assert_eq!(stored_json, &json_value);
2102 assert_eq!(
2103 stored_json.get("author").unwrap().as_str(),
2104 Some("John Doe")
2105 );
2106
2107 let doc_json = doc.to_json(&schema);
2109 let metadata_out = doc_json.get("metadata").unwrap();
2110 assert_eq!(metadata_out, &json_value);
2111
2112 let doc2 = Document::from_json(&doc_json, &schema).unwrap();
2114 let stored_json2 = doc2.get_first(metadata).unwrap().as_json().unwrap();
2115 assert_eq!(stored_json2, &json_value);
2116 }
2117
2118 #[test]
2119 fn test_json_field_various_types() {
2120 let mut builder = Schema::builder();
2121 let data = builder.add_json_field("data", true);
2122 let _schema = builder.build();
2123
2124 let arr_value = serde_json::json!([1, 2, 3, "four", null]);
2126 let mut doc = Document::new();
2127 doc.add_json(data, arr_value.clone());
2128 assert_eq!(doc.get_first(data).unwrap().as_json().unwrap(), &arr_value);
2129
2130 let str_value = serde_json::json!("just a string");
2132 let mut doc2 = Document::new();
2133 doc2.add_json(data, str_value.clone());
2134 assert_eq!(doc2.get_first(data).unwrap().as_json().unwrap(), &str_value);
2135
2136 let num_value = serde_json::json!(42.5);
2138 let mut doc3 = Document::new();
2139 doc3.add_json(data, num_value.clone());
2140 assert_eq!(doc3.get_first(data).unwrap().as_json().unwrap(), &num_value);
2141
2142 let null_value = serde_json::Value::Null;
2144 let mut doc4 = Document::new();
2145 doc4.add_json(data, null_value.clone());
2146 assert_eq!(
2147 doc4.get_first(data).unwrap().as_json().unwrap(),
2148 &null_value
2149 );
2150
2151 let bool_value = serde_json::json!(true);
2153 let mut doc5 = Document::new();
2154 doc5.add_json(data, bool_value.clone());
2155 assert_eq!(
2156 doc5.get_first(data).unwrap().as_json().unwrap(),
2157 &bool_value
2158 );
2159 }
2160}