1use chroma_error::{ChromaError, ErrorCodes};
2use regex::Regex;
3use serde::{Deserialize, Serialize};
4use std::collections::HashMap;
5use std::sync::{Arc, LazyLock};
6use thiserror::Error;
7use validator::Validate;
8
9use crate::chroma_proto;
10use crate::collection_configuration::{
11 EmbeddingFunctionConfiguration, InternalCollectionConfiguration,
12 UpdateVectorIndexConfiguration, VectorIndexConfiguration,
13};
14use crate::hnsw_configuration::Space;
15use crate::metadata::{MetadataComparison, MetadataValueType, Where};
16use crate::operator::QueryVector;
17use crate::{
18 default_batch_size, default_center_drift_threshold, default_construction_ef,
19 default_construction_ef_spann, default_initial_lambda, default_m, default_m_spann,
20 default_merge_threshold, default_nreplica_count, default_num_centers_to_merge_to,
21 default_num_samples_kmeans, default_num_threads, default_reassign_neighbor_count,
22 default_resize_factor, default_search_ef, default_search_ef_spann, default_search_nprobe,
23 default_search_rng_epsilon, default_search_rng_factor, default_space, default_split_threshold,
24 default_sync_threshold, default_write_nprobe, default_write_rng_epsilon,
25 default_write_rng_factor, ConversionError, HnswParametersFromSegmentError,
26 InternalHnswConfiguration, InternalSpannConfiguration, InternalUpdateCollectionConfiguration,
27 KnnIndex, Segment, UpdateCollectionConfiguration, CHROMA_KEY,
28};
29
30impl ChromaError for SchemaError {
31 fn code(&self) -> ErrorCodes {
32 match self {
33 SchemaError::MissingIndexConfiguration { .. } => ErrorCodes::Internal,
36 SchemaError::InvalidSchema { .. } => ErrorCodes::Internal,
37 SchemaError::DefaultsMismatch => ErrorCodes::Internal,
40 SchemaError::ConfigurationConflict { .. } => ErrorCodes::Internal,
41 SchemaError::InvalidConfigurationUpdate { .. } => ErrorCodes::Internal,
42
43 SchemaError::InvalidUserInput { .. } => ErrorCodes::InvalidArgument,
46 SchemaError::ConfigAndSchemaConflict => ErrorCodes::InvalidArgument,
47 SchemaError::InvalidHnswConfig(_) => ErrorCodes::InvalidArgument,
48 SchemaError::InvalidSpannConfig(_) => ErrorCodes::InvalidArgument,
49 SchemaError::Builder(e) => e.code(),
50 }
51 }
52}
53
54#[derive(Debug, Error)]
55pub enum SchemaError {
56 #[error("Schema is malformed: missing index configuration for metadata key '{key}' with type '{value_type}'")]
57 MissingIndexConfiguration { key: String, value_type: String },
58 #[error("Schema reconciliation failed: {reason}")]
59 InvalidSchema { reason: String },
60 #[error("Cannot set both collection config and schema simultaneously")]
61 ConfigAndSchemaConflict,
62 #[error("Cannot merge schemas with differing defaults")]
63 DefaultsMismatch,
64 #[error("Conflicting configuration for {context}")]
65 ConfigurationConflict { context: String },
66 #[error("Invalid HNSW configuration: {0}")]
67 InvalidHnswConfig(validator::ValidationErrors),
68 #[error("Invalid SPANN configuration: {0}")]
69 InvalidSpannConfig(validator::ValidationErrors),
70 #[error("Invalid schema input: {reason}")]
71 InvalidUserInput { reason: String },
72 #[error("Invalid configuration update: {message}")]
73 InvalidConfigurationUpdate { message: String },
74 #[error(transparent)]
75 Builder(#[from] SchemaBuilderError),
76}
77
78#[derive(Debug, Error)]
79pub enum SchemaBuilderError {
80 #[error("Vector index must be configured globally using create_index(None, config), not on specific key '{key}'")]
81 VectorIndexMustBeGlobal { key: String },
82 #[error("Cannot modify special key '{key}' - it is managed automatically by the system.")]
83 SpecialKeyModificationNotAllowed { key: String },
84 #[error("Sparse vector index requires a specific key. Use create_index(Some(\"key_name\"), config) instead of create_index(None, config)")]
85 SparseVectorRequiresKey,
86 #[error("Vector index deletion not supported. The vector index is always enabled on #embedding. To disable vector search, disable the collection instead.")]
87 VectorIndexDeletionNotSupported,
88 #[error("Sparse vector index deletion not supported yet. Sparse vector indexes cannot be removed once created.")]
89 SparseVectorIndexDeletionNotSupported,
90 #[error(
91 "Key '{key}' cannot begin with '#'. Keys starting with '#' are reserved for system use."
92 )]
93 ReservedKeyPrefix { key: String },
94 #[error("FTS index deletion is only supported on #document key.")]
95 FtsIndexDeletionOnlyOnDocument,
96 #[error("FTS index can only be enabled on #document key. Use create_index(Some(\"#document\"), FtsIndexConfig) to enable FTS.")]
97 FtsIndexOnlyOnDocument,
98}
99
100#[derive(Debug, Error)]
101pub enum FilterValidationError {
102 #[error(
103 "Cannot filter using metadata key '{key}' with type '{value_type:?}' because indexing is disabled"
104 )]
105 IndexingDisabled {
106 key: String,
107 value_type: MetadataValueType,
108 },
109 #[error("Cannot filter using full-text search because FTS indexing is disabled")]
110 FtsDisabled,
111 #[error(transparent)]
112 Schema(#[from] SchemaError),
113}
114
115impl ChromaError for SchemaBuilderError {
116 fn code(&self) -> ErrorCodes {
117 ErrorCodes::InvalidArgument
118 }
119}
120
121impl ChromaError for FilterValidationError {
122 fn code(&self) -> ErrorCodes {
123 match self {
124 FilterValidationError::IndexingDisabled { .. } => ErrorCodes::InvalidArgument,
125 FilterValidationError::FtsDisabled => ErrorCodes::InvalidArgument,
126 FilterValidationError::Schema(_) => ErrorCodes::Internal,
127 }
128 }
129}
130
131pub const STRING_VALUE_NAME: &str = "string";
138pub const INT_VALUE_NAME: &str = "int";
139pub const BOOL_VALUE_NAME: &str = "bool";
140pub const FLOAT_VALUE_NAME: &str = "float";
141pub const FLOAT_LIST_VALUE_NAME: &str = "float_list";
142pub const SPARSE_VECTOR_VALUE_NAME: &str = "sparse_vector";
143
144pub const FTS_INDEX_NAME: &str = "fts_index";
146pub const VECTOR_INDEX_NAME: &str = "vector_index";
147pub const SPARSE_VECTOR_INDEX_NAME: &str = "sparse_vector_index";
148pub const STRING_INVERTED_INDEX_NAME: &str = "string_inverted_index";
149pub const INT_INVERTED_INDEX_NAME: &str = "int_inverted_index";
150pub const FLOAT_INVERTED_INDEX_NAME: &str = "float_inverted_index";
151pub const BOOL_INVERTED_INDEX_NAME: &str = "bool_inverted_index";
152
153pub const DOCUMENT_KEY: &str = "#document";
155pub const EMBEDDING_KEY: &str = "#embedding";
156
157static CMEK_GCP_RE: LazyLock<Regex> = LazyLock::new(|| {
159 Regex::new(r"^projects/.+/locations/.+/keyRings/.+/cryptoKeys/.+$")
160 .expect("The CMEK pattern for GCP should be valid")
161});
162
163#[derive(Clone, Debug, Eq, PartialEq, serde::Deserialize, serde::Serialize)]
168#[serde(rename_all = "snake_case")]
169pub enum Cmek {
170 Gcp(Arc<String>),
174}
175
176impl Cmek {
177 pub fn gcp(resource: String) -> Self {
187 Cmek::Gcp(Arc::new(resource))
188 }
189
190 pub fn validate_pattern(&self) -> bool {
196 match self {
197 Cmek::Gcp(resource) => CMEK_GCP_RE.is_match(resource),
198 }
199 }
200}
201
202impl TryFrom<chroma_proto::Cmek> for Cmek {
203 type Error = ConversionError;
204
205 fn try_from(proto: chroma_proto::Cmek) -> Result<Self, Self::Error> {
206 match proto.provider {
207 Some(chroma_proto::cmek::Provider::Gcp(resource)) => Ok(Cmek::gcp(resource)),
208 None => Err(ConversionError::DecodeError),
209 }
210 }
211}
212
213impl From<Cmek> for chroma_proto::Cmek {
214 fn from(cmek: Cmek) -> Self {
215 match cmek {
216 Cmek::Gcp(resource) => chroma_proto::Cmek {
217 provider: Some(chroma_proto::cmek::Provider::Gcp((*resource).clone())),
218 },
219 }
220 }
221}
222
223#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
232#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
233pub struct Schema {
234 pub defaults: ValueTypes,
236 #[serde(rename = "keys", alias = "key_overrides")]
239 pub keys: HashMap<String, ValueTypes>,
240 #[serde(skip_serializing_if = "Option::is_none")]
242 #[cfg_attr(feature = "utoipa", schema(value_type = Option<Object>))]
243 pub cmek: Option<Cmek>,
244}
245
246impl Schema {
247 pub fn update(&mut self, configuration: &InternalUpdateCollectionConfiguration) {
248 if let Some(vector_update) = &configuration.vector_index {
249 if let Some(default_vector_index) = self.defaults_vector_index_mut() {
250 Self::apply_vector_index_update(default_vector_index, vector_update);
251 }
252 if let Some(embedding_vector_index) = self.embedding_vector_index_mut() {
253 Self::apply_vector_index_update(embedding_vector_index, vector_update);
254 }
255 }
256
257 if let Some(embedding_function) = configuration.embedding_function.as_ref() {
258 if let Some(default_vector_index) = self.defaults_vector_index_mut() {
259 default_vector_index.config.embedding_function = Some(embedding_function.clone());
260 }
261 if let Some(embedding_vector_index) = self.embedding_vector_index_mut() {
262 embedding_vector_index.config.embedding_function = Some(embedding_function.clone());
263 }
264 }
265 }
266
267 pub fn apply_update_configuration(
277 &mut self,
278 config: &UpdateCollectionConfiguration,
279 ) -> Result<(), SchemaError> {
280 if config.hnsw.is_some() {
282 return Err(SchemaError::InvalidConfigurationUpdate {
283 message: "HNSW configuration updates are not supported".to_string(),
284 });
285 }
286
287 if let Some(ref spann_update) = config.spann {
289 let defaults_spann = self
290 .defaults_vector_index_mut()
291 .ok_or_else(|| SchemaError::InvalidConfigurationUpdate {
292 message: "schema missing defaults.float_list.vector_index".to_string(),
293 })?
294 .config
295 .spann
296 .as_mut()
297 .ok_or_else(|| SchemaError::InvalidConfigurationUpdate {
298 message: "schema missing defaults spann config".to_string(),
299 })?;
300
301 if let Some(search_nprobe) = spann_update.search_nprobe {
302 defaults_spann.search_nprobe = Some(search_nprobe);
303 }
304 if let Some(ef_search) = spann_update.ef_search {
305 defaults_spann.ef_search = Some(ef_search);
306 }
307
308 let embedding_spann = self
309 .embedding_vector_index_mut()
310 .ok_or_else(|| SchemaError::InvalidConfigurationUpdate {
311 message: "schema missing keys[#embedding].float_list.vector_index".to_string(),
312 })?
313 .config
314 .spann
315 .as_mut()
316 .ok_or_else(|| SchemaError::InvalidConfigurationUpdate {
317 message: "schema missing #embedding spann config".to_string(),
318 })?;
319
320 if let Some(search_nprobe) = spann_update.search_nprobe {
321 embedding_spann.search_nprobe = Some(search_nprobe);
322 }
323 if let Some(ef_search) = spann_update.ef_search {
324 embedding_spann.ef_search = Some(ef_search);
325 }
326 }
327
328 if let Some(ref ef) = config.embedding_function {
330 self.defaults_vector_index_mut()
331 .ok_or_else(|| SchemaError::InvalidConfigurationUpdate {
332 message: "schema missing defaults.float_list.vector_index".to_string(),
333 })?
334 .config
335 .embedding_function = Some(ef.clone());
336
337 self.embedding_vector_index_mut()
338 .ok_or_else(|| SchemaError::InvalidConfigurationUpdate {
339 message: "schema missing keys[#embedding].float_list.vector_index".to_string(),
340 })?
341 .config
342 .embedding_function = Some(ef.clone());
343 }
344
345 Ok(())
346 }
347
348 fn defaults_vector_index_mut(&mut self) -> Option<&mut VectorIndexType> {
349 self.defaults
350 .float_list
351 .as_mut()
352 .and_then(|float_list| float_list.vector_index.as_mut())
353 }
354
355 fn embedding_vector_index_mut(&mut self) -> Option<&mut VectorIndexType> {
356 self.keys
357 .get_mut(EMBEDDING_KEY)
358 .and_then(|value_types| value_types.float_list.as_mut())
359 .and_then(|float_list| float_list.vector_index.as_mut())
360 }
361
362 fn apply_vector_index_update(
363 vector_index: &mut VectorIndexType,
364 update: &UpdateVectorIndexConfiguration,
365 ) {
366 match update {
367 UpdateVectorIndexConfiguration::Hnsw(Some(hnsw_update)) => {
368 if let Some(hnsw_config) = vector_index.config.hnsw.as_mut() {
369 if let Some(ef_search) = hnsw_update.ef_search {
370 hnsw_config.ef_search = Some(ef_search);
371 }
372 if let Some(max_neighbors) = hnsw_update.max_neighbors {
373 hnsw_config.max_neighbors = Some(max_neighbors);
374 }
375 if let Some(num_threads) = hnsw_update.num_threads {
376 hnsw_config.num_threads = Some(num_threads);
377 }
378 if let Some(resize_factor) = hnsw_update.resize_factor {
379 hnsw_config.resize_factor = Some(resize_factor);
380 }
381 if let Some(sync_threshold) = hnsw_update.sync_threshold {
382 hnsw_config.sync_threshold = Some(sync_threshold);
383 }
384 if let Some(batch_size) = hnsw_update.batch_size {
385 hnsw_config.batch_size = Some(batch_size);
386 }
387 }
388 }
389 UpdateVectorIndexConfiguration::Hnsw(None) => {}
390 UpdateVectorIndexConfiguration::Spann(Some(spann_update)) => {
391 if let Some(spann_config) = vector_index.config.spann.as_mut() {
392 if let Some(search_nprobe) = spann_update.search_nprobe {
393 spann_config.search_nprobe = Some(search_nprobe);
394 }
395 if let Some(ef_search) = spann_update.ef_search {
396 spann_config.ef_search = Some(ef_search);
397 }
398 }
399 }
400 UpdateVectorIndexConfiguration::Spann(None) => {}
401 }
402 }
403
404 pub fn is_sparse_index_enabled(&self) -> bool {
405 let defaults_enabled = self
406 .defaults
407 .sparse_vector
408 .as_ref()
409 .and_then(|sv| sv.sparse_vector_index.as_ref())
410 .is_some_and(|idx| idx.enabled);
411 let key_enabled = self.keys.values().any(|value_types| {
412 value_types
413 .sparse_vector
414 .as_ref()
415 .and_then(|sv| sv.sparse_vector_index.as_ref())
416 .is_some_and(|idx| idx.enabled)
417 });
418 defaults_enabled || key_enabled
419 }
420
421 pub fn is_maxscore_enabled(&self) -> bool {
423 let check = |sv: &SparseVectorValueType| -> bool {
424 sv.sparse_vector_index.as_ref().is_some_and(|idx| {
425 idx.enabled && matches!(idx.config.algorithm, SparseIndexAlgorithm::MaxScore)
426 })
427 };
428 self.defaults.sparse_vector.as_ref().is_some_and(check)
429 || self
430 .keys
431 .values()
432 .any(|vt| vt.sparse_vector.as_ref().is_some_and(check))
433 }
434
435 pub fn enabled_sparse_keys(&self) -> Vec<String> {
438 let mut keys: Vec<String> = self
439 .keys
440 .iter()
441 .filter(|(_, value_types)| {
442 value_types
443 .sparse_vector
444 .as_ref()
445 .and_then(|sv| sv.sparse_vector_index.as_ref())
446 .is_some_and(|idx| idx.enabled)
447 })
448 .map(|(key, _)| key.clone())
449 .collect();
450 keys.sort();
451 keys
452 }
453
454 pub fn is_key_maxscore_enabled(&self, key: &str) -> bool {
458 let algorithm = self
459 .keys
460 .get(key)
461 .and_then(|vt| vt.sparse_vector.as_ref())
462 .and_then(|sv| sv.sparse_vector_index.as_ref())
463 .map(|idx| &idx.config.algorithm)
464 .or_else(|| {
465 self.defaults
466 .sparse_vector
467 .as_ref()
468 .and_then(|sv| sv.sparse_vector_index.as_ref())
469 .map(|idx| &idx.config.algorithm)
470 });
471 matches!(algorithm, Some(SparseIndexAlgorithm::MaxScore))
472 }
473
474 pub fn set_sparse_algorithm(&mut self, algorithm: SparseIndexAlgorithm) {
477 if let Some(sv) = &mut self.defaults.sparse_vector {
478 if let Some(idx) = &mut sv.sparse_vector_index {
479 idx.config.algorithm = algorithm.clone();
480 }
481 }
482 for vt in self.keys.values_mut() {
483 if let Some(sv) = &mut vt.sparse_vector {
484 if let Some(idx) = &mut sv.sparse_vector_index {
485 idx.config.algorithm = algorithm.clone();
486 }
487 }
488 }
489 }
490
491 pub fn is_fts_enabled(&self) -> bool {
492 self.keys
494 .get(DOCUMENT_KEY)
495 .and_then(|vt| vt.string.as_ref())
496 .and_then(|s| s.fts_index.as_ref())
497 .or_else(|| {
498 self.defaults
499 .string
500 .as_ref()
501 .and_then(|s| s.fts_index.as_ref())
502 })
503 .is_none_or(|idx| idx.enabled)
504 }
505
506 pub fn is_token_bitmap_fts_enabled(&self) -> bool {
508 let check = |s: &StringValueType| -> bool {
509 s.fts_index.as_ref().is_some_and(|idx| {
510 idx.enabled && matches!(idx.config.algorithm, FtsAlgorithm::TokenBitmap)
511 })
512 };
513 self.keys
514 .get(DOCUMENT_KEY)
515 .and_then(|vt| vt.string.as_ref())
516 .is_some_and(check)
517 || self.defaults.string.as_ref().is_some_and(check)
518 }
519
520 pub fn set_fts_algorithm(&mut self, algorithm: FtsAlgorithm) {
523 if let Some(s) = &mut self.defaults.string {
524 if let Some(idx) = &mut s.fts_index {
525 idx.config.algorithm = algorithm.clone();
526 }
527 }
528 for vt in self.keys.values_mut() {
529 if let Some(s) = &mut vt.string {
530 if let Some(idx) = &mut s.fts_index {
531 idx.config.algorithm = algorithm.clone();
532 }
533 }
534 }
535 }
536}
537
538impl Default for Schema {
539 fn default() -> Self {
556 let defaults = ValueTypes {
558 string: Some(StringValueType {
559 fts_index: Some(FtsIndexType {
560 enabled: false,
561 config: FtsIndexConfig::default(),
562 }),
563 string_inverted_index: Some(StringInvertedIndexType {
564 enabled: true,
565 config: StringInvertedIndexConfig {},
566 }),
567 }),
568 float_list: Some(FloatListValueType {
569 vector_index: Some(VectorIndexType {
570 enabled: false,
571 config: VectorIndexConfig {
572 space: None, embedding_function: Some(EmbeddingFunctionConfiguration::Legacy),
574 source_key: None,
575 hnsw: None, spann: None, },
578 }),
579 }),
580 sparse_vector: Some(SparseVectorValueType {
581 sparse_vector_index: Some(SparseVectorIndexType {
582 enabled: false,
583 config: SparseVectorIndexConfig {
584 embedding_function: None,
585 source_key: None,
586 bm25: None,
587 algorithm: SparseIndexAlgorithm::Wand,
588 },
589 }),
590 }),
591 int: Some(IntValueType {
592 int_inverted_index: Some(IntInvertedIndexType {
593 enabled: true,
594 config: IntInvertedIndexConfig {},
595 }),
596 }),
597 float: Some(FloatValueType {
598 float_inverted_index: Some(FloatInvertedIndexType {
599 enabled: true,
600 config: FloatInvertedIndexConfig {},
601 }),
602 }),
603 boolean: Some(BoolValueType {
604 bool_inverted_index: Some(BoolInvertedIndexType {
605 enabled: true,
606 config: BoolInvertedIndexConfig {},
607 }),
608 }),
609 };
610
611 let mut keys = HashMap::new();
613
614 keys.insert(
616 DOCUMENT_KEY.to_string(),
617 ValueTypes {
618 string: Some(StringValueType {
619 fts_index: Some(FtsIndexType {
620 enabled: true,
621 config: FtsIndexConfig::default(),
622 }),
623 string_inverted_index: Some(StringInvertedIndexType {
624 enabled: false,
625 config: StringInvertedIndexConfig {},
626 }),
627 }),
628 ..Default::default()
629 },
630 );
631
632 keys.insert(
634 EMBEDDING_KEY.to_string(),
635 ValueTypes {
636 float_list: Some(FloatListValueType {
637 vector_index: Some(VectorIndexType {
638 enabled: true,
639 config: VectorIndexConfig {
640 space: None, embedding_function: Some(EmbeddingFunctionConfiguration::Legacy),
642 source_key: Some(DOCUMENT_KEY.to_string()),
643 hnsw: None, spann: None, },
646 }),
647 }),
648 ..Default::default()
649 },
650 );
651
652 Schema {
653 defaults,
654 keys,
655 cmek: None,
656 }
657 }
658}
659
660pub fn is_embedding_function_default(
661 embedding_function: &Option<EmbeddingFunctionConfiguration>,
662) -> bool {
663 match embedding_function {
664 None => true,
665 Some(embedding_function) => embedding_function.is_default(),
666 }
667}
668
669pub fn is_space_default(space: &Option<Space>) -> bool {
671 match space {
672 None => true, Some(s) => *s == default_space(), }
675}
676
677pub fn is_hnsw_config_default(hnsw_config: &HnswIndexConfig) -> bool {
679 hnsw_config.ef_construction == Some(default_construction_ef())
680 && hnsw_config.ef_search == Some(default_search_ef())
681 && hnsw_config.max_neighbors == Some(default_m())
682 && hnsw_config.num_threads == Some(default_num_threads())
683 && hnsw_config.batch_size == Some(default_batch_size())
684 && hnsw_config.sync_threshold == Some(default_sync_threshold())
685 && hnsw_config.resize_factor == Some(default_resize_factor())
686}
687
688#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Default)]
695#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
696pub struct ValueTypes {
697 #[serde(
698 rename = "string",
699 alias = "#string",
700 skip_serializing_if = "Option::is_none"
701 )] pub string: Option<StringValueType>,
703
704 #[serde(
705 rename = "float_list",
706 alias = "#float_list",
707 skip_serializing_if = "Option::is_none"
708 )]
709 pub float_list: Option<FloatListValueType>,
711
712 #[serde(
713 rename = "sparse_vector",
714 alias = "#sparse_vector",
715 skip_serializing_if = "Option::is_none"
716 )]
717 pub sparse_vector: Option<SparseVectorValueType>,
719
720 #[serde(
721 rename = "int",
722 alias = "#int",
723 skip_serializing_if = "Option::is_none"
724 )] pub int: Option<IntValueType>,
726
727 #[serde(
728 rename = "float",
729 alias = "#float",
730 skip_serializing_if = "Option::is_none"
731 )] pub float: Option<FloatValueType>,
733
734 #[serde(
735 rename = "bool",
736 alias = "#bool",
737 skip_serializing_if = "Option::is_none"
738 )] pub boolean: Option<BoolValueType>,
740}
741
742#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
744#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
745pub struct StringValueType {
746 #[serde(
747 rename = "fts_index",
748 alias = "$fts_index",
749 skip_serializing_if = "Option::is_none"
750 )] pub fts_index: Option<FtsIndexType>,
752
753 #[serde(
754 rename = "string_inverted_index", alias = "$string_inverted_index",
756 skip_serializing_if = "Option::is_none"
757 )]
758 pub string_inverted_index: Option<StringInvertedIndexType>,
759}
760
761#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
763#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
764pub struct FloatListValueType {
765 #[serde(
766 rename = "vector_index",
767 alias = "$vector_index",
768 skip_serializing_if = "Option::is_none"
769 )] pub vector_index: Option<VectorIndexType>,
771}
772
773#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
775#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
776pub struct SparseVectorValueType {
777 #[serde(
778 rename = "sparse_vector_index", alias = "$sparse_vector_index",
780 skip_serializing_if = "Option::is_none"
781 )]
782 pub sparse_vector_index: Option<SparseVectorIndexType>,
783}
784
785#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
787#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
788pub struct IntValueType {
789 #[serde(
790 rename = "int_inverted_index",
791 alias = "$int_inverted_index",
792 skip_serializing_if = "Option::is_none"
793 )]
794 pub int_inverted_index: Option<IntInvertedIndexType>,
796}
797
798#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
800#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
801pub struct FloatValueType {
802 #[serde(
803 rename = "float_inverted_index", alias = "$float_inverted_index",
805 skip_serializing_if = "Option::is_none"
806 )]
807 pub float_inverted_index: Option<FloatInvertedIndexType>,
808}
809
810#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
812#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
813pub struct BoolValueType {
814 #[serde(
815 rename = "bool_inverted_index", alias = "$bool_inverted_index",
817 skip_serializing_if = "Option::is_none"
818 )]
819 pub bool_inverted_index: Option<BoolInvertedIndexType>,
820}
821
822#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
824#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
825pub struct FtsIndexType {
826 pub enabled: bool,
827 pub config: FtsIndexConfig,
828}
829
830#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
831#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
832pub struct VectorIndexType {
833 pub enabled: bool,
834 pub config: VectorIndexConfig,
835}
836
837#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
838#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
839pub struct SparseVectorIndexType {
840 pub enabled: bool,
841 pub config: SparseVectorIndexConfig,
842}
843
844#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
845#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
846pub struct StringInvertedIndexType {
847 pub enabled: bool,
848 pub config: StringInvertedIndexConfig,
849}
850
851#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
852#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
853pub struct IntInvertedIndexType {
854 pub enabled: bool,
855 pub config: IntInvertedIndexConfig,
856}
857
858#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
859#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
860pub struct FloatInvertedIndexType {
861 pub enabled: bool,
862 pub config: FloatInvertedIndexConfig,
863}
864
865#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
866#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
867pub struct BoolInvertedIndexType {
868 pub enabled: bool,
869 pub config: BoolInvertedIndexConfig,
870}
871
872impl Schema {
873 pub fn new_default(default_knn_index: KnnIndex) -> Self {
875 let vector_config = VectorIndexType {
877 enabled: false,
878 config: VectorIndexConfig {
879 space: Some(default_space()),
880 embedding_function: None,
881 source_key: None,
882 hnsw: match default_knn_index {
883 KnnIndex::Hnsw => Some(HnswIndexConfig {
884 ef_construction: Some(default_construction_ef()),
885 max_neighbors: Some(default_m()),
886 ef_search: Some(default_search_ef()),
887 num_threads: Some(default_num_threads()),
888 batch_size: Some(default_batch_size()),
889 sync_threshold: Some(default_sync_threshold()),
890 resize_factor: Some(default_resize_factor()),
891 }),
892 KnnIndex::Spann => None,
893 },
894 spann: match default_knn_index {
895 KnnIndex::Hnsw => None,
896 KnnIndex::Spann => Some(SpannIndexConfig {
897 search_nprobe: Some(default_search_nprobe()),
898 search_rng_factor: Some(default_search_rng_factor()),
899 search_rng_epsilon: Some(default_search_rng_epsilon()),
900 nreplica_count: Some(default_nreplica_count()),
901 write_rng_factor: Some(default_write_rng_factor()),
902 write_rng_epsilon: Some(default_write_rng_epsilon()),
903 split_threshold: Some(default_split_threshold()),
904 num_samples_kmeans: Some(default_num_samples_kmeans()),
905 initial_lambda: Some(default_initial_lambda()),
906 reassign_neighbor_count: Some(default_reassign_neighbor_count()),
907 merge_threshold: Some(default_merge_threshold()),
908 num_centers_to_merge_to: Some(default_num_centers_to_merge_to()),
909 write_nprobe: Some(default_write_nprobe()),
910 ef_construction: Some(default_construction_ef_spann()),
911 ef_search: Some(default_search_ef_spann()),
912 max_neighbors: Some(default_m_spann()),
913 center_drift_threshold: None,
914 quantize: Quantization::None,
915 }),
916 },
917 },
918 };
919
920 let defaults = ValueTypes {
922 string: Some(StringValueType {
923 string_inverted_index: Some(StringInvertedIndexType {
924 enabled: true,
925 config: StringInvertedIndexConfig {},
926 }),
927 fts_index: Some(FtsIndexType {
928 enabled: false,
929 config: FtsIndexConfig::default(),
930 }),
931 }),
932 float: Some(FloatValueType {
933 float_inverted_index: Some(FloatInvertedIndexType {
934 enabled: true,
935 config: FloatInvertedIndexConfig {},
936 }),
937 }),
938 int: Some(IntValueType {
939 int_inverted_index: Some(IntInvertedIndexType {
940 enabled: true,
941 config: IntInvertedIndexConfig {},
942 }),
943 }),
944 boolean: Some(BoolValueType {
945 bool_inverted_index: Some(BoolInvertedIndexType {
946 enabled: true,
947 config: BoolInvertedIndexConfig {},
948 }),
949 }),
950 float_list: Some(FloatListValueType {
951 vector_index: Some(vector_config),
952 }),
953 sparse_vector: Some(SparseVectorValueType {
954 sparse_vector_index: Some(SparseVectorIndexType {
955 enabled: false,
956 config: SparseVectorIndexConfig {
957 embedding_function: Some(EmbeddingFunctionConfiguration::Unknown),
958 source_key: None,
959 bm25: Some(false),
960 algorithm: SparseIndexAlgorithm::Wand,
961 },
962 }),
963 }),
964 };
965
966 let mut keys = HashMap::new();
968
969 let embedding_defaults = ValueTypes {
971 float_list: Some(FloatListValueType {
972 vector_index: Some(VectorIndexType {
973 enabled: true,
974 config: VectorIndexConfig {
975 space: Some(default_space()),
976 embedding_function: None,
977 source_key: Some(DOCUMENT_KEY.to_string()),
978 hnsw: match default_knn_index {
979 KnnIndex::Hnsw => Some(HnswIndexConfig {
980 ef_construction: Some(default_construction_ef()),
981 max_neighbors: Some(default_m()),
982 ef_search: Some(default_search_ef()),
983 num_threads: Some(default_num_threads()),
984 batch_size: Some(default_batch_size()),
985 sync_threshold: Some(default_sync_threshold()),
986 resize_factor: Some(default_resize_factor()),
987 }),
988 KnnIndex::Spann => None,
989 },
990 spann: match default_knn_index {
991 KnnIndex::Hnsw => None,
992 KnnIndex::Spann => Some(SpannIndexConfig {
993 search_nprobe: Some(default_search_nprobe()),
994 search_rng_factor: Some(default_search_rng_factor()),
995 search_rng_epsilon: Some(default_search_rng_epsilon()),
996 nreplica_count: Some(default_nreplica_count()),
997 write_rng_factor: Some(default_write_rng_factor()),
998 write_rng_epsilon: Some(default_write_rng_epsilon()),
999 split_threshold: Some(default_split_threshold()),
1000 num_samples_kmeans: Some(default_num_samples_kmeans()),
1001 initial_lambda: Some(default_initial_lambda()),
1002 reassign_neighbor_count: Some(default_reassign_neighbor_count()),
1003 merge_threshold: Some(default_merge_threshold()),
1004 num_centers_to_merge_to: Some(default_num_centers_to_merge_to()),
1005 write_nprobe: Some(default_write_nprobe()),
1006 ef_construction: Some(default_construction_ef_spann()),
1007 ef_search: Some(default_search_ef_spann()),
1008 max_neighbors: Some(default_m_spann()),
1009 center_drift_threshold: None,
1010 quantize: Quantization::None,
1011 }),
1012 },
1013 },
1014 }),
1015 }),
1016 ..Default::default()
1017 };
1018 keys.insert(EMBEDDING_KEY.to_string(), embedding_defaults);
1019
1020 let document_defaults = ValueTypes {
1022 string: Some(StringValueType {
1023 fts_index: Some(FtsIndexType {
1024 enabled: true,
1025 config: FtsIndexConfig::default(),
1026 }),
1027 string_inverted_index: Some(StringInvertedIndexType {
1028 enabled: false,
1029 config: StringInvertedIndexConfig {},
1030 }),
1031 }),
1032 ..Default::default()
1033 };
1034 keys.insert(DOCUMENT_KEY.to_string(), document_defaults);
1035
1036 Schema {
1037 defaults,
1038 keys,
1039 cmek: None,
1040 }
1041 }
1042
1043 pub fn new_record_only() -> Self {
1050 let defaults = ValueTypes {
1051 string: Some(StringValueType {
1052 string_inverted_index: Some(StringInvertedIndexType {
1053 enabled: false,
1054 config: StringInvertedIndexConfig {},
1055 }),
1056 fts_index: Some(FtsIndexType {
1057 enabled: false,
1058 config: FtsIndexConfig::default(),
1059 }),
1060 }),
1061 float: Some(FloatValueType {
1062 float_inverted_index: Some(FloatInvertedIndexType {
1063 enabled: false,
1064 config: FloatInvertedIndexConfig {},
1065 }),
1066 }),
1067 int: Some(IntValueType {
1068 int_inverted_index: Some(IntInvertedIndexType {
1069 enabled: false,
1070 config: IntInvertedIndexConfig {},
1071 }),
1072 }),
1073 boolean: Some(BoolValueType {
1074 bool_inverted_index: Some(BoolInvertedIndexType {
1075 enabled: false,
1076 config: BoolInvertedIndexConfig {},
1077 }),
1078 }),
1079 float_list: Some(FloatListValueType {
1080 vector_index: Some(VectorIndexType {
1081 enabled: false,
1082 config: VectorIndexConfig {
1083 space: Some(default_space()),
1084 embedding_function: None,
1085 source_key: None,
1086 hnsw: Some(HnswIndexConfig {
1087 ef_construction: Some(default_construction_ef()),
1088 max_neighbors: Some(default_m()),
1089 ef_search: Some(default_search_ef()),
1090 num_threads: Some(default_num_threads()),
1091 batch_size: Some(default_batch_size()),
1092 sync_threshold: Some(default_sync_threshold()),
1093 resize_factor: Some(default_resize_factor()),
1094 }),
1095 spann: None,
1096 },
1097 }),
1098 }),
1099 sparse_vector: Some(SparseVectorValueType {
1100 sparse_vector_index: Some(SparseVectorIndexType {
1101 enabled: false,
1102 config: SparseVectorIndexConfig {
1103 embedding_function: Some(EmbeddingFunctionConfiguration::Unknown),
1104 source_key: None,
1105 bm25: Some(false),
1106 algorithm: SparseIndexAlgorithm::Wand,
1107 },
1108 }),
1109 }),
1110 };
1111
1112 let mut keys = HashMap::new();
1113
1114 let embedding_defaults = ValueTypes {
1116 float_list: Some(FloatListValueType {
1117 vector_index: Some(VectorIndexType {
1118 enabled: false,
1119 config: VectorIndexConfig {
1120 space: Some(default_space()),
1121 embedding_function: None,
1122 source_key: Some(DOCUMENT_KEY.to_string()),
1123 hnsw: Some(HnswIndexConfig {
1124 ef_construction: Some(default_construction_ef()),
1125 max_neighbors: Some(default_m()),
1126 ef_search: Some(default_search_ef()),
1127 num_threads: Some(default_num_threads()),
1128 batch_size: Some(default_batch_size()),
1129 sync_threshold: Some(default_sync_threshold()),
1130 resize_factor: Some(default_resize_factor()),
1131 }),
1132 spann: None,
1133 },
1134 }),
1135 }),
1136 ..Default::default()
1137 };
1138 keys.insert(EMBEDDING_KEY.to_string(), embedding_defaults);
1139
1140 let document_defaults = ValueTypes {
1142 string: Some(StringValueType {
1143 fts_index: Some(FtsIndexType {
1144 enabled: false,
1145 config: FtsIndexConfig::default(),
1146 }),
1147 string_inverted_index: Some(StringInvertedIndexType {
1148 enabled: false,
1149 config: StringInvertedIndexConfig {},
1150 }),
1151 }),
1152 ..Default::default()
1153 };
1154 keys.insert(DOCUMENT_KEY.to_string(), document_defaults);
1155
1156 Schema {
1157 defaults,
1158 keys,
1159 cmek: None,
1160 }
1161 }
1162
1163 pub fn get_spann_config(&self) -> Option<(SpannIndexConfig, Space)> {
1164 let extract = |vector_index: &VectorIndexType| {
1165 let space = vector_index.config.space.clone().unwrap_or_default();
1166 vector_index
1167 .config
1168 .spann
1169 .clone()
1170 .map(|config| (config, space))
1171 };
1172
1173 self.keys
1174 .get(EMBEDDING_KEY)
1175 .and_then(|value_types| value_types.float_list.as_ref())
1176 .and_then(|float_list| float_list.vector_index.as_ref())
1177 .and_then(extract)
1178 .or_else(|| {
1179 self.defaults
1180 .float_list
1181 .as_ref()
1182 .and_then(|float_list| float_list.vector_index.as_ref())
1183 .and_then(extract)
1184 })
1185 }
1186
1187 pub fn get_internal_spann_config(&self) -> Option<InternalSpannConfiguration> {
1188 let to_internal = |vector_index: &VectorIndexType| {
1189 let space = vector_index.config.space.clone();
1190 vector_index
1191 .config
1192 .spann
1193 .clone()
1194 .map(|config| (space.as_ref(), &config).into())
1195 };
1196
1197 self.keys
1198 .get(EMBEDDING_KEY)
1199 .and_then(|value_types| value_types.float_list.as_ref())
1200 .and_then(|float_list| float_list.vector_index.as_ref())
1201 .and_then(to_internal)
1202 .or_else(|| {
1203 self.defaults
1204 .float_list
1205 .as_ref()
1206 .and_then(|float_list| float_list.vector_index.as_ref())
1207 .and_then(to_internal)
1208 })
1209 }
1210
1211 pub fn is_quantization_enabled(&self) -> bool {
1213 let check_spann = |vector_index: &VectorIndexType| {
1214 vector_index
1215 .config
1216 .spann
1217 .as_ref()
1218 .is_some_and(|config| !matches!(config.quantize, Quantization::None))
1219 };
1220
1221 self.keys
1222 .get(EMBEDDING_KEY)
1223 .and_then(|value_types| value_types.float_list.as_ref())
1224 .and_then(|float_list| float_list.vector_index.as_ref())
1225 .map(check_spann)
1226 .unwrap_or_else(|| {
1227 self.defaults
1228 .float_list
1229 .as_ref()
1230 .and_then(|float_list| float_list.vector_index.as_ref())
1231 .map(check_spann)
1232 .unwrap_or(false)
1233 })
1234 }
1235
1236 pub fn get_spann_config_mut(&mut self) -> Option<&mut SpannIndexConfig> {
1239 if let Some(value_types) = self.keys.get_mut(EMBEDDING_KEY) {
1241 if let Some(float_list) = &mut value_types.float_list {
1242 if let Some(vector_index) = &mut float_list.vector_index {
1243 if let Some(spann_config) = &mut vector_index.config.spann {
1244 return Some(spann_config);
1245 }
1246 }
1247 }
1248 }
1249
1250 if let Some(float_list) = &mut self.defaults.float_list {
1252 if let Some(vector_index) = &mut float_list.vector_index {
1253 if let Some(spann_config) = &mut vector_index.config.spann {
1254 return Some(spann_config);
1255 }
1256 }
1257 }
1258
1259 None
1260 }
1261
1262 pub fn quantize(&mut self, variant: Quantization) {
1268 if let Some(spann_config) = self.get_spann_config_mut() {
1269 *spann_config = match variant {
1270 Quantization::None => SpannIndexConfig {
1271 quantize: variant,
1272 ..*spann_config
1273 },
1274 Quantization::FourBitRabitQWithUSearch => SpannIndexConfig {
1275 search_nprobe: Some(64),
1276 nreplica_count: Some(2),
1277 write_rng_factor: Some(4.0),
1278 write_rng_epsilon: Some(8.0),
1279 split_threshold: Some(512),
1280 reassign_neighbor_count: Some(32),
1281 merge_threshold: Some(128),
1282 write_nprobe: Some(64),
1283 ef_construction: Some(256),
1284 ef_search: Some(128),
1285 max_neighbors: Some(24),
1286 center_drift_threshold: Some(0.125),
1287 quantize: variant,
1288 ..*spann_config
1289 },
1290 };
1291 }
1292 }
1293
1294 pub fn get_internal_hnsw_config(&self) -> Option<InternalHnswConfiguration> {
1295 let to_internal = |vector_index: &VectorIndexType| {
1296 if vector_index.config.spann.is_some() {
1297 return None;
1298 }
1299 let space = vector_index.config.space.as_ref();
1300 let hnsw_config = vector_index.config.hnsw.as_ref();
1301 Some((space, hnsw_config).into())
1302 };
1303
1304 self.keys
1305 .get(EMBEDDING_KEY)
1306 .and_then(|value_types| value_types.float_list.as_ref())
1307 .and_then(|float_list| float_list.vector_index.as_ref())
1308 .and_then(to_internal)
1309 .or_else(|| {
1310 self.defaults
1311 .float_list
1312 .as_ref()
1313 .and_then(|float_list| float_list.vector_index.as_ref())
1314 .and_then(to_internal)
1315 })
1316 }
1317
1318 pub fn get_internal_hnsw_config_with_legacy_fallback(
1319 &self,
1320 segment: &Segment,
1321 ) -> Result<Option<InternalHnswConfiguration>, HnswParametersFromSegmentError> {
1322 if let Some(config) = self.get_internal_hnsw_config() {
1323 let config_from_metadata =
1324 InternalHnswConfiguration::from_legacy_segment_metadata(&segment.metadata)?;
1325
1326 if config == InternalHnswConfiguration::default() && config != config_from_metadata {
1327 return Ok(Some(config_from_metadata));
1328 }
1329
1330 return Ok(Some(config));
1331 }
1332
1333 Ok(None)
1334 }
1335
1336 pub fn reconcile_with_defaults(
1343 user_schema: Option<&Schema>,
1344 knn_index: KnnIndex,
1345 ) -> Result<Self, SchemaError> {
1346 let default_schema = Schema::new_default(knn_index);
1347
1348 match user_schema {
1349 Some(user) => {
1350 let merged_defaults =
1352 Self::merge_value_types(&default_schema.defaults, &user.defaults, knn_index)?;
1353
1354 let mut merged_keys = default_schema.keys.clone();
1356 for (key, user_value_types) in &user.keys {
1357 if let Some(default_value_types) = merged_keys.get(key) {
1358 let merged_value_types = Self::merge_value_types(
1360 default_value_types,
1361 user_value_types,
1362 knn_index,
1363 )?;
1364 merged_keys.insert(key.clone(), merged_value_types);
1365 } else {
1366 merged_keys.insert(key.clone(), user_value_types.clone());
1368 }
1369 }
1370
1371 Ok(Schema {
1372 defaults: merged_defaults,
1373 keys: merged_keys,
1374 cmek: user.cmek.clone().or(default_schema.cmek.clone()),
1375 })
1376 }
1377 None => Ok(default_schema),
1378 }
1379 }
1380
1381 pub fn merge(&self, other: &Schema) -> Result<Schema, SchemaError> {
1383 if self.defaults != other.defaults {
1384 return Err(SchemaError::DefaultsMismatch);
1385 }
1386
1387 let mut keys = self.keys.clone();
1388
1389 for (key, other_value_types) in &other.keys {
1390 if let Some(existing) = keys.get(key).cloned() {
1391 let merged = Self::merge_override_value_types(key, &existing, other_value_types)?;
1392 keys.insert(key.clone(), merged);
1393 } else {
1394 keys.insert(key.clone(), other_value_types.clone());
1395 }
1396 }
1397
1398 Ok(Schema {
1399 defaults: self.defaults.clone(),
1400 keys,
1401 cmek: other.cmek.clone().or(self.cmek.clone()),
1402 })
1403 }
1404
1405 fn merge_override_value_types(
1406 key: &str,
1407 left: &ValueTypes,
1408 right: &ValueTypes,
1409 ) -> Result<ValueTypes, SchemaError> {
1410 Ok(ValueTypes {
1411 string: Self::merge_string_override(key, left.string.as_ref(), right.string.as_ref())?,
1412 float: Self::merge_float_override(key, left.float.as_ref(), right.float.as_ref())?,
1413 int: Self::merge_int_override(key, left.int.as_ref(), right.int.as_ref())?,
1414 boolean: Self::merge_bool_override(key, left.boolean.as_ref(), right.boolean.as_ref())?,
1415 float_list: Self::merge_float_list_override(
1416 key,
1417 left.float_list.as_ref(),
1418 right.float_list.as_ref(),
1419 )?,
1420 sparse_vector: Self::merge_sparse_vector_override(
1421 key,
1422 left.sparse_vector.as_ref(),
1423 right.sparse_vector.as_ref(),
1424 )?,
1425 })
1426 }
1427
1428 fn merge_string_override(
1429 key: &str,
1430 left: Option<&StringValueType>,
1431 right: Option<&StringValueType>,
1432 ) -> Result<Option<StringValueType>, SchemaError> {
1433 match (left, right) {
1434 (Some(l), Some(r)) => Ok(Some(StringValueType {
1435 string_inverted_index: Self::merge_index_or_error(
1436 l.string_inverted_index.as_ref(),
1437 r.string_inverted_index.as_ref(),
1438 &format!("key '{key}' string.string_inverted_index"),
1439 )?,
1440 fts_index: Self::merge_index_or_error(
1441 l.fts_index.as_ref(),
1442 r.fts_index.as_ref(),
1443 &format!("key '{key}' string.fts_index"),
1444 )?,
1445 })),
1446 (Some(l), None) => Ok(Some(l.clone())),
1447 (None, Some(r)) => Ok(Some(r.clone())),
1448 (None, None) => Ok(None),
1449 }
1450 }
1451
1452 fn merge_float_override(
1453 key: &str,
1454 left: Option<&FloatValueType>,
1455 right: Option<&FloatValueType>,
1456 ) -> Result<Option<FloatValueType>, SchemaError> {
1457 match (left, right) {
1458 (Some(l), Some(r)) => Ok(Some(FloatValueType {
1459 float_inverted_index: Self::merge_index_or_error(
1460 l.float_inverted_index.as_ref(),
1461 r.float_inverted_index.as_ref(),
1462 &format!("key '{key}' float.float_inverted_index"),
1463 )?,
1464 })),
1465 (Some(l), None) => Ok(Some(l.clone())),
1466 (None, Some(r)) => Ok(Some(r.clone())),
1467 (None, None) => Ok(None),
1468 }
1469 }
1470
1471 fn merge_int_override(
1472 key: &str,
1473 left: Option<&IntValueType>,
1474 right: Option<&IntValueType>,
1475 ) -> Result<Option<IntValueType>, SchemaError> {
1476 match (left, right) {
1477 (Some(l), Some(r)) => Ok(Some(IntValueType {
1478 int_inverted_index: Self::merge_index_or_error(
1479 l.int_inverted_index.as_ref(),
1480 r.int_inverted_index.as_ref(),
1481 &format!("key '{key}' int.int_inverted_index"),
1482 )?,
1483 })),
1484 (Some(l), None) => Ok(Some(l.clone())),
1485 (None, Some(r)) => Ok(Some(r.clone())),
1486 (None, None) => Ok(None),
1487 }
1488 }
1489
1490 fn merge_bool_override(
1491 key: &str,
1492 left: Option<&BoolValueType>,
1493 right: Option<&BoolValueType>,
1494 ) -> Result<Option<BoolValueType>, SchemaError> {
1495 match (left, right) {
1496 (Some(l), Some(r)) => Ok(Some(BoolValueType {
1497 bool_inverted_index: Self::merge_index_or_error(
1498 l.bool_inverted_index.as_ref(),
1499 r.bool_inverted_index.as_ref(),
1500 &format!("key '{key}' bool.bool_inverted_index"),
1501 )?,
1502 })),
1503 (Some(l), None) => Ok(Some(l.clone())),
1504 (None, Some(r)) => Ok(Some(r.clone())),
1505 (None, None) => Ok(None),
1506 }
1507 }
1508
1509 fn merge_float_list_override(
1510 key: &str,
1511 left: Option<&FloatListValueType>,
1512 right: Option<&FloatListValueType>,
1513 ) -> Result<Option<FloatListValueType>, SchemaError> {
1514 match (left, right) {
1515 (Some(l), Some(r)) => Ok(Some(FloatListValueType {
1516 vector_index: Self::merge_index_or_error(
1517 l.vector_index.as_ref(),
1518 r.vector_index.as_ref(),
1519 &format!("key '{key}' float_list.vector_index"),
1520 )?,
1521 })),
1522 (Some(l), None) => Ok(Some(l.clone())),
1523 (None, Some(r)) => Ok(Some(r.clone())),
1524 (None, None) => Ok(None),
1525 }
1526 }
1527
1528 fn merge_sparse_vector_override(
1529 key: &str,
1530 left: Option<&SparseVectorValueType>,
1531 right: Option<&SparseVectorValueType>,
1532 ) -> Result<Option<SparseVectorValueType>, SchemaError> {
1533 match (left, right) {
1534 (Some(l), Some(r)) => Ok(Some(SparseVectorValueType {
1535 sparse_vector_index: Self::merge_index_or_error(
1536 l.sparse_vector_index.as_ref(),
1537 r.sparse_vector_index.as_ref(),
1538 &format!("key '{key}' sparse_vector.sparse_vector_index"),
1539 )?,
1540 })),
1541 (Some(l), None) => Ok(Some(l.clone())),
1542 (None, Some(r)) => Ok(Some(r.clone())),
1543 (None, None) => Ok(None),
1544 }
1545 }
1546
1547 fn merge_index_or_error<T: Clone + PartialEq>(
1548 left: Option<&T>,
1549 right: Option<&T>,
1550 context: &str,
1551 ) -> Result<Option<T>, SchemaError> {
1552 match (left, right) {
1553 (Some(l), Some(r)) => {
1554 if l == r {
1555 Ok(Some(l.clone()))
1556 } else {
1557 Err(SchemaError::ConfigurationConflict {
1558 context: context.to_string(),
1559 })
1560 }
1561 }
1562 (Some(l), None) => Ok(Some(l.clone())),
1563 (None, Some(r)) => Ok(Some(r.clone())),
1564 (None, None) => Ok(None),
1565 }
1566 }
1567
1568 fn merge_value_types(
1571 default: &ValueTypes,
1572 user: &ValueTypes,
1573 knn_index: KnnIndex,
1574 ) -> Result<ValueTypes, SchemaError> {
1575 let float_list = Self::merge_float_list_type(
1577 default.float_list.as_ref(),
1578 user.float_list.as_ref(),
1579 knn_index,
1580 )?;
1581
1582 if let Some(ref fl) = float_list {
1584 Self::validate_float_list_value_type(fl)?;
1585 }
1586
1587 Ok(ValueTypes {
1588 string: Self::merge_string_type(default.string.as_ref(), user.string.as_ref())?,
1589 float: Self::merge_float_type(default.float.as_ref(), user.float.as_ref())?,
1590 int: Self::merge_int_type(default.int.as_ref(), user.int.as_ref())?,
1591 boolean: Self::merge_bool_type(default.boolean.as_ref(), user.boolean.as_ref())?,
1592 float_list,
1593 sparse_vector: Self::merge_sparse_vector_type(
1594 default.sparse_vector.as_ref(),
1595 user.sparse_vector.as_ref(),
1596 )?,
1597 })
1598 }
1599
1600 fn merge_string_type(
1602 default: Option<&StringValueType>,
1603 user: Option<&StringValueType>,
1604 ) -> Result<Option<StringValueType>, SchemaError> {
1605 match (default, user) {
1606 (Some(default), Some(user)) => Ok(Some(StringValueType {
1607 string_inverted_index: Self::merge_string_inverted_index_type(
1608 default.string_inverted_index.as_ref(),
1609 user.string_inverted_index.as_ref(),
1610 )?,
1611 fts_index: Self::merge_fts_index_type(
1612 default.fts_index.as_ref(),
1613 user.fts_index.as_ref(),
1614 )?,
1615 })),
1616 (Some(default), None) => Ok(Some(default.clone())),
1617 (None, Some(user)) => Ok(Some(user.clone())),
1618 (None, None) => Ok(None),
1619 }
1620 }
1621
1622 fn merge_float_type(
1624 default: Option<&FloatValueType>,
1625 user: Option<&FloatValueType>,
1626 ) -> Result<Option<FloatValueType>, SchemaError> {
1627 match (default, user) {
1628 (Some(default), Some(user)) => Ok(Some(FloatValueType {
1629 float_inverted_index: Self::merge_float_inverted_index_type(
1630 default.float_inverted_index.as_ref(),
1631 user.float_inverted_index.as_ref(),
1632 )?,
1633 })),
1634 (Some(default), None) => Ok(Some(default.clone())),
1635 (None, Some(user)) => Ok(Some(user.clone())),
1636 (None, None) => Ok(None),
1637 }
1638 }
1639
1640 fn merge_int_type(
1642 default: Option<&IntValueType>,
1643 user: Option<&IntValueType>,
1644 ) -> Result<Option<IntValueType>, SchemaError> {
1645 match (default, user) {
1646 (Some(default), Some(user)) => Ok(Some(IntValueType {
1647 int_inverted_index: Self::merge_int_inverted_index_type(
1648 default.int_inverted_index.as_ref(),
1649 user.int_inverted_index.as_ref(),
1650 )?,
1651 })),
1652 (Some(default), None) => Ok(Some(default.clone())),
1653 (None, Some(user)) => Ok(Some(user.clone())),
1654 (None, None) => Ok(None),
1655 }
1656 }
1657
1658 fn merge_bool_type(
1660 default: Option<&BoolValueType>,
1661 user: Option<&BoolValueType>,
1662 ) -> Result<Option<BoolValueType>, SchemaError> {
1663 match (default, user) {
1664 (Some(default), Some(user)) => Ok(Some(BoolValueType {
1665 bool_inverted_index: Self::merge_bool_inverted_index_type(
1666 default.bool_inverted_index.as_ref(),
1667 user.bool_inverted_index.as_ref(),
1668 )?,
1669 })),
1670 (Some(default), None) => Ok(Some(default.clone())),
1671 (None, Some(user)) => Ok(Some(user.clone())),
1672 (None, None) => Ok(None),
1673 }
1674 }
1675
1676 fn merge_float_list_type(
1678 default: Option<&FloatListValueType>,
1679 user: Option<&FloatListValueType>,
1680 knn_index: KnnIndex,
1681 ) -> Result<Option<FloatListValueType>, SchemaError> {
1682 match (default, user) {
1683 (Some(default), Some(user)) => Ok(Some(FloatListValueType {
1684 vector_index: Self::merge_vector_index_type(
1685 default.vector_index.as_ref(),
1686 user.vector_index.as_ref(),
1687 knn_index,
1688 )?,
1689 })),
1690 (Some(default), None) => Ok(Some(default.clone())),
1691 (None, Some(user)) => Ok(Some(user.clone())),
1692 (None, None) => Ok(None),
1693 }
1694 }
1695
1696 fn merge_sparse_vector_type(
1698 default: Option<&SparseVectorValueType>,
1699 user: Option<&SparseVectorValueType>,
1700 ) -> Result<Option<SparseVectorValueType>, SchemaError> {
1701 match (default, user) {
1702 (Some(default), Some(user)) => Ok(Some(SparseVectorValueType {
1703 sparse_vector_index: Self::merge_sparse_vector_index_type(
1704 default.sparse_vector_index.as_ref(),
1705 user.sparse_vector_index.as_ref(),
1706 )?,
1707 })),
1708 (Some(default), None) => Ok(Some(default.clone())),
1709 (None, Some(user)) => Ok(Some(user.clone())),
1710 (None, None) => Ok(None),
1711 }
1712 }
1713
1714 fn merge_string_inverted_index_type(
1716 default: Option<&StringInvertedIndexType>,
1717 user: Option<&StringInvertedIndexType>,
1718 ) -> Result<Option<StringInvertedIndexType>, SchemaError> {
1719 match (default, user) {
1720 (Some(_default), Some(user)) => {
1721 Ok(Some(StringInvertedIndexType {
1722 enabled: user.enabled, config: user.config.clone(), }))
1725 }
1726 (Some(default), None) => Ok(Some(default.clone())),
1727 (None, Some(user)) => Ok(Some(user.clone())),
1728 (None, None) => Ok(None),
1729 }
1730 }
1731
1732 fn merge_fts_index_type(
1733 default: Option<&FtsIndexType>,
1734 user: Option<&FtsIndexType>,
1735 ) -> Result<Option<FtsIndexType>, SchemaError> {
1736 match (default, user) {
1737 (Some(default), Some(user)) => Ok(Some(FtsIndexType {
1738 enabled: user.enabled,
1739 config: FtsIndexConfig {
1740 algorithm: if !is_default_fts_algorithm(&user.config.algorithm) {
1741 user.config.algorithm.clone()
1742 } else {
1743 default.config.algorithm.clone()
1744 },
1745 },
1746 })),
1747 (Some(default), None) => Ok(Some(default.clone())),
1748 (None, Some(user)) => Ok(Some(user.clone())),
1749 (None, None) => Ok(None),
1750 }
1751 }
1752
1753 fn merge_float_inverted_index_type(
1754 default: Option<&FloatInvertedIndexType>,
1755 user: Option<&FloatInvertedIndexType>,
1756 ) -> Result<Option<FloatInvertedIndexType>, SchemaError> {
1757 match (default, user) {
1758 (Some(_default), Some(user)) => Ok(Some(FloatInvertedIndexType {
1759 enabled: user.enabled,
1760 config: user.config.clone(),
1761 })),
1762 (Some(default), None) => Ok(Some(default.clone())),
1763 (None, Some(user)) => Ok(Some(user.clone())),
1764 (None, None) => Ok(None),
1765 }
1766 }
1767
1768 fn merge_int_inverted_index_type(
1769 default: Option<&IntInvertedIndexType>,
1770 user: Option<&IntInvertedIndexType>,
1771 ) -> Result<Option<IntInvertedIndexType>, SchemaError> {
1772 match (default, user) {
1773 (Some(_default), Some(user)) => Ok(Some(IntInvertedIndexType {
1774 enabled: user.enabled,
1775 config: user.config.clone(),
1776 })),
1777 (Some(default), None) => Ok(Some(default.clone())),
1778 (None, Some(user)) => Ok(Some(user.clone())),
1779 (None, None) => Ok(None),
1780 }
1781 }
1782
1783 fn merge_bool_inverted_index_type(
1784 default: Option<&BoolInvertedIndexType>,
1785 user: Option<&BoolInvertedIndexType>,
1786 ) -> Result<Option<BoolInvertedIndexType>, SchemaError> {
1787 match (default, user) {
1788 (Some(_default), Some(user)) => Ok(Some(BoolInvertedIndexType {
1789 enabled: user.enabled,
1790 config: user.config.clone(),
1791 })),
1792 (Some(default), None) => Ok(Some(default.clone())),
1793 (None, Some(user)) => Ok(Some(user.clone())),
1794 (None, None) => Ok(None),
1795 }
1796 }
1797
1798 fn merge_vector_index_type(
1799 default: Option<&VectorIndexType>,
1800 user: Option<&VectorIndexType>,
1801 knn_index: KnnIndex,
1802 ) -> Result<Option<VectorIndexType>, SchemaError> {
1803 match (default, user) {
1804 (Some(default), Some(user)) => Ok(Some(VectorIndexType {
1805 enabled: user.enabled,
1806 config: Self::merge_vector_index_config(&default.config, &user.config, knn_index)?,
1807 })),
1808 (Some(default), None) => Ok(Some(default.clone())),
1809 (None, Some(user)) => Ok(Some(user.clone())),
1810 (None, None) => Ok(None),
1811 }
1812 }
1813
1814 fn merge_sparse_vector_index_type(
1815 default: Option<&SparseVectorIndexType>,
1816 user: Option<&SparseVectorIndexType>,
1817 ) -> Result<Option<SparseVectorIndexType>, SchemaError> {
1818 match (default, user) {
1819 (Some(default), Some(user)) => Ok(Some(SparseVectorIndexType {
1820 enabled: user.enabled,
1821 config: Self::merge_sparse_vector_index_config(&default.config, &user.config),
1822 })),
1823 (Some(default), None) => Ok(Some(default.clone())),
1824 (None, Some(user)) => Ok(Some(user.clone())),
1825 (None, None) => Ok(None),
1826 }
1827 }
1828
1829 fn validate_float_list_value_type(float_list: &FloatListValueType) -> Result<(), SchemaError> {
1832 if let Some(vector_index) = &float_list.vector_index {
1833 if let Some(hnsw) = &vector_index.config.hnsw {
1834 hnsw.validate().map_err(SchemaError::InvalidHnswConfig)?;
1835 }
1836 if let Some(spann) = &vector_index.config.spann {
1837 spann.validate().map_err(SchemaError::InvalidSpannConfig)?;
1838 }
1839 }
1840 Ok(())
1841 }
1842
1843 fn merge_vector_index_config(
1845 default: &VectorIndexConfig,
1846 user: &VectorIndexConfig,
1847 knn_index: KnnIndex,
1848 ) -> Result<VectorIndexConfig, SchemaError> {
1849 match knn_index {
1850 KnnIndex::Hnsw => Ok(VectorIndexConfig {
1851 space: user.space.clone().or(default.space.clone()),
1852 embedding_function: user
1853 .embedding_function
1854 .clone()
1855 .or(default.embedding_function.clone()),
1856 source_key: user.source_key.clone().or(default.source_key.clone()),
1857 hnsw: Self::merge_hnsw_configs(default.hnsw.as_ref(), user.hnsw.as_ref()),
1858 spann: None,
1859 }),
1860 KnnIndex::Spann => Ok(VectorIndexConfig {
1861 space: user.space.clone().or(default.space.clone()),
1862 embedding_function: user
1863 .embedding_function
1864 .clone()
1865 .or(default.embedding_function.clone()),
1866 source_key: user.source_key.clone().or(default.source_key.clone()),
1867 hnsw: None,
1868 spann: Self::merge_spann_configs(default.spann.as_ref(), user.spann.as_ref())?,
1869 }),
1870 }
1871 }
1872
1873 fn merge_sparse_vector_index_config(
1875 default: &SparseVectorIndexConfig,
1876 user: &SparseVectorIndexConfig,
1877 ) -> SparseVectorIndexConfig {
1878 SparseVectorIndexConfig {
1879 embedding_function: user
1880 .embedding_function
1881 .clone()
1882 .or(default.embedding_function.clone()),
1883 source_key: user.source_key.clone().or(default.source_key.clone()),
1884 bm25: user.bm25.or(default.bm25),
1885 algorithm: if !is_default_sparse_algorithm(&user.algorithm) {
1886 user.algorithm.clone()
1887 } else {
1888 default.algorithm.clone()
1889 },
1890 }
1891 }
1892
1893 fn merge_hnsw_configs(
1895 default_hnsw: Option<&HnswIndexConfig>,
1896 user_hnsw: Option<&HnswIndexConfig>,
1897 ) -> Option<HnswIndexConfig> {
1898 match (default_hnsw, user_hnsw) {
1899 (Some(default), Some(user)) => Some(HnswIndexConfig {
1900 ef_construction: user.ef_construction.or(default.ef_construction),
1901 max_neighbors: user.max_neighbors.or(default.max_neighbors),
1902 ef_search: user.ef_search.or(default.ef_search),
1903 num_threads: user.num_threads.or(default.num_threads),
1904 batch_size: user.batch_size.or(default.batch_size),
1905 sync_threshold: user.sync_threshold.or(default.sync_threshold),
1906 resize_factor: user.resize_factor.or(default.resize_factor),
1907 }),
1908 (Some(default), None) => Some(default.clone()),
1909 (None, Some(user)) => Some(user.clone()),
1910 (None, None) => None,
1911 }
1912 }
1913
1914 fn merge_spann_configs(
1916 default_spann: Option<&SpannIndexConfig>,
1917 user_spann: Option<&SpannIndexConfig>,
1918 ) -> Result<Option<SpannIndexConfig>, SchemaError> {
1919 match (default_spann, user_spann) {
1920 (Some(default), Some(user)) => {
1921 if !matches!(user.quantize, Quantization::None)
1923 || !matches!(default.quantize, Quantization::None)
1924 {
1925 return Err(SchemaError::InvalidUserInput {
1926 reason: "quantize field cannot be set in user schema. Quantization can only be enabled via frontend configuration.".to_string(),
1927 });
1928 }
1929 Ok(Some(SpannIndexConfig {
1930 search_nprobe: user.search_nprobe.or(default.search_nprobe),
1931 search_rng_factor: user.search_rng_factor.or(default.search_rng_factor),
1932 search_rng_epsilon: user.search_rng_epsilon.or(default.search_rng_epsilon),
1933 nreplica_count: user.nreplica_count.or(default.nreplica_count),
1934 write_rng_factor: user.write_rng_factor.or(default.write_rng_factor),
1935 write_rng_epsilon: user.write_rng_epsilon.or(default.write_rng_epsilon),
1936 split_threshold: user.split_threshold.or(default.split_threshold),
1937 num_samples_kmeans: user.num_samples_kmeans.or(default.num_samples_kmeans),
1938 initial_lambda: user.initial_lambda.or(default.initial_lambda),
1939 reassign_neighbor_count: user
1940 .reassign_neighbor_count
1941 .or(default.reassign_neighbor_count),
1942 merge_threshold: user.merge_threshold.or(default.merge_threshold),
1943 num_centers_to_merge_to: user
1944 .num_centers_to_merge_to
1945 .or(default.num_centers_to_merge_to),
1946 write_nprobe: user.write_nprobe.or(default.write_nprobe),
1947 ef_construction: user.ef_construction.or(default.ef_construction),
1948 ef_search: user.ef_search.or(default.ef_search),
1949 max_neighbors: user.max_neighbors.or(default.max_neighbors),
1950 center_drift_threshold: user
1951 .center_drift_threshold
1952 .or(default.center_drift_threshold),
1953 quantize: Quantization::None, }))
1955 }
1956 (Some(default), None) => {
1957 if !matches!(default.quantize, Quantization::None) {
1959 return Err(SchemaError::InvalidUserInput {
1960 reason: "quantize field cannot be set in default schema. Quantization can only be enabled via frontend configuration.".to_string(),
1961 });
1962 }
1963 Ok(Some(default.clone()))
1964 }
1965 (None, Some(user)) => {
1966 if !matches!(user.quantize, Quantization::None) {
1968 return Err(SchemaError::InvalidUserInput {
1969 reason: "quantize field cannot be set in user schema. Quantization can only be enabled via frontend configuration.".to_string(),
1970 });
1971 }
1972 Ok(Some(user.clone()))
1973 }
1974 (None, None) => Ok(None),
1975 }
1976 }
1977
1978 pub fn reconcile_with_collection_config(
1986 schema: &Schema,
1987 collection_config: &InternalCollectionConfiguration,
1988 default_knn_index: KnnIndex,
1989 ) -> Result<Schema, SchemaError> {
1990 if collection_config.is_default() {
1992 if schema.is_default() {
1993 let mut new_schema = Schema::new_default(default_knn_index);
1996
1997 if collection_config.embedding_function.is_some() {
1998 if let Some(float_list) = &mut new_schema.defaults.float_list {
1999 if let Some(vector_index) = &mut float_list.vector_index {
2000 vector_index.config.embedding_function =
2001 collection_config.embedding_function.clone();
2002 }
2003 }
2004 if let Some(embedding_types) = new_schema.keys.get_mut(EMBEDDING_KEY) {
2005 if let Some(float_list) = &mut embedding_types.float_list {
2006 if let Some(vector_index) = &mut float_list.vector_index {
2007 vector_index.config.embedding_function =
2008 collection_config.embedding_function.clone();
2009 }
2010 }
2011 }
2012 }
2013 return Ok(new_schema);
2014 } else {
2015 return Ok(schema.clone());
2017 }
2018 }
2019
2020 Self::try_from(collection_config)
2023 }
2024
2025 pub fn reconcile_schema_and_config(
2026 schema: Option<&Schema>,
2027 configuration: Option<&InternalCollectionConfiguration>,
2028 knn_index: KnnIndex,
2029 ) -> Result<Schema, SchemaError> {
2030 if let (Some(user_schema), Some(config)) = (schema, configuration) {
2032 if !user_schema.is_default() && !config.is_default() {
2033 return Err(SchemaError::ConfigAndSchemaConflict);
2034 }
2035 }
2036
2037 let reconciled_schema = Self::reconcile_with_defaults(schema, knn_index)?;
2038 if let Some(config) = configuration {
2039 Self::reconcile_with_collection_config(&reconciled_schema, config, knn_index)
2040 } else {
2041 Ok(reconciled_schema)
2042 }
2043 }
2044
2045 pub fn default_with_embedding_function(
2046 embedding_function: EmbeddingFunctionConfiguration,
2047 ) -> Schema {
2048 let mut schema = Schema::new_default(KnnIndex::Spann);
2049 if let Some(float_list) = &mut schema.defaults.float_list {
2050 if let Some(vector_index) = &mut float_list.vector_index {
2051 vector_index.config.embedding_function = Some(embedding_function.clone());
2052 }
2053 }
2054 if let Some(embedding_types) = schema.keys.get_mut(EMBEDDING_KEY) {
2055 if let Some(float_list) = &mut embedding_types.float_list {
2056 if let Some(vector_index) = &mut float_list.vector_index {
2057 vector_index.config.embedding_function = Some(embedding_function);
2058 }
2059 }
2060 }
2061 schema
2062 }
2063
2064 pub fn is_default(&self) -> bool {
2066 if !Self::is_value_types_default(&self.defaults) {
2068 return false;
2069 }
2070
2071 for key in self.keys.keys() {
2072 if key != EMBEDDING_KEY && key != DOCUMENT_KEY {
2073 return false;
2074 }
2075 }
2076
2077 if let Some(embedding_value) = self.keys.get(EMBEDDING_KEY) {
2079 if !Self::is_embedding_value_types_default(embedding_value) {
2080 return false;
2081 }
2082 }
2083
2084 if let Some(document_value) = self.keys.get(DOCUMENT_KEY) {
2086 if !Self::is_document_value_types_default(document_value) {
2087 return false;
2088 }
2089 }
2090
2091 if self.cmek.is_some() {
2093 return false;
2094 }
2095
2096 true
2097 }
2098
2099 fn is_value_types_default(value_types: &ValueTypes) -> bool {
2101 if let Some(string) = &value_types.string {
2103 if let Some(string_inverted) = &string.string_inverted_index {
2104 if !string_inverted.enabled {
2105 return false;
2106 }
2107 }
2109 if let Some(fts) = &string.fts_index {
2110 if fts.enabled {
2111 return false;
2112 }
2113 }
2115 }
2116
2117 if let Some(float) = &value_types.float {
2119 if let Some(float_inverted) = &float.float_inverted_index {
2120 if !float_inverted.enabled {
2121 return false;
2122 }
2123 }
2125 }
2126
2127 if let Some(int) = &value_types.int {
2129 if let Some(int_inverted) = &int.int_inverted_index {
2130 if !int_inverted.enabled {
2131 return false;
2132 }
2133 }
2135 }
2136
2137 if let Some(boolean) = &value_types.boolean {
2139 if let Some(bool_inverted) = &boolean.bool_inverted_index {
2140 if !bool_inverted.enabled {
2141 return false;
2142 }
2143 }
2145 }
2146
2147 if let Some(float_list) = &value_types.float_list {
2149 if let Some(vector_index) = &float_list.vector_index {
2150 if vector_index.enabled {
2151 return false;
2152 }
2153 if !is_embedding_function_default(&vector_index.config.embedding_function) {
2154 return false;
2155 }
2156 if !is_space_default(&vector_index.config.space) {
2157 return false;
2158 }
2159 if vector_index.config.source_key.is_some() {
2161 return false;
2162 }
2163 match (&vector_index.config.hnsw, &vector_index.config.spann) {
2166 (Some(hnsw_config), None) => {
2167 if !hnsw_config.is_default() {
2168 return false;
2169 }
2170 }
2171 (None, Some(spann_config)) => {
2172 if !spann_config.is_default() {
2173 return false;
2174 }
2175 }
2176 (Some(_), Some(_)) => return false, (None, None) => {}
2178 }
2179 }
2180 }
2181
2182 if let Some(sparse_vector) = &value_types.sparse_vector {
2184 if let Some(sparse_index) = &sparse_vector.sparse_vector_index {
2185 if sparse_index.enabled {
2186 return false;
2187 }
2188 if !is_embedding_function_default(&sparse_index.config.embedding_function) {
2190 return false;
2191 }
2192 if sparse_index.config.source_key.is_some() {
2193 return false;
2194 }
2195 if let Some(bm25) = &sparse_index.config.bm25 {
2196 if bm25 != &false {
2197 return false;
2198 }
2199 }
2200 }
2201 }
2202
2203 true
2204 }
2205
2206 fn is_embedding_value_types_default(value_types: &ValueTypes) -> bool {
2208 if value_types.string.is_some()
2210 || value_types.float.is_some()
2211 || value_types.int.is_some()
2212 || value_types.boolean.is_some()
2213 || value_types.sparse_vector.is_some()
2214 {
2215 return false;
2216 }
2217
2218 if let Some(float_list) = &value_types.float_list {
2220 if let Some(vector_index) = &float_list.vector_index {
2221 if !vector_index.enabled {
2222 return false;
2223 }
2224 if !is_space_default(&vector_index.config.space) {
2225 return false;
2226 }
2227 if !is_embedding_function_default(&vector_index.config.embedding_function) {
2229 return false;
2230 }
2231 if vector_index.config.source_key.as_deref() != Some(DOCUMENT_KEY) {
2233 return false;
2234 }
2235 match (&vector_index.config.hnsw, &vector_index.config.spann) {
2238 (Some(hnsw_config), None) => {
2239 if !hnsw_config.is_default() {
2240 return false;
2241 }
2242 }
2243 (None, Some(spann_config)) => {
2244 if !spann_config.is_default() {
2245 return false;
2246 }
2247 }
2248 (Some(_), Some(_)) => return false, (None, None) => {}
2250 }
2251 }
2252 }
2253
2254 true
2255 }
2256
2257 fn is_document_value_types_default(value_types: &ValueTypes) -> bool {
2259 if value_types.float_list.is_some()
2261 || value_types.float.is_some()
2262 || value_types.int.is_some()
2263 || value_types.boolean.is_some()
2264 || value_types.sparse_vector.is_some()
2265 {
2266 return false;
2267 }
2268
2269 if let Some(string) = &value_types.string {
2271 if let Some(fts) = &string.fts_index {
2272 if !fts.enabled {
2273 return false;
2274 }
2275 }
2277 if let Some(string_inverted) = &string.string_inverted_index {
2278 if string_inverted.enabled {
2279 return false;
2280 }
2281 }
2283 }
2284
2285 true
2286 }
2287
2288 pub fn is_metadata_type_index_enabled(
2290 &self,
2291 key: &str,
2292 value_type: MetadataValueType,
2293 ) -> Result<bool, SchemaError> {
2294 let v_type = self.keys.get(key).unwrap_or(&self.defaults);
2295
2296 match value_type {
2297 MetadataValueType::Bool => match &v_type.boolean {
2298 Some(bool_type) => match &bool_type.bool_inverted_index {
2299 Some(bool_inverted_index) => Ok(bool_inverted_index.enabled),
2300 None => Err(SchemaError::MissingIndexConfiguration {
2301 key: key.to_string(),
2302 value_type: "bool".to_string(),
2303 }),
2304 },
2305 None => match &self.defaults.boolean {
2306 Some(bool_type) => match &bool_type.bool_inverted_index {
2307 Some(bool_inverted_index) => Ok(bool_inverted_index.enabled),
2308 None => Err(SchemaError::MissingIndexConfiguration {
2309 key: key.to_string(),
2310 value_type: "bool".to_string(),
2311 }),
2312 },
2313 None => Err(SchemaError::MissingIndexConfiguration {
2314 key: key.to_string(),
2315 value_type: "bool".to_string(),
2316 }),
2317 },
2318 },
2319 MetadataValueType::Int => match &v_type.int {
2320 Some(int_type) => match &int_type.int_inverted_index {
2321 Some(int_inverted_index) => Ok(int_inverted_index.enabled),
2322 None => Err(SchemaError::MissingIndexConfiguration {
2323 key: key.to_string(),
2324 value_type: "int".to_string(),
2325 }),
2326 },
2327 None => match &self.defaults.int {
2328 Some(int_type) => match &int_type.int_inverted_index {
2329 Some(int_inverted_index) => Ok(int_inverted_index.enabled),
2330 None => Err(SchemaError::MissingIndexConfiguration {
2331 key: key.to_string(),
2332 value_type: "int".to_string(),
2333 }),
2334 },
2335 None => Err(SchemaError::MissingIndexConfiguration {
2336 key: key.to_string(),
2337 value_type: "int".to_string(),
2338 }),
2339 },
2340 },
2341 MetadataValueType::Float => match &v_type.float {
2342 Some(float_type) => match &float_type.float_inverted_index {
2343 Some(float_inverted_index) => Ok(float_inverted_index.enabled),
2344 None => Err(SchemaError::MissingIndexConfiguration {
2345 key: key.to_string(),
2346 value_type: "float".to_string(),
2347 }),
2348 },
2349 None => match &self.defaults.float {
2350 Some(float_type) => match &float_type.float_inverted_index {
2351 Some(float_inverted_index) => Ok(float_inverted_index.enabled),
2352 None => Err(SchemaError::MissingIndexConfiguration {
2353 key: key.to_string(),
2354 value_type: "float".to_string(),
2355 }),
2356 },
2357 None => Err(SchemaError::MissingIndexConfiguration {
2358 key: key.to_string(),
2359 value_type: "float".to_string(),
2360 }),
2361 },
2362 },
2363 MetadataValueType::Str => match &v_type.string {
2364 Some(string_type) => match &string_type.string_inverted_index {
2365 Some(string_inverted_index) => Ok(string_inverted_index.enabled),
2366 None => Err(SchemaError::MissingIndexConfiguration {
2367 key: key.to_string(),
2368 value_type: "string".to_string(),
2369 }),
2370 },
2371 None => match &self.defaults.string {
2372 Some(string_type) => match &string_type.string_inverted_index {
2373 Some(string_inverted_index) => Ok(string_inverted_index.enabled),
2374 None => Err(SchemaError::MissingIndexConfiguration {
2375 key: key.to_string(),
2376 value_type: "string".to_string(),
2377 }),
2378 },
2379 None => Err(SchemaError::MissingIndexConfiguration {
2380 key: key.to_string(),
2381 value_type: "string".to_string(),
2382 }),
2383 },
2384 },
2385 MetadataValueType::SparseVector => match &v_type.sparse_vector {
2386 Some(sparse_vector_type) => match &sparse_vector_type.sparse_vector_index {
2387 Some(sparse_vector_index) => Ok(sparse_vector_index.enabled),
2388 None => Err(SchemaError::MissingIndexConfiguration {
2389 key: key.to_string(),
2390 value_type: "sparse_vector".to_string(),
2391 }),
2392 },
2393 None => match &self.defaults.sparse_vector {
2394 Some(sparse_vector_type) => match &sparse_vector_type.sparse_vector_index {
2395 Some(sparse_vector_index) => Ok(sparse_vector_index.enabled),
2396 None => Err(SchemaError::MissingIndexConfiguration {
2397 key: key.to_string(),
2398 value_type: "sparse_vector".to_string(),
2399 }),
2400 },
2401 None => Err(SchemaError::MissingIndexConfiguration {
2402 key: key.to_string(),
2403 value_type: "sparse_vector".to_string(),
2404 }),
2405 },
2406 },
2407 MetadataValueType::BoolArray => {
2409 self.is_metadata_type_index_enabled(key, MetadataValueType::Bool)
2410 }
2411 MetadataValueType::IntArray => {
2412 self.is_metadata_type_index_enabled(key, MetadataValueType::Int)
2413 }
2414 MetadataValueType::FloatArray => {
2415 self.is_metadata_type_index_enabled(key, MetadataValueType::Float)
2416 }
2417 MetadataValueType::StringArray => {
2418 self.is_metadata_type_index_enabled(key, MetadataValueType::Str)
2419 }
2420 }
2421 }
2422
2423 pub fn is_metadata_key_unindexed(&self, key: &str, value_type: MetadataValueType) -> bool {
2426 match self.is_metadata_type_index_enabled(key, value_type) {
2427 Ok(enabled) => !enabled,
2428 Err(_) => false,
2429 }
2430 }
2431
2432 pub fn is_metadata_where_indexing_enabled(
2433 &self,
2434 where_clause: &Where,
2435 ) -> Result<(), FilterValidationError> {
2436 match where_clause {
2437 Where::Composite(composite) => {
2438 for child in &composite.children {
2439 self.is_metadata_where_indexing_enabled(child)?;
2440 }
2441 Ok(())
2442 }
2443 Where::Document(_) => {
2444 if !self.is_fts_enabled() {
2445 return Err(FilterValidationError::FtsDisabled);
2446 }
2447 Ok(())
2448 }
2449 Where::Metadata(expression) => {
2450 let value_type = match &expression.comparison {
2451 MetadataComparison::Primitive(_, value) => value.value_type(),
2452 MetadataComparison::Set(_, set_value) => set_value.value_type(),
2453 MetadataComparison::ArrayContains(_, value) => value.value_type(),
2454 };
2455 let is_enabled = self
2456 .is_metadata_type_index_enabled(expression.key.as_str(), value_type)
2457 .map_err(FilterValidationError::Schema)?;
2458 if !is_enabled {
2459 return Err(FilterValidationError::IndexingDisabled {
2460 key: expression.key.clone(),
2461 value_type,
2462 });
2463 }
2464 Ok(())
2465 }
2466 }
2467 }
2468
2469 pub fn is_knn_key_indexing_enabled(
2470 &self,
2471 key: &str,
2472 query: &QueryVector,
2473 ) -> Result<(), FilterValidationError> {
2474 match query {
2475 QueryVector::Sparse(_) => {
2476 let is_enabled = self
2477 .is_metadata_type_index_enabled(key, MetadataValueType::SparseVector)
2478 .map_err(FilterValidationError::Schema)?;
2479 if !is_enabled {
2480 return Err(FilterValidationError::IndexingDisabled {
2481 key: key.to_string(),
2482 value_type: MetadataValueType::SparseVector,
2483 });
2484 }
2485 Ok(())
2486 }
2487 QueryVector::Dense(_) => {
2488 Ok(())
2491 }
2492 }
2493 }
2494
2495 pub fn ensure_key_from_metadata(&mut self, key: &str, value_type: MetadataValueType) -> bool {
2496 if key.starts_with(CHROMA_KEY) {
2497 return false;
2498 }
2499 let value_types = self.keys.entry(key.to_string()).or_default();
2500 match value_type {
2501 MetadataValueType::Bool => {
2502 if value_types.boolean.is_none() {
2503 value_types.boolean = self.defaults.boolean.clone();
2504 return true;
2505 }
2506 }
2507 MetadataValueType::Int => {
2508 if value_types.int.is_none() {
2509 value_types.int = self.defaults.int.clone();
2510 return true;
2511 }
2512 }
2513 MetadataValueType::Float => {
2514 if value_types.float.is_none() {
2515 value_types.float = self.defaults.float.clone();
2516 return true;
2517 }
2518 }
2519 MetadataValueType::Str => {
2520 if value_types.string.is_none() {
2521 value_types.string = self.defaults.string.clone();
2522 return true;
2523 }
2524 }
2525 MetadataValueType::SparseVector => {
2526 if value_types.sparse_vector.is_none() {
2527 value_types.sparse_vector = self.defaults.sparse_vector.clone();
2528 return true;
2529 }
2530 }
2531 MetadataValueType::BoolArray => {
2533 if value_types.boolean.is_none() {
2534 value_types.boolean = self.defaults.boolean.clone();
2535 return true;
2536 }
2537 }
2538 MetadataValueType::IntArray => {
2539 if value_types.int.is_none() {
2540 value_types.int = self.defaults.int.clone();
2541 return true;
2542 }
2543 }
2544 MetadataValueType::FloatArray => {
2545 if value_types.float.is_none() {
2546 value_types.float = self.defaults.float.clone();
2547 return true;
2548 }
2549 }
2550 MetadataValueType::StringArray => {
2551 if value_types.string.is_none() {
2552 value_types.string = self.defaults.string.clone();
2553 return true;
2554 }
2555 }
2556 }
2557 false
2558 }
2559
2560 pub fn create_index(
2600 mut self,
2601 key: Option<&str>,
2602 config: IndexConfig,
2603 ) -> Result<Self, SchemaBuilderError> {
2604 match &config {
2606 IndexConfig::Vector(cfg) => {
2607 if let Some(k) = key {
2609 return Err(SchemaBuilderError::VectorIndexMustBeGlobal { key: k.to_string() });
2610 }
2611 self._set_vector_index_config_builder(cfg.clone());
2612 return Ok(self);
2613 }
2614 IndexConfig::Fts(_) if key != Some(DOCUMENT_KEY) => {
2615 return Err(SchemaBuilderError::FtsIndexOnlyOnDocument);
2617 }
2618 IndexConfig::Fts(_) => {
2619 }
2621 IndexConfig::SparseVector(_) if key.is_none() => {
2622 return Err(SchemaBuilderError::SparseVectorRequiresKey);
2624 }
2625 IndexConfig::SparseVector(_) => {
2626 }
2628 _ => {}
2629 }
2630
2631 if let Some(k) = key {
2633 if k == EMBEDDING_KEY {
2634 return Err(SchemaBuilderError::SpecialKeyModificationNotAllowed {
2635 key: k.to_string(),
2636 });
2637 }
2638 if k == DOCUMENT_KEY && !matches!(config, IndexConfig::Fts(_)) {
2639 return Err(SchemaBuilderError::SpecialKeyModificationNotAllowed {
2640 key: k.to_string(),
2641 });
2642 }
2643 if k.starts_with('#') && k != DOCUMENT_KEY {
2644 return Err(SchemaBuilderError::ReservedKeyPrefix { key: k.to_string() });
2645 }
2646 }
2647
2648 match key {
2650 Some(k) => self._set_index_for_key_builder(k, config, true)?,
2651 None => self._set_index_in_defaults_builder(config, true)?,
2652 }
2653
2654 Ok(self)
2655 }
2656
2657 pub fn delete_index(
2685 mut self,
2686 key: Option<&str>,
2687 config: IndexConfig,
2688 ) -> Result<Self, SchemaBuilderError> {
2689 match &config {
2691 IndexConfig::Vector(_) => {
2692 return Err(SchemaBuilderError::VectorIndexDeletionNotSupported);
2694 }
2695 IndexConfig::Fts(_) if key != Some(DOCUMENT_KEY) => {
2696 return Err(SchemaBuilderError::FtsIndexDeletionOnlyOnDocument);
2698 }
2699 IndexConfig::Fts(_) => {
2700 }
2702 IndexConfig::SparseVector(_) => {
2703 return Err(SchemaBuilderError::SparseVectorIndexDeletionNotSupported);
2705 }
2706 _ => {}
2707 }
2708
2709 if let Some(k) = key {
2711 if k == EMBEDDING_KEY {
2712 return Err(SchemaBuilderError::SpecialKeyModificationNotAllowed {
2713 key: k.to_string(),
2714 });
2715 }
2716 if k == DOCUMENT_KEY && !matches!(config, IndexConfig::Fts(_)) {
2717 return Err(SchemaBuilderError::SpecialKeyModificationNotAllowed {
2718 key: k.to_string(),
2719 });
2720 }
2721 if k.starts_with('#') && k != DOCUMENT_KEY {
2722 return Err(SchemaBuilderError::ReservedKeyPrefix { key: k.to_string() });
2723 }
2724 }
2725
2726 match key {
2728 Some(k) => self._set_index_for_key_builder(k, config, false)?,
2729 None => self._set_index_in_defaults_builder(config, false)?,
2730 }
2731
2732 Ok(self)
2733 }
2734
2735 pub fn with_cmek(mut self, cmek: Cmek) -> Self {
2753 self.cmek = Some(cmek);
2754 self
2755 }
2756
2757 fn _set_vector_index_config_builder(&mut self, config: VectorIndexConfig) {
2759 if let Some(float_list) = &mut self.defaults.float_list {
2761 if let Some(vector_index) = &mut float_list.vector_index {
2762 vector_index.config = config.clone();
2763 }
2764 }
2765
2766 if let Some(embedding_types) = self.keys.get_mut(EMBEDDING_KEY) {
2768 if let Some(float_list) = &mut embedding_types.float_list {
2769 if let Some(vector_index) = &mut float_list.vector_index {
2770 let mut updated_config = config;
2771 updated_config.source_key = Some(DOCUMENT_KEY.to_string());
2773 vector_index.config = updated_config;
2774 }
2775 }
2776 }
2777 }
2778
2779 fn _set_fts_index_config_builder(&mut self, config: FtsIndexConfig) {
2781 if let Some(string) = &mut self.defaults.string {
2783 if let Some(fts_index) = &mut string.fts_index {
2784 fts_index.config = config.clone();
2785 }
2786 }
2787
2788 if let Some(document_types) = self.keys.get_mut(DOCUMENT_KEY) {
2790 if let Some(string) = &mut document_types.string {
2791 if let Some(fts_index) = &mut string.fts_index {
2792 fts_index.config = config;
2793 }
2794 }
2795 }
2796 }
2797
2798 fn _set_index_for_key_builder(
2800 &mut self,
2801 key: &str,
2802 config: IndexConfig,
2803 enabled: bool,
2804 ) -> Result<(), SchemaBuilderError> {
2805 let value_types = self.keys.entry(key.to_string()).or_default();
2807
2808 match config {
2810 IndexConfig::Vector(_) => {
2811 return Err(SchemaBuilderError::VectorIndexMustBeGlobal {
2812 key: key.to_string(),
2813 });
2814 }
2815 IndexConfig::Fts(cfg) => {
2816 if let Some(string) = value_types.string.as_mut() {
2818 if let Some(fts_index) = string.fts_index.as_mut() {
2819 fts_index.enabled = enabled;
2820 fts_index.config = cfg;
2821 }
2822 }
2823 }
2824 IndexConfig::SparseVector(cfg) => {
2825 value_types.sparse_vector = Some(SparseVectorValueType {
2826 sparse_vector_index: Some(SparseVectorIndexType {
2827 enabled,
2828 config: cfg,
2829 }),
2830 });
2831 }
2832 IndexConfig::StringInverted(cfg) => {
2833 if value_types.string.is_none() {
2834 value_types.string = Some(StringValueType {
2835 fts_index: None,
2836 string_inverted_index: None,
2837 });
2838 }
2839 if let Some(string) = &mut value_types.string {
2840 string.string_inverted_index = Some(StringInvertedIndexType {
2841 enabled,
2842 config: cfg,
2843 });
2844 }
2845 }
2846 IndexConfig::IntInverted(cfg) => {
2847 value_types.int = Some(IntValueType {
2848 int_inverted_index: Some(IntInvertedIndexType {
2849 enabled,
2850 config: cfg,
2851 }),
2852 });
2853 }
2854 IndexConfig::FloatInverted(cfg) => {
2855 value_types.float = Some(FloatValueType {
2856 float_inverted_index: Some(FloatInvertedIndexType {
2857 enabled,
2858 config: cfg,
2859 }),
2860 });
2861 }
2862 IndexConfig::BoolInverted(cfg) => {
2863 value_types.boolean = Some(BoolValueType {
2864 bool_inverted_index: Some(BoolInvertedIndexType {
2865 enabled,
2866 config: cfg,
2867 }),
2868 });
2869 }
2870 }
2871
2872 Ok(())
2873 }
2874
2875 fn _set_index_in_defaults_builder(
2877 &mut self,
2878 config: IndexConfig,
2879 enabled: bool,
2880 ) -> Result<(), SchemaBuilderError> {
2881 match config {
2882 IndexConfig::Vector(_) => {
2883 return Err(SchemaBuilderError::VectorIndexMustBeGlobal {
2884 key: "defaults".to_string(),
2885 });
2886 }
2887 IndexConfig::Fts(_) => {
2888 return Err(SchemaBuilderError::FtsIndexOnlyOnDocument);
2890 }
2891 IndexConfig::SparseVector(cfg) => {
2892 self.defaults.sparse_vector = Some(SparseVectorValueType {
2893 sparse_vector_index: Some(SparseVectorIndexType {
2894 enabled,
2895 config: cfg,
2896 }),
2897 });
2898 }
2899 IndexConfig::StringInverted(cfg) => {
2900 if self.defaults.string.is_none() {
2901 self.defaults.string = Some(StringValueType {
2902 fts_index: None,
2903 string_inverted_index: None,
2904 });
2905 }
2906 if let Some(string) = &mut self.defaults.string {
2907 string.string_inverted_index = Some(StringInvertedIndexType {
2908 enabled,
2909 config: cfg,
2910 });
2911 }
2912 }
2913 IndexConfig::IntInverted(cfg) => {
2914 self.defaults.int = Some(IntValueType {
2915 int_inverted_index: Some(IntInvertedIndexType {
2916 enabled,
2917 config: cfg,
2918 }),
2919 });
2920 }
2921 IndexConfig::FloatInverted(cfg) => {
2922 self.defaults.float = Some(FloatValueType {
2923 float_inverted_index: Some(FloatInvertedIndexType {
2924 enabled,
2925 config: cfg,
2926 }),
2927 });
2928 }
2929 IndexConfig::BoolInverted(cfg) => {
2930 self.defaults.boolean = Some(BoolValueType {
2931 bool_inverted_index: Some(BoolInvertedIndexType {
2932 enabled,
2933 config: cfg,
2934 }),
2935 });
2936 }
2937 }
2938
2939 Ok(())
2940 }
2941}
2942
2943#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
2948#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
2949#[serde(deny_unknown_fields)]
2950pub struct VectorIndexConfig {
2951 #[serde(skip_serializing_if = "Option::is_none")]
2953 pub space: Option<Space>,
2954 #[serde(skip_serializing_if = "Option::is_none")]
2956 pub embedding_function: Option<EmbeddingFunctionConfiguration>,
2957 #[serde(skip_serializing_if = "Option::is_none")]
2959 pub source_key: Option<String>,
2960 #[serde(skip_serializing_if = "Option::is_none")]
2962 pub hnsw: Option<HnswIndexConfig>,
2963 #[serde(skip_serializing_if = "Option::is_none")]
2965 pub spann: Option<SpannIndexConfig>,
2966}
2967
2968#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Validate, Default)]
2970#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
2971#[serde(deny_unknown_fields)]
2972pub struct HnswIndexConfig {
2973 #[serde(skip_serializing_if = "Option::is_none")]
2974 pub ef_construction: Option<usize>,
2975 #[serde(skip_serializing_if = "Option::is_none")]
2976 pub max_neighbors: Option<usize>,
2977 #[serde(skip_serializing_if = "Option::is_none")]
2978 pub ef_search: Option<usize>,
2979 #[serde(skip_serializing_if = "Option::is_none")]
2980 pub num_threads: Option<usize>,
2981 #[serde(skip_serializing_if = "Option::is_none")]
2982 #[validate(range(min = 2))]
2983 pub batch_size: Option<usize>,
2984 #[serde(skip_serializing_if = "Option::is_none")]
2985 #[validate(range(min = 2))]
2986 pub sync_threshold: Option<usize>,
2987 #[serde(skip_serializing_if = "Option::is_none")]
2988 pub resize_factor: Option<f64>,
2989}
2990
2991impl HnswIndexConfig {
2992 pub fn is_default(&self) -> bool {
2996 if let Some(ef_construction) = self.ef_construction {
2997 if ef_construction != default_construction_ef() {
2998 return false;
2999 }
3000 }
3001 if let Some(max_neighbors) = self.max_neighbors {
3002 if max_neighbors != default_m() {
3003 return false;
3004 }
3005 }
3006 if let Some(ef_search) = self.ef_search {
3007 if ef_search != default_search_ef() {
3008 return false;
3009 }
3010 }
3011 if let Some(batch_size) = self.batch_size {
3012 if batch_size != default_batch_size() {
3013 return false;
3014 }
3015 }
3016 if let Some(sync_threshold) = self.sync_threshold {
3017 if sync_threshold != default_sync_threshold() {
3018 return false;
3019 }
3020 }
3021 if let Some(resize_factor) = self.resize_factor {
3022 if resize_factor != default_resize_factor() {
3023 return false;
3024 }
3025 }
3026 true
3028 }
3029}
3030
3031#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
3033#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
3034#[serde(rename_all = "snake_case")]
3035pub enum Quantization {
3036 #[default]
3037 None,
3038 FourBitRabitQWithUSearch,
3039}
3040
3041fn is_default_quantization(v: &Quantization) -> bool {
3042 matches!(v, Quantization::None)
3043}
3044
3045#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Validate, Default)]
3047#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
3048#[serde(deny_unknown_fields)]
3049pub struct SpannIndexConfig {
3050 #[serde(skip_serializing_if = "Option::is_none")]
3051 #[validate(range(max = 128))]
3052 pub search_nprobe: Option<u32>,
3053 #[serde(skip_serializing_if = "Option::is_none")]
3054 #[validate(range(min = 1.0, max = 1.0))]
3055 pub search_rng_factor: Option<f32>,
3056 #[serde(skip_serializing_if = "Option::is_none")]
3057 #[validate(range(min = 5.0, max = 10.0))]
3058 pub search_rng_epsilon: Option<f32>,
3059 #[serde(skip_serializing_if = "Option::is_none")]
3060 #[validate(range(max = 8))]
3061 pub nreplica_count: Option<u32>,
3062 #[serde(skip_serializing_if = "Option::is_none")]
3063 #[validate(range(min = 1.0, max = 1.0))]
3064 pub write_rng_factor: Option<f32>,
3065 #[serde(skip_serializing_if = "Option::is_none")]
3066 #[validate(range(min = 5.0, max = 10.0))]
3067 pub write_rng_epsilon: Option<f32>,
3068 #[serde(skip_serializing_if = "Option::is_none")]
3069 #[validate(range(min = 50, max = 200))]
3070 pub split_threshold: Option<u32>,
3071 #[serde(skip_serializing_if = "Option::is_none")]
3072 #[validate(range(max = 1000))]
3073 pub num_samples_kmeans: Option<usize>,
3074 #[serde(skip_serializing_if = "Option::is_none")]
3075 #[validate(range(min = 100.0, max = 100.0))]
3076 pub initial_lambda: Option<f32>,
3077 #[serde(skip_serializing_if = "Option::is_none")]
3078 #[validate(range(max = 64))]
3079 pub reassign_neighbor_count: Option<u32>,
3080 #[serde(skip_serializing_if = "Option::is_none")]
3081 #[validate(range(min = 25, max = 100))]
3082 pub merge_threshold: Option<u32>,
3083 #[serde(skip_serializing_if = "Option::is_none")]
3084 #[validate(range(max = 8))]
3085 pub num_centers_to_merge_to: Option<u32>,
3086 #[serde(skip_serializing_if = "Option::is_none")]
3087 #[validate(range(max = 64))]
3088 pub write_nprobe: Option<u32>,
3089 #[serde(skip_serializing_if = "Option::is_none")]
3090 #[validate(range(max = 200))]
3091 pub ef_construction: Option<usize>,
3092 #[serde(skip_serializing_if = "Option::is_none")]
3093 #[validate(range(max = 200))]
3094 pub ef_search: Option<usize>,
3095 #[serde(skip_serializing_if = "Option::is_none")]
3096 #[validate(range(max = 64))]
3097 pub max_neighbors: Option<usize>,
3098 #[serde(skip_serializing_if = "Option::is_none")]
3099 #[validate(range(min = 0.1, max = 1.0))]
3100 pub center_drift_threshold: Option<f32>,
3101 #[serde(default, skip_serializing_if = "is_default_quantization")]
3103 pub quantize: Quantization,
3104}
3105
3106impl SpannIndexConfig {
3107 pub fn is_default(&self) -> bool {
3110 if let Some(search_nprobe) = self.search_nprobe {
3111 if search_nprobe != default_search_nprobe() {
3112 return false;
3113 }
3114 }
3115 if let Some(search_rng_factor) = self.search_rng_factor {
3116 if search_rng_factor != default_search_rng_factor() {
3117 return false;
3118 }
3119 }
3120 if let Some(search_rng_epsilon) = self.search_rng_epsilon {
3121 if search_rng_epsilon != default_search_rng_epsilon() {
3122 return false;
3123 }
3124 }
3125 if let Some(nreplica_count) = self.nreplica_count {
3126 if nreplica_count != default_nreplica_count() {
3127 return false;
3128 }
3129 }
3130 if let Some(write_rng_factor) = self.write_rng_factor {
3131 if write_rng_factor != default_write_rng_factor() {
3132 return false;
3133 }
3134 }
3135 if let Some(write_rng_epsilon) = self.write_rng_epsilon {
3136 if write_rng_epsilon != default_write_rng_epsilon() {
3137 return false;
3138 }
3139 }
3140 if let Some(split_threshold) = self.split_threshold {
3141 if split_threshold != default_split_threshold() {
3142 return false;
3143 }
3144 }
3145 if let Some(num_samples_kmeans) = self.num_samples_kmeans {
3146 if num_samples_kmeans != default_num_samples_kmeans() {
3147 return false;
3148 }
3149 }
3150 if let Some(initial_lambda) = self.initial_lambda {
3151 if initial_lambda != default_initial_lambda() {
3152 return false;
3153 }
3154 }
3155 if let Some(reassign_neighbor_count) = self.reassign_neighbor_count {
3156 if reassign_neighbor_count != default_reassign_neighbor_count() {
3157 return false;
3158 }
3159 }
3160 if let Some(merge_threshold) = self.merge_threshold {
3161 if merge_threshold != default_merge_threshold() {
3162 return false;
3163 }
3164 }
3165 if let Some(num_centers_to_merge_to) = self.num_centers_to_merge_to {
3166 if num_centers_to_merge_to != default_num_centers_to_merge_to() {
3167 return false;
3168 }
3169 }
3170 if let Some(write_nprobe) = self.write_nprobe {
3171 if write_nprobe != default_write_nprobe() {
3172 return false;
3173 }
3174 }
3175 if let Some(ef_construction) = self.ef_construction {
3176 if ef_construction != default_construction_ef_spann() {
3177 return false;
3178 }
3179 }
3180 if let Some(ef_search) = self.ef_search {
3181 if ef_search != default_search_ef_spann() {
3182 return false;
3183 }
3184 }
3185 if let Some(max_neighbors) = self.max_neighbors {
3186 if max_neighbors != default_m_spann() {
3187 return false;
3188 }
3189 }
3190 if let Some(center_drift_threshold) = self.center_drift_threshold {
3191 if center_drift_threshold != default_center_drift_threshold() {
3192 return false;
3193 }
3194 }
3195 if !matches!(self.quantize, Quantization::None) {
3196 return false;
3197 }
3198 true
3199 }
3200}
3201
3202#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
3207#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
3208#[serde(rename_all = "snake_case")]
3209pub enum SparseIndexAlgorithm {
3210 #[default]
3211 Wand,
3212 MaxScore,
3213}
3214
3215fn is_default_sparse_algorithm(v: &SparseIndexAlgorithm) -> bool {
3216 v == &SparseIndexAlgorithm::default()
3217}
3218
3219#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
3220#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
3221#[serde(deny_unknown_fields)]
3222pub struct SparseVectorIndexConfig {
3223 #[serde(skip_serializing_if = "Option::is_none")]
3225 pub embedding_function: Option<EmbeddingFunctionConfiguration>,
3226 #[serde(skip_serializing_if = "Option::is_none")]
3228 pub source_key: Option<String>,
3229 #[serde(skip_serializing_if = "Option::is_none")]
3231 pub bm25: Option<bool>,
3232 #[serde(default, skip_serializing_if = "is_default_sparse_algorithm")]
3237 pub algorithm: SparseIndexAlgorithm,
3238}
3239
3240#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
3245#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
3246#[serde(rename_all = "snake_case")]
3247pub enum FtsAlgorithm {
3248 #[default]
3249 Trigram,
3250 TokenBitmap,
3251}
3252
3253fn is_default_fts_algorithm(v: &FtsAlgorithm) -> bool {
3254 v == &FtsAlgorithm::default()
3255}
3256
3257#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
3258#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
3259#[serde(deny_unknown_fields)]
3260pub struct FtsIndexConfig {
3261 #[serde(default, skip_serializing_if = "is_default_fts_algorithm")]
3266 pub algorithm: FtsAlgorithm,
3267}
3268
3269impl Default for FtsIndexConfig {
3270 fn default() -> Self {
3271 Self {
3272 algorithm: FtsAlgorithm::Trigram,
3273 }
3274 }
3275}
3276
3277#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
3278#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
3279#[serde(deny_unknown_fields)]
3280pub struct StringInvertedIndexConfig {
3281 }
3283
3284#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
3285#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
3286#[serde(deny_unknown_fields)]
3287pub struct IntInvertedIndexConfig {
3288 }
3290
3291#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
3292#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
3293#[serde(deny_unknown_fields)]
3294pub struct FloatInvertedIndexConfig {
3295 }
3297
3298#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
3299#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
3300#[serde(deny_unknown_fields)]
3301pub struct BoolInvertedIndexConfig {
3302 }
3304
3305#[derive(Clone, Debug)]
3311#[allow(clippy::large_enum_variant)]
3312pub enum IndexConfig {
3313 Vector(VectorIndexConfig),
3314 SparseVector(SparseVectorIndexConfig),
3315 Fts(FtsIndexConfig),
3316 StringInverted(StringInvertedIndexConfig),
3317 IntInverted(IntInvertedIndexConfig),
3318 FloatInverted(FloatInvertedIndexConfig),
3319 BoolInverted(BoolInvertedIndexConfig),
3320}
3321
3322impl From<VectorIndexConfig> for IndexConfig {
3324 fn from(config: VectorIndexConfig) -> Self {
3325 IndexConfig::Vector(config)
3326 }
3327}
3328
3329impl From<SparseVectorIndexConfig> for IndexConfig {
3330 fn from(config: SparseVectorIndexConfig) -> Self {
3331 IndexConfig::SparseVector(config)
3332 }
3333}
3334
3335impl From<FtsIndexConfig> for IndexConfig {
3336 fn from(config: FtsIndexConfig) -> Self {
3337 IndexConfig::Fts(config)
3338 }
3339}
3340
3341impl From<StringInvertedIndexConfig> for IndexConfig {
3342 fn from(config: StringInvertedIndexConfig) -> Self {
3343 IndexConfig::StringInverted(config)
3344 }
3345}
3346
3347impl From<IntInvertedIndexConfig> for IndexConfig {
3348 fn from(config: IntInvertedIndexConfig) -> Self {
3349 IndexConfig::IntInverted(config)
3350 }
3351}
3352
3353impl From<FloatInvertedIndexConfig> for IndexConfig {
3354 fn from(config: FloatInvertedIndexConfig) -> Self {
3355 IndexConfig::FloatInverted(config)
3356 }
3357}
3358
3359impl From<BoolInvertedIndexConfig> for IndexConfig {
3360 fn from(config: BoolInvertedIndexConfig) -> Self {
3361 IndexConfig::BoolInverted(config)
3362 }
3363}
3364
3365impl TryFrom<&InternalCollectionConfiguration> for Schema {
3366 type Error = SchemaError;
3367
3368 fn try_from(config: &InternalCollectionConfiguration) -> Result<Self, Self::Error> {
3369 let mut schema = match &config.vector_index {
3371 VectorIndexConfiguration::Hnsw(_) => Schema::new_default(KnnIndex::Hnsw),
3372 VectorIndexConfiguration::Spann(_) => Schema::new_default(KnnIndex::Spann),
3373 };
3374 let vector_config = match &config.vector_index {
3376 VectorIndexConfiguration::Hnsw(hnsw_config) => VectorIndexConfig {
3377 space: Some(hnsw_config.space.clone()),
3378 embedding_function: config.embedding_function.clone(),
3379 source_key: None,
3380 hnsw: Some(HnswIndexConfig {
3381 ef_construction: Some(hnsw_config.ef_construction),
3382 max_neighbors: Some(hnsw_config.max_neighbors),
3383 ef_search: Some(hnsw_config.ef_search),
3384 num_threads: Some(hnsw_config.num_threads),
3385 batch_size: Some(hnsw_config.batch_size),
3386 sync_threshold: Some(hnsw_config.sync_threshold),
3387 resize_factor: Some(hnsw_config.resize_factor),
3388 }),
3389 spann: None,
3390 },
3391 VectorIndexConfiguration::Spann(spann_config) => VectorIndexConfig {
3392 space: Some(spann_config.space.clone()),
3393 embedding_function: config.embedding_function.clone(),
3394 source_key: None,
3395 hnsw: None,
3396 spann: Some(SpannIndexConfig {
3397 search_nprobe: Some(spann_config.search_nprobe),
3398 search_rng_factor: Some(spann_config.search_rng_factor),
3399 search_rng_epsilon: Some(spann_config.search_rng_epsilon),
3400 nreplica_count: Some(spann_config.nreplica_count),
3401 write_rng_factor: Some(spann_config.write_rng_factor),
3402 write_rng_epsilon: Some(spann_config.write_rng_epsilon),
3403 split_threshold: Some(spann_config.split_threshold),
3404 num_samples_kmeans: Some(spann_config.num_samples_kmeans),
3405 initial_lambda: Some(spann_config.initial_lambda),
3406 reassign_neighbor_count: Some(spann_config.reassign_neighbor_count),
3407 merge_threshold: Some(spann_config.merge_threshold),
3408 num_centers_to_merge_to: Some(spann_config.num_centers_to_merge_to),
3409 write_nprobe: Some(spann_config.write_nprobe),
3410 ef_construction: Some(spann_config.ef_construction),
3411 ef_search: Some(spann_config.ef_search),
3412 max_neighbors: Some(spann_config.max_neighbors),
3413 center_drift_threshold: None,
3414 quantize: Quantization::None,
3415 }),
3416 },
3417 };
3418
3419 if let Some(float_list) = &mut schema.defaults.float_list {
3422 if let Some(vector_index) = &mut float_list.vector_index {
3423 vector_index.config = vector_config.clone();
3424 }
3425 }
3426
3427 if let Some(embedding_types) = schema.keys.get_mut(EMBEDDING_KEY) {
3431 if let Some(float_list) = &mut embedding_types.float_list {
3432 if let Some(vector_index) = &mut float_list.vector_index {
3433 let mut vector_config = vector_config;
3434 vector_config.source_key = Some(DOCUMENT_KEY.to_string());
3435 vector_index.config = vector_config;
3436 }
3437 }
3438 }
3439
3440 Ok(schema)
3441 }
3442}
3443
3444#[cfg(test)]
3445mod tests {
3446 use super::*;
3447 use crate::hnsw_configuration::Space;
3448 use crate::metadata::SparseVector;
3449 use crate::{
3450 EmbeddingFunctionNewConfiguration, InternalHnswConfiguration, InternalSpannConfiguration,
3451 };
3452 use serde_json::json;
3453
3454 #[test]
3455 fn test_record_only_schema_disables_all_indexes() {
3456 let schema = Schema::new_record_only();
3457
3458 let string = schema.defaults.string.as_ref().unwrap();
3460 assert!(!string.string_inverted_index.as_ref().unwrap().enabled);
3461 assert!(!string.fts_index.as_ref().unwrap().enabled);
3462
3463 let float = schema.defaults.float.as_ref().unwrap();
3464 assert!(!float.float_inverted_index.as_ref().unwrap().enabled);
3465
3466 let int = schema.defaults.int.as_ref().unwrap();
3467 assert!(!int.int_inverted_index.as_ref().unwrap().enabled);
3468
3469 let boolean = schema.defaults.boolean.as_ref().unwrap();
3470 assert!(!boolean.bool_inverted_index.as_ref().unwrap().enabled);
3471
3472 let float_list = schema.defaults.float_list.as_ref().unwrap();
3474 assert!(!float_list.vector_index.as_ref().unwrap().enabled);
3475
3476 let sparse = schema.defaults.sparse_vector.as_ref().unwrap();
3478 assert!(!sparse.sparse_vector_index.as_ref().unwrap().enabled);
3479
3480 let embedding = schema.keys.get(EMBEDDING_KEY).unwrap();
3482 let emb_vector = embedding.float_list.as_ref().unwrap();
3483 assert!(!emb_vector.vector_index.as_ref().unwrap().enabled);
3484
3485 let document = schema.keys.get(DOCUMENT_KEY).unwrap();
3487 let doc_string = document.string.as_ref().unwrap();
3488 assert!(!doc_string.fts_index.as_ref().unwrap().enabled);
3489 assert!(!doc_string.string_inverted_index.as_ref().unwrap().enabled);
3490 }
3491
3492 #[test]
3493 fn test_reconcile_with_defaults_none_user_schema() {
3494 let result = Schema::reconcile_with_defaults(None, KnnIndex::Spann).unwrap();
3496 let expected = Schema::new_default(KnnIndex::Spann);
3497 assert_eq!(result, expected);
3498 }
3499
3500 #[test]
3501 fn test_reconcile_with_defaults_empty_user_schema() {
3502 let user_schema = Schema {
3504 defaults: ValueTypes::default(),
3505 keys: HashMap::new(),
3506 cmek: None,
3507 };
3508
3509 let result = Schema::reconcile_with_defaults(Some(&user_schema), KnnIndex::Spann).unwrap();
3510 let expected = Schema::new_default(KnnIndex::Spann);
3511 assert_eq!(result, expected);
3512 }
3513
3514 #[test]
3515 fn test_reconcile_with_defaults_user_overrides_string_enabled() {
3516 let mut user_schema = Schema {
3518 defaults: ValueTypes::default(),
3519 keys: HashMap::new(),
3520 cmek: None,
3521 };
3522
3523 user_schema.defaults.string = Some(StringValueType {
3524 string_inverted_index: Some(StringInvertedIndexType {
3525 enabled: false, config: StringInvertedIndexConfig {},
3527 }),
3528 fts_index: None,
3529 });
3530
3531 let result = Schema::reconcile_with_defaults(Some(&user_schema), KnnIndex::Spann).unwrap();
3532
3533 assert!(
3535 !result
3536 .defaults
3537 .string
3538 .as_ref()
3539 .unwrap()
3540 .string_inverted_index
3541 .as_ref()
3542 .unwrap()
3543 .enabled
3544 );
3545 assert!(result.defaults.float.is_some());
3547 assert!(result.defaults.int.is_some());
3548 }
3549
3550 #[test]
3551 fn test_is_metadata_key_unindexed() {
3552 let mut schema = Schema::new_default(KnnIndex::Spann);
3554 schema.defaults.string = Some(StringValueType {
3555 string_inverted_index: Some(StringInvertedIndexType {
3556 enabled: false,
3557 config: StringInvertedIndexConfig {},
3558 }),
3559 fts_index: None,
3560 });
3561
3562 assert!(schema.is_metadata_key_unindexed("some_key", MetadataValueType::Str));
3564
3565 schema.keys.insert(
3567 "indexed_key".to_string(),
3568 ValueTypes {
3569 string: Some(StringValueType {
3570 string_inverted_index: Some(StringInvertedIndexType {
3571 enabled: true,
3572 config: StringInvertedIndexConfig {},
3573 }),
3574 fts_index: None,
3575 }),
3576 ..Default::default()
3577 },
3578 );
3579
3580 assert!(!schema.is_metadata_key_unindexed("indexed_key", MetadataValueType::Str));
3582
3583 schema.defaults.int = Some(IntValueType {
3585 int_inverted_index: Some(IntInvertedIndexType {
3586 enabled: false,
3587 config: IntInvertedIndexConfig {},
3588 }),
3589 });
3590 assert!(schema.is_metadata_key_unindexed("some_key", MetadataValueType::Int));
3591
3592 schema.defaults.int = Some(IntValueType {
3594 int_inverted_index: Some(IntInvertedIndexType {
3595 enabled: true,
3596 config: IntInvertedIndexConfig {},
3597 }),
3598 });
3599 assert!(!schema.is_metadata_key_unindexed("some_key", MetadataValueType::Int));
3600 }
3601
3602 #[test]
3603 fn test_reconcile_with_defaults_user_overrides_vector_config() {
3604 let mut user_schema = Schema {
3606 defaults: ValueTypes::default(),
3607 keys: HashMap::new(),
3608 cmek: None,
3609 };
3610
3611 user_schema.defaults.float_list = Some(FloatListValueType {
3612 vector_index: Some(VectorIndexType {
3613 enabled: true, config: VectorIndexConfig {
3615 space: Some(Space::L2), embedding_function: None, source_key: Some("custom_key".to_string()), hnsw: Some(HnswIndexConfig {
3619 ef_construction: Some(500), max_neighbors: None, ef_search: None, num_threads: None,
3623 batch_size: None,
3624 sync_threshold: None,
3625 resize_factor: None,
3626 }),
3627 spann: None,
3628 },
3629 }),
3630 });
3631
3632 let result = {
3634 let default_schema = Schema::new_default(KnnIndex::Hnsw);
3635 let merged_defaults = Schema::merge_value_types(
3636 &default_schema.defaults,
3637 &user_schema.defaults,
3638 KnnIndex::Hnsw,
3639 )
3640 .unwrap();
3641 let mut merged_keys = default_schema.keys.clone();
3642 for (key, user_value_types) in user_schema.keys {
3643 if let Some(default_value_types) = merged_keys.get(&key) {
3644 let merged_value_types = Schema::merge_value_types(
3645 default_value_types,
3646 &user_value_types,
3647 KnnIndex::Hnsw,
3648 )
3649 .unwrap();
3650 merged_keys.insert(key, merged_value_types);
3651 } else {
3652 merged_keys.insert(key, user_value_types);
3653 }
3654 }
3655 Schema {
3656 defaults: merged_defaults,
3657 keys: merged_keys,
3658 cmek: None,
3659 }
3660 };
3661
3662 let vector_config = &result
3663 .defaults
3664 .float_list
3665 .as_ref()
3666 .unwrap()
3667 .vector_index
3668 .as_ref()
3669 .unwrap()
3670 .config;
3671
3672 assert_eq!(vector_config.space, Some(Space::L2));
3674 assert_eq!(vector_config.source_key, Some("custom_key".to_string()));
3675 assert_eq!(
3676 vector_config.hnsw.as_ref().unwrap().ef_construction,
3677 Some(500)
3678 );
3679
3680 assert_eq!(vector_config.embedding_function, None);
3682 assert_eq!(
3684 vector_config.hnsw.as_ref().unwrap().max_neighbors,
3685 Some(default_m())
3686 );
3687 }
3688
3689 #[test]
3690 fn test_reconcile_with_defaults_keys() {
3691 let mut user_schema = Schema {
3693 defaults: ValueTypes::default(),
3694 keys: HashMap::new(),
3695 cmek: None,
3696 };
3697
3698 let custom_key_types = ValueTypes {
3700 string: Some(StringValueType {
3701 fts_index: Some(FtsIndexType {
3702 enabled: true,
3703 config: FtsIndexConfig::default(),
3704 }),
3705 string_inverted_index: Some(StringInvertedIndexType {
3706 enabled: false,
3707 config: StringInvertedIndexConfig {},
3708 }),
3709 }),
3710 ..Default::default()
3711 };
3712 user_schema
3713 .keys
3714 .insert("custom_key".to_string(), custom_key_types);
3715
3716 let result = Schema::reconcile_with_defaults(Some(&user_schema), KnnIndex::Spann).unwrap();
3717
3718 assert!(result.keys.contains_key(EMBEDDING_KEY));
3720 assert!(result.keys.contains_key(DOCUMENT_KEY));
3721
3722 assert!(result.keys.contains_key("custom_key"));
3724 let custom_override = result.keys.get("custom_key").unwrap();
3725 assert!(
3726 custom_override
3727 .string
3728 .as_ref()
3729 .unwrap()
3730 .fts_index
3731 .as_ref()
3732 .unwrap()
3733 .enabled
3734 );
3735 }
3736
3737 #[test]
3738 fn test_reconcile_with_defaults_override_existing_key() {
3739 let mut user_schema = Schema {
3741 defaults: ValueTypes::default(),
3742 keys: HashMap::new(),
3743 cmek: None,
3744 };
3745
3746 let embedding_override = ValueTypes {
3748 float_list: Some(FloatListValueType {
3749 vector_index: Some(VectorIndexType {
3750 enabled: false, config: VectorIndexConfig {
3752 space: Some(Space::Ip), embedding_function: Some(EmbeddingFunctionConfiguration::Legacy),
3754 source_key: Some("custom_embedding_key".to_string()),
3755 hnsw: None,
3756 spann: None,
3757 },
3758 }),
3759 }),
3760 ..Default::default()
3761 };
3762 user_schema
3763 .keys
3764 .insert(EMBEDDING_KEY.to_string(), embedding_override);
3765
3766 let result = Schema::reconcile_with_defaults(Some(&user_schema), KnnIndex::Spann).unwrap();
3767
3768 let embedding_config = result.keys.get(EMBEDDING_KEY).unwrap();
3769 let vector_config = &embedding_config
3770 .float_list
3771 .as_ref()
3772 .unwrap()
3773 .vector_index
3774 .as_ref()
3775 .unwrap();
3776
3777 assert!(!vector_config.enabled);
3779 assert_eq!(vector_config.config.space, Some(Space::Ip));
3780 assert_eq!(
3781 vector_config.config.source_key,
3782 Some("custom_embedding_key".to_string())
3783 );
3784 }
3785
3786 #[test]
3787 fn test_convert_schema_to_collection_config_hnsw_roundtrip() {
3788 let collection_config = InternalCollectionConfiguration {
3789 vector_index: VectorIndexConfiguration::Hnsw(InternalHnswConfiguration {
3790 space: Space::Cosine,
3791 ef_construction: 128,
3792 ef_search: 96,
3793 max_neighbors: 42,
3794 num_threads: 8,
3795 resize_factor: 1.5,
3796 sync_threshold: 2_000,
3797 batch_size: 256,
3798 }),
3799 embedding_function: Some(EmbeddingFunctionConfiguration::Known(
3800 EmbeddingFunctionNewConfiguration {
3801 name: "custom".to_string(),
3802 config: json!({"alpha": 1}),
3803 },
3804 )),
3805 };
3806
3807 let schema = Schema::try_from(&collection_config).unwrap();
3808 let reconstructed = InternalCollectionConfiguration::try_from(&schema).unwrap();
3809
3810 assert_eq!(reconstructed, collection_config);
3811 }
3812
3813 #[test]
3814 fn test_convert_schema_to_collection_config_spann_roundtrip() {
3815 let spann_config = InternalSpannConfiguration {
3816 space: Space::Cosine,
3817 search_nprobe: 11,
3818 search_rng_factor: 1.7,
3819 write_nprobe: 5,
3820 nreplica_count: 3,
3821 split_threshold: 150,
3822 merge_threshold: 80,
3823 ef_construction: 120,
3824 ef_search: 90,
3825 max_neighbors: 40,
3826 ..Default::default()
3827 };
3828
3829 let collection_config = InternalCollectionConfiguration {
3830 vector_index: VectorIndexConfiguration::Spann(spann_config.clone()),
3831 embedding_function: Some(EmbeddingFunctionConfiguration::Known(
3832 EmbeddingFunctionNewConfiguration {
3833 name: "custom".to_string(),
3834 config: json!({"beta": true}),
3835 },
3836 )),
3837 };
3838
3839 let schema = Schema::try_from(&collection_config).unwrap();
3840 let reconstructed = InternalCollectionConfiguration::try_from(&schema).unwrap();
3841
3842 assert_eq!(reconstructed, collection_config);
3843 }
3844
3845 #[test]
3846 fn test_convert_schema_to_collection_config_rejects_mixed_index() {
3847 let mut schema = Schema::new_default(KnnIndex::Hnsw);
3848 if let Some(embedding) = schema.keys.get_mut(EMBEDDING_KEY) {
3849 if let Some(float_list) = &mut embedding.float_list {
3850 if let Some(vector_index) = &mut float_list.vector_index {
3851 vector_index.config.spann = Some(SpannIndexConfig {
3852 search_nprobe: Some(1),
3853 search_rng_factor: Some(1.0),
3854 search_rng_epsilon: Some(0.1),
3855 nreplica_count: Some(1),
3856 write_rng_factor: Some(1.0),
3857 write_rng_epsilon: Some(0.1),
3858 split_threshold: Some(100),
3859 num_samples_kmeans: Some(10),
3860 initial_lambda: Some(0.5),
3861 reassign_neighbor_count: Some(10),
3862 merge_threshold: Some(50),
3863 num_centers_to_merge_to: Some(3),
3864 write_nprobe: Some(1),
3865 ef_construction: Some(50),
3866 ef_search: Some(40),
3867 max_neighbors: Some(20),
3868 center_drift_threshold: None,
3869 quantize: Quantization::None,
3870 });
3871 }
3872 }
3873 }
3874
3875 let result = InternalCollectionConfiguration::try_from(&schema);
3876 assert!(result.is_err());
3877 }
3878
3879 #[test]
3880 fn test_ensure_key_from_metadata_no_changes_for_existing_key() {
3881 let mut schema = Schema::new_default(KnnIndex::Hnsw);
3882 let before = schema.clone();
3883 let modified = schema.ensure_key_from_metadata(DOCUMENT_KEY, MetadataValueType::Str);
3884 assert!(!modified);
3885 assert_eq!(schema, before);
3886 }
3887
3888 #[test]
3889 fn test_ensure_key_from_metadata_populates_new_key_with_default_value_type() {
3890 let mut schema = Schema::new_default(KnnIndex::Hnsw);
3891 assert!(!schema.keys.contains_key("custom_field"));
3892
3893 let modified = schema.ensure_key_from_metadata("custom_field", MetadataValueType::Bool);
3894
3895 assert!(modified);
3896 let entry = schema
3897 .keys
3898 .get("custom_field")
3899 .expect("expected new key override to be inserted");
3900 assert_eq!(entry.boolean, schema.defaults.boolean);
3901 assert!(entry.string.is_none());
3902 assert!(entry.int.is_none());
3903 assert!(entry.float.is_none());
3904 assert!(entry.float_list.is_none());
3905 assert!(entry.sparse_vector.is_none());
3906 }
3907
3908 #[test]
3909 fn test_ensure_key_from_metadata_adds_missing_value_type_to_existing_key() {
3910 let mut schema = Schema::new_default(KnnIndex::Hnsw);
3911 let initial_len = schema.keys.len();
3912 schema.keys.insert(
3913 "custom_field".to_string(),
3914 ValueTypes {
3915 string: schema.defaults.string.clone(),
3916 ..Default::default()
3917 },
3918 );
3919
3920 let modified = schema.ensure_key_from_metadata("custom_field", MetadataValueType::Bool);
3921
3922 assert!(modified);
3923 assert_eq!(schema.keys.len(), initial_len + 1);
3924 let entry = schema
3925 .keys
3926 .get("custom_field")
3927 .expect("expected key override to exist after ensure call");
3928 assert!(entry.string.is_some());
3929 assert_eq!(entry.boolean, schema.defaults.boolean);
3930 }
3931
3932 #[test]
3933 fn test_is_knn_key_indexing_enabled_sparse_disabled_errors() {
3934 let schema = Schema::new_default(KnnIndex::Spann);
3935 let result = schema.is_knn_key_indexing_enabled(
3936 "custom_sparse",
3937 &QueryVector::Sparse(SparseVector::new(vec![0_u32], vec![1.0_f32]).unwrap()),
3938 );
3939
3940 let err = result.expect_err("expected indexing disabled error");
3941 match err {
3942 FilterValidationError::IndexingDisabled { key, value_type } => {
3943 assert_eq!(key, "custom_sparse");
3944 assert_eq!(value_type, crate::metadata::MetadataValueType::SparseVector);
3945 }
3946 other => panic!("unexpected error variant: {other:?}"),
3947 }
3948 }
3949
3950 #[test]
3951 fn test_is_knn_key_indexing_enabled_sparse_enabled_succeeds() {
3952 let mut schema = Schema::new_default(KnnIndex::Spann);
3953 schema.keys.insert(
3954 "sparse_enabled".to_string(),
3955 ValueTypes {
3956 sparse_vector: Some(SparseVectorValueType {
3957 sparse_vector_index: Some(SparseVectorIndexType {
3958 enabled: true,
3959 config: SparseVectorIndexConfig {
3960 embedding_function: Some(EmbeddingFunctionConfiguration::Legacy),
3961 source_key: None,
3962 bm25: None,
3963 algorithm: SparseIndexAlgorithm::Wand,
3964 },
3965 }),
3966 }),
3967 ..Default::default()
3968 },
3969 );
3970
3971 let result = schema.is_knn_key_indexing_enabled(
3972 "sparse_enabled",
3973 &QueryVector::Sparse(SparseVector::new(vec![0_u32], vec![1.0_f32]).unwrap()),
3974 );
3975
3976 assert!(result.is_ok());
3977 }
3978
3979 #[test]
3980 fn test_is_knn_key_indexing_enabled_dense_succeeds() {
3981 let schema = Schema::new_default(KnnIndex::Spann);
3982 let result = schema.is_knn_key_indexing_enabled(
3983 EMBEDDING_KEY,
3984 &QueryVector::Dense(vec![0.1_f32, 0.2_f32]),
3985 );
3986
3987 assert!(result.is_ok());
3988 }
3989
3990 #[test]
3991 fn test_combined_knn_rejects_unindexed_sparse_key() {
3992 use crate::operator::{Key, RankExpr};
3997
3998 let mut schema = Schema::new_default(KnnIndex::Spann);
3999 schema.keys.insert(
4000 "sparse_indexed".to_string(),
4001 ValueTypes {
4002 sparse_vector: Some(SparseVectorValueType {
4003 sparse_vector_index: Some(SparseVectorIndexType {
4004 enabled: true,
4005 config: SparseVectorIndexConfig {
4006 embedding_function: Some(EmbeddingFunctionConfiguration::Legacy),
4007 source_key: None,
4008 bm25: None,
4009 algorithm: SparseIndexAlgorithm::Wand,
4010 },
4011 }),
4012 }),
4013 ..Default::default()
4014 },
4015 );
4016
4017 let sparse_leaf = |key: &str| RankExpr::Knn {
4018 query: QueryVector::Sparse(SparseVector::new(vec![0_u32], vec![1.0_f32]).unwrap()),
4019 key: Key::field(key),
4020 limit: 10,
4021 default: None,
4022 return_rank: false,
4023 };
4024
4025 let expr = RankExpr::Summation(vec![
4027 sparse_leaf("sparse_indexed"),
4028 sparse_leaf("sparse_unindexed"),
4029 ]);
4030
4031 let mut results: Vec<(String, Result<(), FilterValidationError>)> = expr
4033 .knn_queries()
4034 .into_iter()
4035 .map(|knn| {
4036 let key = knn.key.to_string();
4037 let result = schema.is_knn_key_indexing_enabled(&key, &knn.query);
4038 (key, result)
4039 })
4040 .collect();
4041
4042 assert_eq!(results.len(), 2);
4043
4044 let (indexed_key, indexed_result) = results.remove(0);
4046 assert_eq!(indexed_key, "sparse_indexed");
4047 assert!(indexed_result.is_ok());
4048
4049 let (unindexed_key, unindexed_result) = results.remove(0);
4051 assert_eq!(unindexed_key, "sparse_unindexed");
4052 match unindexed_result.expect_err("expected unindexed sparse leaf to be rejected") {
4053 FilterValidationError::IndexingDisabled { key, value_type } => {
4054 assert_eq!(key, "sparse_unindexed");
4055 assert_eq!(value_type, crate::metadata::MetadataValueType::SparseVector);
4056 }
4057 other => panic!("unexpected error variant: {other:?}"),
4058 }
4059 }
4060
4061 #[test]
4062 fn test_merge_hnsw_configs_field_level() {
4063 let default_hnsw = HnswIndexConfig {
4065 ef_construction: Some(200),
4066 max_neighbors: Some(16),
4067 ef_search: Some(10),
4068 num_threads: Some(4),
4069 batch_size: Some(100),
4070 sync_threshold: Some(1000),
4071 resize_factor: Some(1.2),
4072 };
4073
4074 let user_hnsw = HnswIndexConfig {
4075 ef_construction: Some(300), max_neighbors: None, ef_search: Some(20), num_threads: None, batch_size: None, sync_threshold: Some(2000), resize_factor: None, };
4083
4084 let result = Schema::merge_hnsw_configs(Some(&default_hnsw), Some(&user_hnsw)).unwrap();
4085
4086 assert_eq!(result.ef_construction, Some(300));
4088 assert_eq!(result.ef_search, Some(20));
4089 assert_eq!(result.sync_threshold, Some(2000));
4090
4091 assert_eq!(result.max_neighbors, Some(16));
4093 assert_eq!(result.num_threads, Some(4));
4094 assert_eq!(result.batch_size, Some(100));
4095 assert_eq!(result.resize_factor, Some(1.2));
4096 }
4097
4098 #[test]
4099 fn test_merge_spann_configs_field_level() {
4100 let default_spann = SpannIndexConfig {
4102 search_nprobe: Some(10),
4103 search_rng_factor: Some(1.0), search_rng_epsilon: Some(7.0), nreplica_count: Some(3),
4106 write_rng_factor: Some(1.0), write_rng_epsilon: Some(6.0), split_threshold: Some(100), num_samples_kmeans: Some(100),
4110 initial_lambda: Some(100.0), reassign_neighbor_count: Some(50),
4112 merge_threshold: Some(50), num_centers_to_merge_to: Some(4), write_nprobe: Some(5),
4115 ef_construction: Some(100),
4116 ef_search: Some(10),
4117 max_neighbors: Some(16),
4118 center_drift_threshold: None,
4119 quantize: Quantization::None,
4120 };
4121
4122 let user_spann = SpannIndexConfig {
4123 search_nprobe: Some(20), search_rng_factor: None, search_rng_epsilon: Some(8.0), nreplica_count: None, write_rng_factor: None,
4128 write_rng_epsilon: None,
4129 split_threshold: Some(150), num_samples_kmeans: None,
4131 initial_lambda: None,
4132 reassign_neighbor_count: None,
4133 merge_threshold: None,
4134 num_centers_to_merge_to: None,
4135 write_nprobe: None,
4136 ef_construction: None,
4137 ef_search: None,
4138 max_neighbors: None,
4139 center_drift_threshold: None,
4140 quantize: Quantization::None,
4141 };
4142
4143 let result = Schema::merge_spann_configs(Some(&default_spann), Some(&user_spann))
4144 .unwrap()
4145 .unwrap();
4146
4147 assert_eq!(result.search_nprobe, Some(20));
4149 assert_eq!(result.search_rng_epsilon, Some(8.0));
4150 assert_eq!(result.split_threshold, Some(150));
4151
4152 assert_eq!(result.search_rng_factor, Some(1.0));
4154 assert_eq!(result.nreplica_count, Some(3));
4155 assert_eq!(result.initial_lambda, Some(100.0));
4156 }
4157
4158 #[test]
4159 fn test_merge_spann_configs_rejects_quantize_true() {
4160 let default_spann = SpannIndexConfig {
4162 search_nprobe: Some(10),
4163 search_rng_factor: Some(1.0),
4164 search_rng_epsilon: Some(7.0),
4165 nreplica_count: Some(3),
4166 write_rng_factor: Some(1.0),
4167 write_rng_epsilon: Some(6.0),
4168 split_threshold: Some(100),
4169 num_samples_kmeans: Some(100),
4170 initial_lambda: Some(100.0),
4171 reassign_neighbor_count: Some(50),
4172 merge_threshold: Some(50),
4173 num_centers_to_merge_to: Some(4),
4174 write_nprobe: Some(5),
4175 ef_construction: Some(100),
4176 ef_search: Some(10),
4177 max_neighbors: Some(16),
4178 center_drift_threshold: None,
4179 quantize: Quantization::None,
4180 };
4181
4182 let user_spann_with_quantize = SpannIndexConfig {
4183 search_nprobe: Some(20),
4184 search_rng_factor: None,
4185 search_rng_epsilon: Some(8.0),
4186 nreplica_count: None,
4187 write_rng_factor: None,
4188 write_rng_epsilon: None,
4189 split_threshold: Some(150),
4190 num_samples_kmeans: None,
4191 initial_lambda: None,
4192 reassign_neighbor_count: None,
4193 merge_threshold: None,
4194 num_centers_to_merge_to: None,
4195 write_nprobe: None,
4196 ef_construction: None,
4197 ef_search: None,
4198 max_neighbors: None,
4199 center_drift_threshold: None,
4200 quantize: Quantization::FourBitRabitQWithUSearch, };
4202
4203 let result =
4205 Schema::merge_spann_configs(Some(&default_spann), Some(&user_spann_with_quantize));
4206 assert!(result.is_err());
4207 match result {
4208 Err(SchemaError::InvalidUserInput { reason }) => {
4209 assert!(reason.contains("quantize field cannot be set"));
4210 }
4211 _ => panic!("Expected InvalidUserInput error"),
4212 }
4213
4214 let default_spann_with_quantize = SpannIndexConfig {
4216 search_nprobe: Some(10),
4217 search_rng_factor: Some(1.0),
4218 search_rng_epsilon: Some(7.0),
4219 nreplica_count: Some(3),
4220 write_rng_factor: Some(1.0),
4221 write_rng_epsilon: Some(6.0),
4222 split_threshold: Some(100),
4223 num_samples_kmeans: Some(100),
4224 initial_lambda: Some(100.0),
4225 reassign_neighbor_count: Some(50),
4226 merge_threshold: Some(50),
4227 num_centers_to_merge_to: Some(4),
4228 write_nprobe: Some(5),
4229 ef_construction: Some(100),
4230 ef_search: Some(10),
4231 max_neighbors: Some(16),
4232 center_drift_threshold: None,
4233 quantize: Quantization::FourBitRabitQWithUSearch, };
4235
4236 let result = Schema::merge_spann_configs(Some(&default_spann_with_quantize), None);
4237 assert!(result.is_err());
4238 match result {
4239 Err(SchemaError::InvalidUserInput { reason }) => {
4240 assert!(reason.contains("quantize field cannot be set"));
4241 }
4242 _ => panic!("Expected InvalidUserInput error"),
4243 }
4244
4245 let result = Schema::merge_spann_configs(None, Some(&user_spann_with_quantize));
4247 assert!(result.is_err());
4248 match result {
4249 Err(SchemaError::InvalidUserInput { reason }) => {
4250 assert!(reason.contains("quantize field cannot be set"));
4251 }
4252 _ => panic!("Expected InvalidUserInput error"),
4253 }
4254 }
4255
4256 #[test]
4257 fn test_spann_index_config_into_internal_configuration() {
4258 let config = SpannIndexConfig {
4259 search_nprobe: Some(33),
4260 search_rng_factor: Some(1.2),
4261 search_rng_epsilon: None,
4262 nreplica_count: None,
4263 write_rng_factor: Some(1.5),
4264 write_rng_epsilon: None,
4265 split_threshold: Some(75),
4266 num_samples_kmeans: None,
4267 initial_lambda: Some(0.9),
4268 reassign_neighbor_count: Some(40),
4269 merge_threshold: None,
4270 num_centers_to_merge_to: Some(4),
4271 write_nprobe: Some(60),
4272 ef_construction: Some(180),
4273 ef_search: Some(170),
4274 max_neighbors: Some(32),
4275 center_drift_threshold: None,
4276 quantize: Quantization::None,
4277 };
4278
4279 let with_space: InternalSpannConfiguration = (Some(&Space::Cosine), &config).into();
4280 assert_eq!(with_space.space, Space::Cosine);
4281 assert_eq!(with_space.search_nprobe, 33);
4282 assert_eq!(with_space.search_rng_factor, 1.2);
4283 assert_eq!(with_space.search_rng_epsilon, default_search_rng_epsilon());
4284 assert_eq!(with_space.write_rng_factor, 1.5);
4285 assert_eq!(with_space.write_nprobe, 60);
4286 assert_eq!(with_space.ef_construction, 180);
4287 assert_eq!(with_space.ef_search, 170);
4288 assert_eq!(with_space.max_neighbors, 32);
4289 assert_eq!(with_space.merge_threshold, default_merge_threshold());
4290
4291 let default_space_config: InternalSpannConfiguration = (None, &config).into();
4292 assert_eq!(default_space_config.space, default_space());
4293 }
4294
4295 #[test]
4296 fn test_merge_string_type_combinations() {
4297 let default = StringValueType {
4301 string_inverted_index: Some(StringInvertedIndexType {
4302 enabled: true,
4303 config: StringInvertedIndexConfig {},
4304 }),
4305 fts_index: Some(FtsIndexType {
4306 enabled: false,
4307 config: FtsIndexConfig::default(),
4308 }),
4309 };
4310
4311 let user = StringValueType {
4312 string_inverted_index: Some(StringInvertedIndexType {
4313 enabled: false, config: StringInvertedIndexConfig {},
4315 }),
4316 fts_index: None, };
4318
4319 let result = Schema::merge_string_type(Some(&default), Some(&user))
4320 .unwrap()
4321 .unwrap();
4322 assert!(!result.string_inverted_index.as_ref().unwrap().enabled); assert!(!result.fts_index.as_ref().unwrap().enabled); let result = Schema::merge_string_type(Some(&default), None)
4327 .unwrap()
4328 .unwrap();
4329 assert!(result.string_inverted_index.as_ref().unwrap().enabled);
4330
4331 let result = Schema::merge_string_type(None, Some(&user))
4333 .unwrap()
4334 .unwrap();
4335 assert!(!result.string_inverted_index.as_ref().unwrap().enabled);
4336
4337 let result = Schema::merge_string_type(None, None).unwrap();
4339 assert!(result.is_none());
4340 }
4341
4342 #[test]
4343 fn test_merge_vector_index_config_comprehensive() {
4344 let default_config = VectorIndexConfig {
4346 space: Some(Space::Cosine),
4347 embedding_function: Some(EmbeddingFunctionConfiguration::Legacy),
4348 source_key: Some("default_key".to_string()),
4349 hnsw: Some(HnswIndexConfig {
4350 ef_construction: Some(200),
4351 max_neighbors: Some(16),
4352 ef_search: Some(10),
4353 num_threads: Some(4),
4354 batch_size: Some(100),
4355 sync_threshold: Some(1000),
4356 resize_factor: Some(1.2),
4357 }),
4358 spann: None,
4359 };
4360
4361 let user_config = VectorIndexConfig {
4362 space: Some(Space::L2), embedding_function: None, source_key: Some("user_key".to_string()), hnsw: Some(HnswIndexConfig {
4366 ef_construction: Some(300), max_neighbors: None, ef_search: None, num_threads: None,
4370 batch_size: None,
4371 sync_threshold: None,
4372 resize_factor: None,
4373 }),
4374 spann: Some(SpannIndexConfig {
4375 search_nprobe: Some(15),
4376 search_rng_factor: None,
4377 search_rng_epsilon: None,
4378 nreplica_count: None,
4379 write_rng_factor: None,
4380 write_rng_epsilon: None,
4381 split_threshold: None,
4382 num_samples_kmeans: None,
4383 initial_lambda: None,
4384 reassign_neighbor_count: None,
4385 merge_threshold: None,
4386 num_centers_to_merge_to: None,
4387 write_nprobe: None,
4388 ef_construction: None,
4389 ef_search: None,
4390 max_neighbors: None,
4391 center_drift_threshold: None,
4392 quantize: Quantization::None,
4393 }), };
4395
4396 let result =
4397 Schema::merge_vector_index_config(&default_config, &user_config, KnnIndex::Hnsw)
4398 .expect("merge should succeed");
4399
4400 assert_eq!(result.space, Some(Space::L2)); assert_eq!(
4403 result.embedding_function,
4404 Some(EmbeddingFunctionConfiguration::Legacy)
4405 ); assert_eq!(result.source_key, Some("user_key".to_string())); assert_eq!(result.hnsw.as_ref().unwrap().ef_construction, Some(300)); assert_eq!(result.hnsw.as_ref().unwrap().max_neighbors, Some(16)); assert!(result.spann.is_none());
4414 }
4415
4416 #[test]
4417 fn test_merge_sparse_vector_index_config() {
4418 let default_config = SparseVectorIndexConfig {
4420 embedding_function: Some(EmbeddingFunctionConfiguration::Legacy),
4421 source_key: Some("default_sparse_key".to_string()),
4422 bm25: None,
4423 algorithm: SparseIndexAlgorithm::Wand,
4424 };
4425
4426 let user_config = SparseVectorIndexConfig {
4427 embedding_function: None, source_key: Some("user_sparse_key".to_string()), bm25: None,
4430 algorithm: SparseIndexAlgorithm::Wand,
4431 };
4432
4433 let result = Schema::merge_sparse_vector_index_config(&default_config, &user_config);
4434
4435 assert_eq!(result.source_key, Some("user_sparse_key".to_string()));
4437 assert_eq!(
4439 result.embedding_function,
4440 Some(EmbeddingFunctionConfiguration::Legacy)
4441 );
4442 }
4443
4444 #[test]
4445 fn test_sparse_algorithm_serde_roundtrip() {
4446 let json_no_algorithm = r#"{
4448 "embedding_function": null,
4449 "source_key": null,
4450 "bm25": null
4451 }"#;
4452 let config: SparseVectorIndexConfig = serde_json::from_str(json_no_algorithm).unwrap();
4453 assert_eq!(config.algorithm, SparseIndexAlgorithm::Wand);
4454
4455 let serialized = serde_json::to_string(&config).unwrap();
4457 assert!(
4458 !serialized.contains("algorithm"),
4459 "Wand (default) must not appear in JSON: {serialized}"
4460 );
4461
4462 let json_maxscore = r#"{
4464 "embedding_function": null,
4465 "source_key": null,
4466 "bm25": null,
4467 "algorithm": "max_score"
4468 }"#;
4469 let config: SparseVectorIndexConfig = serde_json::from_str(json_maxscore).unwrap();
4470 assert_eq!(config.algorithm, SparseIndexAlgorithm::MaxScore);
4471
4472 let serialized = serde_json::to_string(&config).unwrap();
4474 assert!(
4475 serialized.contains(r#""algorithm":"max_score""#),
4476 "MaxScore must appear in JSON: {serialized}"
4477 );
4478 }
4479
4480 #[test]
4481 fn test_is_maxscore_enabled() {
4482 let schema = Schema::default();
4484 assert!(!schema.is_maxscore_enabled());
4485
4486 let mut schema = Schema::new_default(KnnIndex::Hnsw);
4488 schema
4489 .keys
4490 .entry("sparse_key".to_string())
4491 .or_default()
4492 .sparse_vector = Some(SparseVectorValueType {
4493 sparse_vector_index: Some(SparseVectorIndexType {
4494 enabled: true,
4495 config: SparseVectorIndexConfig {
4496 embedding_function: None,
4497 source_key: None,
4498 bm25: None,
4499 algorithm: SparseIndexAlgorithm::Wand,
4500 },
4501 }),
4502 });
4503 assert!(!schema.is_maxscore_enabled());
4504
4505 schema.set_sparse_algorithm(SparseIndexAlgorithm::MaxScore);
4507 assert!(schema.is_maxscore_enabled());
4508
4509 schema
4511 .keys
4512 .get_mut("sparse_key")
4513 .unwrap()
4514 .sparse_vector
4515 .as_mut()
4516 .unwrap()
4517 .sparse_vector_index
4518 .as_mut()
4519 .unwrap()
4520 .enabled = false;
4521 assert!(!schema.is_maxscore_enabled());
4522 }
4523
4524 fn enable_sparse_key(schema: &mut Schema, key: &str, algorithm: SparseIndexAlgorithm) {
4525 schema
4526 .keys
4527 .entry(key.to_string())
4528 .or_default()
4529 .sparse_vector = Some(SparseVectorValueType {
4530 sparse_vector_index: Some(SparseVectorIndexType {
4531 enabled: true,
4532 config: SparseVectorIndexConfig {
4533 embedding_function: None,
4534 source_key: None,
4535 bm25: None,
4536 algorithm,
4537 },
4538 }),
4539 });
4540 }
4541
4542 #[test]
4543 fn test_enabled_sparse_keys() {
4544 let mut schema = Schema::new_default(KnnIndex::Hnsw);
4545 assert!(schema.enabled_sparse_keys().is_empty());
4546
4547 enable_sparse_key(&mut schema, "beta", SparseIndexAlgorithm::Wand);
4548 enable_sparse_key(&mut schema, "alpha", SparseIndexAlgorithm::MaxScore);
4549 enable_sparse_key(&mut schema, "gamma", SparseIndexAlgorithm::Wand);
4551 schema
4552 .keys
4553 .get_mut("gamma")
4554 .unwrap()
4555 .sparse_vector
4556 .as_mut()
4557 .unwrap()
4558 .sparse_vector_index
4559 .as_mut()
4560 .unwrap()
4561 .enabled = false;
4562
4563 assert_eq!(
4565 schema.enabled_sparse_keys(),
4566 vec!["alpha".to_string(), "beta".to_string()]
4567 );
4568 }
4569
4570 #[test]
4571 fn test_is_key_maxscore_enabled() {
4572 let mut schema = Schema::new_default(KnnIndex::Hnsw);
4573 enable_sparse_key(&mut schema, "ms", SparseIndexAlgorithm::MaxScore);
4574 enable_sparse_key(&mut schema, "wand", SparseIndexAlgorithm::Wand);
4575
4576 assert!(schema.is_key_maxscore_enabled("ms"));
4577 assert!(!schema.is_key_maxscore_enabled("wand"));
4578 assert!(!schema.is_key_maxscore_enabled("unknown"));
4580 }
4581
4582 #[test]
4583 fn test_fts_algorithm_serde_roundtrip() {
4584 let json_no_algorithm = r#"{}"#;
4586 let config: FtsIndexConfig = serde_json::from_str(json_no_algorithm).unwrap();
4587 assert_eq!(config.algorithm, FtsAlgorithm::Trigram);
4588
4589 let serialized = serde_json::to_string(&config).unwrap();
4591 assert!(
4592 !serialized.contains("algorithm"),
4593 "Trigram (default) must not appear in JSON: {serialized}"
4594 );
4595
4596 let json_token_bitmap = r#"{"algorithm": "token_bitmap"}"#;
4598 let config: FtsIndexConfig = serde_json::from_str(json_token_bitmap).unwrap();
4599 assert_eq!(config.algorithm, FtsAlgorithm::TokenBitmap);
4600
4601 let serialized = serde_json::to_string(&config).unwrap();
4603 assert!(
4604 serialized.contains(r#""algorithm":"token_bitmap""#),
4605 "TokenBitmap must appear in JSON: {serialized}"
4606 );
4607 }
4608
4609 #[test]
4610 fn test_is_token_bitmap_fts_enabled() {
4611 let schema = Schema::default();
4613 assert!(!schema.is_token_bitmap_fts_enabled());
4614
4615 let mut schema = Schema::new_default(KnnIndex::Hnsw);
4617 assert!(!schema.is_token_bitmap_fts_enabled());
4618
4619 schema.set_fts_algorithm(FtsAlgorithm::TokenBitmap);
4621 assert!(schema.is_token_bitmap_fts_enabled());
4622
4623 schema
4625 .keys
4626 .get_mut(DOCUMENT_KEY)
4627 .unwrap()
4628 .string
4629 .as_mut()
4630 .unwrap()
4631 .fts_index
4632 .as_mut()
4633 .unwrap()
4634 .enabled = false;
4635 assert!(!schema.is_token_bitmap_fts_enabled());
4636 }
4637
4638 #[test]
4639 fn test_complex_nested_merging_scenario() {
4640 let mut user_schema = Schema {
4642 defaults: ValueTypes::default(),
4643 keys: HashMap::new(),
4644 cmek: None,
4645 };
4646
4647 user_schema.defaults.string = Some(StringValueType {
4649 string_inverted_index: Some(StringInvertedIndexType {
4650 enabled: false,
4651 config: StringInvertedIndexConfig {},
4652 }),
4653 fts_index: Some(FtsIndexType {
4654 enabled: true,
4655 config: FtsIndexConfig::default(),
4656 }),
4657 });
4658
4659 user_schema.defaults.float_list = Some(FloatListValueType {
4660 vector_index: Some(VectorIndexType {
4661 enabled: true,
4662 config: VectorIndexConfig {
4663 space: Some(Space::Ip),
4664 embedding_function: None, source_key: Some("custom_vector_key".to_string()),
4666 hnsw: Some(HnswIndexConfig {
4667 ef_construction: Some(400),
4668 max_neighbors: Some(32),
4669 ef_search: None, num_threads: None,
4671 batch_size: None,
4672 sync_threshold: None,
4673 resize_factor: None,
4674 }),
4675 spann: None,
4676 },
4677 }),
4678 });
4679
4680 let custom_key_override = ValueTypes {
4682 string: Some(StringValueType {
4683 fts_index: Some(FtsIndexType {
4684 enabled: true,
4685 config: FtsIndexConfig::default(),
4686 }),
4687 string_inverted_index: None,
4688 }),
4689 ..Default::default()
4690 };
4691 user_schema
4692 .keys
4693 .insert("custom_field".to_string(), custom_key_override);
4694
4695 let result = {
4697 let default_schema = Schema::new_default(KnnIndex::Hnsw);
4698 let merged_defaults = Schema::merge_value_types(
4699 &default_schema.defaults,
4700 &user_schema.defaults,
4701 KnnIndex::Hnsw,
4702 )
4703 .unwrap();
4704 let mut merged_keys = default_schema.keys.clone();
4705 for (key, user_value_types) in user_schema.keys {
4706 if let Some(default_value_types) = merged_keys.get(&key) {
4707 let merged_value_types = Schema::merge_value_types(
4708 default_value_types,
4709 &user_value_types,
4710 KnnIndex::Hnsw,
4711 )
4712 .unwrap();
4713 merged_keys.insert(key, merged_value_types);
4714 } else {
4715 merged_keys.insert(key, user_value_types);
4716 }
4717 }
4718 Schema {
4719 defaults: merged_defaults,
4720 keys: merged_keys,
4721 cmek: None,
4722 }
4723 };
4724
4725 assert!(
4729 !result
4730 .defaults
4731 .string
4732 .as_ref()
4733 .unwrap()
4734 .string_inverted_index
4735 .as_ref()
4736 .unwrap()
4737 .enabled
4738 );
4739 assert!(
4740 result
4741 .defaults
4742 .string
4743 .as_ref()
4744 .unwrap()
4745 .fts_index
4746 .as_ref()
4747 .unwrap()
4748 .enabled
4749 );
4750
4751 let vector_config = &result
4752 .defaults
4753 .float_list
4754 .as_ref()
4755 .unwrap()
4756 .vector_index
4757 .as_ref()
4758 .unwrap()
4759 .config;
4760 assert_eq!(vector_config.space, Some(Space::Ip));
4761 assert_eq!(vector_config.embedding_function, None); assert_eq!(
4763 vector_config.source_key,
4764 Some("custom_vector_key".to_string())
4765 );
4766 assert_eq!(
4767 vector_config.hnsw.as_ref().unwrap().ef_construction,
4768 Some(400)
4769 );
4770 assert_eq!(vector_config.hnsw.as_ref().unwrap().max_neighbors, Some(32));
4771 assert_eq!(
4772 vector_config.hnsw.as_ref().unwrap().ef_search,
4773 Some(default_search_ef())
4774 ); assert!(result.keys.contains_key(EMBEDDING_KEY)); assert!(result.keys.contains_key(DOCUMENT_KEY)); assert!(result.keys.contains_key("custom_field")); let custom_override = result.keys.get("custom_field").unwrap();
4782 assert!(
4783 custom_override
4784 .string
4785 .as_ref()
4786 .unwrap()
4787 .fts_index
4788 .as_ref()
4789 .unwrap()
4790 .enabled
4791 );
4792 assert!(custom_override
4793 .string
4794 .as_ref()
4795 .unwrap()
4796 .string_inverted_index
4797 .is_none());
4798 }
4799
4800 #[test]
4801 fn test_reconcile_with_collection_config_default_config() {
4802 let collection_config = InternalCollectionConfiguration::default_hnsw();
4804 let schema = Schema::try_from(&collection_config).unwrap();
4805
4806 let result =
4807 Schema::reconcile_with_collection_config(&schema, &collection_config, KnnIndex::Hnsw)
4808 .unwrap();
4809 assert_eq!(result, schema);
4810 }
4811
4812 #[test]
4814 fn test_reconcile_double_default_hnsw_config_hnsw_schema_default_knn_hnsw() {
4815 let collection_config = InternalCollectionConfiguration::default_hnsw();
4816 let schema = Schema::new_default(KnnIndex::Hnsw);
4817 let result =
4818 Schema::reconcile_with_collection_config(&schema, &collection_config, KnnIndex::Hnsw)
4819 .unwrap();
4820
4821 assert!(result.defaults.float_list.is_some());
4823 assert!(result
4824 .defaults
4825 .float_list
4826 .as_ref()
4827 .unwrap()
4828 .vector_index
4829 .as_ref()
4830 .unwrap()
4831 .config
4832 .hnsw
4833 .is_some());
4834 assert!(result
4835 .defaults
4836 .float_list
4837 .as_ref()
4838 .unwrap()
4839 .vector_index
4840 .as_ref()
4841 .unwrap()
4842 .config
4843 .spann
4844 .is_none());
4845 }
4846
4847 #[test]
4848 fn test_reconcile_double_default_hnsw_config_hnsw_schema_default_knn_spann() {
4849 let collection_config = InternalCollectionConfiguration::default_hnsw();
4850 let schema = Schema::new_default(KnnIndex::Hnsw);
4851 let result =
4852 Schema::reconcile_with_collection_config(&schema, &collection_config, KnnIndex::Spann)
4853 .unwrap();
4854
4855 assert!(result.defaults.float_list.is_some());
4857 assert!(result
4858 .defaults
4859 .float_list
4860 .as_ref()
4861 .unwrap()
4862 .vector_index
4863 .as_ref()
4864 .unwrap()
4865 .config
4866 .spann
4867 .is_some());
4868 assert!(result
4869 .defaults
4870 .float_list
4871 .as_ref()
4872 .unwrap()
4873 .vector_index
4874 .as_ref()
4875 .unwrap()
4876 .config
4877 .hnsw
4878 .is_none());
4879 }
4880
4881 #[test]
4882 fn test_reconcile_double_default_hnsw_config_spann_schema_default_knn_hnsw() {
4883 let collection_config = InternalCollectionConfiguration::default_hnsw();
4884 let schema = Schema::new_default(KnnIndex::Spann);
4885 let result =
4886 Schema::reconcile_with_collection_config(&schema, &collection_config, KnnIndex::Hnsw)
4887 .unwrap();
4888
4889 assert!(result.defaults.float_list.is_some());
4891 assert!(result
4892 .defaults
4893 .float_list
4894 .as_ref()
4895 .unwrap()
4896 .vector_index
4897 .as_ref()
4898 .unwrap()
4899 .config
4900 .hnsw
4901 .is_some());
4902 assert!(result
4903 .defaults
4904 .float_list
4905 .as_ref()
4906 .unwrap()
4907 .vector_index
4908 .as_ref()
4909 .unwrap()
4910 .config
4911 .spann
4912 .is_none());
4913 }
4914
4915 #[test]
4916 fn test_reconcile_double_default_hnsw_config_spann_schema_default_knn_spann() {
4917 let collection_config = InternalCollectionConfiguration::default_hnsw();
4918 let schema = Schema::new_default(KnnIndex::Spann);
4919 let result =
4920 Schema::reconcile_with_collection_config(&schema, &collection_config, KnnIndex::Spann)
4921 .unwrap();
4922
4923 assert!(result.defaults.float_list.is_some());
4925 assert!(result
4926 .defaults
4927 .float_list
4928 .as_ref()
4929 .unwrap()
4930 .vector_index
4931 .as_ref()
4932 .unwrap()
4933 .config
4934 .spann
4935 .is_some());
4936 assert!(result
4937 .defaults
4938 .float_list
4939 .as_ref()
4940 .unwrap()
4941 .vector_index
4942 .as_ref()
4943 .unwrap()
4944 .config
4945 .hnsw
4946 .is_none());
4947 }
4948
4949 #[test]
4950 fn test_reconcile_double_default_spann_config_spann_schema_default_knn_hnsw() {
4951 let collection_config = InternalCollectionConfiguration::default_spann();
4952 let schema = Schema::new_default(KnnIndex::Spann);
4953 let result =
4954 Schema::reconcile_with_collection_config(&schema, &collection_config, KnnIndex::Hnsw)
4955 .unwrap();
4956
4957 assert!(result.defaults.float_list.is_some());
4959 assert!(result
4960 .defaults
4961 .float_list
4962 .as_ref()
4963 .unwrap()
4964 .vector_index
4965 .as_ref()
4966 .unwrap()
4967 .config
4968 .hnsw
4969 .is_some());
4970 assert!(result
4971 .defaults
4972 .float_list
4973 .as_ref()
4974 .unwrap()
4975 .vector_index
4976 .as_ref()
4977 .unwrap()
4978 .config
4979 .spann
4980 .is_none());
4981 }
4982
4983 #[test]
4984 fn test_reconcile_double_default_spann_config_spann_schema_default_knn_spann() {
4985 let collection_config = InternalCollectionConfiguration::default_spann();
4986 let schema = Schema::new_default(KnnIndex::Spann);
4987 let result =
4988 Schema::reconcile_with_collection_config(&schema, &collection_config, KnnIndex::Spann)
4989 .unwrap();
4990
4991 assert!(result.defaults.float_list.is_some());
4993 assert!(result
4994 .defaults
4995 .float_list
4996 .as_ref()
4997 .unwrap()
4998 .vector_index
4999 .as_ref()
5000 .unwrap()
5001 .config
5002 .spann
5003 .is_some());
5004 assert!(result
5005 .defaults
5006 .float_list
5007 .as_ref()
5008 .unwrap()
5009 .vector_index
5010 .as_ref()
5011 .unwrap()
5012 .config
5013 .hnsw
5014 .is_none());
5015 assert_eq!(
5017 result
5018 .defaults
5019 .float_list
5020 .as_ref()
5021 .unwrap()
5022 .vector_index
5023 .as_ref()
5024 .unwrap()
5025 .config
5026 .source_key,
5027 None
5028 );
5029 }
5030
5031 #[test]
5032 fn test_reconcile_double_default_spann_config_hnsw_schema_default_knn_hnsw() {
5033 let collection_config = InternalCollectionConfiguration::default_spann();
5034 let schema = Schema::new_default(KnnIndex::Hnsw);
5035 let result =
5036 Schema::reconcile_with_collection_config(&schema, &collection_config, KnnIndex::Hnsw)
5037 .unwrap();
5038
5039 assert!(result.defaults.float_list.is_some());
5041 assert!(result
5042 .defaults
5043 .float_list
5044 .as_ref()
5045 .unwrap()
5046 .vector_index
5047 .as_ref()
5048 .unwrap()
5049 .config
5050 .hnsw
5051 .is_some());
5052 assert!(result
5053 .defaults
5054 .float_list
5055 .as_ref()
5056 .unwrap()
5057 .vector_index
5058 .as_ref()
5059 .unwrap()
5060 .config
5061 .spann
5062 .is_none());
5063 }
5064
5065 #[test]
5066 fn test_reconcile_double_default_spann_config_hnsw_schema_default_knn_spann() {
5067 let collection_config = InternalCollectionConfiguration::default_spann();
5068 let schema = Schema::new_default(KnnIndex::Hnsw);
5069 let result =
5070 Schema::reconcile_with_collection_config(&schema, &collection_config, KnnIndex::Spann)
5071 .unwrap();
5072
5073 assert!(result.defaults.float_list.is_some());
5075 assert!(result
5076 .defaults
5077 .float_list
5078 .as_ref()
5079 .unwrap()
5080 .vector_index
5081 .as_ref()
5082 .unwrap()
5083 .config
5084 .spann
5085 .is_some());
5086 assert!(result
5087 .defaults
5088 .float_list
5089 .as_ref()
5090 .unwrap()
5091 .vector_index
5092 .as_ref()
5093 .unwrap()
5094 .config
5095 .hnsw
5096 .is_none());
5097 }
5098
5099 #[test]
5100 fn test_defaults_source_key_not_document() {
5101 let schema_hnsw = Schema::new_default(KnnIndex::Hnsw);
5103 let schema_spann = Schema::new_default(KnnIndex::Spann);
5104
5105 let defaults_hnsw = schema_hnsw
5107 .defaults
5108 .float_list
5109 .as_ref()
5110 .unwrap()
5111 .vector_index
5112 .as_ref()
5113 .unwrap();
5114 assert_eq!(defaults_hnsw.config.source_key, None);
5115
5116 let defaults_spann = schema_spann
5118 .defaults
5119 .float_list
5120 .as_ref()
5121 .unwrap()
5122 .vector_index
5123 .as_ref()
5124 .unwrap();
5125 assert_eq!(defaults_spann.config.source_key, None);
5126
5127 let collection_config_hnsw = InternalCollectionConfiguration {
5130 vector_index: VectorIndexConfiguration::Hnsw(InternalHnswConfiguration {
5131 ef_construction: 300,
5132 max_neighbors: 32,
5133 ef_search: 50,
5134 num_threads: 8,
5135 batch_size: 200,
5136 sync_threshold: 2000,
5137 resize_factor: 1.5,
5138 space: Space::L2,
5139 }),
5140 embedding_function: Some(EmbeddingFunctionConfiguration::Legacy),
5141 };
5142 let result_hnsw = Schema::reconcile_with_collection_config(
5143 &schema_hnsw,
5144 &collection_config_hnsw,
5145 KnnIndex::Hnsw,
5146 )
5147 .unwrap();
5148 let reconciled_defaults_hnsw = result_hnsw
5149 .defaults
5150 .float_list
5151 .as_ref()
5152 .unwrap()
5153 .vector_index
5154 .as_ref()
5155 .unwrap();
5156 assert_eq!(reconciled_defaults_hnsw.config.source_key, None);
5157
5158 let collection_config_spann = InternalCollectionConfiguration {
5159 vector_index: VectorIndexConfiguration::Spann(InternalSpannConfiguration {
5160 search_nprobe: 20,
5161 search_rng_factor: 3.0,
5162 search_rng_epsilon: 0.2,
5163 nreplica_count: 5,
5164 write_rng_factor: 2.0,
5165 write_rng_epsilon: 0.1,
5166 split_threshold: 2000,
5167 num_samples_kmeans: 200,
5168 initial_lambda: 0.8,
5169 reassign_neighbor_count: 100,
5170 merge_threshold: 800,
5171 num_centers_to_merge_to: 20,
5172 write_nprobe: 10,
5173 ef_construction: 400,
5174 ef_search: 60,
5175 max_neighbors: 24,
5176 space: Space::Cosine,
5177 }),
5178 embedding_function: None,
5179 };
5180 let result_spann = Schema::reconcile_with_collection_config(
5181 &schema_spann,
5182 &collection_config_spann,
5183 KnnIndex::Spann,
5184 )
5185 .unwrap();
5186 let reconciled_defaults_spann = result_spann
5187 .defaults
5188 .float_list
5189 .as_ref()
5190 .unwrap()
5191 .vector_index
5192 .as_ref()
5193 .unwrap();
5194 assert_eq!(reconciled_defaults_spann.config.source_key, None);
5195
5196 let embedding_hnsw = result_hnsw.keys.get(EMBEDDING_KEY).unwrap();
5198 let embedding_vector_index_hnsw = embedding_hnsw
5199 .float_list
5200 .as_ref()
5201 .unwrap()
5202 .vector_index
5203 .as_ref()
5204 .unwrap();
5205 assert_eq!(
5206 embedding_vector_index_hnsw.config.source_key,
5207 Some(DOCUMENT_KEY.to_string())
5208 );
5209
5210 let embedding_spann = result_spann.keys.get(EMBEDDING_KEY).unwrap();
5211 let embedding_vector_index_spann = embedding_spann
5212 .float_list
5213 .as_ref()
5214 .unwrap()
5215 .vector_index
5216 .as_ref()
5217 .unwrap();
5218 assert_eq!(
5219 embedding_vector_index_spann.config.source_key,
5220 Some(DOCUMENT_KEY.to_string())
5221 );
5222 }
5223
5224 #[test]
5225 fn test_try_from_source_key() {
5226 let collection_config_hnsw = InternalCollectionConfiguration {
5231 vector_index: VectorIndexConfiguration::Hnsw(InternalHnswConfiguration {
5232 ef_construction: 300,
5233 max_neighbors: 32,
5234 ef_search: 50,
5235 num_threads: 8,
5236 batch_size: 200,
5237 sync_threshold: 2000,
5238 resize_factor: 1.5,
5239 space: Space::L2,
5240 }),
5241 embedding_function: Some(EmbeddingFunctionConfiguration::Legacy),
5242 };
5243 let schema_hnsw = Schema::try_from(&collection_config_hnsw).unwrap();
5244
5245 let defaults_hnsw = schema_hnsw
5247 .defaults
5248 .float_list
5249 .as_ref()
5250 .unwrap()
5251 .vector_index
5252 .as_ref()
5253 .unwrap();
5254 assert_eq!(defaults_hnsw.config.source_key, None);
5255
5256 let embedding_hnsw = schema_hnsw.keys.get(EMBEDDING_KEY).unwrap();
5258 let embedding_vector_index_hnsw = embedding_hnsw
5259 .float_list
5260 .as_ref()
5261 .unwrap()
5262 .vector_index
5263 .as_ref()
5264 .unwrap();
5265 assert_eq!(
5266 embedding_vector_index_hnsw.config.source_key,
5267 Some(DOCUMENT_KEY.to_string())
5268 );
5269
5270 let collection_config_spann = InternalCollectionConfiguration {
5272 vector_index: VectorIndexConfiguration::Spann(InternalSpannConfiguration {
5273 search_nprobe: 20,
5274 search_rng_factor: 3.0,
5275 search_rng_epsilon: 0.2,
5276 nreplica_count: 5,
5277 write_rng_factor: 2.0,
5278 write_rng_epsilon: 0.1,
5279 split_threshold: 2000,
5280 num_samples_kmeans: 200,
5281 initial_lambda: 0.8,
5282 reassign_neighbor_count: 100,
5283 merge_threshold: 800,
5284 num_centers_to_merge_to: 20,
5285 write_nprobe: 10,
5286 ef_construction: 400,
5287 ef_search: 60,
5288 max_neighbors: 24,
5289 space: Space::Cosine,
5290 }),
5291 embedding_function: None,
5292 };
5293 let schema_spann = Schema::try_from(&collection_config_spann).unwrap();
5294
5295 let defaults_spann = schema_spann
5297 .defaults
5298 .float_list
5299 .as_ref()
5300 .unwrap()
5301 .vector_index
5302 .as_ref()
5303 .unwrap();
5304 assert_eq!(defaults_spann.config.source_key, None);
5305
5306 let embedding_spann = schema_spann.keys.get(EMBEDDING_KEY).unwrap();
5308 let embedding_vector_index_spann = embedding_spann
5309 .float_list
5310 .as_ref()
5311 .unwrap()
5312 .vector_index
5313 .as_ref()
5314 .unwrap();
5315 assert_eq!(
5316 embedding_vector_index_spann.config.source_key,
5317 Some(DOCUMENT_KEY.to_string())
5318 );
5319 }
5320
5321 #[test]
5322 fn test_default_hnsw_with_default_embedding_function() {
5323 use crate::collection_configuration::EmbeddingFunctionNewConfiguration;
5327
5328 let collection_config = InternalCollectionConfiguration {
5329 vector_index: VectorIndexConfiguration::Hnsw(InternalHnswConfiguration::default()),
5330 embedding_function: Some(EmbeddingFunctionConfiguration::Known(
5331 EmbeddingFunctionNewConfiguration {
5332 name: "default".to_string(),
5333 config: serde_json::json!({}),
5334 },
5335 )),
5336 };
5337
5338 assert!(collection_config.is_default());
5340
5341 let schema = Schema::new_default(KnnIndex::Hnsw);
5342 let result =
5343 Schema::reconcile_with_collection_config(&schema, &collection_config, KnnIndex::Spann)
5344 .unwrap();
5345
5346 let defaults = result
5348 .defaults
5349 .float_list
5350 .as_ref()
5351 .unwrap()
5352 .vector_index
5353 .as_ref()
5354 .unwrap();
5355 assert_eq!(defaults.config.source_key, None);
5356
5357 let embedding = result.keys.get(EMBEDDING_KEY).unwrap();
5359 let embedding_vector_index = embedding
5360 .float_list
5361 .as_ref()
5362 .unwrap()
5363 .vector_index
5364 .as_ref()
5365 .unwrap();
5366 assert_eq!(
5367 embedding_vector_index.config.source_key,
5368 Some(DOCUMENT_KEY.to_string())
5369 );
5370
5371 let vector_index_config = defaults.config.clone();
5373 assert!(vector_index_config.spann.is_some());
5374 assert!(vector_index_config.hnsw.is_none());
5375
5376 assert_eq!(
5378 embedding_vector_index.config.embedding_function,
5379 Some(EmbeddingFunctionConfiguration::Known(
5380 EmbeddingFunctionNewConfiguration {
5381 name: "default".to_string(),
5382 config: serde_json::json!({}),
5383 },
5384 ))
5385 );
5386 assert_eq!(
5387 defaults.config.embedding_function,
5388 Some(EmbeddingFunctionConfiguration::Known(
5389 EmbeddingFunctionNewConfiguration {
5390 name: "default".to_string(),
5391 config: serde_json::json!({}),
5392 },
5393 ))
5394 );
5395 }
5396
5397 #[test]
5398 fn test_reconcile_with_collection_config_both_non_default() {
5399 let mut schema = Schema::new_default(KnnIndex::Hnsw);
5401 schema.defaults.string = Some(StringValueType {
5402 fts_index: Some(FtsIndexType {
5403 enabled: true,
5404 config: FtsIndexConfig::default(),
5405 }),
5406 string_inverted_index: None,
5407 });
5408
5409 let mut collection_config = InternalCollectionConfiguration::default_hnsw();
5410 if let VectorIndexConfiguration::Hnsw(ref mut hnsw_config) = collection_config.vector_index
5412 {
5413 hnsw_config.ef_construction = 500; }
5415
5416 let result = Schema::reconcile_schema_and_config(
5418 Some(&schema),
5419 Some(&collection_config),
5420 KnnIndex::Spann,
5421 );
5422 assert!(result.is_err());
5423 assert!(matches!(
5424 result.unwrap_err(),
5425 SchemaError::ConfigAndSchemaConflict
5426 ));
5427 }
5428
5429 #[test]
5430 fn test_reconcile_with_collection_config_hnsw_override() {
5431 let schema = Schema::new_default(KnnIndex::Hnsw); let collection_config = InternalCollectionConfiguration {
5435 vector_index: VectorIndexConfiguration::Hnsw(InternalHnswConfiguration {
5436 ef_construction: 300,
5437 max_neighbors: 32,
5438 ef_search: 50,
5439 num_threads: 8,
5440 batch_size: 200,
5441 sync_threshold: 2000,
5442 resize_factor: 1.5,
5443 space: Space::L2,
5444 }),
5445 embedding_function: Some(EmbeddingFunctionConfiguration::Legacy),
5446 };
5447
5448 let result =
5449 Schema::reconcile_with_collection_config(&schema, &collection_config, KnnIndex::Hnsw)
5450 .unwrap();
5451
5452 let embedding_override = result.keys.get(EMBEDDING_KEY).unwrap();
5454 let vector_index = embedding_override
5455 .float_list
5456 .as_ref()
5457 .unwrap()
5458 .vector_index
5459 .as_ref()
5460 .unwrap();
5461
5462 assert!(vector_index.enabled);
5463 assert_eq!(vector_index.config.space, Some(Space::L2));
5464 assert_eq!(
5465 vector_index.config.embedding_function,
5466 Some(EmbeddingFunctionConfiguration::Legacy)
5467 );
5468 assert_eq!(
5469 vector_index.config.source_key,
5470 Some(DOCUMENT_KEY.to_string())
5471 );
5472
5473 let hnsw_config = vector_index.config.hnsw.as_ref().unwrap();
5474 assert_eq!(hnsw_config.ef_construction, Some(300));
5475 assert_eq!(hnsw_config.max_neighbors, Some(32));
5476 assert_eq!(hnsw_config.ef_search, Some(50));
5477 assert_eq!(hnsw_config.num_threads, Some(8));
5478 assert_eq!(hnsw_config.batch_size, Some(200));
5479 assert_eq!(hnsw_config.sync_threshold, Some(2000));
5480 assert_eq!(hnsw_config.resize_factor, Some(1.5));
5481
5482 assert!(vector_index.config.spann.is_none());
5483 }
5484
5485 #[test]
5486 fn test_reconcile_with_collection_config_spann_override() {
5487 let schema = Schema::new_default(KnnIndex::Spann); let collection_config = InternalCollectionConfiguration {
5491 vector_index: VectorIndexConfiguration::Spann(InternalSpannConfiguration {
5492 search_nprobe: 20,
5493 search_rng_factor: 3.0,
5494 search_rng_epsilon: 0.2,
5495 nreplica_count: 5,
5496 write_rng_factor: 2.0,
5497 write_rng_epsilon: 0.1,
5498 split_threshold: 2000,
5499 num_samples_kmeans: 200,
5500 initial_lambda: 0.8,
5501 reassign_neighbor_count: 100,
5502 merge_threshold: 800,
5503 num_centers_to_merge_to: 20,
5504 write_nprobe: 10,
5505 ef_construction: 400,
5506 ef_search: 60,
5507 max_neighbors: 24,
5508 space: Space::Cosine,
5509 }),
5510 embedding_function: None,
5511 };
5512
5513 let result =
5514 Schema::reconcile_with_collection_config(&schema, &collection_config, KnnIndex::Spann)
5515 .unwrap();
5516
5517 let embedding_override = result.keys.get(EMBEDDING_KEY).unwrap();
5519 let vector_index = embedding_override
5520 .float_list
5521 .as_ref()
5522 .unwrap()
5523 .vector_index
5524 .as_ref()
5525 .unwrap();
5526
5527 assert!(vector_index.enabled);
5528 assert_eq!(vector_index.config.space, Some(Space::Cosine));
5529 assert_eq!(vector_index.config.embedding_function, None);
5530 assert_eq!(
5531 vector_index.config.source_key,
5532 Some(DOCUMENT_KEY.to_string())
5533 );
5534
5535 assert!(vector_index.config.hnsw.is_none());
5536
5537 let spann_config = vector_index.config.spann.as_ref().unwrap();
5538 assert_eq!(spann_config.search_nprobe, Some(20));
5539 assert_eq!(spann_config.search_rng_factor, Some(3.0));
5540 assert_eq!(spann_config.search_rng_epsilon, Some(0.2));
5541 assert_eq!(spann_config.nreplica_count, Some(5));
5542 assert_eq!(spann_config.write_rng_factor, Some(2.0));
5543 assert_eq!(spann_config.write_rng_epsilon, Some(0.1));
5544 assert_eq!(spann_config.split_threshold, Some(2000));
5545 assert_eq!(spann_config.num_samples_kmeans, Some(200));
5546 assert_eq!(spann_config.initial_lambda, Some(0.8));
5547 assert_eq!(spann_config.reassign_neighbor_count, Some(100));
5548 assert_eq!(spann_config.merge_threshold, Some(800));
5549 assert_eq!(spann_config.num_centers_to_merge_to, Some(20));
5550 assert_eq!(spann_config.write_nprobe, Some(10));
5551 assert_eq!(spann_config.ef_construction, Some(400));
5552 assert_eq!(spann_config.ef_search, Some(60));
5553 assert_eq!(spann_config.max_neighbors, Some(24));
5554 }
5555
5556 #[test]
5557 fn test_reconcile_with_collection_config_updates_both_defaults_and_embedding() {
5558 let schema = Schema::new_default(KnnIndex::Hnsw);
5561
5562 let collection_config = InternalCollectionConfiguration {
5563 vector_index: VectorIndexConfiguration::Hnsw(InternalHnswConfiguration {
5564 ef_construction: 300,
5565 max_neighbors: 32,
5566 ef_search: 50,
5567 num_threads: 8,
5568 batch_size: 200,
5569 sync_threshold: 2000,
5570 resize_factor: 1.5,
5571 space: Space::L2,
5572 }),
5573 embedding_function: Some(EmbeddingFunctionConfiguration::Legacy),
5574 };
5575
5576 let result =
5577 Schema::reconcile_with_collection_config(&schema, &collection_config, KnnIndex::Hnsw)
5578 .unwrap();
5579
5580 let defaults_vector_index = result
5582 .defaults
5583 .float_list
5584 .as_ref()
5585 .unwrap()
5586 .vector_index
5587 .as_ref()
5588 .unwrap();
5589
5590 assert!(!defaults_vector_index.enabled);
5592 assert_eq!(defaults_vector_index.config.space, Some(Space::L2));
5594 assert_eq!(
5595 defaults_vector_index.config.embedding_function,
5596 Some(EmbeddingFunctionConfiguration::Legacy)
5597 );
5598 assert_eq!(defaults_vector_index.config.source_key, None);
5599 let defaults_hnsw = defaults_vector_index.config.hnsw.as_ref().unwrap();
5600 assert_eq!(defaults_hnsw.ef_construction, Some(300));
5601 assert_eq!(defaults_hnsw.max_neighbors, Some(32));
5602
5603 let embedding_override = result.keys.get(EMBEDDING_KEY).unwrap();
5605 let embedding_vector_index = embedding_override
5606 .float_list
5607 .as_ref()
5608 .unwrap()
5609 .vector_index
5610 .as_ref()
5611 .unwrap();
5612
5613 assert!(embedding_vector_index.enabled);
5615 assert_eq!(embedding_vector_index.config.space, Some(Space::L2));
5617 assert_eq!(
5618 embedding_vector_index.config.embedding_function,
5619 Some(EmbeddingFunctionConfiguration::Legacy)
5620 );
5621 assert_eq!(
5622 embedding_vector_index.config.source_key,
5623 Some(DOCUMENT_KEY.to_string())
5624 );
5625 let embedding_hnsw = embedding_vector_index.config.hnsw.as_ref().unwrap();
5626 assert_eq!(embedding_hnsw.ef_construction, Some(300));
5627 assert_eq!(embedding_hnsw.max_neighbors, Some(32));
5628 }
5629
5630 #[test]
5631 fn test_is_schema_default() {
5632 let default_hnsw_schema = Schema::new_default(KnnIndex::Hnsw);
5634 assert!(default_hnsw_schema.is_default());
5635
5636 let default_spann_schema = Schema::new_default(KnnIndex::Spann);
5637 assert!(default_spann_schema.is_default());
5638
5639 let mut modified_schema = Schema::new_default(KnnIndex::Hnsw);
5641 if let Some(ref mut string_type) = modified_schema.defaults.string {
5643 if let Some(ref mut string_inverted) = string_type.string_inverted_index {
5644 string_inverted.enabled = false; }
5646 }
5647 assert!(!modified_schema.is_default());
5648
5649 let mut schema_with_extra_overrides = Schema::new_default(KnnIndex::Hnsw);
5651 schema_with_extra_overrides
5652 .keys
5653 .insert("custom_key".to_string(), ValueTypes::default());
5654 assert!(!schema_with_extra_overrides.is_default());
5655 }
5656
5657 #[test]
5658 fn test_is_schema_default_with_space() {
5659 let schema = Schema::new_default(KnnIndex::Hnsw);
5660 assert!(schema.is_default());
5661
5662 let mut schema_with_space = Schema::new_default(KnnIndex::Hnsw);
5663 if let Some(ref mut float_list) = schema_with_space.defaults.float_list {
5664 if let Some(ref mut vector_index) = float_list.vector_index {
5665 vector_index.config.space = Some(Space::Cosine);
5666 }
5667 }
5668 assert!(!schema_with_space.is_default());
5669
5670 let mut schema_with_space_in_embedding_key = Schema::new_default(KnnIndex::Spann);
5671 if let Some(ref mut embedding_key) = schema_with_space_in_embedding_key
5672 .keys
5673 .get_mut(EMBEDDING_KEY)
5674 {
5675 if let Some(ref mut float_list) = embedding_key.float_list {
5676 if let Some(ref mut vector_index) = float_list.vector_index {
5677 vector_index.config.space = Some(Space::Cosine);
5678 }
5679 }
5680 }
5681 assert!(!schema_with_space_in_embedding_key.is_default());
5682 }
5683
5684 #[test]
5685 fn test_is_schema_default_with_embedding_function() {
5686 let schema = Schema::new_default(KnnIndex::Hnsw);
5687 assert!(schema.is_default());
5688
5689 let mut schema_with_embedding_function = Schema::new_default(KnnIndex::Hnsw);
5690 if let Some(ref mut float_list) = schema_with_embedding_function.defaults.float_list {
5691 if let Some(ref mut vector_index) = float_list.vector_index {
5692 vector_index.config.embedding_function =
5693 Some(EmbeddingFunctionConfiguration::Legacy);
5694 }
5695 }
5696 assert!(!schema_with_embedding_function.is_default());
5697
5698 let mut schema_with_embedding_function_in_embedding_key =
5699 Schema::new_default(KnnIndex::Spann);
5700 if let Some(ref mut embedding_key) = schema_with_embedding_function_in_embedding_key
5701 .keys
5702 .get_mut(EMBEDDING_KEY)
5703 {
5704 if let Some(ref mut float_list) = embedding_key.float_list {
5705 if let Some(ref mut vector_index) = float_list.vector_index {
5706 vector_index.config.embedding_function =
5707 Some(EmbeddingFunctionConfiguration::Legacy);
5708 }
5709 }
5710 }
5711 assert!(!schema_with_embedding_function_in_embedding_key.is_default());
5712 }
5713
5714 #[test]
5715 fn test_add_merges_keys_by_value_type() {
5716 let mut schema_a = Schema::new_default(KnnIndex::Hnsw);
5717 let mut schema_b = Schema::new_default(KnnIndex::Hnsw);
5718
5719 let string_override = ValueTypes {
5720 string: Some(StringValueType {
5721 string_inverted_index: Some(StringInvertedIndexType {
5722 enabled: true,
5723 config: StringInvertedIndexConfig {},
5724 }),
5725 fts_index: None,
5726 }),
5727 ..Default::default()
5728 };
5729 schema_a
5730 .keys
5731 .insert("custom_field".to_string(), string_override);
5732
5733 let float_override = ValueTypes {
5734 float: Some(FloatValueType {
5735 float_inverted_index: Some(FloatInvertedIndexType {
5736 enabled: true,
5737 config: FloatInvertedIndexConfig {},
5738 }),
5739 }),
5740 ..Default::default()
5741 };
5742 schema_b
5743 .keys
5744 .insert("custom_field".to_string(), float_override);
5745
5746 let merged = schema_a.merge(&schema_b).unwrap();
5747 let merged_override = merged.keys.get("custom_field").unwrap();
5748
5749 assert!(merged_override.string.is_some());
5750 assert!(merged_override.float.is_some());
5751 assert!(
5752 merged_override
5753 .string
5754 .as_ref()
5755 .unwrap()
5756 .string_inverted_index
5757 .as_ref()
5758 .unwrap()
5759 .enabled
5760 );
5761 assert!(
5762 merged_override
5763 .float
5764 .as_ref()
5765 .unwrap()
5766 .float_inverted_index
5767 .as_ref()
5768 .unwrap()
5769 .enabled
5770 );
5771 }
5772
5773 #[test]
5774 fn test_add_rejects_different_defaults() {
5775 let schema_a = Schema::new_default(KnnIndex::Hnsw);
5776 let mut schema_b = Schema::new_default(KnnIndex::Hnsw);
5777
5778 if let Some(string_type) = schema_b.defaults.string.as_mut() {
5779 if let Some(string_index) = string_type.string_inverted_index.as_mut() {
5780 string_index.enabled = false;
5781 }
5782 }
5783
5784 let err = schema_a.merge(&schema_b).unwrap_err();
5785 assert!(matches!(err, SchemaError::DefaultsMismatch));
5786 }
5787
5788 #[test]
5789 fn test_add_detects_conflicting_value_type_configuration() {
5790 let mut schema_a = Schema::new_default(KnnIndex::Hnsw);
5791 let mut schema_b = Schema::new_default(KnnIndex::Hnsw);
5792
5793 let string_override_enabled = ValueTypes {
5794 string: Some(StringValueType {
5795 string_inverted_index: Some(StringInvertedIndexType {
5796 enabled: true,
5797 config: StringInvertedIndexConfig {},
5798 }),
5799 fts_index: None,
5800 }),
5801 ..Default::default()
5802 };
5803 schema_a
5804 .keys
5805 .insert("custom_field".to_string(), string_override_enabled);
5806
5807 let string_override_disabled = ValueTypes {
5808 string: Some(StringValueType {
5809 string_inverted_index: Some(StringInvertedIndexType {
5810 enabled: false,
5811 config: StringInvertedIndexConfig {},
5812 }),
5813 fts_index: None,
5814 }),
5815 ..Default::default()
5816 };
5817 schema_b
5818 .keys
5819 .insert("custom_field".to_string(), string_override_disabled);
5820
5821 let err = schema_a.merge(&schema_b).unwrap_err();
5822 assert!(matches!(err, SchemaError::ConfigurationConflict { .. }));
5823 }
5824
5825 #[test]
5827 fn test_backward_compatibility_aliases() {
5828 let old_format_json = r###"{
5830 "defaults": {
5831 "#string": {
5832 "$fts_index": {
5833 "enabled": true,
5834 "config": {}
5835 }
5836 },
5837 "#int": {
5838 "$int_inverted_index": {
5839 "enabled": true,
5840 "config": {}
5841 }
5842 },
5843 "#float_list": {
5844 "$vector_index": {
5845 "enabled": true,
5846 "config": {
5847 "spann": {
5848 "search_nprobe": 10
5849 }
5850 }
5851 }
5852 }
5853 },
5854 "key_overrides": {
5855 "#document": {
5856 "#string": {
5857 "$fts_index": {
5858 "enabled": false,
5859 "config": {}
5860 }
5861 }
5862 }
5863 }
5864 }"###;
5865
5866 let schema_from_old: Schema = serde_json::from_str(old_format_json).unwrap();
5867
5868 let new_format_json = r###"{
5870 "defaults": {
5871 "string": {
5872 "fts_index": {
5873 "enabled": true,
5874 "config": {}
5875 }
5876 },
5877 "int": {
5878 "int_inverted_index": {
5879 "enabled": true,
5880 "config": {}
5881 }
5882 },
5883 "float_list": {
5884 "vector_index": {
5885 "enabled": true,
5886 "config": {
5887 "spann": {
5888 "search_nprobe": 10
5889 }
5890 }
5891 }
5892 }
5893 },
5894 "keys": {
5895 "#document": {
5896 "string": {
5897 "fts_index": {
5898 "enabled": false,
5899 "config": {}
5900 }
5901 }
5902 }
5903 }
5904 }"###;
5905
5906 let schema_from_new: Schema = serde_json::from_str(new_format_json).unwrap();
5907
5908 assert_eq!(schema_from_old, schema_from_new);
5910
5911 assert!(schema_from_old.defaults.string.is_some());
5913 assert!(schema_from_old
5914 .defaults
5915 .string
5916 .as_ref()
5917 .unwrap()
5918 .fts_index
5919 .is_some());
5920 assert!(
5921 schema_from_old
5922 .defaults
5923 .string
5924 .as_ref()
5925 .unwrap()
5926 .fts_index
5927 .as_ref()
5928 .unwrap()
5929 .enabled
5930 );
5931
5932 assert!(schema_from_old.defaults.int.is_some());
5933 assert!(schema_from_old
5934 .defaults
5935 .int
5936 .as_ref()
5937 .unwrap()
5938 .int_inverted_index
5939 .is_some());
5940
5941 assert!(schema_from_old.defaults.float_list.is_some());
5942 assert!(schema_from_old
5943 .defaults
5944 .float_list
5945 .as_ref()
5946 .unwrap()
5947 .vector_index
5948 .is_some());
5949
5950 assert!(schema_from_old.keys.contains_key(DOCUMENT_KEY));
5951 let doc_override = schema_from_old.keys.get(DOCUMENT_KEY).unwrap();
5952 assert!(doc_override.string.is_some());
5953 assert!(
5954 !doc_override
5955 .string
5956 .as_ref()
5957 .unwrap()
5958 .fts_index
5959 .as_ref()
5960 .unwrap()
5961 .enabled
5962 );
5963
5964 let serialized = serde_json::to_string(&schema_from_old).unwrap();
5966
5967 assert!(serialized.contains(r#""keys":"#));
5969 assert!(serialized.contains(r#""string":"#));
5970 assert!(serialized.contains(r#""fts_index":"#));
5971 assert!(serialized.contains(r#""int_inverted_index":"#));
5972 assert!(serialized.contains(r#""vector_index":"#));
5973
5974 assert!(!serialized.contains(r#""key_overrides":"#));
5976 assert!(!serialized.contains(r###""#string":"###));
5977 assert!(!serialized.contains(r###""$fts_index":"###));
5978 assert!(!serialized.contains(r###""$int_inverted_index":"###));
5979 assert!(!serialized.contains(r###""$vector_index":"###));
5980 }
5981
5982 #[test]
5983 fn test_hnsw_index_config_validation() {
5984 use validator::Validate;
5985
5986 let valid_config = HnswIndexConfig {
5988 batch_size: Some(10),
5989 sync_threshold: Some(100),
5990 ef_construction: Some(100),
5991 max_neighbors: Some(16),
5992 ..Default::default()
5993 };
5994 assert!(valid_config.validate().is_ok());
5995
5996 let invalid_batch_size = HnswIndexConfig {
5998 batch_size: Some(1),
5999 ..Default::default()
6000 };
6001 assert!(invalid_batch_size.validate().is_err());
6002
6003 let invalid_sync_threshold = HnswIndexConfig {
6005 sync_threshold: Some(1),
6006 ..Default::default()
6007 };
6008 assert!(invalid_sync_threshold.validate().is_err());
6009
6010 let boundary_config = HnswIndexConfig {
6012 batch_size: Some(2),
6013 sync_threshold: Some(2),
6014 ..Default::default()
6015 };
6016 assert!(boundary_config.validate().is_ok());
6017
6018 let all_none_config = HnswIndexConfig {
6020 ..Default::default()
6021 };
6022 assert!(all_none_config.validate().is_ok());
6023
6024 let other_fields_config = HnswIndexConfig {
6026 ef_construction: Some(1),
6027 max_neighbors: Some(1),
6028 ef_search: Some(1),
6029 num_threads: Some(1),
6030 resize_factor: Some(0.1),
6031 ..Default::default()
6032 };
6033 assert!(other_fields_config.validate().is_ok());
6034 }
6035
6036 #[test]
6037 fn test_spann_index_config_validation() {
6038 use validator::Validate;
6039
6040 let valid_config = SpannIndexConfig {
6042 write_nprobe: Some(32),
6043 nreplica_count: Some(4),
6044 split_threshold: Some(100),
6045 merge_threshold: Some(50),
6046 reassign_neighbor_count: Some(32),
6047 num_centers_to_merge_to: Some(4),
6048 ef_construction: Some(100),
6049 ef_search: Some(100),
6050 max_neighbors: Some(32),
6051 search_rng_factor: Some(1.0),
6052 write_rng_factor: Some(1.0),
6053 search_rng_epsilon: Some(7.5),
6054 write_rng_epsilon: Some(7.5),
6055 ..Default::default()
6056 };
6057 assert!(valid_config.validate().is_ok());
6058
6059 let invalid_write_nprobe = SpannIndexConfig {
6061 write_nprobe: Some(200),
6062 ..Default::default()
6063 };
6064 assert!(invalid_write_nprobe.validate().is_err());
6065
6066 let invalid_split_threshold = SpannIndexConfig {
6068 split_threshold: Some(10),
6069 ..Default::default()
6070 };
6071 assert!(invalid_split_threshold.validate().is_err());
6072
6073 let invalid_split_threshold_high = SpannIndexConfig {
6075 split_threshold: Some(250),
6076 ..Default::default()
6077 };
6078 assert!(invalid_split_threshold_high.validate().is_err());
6079
6080 let invalid_nreplica = SpannIndexConfig {
6082 nreplica_count: Some(10),
6083 ..Default::default()
6084 };
6085 assert!(invalid_nreplica.validate().is_err());
6086
6087 let invalid_reassign = SpannIndexConfig {
6089 reassign_neighbor_count: Some(100),
6090 ..Default::default()
6091 };
6092 assert!(invalid_reassign.validate().is_err());
6093
6094 let invalid_merge_threshold_low = SpannIndexConfig {
6096 merge_threshold: Some(5),
6097 ..Default::default()
6098 };
6099 assert!(invalid_merge_threshold_low.validate().is_err());
6100
6101 let invalid_merge_threshold_high = SpannIndexConfig {
6102 merge_threshold: Some(150),
6103 ..Default::default()
6104 };
6105 assert!(invalid_merge_threshold_high.validate().is_err());
6106
6107 let invalid_num_centers = SpannIndexConfig {
6109 num_centers_to_merge_to: Some(10),
6110 ..Default::default()
6111 };
6112 assert!(invalid_num_centers.validate().is_err());
6113
6114 let invalid_ef_construction = SpannIndexConfig {
6116 ef_construction: Some(300),
6117 ..Default::default()
6118 };
6119 assert!(invalid_ef_construction.validate().is_err());
6120
6121 let invalid_ef_search = SpannIndexConfig {
6123 ef_search: Some(300),
6124 ..Default::default()
6125 };
6126 assert!(invalid_ef_search.validate().is_err());
6127
6128 let invalid_max_neighbors = SpannIndexConfig {
6130 max_neighbors: Some(100),
6131 ..Default::default()
6132 };
6133 assert!(invalid_max_neighbors.validate().is_err());
6134
6135 let invalid_search_nprobe = SpannIndexConfig {
6137 search_nprobe: Some(200),
6138 ..Default::default()
6139 };
6140 assert!(invalid_search_nprobe.validate().is_err());
6141
6142 let invalid_search_rng_factor_low = SpannIndexConfig {
6144 search_rng_factor: Some(0.9),
6145 ..Default::default()
6146 };
6147 assert!(invalid_search_rng_factor_low.validate().is_err());
6148
6149 let invalid_search_rng_factor_high = SpannIndexConfig {
6150 search_rng_factor: Some(1.1),
6151 ..Default::default()
6152 };
6153 assert!(invalid_search_rng_factor_high.validate().is_err());
6154
6155 let valid_search_rng_factor = SpannIndexConfig {
6157 search_rng_factor: Some(1.0),
6158 ..Default::default()
6159 };
6160 assert!(valid_search_rng_factor.validate().is_ok());
6161
6162 let invalid_search_rng_epsilon_low = SpannIndexConfig {
6164 search_rng_epsilon: Some(4.0),
6165 ..Default::default()
6166 };
6167 assert!(invalid_search_rng_epsilon_low.validate().is_err());
6168
6169 let invalid_search_rng_epsilon_high = SpannIndexConfig {
6170 search_rng_epsilon: Some(11.0),
6171 ..Default::default()
6172 };
6173 assert!(invalid_search_rng_epsilon_high.validate().is_err());
6174
6175 let valid_search_rng_epsilon = SpannIndexConfig {
6177 search_rng_epsilon: Some(7.5),
6178 ..Default::default()
6179 };
6180 assert!(valid_search_rng_epsilon.validate().is_ok());
6181
6182 let invalid_write_rng_factor_low = SpannIndexConfig {
6184 write_rng_factor: Some(0.9),
6185 ..Default::default()
6186 };
6187 assert!(invalid_write_rng_factor_low.validate().is_err());
6188
6189 let invalid_write_rng_factor_high = SpannIndexConfig {
6190 write_rng_factor: Some(1.1),
6191 ..Default::default()
6192 };
6193 assert!(invalid_write_rng_factor_high.validate().is_err());
6194
6195 let valid_write_rng_factor = SpannIndexConfig {
6197 write_rng_factor: Some(1.0),
6198 ..Default::default()
6199 };
6200 assert!(valid_write_rng_factor.validate().is_ok());
6201
6202 let invalid_write_rng_epsilon_low = SpannIndexConfig {
6204 write_rng_epsilon: Some(4.0),
6205 ..Default::default()
6206 };
6207 assert!(invalid_write_rng_epsilon_low.validate().is_err());
6208
6209 let invalid_write_rng_epsilon_high = SpannIndexConfig {
6210 write_rng_epsilon: Some(11.0),
6211 ..Default::default()
6212 };
6213 assert!(invalid_write_rng_epsilon_high.validate().is_err());
6214
6215 let valid_write_rng_epsilon = SpannIndexConfig {
6217 write_rng_epsilon: Some(7.5),
6218 ..Default::default()
6219 };
6220 assert!(valid_write_rng_epsilon.validate().is_ok());
6221
6222 let invalid_num_samples_kmeans = SpannIndexConfig {
6224 num_samples_kmeans: Some(1500),
6225 ..Default::default()
6226 };
6227 assert!(invalid_num_samples_kmeans.validate().is_err());
6228
6229 let valid_num_samples_kmeans = SpannIndexConfig {
6231 num_samples_kmeans: Some(500),
6232 ..Default::default()
6233 };
6234 assert!(valid_num_samples_kmeans.validate().is_ok());
6235
6236 let invalid_initial_lambda_high = SpannIndexConfig {
6238 initial_lambda: Some(150.0),
6239 ..Default::default()
6240 };
6241 assert!(invalid_initial_lambda_high.validate().is_err());
6242
6243 let invalid_initial_lambda_low = SpannIndexConfig {
6244 initial_lambda: Some(50.0),
6245 ..Default::default()
6246 };
6247 assert!(invalid_initial_lambda_low.validate().is_err());
6248
6249 let valid_initial_lambda = SpannIndexConfig {
6251 initial_lambda: Some(100.0),
6252 ..Default::default()
6253 };
6254 assert!(valid_initial_lambda.validate().is_ok());
6255
6256 let all_none_config = SpannIndexConfig {
6258 ..Default::default()
6259 };
6260 assert!(all_none_config.validate().is_ok());
6261 }
6262
6263 #[test]
6264 fn test_builder_pattern_crud_workflow() {
6265 let schema = Schema::new_default(KnnIndex::Hnsw)
6269 .create_index(
6270 None,
6271 IndexConfig::Vector(VectorIndexConfig {
6272 space: Some(Space::Cosine),
6273 embedding_function: None,
6274 source_key: None,
6275 hnsw: Some(HnswIndexConfig {
6276 ef_construction: Some(200),
6277 max_neighbors: Some(32),
6278 ef_search: Some(50),
6279 num_threads: None,
6280 batch_size: None,
6281 sync_threshold: None,
6282 resize_factor: None,
6283 }),
6284 spann: None,
6285 }),
6286 )
6287 .expect("vector config should succeed")
6288 .create_index(
6289 Some("category"),
6290 IndexConfig::StringInverted(StringInvertedIndexConfig {}),
6291 )
6292 .expect("string inverted on key should succeed")
6293 .create_index(
6294 Some("year"),
6295 IndexConfig::IntInverted(IntInvertedIndexConfig {}),
6296 )
6297 .expect("int inverted on key should succeed")
6298 .create_index(
6299 Some("rating"),
6300 IndexConfig::FloatInverted(FloatInvertedIndexConfig {}),
6301 )
6302 .expect("float inverted on key should succeed")
6303 .create_index(
6304 Some("is_active"),
6305 IndexConfig::BoolInverted(BoolInvertedIndexConfig {}),
6306 )
6307 .expect("bool inverted on key should succeed");
6308
6309 assert!(schema.keys.contains_key(EMBEDDING_KEY));
6312 let embedding = schema.keys.get(EMBEDDING_KEY).unwrap();
6313 assert!(embedding.float_list.is_some());
6314 let vector_index = embedding
6315 .float_list
6316 .as_ref()
6317 .unwrap()
6318 .vector_index
6319 .as_ref()
6320 .unwrap();
6321 assert!(vector_index.enabled);
6322 assert_eq!(vector_index.config.space, Some(Space::Cosine));
6323 assert_eq!(
6324 vector_index.config.hnsw.as_ref().unwrap().ef_construction,
6325 Some(200)
6326 );
6327
6328 assert!(schema.keys.contains_key("category"));
6330 assert!(schema.keys.contains_key("year"));
6331 assert!(schema.keys.contains_key("rating"));
6332 assert!(schema.keys.contains_key("is_active"));
6333
6334 let category = schema.keys.get("category").unwrap();
6336 assert!(category.string.is_some());
6337 let string_idx = category
6338 .string
6339 .as_ref()
6340 .unwrap()
6341 .string_inverted_index
6342 .as_ref()
6343 .unwrap();
6344 assert!(string_idx.enabled);
6345
6346 let year = schema.keys.get("year").unwrap();
6348 assert!(year.int.is_some());
6349 let int_idx = year
6350 .int
6351 .as_ref()
6352 .unwrap()
6353 .int_inverted_index
6354 .as_ref()
6355 .unwrap();
6356 assert!(int_idx.enabled);
6357
6358 let schema = schema
6360 .delete_index(
6361 Some("category"),
6362 IndexConfig::StringInverted(StringInvertedIndexConfig {}),
6363 )
6364 .expect("delete string inverted should succeed")
6365 .delete_index(
6366 Some("year"),
6367 IndexConfig::IntInverted(IntInvertedIndexConfig {}),
6368 )
6369 .expect("delete int inverted should succeed");
6370
6371 let category = schema.keys.get("category").unwrap();
6373 let string_idx = category
6374 .string
6375 .as_ref()
6376 .unwrap()
6377 .string_inverted_index
6378 .as_ref()
6379 .unwrap();
6380 assert!(!string_idx.enabled); let year = schema.keys.get("year").unwrap();
6383 let int_idx = year
6384 .int
6385 .as_ref()
6386 .unwrap()
6387 .int_inverted_index
6388 .as_ref()
6389 .unwrap();
6390 assert!(!int_idx.enabled); let rating = schema.keys.get("rating").unwrap();
6394 let float_idx = rating
6395 .float
6396 .as_ref()
6397 .unwrap()
6398 .float_inverted_index
6399 .as_ref()
6400 .unwrap();
6401 assert!(float_idx.enabled); let is_active = schema.keys.get("is_active").unwrap();
6404 let bool_idx = is_active
6405 .boolean
6406 .as_ref()
6407 .unwrap()
6408 .bool_inverted_index
6409 .as_ref()
6410 .unwrap();
6411 assert!(bool_idx.enabled); }
6413
6414 #[test]
6415 fn test_builder_create_index_validation_errors() {
6416 let result = Schema::new_default(KnnIndex::Hnsw).create_index(
6423 Some("my_vectors"),
6424 IndexConfig::Vector(VectorIndexConfig {
6425 space: Some(Space::L2),
6426 embedding_function: None,
6427 source_key: None,
6428 hnsw: None,
6429 spann: None,
6430 }),
6431 );
6432 assert!(result.is_err());
6433 assert!(matches!(
6434 result.unwrap_err(),
6435 SchemaBuilderError::VectorIndexMustBeGlobal { key } if key == "my_vectors"
6436 ));
6437
6438 let result = Schema::new_default(KnnIndex::Hnsw)
6440 .create_index(Some("my_text"), IndexConfig::Fts(FtsIndexConfig::default()));
6441 assert!(result.is_err());
6442 assert!(matches!(
6443 result.unwrap_err(),
6444 SchemaBuilderError::FtsIndexOnlyOnDocument
6445 ));
6446
6447 let schema = Schema::new_default(KnnIndex::Hnsw)
6449 .create_index(
6450 Some(DOCUMENT_KEY),
6451 IndexConfig::Fts(FtsIndexConfig::default()),
6452 )
6453 .expect("FTS on #document should succeed");
6454 assert!(schema.is_fts_enabled());
6455
6456 let result = Schema::new_default(KnnIndex::Hnsw).create_index(
6458 Some(DOCUMENT_KEY),
6459 IndexConfig::StringInverted(StringInvertedIndexConfig {}),
6460 );
6461 assert!(result.is_err());
6462 assert!(matches!(
6463 result.unwrap_err(),
6464 SchemaBuilderError::SpecialKeyModificationNotAllowed { .. }
6465 ));
6466
6467 let result = Schema::new_default(KnnIndex::Hnsw).create_index(
6469 Some(EMBEDDING_KEY),
6470 IndexConfig::IntInverted(IntInvertedIndexConfig {}),
6471 );
6472 assert!(result.is_err());
6473 assert!(matches!(
6474 result.unwrap_err(),
6475 SchemaBuilderError::SpecialKeyModificationNotAllowed { .. }
6476 ));
6477
6478 let result = Schema::new_default(KnnIndex::Hnsw).create_index(
6480 None,
6481 IndexConfig::SparseVector(SparseVectorIndexConfig {
6482 embedding_function: None,
6483 source_key: None,
6484 bm25: None,
6485 algorithm: SparseIndexAlgorithm::Wand,
6486 }),
6487 );
6488 assert!(result.is_err());
6489 assert!(matches!(
6490 result.unwrap_err(),
6491 SchemaBuilderError::SparseVectorRequiresKey
6492 ));
6493
6494 let schema = Schema::new_default(KnnIndex::Hnsw)
6496 .create_index(
6497 Some("sparse1"),
6498 IndexConfig::SparseVector(SparseVectorIndexConfig {
6499 embedding_function: None,
6500 source_key: None,
6501 bm25: None,
6502 algorithm: SparseIndexAlgorithm::Wand,
6503 }),
6504 )
6505 .expect("first sparse should succeed")
6506 .create_index(
6507 Some("sparse2"),
6508 IndexConfig::SparseVector(SparseVectorIndexConfig {
6509 embedding_function: None,
6510 source_key: None,
6511 bm25: None,
6512 algorithm: SparseIndexAlgorithm::MaxScore,
6513 }),
6514 )
6515 .expect("second sparse should succeed");
6516 assert_eq!(
6517 schema.enabled_sparse_keys(),
6518 vec!["sparse1".to_string(), "sparse2".to_string()]
6519 );
6520 }
6521
6522 #[test]
6523 fn test_builder_delete_index_validation_errors() {
6524 let result = Schema::new_default(KnnIndex::Hnsw).delete_index(
6530 Some(EMBEDDING_KEY),
6531 IndexConfig::StringInverted(StringInvertedIndexConfig {}),
6532 );
6533 assert!(result.is_err());
6534 assert!(matches!(
6535 result.unwrap_err(),
6536 SchemaBuilderError::SpecialKeyModificationNotAllowed { .. }
6537 ));
6538
6539 let result = Schema::new_default(KnnIndex::Hnsw).delete_index(
6541 Some(DOCUMENT_KEY),
6542 IndexConfig::IntInverted(IntInvertedIndexConfig {}),
6543 );
6544 assert!(result.is_err());
6545 assert!(matches!(
6546 result.unwrap_err(),
6547 SchemaBuilderError::SpecialKeyModificationNotAllowed { .. }
6548 ));
6549
6550 let result = Schema::new_default(KnnIndex::Hnsw).delete_index(
6552 None,
6553 IndexConfig::Vector(VectorIndexConfig {
6554 space: None,
6555 embedding_function: None,
6556 source_key: None,
6557 hnsw: None,
6558 spann: None,
6559 }),
6560 );
6561 assert!(result.is_err());
6562 assert!(matches!(
6563 result.unwrap_err(),
6564 SchemaBuilderError::VectorIndexDeletionNotSupported
6565 ));
6566
6567 let schema = Schema::new_default(KnnIndex::Hnsw)
6569 .delete_index(
6570 Some(DOCUMENT_KEY),
6571 IndexConfig::Fts(FtsIndexConfig::default()),
6572 )
6573 .expect("FTS deletion should succeed");
6574 assert!(!schema.is_fts_enabled());
6575
6576 let result = Schema::new_default(KnnIndex::Hnsw)
6578 .create_index(
6579 Some("sparse"),
6580 IndexConfig::SparseVector(SparseVectorIndexConfig {
6581 embedding_function: None,
6582 source_key: None,
6583 bm25: None,
6584 algorithm: SparseIndexAlgorithm::Wand,
6585 }),
6586 )
6587 .expect("create should succeed")
6588 .delete_index(
6589 Some("sparse"),
6590 IndexConfig::SparseVector(SparseVectorIndexConfig {
6591 embedding_function: None,
6592 source_key: None,
6593 bm25: None,
6594 algorithm: SparseIndexAlgorithm::Wand,
6595 }),
6596 );
6597 assert!(result.is_err());
6598 assert!(matches!(
6599 result.unwrap_err(),
6600 SchemaBuilderError::SparseVectorIndexDeletionNotSupported
6601 ));
6602 }
6603
6604 #[test]
6605 fn test_fts_create_global_without_key_rejected() {
6606 let result = Schema::new_default(KnnIndex::Hnsw)
6608 .create_index(None, IndexConfig::Fts(FtsIndexConfig::default()));
6609 assert!(result.is_err());
6610 assert!(matches!(
6611 result.unwrap_err(),
6612 SchemaBuilderError::FtsIndexOnlyOnDocument
6613 ));
6614 }
6615
6616 #[test]
6617 fn test_fts_delete_global_without_key_rejected() {
6618 let result = Schema::new_default(KnnIndex::Hnsw)
6620 .delete_index(None, IndexConfig::Fts(FtsIndexConfig::default()));
6621 assert!(result.is_err());
6622 assert!(matches!(
6623 result.unwrap_err(),
6624 SchemaBuilderError::FtsIndexDeletionOnlyOnDocument
6625 ));
6626 }
6627
6628 #[test]
6629 fn test_fts_delete_on_custom_key_rejected() {
6630 let result = Schema::new_default(KnnIndex::Hnsw)
6632 .delete_index(Some("my_text"), IndexConfig::Fts(FtsIndexConfig::default()));
6633 assert!(result.is_err());
6634 assert!(matches!(
6635 result.unwrap_err(),
6636 SchemaBuilderError::FtsIndexDeletionOnlyOnDocument
6637 ));
6638 }
6639
6640 #[test]
6641 fn test_reserved_key_prefix_create_index() {
6642 let result = Schema::new_default(KnnIndex::Hnsw).create_index(
6644 Some("#custom_field"),
6645 IndexConfig::StringInverted(StringInvertedIndexConfig {}),
6646 );
6647 assert!(result.is_err());
6648 assert!(matches!(
6649 result.unwrap_err(),
6650 SchemaBuilderError::ReservedKeyPrefix { key } if key == "#custom_field"
6651 ));
6652 }
6653
6654 #[test]
6655 fn test_reserved_key_prefix_delete_index() {
6656 let result = Schema::new_default(KnnIndex::Hnsw).delete_index(
6658 Some("#custom_field"),
6659 IndexConfig::StringInverted(StringInvertedIndexConfig {}),
6660 );
6661 assert!(result.is_err());
6662 assert!(matches!(
6663 result.unwrap_err(),
6664 SchemaBuilderError::ReservedKeyPrefix { key } if key == "#custom_field"
6665 ));
6666 }
6667
6668 #[test]
6669 fn test_is_fts_enabled_backward_compatibility() {
6670 let schema = Schema::new_default(KnnIndex::Hnsw);
6672 assert!(schema.is_fts_enabled());
6673
6674 let empty_schema = Schema {
6676 defaults: ValueTypes::default(),
6677 keys: HashMap::new(),
6678 cmek: None,
6679 };
6680 assert!(empty_schema.is_fts_enabled());
6681 }
6682
6683 #[test]
6684 fn test_is_fts_enabled_after_disable() {
6685 let schema = Schema::new_default(KnnIndex::Hnsw)
6687 .delete_index(
6688 Some(DOCUMENT_KEY),
6689 IndexConfig::Fts(FtsIndexConfig::default()),
6690 )
6691 .expect("FTS deletion should succeed");
6692 assert!(!schema.is_fts_enabled());
6693 }
6694
6695 #[test]
6696 fn test_is_fts_enabled_after_reenable() {
6697 let schema = Schema::new_default(KnnIndex::Hnsw)
6699 .delete_index(
6700 Some(DOCUMENT_KEY),
6701 IndexConfig::Fts(FtsIndexConfig::default()),
6702 )
6703 .expect("FTS deletion should succeed")
6704 .create_index(
6705 Some(DOCUMENT_KEY),
6706 IndexConfig::Fts(FtsIndexConfig::default()),
6707 )
6708 .expect("FTS creation should succeed");
6709 assert!(schema.is_fts_enabled());
6710 }
6711
6712 #[test]
6713 fn test_fts_disabled_blocks_where_document_validation() {
6714 use crate::{DocumentExpression, DocumentOperator};
6715
6716 let schema = Schema::new_default(KnnIndex::Hnsw)
6718 .delete_index(
6719 Some(DOCUMENT_KEY),
6720 IndexConfig::Fts(FtsIndexConfig::default()),
6721 )
6722 .expect("FTS deletion should succeed");
6723
6724 let where_clause = Where::Document(DocumentExpression {
6726 operator: DocumentOperator::Contains,
6727 pattern: "test query".to_string(),
6728 });
6729 let result = schema.is_metadata_where_indexing_enabled(&where_clause);
6730 assert!(result.is_err());
6731 assert!(matches!(
6732 result.unwrap_err(),
6733 FilterValidationError::FtsDisabled
6734 ));
6735 }
6736
6737 #[test]
6738 fn test_fts_enabled_allows_where_document_validation() {
6739 use crate::{DocumentExpression, DocumentOperator};
6740
6741 let schema = Schema::new_default(KnnIndex::Hnsw);
6743
6744 let where_clause = Where::Document(DocumentExpression {
6746 operator: DocumentOperator::Contains,
6747 pattern: "test query".to_string(),
6748 });
6749 let result = schema.is_metadata_where_indexing_enabled(&where_clause);
6750 assert!(result.is_ok());
6751 }
6752
6753 #[test]
6754 fn test_builder_pattern_chaining() {
6755 let schema = Schema::new_default(KnnIndex::Hnsw)
6757 .create_index(Some("tag1"), StringInvertedIndexConfig {}.into())
6758 .unwrap()
6759 .create_index(Some("tag2"), StringInvertedIndexConfig {}.into())
6760 .unwrap()
6761 .create_index(Some("tag3"), StringInvertedIndexConfig {}.into())
6762 .unwrap()
6763 .create_index(Some("count"), IntInvertedIndexConfig {}.into())
6764 .unwrap()
6765 .delete_index(Some("tag2"), StringInvertedIndexConfig {}.into())
6766 .unwrap()
6767 .create_index(Some("score"), FloatInvertedIndexConfig {}.into())
6768 .unwrap();
6769
6770 assert!(
6772 schema
6773 .keys
6774 .get("tag1")
6775 .unwrap()
6776 .string
6777 .as_ref()
6778 .unwrap()
6779 .string_inverted_index
6780 .as_ref()
6781 .unwrap()
6782 .enabled
6783 );
6784
6785 assert!(
6787 !schema
6788 .keys
6789 .get("tag2")
6790 .unwrap()
6791 .string
6792 .as_ref()
6793 .unwrap()
6794 .string_inverted_index
6795 .as_ref()
6796 .unwrap()
6797 .enabled
6798 );
6799
6800 assert!(
6802 schema
6803 .keys
6804 .get("tag3")
6805 .unwrap()
6806 .string
6807 .as_ref()
6808 .unwrap()
6809 .string_inverted_index
6810 .as_ref()
6811 .unwrap()
6812 .enabled
6813 );
6814
6815 assert!(
6817 schema
6818 .keys
6819 .get("count")
6820 .unwrap()
6821 .int
6822 .as_ref()
6823 .unwrap()
6824 .int_inverted_index
6825 .as_ref()
6826 .unwrap()
6827 .enabled
6828 );
6829
6830 assert!(
6832 schema
6833 .keys
6834 .get("score")
6835 .unwrap()
6836 .float
6837 .as_ref()
6838 .unwrap()
6839 .float_inverted_index
6840 .as_ref()
6841 .unwrap()
6842 .enabled
6843 );
6844 }
6845
6846 #[test]
6847 fn test_schema_default_matches_python() {
6848 let schema = Schema::default();
6850
6851 assert!(schema.defaults.string.is_some());
6857 let string = schema.defaults.string.as_ref().unwrap();
6858 assert!(!string.fts_index.as_ref().unwrap().enabled);
6859 assert!(string.string_inverted_index.as_ref().unwrap().enabled);
6860
6861 assert!(schema.defaults.float_list.is_some());
6863 let float_list = schema.defaults.float_list.as_ref().unwrap();
6864 assert!(!float_list.vector_index.as_ref().unwrap().enabled);
6865 let vector_config = &float_list.vector_index.as_ref().unwrap().config;
6866 assert_eq!(vector_config.space, None); assert_eq!(vector_config.hnsw, None); assert_eq!(vector_config.spann, None); assert_eq!(vector_config.source_key, None);
6870
6871 assert!(schema.defaults.sparse_vector.is_some());
6873 let sparse = schema.defaults.sparse_vector.as_ref().unwrap();
6874 assert!(!sparse.sparse_vector_index.as_ref().unwrap().enabled);
6875
6876 assert!(schema.defaults.int.is_some());
6878 assert!(
6879 schema
6880 .defaults
6881 .int
6882 .as_ref()
6883 .unwrap()
6884 .int_inverted_index
6885 .as_ref()
6886 .unwrap()
6887 .enabled
6888 );
6889
6890 assert!(schema.defaults.float.is_some());
6892 assert!(
6893 schema
6894 .defaults
6895 .float
6896 .as_ref()
6897 .unwrap()
6898 .float_inverted_index
6899 .as_ref()
6900 .unwrap()
6901 .enabled
6902 );
6903
6904 assert!(schema.defaults.boolean.is_some());
6906 assert!(
6907 schema
6908 .defaults
6909 .boolean
6910 .as_ref()
6911 .unwrap()
6912 .bool_inverted_index
6913 .as_ref()
6914 .unwrap()
6915 .enabled
6916 );
6917
6918 assert!(schema.keys.contains_key(DOCUMENT_KEY));
6924 let doc = schema.keys.get(DOCUMENT_KEY).unwrap();
6925 assert!(doc.string.is_some());
6926 assert!(
6927 doc.string
6928 .as_ref()
6929 .unwrap()
6930 .fts_index
6931 .as_ref()
6932 .unwrap()
6933 .enabled
6934 );
6935 assert!(
6936 !doc.string
6937 .as_ref()
6938 .unwrap()
6939 .string_inverted_index
6940 .as_ref()
6941 .unwrap()
6942 .enabled
6943 );
6944
6945 assert!(schema.keys.contains_key(EMBEDDING_KEY));
6947 let embedding = schema.keys.get(EMBEDDING_KEY).unwrap();
6948 assert!(embedding.float_list.is_some());
6949 let vec_idx = embedding
6950 .float_list
6951 .as_ref()
6952 .unwrap()
6953 .vector_index
6954 .as_ref()
6955 .unwrap();
6956 assert!(vec_idx.enabled);
6957 assert_eq!(vec_idx.config.source_key, Some(DOCUMENT_KEY.to_string()));
6958 assert_eq!(vec_idx.config.space, None); assert_eq!(vec_idx.config.hnsw, None); assert_eq!(vec_idx.config.spann, None); assert_eq!(schema.keys.len(), 2);
6964 }
6965
6966 #[test]
6967 fn test_schema_default_works_with_builder() {
6968 let schema = Schema::default()
6970 .create_index(Some("category"), StringInvertedIndexConfig {}.into())
6971 .expect("should succeed");
6972
6973 assert!(schema.keys.contains_key("category"));
6975 assert!(schema.keys.contains_key(DOCUMENT_KEY));
6976 assert!(schema.keys.contains_key(EMBEDDING_KEY));
6977 assert_eq!(schema.keys.len(), 3);
6978 }
6979
6980 #[cfg(feature = "testing")]
6981 mod proptests {
6982 use super::*;
6983 use crate::strategies::{
6984 embedding_function_strategy, internal_collection_configuration_strategy,
6985 internal_hnsw_configuration_strategy, internal_spann_configuration_strategy,
6986 knn_index_strategy, space_strategy, TEST_NAME_PATTERN,
6987 };
6988 use crate::{
6989 HnswIndexConfig, SpannIndexConfig, VectorIndexConfig, DOCUMENT_KEY, EMBEDDING_KEY,
6990 };
6991 use proptest::prelude::*;
6992 use proptest::strategy::BoxedStrategy;
6993 use proptest::string::string_regex;
6994 use serde_json::json;
6995
6996 fn default_embedding_function_strategy(
6997 ) -> impl Strategy<Value = Option<EmbeddingFunctionConfiguration>> {
6998 proptest::option::of(prop_oneof![
6999 Just(EmbeddingFunctionConfiguration::Unknown),
7000 Just(EmbeddingFunctionConfiguration::Known(
7001 EmbeddingFunctionNewConfiguration {
7002 name: "default".to_string(),
7003 config: json!({ "alpha": 1 }),
7004 }
7005 )),
7006 ])
7007 }
7008
7009 fn sparse_embedding_function_strategy(
7010 ) -> impl Strategy<Value = Option<EmbeddingFunctionConfiguration>> {
7011 let known_strategy = string_regex(TEST_NAME_PATTERN).unwrap().prop_map(|name| {
7012 EmbeddingFunctionConfiguration::Known(EmbeddingFunctionNewConfiguration {
7013 name,
7014 config: json!({ "alpha": 1 }),
7015 })
7016 });
7017
7018 proptest::option::of(prop_oneof![
7019 Just(EmbeddingFunctionConfiguration::Unknown),
7020 known_strategy,
7021 ])
7022 }
7023
7024 fn non_default_internal_collection_configuration_strategy(
7025 ) -> impl Strategy<Value = InternalCollectionConfiguration> {
7026 internal_collection_configuration_strategy()
7027 .prop_filter("non-default configuration", |config| !config.is_default())
7028 }
7029
7030 fn partial_hnsw_index_config_strategy() -> impl Strategy<Value = HnswIndexConfig> {
7031 (
7032 proptest::option::of(1usize..=512),
7033 proptest::option::of(1usize..=128),
7034 proptest::option::of(1usize..=512),
7035 proptest::option::of(1usize..=64),
7036 proptest::option::of(2usize..=4096),
7037 proptest::option::of(2usize..=4096),
7038 proptest::option::of(prop_oneof![
7039 Just(0.5f64),
7040 Just(1.0f64),
7041 Just(1.5f64),
7042 Just(2.0f64)
7043 ]),
7044 )
7045 .prop_map(
7046 |(
7047 ef_construction,
7048 max_neighbors,
7049 ef_search,
7050 num_threads,
7051 batch_size,
7052 sync_threshold,
7053 resize_factor,
7054 )| HnswIndexConfig {
7055 ef_construction,
7056 max_neighbors,
7057 ef_search,
7058 num_threads,
7059 batch_size,
7060 sync_threshold,
7061 resize_factor,
7062 },
7063 )
7064 }
7065
7066 fn partial_spann_index_config_strategy() -> impl Strategy<Value = SpannIndexConfig> {
7067 let epsilon_strategy = prop_oneof![Just(5.0f32), Just(7.5f32), Just(10.0f32)];
7068 (
7069 (
7070 proptest::option::of(1u32..=128), proptest::option::of(Just(1.0f32)), proptest::option::of(epsilon_strategy.clone()), proptest::option::of(1u32..=8), proptest::option::of(Just(1.0f32)), proptest::option::of(epsilon_strategy), proptest::option::of(50u32..=200), proptest::option::of(1usize..=1000), ),
7079 (
7080 proptest::option::of(Just(100.0f32)), proptest::option::of(1u32..=64), proptest::option::of(25u32..=100), proptest::option::of(1u32..=8), proptest::option::of(1u32..=64), proptest::option::of(1usize..=200), proptest::option::of(1usize..=200), proptest::option::of(1usize..=64), ),
7089 )
7090 .prop_map(
7091 |(
7092 (
7093 search_nprobe,
7094 search_rng_factor,
7095 search_rng_epsilon,
7096 nreplica_count,
7097 write_rng_factor,
7098 write_rng_epsilon,
7099 split_threshold,
7100 num_samples_kmeans,
7101 ),
7102 (
7103 initial_lambda,
7104 reassign_neighbor_count,
7105 merge_threshold,
7106 num_centers_to_merge_to,
7107 write_nprobe,
7108 ef_construction,
7109 ef_search,
7110 max_neighbors,
7111 ),
7112 )| SpannIndexConfig {
7113 search_nprobe,
7114 search_rng_factor,
7115 search_rng_epsilon,
7116 nreplica_count,
7117 write_rng_factor,
7118 write_rng_epsilon,
7119 split_threshold,
7120 num_samples_kmeans,
7121 initial_lambda,
7122 reassign_neighbor_count,
7123 merge_threshold,
7124 num_centers_to_merge_to,
7125 write_nprobe,
7126 ef_construction,
7127 ef_search,
7128 max_neighbors,
7129 center_drift_threshold: None,
7130 quantize: Quantization::None,
7131 },
7132 )
7133 }
7134
7135 proptest! {
7136 #[test]
7137 fn merge_hnsw_configs_preserves_user_overrides(
7138 base in partial_hnsw_index_config_strategy(),
7139 user in partial_hnsw_index_config_strategy(),
7140 ) {
7141 let merged = Schema::merge_hnsw_configs(Some(&base), Some(&user))
7142 .expect("merge should return Some when both are Some");
7143
7144 if user.ef_construction.is_some() {
7146 prop_assert_eq!(merged.ef_construction, user.ef_construction);
7147 }
7148 if user.max_neighbors.is_some() {
7149 prop_assert_eq!(merged.max_neighbors, user.max_neighbors);
7150 }
7151 if user.ef_search.is_some() {
7152 prop_assert_eq!(merged.ef_search, user.ef_search);
7153 }
7154 if user.num_threads.is_some() {
7155 prop_assert_eq!(merged.num_threads, user.num_threads);
7156 }
7157 if user.batch_size.is_some() {
7158 prop_assert_eq!(merged.batch_size, user.batch_size);
7159 }
7160 if user.sync_threshold.is_some() {
7161 prop_assert_eq!(merged.sync_threshold, user.sync_threshold);
7162 }
7163 if user.resize_factor.is_some() {
7164 prop_assert_eq!(merged.resize_factor, user.resize_factor);
7165 }
7166 }
7167
7168 #[test]
7169 fn merge_hnsw_configs_falls_back_to_base_when_user_is_none(
7170 base in partial_hnsw_index_config_strategy(),
7171 ) {
7172 let merged = Schema::merge_hnsw_configs(Some(&base), None)
7173 .expect("merge should return Some when base is Some");
7174
7175 prop_assert_eq!(merged, base);
7177 }
7178
7179 #[test]
7180 fn merge_hnsw_configs_returns_user_when_base_is_none(
7181 user in partial_hnsw_index_config_strategy(),
7182 ) {
7183 let merged = Schema::merge_hnsw_configs(None, Some(&user))
7184 .expect("merge should return Some when user is Some");
7185
7186 prop_assert_eq!(merged, user);
7188 }
7189
7190 #[test]
7191 fn merge_spann_configs_preserves_user_overrides(
7192 base in partial_spann_index_config_strategy(),
7193 user in partial_spann_index_config_strategy(),
7194 ) {
7195 let merged = Schema::merge_spann_configs(Some(&base), Some(&user))
7196 .expect("merge should return Ok")
7197 .expect("merge should return Some when both are Some");
7198
7199 if user.search_nprobe.is_some() {
7201 prop_assert_eq!(merged.search_nprobe, user.search_nprobe);
7202 }
7203 if user.search_rng_epsilon.is_some() {
7204 prop_assert_eq!(merged.search_rng_epsilon, user.search_rng_epsilon);
7205 }
7206 if user.split_threshold.is_some() {
7207 prop_assert_eq!(merged.split_threshold, user.split_threshold);
7208 }
7209 if user.ef_construction.is_some() {
7210 prop_assert_eq!(merged.ef_construction, user.ef_construction);
7211 }
7212 if user.ef_search.is_some() {
7213 prop_assert_eq!(merged.ef_search, user.ef_search);
7214 }
7215 if user.max_neighbors.is_some() {
7216 prop_assert_eq!(merged.max_neighbors, user.max_neighbors);
7217 }
7218 }
7219
7220 #[test]
7221 fn merge_spann_configs_falls_back_to_base_when_user_is_none(
7222 base in partial_spann_index_config_strategy(),
7223 ) {
7224 let merged = Schema::merge_spann_configs(Some(&base), None)
7225 .expect("merge should return Ok")
7226 .expect("merge should return Some when base is Some");
7227
7228 prop_assert_eq!(merged, base);
7230 }
7231
7232 #[test]
7233 fn merge_vector_index_config_preserves_user_overrides(
7234 base in vector_index_config_strategy(),
7235 user in vector_index_config_strategy(),
7236 knn in knn_index_strategy(),
7237 ) {
7238 let merged = Schema::merge_vector_index_config(&base, &user, knn)
7239 .expect("merge should succeed");
7240
7241 if user.space.is_some() {
7243 prop_assert_eq!(merged.space, user.space);
7244 }
7245 if user.embedding_function.is_some() {
7246 prop_assert_eq!(merged.embedding_function, user.embedding_function);
7247 }
7248 if user.source_key.is_some() {
7249 prop_assert_eq!(merged.source_key, user.source_key);
7250 }
7251
7252 match knn {
7254 KnnIndex::Hnsw => {
7255 if let (Some(_base_hnsw), Some(user_hnsw)) = (&base.hnsw, &user.hnsw) {
7256 let merged_hnsw = merged.hnsw.as_ref().expect("hnsw should be Some");
7257 if user_hnsw.ef_construction.is_some() {
7258 prop_assert_eq!(merged_hnsw.ef_construction, user_hnsw.ef_construction);
7259 }
7260 }
7261 }
7262 KnnIndex::Spann => {
7263 if let (Some(_base_spann), Some(user_spann)) = (&base.spann, &user.spann) {
7264 let merged_spann = merged.spann.as_ref().expect("spann should be Some");
7265 if user_spann.search_nprobe.is_some() {
7266 prop_assert_eq!(merged_spann.search_nprobe, user_spann.search_nprobe);
7267 }
7268 }
7269 }
7270 }
7271 }
7272 }
7273
7274 fn expected_vector_index_config(
7275 config: &InternalCollectionConfiguration,
7276 ) -> VectorIndexConfig {
7277 match &config.vector_index {
7278 VectorIndexConfiguration::Hnsw(hnsw_config) => VectorIndexConfig {
7279 space: Some(hnsw_config.space.clone()),
7280 embedding_function: config.embedding_function.clone(),
7281 source_key: None,
7282 hnsw: Some(HnswIndexConfig {
7283 ef_construction: Some(hnsw_config.ef_construction),
7284 max_neighbors: Some(hnsw_config.max_neighbors),
7285 ef_search: Some(hnsw_config.ef_search),
7286 num_threads: Some(hnsw_config.num_threads),
7287 batch_size: Some(hnsw_config.batch_size),
7288 sync_threshold: Some(hnsw_config.sync_threshold),
7289 resize_factor: Some(hnsw_config.resize_factor),
7290 }),
7291 spann: None,
7292 },
7293 VectorIndexConfiguration::Spann(spann_config) => VectorIndexConfig {
7294 space: Some(spann_config.space.clone()),
7295 embedding_function: config.embedding_function.clone(),
7296 source_key: None,
7297 hnsw: None,
7298 spann: Some(SpannIndexConfig {
7299 search_nprobe: Some(spann_config.search_nprobe),
7300 search_rng_factor: Some(spann_config.search_rng_factor),
7301 search_rng_epsilon: Some(spann_config.search_rng_epsilon),
7302 nreplica_count: Some(spann_config.nreplica_count),
7303 write_rng_factor: Some(spann_config.write_rng_factor),
7304 write_rng_epsilon: Some(spann_config.write_rng_epsilon),
7305 split_threshold: Some(spann_config.split_threshold),
7306 num_samples_kmeans: Some(spann_config.num_samples_kmeans),
7307 initial_lambda: Some(spann_config.initial_lambda),
7308 reassign_neighbor_count: Some(spann_config.reassign_neighbor_count),
7309 merge_threshold: Some(spann_config.merge_threshold),
7310 num_centers_to_merge_to: Some(spann_config.num_centers_to_merge_to),
7311 write_nprobe: Some(spann_config.write_nprobe),
7312 ef_construction: Some(spann_config.ef_construction),
7313 ef_search: Some(spann_config.ef_search),
7314 max_neighbors: Some(spann_config.max_neighbors),
7315 center_drift_threshold: None,
7316 quantize: Quantization::None,
7317 }),
7318 },
7319 }
7320 }
7321
7322 fn non_special_key_strategy() -> BoxedStrategy<String> {
7323 string_regex(TEST_NAME_PATTERN)
7324 .unwrap()
7325 .prop_filter("exclude special keys", |key| {
7326 key != DOCUMENT_KEY && key != EMBEDDING_KEY
7327 })
7328 .boxed()
7329 }
7330
7331 fn source_key_strategy() -> BoxedStrategy<Option<String>> {
7332 proptest::option::of(prop_oneof![
7333 Just(DOCUMENT_KEY.to_string()),
7334 string_regex(TEST_NAME_PATTERN).unwrap(),
7335 ])
7336 .boxed()
7337 }
7338
7339 fn fts_index_type_strategy() -> impl Strategy<Value = FtsIndexType> {
7340 any::<bool>().prop_map(|enabled| FtsIndexType {
7341 enabled,
7342 config: FtsIndexConfig::default(),
7343 })
7344 }
7345
7346 fn string_inverted_index_type_strategy() -> impl Strategy<Value = StringInvertedIndexType> {
7347 any::<bool>().prop_map(|enabled| StringInvertedIndexType {
7348 enabled,
7349 config: StringInvertedIndexConfig {},
7350 })
7351 }
7352
7353 fn string_value_type_strategy() -> BoxedStrategy<Option<StringValueType>> {
7354 proptest::option::of(
7355 (
7356 proptest::option::of(string_inverted_index_type_strategy()),
7357 proptest::option::of(fts_index_type_strategy()),
7358 )
7359 .prop_map(|(string_inverted_index, fts_index)| {
7360 StringValueType {
7361 string_inverted_index,
7362 fts_index,
7363 }
7364 }),
7365 )
7366 .boxed()
7367 }
7368
7369 fn float_inverted_index_type_strategy() -> impl Strategy<Value = FloatInvertedIndexType> {
7370 any::<bool>().prop_map(|enabled| FloatInvertedIndexType {
7371 enabled,
7372 config: FloatInvertedIndexConfig {},
7373 })
7374 }
7375
7376 fn float_value_type_strategy() -> BoxedStrategy<Option<FloatValueType>> {
7377 proptest::option::of(
7378 proptest::option::of(float_inverted_index_type_strategy()).prop_map(
7379 |float_inverted_index| FloatValueType {
7380 float_inverted_index,
7381 },
7382 ),
7383 )
7384 .boxed()
7385 }
7386
7387 fn int_inverted_index_type_strategy() -> impl Strategy<Value = IntInvertedIndexType> {
7388 any::<bool>().prop_map(|enabled| IntInvertedIndexType {
7389 enabled,
7390 config: IntInvertedIndexConfig {},
7391 })
7392 }
7393
7394 fn int_value_type_strategy() -> BoxedStrategy<Option<IntValueType>> {
7395 proptest::option::of(
7396 proptest::option::of(int_inverted_index_type_strategy())
7397 .prop_map(|int_inverted_index| IntValueType { int_inverted_index }),
7398 )
7399 .boxed()
7400 }
7401
7402 fn bool_inverted_index_type_strategy() -> impl Strategy<Value = BoolInvertedIndexType> {
7403 any::<bool>().prop_map(|enabled| BoolInvertedIndexType {
7404 enabled,
7405 config: BoolInvertedIndexConfig {},
7406 })
7407 }
7408
7409 fn bool_value_type_strategy() -> BoxedStrategy<Option<BoolValueType>> {
7410 proptest::option::of(
7411 proptest::option::of(bool_inverted_index_type_strategy()).prop_map(
7412 |bool_inverted_index| BoolValueType {
7413 bool_inverted_index,
7414 },
7415 ),
7416 )
7417 .boxed()
7418 }
7419
7420 fn sparse_vector_index_config_strategy() -> impl Strategy<Value = SparseVectorIndexConfig> {
7421 (
7422 sparse_embedding_function_strategy(),
7423 source_key_strategy(),
7424 proptest::option::of(any::<bool>()),
7425 )
7426 .prop_map(|(embedding_function, source_key, bm25)| {
7427 SparseVectorIndexConfig {
7428 embedding_function,
7429 source_key,
7430 bm25,
7431 algorithm: SparseIndexAlgorithm::Wand,
7432 }
7433 })
7434 }
7435
7436 fn sparse_vector_value_type_strategy() -> BoxedStrategy<Option<SparseVectorValueType>> {
7437 proptest::option::of(
7438 (
7439 any::<bool>(),
7440 proptest::option::of(sparse_vector_index_config_strategy()),
7441 )
7442 .prop_map(|(enabled, config)| SparseVectorValueType {
7443 sparse_vector_index: config.map(|cfg| SparseVectorIndexType {
7444 enabled,
7445 config: cfg,
7446 }),
7447 }),
7448 )
7449 .boxed()
7450 }
7451
7452 fn hnsw_index_config_strategy() -> impl Strategy<Value = HnswIndexConfig> {
7453 internal_hnsw_configuration_strategy().prop_map(|config| HnswIndexConfig {
7454 ef_construction: Some(config.ef_construction),
7455 max_neighbors: Some(config.max_neighbors),
7456 ef_search: Some(config.ef_search),
7457 num_threads: Some(config.num_threads),
7458 batch_size: Some(config.batch_size),
7459 sync_threshold: Some(config.sync_threshold),
7460 resize_factor: Some(config.resize_factor),
7461 })
7462 }
7463
7464 fn spann_index_config_strategy() -> impl Strategy<Value = SpannIndexConfig> {
7465 internal_spann_configuration_strategy().prop_map(|config| SpannIndexConfig {
7466 search_nprobe: Some(config.search_nprobe),
7467 search_rng_factor: Some(config.search_rng_factor),
7468 search_rng_epsilon: Some(config.search_rng_epsilon),
7469 nreplica_count: Some(config.nreplica_count),
7470 write_rng_factor: Some(config.write_rng_factor),
7471 write_rng_epsilon: Some(config.write_rng_epsilon),
7472 split_threshold: Some(config.split_threshold),
7473 num_samples_kmeans: Some(config.num_samples_kmeans),
7474 initial_lambda: Some(config.initial_lambda),
7475 reassign_neighbor_count: Some(config.reassign_neighbor_count),
7476 merge_threshold: Some(config.merge_threshold),
7477 num_centers_to_merge_to: Some(config.num_centers_to_merge_to),
7478 write_nprobe: Some(config.write_nprobe),
7479 ef_construction: Some(config.ef_construction),
7480 ef_search: Some(config.ef_search),
7481 max_neighbors: Some(config.max_neighbors),
7482 center_drift_threshold: None,
7483 quantize: Quantization::None,
7484 })
7485 }
7486
7487 fn vector_index_config_strategy() -> impl Strategy<Value = VectorIndexConfig> {
7488 (
7489 proptest::option::of(space_strategy()),
7490 embedding_function_strategy(),
7491 source_key_strategy(),
7492 proptest::option::of(hnsw_index_config_strategy()),
7493 proptest::option::of(spann_index_config_strategy()),
7494 )
7495 .prop_map(|(space, embedding_function, source_key, hnsw, spann)| {
7496 VectorIndexConfig {
7497 space,
7498 embedding_function,
7499 source_key,
7500 hnsw,
7501 spann,
7502 }
7503 })
7504 }
7505
7506 fn vector_index_type_strategy() -> impl Strategy<Value = VectorIndexType> {
7507 (any::<bool>(), vector_index_config_strategy())
7508 .prop_map(|(enabled, config)| VectorIndexType { enabled, config })
7509 }
7510
7511 fn float_list_value_type_strategy() -> BoxedStrategy<Option<FloatListValueType>> {
7512 proptest::option::of(
7513 proptest::option::of(vector_index_type_strategy())
7514 .prop_map(|vector_index| FloatListValueType { vector_index }),
7515 )
7516 .boxed()
7517 }
7518
7519 fn value_types_strategy() -> BoxedStrategy<ValueTypes> {
7520 (
7521 string_value_type_strategy(),
7522 float_list_value_type_strategy(),
7523 sparse_vector_value_type_strategy(),
7524 int_value_type_strategy(),
7525 float_value_type_strategy(),
7526 bool_value_type_strategy(),
7527 )
7528 .prop_map(
7529 |(string, float_list, sparse_vector, int, float, boolean)| ValueTypes {
7530 string,
7531 float_list,
7532 sparse_vector,
7533 int,
7534 float,
7535 boolean,
7536 },
7537 )
7538 .boxed()
7539 }
7540
7541 fn schema_strategy() -> BoxedStrategy<Schema> {
7542 (
7543 value_types_strategy(),
7544 proptest::collection::hash_map(
7545 non_special_key_strategy(),
7546 value_types_strategy(),
7547 0..=3,
7548 ),
7549 proptest::option::of(value_types_strategy()),
7550 proptest::option::of(value_types_strategy()),
7551 )
7552 .prop_map(
7553 |(defaults, mut extra_keys, document_override, embedding_override)| {
7554 if let Some(doc) = document_override {
7555 extra_keys.insert(DOCUMENT_KEY.to_string(), doc);
7556 }
7557 if let Some(embed) = embedding_override {
7558 extra_keys.insert(EMBEDDING_KEY.to_string(), embed);
7559 }
7560 Schema {
7561 defaults,
7562 keys: extra_keys,
7563 cmek: None,
7564 }
7565 },
7566 )
7567 .boxed()
7568 }
7569
7570 fn force_non_default_schema(mut schema: Schema) -> Schema {
7571 if schema.is_default() {
7572 if let Some(string_value) = schema
7573 .defaults
7574 .string
7575 .as_mut()
7576 .and_then(|string_value| string_value.string_inverted_index.as_mut())
7577 {
7578 string_value.enabled = !string_value.enabled;
7579 } else {
7580 schema.defaults.string = Some(StringValueType {
7581 string_inverted_index: Some(StringInvertedIndexType {
7582 enabled: false,
7583 config: StringInvertedIndexConfig {},
7584 }),
7585 fts_index: None,
7586 });
7587 }
7588 }
7589 schema
7590 }
7591
7592 fn non_default_schema_strategy() -> BoxedStrategy<Schema> {
7593 schema_strategy().prop_map(force_non_default_schema).boxed()
7594 }
7595
7596 fn extract_vector_configs(schema: &Schema) -> (VectorIndexConfig, VectorIndexConfig) {
7597 let defaults = schema
7598 .defaults
7599 .float_list
7600 .as_ref()
7601 .and_then(|fl| fl.vector_index.as_ref())
7602 .map(|vi| vi.config.clone())
7603 .expect("defaults vector index missing");
7604
7605 let embedding = schema
7606 .keys
7607 .get(EMBEDDING_KEY)
7608 .and_then(|value_types| value_types.float_list.as_ref())
7609 .and_then(|fl| fl.vector_index.as_ref())
7610 .map(|vi| vi.config.clone())
7611 .expect("#embedding vector index missing");
7612
7613 (defaults, embedding)
7614 }
7615
7616 proptest! {
7617 #[test]
7618 fn reconcile_schema_and_config_matches_convert_for_config_only(
7619 config in internal_collection_configuration_strategy(),
7620 knn in knn_index_strategy(),
7621 ) {
7622 let result = Schema::reconcile_schema_and_config(None, Some(&config), knn)
7623 .expect("reconciliation should succeed");
7624
7625 let (defaults_vi, embedding_vi) = extract_vector_configs(&result);
7626 let expected_config = expected_vector_index_config(&config);
7627
7628 prop_assert_eq!(defaults_vi, expected_config.clone());
7629
7630 let mut expected_embedding_config = expected_config;
7631 expected_embedding_config.source_key = Some(DOCUMENT_KEY.to_string());
7632 prop_assert_eq!(embedding_vi, expected_embedding_config);
7633
7634 prop_assert_eq!(result.keys.len(), 2);
7635 }
7636 }
7637
7638 proptest! {
7639 #[test]
7640 fn reconcile_schema_and_config_errors_when_both_non_default(
7641 config in non_default_internal_collection_configuration_strategy(),
7642 knn in knn_index_strategy(),
7643 ) {
7644 let schema = Schema::try_from(&config)
7645 .expect("conversion should succeed");
7646 prop_assume!(!schema.is_default());
7647
7648 let result = Schema::reconcile_schema_and_config(Some(&schema), Some(&config), knn);
7649
7650 prop_assert!(matches!(result, Err(SchemaError::ConfigAndSchemaConflict)));
7651 }
7652 }
7653
7654 proptest! {
7655 #[test]
7656 fn reconcile_schema_and_config_matches_schema_only_path(
7657 schema in schema_strategy(),
7658 knn in knn_index_strategy(),
7659 ) {
7660 let result = Schema::reconcile_schema_and_config(Some(&schema), None, knn)
7661 .expect("reconciliation should succeed");
7662
7663 let (defaults_vi, embedding_vi) = extract_vector_configs(&result);
7664
7665 if let Some(schema_float_list) = schema.defaults.float_list.as_ref() {
7667 if let Some(schema_vi) = schema_float_list.vector_index.as_ref() {
7668 if let Some(schema_space) = &schema_vi.config.space {
7670 prop_assert_eq!(defaults_vi.space, Some(schema_space.clone()));
7671 }
7672 if let Some(schema_ef) = &schema_vi.config.embedding_function {
7673 prop_assert_eq!(defaults_vi.embedding_function, Some(schema_ef.clone()));
7674 }
7675 match knn {
7677 KnnIndex::Hnsw => {
7678 if let Some(schema_hnsw) = &schema_vi.config.hnsw {
7679 if let Some(merged_hnsw) = &defaults_vi.hnsw {
7680 if let Some(schema_ef_construction) = schema_hnsw.ef_construction {
7681 prop_assert_eq!(merged_hnsw.ef_construction, Some(schema_ef_construction));
7682 }
7683 }
7684 }
7685 }
7686 KnnIndex::Spann => {
7687 if let Some(schema_spann) = &schema_vi.config.spann {
7688 if let Some(merged_spann) = &defaults_vi.spann {
7689 if let Some(schema_search_nprobe) = schema_spann.search_nprobe {
7690 prop_assert_eq!(merged_spann.search_nprobe, Some(schema_search_nprobe));
7691 }
7692 }
7693 }
7694 }
7695 }
7696 }
7697 }
7698
7699 if let Some(embedding_values) = schema.keys.get(EMBEDDING_KEY) {
7701 if let Some(embedding_float_list) = embedding_values.float_list.as_ref() {
7702 if let Some(embedding_vi_type) = embedding_float_list.vector_index.as_ref() {
7703 if let Some(schema_space) = &embedding_vi_type.config.space {
7704 prop_assert_eq!(embedding_vi.space, Some(schema_space.clone()));
7705 }
7706 }
7707 }
7708 }
7709 }
7710 }
7711
7712 proptest! {
7713 #[test]
7714 fn reconcile_schema_and_config_with_default_schema_and_default_config_applies_embedding_function(
7715 embedding_function in default_embedding_function_strategy(),
7716 knn in knn_index_strategy(),
7717 ) {
7718 let schema = Schema::new_default(knn);
7719 let mut config = match knn {
7720 KnnIndex::Hnsw => InternalCollectionConfiguration::default_hnsw(),
7721 KnnIndex::Spann => InternalCollectionConfiguration::default_spann(),
7722 };
7723 config.embedding_function = embedding_function.clone();
7724
7725 let result = Schema::reconcile_schema_and_config(
7726 Some(&schema),
7727 Some(&config),
7728 knn,
7729 )
7730 .expect("reconciliation should succeed");
7731
7732 let (defaults_vi, embedding_vi) = extract_vector_configs(&result);
7733
7734 if let Some(ef) = embedding_function {
7736 prop_assert_eq!(defaults_vi.embedding_function, Some(ef.clone()));
7737 prop_assert_eq!(embedding_vi.embedding_function, Some(ef));
7738 } else {
7739 prop_assert_eq!(defaults_vi.embedding_function, None);
7741 prop_assert_eq!(embedding_vi.embedding_function, None);
7742 }
7743 }
7744 }
7745
7746 proptest! {
7747 #[test]
7748 fn reconcile_schema_and_config_with_default_config_keeps_non_default_schema(
7749 schema in non_default_schema_strategy(),
7750 knn in knn_index_strategy(),
7751 ) {
7752 let default_config = match knn {
7753 KnnIndex::Hnsw => InternalCollectionConfiguration::default_hnsw(),
7754 KnnIndex::Spann => InternalCollectionConfiguration::default_spann(),
7755 };
7756
7757 let result = Schema::reconcile_schema_and_config(
7758 Some(&schema),
7759 Some(&default_config),
7760 knn,
7761 )
7762 .expect("reconciliation should succeed");
7763
7764 let (defaults_vi, embedding_vi) = extract_vector_configs(&result);
7765
7766 if let Some(schema_float_list) = schema.defaults.float_list.as_ref() {
7769 if let Some(schema_vi) = schema_float_list.vector_index.as_ref() {
7770 if let Some(schema_space) = &schema_vi.config.space {
7771 prop_assert_eq!(defaults_vi.space, Some(schema_space.clone()));
7772 }
7773 if let Some(schema_ef) = &schema_vi.config.embedding_function {
7774 prop_assert_eq!(defaults_vi.embedding_function, Some(schema_ef.clone()));
7775 }
7776 }
7777 }
7778
7779 if let Some(embedding_values) = schema.keys.get(EMBEDDING_KEY) {
7781 if let Some(embedding_float_list) = embedding_values.float_list.as_ref() {
7782 if let Some(embedding_vi_type) = embedding_float_list.vector_index.as_ref() {
7783 if let Some(schema_space) = &embedding_vi_type.config.space {
7784 prop_assert_eq!(embedding_vi.space, Some(schema_space.clone()));
7785 }
7786 }
7787 }
7788 }
7789 }
7790 }
7791 }
7792}