use chroma_error::{ChromaError, ErrorCodes};
use regex::Regex;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::{Arc, LazyLock};
use thiserror::Error;
use validator::Validate;
use crate::chroma_proto;
use crate::collection_configuration::{
EmbeddingFunctionConfiguration, InternalCollectionConfiguration,
UpdateVectorIndexConfiguration, VectorIndexConfiguration,
};
use crate::hnsw_configuration::Space;
use crate::metadata::{MetadataComparison, MetadataValueType, Where};
use crate::operator::QueryVector;
use crate::{
default_batch_size, default_center_drift_threshold, default_construction_ef,
default_construction_ef_spann, default_initial_lambda, default_m, default_m_spann,
default_merge_threshold, default_nreplica_count, default_num_centers_to_merge_to,
default_num_samples_kmeans, default_num_threads, default_reassign_neighbor_count,
default_resize_factor, default_search_ef, default_search_ef_spann, default_search_nprobe,
default_search_rng_epsilon, default_search_rng_factor, default_space, default_split_threshold,
default_sync_threshold, default_write_nprobe, default_write_rng_epsilon,
default_write_rng_factor, ConversionError, HnswParametersFromSegmentError,
InternalHnswConfiguration, InternalSpannConfiguration, InternalUpdateCollectionConfiguration,
KnnIndex, Segment, UpdateCollectionConfiguration, CHROMA_KEY,
};
impl ChromaError for SchemaError {
fn code(&self) -> ErrorCodes {
match self {
SchemaError::MissingIndexConfiguration { .. } => ErrorCodes::Internal,
SchemaError::InvalidSchema { .. } => ErrorCodes::Internal,
SchemaError::DefaultsMismatch => ErrorCodes::Internal,
SchemaError::ConfigurationConflict { .. } => ErrorCodes::Internal,
SchemaError::InvalidConfigurationUpdate { .. } => ErrorCodes::Internal,
SchemaError::InvalidUserInput { .. } => ErrorCodes::InvalidArgument,
SchemaError::ConfigAndSchemaConflict => ErrorCodes::InvalidArgument,
SchemaError::InvalidHnswConfig(_) => ErrorCodes::InvalidArgument,
SchemaError::InvalidSpannConfig(_) => ErrorCodes::InvalidArgument,
SchemaError::Builder(e) => e.code(),
}
}
}
#[derive(Debug, Error)]
pub enum SchemaError {
#[error("Schema is malformed: missing index configuration for metadata key '{key}' with type '{value_type}'")]
MissingIndexConfiguration { key: String, value_type: String },
#[error("Schema reconciliation failed: {reason}")]
InvalidSchema { reason: String },
#[error("Cannot set both collection config and schema simultaneously")]
ConfigAndSchemaConflict,
#[error("Cannot merge schemas with differing defaults")]
DefaultsMismatch,
#[error("Conflicting configuration for {context}")]
ConfigurationConflict { context: String },
#[error("Invalid HNSW configuration: {0}")]
InvalidHnswConfig(validator::ValidationErrors),
#[error("Invalid SPANN configuration: {0}")]
InvalidSpannConfig(validator::ValidationErrors),
#[error("Invalid schema input: {reason}")]
InvalidUserInput { reason: String },
#[error("Invalid configuration update: {message}")]
InvalidConfigurationUpdate { message: String },
#[error(transparent)]
Builder(#[from] SchemaBuilderError),
}
#[derive(Debug, Error)]
pub enum SchemaBuilderError {
#[error("Vector index must be configured globally using create_index(None, config), not on specific key '{key}'")]
VectorIndexMustBeGlobal { key: String },
#[error("Cannot modify special key '{key}' - it is managed automatically by the system.")]
SpecialKeyModificationNotAllowed { key: String },
#[error("Sparse vector index requires a specific key. Use create_index(Some(\"key_name\"), config) instead of create_index(None, config)")]
SparseVectorRequiresKey,
#[error("Only one sparse vector index allowed per collection. Key '{existing_key}' already has a sparse vector index. Remove it first or use that key.")]
MultipleSparseVectorIndexes { existing_key: String },
#[error("Vector index deletion not supported. The vector index is always enabled on #embedding. To disable vector search, disable the collection instead.")]
VectorIndexDeletionNotSupported,
#[error("Sparse vector index deletion not supported yet. Sparse vector indexes cannot be removed once created.")]
SparseVectorIndexDeletionNotSupported,
#[error(
"Key '{key}' cannot begin with '#'. Keys starting with '#' are reserved for system use."
)]
ReservedKeyPrefix { key: String },
#[error("FTS index deletion is only supported on #document key.")]
FtsIndexDeletionOnlyOnDocument,
#[error("FTS index can only be enabled on #document key. Use create_index(Some(\"#document\"), FtsIndexConfig) to enable FTS.")]
FtsIndexOnlyOnDocument,
}
#[derive(Debug, Error)]
pub enum FilterValidationError {
#[error(
"Cannot filter using metadata key '{key}' with type '{value_type:?}' because indexing is disabled"
)]
IndexingDisabled {
key: String,
value_type: MetadataValueType,
},
#[error("Cannot filter using full-text search because FTS indexing is disabled")]
FtsDisabled,
#[error(transparent)]
Schema(#[from] SchemaError),
}
impl ChromaError for SchemaBuilderError {
fn code(&self) -> ErrorCodes {
ErrorCodes::InvalidArgument
}
}
impl ChromaError for FilterValidationError {
fn code(&self) -> ErrorCodes {
match self {
FilterValidationError::IndexingDisabled { .. } => ErrorCodes::InvalidArgument,
FilterValidationError::FtsDisabled => ErrorCodes::InvalidArgument,
FilterValidationError::Schema(_) => ErrorCodes::Internal,
}
}
}
pub const STRING_VALUE_NAME: &str = "string";
pub const INT_VALUE_NAME: &str = "int";
pub const BOOL_VALUE_NAME: &str = "bool";
pub const FLOAT_VALUE_NAME: &str = "float";
pub const FLOAT_LIST_VALUE_NAME: &str = "float_list";
pub const SPARSE_VECTOR_VALUE_NAME: &str = "sparse_vector";
pub const FTS_INDEX_NAME: &str = "fts_index";
pub const VECTOR_INDEX_NAME: &str = "vector_index";
pub const SPARSE_VECTOR_INDEX_NAME: &str = "sparse_vector_index";
pub const STRING_INVERTED_INDEX_NAME: &str = "string_inverted_index";
pub const INT_INVERTED_INDEX_NAME: &str = "int_inverted_index";
pub const FLOAT_INVERTED_INDEX_NAME: &str = "float_inverted_index";
pub const BOOL_INVERTED_INDEX_NAME: &str = "bool_inverted_index";
pub const DOCUMENT_KEY: &str = "#document";
pub const EMBEDDING_KEY: &str = "#embedding";
static CMEK_GCP_RE: LazyLock<Regex> = LazyLock::new(|| {
Regex::new(r"^projects/.+/locations/.+/keyRings/.+/cryptoKeys/.+$")
.expect("The CMEK pattern for GCP should be valid")
});
#[derive(Clone, Debug, Eq, PartialEq, serde::Deserialize, serde::Serialize)]
#[serde(rename_all = "snake_case")]
pub enum Cmek {
Gcp(Arc<String>),
}
impl Cmek {
pub fn gcp(resource: String) -> Self {
Cmek::Gcp(Arc::new(resource))
}
pub fn validate_pattern(&self) -> bool {
match self {
Cmek::Gcp(resource) => CMEK_GCP_RE.is_match(resource),
}
}
}
impl TryFrom<chroma_proto::Cmek> for Cmek {
type Error = ConversionError;
fn try_from(proto: chroma_proto::Cmek) -> Result<Self, Self::Error> {
match proto.provider {
Some(chroma_proto::cmek::Provider::Gcp(resource)) => Ok(Cmek::gcp(resource)),
None => Err(ConversionError::DecodeError),
}
}
}
impl From<Cmek> for chroma_proto::Cmek {
fn from(cmek: Cmek) -> Self {
match cmek {
Cmek::Gcp(resource) => chroma_proto::Cmek {
provider: Some(chroma_proto::cmek::Provider::Gcp((*resource).clone())),
},
}
}
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
pub struct Schema {
pub defaults: ValueTypes,
#[serde(rename = "keys", alias = "key_overrides")]
pub keys: HashMap<String, ValueTypes>,
#[serde(skip_serializing_if = "Option::is_none")]
#[cfg_attr(feature = "utoipa", schema(value_type = Option<Object>))]
pub cmek: Option<Cmek>,
#[serde(skip_serializing_if = "Option::is_none")]
pub source_attached_function_id: Option<String>,
}
impl Schema {
pub fn update(&mut self, configuration: &InternalUpdateCollectionConfiguration) {
if let Some(vector_update) = &configuration.vector_index {
if let Some(default_vector_index) = self.defaults_vector_index_mut() {
Self::apply_vector_index_update(default_vector_index, vector_update);
}
if let Some(embedding_vector_index) = self.embedding_vector_index_mut() {
Self::apply_vector_index_update(embedding_vector_index, vector_update);
}
}
if let Some(embedding_function) = configuration.embedding_function.as_ref() {
if let Some(default_vector_index) = self.defaults_vector_index_mut() {
default_vector_index.config.embedding_function = Some(embedding_function.clone());
}
if let Some(embedding_vector_index) = self.embedding_vector_index_mut() {
embedding_vector_index.config.embedding_function = Some(embedding_function.clone());
}
}
}
pub fn apply_update_configuration(
&mut self,
config: &UpdateCollectionConfiguration,
) -> Result<(), SchemaError> {
if config.hnsw.is_some() {
return Err(SchemaError::InvalidConfigurationUpdate {
message: "HNSW configuration updates are not supported".to_string(),
});
}
if let Some(ref spann_update) = config.spann {
let defaults_spann = self
.defaults_vector_index_mut()
.ok_or_else(|| SchemaError::InvalidConfigurationUpdate {
message: "schema missing defaults.float_list.vector_index".to_string(),
})?
.config
.spann
.as_mut()
.ok_or_else(|| SchemaError::InvalidConfigurationUpdate {
message: "schema missing defaults spann config".to_string(),
})?;
if let Some(search_nprobe) = spann_update.search_nprobe {
defaults_spann.search_nprobe = Some(search_nprobe);
}
if let Some(ef_search) = spann_update.ef_search {
defaults_spann.ef_search = Some(ef_search);
}
let embedding_spann = self
.embedding_vector_index_mut()
.ok_or_else(|| SchemaError::InvalidConfigurationUpdate {
message: "schema missing keys[#embedding].float_list.vector_index".to_string(),
})?
.config
.spann
.as_mut()
.ok_or_else(|| SchemaError::InvalidConfigurationUpdate {
message: "schema missing #embedding spann config".to_string(),
})?;
if let Some(search_nprobe) = spann_update.search_nprobe {
embedding_spann.search_nprobe = Some(search_nprobe);
}
if let Some(ef_search) = spann_update.ef_search {
embedding_spann.ef_search = Some(ef_search);
}
}
if let Some(ref ef) = config.embedding_function {
self.defaults_vector_index_mut()
.ok_or_else(|| SchemaError::InvalidConfigurationUpdate {
message: "schema missing defaults.float_list.vector_index".to_string(),
})?
.config
.embedding_function = Some(ef.clone());
self.embedding_vector_index_mut()
.ok_or_else(|| SchemaError::InvalidConfigurationUpdate {
message: "schema missing keys[#embedding].float_list.vector_index".to_string(),
})?
.config
.embedding_function = Some(ef.clone());
}
Ok(())
}
fn defaults_vector_index_mut(&mut self) -> Option<&mut VectorIndexType> {
self.defaults
.float_list
.as_mut()
.and_then(|float_list| float_list.vector_index.as_mut())
}
fn embedding_vector_index_mut(&mut self) -> Option<&mut VectorIndexType> {
self.keys
.get_mut(EMBEDDING_KEY)
.and_then(|value_types| value_types.float_list.as_mut())
.and_then(|float_list| float_list.vector_index.as_mut())
}
fn apply_vector_index_update(
vector_index: &mut VectorIndexType,
update: &UpdateVectorIndexConfiguration,
) {
match update {
UpdateVectorIndexConfiguration::Hnsw(Some(hnsw_update)) => {
if let Some(hnsw_config) = vector_index.config.hnsw.as_mut() {
if let Some(ef_search) = hnsw_update.ef_search {
hnsw_config.ef_search = Some(ef_search);
}
if let Some(max_neighbors) = hnsw_update.max_neighbors {
hnsw_config.max_neighbors = Some(max_neighbors);
}
if let Some(num_threads) = hnsw_update.num_threads {
hnsw_config.num_threads = Some(num_threads);
}
if let Some(resize_factor) = hnsw_update.resize_factor {
hnsw_config.resize_factor = Some(resize_factor);
}
if let Some(sync_threshold) = hnsw_update.sync_threshold {
hnsw_config.sync_threshold = Some(sync_threshold);
}
if let Some(batch_size) = hnsw_update.batch_size {
hnsw_config.batch_size = Some(batch_size);
}
}
}
UpdateVectorIndexConfiguration::Hnsw(None) => {}
UpdateVectorIndexConfiguration::Spann(Some(spann_update)) => {
if let Some(spann_config) = vector_index.config.spann.as_mut() {
if let Some(search_nprobe) = spann_update.search_nprobe {
spann_config.search_nprobe = Some(search_nprobe);
}
if let Some(ef_search) = spann_update.ef_search {
spann_config.ef_search = Some(ef_search);
}
}
}
UpdateVectorIndexConfiguration::Spann(None) => {}
}
}
pub fn is_sparse_index_enabled(&self) -> bool {
let defaults_enabled = self
.defaults
.sparse_vector
.as_ref()
.and_then(|sv| sv.sparse_vector_index.as_ref())
.is_some_and(|idx| idx.enabled);
let key_enabled = self.keys.values().any(|value_types| {
value_types
.sparse_vector
.as_ref()
.and_then(|sv| sv.sparse_vector_index.as_ref())
.is_some_and(|idx| idx.enabled)
});
defaults_enabled || key_enabled
}
pub fn is_fts_enabled(&self) -> bool {
self.keys
.get(DOCUMENT_KEY)
.and_then(|vt| vt.string.as_ref())
.and_then(|s| s.fts_index.as_ref())
.or_else(|| {
self.defaults
.string
.as_ref()
.and_then(|s| s.fts_index.as_ref())
})
.is_none_or(|idx| idx.enabled)
}
}
impl Default for Schema {
fn default() -> Self {
let defaults = ValueTypes {
string: Some(StringValueType {
fts_index: Some(FtsIndexType {
enabled: false,
config: FtsIndexConfig {},
}),
string_inverted_index: Some(StringInvertedIndexType {
enabled: true,
config: StringInvertedIndexConfig {},
}),
}),
float_list: Some(FloatListValueType {
vector_index: Some(VectorIndexType {
enabled: false,
config: VectorIndexConfig {
space: None, embedding_function: Some(EmbeddingFunctionConfiguration::Legacy),
source_key: None,
hnsw: None, spann: None, },
}),
}),
sparse_vector: Some(SparseVectorValueType {
sparse_vector_index: Some(SparseVectorIndexType {
enabled: false,
config: SparseVectorIndexConfig {
embedding_function: None,
source_key: None,
bm25: None,
},
}),
}),
int: Some(IntValueType {
int_inverted_index: Some(IntInvertedIndexType {
enabled: true,
config: IntInvertedIndexConfig {},
}),
}),
float: Some(FloatValueType {
float_inverted_index: Some(FloatInvertedIndexType {
enabled: true,
config: FloatInvertedIndexConfig {},
}),
}),
boolean: Some(BoolValueType {
bool_inverted_index: Some(BoolInvertedIndexType {
enabled: true,
config: BoolInvertedIndexConfig {},
}),
}),
};
let mut keys = HashMap::new();
keys.insert(
DOCUMENT_KEY.to_string(),
ValueTypes {
string: Some(StringValueType {
fts_index: Some(FtsIndexType {
enabled: true,
config: FtsIndexConfig {},
}),
string_inverted_index: Some(StringInvertedIndexType {
enabled: false,
config: StringInvertedIndexConfig {},
}),
}),
..Default::default()
},
);
keys.insert(
EMBEDDING_KEY.to_string(),
ValueTypes {
float_list: Some(FloatListValueType {
vector_index: Some(VectorIndexType {
enabled: true,
config: VectorIndexConfig {
space: None, embedding_function: Some(EmbeddingFunctionConfiguration::Legacy),
source_key: Some(DOCUMENT_KEY.to_string()),
hnsw: None, spann: None, },
}),
}),
..Default::default()
},
);
Schema {
defaults,
keys,
cmek: None,
source_attached_function_id: None,
}
}
}
pub fn is_embedding_function_default(
embedding_function: &Option<EmbeddingFunctionConfiguration>,
) -> bool {
match embedding_function {
None => true,
Some(embedding_function) => embedding_function.is_default(),
}
}
pub fn is_space_default(space: &Option<Space>) -> bool {
match space {
None => true, Some(s) => *s == default_space(), }
}
pub fn is_hnsw_config_default(hnsw_config: &HnswIndexConfig) -> bool {
hnsw_config.ef_construction == Some(default_construction_ef())
&& hnsw_config.ef_search == Some(default_search_ef())
&& hnsw_config.max_neighbors == Some(default_m())
&& hnsw_config.num_threads == Some(default_num_threads())
&& hnsw_config.batch_size == Some(default_batch_size())
&& hnsw_config.sync_threshold == Some(default_sync_threshold())
&& hnsw_config.resize_factor == Some(default_resize_factor())
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Default)]
#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
pub struct ValueTypes {
#[serde(
rename = "string",
alias = "#string",
skip_serializing_if = "Option::is_none"
)] pub string: Option<StringValueType>,
#[serde(
rename = "float_list",
alias = "#float_list",
skip_serializing_if = "Option::is_none"
)]
pub float_list: Option<FloatListValueType>,
#[serde(
rename = "sparse_vector",
alias = "#sparse_vector",
skip_serializing_if = "Option::is_none"
)]
pub sparse_vector: Option<SparseVectorValueType>,
#[serde(
rename = "int",
alias = "#int",
skip_serializing_if = "Option::is_none"
)] pub int: Option<IntValueType>,
#[serde(
rename = "float",
alias = "#float",
skip_serializing_if = "Option::is_none"
)] pub float: Option<FloatValueType>,
#[serde(
rename = "bool",
alias = "#bool",
skip_serializing_if = "Option::is_none"
)] pub boolean: Option<BoolValueType>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
pub struct StringValueType {
#[serde(
rename = "fts_index",
alias = "$fts_index",
skip_serializing_if = "Option::is_none"
)] pub fts_index: Option<FtsIndexType>,
#[serde(
rename = "string_inverted_index", // STRING_INVERTED_INDEX_NAME
alias = "$string_inverted_index",
skip_serializing_if = "Option::is_none"
)]
pub string_inverted_index: Option<StringInvertedIndexType>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
pub struct FloatListValueType {
#[serde(
rename = "vector_index",
alias = "$vector_index",
skip_serializing_if = "Option::is_none"
)] pub vector_index: Option<VectorIndexType>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
pub struct SparseVectorValueType {
#[serde(
rename = "sparse_vector_index", // SPARSE_VECTOR_INDEX_NAME
alias = "$sparse_vector_index",
skip_serializing_if = "Option::is_none"
)]
pub sparse_vector_index: Option<SparseVectorIndexType>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
pub struct IntValueType {
#[serde(
rename = "int_inverted_index",
alias = "$int_inverted_index",
skip_serializing_if = "Option::is_none"
)]
pub int_inverted_index: Option<IntInvertedIndexType>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
pub struct FloatValueType {
#[serde(
rename = "float_inverted_index", // FLOAT_INVERTED_INDEX_NAME
alias = "$float_inverted_index",
skip_serializing_if = "Option::is_none"
)]
pub float_inverted_index: Option<FloatInvertedIndexType>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
pub struct BoolValueType {
#[serde(
rename = "bool_inverted_index", // BOOL_INVERTED_INDEX_NAME
alias = "$bool_inverted_index",
skip_serializing_if = "Option::is_none"
)]
pub bool_inverted_index: Option<BoolInvertedIndexType>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
pub struct FtsIndexType {
pub enabled: bool,
pub config: FtsIndexConfig,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
pub struct VectorIndexType {
pub enabled: bool,
pub config: VectorIndexConfig,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
pub struct SparseVectorIndexType {
pub enabled: bool,
pub config: SparseVectorIndexConfig,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
pub struct StringInvertedIndexType {
pub enabled: bool,
pub config: StringInvertedIndexConfig,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
pub struct IntInvertedIndexType {
pub enabled: bool,
pub config: IntInvertedIndexConfig,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
pub struct FloatInvertedIndexType {
pub enabled: bool,
pub config: FloatInvertedIndexConfig,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
pub struct BoolInvertedIndexType {
pub enabled: bool,
pub config: BoolInvertedIndexConfig,
}
impl Schema {
pub fn new_default(default_knn_index: KnnIndex) -> Self {
let vector_config = VectorIndexType {
enabled: false,
config: VectorIndexConfig {
space: Some(default_space()),
embedding_function: None,
source_key: None,
hnsw: match default_knn_index {
KnnIndex::Hnsw => Some(HnswIndexConfig {
ef_construction: Some(default_construction_ef()),
max_neighbors: Some(default_m()),
ef_search: Some(default_search_ef()),
num_threads: Some(default_num_threads()),
batch_size: Some(default_batch_size()),
sync_threshold: Some(default_sync_threshold()),
resize_factor: Some(default_resize_factor()),
}),
KnnIndex::Spann => None,
},
spann: match default_knn_index {
KnnIndex::Hnsw => None,
KnnIndex::Spann => Some(SpannIndexConfig {
search_nprobe: Some(default_search_nprobe()),
search_rng_factor: Some(default_search_rng_factor()),
search_rng_epsilon: Some(default_search_rng_epsilon()),
nreplica_count: Some(default_nreplica_count()),
write_rng_factor: Some(default_write_rng_factor()),
write_rng_epsilon: Some(default_write_rng_epsilon()),
split_threshold: Some(default_split_threshold()),
num_samples_kmeans: Some(default_num_samples_kmeans()),
initial_lambda: Some(default_initial_lambda()),
reassign_neighbor_count: Some(default_reassign_neighbor_count()),
merge_threshold: Some(default_merge_threshold()),
num_centers_to_merge_to: Some(default_num_centers_to_merge_to()),
write_nprobe: Some(default_write_nprobe()),
ef_construction: Some(default_construction_ef_spann()),
ef_search: Some(default_search_ef_spann()),
max_neighbors: Some(default_m_spann()),
center_drift_threshold: None,
quantize: Quantization::None,
}),
},
},
};
let defaults = ValueTypes {
string: Some(StringValueType {
string_inverted_index: Some(StringInvertedIndexType {
enabled: true,
config: StringInvertedIndexConfig {},
}),
fts_index: Some(FtsIndexType {
enabled: false,
config: FtsIndexConfig {},
}),
}),
float: Some(FloatValueType {
float_inverted_index: Some(FloatInvertedIndexType {
enabled: true,
config: FloatInvertedIndexConfig {},
}),
}),
int: Some(IntValueType {
int_inverted_index: Some(IntInvertedIndexType {
enabled: true,
config: IntInvertedIndexConfig {},
}),
}),
boolean: Some(BoolValueType {
bool_inverted_index: Some(BoolInvertedIndexType {
enabled: true,
config: BoolInvertedIndexConfig {},
}),
}),
float_list: Some(FloatListValueType {
vector_index: Some(vector_config),
}),
sparse_vector: Some(SparseVectorValueType {
sparse_vector_index: Some(SparseVectorIndexType {
enabled: false,
config: SparseVectorIndexConfig {
embedding_function: Some(EmbeddingFunctionConfiguration::Unknown),
source_key: None,
bm25: Some(false),
},
}),
}),
};
let mut keys = HashMap::new();
let embedding_defaults = ValueTypes {
float_list: Some(FloatListValueType {
vector_index: Some(VectorIndexType {
enabled: true,
config: VectorIndexConfig {
space: Some(default_space()),
embedding_function: None,
source_key: Some(DOCUMENT_KEY.to_string()),
hnsw: match default_knn_index {
KnnIndex::Hnsw => Some(HnswIndexConfig {
ef_construction: Some(default_construction_ef()),
max_neighbors: Some(default_m()),
ef_search: Some(default_search_ef()),
num_threads: Some(default_num_threads()),
batch_size: Some(default_batch_size()),
sync_threshold: Some(default_sync_threshold()),
resize_factor: Some(default_resize_factor()),
}),
KnnIndex::Spann => None,
},
spann: match default_knn_index {
KnnIndex::Hnsw => None,
KnnIndex::Spann => Some(SpannIndexConfig {
search_nprobe: Some(default_search_nprobe()),
search_rng_factor: Some(default_search_rng_factor()),
search_rng_epsilon: Some(default_search_rng_epsilon()),
nreplica_count: Some(default_nreplica_count()),
write_rng_factor: Some(default_write_rng_factor()),
write_rng_epsilon: Some(default_write_rng_epsilon()),
split_threshold: Some(default_split_threshold()),
num_samples_kmeans: Some(default_num_samples_kmeans()),
initial_lambda: Some(default_initial_lambda()),
reassign_neighbor_count: Some(default_reassign_neighbor_count()),
merge_threshold: Some(default_merge_threshold()),
num_centers_to_merge_to: Some(default_num_centers_to_merge_to()),
write_nprobe: Some(default_write_nprobe()),
ef_construction: Some(default_construction_ef_spann()),
ef_search: Some(default_search_ef_spann()),
max_neighbors: Some(default_m_spann()),
center_drift_threshold: None,
quantize: Quantization::None,
}),
},
},
}),
}),
..Default::default()
};
keys.insert(EMBEDDING_KEY.to_string(), embedding_defaults);
let document_defaults = ValueTypes {
string: Some(StringValueType {
fts_index: Some(FtsIndexType {
enabled: true,
config: FtsIndexConfig {},
}),
string_inverted_index: Some(StringInvertedIndexType {
enabled: false,
config: StringInvertedIndexConfig {},
}),
}),
..Default::default()
};
keys.insert(DOCUMENT_KEY.to_string(), document_defaults);
Schema {
defaults,
keys,
cmek: None,
source_attached_function_id: None,
}
}
pub fn get_spann_config(&self) -> Option<(SpannIndexConfig, Space)> {
let extract = |vector_index: &VectorIndexType| {
let space = vector_index.config.space.clone().unwrap_or_default();
vector_index
.config
.spann
.clone()
.map(|config| (config, space))
};
self.keys
.get(EMBEDDING_KEY)
.and_then(|value_types| value_types.float_list.as_ref())
.and_then(|float_list| float_list.vector_index.as_ref())
.and_then(extract)
.or_else(|| {
self.defaults
.float_list
.as_ref()
.and_then(|float_list| float_list.vector_index.as_ref())
.and_then(extract)
})
}
pub fn get_internal_spann_config(&self) -> Option<InternalSpannConfiguration> {
let to_internal = |vector_index: &VectorIndexType| {
let space = vector_index.config.space.clone();
vector_index
.config
.spann
.clone()
.map(|config| (space.as_ref(), &config).into())
};
self.keys
.get(EMBEDDING_KEY)
.and_then(|value_types| value_types.float_list.as_ref())
.and_then(|float_list| float_list.vector_index.as_ref())
.and_then(to_internal)
.or_else(|| {
self.defaults
.float_list
.as_ref()
.and_then(|float_list| float_list.vector_index.as_ref())
.and_then(to_internal)
})
}
pub fn is_quantization_enabled(&self) -> bool {
let check_spann = |vector_index: &VectorIndexType| {
vector_index
.config
.spann
.as_ref()
.is_some_and(|config| !matches!(config.quantize, Quantization::None))
};
self.keys
.get(EMBEDDING_KEY)
.and_then(|value_types| value_types.float_list.as_ref())
.and_then(|float_list| float_list.vector_index.as_ref())
.map(check_spann)
.unwrap_or_else(|| {
self.defaults
.float_list
.as_ref()
.and_then(|float_list| float_list.vector_index.as_ref())
.map(check_spann)
.unwrap_or(false)
})
}
pub fn get_spann_config_mut(&mut self) -> Option<&mut SpannIndexConfig> {
if let Some(value_types) = self.keys.get_mut(EMBEDDING_KEY) {
if let Some(float_list) = &mut value_types.float_list {
if let Some(vector_index) = &mut float_list.vector_index {
if let Some(spann_config) = &mut vector_index.config.spann {
return Some(spann_config);
}
}
}
}
if let Some(float_list) = &mut self.defaults.float_list {
if let Some(vector_index) = &mut float_list.vector_index {
if let Some(spann_config) = &mut vector_index.config.spann {
return Some(spann_config);
}
}
}
None
}
pub fn quantize(&mut self, variant: Quantization) {
if let Some(spann_config) = self.get_spann_config_mut() {
*spann_config = match variant {
Quantization::None => SpannIndexConfig {
quantize: variant,
..*spann_config
},
Quantization::FourBitRabitQWithUSearch => SpannIndexConfig {
search_nprobe: Some(64),
nreplica_count: Some(2),
write_rng_factor: Some(4.0),
write_rng_epsilon: Some(8.0),
split_threshold: Some(512),
reassign_neighbor_count: Some(32),
merge_threshold: Some(128),
write_nprobe: Some(64),
ef_construction: Some(256),
ef_search: Some(128),
max_neighbors: Some(24),
center_drift_threshold: Some(0.125),
quantize: variant,
..*spann_config
},
};
}
}
pub fn get_internal_hnsw_config(&self) -> Option<InternalHnswConfiguration> {
let to_internal = |vector_index: &VectorIndexType| {
if vector_index.config.spann.is_some() {
return None;
}
let space = vector_index.config.space.as_ref();
let hnsw_config = vector_index.config.hnsw.as_ref();
Some((space, hnsw_config).into())
};
self.keys
.get(EMBEDDING_KEY)
.and_then(|value_types| value_types.float_list.as_ref())
.and_then(|float_list| float_list.vector_index.as_ref())
.and_then(to_internal)
.or_else(|| {
self.defaults
.float_list
.as_ref()
.and_then(|float_list| float_list.vector_index.as_ref())
.and_then(to_internal)
})
}
pub fn get_internal_hnsw_config_with_legacy_fallback(
&self,
segment: &Segment,
) -> Result<Option<InternalHnswConfiguration>, HnswParametersFromSegmentError> {
if let Some(config) = self.get_internal_hnsw_config() {
let config_from_metadata =
InternalHnswConfiguration::from_legacy_segment_metadata(&segment.metadata)?;
if config == InternalHnswConfiguration::default() && config != config_from_metadata {
return Ok(Some(config_from_metadata));
}
return Ok(Some(config));
}
Ok(None)
}
pub fn reconcile_with_defaults(
user_schema: Option<&Schema>,
knn_index: KnnIndex,
) -> Result<Self, SchemaError> {
let default_schema = Schema::new_default(knn_index);
match user_schema {
Some(user) => {
let merged_defaults =
Self::merge_value_types(&default_schema.defaults, &user.defaults, knn_index)?;
let mut merged_keys = default_schema.keys.clone();
for (key, user_value_types) in &user.keys {
if let Some(default_value_types) = merged_keys.get(key) {
let merged_value_types = Self::merge_value_types(
default_value_types,
user_value_types,
knn_index,
)?;
merged_keys.insert(key.clone(), merged_value_types);
} else {
merged_keys.insert(key.clone(), user_value_types.clone());
}
}
Ok(Schema {
defaults: merged_defaults,
keys: merged_keys,
cmek: user.cmek.clone().or(default_schema.cmek.clone()),
source_attached_function_id: user
.source_attached_function_id
.clone()
.or(default_schema.source_attached_function_id.clone()),
})
}
None => Ok(default_schema),
}
}
pub fn merge(&self, other: &Schema) -> Result<Schema, SchemaError> {
if self.defaults != other.defaults {
return Err(SchemaError::DefaultsMismatch);
}
let mut keys = self.keys.clone();
for (key, other_value_types) in &other.keys {
if let Some(existing) = keys.get(key).cloned() {
let merged = Self::merge_override_value_types(key, &existing, other_value_types)?;
keys.insert(key.clone(), merged);
} else {
keys.insert(key.clone(), other_value_types.clone());
}
}
Ok(Schema {
defaults: self.defaults.clone(),
keys,
cmek: other.cmek.clone().or(self.cmek.clone()),
source_attached_function_id: other
.source_attached_function_id
.clone()
.or(self.source_attached_function_id.clone()),
})
}
fn merge_override_value_types(
key: &str,
left: &ValueTypes,
right: &ValueTypes,
) -> Result<ValueTypes, SchemaError> {
Ok(ValueTypes {
string: Self::merge_string_override(key, left.string.as_ref(), right.string.as_ref())?,
float: Self::merge_float_override(key, left.float.as_ref(), right.float.as_ref())?,
int: Self::merge_int_override(key, left.int.as_ref(), right.int.as_ref())?,
boolean: Self::merge_bool_override(key, left.boolean.as_ref(), right.boolean.as_ref())?,
float_list: Self::merge_float_list_override(
key,
left.float_list.as_ref(),
right.float_list.as_ref(),
)?,
sparse_vector: Self::merge_sparse_vector_override(
key,
left.sparse_vector.as_ref(),
right.sparse_vector.as_ref(),
)?,
})
}
fn merge_string_override(
key: &str,
left: Option<&StringValueType>,
right: Option<&StringValueType>,
) -> Result<Option<StringValueType>, SchemaError> {
match (left, right) {
(Some(l), Some(r)) => Ok(Some(StringValueType {
string_inverted_index: Self::merge_index_or_error(
l.string_inverted_index.as_ref(),
r.string_inverted_index.as_ref(),
&format!("key '{key}' string.string_inverted_index"),
)?,
fts_index: Self::merge_index_or_error(
l.fts_index.as_ref(),
r.fts_index.as_ref(),
&format!("key '{key}' string.fts_index"),
)?,
})),
(Some(l), None) => Ok(Some(l.clone())),
(None, Some(r)) => Ok(Some(r.clone())),
(None, None) => Ok(None),
}
}
fn merge_float_override(
key: &str,
left: Option<&FloatValueType>,
right: Option<&FloatValueType>,
) -> Result<Option<FloatValueType>, SchemaError> {
match (left, right) {
(Some(l), Some(r)) => Ok(Some(FloatValueType {
float_inverted_index: Self::merge_index_or_error(
l.float_inverted_index.as_ref(),
r.float_inverted_index.as_ref(),
&format!("key '{key}' float.float_inverted_index"),
)?,
})),
(Some(l), None) => Ok(Some(l.clone())),
(None, Some(r)) => Ok(Some(r.clone())),
(None, None) => Ok(None),
}
}
fn merge_int_override(
key: &str,
left: Option<&IntValueType>,
right: Option<&IntValueType>,
) -> Result<Option<IntValueType>, SchemaError> {
match (left, right) {
(Some(l), Some(r)) => Ok(Some(IntValueType {
int_inverted_index: Self::merge_index_or_error(
l.int_inverted_index.as_ref(),
r.int_inverted_index.as_ref(),
&format!("key '{key}' int.int_inverted_index"),
)?,
})),
(Some(l), None) => Ok(Some(l.clone())),
(None, Some(r)) => Ok(Some(r.clone())),
(None, None) => Ok(None),
}
}
fn merge_bool_override(
key: &str,
left: Option<&BoolValueType>,
right: Option<&BoolValueType>,
) -> Result<Option<BoolValueType>, SchemaError> {
match (left, right) {
(Some(l), Some(r)) => Ok(Some(BoolValueType {
bool_inverted_index: Self::merge_index_or_error(
l.bool_inverted_index.as_ref(),
r.bool_inverted_index.as_ref(),
&format!("key '{key}' bool.bool_inverted_index"),
)?,
})),
(Some(l), None) => Ok(Some(l.clone())),
(None, Some(r)) => Ok(Some(r.clone())),
(None, None) => Ok(None),
}
}
fn merge_float_list_override(
key: &str,
left: Option<&FloatListValueType>,
right: Option<&FloatListValueType>,
) -> Result<Option<FloatListValueType>, SchemaError> {
match (left, right) {
(Some(l), Some(r)) => Ok(Some(FloatListValueType {
vector_index: Self::merge_index_or_error(
l.vector_index.as_ref(),
r.vector_index.as_ref(),
&format!("key '{key}' float_list.vector_index"),
)?,
})),
(Some(l), None) => Ok(Some(l.clone())),
(None, Some(r)) => Ok(Some(r.clone())),
(None, None) => Ok(None),
}
}
fn merge_sparse_vector_override(
key: &str,
left: Option<&SparseVectorValueType>,
right: Option<&SparseVectorValueType>,
) -> Result<Option<SparseVectorValueType>, SchemaError> {
match (left, right) {
(Some(l), Some(r)) => Ok(Some(SparseVectorValueType {
sparse_vector_index: Self::merge_index_or_error(
l.sparse_vector_index.as_ref(),
r.sparse_vector_index.as_ref(),
&format!("key '{key}' sparse_vector.sparse_vector_index"),
)?,
})),
(Some(l), None) => Ok(Some(l.clone())),
(None, Some(r)) => Ok(Some(r.clone())),
(None, None) => Ok(None),
}
}
fn merge_index_or_error<T: Clone + PartialEq>(
left: Option<&T>,
right: Option<&T>,
context: &str,
) -> Result<Option<T>, SchemaError> {
match (left, right) {
(Some(l), Some(r)) => {
if l == r {
Ok(Some(l.clone()))
} else {
Err(SchemaError::ConfigurationConflict {
context: context.to_string(),
})
}
}
(Some(l), None) => Ok(Some(l.clone())),
(None, Some(r)) => Ok(Some(r.clone())),
(None, None) => Ok(None),
}
}
fn merge_value_types(
default: &ValueTypes,
user: &ValueTypes,
knn_index: KnnIndex,
) -> Result<ValueTypes, SchemaError> {
let float_list = Self::merge_float_list_type(
default.float_list.as_ref(),
user.float_list.as_ref(),
knn_index,
)?;
if let Some(ref fl) = float_list {
Self::validate_float_list_value_type(fl)?;
}
Ok(ValueTypes {
string: Self::merge_string_type(default.string.as_ref(), user.string.as_ref())?,
float: Self::merge_float_type(default.float.as_ref(), user.float.as_ref())?,
int: Self::merge_int_type(default.int.as_ref(), user.int.as_ref())?,
boolean: Self::merge_bool_type(default.boolean.as_ref(), user.boolean.as_ref())?,
float_list,
sparse_vector: Self::merge_sparse_vector_type(
default.sparse_vector.as_ref(),
user.sparse_vector.as_ref(),
)?,
})
}
fn merge_string_type(
default: Option<&StringValueType>,
user: Option<&StringValueType>,
) -> Result<Option<StringValueType>, SchemaError> {
match (default, user) {
(Some(default), Some(user)) => Ok(Some(StringValueType {
string_inverted_index: Self::merge_string_inverted_index_type(
default.string_inverted_index.as_ref(),
user.string_inverted_index.as_ref(),
)?,
fts_index: Self::merge_fts_index_type(
default.fts_index.as_ref(),
user.fts_index.as_ref(),
)?,
})),
(Some(default), None) => Ok(Some(default.clone())),
(None, Some(user)) => Ok(Some(user.clone())),
(None, None) => Ok(None),
}
}
fn merge_float_type(
default: Option<&FloatValueType>,
user: Option<&FloatValueType>,
) -> Result<Option<FloatValueType>, SchemaError> {
match (default, user) {
(Some(default), Some(user)) => Ok(Some(FloatValueType {
float_inverted_index: Self::merge_float_inverted_index_type(
default.float_inverted_index.as_ref(),
user.float_inverted_index.as_ref(),
)?,
})),
(Some(default), None) => Ok(Some(default.clone())),
(None, Some(user)) => Ok(Some(user.clone())),
(None, None) => Ok(None),
}
}
fn merge_int_type(
default: Option<&IntValueType>,
user: Option<&IntValueType>,
) -> Result<Option<IntValueType>, SchemaError> {
match (default, user) {
(Some(default), Some(user)) => Ok(Some(IntValueType {
int_inverted_index: Self::merge_int_inverted_index_type(
default.int_inverted_index.as_ref(),
user.int_inverted_index.as_ref(),
)?,
})),
(Some(default), None) => Ok(Some(default.clone())),
(None, Some(user)) => Ok(Some(user.clone())),
(None, None) => Ok(None),
}
}
fn merge_bool_type(
default: Option<&BoolValueType>,
user: Option<&BoolValueType>,
) -> Result<Option<BoolValueType>, SchemaError> {
match (default, user) {
(Some(default), Some(user)) => Ok(Some(BoolValueType {
bool_inverted_index: Self::merge_bool_inverted_index_type(
default.bool_inverted_index.as_ref(),
user.bool_inverted_index.as_ref(),
)?,
})),
(Some(default), None) => Ok(Some(default.clone())),
(None, Some(user)) => Ok(Some(user.clone())),
(None, None) => Ok(None),
}
}
fn merge_float_list_type(
default: Option<&FloatListValueType>,
user: Option<&FloatListValueType>,
knn_index: KnnIndex,
) -> Result<Option<FloatListValueType>, SchemaError> {
match (default, user) {
(Some(default), Some(user)) => Ok(Some(FloatListValueType {
vector_index: Self::merge_vector_index_type(
default.vector_index.as_ref(),
user.vector_index.as_ref(),
knn_index,
)?,
})),
(Some(default), None) => Ok(Some(default.clone())),
(None, Some(user)) => Ok(Some(user.clone())),
(None, None) => Ok(None),
}
}
fn merge_sparse_vector_type(
default: Option<&SparseVectorValueType>,
user: Option<&SparseVectorValueType>,
) -> Result<Option<SparseVectorValueType>, SchemaError> {
match (default, user) {
(Some(default), Some(user)) => Ok(Some(SparseVectorValueType {
sparse_vector_index: Self::merge_sparse_vector_index_type(
default.sparse_vector_index.as_ref(),
user.sparse_vector_index.as_ref(),
)?,
})),
(Some(default), None) => Ok(Some(default.clone())),
(None, Some(user)) => Ok(Some(user.clone())),
(None, None) => Ok(None),
}
}
fn merge_string_inverted_index_type(
default: Option<&StringInvertedIndexType>,
user: Option<&StringInvertedIndexType>,
) -> Result<Option<StringInvertedIndexType>, SchemaError> {
match (default, user) {
(Some(_default), Some(user)) => {
Ok(Some(StringInvertedIndexType {
enabled: user.enabled, config: user.config.clone(), }))
}
(Some(default), None) => Ok(Some(default.clone())),
(None, Some(user)) => Ok(Some(user.clone())),
(None, None) => Ok(None),
}
}
fn merge_fts_index_type(
default: Option<&FtsIndexType>,
user: Option<&FtsIndexType>,
) -> Result<Option<FtsIndexType>, SchemaError> {
match (default, user) {
(Some(_default), Some(user)) => Ok(Some(FtsIndexType {
enabled: user.enabled,
config: user.config.clone(),
})),
(Some(default), None) => Ok(Some(default.clone())),
(None, Some(user)) => Ok(Some(user.clone())),
(None, None) => Ok(None),
}
}
fn merge_float_inverted_index_type(
default: Option<&FloatInvertedIndexType>,
user: Option<&FloatInvertedIndexType>,
) -> Result<Option<FloatInvertedIndexType>, SchemaError> {
match (default, user) {
(Some(_default), Some(user)) => Ok(Some(FloatInvertedIndexType {
enabled: user.enabled,
config: user.config.clone(),
})),
(Some(default), None) => Ok(Some(default.clone())),
(None, Some(user)) => Ok(Some(user.clone())),
(None, None) => Ok(None),
}
}
fn merge_int_inverted_index_type(
default: Option<&IntInvertedIndexType>,
user: Option<&IntInvertedIndexType>,
) -> Result<Option<IntInvertedIndexType>, SchemaError> {
match (default, user) {
(Some(_default), Some(user)) => Ok(Some(IntInvertedIndexType {
enabled: user.enabled,
config: user.config.clone(),
})),
(Some(default), None) => Ok(Some(default.clone())),
(None, Some(user)) => Ok(Some(user.clone())),
(None, None) => Ok(None),
}
}
fn merge_bool_inverted_index_type(
default: Option<&BoolInvertedIndexType>,
user: Option<&BoolInvertedIndexType>,
) -> Result<Option<BoolInvertedIndexType>, SchemaError> {
match (default, user) {
(Some(_default), Some(user)) => Ok(Some(BoolInvertedIndexType {
enabled: user.enabled,
config: user.config.clone(),
})),
(Some(default), None) => Ok(Some(default.clone())),
(None, Some(user)) => Ok(Some(user.clone())),
(None, None) => Ok(None),
}
}
fn merge_vector_index_type(
default: Option<&VectorIndexType>,
user: Option<&VectorIndexType>,
knn_index: KnnIndex,
) -> Result<Option<VectorIndexType>, SchemaError> {
match (default, user) {
(Some(default), Some(user)) => Ok(Some(VectorIndexType {
enabled: user.enabled,
config: Self::merge_vector_index_config(&default.config, &user.config, knn_index)?,
})),
(Some(default), None) => Ok(Some(default.clone())),
(None, Some(user)) => Ok(Some(user.clone())),
(None, None) => Ok(None),
}
}
fn merge_sparse_vector_index_type(
default: Option<&SparseVectorIndexType>,
user: Option<&SparseVectorIndexType>,
) -> Result<Option<SparseVectorIndexType>, SchemaError> {
match (default, user) {
(Some(default), Some(user)) => Ok(Some(SparseVectorIndexType {
enabled: user.enabled,
config: Self::merge_sparse_vector_index_config(&default.config, &user.config),
})),
(Some(default), None) => Ok(Some(default.clone())),
(None, Some(user)) => Ok(Some(user.clone())),
(None, None) => Ok(None),
}
}
fn validate_float_list_value_type(float_list: &FloatListValueType) -> Result<(), SchemaError> {
if let Some(vector_index) = &float_list.vector_index {
if let Some(hnsw) = &vector_index.config.hnsw {
hnsw.validate().map_err(SchemaError::InvalidHnswConfig)?;
}
if let Some(spann) = &vector_index.config.spann {
spann.validate().map_err(SchemaError::InvalidSpannConfig)?;
}
}
Ok(())
}
fn merge_vector_index_config(
default: &VectorIndexConfig,
user: &VectorIndexConfig,
knn_index: KnnIndex,
) -> Result<VectorIndexConfig, SchemaError> {
match knn_index {
KnnIndex::Hnsw => Ok(VectorIndexConfig {
space: user.space.clone().or(default.space.clone()),
embedding_function: user
.embedding_function
.clone()
.or(default.embedding_function.clone()),
source_key: user.source_key.clone().or(default.source_key.clone()),
hnsw: Self::merge_hnsw_configs(default.hnsw.as_ref(), user.hnsw.as_ref()),
spann: None,
}),
KnnIndex::Spann => Ok(VectorIndexConfig {
space: user.space.clone().or(default.space.clone()),
embedding_function: user
.embedding_function
.clone()
.or(default.embedding_function.clone()),
source_key: user.source_key.clone().or(default.source_key.clone()),
hnsw: None,
spann: Self::merge_spann_configs(default.spann.as_ref(), user.spann.as_ref())?,
}),
}
}
fn merge_sparse_vector_index_config(
default: &SparseVectorIndexConfig,
user: &SparseVectorIndexConfig,
) -> SparseVectorIndexConfig {
SparseVectorIndexConfig {
embedding_function: user
.embedding_function
.clone()
.or(default.embedding_function.clone()),
source_key: user.source_key.clone().or(default.source_key.clone()),
bm25: user.bm25.or(default.bm25),
}
}
fn merge_hnsw_configs(
default_hnsw: Option<&HnswIndexConfig>,
user_hnsw: Option<&HnswIndexConfig>,
) -> Option<HnswIndexConfig> {
match (default_hnsw, user_hnsw) {
(Some(default), Some(user)) => Some(HnswIndexConfig {
ef_construction: user.ef_construction.or(default.ef_construction),
max_neighbors: user.max_neighbors.or(default.max_neighbors),
ef_search: user.ef_search.or(default.ef_search),
num_threads: user.num_threads.or(default.num_threads),
batch_size: user.batch_size.or(default.batch_size),
sync_threshold: user.sync_threshold.or(default.sync_threshold),
resize_factor: user.resize_factor.or(default.resize_factor),
}),
(Some(default), None) => Some(default.clone()),
(None, Some(user)) => Some(user.clone()),
(None, None) => None,
}
}
fn merge_spann_configs(
default_spann: Option<&SpannIndexConfig>,
user_spann: Option<&SpannIndexConfig>,
) -> Result<Option<SpannIndexConfig>, SchemaError> {
match (default_spann, user_spann) {
(Some(default), Some(user)) => {
if !matches!(user.quantize, Quantization::None)
|| !matches!(default.quantize, Quantization::None)
{
return Err(SchemaError::InvalidUserInput {
reason: "quantize field cannot be set in user schema. Quantization can only be enabled via frontend configuration.".to_string(),
});
}
Ok(Some(SpannIndexConfig {
search_nprobe: user.search_nprobe.or(default.search_nprobe),
search_rng_factor: user.search_rng_factor.or(default.search_rng_factor),
search_rng_epsilon: user.search_rng_epsilon.or(default.search_rng_epsilon),
nreplica_count: user.nreplica_count.or(default.nreplica_count),
write_rng_factor: user.write_rng_factor.or(default.write_rng_factor),
write_rng_epsilon: user.write_rng_epsilon.or(default.write_rng_epsilon),
split_threshold: user.split_threshold.or(default.split_threshold),
num_samples_kmeans: user.num_samples_kmeans.or(default.num_samples_kmeans),
initial_lambda: user.initial_lambda.or(default.initial_lambda),
reassign_neighbor_count: user
.reassign_neighbor_count
.or(default.reassign_neighbor_count),
merge_threshold: user.merge_threshold.or(default.merge_threshold),
num_centers_to_merge_to: user
.num_centers_to_merge_to
.or(default.num_centers_to_merge_to),
write_nprobe: user.write_nprobe.or(default.write_nprobe),
ef_construction: user.ef_construction.or(default.ef_construction),
ef_search: user.ef_search.or(default.ef_search),
max_neighbors: user.max_neighbors.or(default.max_neighbors),
center_drift_threshold: user
.center_drift_threshold
.or(default.center_drift_threshold),
quantize: Quantization::None, }))
}
(Some(default), None) => {
if !matches!(default.quantize, Quantization::None) {
return Err(SchemaError::InvalidUserInput {
reason: "quantize field cannot be set in default schema. Quantization can only be enabled via frontend configuration.".to_string(),
});
}
Ok(Some(default.clone()))
}
(None, Some(user)) => {
if !matches!(user.quantize, Quantization::None) {
return Err(SchemaError::InvalidUserInput {
reason: "quantize field cannot be set in user schema. Quantization can only be enabled via frontend configuration.".to_string(),
});
}
Ok(Some(user.clone()))
}
(None, None) => Ok(None),
}
}
pub fn reconcile_with_collection_config(
schema: &Schema,
collection_config: &InternalCollectionConfiguration,
default_knn_index: KnnIndex,
) -> Result<Schema, SchemaError> {
if collection_config.is_default() {
if schema.is_default() {
let mut new_schema = Schema::new_default(default_knn_index);
if collection_config.embedding_function.is_some() {
if let Some(float_list) = &mut new_schema.defaults.float_list {
if let Some(vector_index) = &mut float_list.vector_index {
vector_index.config.embedding_function =
collection_config.embedding_function.clone();
}
}
if let Some(embedding_types) = new_schema.keys.get_mut(EMBEDDING_KEY) {
if let Some(float_list) = &mut embedding_types.float_list {
if let Some(vector_index) = &mut float_list.vector_index {
vector_index.config.embedding_function =
collection_config.embedding_function.clone();
}
}
}
}
return Ok(new_schema);
} else {
return Ok(schema.clone());
}
}
Self::try_from(collection_config)
}
pub fn reconcile_schema_and_config(
schema: Option<&Schema>,
configuration: Option<&InternalCollectionConfiguration>,
knn_index: KnnIndex,
) -> Result<Schema, SchemaError> {
if let (Some(user_schema), Some(config)) = (schema, configuration) {
if !user_schema.is_default() && !config.is_default() {
return Err(SchemaError::ConfigAndSchemaConflict);
}
}
let reconciled_schema = Self::reconcile_with_defaults(schema, knn_index)?;
if let Some(config) = configuration {
Self::reconcile_with_collection_config(&reconciled_schema, config, knn_index)
} else {
Ok(reconciled_schema)
}
}
pub fn default_with_embedding_function(
embedding_function: EmbeddingFunctionConfiguration,
) -> Schema {
let mut schema = Schema::new_default(KnnIndex::Spann);
if let Some(float_list) = &mut schema.defaults.float_list {
if let Some(vector_index) = &mut float_list.vector_index {
vector_index.config.embedding_function = Some(embedding_function.clone());
}
}
if let Some(embedding_types) = schema.keys.get_mut(EMBEDDING_KEY) {
if let Some(float_list) = &mut embedding_types.float_list {
if let Some(vector_index) = &mut float_list.vector_index {
vector_index.config.embedding_function = Some(embedding_function);
}
}
}
schema
}
pub fn is_default(&self) -> bool {
if !Self::is_value_types_default(&self.defaults) {
return false;
}
for key in self.keys.keys() {
if key != EMBEDDING_KEY && key != DOCUMENT_KEY {
return false;
}
}
if let Some(embedding_value) = self.keys.get(EMBEDDING_KEY) {
if !Self::is_embedding_value_types_default(embedding_value) {
return false;
}
}
if let Some(document_value) = self.keys.get(DOCUMENT_KEY) {
if !Self::is_document_value_types_default(document_value) {
return false;
}
}
if self.cmek.is_some() {
return false;
}
true
}
fn is_value_types_default(value_types: &ValueTypes) -> bool {
if let Some(string) = &value_types.string {
if let Some(string_inverted) = &string.string_inverted_index {
if !string_inverted.enabled {
return false;
}
}
if let Some(fts) = &string.fts_index {
if fts.enabled {
return false;
}
}
}
if let Some(float) = &value_types.float {
if let Some(float_inverted) = &float.float_inverted_index {
if !float_inverted.enabled {
return false;
}
}
}
if let Some(int) = &value_types.int {
if let Some(int_inverted) = &int.int_inverted_index {
if !int_inverted.enabled {
return false;
}
}
}
if let Some(boolean) = &value_types.boolean {
if let Some(bool_inverted) = &boolean.bool_inverted_index {
if !bool_inverted.enabled {
return false;
}
}
}
if let Some(float_list) = &value_types.float_list {
if let Some(vector_index) = &float_list.vector_index {
if vector_index.enabled {
return false;
}
if !is_embedding_function_default(&vector_index.config.embedding_function) {
return false;
}
if !is_space_default(&vector_index.config.space) {
return false;
}
if vector_index.config.source_key.is_some() {
return false;
}
match (&vector_index.config.hnsw, &vector_index.config.spann) {
(Some(hnsw_config), None) => {
if !hnsw_config.is_default() {
return false;
}
}
(None, Some(spann_config)) => {
if !spann_config.is_default() {
return false;
}
}
(Some(_), Some(_)) => return false, (None, None) => {}
}
}
}
if let Some(sparse_vector) = &value_types.sparse_vector {
if let Some(sparse_index) = &sparse_vector.sparse_vector_index {
if sparse_index.enabled {
return false;
}
if !is_embedding_function_default(&sparse_index.config.embedding_function) {
return false;
}
if sparse_index.config.source_key.is_some() {
return false;
}
if let Some(bm25) = &sparse_index.config.bm25 {
if bm25 != &false {
return false;
}
}
}
}
true
}
fn is_embedding_value_types_default(value_types: &ValueTypes) -> bool {
if value_types.string.is_some()
|| value_types.float.is_some()
|| value_types.int.is_some()
|| value_types.boolean.is_some()
|| value_types.sparse_vector.is_some()
{
return false;
}
if let Some(float_list) = &value_types.float_list {
if let Some(vector_index) = &float_list.vector_index {
if !vector_index.enabled {
return false;
}
if !is_space_default(&vector_index.config.space) {
return false;
}
if !is_embedding_function_default(&vector_index.config.embedding_function) {
return false;
}
if vector_index.config.source_key.as_deref() != Some(DOCUMENT_KEY) {
return false;
}
match (&vector_index.config.hnsw, &vector_index.config.spann) {
(Some(hnsw_config), None) => {
if !hnsw_config.is_default() {
return false;
}
}
(None, Some(spann_config)) => {
if !spann_config.is_default() {
return false;
}
}
(Some(_), Some(_)) => return false, (None, None) => {}
}
}
}
true
}
fn is_document_value_types_default(value_types: &ValueTypes) -> bool {
if value_types.float_list.is_some()
|| value_types.float.is_some()
|| value_types.int.is_some()
|| value_types.boolean.is_some()
|| value_types.sparse_vector.is_some()
{
return false;
}
if let Some(string) = &value_types.string {
if let Some(fts) = &string.fts_index {
if !fts.enabled {
return false;
}
}
if let Some(string_inverted) = &string.string_inverted_index {
if string_inverted.enabled {
return false;
}
}
}
true
}
pub fn is_metadata_type_index_enabled(
&self,
key: &str,
value_type: MetadataValueType,
) -> Result<bool, SchemaError> {
let v_type = self.keys.get(key).unwrap_or(&self.defaults);
match value_type {
MetadataValueType::Bool => match &v_type.boolean {
Some(bool_type) => match &bool_type.bool_inverted_index {
Some(bool_inverted_index) => Ok(bool_inverted_index.enabled),
None => Err(SchemaError::MissingIndexConfiguration {
key: key.to_string(),
value_type: "bool".to_string(),
}),
},
None => match &self.defaults.boolean {
Some(bool_type) => match &bool_type.bool_inverted_index {
Some(bool_inverted_index) => Ok(bool_inverted_index.enabled),
None => Err(SchemaError::MissingIndexConfiguration {
key: key.to_string(),
value_type: "bool".to_string(),
}),
},
None => Err(SchemaError::MissingIndexConfiguration {
key: key.to_string(),
value_type: "bool".to_string(),
}),
},
},
MetadataValueType::Int => match &v_type.int {
Some(int_type) => match &int_type.int_inverted_index {
Some(int_inverted_index) => Ok(int_inverted_index.enabled),
None => Err(SchemaError::MissingIndexConfiguration {
key: key.to_string(),
value_type: "int".to_string(),
}),
},
None => match &self.defaults.int {
Some(int_type) => match &int_type.int_inverted_index {
Some(int_inverted_index) => Ok(int_inverted_index.enabled),
None => Err(SchemaError::MissingIndexConfiguration {
key: key.to_string(),
value_type: "int".to_string(),
}),
},
None => Err(SchemaError::MissingIndexConfiguration {
key: key.to_string(),
value_type: "int".to_string(),
}),
},
},
MetadataValueType::Float => match &v_type.float {
Some(float_type) => match &float_type.float_inverted_index {
Some(float_inverted_index) => Ok(float_inverted_index.enabled),
None => Err(SchemaError::MissingIndexConfiguration {
key: key.to_string(),
value_type: "float".to_string(),
}),
},
None => match &self.defaults.float {
Some(float_type) => match &float_type.float_inverted_index {
Some(float_inverted_index) => Ok(float_inverted_index.enabled),
None => Err(SchemaError::MissingIndexConfiguration {
key: key.to_string(),
value_type: "float".to_string(),
}),
},
None => Err(SchemaError::MissingIndexConfiguration {
key: key.to_string(),
value_type: "float".to_string(),
}),
},
},
MetadataValueType::Str => match &v_type.string {
Some(string_type) => match &string_type.string_inverted_index {
Some(string_inverted_index) => Ok(string_inverted_index.enabled),
None => Err(SchemaError::MissingIndexConfiguration {
key: key.to_string(),
value_type: "string".to_string(),
}),
},
None => match &self.defaults.string {
Some(string_type) => match &string_type.string_inverted_index {
Some(string_inverted_index) => Ok(string_inverted_index.enabled),
None => Err(SchemaError::MissingIndexConfiguration {
key: key.to_string(),
value_type: "string".to_string(),
}),
},
None => Err(SchemaError::MissingIndexConfiguration {
key: key.to_string(),
value_type: "string".to_string(),
}),
},
},
MetadataValueType::SparseVector => match &v_type.sparse_vector {
Some(sparse_vector_type) => match &sparse_vector_type.sparse_vector_index {
Some(sparse_vector_index) => Ok(sparse_vector_index.enabled),
None => Err(SchemaError::MissingIndexConfiguration {
key: key.to_string(),
value_type: "sparse_vector".to_string(),
}),
},
None => match &self.defaults.sparse_vector {
Some(sparse_vector_type) => match &sparse_vector_type.sparse_vector_index {
Some(sparse_vector_index) => Ok(sparse_vector_index.enabled),
None => Err(SchemaError::MissingIndexConfiguration {
key: key.to_string(),
value_type: "sparse_vector".to_string(),
}),
},
None => Err(SchemaError::MissingIndexConfiguration {
key: key.to_string(),
value_type: "sparse_vector".to_string(),
}),
},
},
MetadataValueType::BoolArray => {
self.is_metadata_type_index_enabled(key, MetadataValueType::Bool)
}
MetadataValueType::IntArray => {
self.is_metadata_type_index_enabled(key, MetadataValueType::Int)
}
MetadataValueType::FloatArray => {
self.is_metadata_type_index_enabled(key, MetadataValueType::Float)
}
MetadataValueType::StringArray => {
self.is_metadata_type_index_enabled(key, MetadataValueType::Str)
}
}
}
pub fn is_metadata_where_indexing_enabled(
&self,
where_clause: &Where,
) -> Result<(), FilterValidationError> {
match where_clause {
Where::Composite(composite) => {
for child in &composite.children {
self.is_metadata_where_indexing_enabled(child)?;
}
Ok(())
}
Where::Document(_) => {
if !self.is_fts_enabled() {
return Err(FilterValidationError::FtsDisabled);
}
Ok(())
}
Where::Metadata(expression) => {
let value_type = match &expression.comparison {
MetadataComparison::Primitive(_, value) => value.value_type(),
MetadataComparison::Set(_, set_value) => set_value.value_type(),
MetadataComparison::ArrayContains(_, value) => value.value_type(),
};
let is_enabled = self
.is_metadata_type_index_enabled(expression.key.as_str(), value_type)
.map_err(FilterValidationError::Schema)?;
if !is_enabled {
return Err(FilterValidationError::IndexingDisabled {
key: expression.key.clone(),
value_type,
});
}
Ok(())
}
}
}
pub fn is_knn_key_indexing_enabled(
&self,
key: &str,
query: &QueryVector,
) -> Result<(), FilterValidationError> {
match query {
QueryVector::Sparse(_) => {
let is_enabled = self
.is_metadata_type_index_enabled(key, MetadataValueType::SparseVector)
.map_err(FilterValidationError::Schema)?;
if !is_enabled {
return Err(FilterValidationError::IndexingDisabled {
key: key.to_string(),
value_type: MetadataValueType::SparseVector,
});
}
Ok(())
}
QueryVector::Dense(_) => {
Ok(())
}
}
}
pub fn ensure_key_from_metadata(&mut self, key: &str, value_type: MetadataValueType) -> bool {
if key.starts_with(CHROMA_KEY) {
return false;
}
let value_types = self.keys.entry(key.to_string()).or_default();
match value_type {
MetadataValueType::Bool => {
if value_types.boolean.is_none() {
value_types.boolean = self.defaults.boolean.clone();
return true;
}
}
MetadataValueType::Int => {
if value_types.int.is_none() {
value_types.int = self.defaults.int.clone();
return true;
}
}
MetadataValueType::Float => {
if value_types.float.is_none() {
value_types.float = self.defaults.float.clone();
return true;
}
}
MetadataValueType::Str => {
if value_types.string.is_none() {
value_types.string = self.defaults.string.clone();
return true;
}
}
MetadataValueType::SparseVector => {
if value_types.sparse_vector.is_none() {
value_types.sparse_vector = self.defaults.sparse_vector.clone();
return true;
}
}
MetadataValueType::BoolArray => {
if value_types.boolean.is_none() {
value_types.boolean = self.defaults.boolean.clone();
return true;
}
}
MetadataValueType::IntArray => {
if value_types.int.is_none() {
value_types.int = self.defaults.int.clone();
return true;
}
}
MetadataValueType::FloatArray => {
if value_types.float.is_none() {
value_types.float = self.defaults.float.clone();
return true;
}
}
MetadataValueType::StringArray => {
if value_types.string.is_none() {
value_types.string = self.defaults.string.clone();
return true;
}
}
}
false
}
pub fn create_index(
mut self,
key: Option<&str>,
config: IndexConfig,
) -> Result<Self, SchemaBuilderError> {
match &config {
IndexConfig::Vector(cfg) => {
if let Some(k) = key {
return Err(SchemaBuilderError::VectorIndexMustBeGlobal { key: k.to_string() });
}
self._set_vector_index_config_builder(cfg.clone());
return Ok(self);
}
IndexConfig::Fts(_) => {
if key != Some(DOCUMENT_KEY) {
return Err(SchemaBuilderError::FtsIndexOnlyOnDocument);
}
}
IndexConfig::SparseVector(_) => {
if key.is_none() {
return Err(SchemaBuilderError::SparseVectorRequiresKey);
}
}
_ => {}
}
if let Some(k) = key {
if k == EMBEDDING_KEY {
return Err(SchemaBuilderError::SpecialKeyModificationNotAllowed {
key: k.to_string(),
});
}
if k == DOCUMENT_KEY && !matches!(config, IndexConfig::Fts(_)) {
return Err(SchemaBuilderError::SpecialKeyModificationNotAllowed {
key: k.to_string(),
});
}
if k.starts_with('#') && k != DOCUMENT_KEY {
return Err(SchemaBuilderError::ReservedKeyPrefix { key: k.to_string() });
}
}
match key {
Some(k) => self._set_index_for_key_builder(k, config, true)?,
None => self._set_index_in_defaults_builder(config, true)?,
}
Ok(self)
}
pub fn delete_index(
mut self,
key: Option<&str>,
config: IndexConfig,
) -> Result<Self, SchemaBuilderError> {
match &config {
IndexConfig::Vector(_) => {
return Err(SchemaBuilderError::VectorIndexDeletionNotSupported);
}
IndexConfig::Fts(_) => {
if key != Some(DOCUMENT_KEY) {
return Err(SchemaBuilderError::FtsIndexDeletionOnlyOnDocument);
}
}
IndexConfig::SparseVector(_) => {
return Err(SchemaBuilderError::SparseVectorIndexDeletionNotSupported);
}
_ => {}
}
if let Some(k) = key {
if k == EMBEDDING_KEY {
return Err(SchemaBuilderError::SpecialKeyModificationNotAllowed {
key: k.to_string(),
});
}
if k == DOCUMENT_KEY && !matches!(config, IndexConfig::Fts(_)) {
return Err(SchemaBuilderError::SpecialKeyModificationNotAllowed {
key: k.to_string(),
});
}
if k.starts_with('#') && k != DOCUMENT_KEY {
return Err(SchemaBuilderError::ReservedKeyPrefix { key: k.to_string() });
}
}
match key {
Some(k) => self._set_index_for_key_builder(k, config, false)?,
None => self._set_index_in_defaults_builder(config, false)?,
}
Ok(self)
}
pub fn with_cmek(mut self, cmek: Cmek) -> Self {
self.cmek = Some(cmek);
self
}
fn _set_vector_index_config_builder(&mut self, config: VectorIndexConfig) {
if let Some(float_list) = &mut self.defaults.float_list {
if let Some(vector_index) = &mut float_list.vector_index {
vector_index.config = config.clone();
}
}
if let Some(embedding_types) = self.keys.get_mut(EMBEDDING_KEY) {
if let Some(float_list) = &mut embedding_types.float_list {
if let Some(vector_index) = &mut float_list.vector_index {
let mut updated_config = config;
updated_config.source_key = Some(DOCUMENT_KEY.to_string());
vector_index.config = updated_config;
}
}
}
}
fn _set_fts_index_config_builder(&mut self, config: FtsIndexConfig) {
if let Some(string) = &mut self.defaults.string {
if let Some(fts_index) = &mut string.fts_index {
fts_index.config = config.clone();
}
}
if let Some(document_types) = self.keys.get_mut(DOCUMENT_KEY) {
if let Some(string) = &mut document_types.string {
if let Some(fts_index) = &mut string.fts_index {
fts_index.config = config;
}
}
}
}
fn _set_index_for_key_builder(
&mut self,
key: &str,
config: IndexConfig,
enabled: bool,
) -> Result<(), SchemaBuilderError> {
if enabled && matches!(config, IndexConfig::SparseVector(_)) {
let existing_key = self
.keys
.iter()
.find(|(k, v)| {
k.as_str() != key
&& v.sparse_vector
.as_ref()
.and_then(|sv| sv.sparse_vector_index.as_ref())
.map(|idx| idx.enabled)
.unwrap_or(false)
})
.map(|(k, _)| k.clone());
if let Some(existing_key) = existing_key {
return Err(SchemaBuilderError::MultipleSparseVectorIndexes { existing_key });
}
}
let value_types = self.keys.entry(key.to_string()).or_default();
match config {
IndexConfig::Vector(_) => {
return Err(SchemaBuilderError::VectorIndexMustBeGlobal {
key: key.to_string(),
});
}
IndexConfig::Fts(cfg) => {
if let Some(string) = value_types.string.as_mut() {
if let Some(fts_index) = string.fts_index.as_mut() {
fts_index.enabled = enabled;
fts_index.config = cfg;
}
}
}
IndexConfig::SparseVector(cfg) => {
value_types.sparse_vector = Some(SparseVectorValueType {
sparse_vector_index: Some(SparseVectorIndexType {
enabled,
config: cfg,
}),
});
}
IndexConfig::StringInverted(cfg) => {
if value_types.string.is_none() {
value_types.string = Some(StringValueType {
fts_index: None,
string_inverted_index: None,
});
}
if let Some(string) = &mut value_types.string {
string.string_inverted_index = Some(StringInvertedIndexType {
enabled,
config: cfg,
});
}
}
IndexConfig::IntInverted(cfg) => {
value_types.int = Some(IntValueType {
int_inverted_index: Some(IntInvertedIndexType {
enabled,
config: cfg,
}),
});
}
IndexConfig::FloatInverted(cfg) => {
value_types.float = Some(FloatValueType {
float_inverted_index: Some(FloatInvertedIndexType {
enabled,
config: cfg,
}),
});
}
IndexConfig::BoolInverted(cfg) => {
value_types.boolean = Some(BoolValueType {
bool_inverted_index: Some(BoolInvertedIndexType {
enabled,
config: cfg,
}),
});
}
}
Ok(())
}
fn _set_index_in_defaults_builder(
&mut self,
config: IndexConfig,
enabled: bool,
) -> Result<(), SchemaBuilderError> {
match config {
IndexConfig::Vector(_) => {
return Err(SchemaBuilderError::VectorIndexMustBeGlobal {
key: "defaults".to_string(),
});
}
IndexConfig::Fts(_) => {
return Err(SchemaBuilderError::FtsIndexOnlyOnDocument);
}
IndexConfig::SparseVector(cfg) => {
self.defaults.sparse_vector = Some(SparseVectorValueType {
sparse_vector_index: Some(SparseVectorIndexType {
enabled,
config: cfg,
}),
});
}
IndexConfig::StringInverted(cfg) => {
if self.defaults.string.is_none() {
self.defaults.string = Some(StringValueType {
fts_index: None,
string_inverted_index: None,
});
}
if let Some(string) = &mut self.defaults.string {
string.string_inverted_index = Some(StringInvertedIndexType {
enabled,
config: cfg,
});
}
}
IndexConfig::IntInverted(cfg) => {
self.defaults.int = Some(IntValueType {
int_inverted_index: Some(IntInvertedIndexType {
enabled,
config: cfg,
}),
});
}
IndexConfig::FloatInverted(cfg) => {
self.defaults.float = Some(FloatValueType {
float_inverted_index: Some(FloatInvertedIndexType {
enabled,
config: cfg,
}),
});
}
IndexConfig::BoolInverted(cfg) => {
self.defaults.boolean = Some(BoolValueType {
bool_inverted_index: Some(BoolInvertedIndexType {
enabled,
config: cfg,
}),
});
}
}
Ok(())
}
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
#[serde(deny_unknown_fields)]
pub struct VectorIndexConfig {
#[serde(skip_serializing_if = "Option::is_none")]
pub space: Option<Space>,
#[serde(skip_serializing_if = "Option::is_none")]
pub embedding_function: Option<EmbeddingFunctionConfiguration>,
#[serde(skip_serializing_if = "Option::is_none")]
pub source_key: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub hnsw: Option<HnswIndexConfig>,
#[serde(skip_serializing_if = "Option::is_none")]
pub spann: Option<SpannIndexConfig>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Validate, Default)]
#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
#[serde(deny_unknown_fields)]
pub struct HnswIndexConfig {
#[serde(skip_serializing_if = "Option::is_none")]
pub ef_construction: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
pub max_neighbors: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
pub ef_search: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
pub num_threads: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
#[validate(range(min = 2))]
pub batch_size: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
#[validate(range(min = 2))]
pub sync_threshold: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
pub resize_factor: Option<f64>,
}
impl HnswIndexConfig {
pub fn is_default(&self) -> bool {
if let Some(ef_construction) = self.ef_construction {
if ef_construction != default_construction_ef() {
return false;
}
}
if let Some(max_neighbors) = self.max_neighbors {
if max_neighbors != default_m() {
return false;
}
}
if let Some(ef_search) = self.ef_search {
if ef_search != default_search_ef() {
return false;
}
}
if let Some(batch_size) = self.batch_size {
if batch_size != default_batch_size() {
return false;
}
}
if let Some(sync_threshold) = self.sync_threshold {
if sync_threshold != default_sync_threshold() {
return false;
}
}
if let Some(resize_factor) = self.resize_factor {
if resize_factor != default_resize_factor() {
return false;
}
}
true
}
}
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
#[serde(rename_all = "snake_case")]
pub enum Quantization {
#[default]
None,
FourBitRabitQWithUSearch,
}
fn is_default_quantization(v: &Quantization) -> bool {
matches!(v, Quantization::None)
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Validate, Default)]
#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
#[serde(deny_unknown_fields)]
pub struct SpannIndexConfig {
#[serde(skip_serializing_if = "Option::is_none")]
#[validate(range(max = 128))]
pub search_nprobe: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
#[validate(range(min = 1.0, max = 1.0))]
pub search_rng_factor: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
#[validate(range(min = 5.0, max = 10.0))]
pub search_rng_epsilon: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
#[validate(range(max = 8))]
pub nreplica_count: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
#[validate(range(min = 1.0, max = 1.0))]
pub write_rng_factor: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
#[validate(range(min = 5.0, max = 10.0))]
pub write_rng_epsilon: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
#[validate(range(min = 50, max = 200))]
pub split_threshold: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
#[validate(range(max = 1000))]
pub num_samples_kmeans: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
#[validate(range(min = 100.0, max = 100.0))]
pub initial_lambda: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
#[validate(range(max = 64))]
pub reassign_neighbor_count: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
#[validate(range(min = 25, max = 100))]
pub merge_threshold: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
#[validate(range(max = 8))]
pub num_centers_to_merge_to: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
#[validate(range(max = 64))]
pub write_nprobe: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
#[validate(range(max = 200))]
pub ef_construction: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
#[validate(range(max = 200))]
pub ef_search: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
#[validate(range(max = 64))]
pub max_neighbors: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
#[validate(range(min = 0.1, max = 1.0))]
pub center_drift_threshold: Option<f32>,
#[serde(default, skip_serializing_if = "is_default_quantization")]
pub quantize: Quantization,
}
impl SpannIndexConfig {
pub fn is_default(&self) -> bool {
if let Some(search_nprobe) = self.search_nprobe {
if search_nprobe != default_search_nprobe() {
return false;
}
}
if let Some(search_rng_factor) = self.search_rng_factor {
if search_rng_factor != default_search_rng_factor() {
return false;
}
}
if let Some(search_rng_epsilon) = self.search_rng_epsilon {
if search_rng_epsilon != default_search_rng_epsilon() {
return false;
}
}
if let Some(nreplica_count) = self.nreplica_count {
if nreplica_count != default_nreplica_count() {
return false;
}
}
if let Some(write_rng_factor) = self.write_rng_factor {
if write_rng_factor != default_write_rng_factor() {
return false;
}
}
if let Some(write_rng_epsilon) = self.write_rng_epsilon {
if write_rng_epsilon != default_write_rng_epsilon() {
return false;
}
}
if let Some(split_threshold) = self.split_threshold {
if split_threshold != default_split_threshold() {
return false;
}
}
if let Some(num_samples_kmeans) = self.num_samples_kmeans {
if num_samples_kmeans != default_num_samples_kmeans() {
return false;
}
}
if let Some(initial_lambda) = self.initial_lambda {
if initial_lambda != default_initial_lambda() {
return false;
}
}
if let Some(reassign_neighbor_count) = self.reassign_neighbor_count {
if reassign_neighbor_count != default_reassign_neighbor_count() {
return false;
}
}
if let Some(merge_threshold) = self.merge_threshold {
if merge_threshold != default_merge_threshold() {
return false;
}
}
if let Some(num_centers_to_merge_to) = self.num_centers_to_merge_to {
if num_centers_to_merge_to != default_num_centers_to_merge_to() {
return false;
}
}
if let Some(write_nprobe) = self.write_nprobe {
if write_nprobe != default_write_nprobe() {
return false;
}
}
if let Some(ef_construction) = self.ef_construction {
if ef_construction != default_construction_ef_spann() {
return false;
}
}
if let Some(ef_search) = self.ef_search {
if ef_search != default_search_ef_spann() {
return false;
}
}
if let Some(max_neighbors) = self.max_neighbors {
if max_neighbors != default_m_spann() {
return false;
}
}
if let Some(center_drift_threshold) = self.center_drift_threshold {
if center_drift_threshold != default_center_drift_threshold() {
return false;
}
}
if !matches!(self.quantize, Quantization::None) {
return false;
}
true
}
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
#[serde(deny_unknown_fields)]
pub struct SparseVectorIndexConfig {
#[serde(skip_serializing_if = "Option::is_none")]
pub embedding_function: Option<EmbeddingFunctionConfiguration>,
#[serde(skip_serializing_if = "Option::is_none")]
pub source_key: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub bm25: Option<bool>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
#[serde(deny_unknown_fields)]
pub struct FtsIndexConfig {
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
#[serde(deny_unknown_fields)]
pub struct StringInvertedIndexConfig {
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
#[serde(deny_unknown_fields)]
pub struct IntInvertedIndexConfig {
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
#[serde(deny_unknown_fields)]
pub struct FloatInvertedIndexConfig {
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
#[serde(deny_unknown_fields)]
pub struct BoolInvertedIndexConfig {
}
#[derive(Clone, Debug)]
#[allow(clippy::large_enum_variant)]
pub enum IndexConfig {
Vector(VectorIndexConfig),
SparseVector(SparseVectorIndexConfig),
Fts(FtsIndexConfig),
StringInverted(StringInvertedIndexConfig),
IntInverted(IntInvertedIndexConfig),
FloatInverted(FloatInvertedIndexConfig),
BoolInverted(BoolInvertedIndexConfig),
}
impl From<VectorIndexConfig> for IndexConfig {
fn from(config: VectorIndexConfig) -> Self {
IndexConfig::Vector(config)
}
}
impl From<SparseVectorIndexConfig> for IndexConfig {
fn from(config: SparseVectorIndexConfig) -> Self {
IndexConfig::SparseVector(config)
}
}
impl From<FtsIndexConfig> for IndexConfig {
fn from(config: FtsIndexConfig) -> Self {
IndexConfig::Fts(config)
}
}
impl From<StringInvertedIndexConfig> for IndexConfig {
fn from(config: StringInvertedIndexConfig) -> Self {
IndexConfig::StringInverted(config)
}
}
impl From<IntInvertedIndexConfig> for IndexConfig {
fn from(config: IntInvertedIndexConfig) -> Self {
IndexConfig::IntInverted(config)
}
}
impl From<FloatInvertedIndexConfig> for IndexConfig {
fn from(config: FloatInvertedIndexConfig) -> Self {
IndexConfig::FloatInverted(config)
}
}
impl From<BoolInvertedIndexConfig> for IndexConfig {
fn from(config: BoolInvertedIndexConfig) -> Self {
IndexConfig::BoolInverted(config)
}
}
impl TryFrom<&InternalCollectionConfiguration> for Schema {
type Error = SchemaError;
fn try_from(config: &InternalCollectionConfiguration) -> Result<Self, Self::Error> {
let mut schema = match &config.vector_index {
VectorIndexConfiguration::Hnsw(_) => Schema::new_default(KnnIndex::Hnsw),
VectorIndexConfiguration::Spann(_) => Schema::new_default(KnnIndex::Spann),
};
let vector_config = match &config.vector_index {
VectorIndexConfiguration::Hnsw(hnsw_config) => VectorIndexConfig {
space: Some(hnsw_config.space.clone()),
embedding_function: config.embedding_function.clone(),
source_key: None,
hnsw: Some(HnswIndexConfig {
ef_construction: Some(hnsw_config.ef_construction),
max_neighbors: Some(hnsw_config.max_neighbors),
ef_search: Some(hnsw_config.ef_search),
num_threads: Some(hnsw_config.num_threads),
batch_size: Some(hnsw_config.batch_size),
sync_threshold: Some(hnsw_config.sync_threshold),
resize_factor: Some(hnsw_config.resize_factor),
}),
spann: None,
},
VectorIndexConfiguration::Spann(spann_config) => VectorIndexConfig {
space: Some(spann_config.space.clone()),
embedding_function: config.embedding_function.clone(),
source_key: None,
hnsw: None,
spann: Some(SpannIndexConfig {
search_nprobe: Some(spann_config.search_nprobe),
search_rng_factor: Some(spann_config.search_rng_factor),
search_rng_epsilon: Some(spann_config.search_rng_epsilon),
nreplica_count: Some(spann_config.nreplica_count),
write_rng_factor: Some(spann_config.write_rng_factor),
write_rng_epsilon: Some(spann_config.write_rng_epsilon),
split_threshold: Some(spann_config.split_threshold),
num_samples_kmeans: Some(spann_config.num_samples_kmeans),
initial_lambda: Some(spann_config.initial_lambda),
reassign_neighbor_count: Some(spann_config.reassign_neighbor_count),
merge_threshold: Some(spann_config.merge_threshold),
num_centers_to_merge_to: Some(spann_config.num_centers_to_merge_to),
write_nprobe: Some(spann_config.write_nprobe),
ef_construction: Some(spann_config.ef_construction),
ef_search: Some(spann_config.ef_search),
max_neighbors: Some(spann_config.max_neighbors),
center_drift_threshold: None,
quantize: Quantization::None,
}),
},
};
if let Some(float_list) = &mut schema.defaults.float_list {
if let Some(vector_index) = &mut float_list.vector_index {
vector_index.config = vector_config.clone();
}
}
if let Some(embedding_types) = schema.keys.get_mut(EMBEDDING_KEY) {
if let Some(float_list) = &mut embedding_types.float_list {
if let Some(vector_index) = &mut float_list.vector_index {
let mut vector_config = vector_config;
vector_config.source_key = Some(DOCUMENT_KEY.to_string());
vector_index.config = vector_config;
}
}
}
Ok(schema)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::hnsw_configuration::Space;
use crate::metadata::SparseVector;
use crate::{
EmbeddingFunctionNewConfiguration, InternalHnswConfiguration, InternalSpannConfiguration,
};
use serde_json::json;
#[test]
fn test_reconcile_with_defaults_none_user_schema() {
let result = Schema::reconcile_with_defaults(None, KnnIndex::Spann).unwrap();
let expected = Schema::new_default(KnnIndex::Spann);
assert_eq!(result, expected);
}
#[test]
fn test_reconcile_with_defaults_empty_user_schema() {
let user_schema = Schema {
defaults: ValueTypes::default(),
keys: HashMap::new(),
cmek: None,
source_attached_function_id: None,
};
let result = Schema::reconcile_with_defaults(Some(&user_schema), KnnIndex::Spann).unwrap();
let expected = Schema::new_default(KnnIndex::Spann);
assert_eq!(result, expected);
}
#[test]
fn test_reconcile_with_defaults_user_overrides_string_enabled() {
let mut user_schema = Schema {
defaults: ValueTypes::default(),
keys: HashMap::new(),
cmek: None,
source_attached_function_id: None,
};
user_schema.defaults.string = Some(StringValueType {
string_inverted_index: Some(StringInvertedIndexType {
enabled: false, config: StringInvertedIndexConfig {},
}),
fts_index: None,
});
let result = Schema::reconcile_with_defaults(Some(&user_schema), KnnIndex::Spann).unwrap();
assert!(
!result
.defaults
.string
.as_ref()
.unwrap()
.string_inverted_index
.as_ref()
.unwrap()
.enabled
);
assert!(result.defaults.float.is_some());
assert!(result.defaults.int.is_some());
}
#[test]
fn test_reconcile_with_defaults_user_overrides_vector_config() {
let mut user_schema = Schema {
defaults: ValueTypes::default(),
keys: HashMap::new(),
cmek: None,
source_attached_function_id: None,
};
user_schema.defaults.float_list = Some(FloatListValueType {
vector_index: Some(VectorIndexType {
enabled: true, config: VectorIndexConfig {
space: Some(Space::L2), embedding_function: None, source_key: Some("custom_key".to_string()), hnsw: Some(HnswIndexConfig {
ef_construction: Some(500), max_neighbors: None, ef_search: None, num_threads: None,
batch_size: None,
sync_threshold: None,
resize_factor: None,
}),
spann: None,
},
}),
});
let result = {
let default_schema = Schema::new_default(KnnIndex::Hnsw);
let merged_defaults = Schema::merge_value_types(
&default_schema.defaults,
&user_schema.defaults,
KnnIndex::Hnsw,
)
.unwrap();
let mut merged_keys = default_schema.keys.clone();
for (key, user_value_types) in user_schema.keys {
if let Some(default_value_types) = merged_keys.get(&key) {
let merged_value_types = Schema::merge_value_types(
default_value_types,
&user_value_types,
KnnIndex::Hnsw,
)
.unwrap();
merged_keys.insert(key, merged_value_types);
} else {
merged_keys.insert(key, user_value_types);
}
}
Schema {
defaults: merged_defaults,
keys: merged_keys,
cmek: None,
source_attached_function_id: None,
}
};
let vector_config = &result
.defaults
.float_list
.as_ref()
.unwrap()
.vector_index
.as_ref()
.unwrap()
.config;
assert_eq!(vector_config.space, Some(Space::L2));
assert_eq!(vector_config.source_key, Some("custom_key".to_string()));
assert_eq!(
vector_config.hnsw.as_ref().unwrap().ef_construction,
Some(500)
);
assert_eq!(vector_config.embedding_function, None);
assert_eq!(
vector_config.hnsw.as_ref().unwrap().max_neighbors,
Some(default_m())
);
}
#[test]
fn test_reconcile_with_defaults_keys() {
let mut user_schema = Schema {
defaults: ValueTypes::default(),
keys: HashMap::new(),
cmek: None,
source_attached_function_id: None,
};
let custom_key_types = ValueTypes {
string: Some(StringValueType {
fts_index: Some(FtsIndexType {
enabled: true,
config: FtsIndexConfig {},
}),
string_inverted_index: Some(StringInvertedIndexType {
enabled: false,
config: StringInvertedIndexConfig {},
}),
}),
..Default::default()
};
user_schema
.keys
.insert("custom_key".to_string(), custom_key_types);
let result = Schema::reconcile_with_defaults(Some(&user_schema), KnnIndex::Spann).unwrap();
assert!(result.keys.contains_key(EMBEDDING_KEY));
assert!(result.keys.contains_key(DOCUMENT_KEY));
assert!(result.keys.contains_key("custom_key"));
let custom_override = result.keys.get("custom_key").unwrap();
assert!(
custom_override
.string
.as_ref()
.unwrap()
.fts_index
.as_ref()
.unwrap()
.enabled
);
}
#[test]
fn test_reconcile_with_defaults_override_existing_key() {
let mut user_schema = Schema {
defaults: ValueTypes::default(),
keys: HashMap::new(),
cmek: None,
source_attached_function_id: None,
};
let embedding_override = ValueTypes {
float_list: Some(FloatListValueType {
vector_index: Some(VectorIndexType {
enabled: false, config: VectorIndexConfig {
space: Some(Space::Ip), embedding_function: Some(EmbeddingFunctionConfiguration::Legacy),
source_key: Some("custom_embedding_key".to_string()),
hnsw: None,
spann: None,
},
}),
}),
..Default::default()
};
user_schema
.keys
.insert(EMBEDDING_KEY.to_string(), embedding_override);
let result = Schema::reconcile_with_defaults(Some(&user_schema), KnnIndex::Spann).unwrap();
let embedding_config = result.keys.get(EMBEDDING_KEY).unwrap();
let vector_config = &embedding_config
.float_list
.as_ref()
.unwrap()
.vector_index
.as_ref()
.unwrap();
assert!(!vector_config.enabled);
assert_eq!(vector_config.config.space, Some(Space::Ip));
assert_eq!(
vector_config.config.source_key,
Some("custom_embedding_key".to_string())
);
}
#[test]
fn test_convert_schema_to_collection_config_hnsw_roundtrip() {
let collection_config = InternalCollectionConfiguration {
vector_index: VectorIndexConfiguration::Hnsw(InternalHnswConfiguration {
space: Space::Cosine,
ef_construction: 128,
ef_search: 96,
max_neighbors: 42,
num_threads: 8,
resize_factor: 1.5,
sync_threshold: 2_000,
batch_size: 256,
}),
embedding_function: Some(EmbeddingFunctionConfiguration::Known(
EmbeddingFunctionNewConfiguration {
name: "custom".to_string(),
config: json!({"alpha": 1}),
},
)),
};
let schema = Schema::try_from(&collection_config).unwrap();
let reconstructed = InternalCollectionConfiguration::try_from(&schema).unwrap();
assert_eq!(reconstructed, collection_config);
}
#[test]
fn test_convert_schema_to_collection_config_spann_roundtrip() {
let spann_config = InternalSpannConfiguration {
space: Space::Cosine,
search_nprobe: 11,
search_rng_factor: 1.7,
write_nprobe: 5,
nreplica_count: 3,
split_threshold: 150,
merge_threshold: 80,
ef_construction: 120,
ef_search: 90,
max_neighbors: 40,
..Default::default()
};
let collection_config = InternalCollectionConfiguration {
vector_index: VectorIndexConfiguration::Spann(spann_config.clone()),
embedding_function: Some(EmbeddingFunctionConfiguration::Known(
EmbeddingFunctionNewConfiguration {
name: "custom".to_string(),
config: json!({"beta": true}),
},
)),
};
let schema = Schema::try_from(&collection_config).unwrap();
let reconstructed = InternalCollectionConfiguration::try_from(&schema).unwrap();
assert_eq!(reconstructed, collection_config);
}
#[test]
fn test_convert_schema_to_collection_config_rejects_mixed_index() {
let mut schema = Schema::new_default(KnnIndex::Hnsw);
if let Some(embedding) = schema.keys.get_mut(EMBEDDING_KEY) {
if let Some(float_list) = &mut embedding.float_list {
if let Some(vector_index) = &mut float_list.vector_index {
vector_index.config.spann = Some(SpannIndexConfig {
search_nprobe: Some(1),
search_rng_factor: Some(1.0),
search_rng_epsilon: Some(0.1),
nreplica_count: Some(1),
write_rng_factor: Some(1.0),
write_rng_epsilon: Some(0.1),
split_threshold: Some(100),
num_samples_kmeans: Some(10),
initial_lambda: Some(0.5),
reassign_neighbor_count: Some(10),
merge_threshold: Some(50),
num_centers_to_merge_to: Some(3),
write_nprobe: Some(1),
ef_construction: Some(50),
ef_search: Some(40),
max_neighbors: Some(20),
center_drift_threshold: None,
quantize: Quantization::None,
});
}
}
}
let result = InternalCollectionConfiguration::try_from(&schema);
assert!(result.is_err());
}
#[test]
fn test_ensure_key_from_metadata_no_changes_for_existing_key() {
let mut schema = Schema::new_default(KnnIndex::Hnsw);
let before = schema.clone();
let modified = schema.ensure_key_from_metadata(DOCUMENT_KEY, MetadataValueType::Str);
assert!(!modified);
assert_eq!(schema, before);
}
#[test]
fn test_ensure_key_from_metadata_populates_new_key_with_default_value_type() {
let mut schema = Schema::new_default(KnnIndex::Hnsw);
assert!(!schema.keys.contains_key("custom_field"));
let modified = schema.ensure_key_from_metadata("custom_field", MetadataValueType::Bool);
assert!(modified);
let entry = schema
.keys
.get("custom_field")
.expect("expected new key override to be inserted");
assert_eq!(entry.boolean, schema.defaults.boolean);
assert!(entry.string.is_none());
assert!(entry.int.is_none());
assert!(entry.float.is_none());
assert!(entry.float_list.is_none());
assert!(entry.sparse_vector.is_none());
}
#[test]
fn test_ensure_key_from_metadata_adds_missing_value_type_to_existing_key() {
let mut schema = Schema::new_default(KnnIndex::Hnsw);
let initial_len = schema.keys.len();
schema.keys.insert(
"custom_field".to_string(),
ValueTypes {
string: schema.defaults.string.clone(),
..Default::default()
},
);
let modified = schema.ensure_key_from_metadata("custom_field", MetadataValueType::Bool);
assert!(modified);
assert_eq!(schema.keys.len(), initial_len + 1);
let entry = schema
.keys
.get("custom_field")
.expect("expected key override to exist after ensure call");
assert!(entry.string.is_some());
assert_eq!(entry.boolean, schema.defaults.boolean);
}
#[test]
fn test_is_knn_key_indexing_enabled_sparse_disabled_errors() {
let schema = Schema::new_default(KnnIndex::Spann);
let result = schema.is_knn_key_indexing_enabled(
"custom_sparse",
&QueryVector::Sparse(SparseVector::new(vec![0_u32], vec![1.0_f32]).unwrap()),
);
let err = result.expect_err("expected indexing disabled error");
match err {
FilterValidationError::IndexingDisabled { key, value_type } => {
assert_eq!(key, "custom_sparse");
assert_eq!(value_type, crate::metadata::MetadataValueType::SparseVector);
}
other => panic!("unexpected error variant: {other:?}"),
}
}
#[test]
fn test_is_knn_key_indexing_enabled_sparse_enabled_succeeds() {
let mut schema = Schema::new_default(KnnIndex::Spann);
schema.keys.insert(
"sparse_enabled".to_string(),
ValueTypes {
sparse_vector: Some(SparseVectorValueType {
sparse_vector_index: Some(SparseVectorIndexType {
enabled: true,
config: SparseVectorIndexConfig {
embedding_function: Some(EmbeddingFunctionConfiguration::Legacy),
source_key: None,
bm25: None,
},
}),
}),
..Default::default()
},
);
let result = schema.is_knn_key_indexing_enabled(
"sparse_enabled",
&QueryVector::Sparse(SparseVector::new(vec![0_u32], vec![1.0_f32]).unwrap()),
);
assert!(result.is_ok());
}
#[test]
fn test_is_knn_key_indexing_enabled_dense_succeeds() {
let schema = Schema::new_default(KnnIndex::Spann);
let result = schema.is_knn_key_indexing_enabled(
EMBEDDING_KEY,
&QueryVector::Dense(vec![0.1_f32, 0.2_f32]),
);
assert!(result.is_ok());
}
#[test]
fn test_merge_hnsw_configs_field_level() {
let default_hnsw = HnswIndexConfig {
ef_construction: Some(200),
max_neighbors: Some(16),
ef_search: Some(10),
num_threads: Some(4),
batch_size: Some(100),
sync_threshold: Some(1000),
resize_factor: Some(1.2),
};
let user_hnsw = HnswIndexConfig {
ef_construction: Some(300), max_neighbors: None, ef_search: Some(20), num_threads: None, batch_size: None, sync_threshold: Some(2000), resize_factor: None, };
let result = Schema::merge_hnsw_configs(Some(&default_hnsw), Some(&user_hnsw)).unwrap();
assert_eq!(result.ef_construction, Some(300));
assert_eq!(result.ef_search, Some(20));
assert_eq!(result.sync_threshold, Some(2000));
assert_eq!(result.max_neighbors, Some(16));
assert_eq!(result.num_threads, Some(4));
assert_eq!(result.batch_size, Some(100));
assert_eq!(result.resize_factor, Some(1.2));
}
#[test]
fn test_merge_spann_configs_field_level() {
let default_spann = SpannIndexConfig {
search_nprobe: Some(10),
search_rng_factor: Some(1.0), search_rng_epsilon: Some(7.0), nreplica_count: Some(3),
write_rng_factor: Some(1.0), write_rng_epsilon: Some(6.0), split_threshold: Some(100), num_samples_kmeans: Some(100),
initial_lambda: Some(100.0), reassign_neighbor_count: Some(50),
merge_threshold: Some(50), num_centers_to_merge_to: Some(4), write_nprobe: Some(5),
ef_construction: Some(100),
ef_search: Some(10),
max_neighbors: Some(16),
center_drift_threshold: None,
quantize: Quantization::None,
};
let user_spann = SpannIndexConfig {
search_nprobe: Some(20), search_rng_factor: None, search_rng_epsilon: Some(8.0), nreplica_count: None, write_rng_factor: None,
write_rng_epsilon: None,
split_threshold: Some(150), num_samples_kmeans: None,
initial_lambda: None,
reassign_neighbor_count: None,
merge_threshold: None,
num_centers_to_merge_to: None,
write_nprobe: None,
ef_construction: None,
ef_search: None,
max_neighbors: None,
center_drift_threshold: None,
quantize: Quantization::None,
};
let result = Schema::merge_spann_configs(Some(&default_spann), Some(&user_spann))
.unwrap()
.unwrap();
assert_eq!(result.search_nprobe, Some(20));
assert_eq!(result.search_rng_epsilon, Some(8.0));
assert_eq!(result.split_threshold, Some(150));
assert_eq!(result.search_rng_factor, Some(1.0));
assert_eq!(result.nreplica_count, Some(3));
assert_eq!(result.initial_lambda, Some(100.0));
}
#[test]
fn test_merge_spann_configs_rejects_quantize_true() {
let default_spann = SpannIndexConfig {
search_nprobe: Some(10),
search_rng_factor: Some(1.0),
search_rng_epsilon: Some(7.0),
nreplica_count: Some(3),
write_rng_factor: Some(1.0),
write_rng_epsilon: Some(6.0),
split_threshold: Some(100),
num_samples_kmeans: Some(100),
initial_lambda: Some(100.0),
reassign_neighbor_count: Some(50),
merge_threshold: Some(50),
num_centers_to_merge_to: Some(4),
write_nprobe: Some(5),
ef_construction: Some(100),
ef_search: Some(10),
max_neighbors: Some(16),
center_drift_threshold: None,
quantize: Quantization::None,
};
let user_spann_with_quantize = SpannIndexConfig {
search_nprobe: Some(20),
search_rng_factor: None,
search_rng_epsilon: Some(8.0),
nreplica_count: None,
write_rng_factor: None,
write_rng_epsilon: None,
split_threshold: Some(150),
num_samples_kmeans: None,
initial_lambda: None,
reassign_neighbor_count: None,
merge_threshold: None,
num_centers_to_merge_to: None,
write_nprobe: None,
ef_construction: None,
ef_search: None,
max_neighbors: None,
center_drift_threshold: None,
quantize: Quantization::FourBitRabitQWithUSearch, };
let result =
Schema::merge_spann_configs(Some(&default_spann), Some(&user_spann_with_quantize));
assert!(result.is_err());
match result {
Err(SchemaError::InvalidUserInput { reason }) => {
assert!(reason.contains("quantize field cannot be set"));
}
_ => panic!("Expected InvalidUserInput error"),
}
let default_spann_with_quantize = SpannIndexConfig {
search_nprobe: Some(10),
search_rng_factor: Some(1.0),
search_rng_epsilon: Some(7.0),
nreplica_count: Some(3),
write_rng_factor: Some(1.0),
write_rng_epsilon: Some(6.0),
split_threshold: Some(100),
num_samples_kmeans: Some(100),
initial_lambda: Some(100.0),
reassign_neighbor_count: Some(50),
merge_threshold: Some(50),
num_centers_to_merge_to: Some(4),
write_nprobe: Some(5),
ef_construction: Some(100),
ef_search: Some(10),
max_neighbors: Some(16),
center_drift_threshold: None,
quantize: Quantization::FourBitRabitQWithUSearch, };
let result = Schema::merge_spann_configs(Some(&default_spann_with_quantize), None);
assert!(result.is_err());
match result {
Err(SchemaError::InvalidUserInput { reason }) => {
assert!(reason.contains("quantize field cannot be set"));
}
_ => panic!("Expected InvalidUserInput error"),
}
let result = Schema::merge_spann_configs(None, Some(&user_spann_with_quantize));
assert!(result.is_err());
match result {
Err(SchemaError::InvalidUserInput { reason }) => {
assert!(reason.contains("quantize field cannot be set"));
}
_ => panic!("Expected InvalidUserInput error"),
}
}
#[test]
fn test_spann_index_config_into_internal_configuration() {
let config = SpannIndexConfig {
search_nprobe: Some(33),
search_rng_factor: Some(1.2),
search_rng_epsilon: None,
nreplica_count: None,
write_rng_factor: Some(1.5),
write_rng_epsilon: None,
split_threshold: Some(75),
num_samples_kmeans: None,
initial_lambda: Some(0.9),
reassign_neighbor_count: Some(40),
merge_threshold: None,
num_centers_to_merge_to: Some(4),
write_nprobe: Some(60),
ef_construction: Some(180),
ef_search: Some(170),
max_neighbors: Some(32),
center_drift_threshold: None,
quantize: Quantization::None,
};
let with_space: InternalSpannConfiguration = (Some(&Space::Cosine), &config).into();
assert_eq!(with_space.space, Space::Cosine);
assert_eq!(with_space.search_nprobe, 33);
assert_eq!(with_space.search_rng_factor, 1.2);
assert_eq!(with_space.search_rng_epsilon, default_search_rng_epsilon());
assert_eq!(with_space.write_rng_factor, 1.5);
assert_eq!(with_space.write_nprobe, 60);
assert_eq!(with_space.ef_construction, 180);
assert_eq!(with_space.ef_search, 170);
assert_eq!(with_space.max_neighbors, 32);
assert_eq!(with_space.merge_threshold, default_merge_threshold());
let default_space_config: InternalSpannConfiguration = (None, &config).into();
assert_eq!(default_space_config.space, default_space());
}
#[test]
fn test_merge_string_type_combinations() {
let default = StringValueType {
string_inverted_index: Some(StringInvertedIndexType {
enabled: true,
config: StringInvertedIndexConfig {},
}),
fts_index: Some(FtsIndexType {
enabled: false,
config: FtsIndexConfig {},
}),
};
let user = StringValueType {
string_inverted_index: Some(StringInvertedIndexType {
enabled: false, config: StringInvertedIndexConfig {},
}),
fts_index: None, };
let result = Schema::merge_string_type(Some(&default), Some(&user))
.unwrap()
.unwrap();
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)
.unwrap()
.unwrap();
assert!(result.string_inverted_index.as_ref().unwrap().enabled);
let result = Schema::merge_string_type(None, Some(&user))
.unwrap()
.unwrap();
assert!(!result.string_inverted_index.as_ref().unwrap().enabled);
let result = Schema::merge_string_type(None, None).unwrap();
assert!(result.is_none());
}
#[test]
fn test_merge_vector_index_config_comprehensive() {
let default_config = VectorIndexConfig {
space: Some(Space::Cosine),
embedding_function: Some(EmbeddingFunctionConfiguration::Legacy),
source_key: Some("default_key".to_string()),
hnsw: Some(HnswIndexConfig {
ef_construction: Some(200),
max_neighbors: Some(16),
ef_search: Some(10),
num_threads: Some(4),
batch_size: Some(100),
sync_threshold: Some(1000),
resize_factor: Some(1.2),
}),
spann: None,
};
let user_config = VectorIndexConfig {
space: Some(Space::L2), embedding_function: None, source_key: Some("user_key".to_string()), hnsw: Some(HnswIndexConfig {
ef_construction: Some(300), max_neighbors: None, ef_search: None, num_threads: None,
batch_size: None,
sync_threshold: None,
resize_factor: None,
}),
spann: Some(SpannIndexConfig {
search_nprobe: Some(15),
search_rng_factor: None,
search_rng_epsilon: None,
nreplica_count: None,
write_rng_factor: None,
write_rng_epsilon: None,
split_threshold: None,
num_samples_kmeans: None,
initial_lambda: None,
reassign_neighbor_count: None,
merge_threshold: None,
num_centers_to_merge_to: None,
write_nprobe: None,
ef_construction: None,
ef_search: None,
max_neighbors: None,
center_drift_threshold: None,
quantize: Quantization::None,
}), };
let result =
Schema::merge_vector_index_config(&default_config, &user_config, KnnIndex::Hnsw)
.expect("merge should succeed");
assert_eq!(result.space, Some(Space::L2)); assert_eq!(
result.embedding_function,
Some(EmbeddingFunctionConfiguration::Legacy)
); 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());
}
#[test]
fn test_merge_sparse_vector_index_config() {
let default_config = SparseVectorIndexConfig {
embedding_function: Some(EmbeddingFunctionConfiguration::Legacy),
source_key: Some("default_sparse_key".to_string()),
bm25: None,
};
let user_config = SparseVectorIndexConfig {
embedding_function: None, source_key: Some("user_sparse_key".to_string()), bm25: None,
};
let result = Schema::merge_sparse_vector_index_config(&default_config, &user_config);
assert_eq!(result.source_key, Some("user_sparse_key".to_string()));
assert_eq!(
result.embedding_function,
Some(EmbeddingFunctionConfiguration::Legacy)
);
}
#[test]
fn test_complex_nested_merging_scenario() {
let mut user_schema = Schema {
defaults: ValueTypes::default(),
keys: HashMap::new(),
cmek: None,
source_attached_function_id: None,
};
user_schema.defaults.string = Some(StringValueType {
string_inverted_index: Some(StringInvertedIndexType {
enabled: false,
config: StringInvertedIndexConfig {},
}),
fts_index: Some(FtsIndexType {
enabled: true,
config: FtsIndexConfig {},
}),
});
user_schema.defaults.float_list = Some(FloatListValueType {
vector_index: Some(VectorIndexType {
enabled: true,
config: VectorIndexConfig {
space: Some(Space::Ip),
embedding_function: None, source_key: Some("custom_vector_key".to_string()),
hnsw: Some(HnswIndexConfig {
ef_construction: Some(400),
max_neighbors: Some(32),
ef_search: None, num_threads: None,
batch_size: None,
sync_threshold: None,
resize_factor: None,
}),
spann: None,
},
}),
});
let custom_key_override = ValueTypes {
string: Some(StringValueType {
fts_index: Some(FtsIndexType {
enabled: true,
config: FtsIndexConfig {},
}),
string_inverted_index: None,
}),
..Default::default()
};
user_schema
.keys
.insert("custom_field".to_string(), custom_key_override);
let result = {
let default_schema = Schema::new_default(KnnIndex::Hnsw);
let merged_defaults = Schema::merge_value_types(
&default_schema.defaults,
&user_schema.defaults,
KnnIndex::Hnsw,
)
.unwrap();
let mut merged_keys = default_schema.keys.clone();
for (key, user_value_types) in user_schema.keys {
if let Some(default_value_types) = merged_keys.get(&key) {
let merged_value_types = Schema::merge_value_types(
default_value_types,
&user_value_types,
KnnIndex::Hnsw,
)
.unwrap();
merged_keys.insert(key, merged_value_types);
} else {
merged_keys.insert(key, user_value_types);
}
}
Schema {
defaults: merged_defaults,
keys: merged_keys,
cmek: None,
source_attached_function_id: None,
}
};
assert!(
!result
.defaults
.string
.as_ref()
.unwrap()
.string_inverted_index
.as_ref()
.unwrap()
.enabled
);
assert!(
result
.defaults
.string
.as_ref()
.unwrap()
.fts_index
.as_ref()
.unwrap()
.enabled
);
let vector_config = &result
.defaults
.float_list
.as_ref()
.unwrap()
.vector_index
.as_ref()
.unwrap()
.config;
assert_eq!(vector_config.space, Some(Space::Ip));
assert_eq!(vector_config.embedding_function, None); assert_eq!(
vector_config.source_key,
Some("custom_vector_key".to_string())
);
assert_eq!(
vector_config.hnsw.as_ref().unwrap().ef_construction,
Some(400)
);
assert_eq!(vector_config.hnsw.as_ref().unwrap().max_neighbors, Some(32));
assert_eq!(
vector_config.hnsw.as_ref().unwrap().ef_search,
Some(default_search_ef())
);
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();
assert!(
custom_override
.string
.as_ref()
.unwrap()
.fts_index
.as_ref()
.unwrap()
.enabled
);
assert!(custom_override
.string
.as_ref()
.unwrap()
.string_inverted_index
.is_none());
}
#[test]
fn test_reconcile_with_collection_config_default_config() {
let collection_config = InternalCollectionConfiguration::default_hnsw();
let schema = Schema::try_from(&collection_config).unwrap();
let result =
Schema::reconcile_with_collection_config(&schema, &collection_config, KnnIndex::Hnsw)
.unwrap();
assert_eq!(result, schema);
}
#[test]
fn test_reconcile_double_default_hnsw_config_hnsw_schema_default_knn_hnsw() {
let collection_config = InternalCollectionConfiguration::default_hnsw();
let schema = Schema::new_default(KnnIndex::Hnsw);
let result =
Schema::reconcile_with_collection_config(&schema, &collection_config, KnnIndex::Hnsw)
.unwrap();
assert!(result.defaults.float_list.is_some());
assert!(result
.defaults
.float_list
.as_ref()
.unwrap()
.vector_index
.as_ref()
.unwrap()
.config
.hnsw
.is_some());
assert!(result
.defaults
.float_list
.as_ref()
.unwrap()
.vector_index
.as_ref()
.unwrap()
.config
.spann
.is_none());
}
#[test]
fn test_reconcile_double_default_hnsw_config_hnsw_schema_default_knn_spann() {
let collection_config = InternalCollectionConfiguration::default_hnsw();
let schema = Schema::new_default(KnnIndex::Hnsw);
let result =
Schema::reconcile_with_collection_config(&schema, &collection_config, KnnIndex::Spann)
.unwrap();
assert!(result.defaults.float_list.is_some());
assert!(result
.defaults
.float_list
.as_ref()
.unwrap()
.vector_index
.as_ref()
.unwrap()
.config
.spann
.is_some());
assert!(result
.defaults
.float_list
.as_ref()
.unwrap()
.vector_index
.as_ref()
.unwrap()
.config
.hnsw
.is_none());
}
#[test]
fn test_reconcile_double_default_hnsw_config_spann_schema_default_knn_hnsw() {
let collection_config = InternalCollectionConfiguration::default_hnsw();
let schema = Schema::new_default(KnnIndex::Spann);
let result =
Schema::reconcile_with_collection_config(&schema, &collection_config, KnnIndex::Hnsw)
.unwrap();
assert!(result.defaults.float_list.is_some());
assert!(result
.defaults
.float_list
.as_ref()
.unwrap()
.vector_index
.as_ref()
.unwrap()
.config
.hnsw
.is_some());
assert!(result
.defaults
.float_list
.as_ref()
.unwrap()
.vector_index
.as_ref()
.unwrap()
.config
.spann
.is_none());
}
#[test]
fn test_reconcile_double_default_hnsw_config_spann_schema_default_knn_spann() {
let collection_config = InternalCollectionConfiguration::default_hnsw();
let schema = Schema::new_default(KnnIndex::Spann);
let result =
Schema::reconcile_with_collection_config(&schema, &collection_config, KnnIndex::Spann)
.unwrap();
assert!(result.defaults.float_list.is_some());
assert!(result
.defaults
.float_list
.as_ref()
.unwrap()
.vector_index
.as_ref()
.unwrap()
.config
.spann
.is_some());
assert!(result
.defaults
.float_list
.as_ref()
.unwrap()
.vector_index
.as_ref()
.unwrap()
.config
.hnsw
.is_none());
}
#[test]
fn test_reconcile_double_default_spann_config_spann_schema_default_knn_hnsw() {
let collection_config = InternalCollectionConfiguration::default_spann();
let schema = Schema::new_default(KnnIndex::Spann);
let result =
Schema::reconcile_with_collection_config(&schema, &collection_config, KnnIndex::Hnsw)
.unwrap();
assert!(result.defaults.float_list.is_some());
assert!(result
.defaults
.float_list
.as_ref()
.unwrap()
.vector_index
.as_ref()
.unwrap()
.config
.hnsw
.is_some());
assert!(result
.defaults
.float_list
.as_ref()
.unwrap()
.vector_index
.as_ref()
.unwrap()
.config
.spann
.is_none());
}
#[test]
fn test_reconcile_double_default_spann_config_spann_schema_default_knn_spann() {
let collection_config = InternalCollectionConfiguration::default_spann();
let schema = Schema::new_default(KnnIndex::Spann);
let result =
Schema::reconcile_with_collection_config(&schema, &collection_config, KnnIndex::Spann)
.unwrap();
assert!(result.defaults.float_list.is_some());
assert!(result
.defaults
.float_list
.as_ref()
.unwrap()
.vector_index
.as_ref()
.unwrap()
.config
.spann
.is_some());
assert!(result
.defaults
.float_list
.as_ref()
.unwrap()
.vector_index
.as_ref()
.unwrap()
.config
.hnsw
.is_none());
assert_eq!(
result
.defaults
.float_list
.as_ref()
.unwrap()
.vector_index
.as_ref()
.unwrap()
.config
.source_key,
None
);
}
#[test]
fn test_reconcile_double_default_spann_config_hnsw_schema_default_knn_hnsw() {
let collection_config = InternalCollectionConfiguration::default_spann();
let schema = Schema::new_default(KnnIndex::Hnsw);
let result =
Schema::reconcile_with_collection_config(&schema, &collection_config, KnnIndex::Hnsw)
.unwrap();
assert!(result.defaults.float_list.is_some());
assert!(result
.defaults
.float_list
.as_ref()
.unwrap()
.vector_index
.as_ref()
.unwrap()
.config
.hnsw
.is_some());
assert!(result
.defaults
.float_list
.as_ref()
.unwrap()
.vector_index
.as_ref()
.unwrap()
.config
.spann
.is_none());
}
#[test]
fn test_reconcile_double_default_spann_config_hnsw_schema_default_knn_spann() {
let collection_config = InternalCollectionConfiguration::default_spann();
let schema = Schema::new_default(KnnIndex::Hnsw);
let result =
Schema::reconcile_with_collection_config(&schema, &collection_config, KnnIndex::Spann)
.unwrap();
assert!(result.defaults.float_list.is_some());
assert!(result
.defaults
.float_list
.as_ref()
.unwrap()
.vector_index
.as_ref()
.unwrap()
.config
.spann
.is_some());
assert!(result
.defaults
.float_list
.as_ref()
.unwrap()
.vector_index
.as_ref()
.unwrap()
.config
.hnsw
.is_none());
}
#[test]
fn test_defaults_source_key_not_document() {
let schema_hnsw = Schema::new_default(KnnIndex::Hnsw);
let schema_spann = Schema::new_default(KnnIndex::Spann);
let defaults_hnsw = schema_hnsw
.defaults
.float_list
.as_ref()
.unwrap()
.vector_index
.as_ref()
.unwrap();
assert_eq!(defaults_hnsw.config.source_key, None);
let defaults_spann = schema_spann
.defaults
.float_list
.as_ref()
.unwrap()
.vector_index
.as_ref()
.unwrap();
assert_eq!(defaults_spann.config.source_key, None);
let collection_config_hnsw = InternalCollectionConfiguration {
vector_index: VectorIndexConfiguration::Hnsw(InternalHnswConfiguration {
ef_construction: 300,
max_neighbors: 32,
ef_search: 50,
num_threads: 8,
batch_size: 200,
sync_threshold: 2000,
resize_factor: 1.5,
space: Space::L2,
}),
embedding_function: Some(EmbeddingFunctionConfiguration::Legacy),
};
let result_hnsw = Schema::reconcile_with_collection_config(
&schema_hnsw,
&collection_config_hnsw,
KnnIndex::Hnsw,
)
.unwrap();
let reconciled_defaults_hnsw = result_hnsw
.defaults
.float_list
.as_ref()
.unwrap()
.vector_index
.as_ref()
.unwrap();
assert_eq!(reconciled_defaults_hnsw.config.source_key, None);
let collection_config_spann = InternalCollectionConfiguration {
vector_index: VectorIndexConfiguration::Spann(InternalSpannConfiguration {
search_nprobe: 20,
search_rng_factor: 3.0,
search_rng_epsilon: 0.2,
nreplica_count: 5,
write_rng_factor: 2.0,
write_rng_epsilon: 0.1,
split_threshold: 2000,
num_samples_kmeans: 200,
initial_lambda: 0.8,
reassign_neighbor_count: 100,
merge_threshold: 800,
num_centers_to_merge_to: 20,
write_nprobe: 10,
ef_construction: 400,
ef_search: 60,
max_neighbors: 24,
space: Space::Cosine,
}),
embedding_function: None,
};
let result_spann = Schema::reconcile_with_collection_config(
&schema_spann,
&collection_config_spann,
KnnIndex::Spann,
)
.unwrap();
let reconciled_defaults_spann = result_spann
.defaults
.float_list
.as_ref()
.unwrap()
.vector_index
.as_ref()
.unwrap();
assert_eq!(reconciled_defaults_spann.config.source_key, None);
let embedding_hnsw = result_hnsw.keys.get(EMBEDDING_KEY).unwrap();
let embedding_vector_index_hnsw = embedding_hnsw
.float_list
.as_ref()
.unwrap()
.vector_index
.as_ref()
.unwrap();
assert_eq!(
embedding_vector_index_hnsw.config.source_key,
Some(DOCUMENT_KEY.to_string())
);
let embedding_spann = result_spann.keys.get(EMBEDDING_KEY).unwrap();
let embedding_vector_index_spann = embedding_spann
.float_list
.as_ref()
.unwrap()
.vector_index
.as_ref()
.unwrap();
assert_eq!(
embedding_vector_index_spann.config.source_key,
Some(DOCUMENT_KEY.to_string())
);
}
#[test]
fn test_try_from_source_key() {
let collection_config_hnsw = InternalCollectionConfiguration {
vector_index: VectorIndexConfiguration::Hnsw(InternalHnswConfiguration {
ef_construction: 300,
max_neighbors: 32,
ef_search: 50,
num_threads: 8,
batch_size: 200,
sync_threshold: 2000,
resize_factor: 1.5,
space: Space::L2,
}),
embedding_function: Some(EmbeddingFunctionConfiguration::Legacy),
};
let schema_hnsw = Schema::try_from(&collection_config_hnsw).unwrap();
let defaults_hnsw = schema_hnsw
.defaults
.float_list
.as_ref()
.unwrap()
.vector_index
.as_ref()
.unwrap();
assert_eq!(defaults_hnsw.config.source_key, None);
let embedding_hnsw = schema_hnsw.keys.get(EMBEDDING_KEY).unwrap();
let embedding_vector_index_hnsw = embedding_hnsw
.float_list
.as_ref()
.unwrap()
.vector_index
.as_ref()
.unwrap();
assert_eq!(
embedding_vector_index_hnsw.config.source_key,
Some(DOCUMENT_KEY.to_string())
);
let collection_config_spann = InternalCollectionConfiguration {
vector_index: VectorIndexConfiguration::Spann(InternalSpannConfiguration {
search_nprobe: 20,
search_rng_factor: 3.0,
search_rng_epsilon: 0.2,
nreplica_count: 5,
write_rng_factor: 2.0,
write_rng_epsilon: 0.1,
split_threshold: 2000,
num_samples_kmeans: 200,
initial_lambda: 0.8,
reassign_neighbor_count: 100,
merge_threshold: 800,
num_centers_to_merge_to: 20,
write_nprobe: 10,
ef_construction: 400,
ef_search: 60,
max_neighbors: 24,
space: Space::Cosine,
}),
embedding_function: None,
};
let schema_spann = Schema::try_from(&collection_config_spann).unwrap();
let defaults_spann = schema_spann
.defaults
.float_list
.as_ref()
.unwrap()
.vector_index
.as_ref()
.unwrap();
assert_eq!(defaults_spann.config.source_key, None);
let embedding_spann = schema_spann.keys.get(EMBEDDING_KEY).unwrap();
let embedding_vector_index_spann = embedding_spann
.float_list
.as_ref()
.unwrap()
.vector_index
.as_ref()
.unwrap();
assert_eq!(
embedding_vector_index_spann.config.source_key,
Some(DOCUMENT_KEY.to_string())
);
}
#[test]
fn test_default_hnsw_with_default_embedding_function() {
use crate::collection_configuration::EmbeddingFunctionNewConfiguration;
let collection_config = InternalCollectionConfiguration {
vector_index: VectorIndexConfiguration::Hnsw(InternalHnswConfiguration::default()),
embedding_function: Some(EmbeddingFunctionConfiguration::Known(
EmbeddingFunctionNewConfiguration {
name: "default".to_string(),
config: serde_json::json!({}),
},
)),
};
assert!(collection_config.is_default());
let schema = Schema::new_default(KnnIndex::Hnsw);
let result =
Schema::reconcile_with_collection_config(&schema, &collection_config, KnnIndex::Spann)
.unwrap();
let defaults = result
.defaults
.float_list
.as_ref()
.unwrap()
.vector_index
.as_ref()
.unwrap();
assert_eq!(defaults.config.source_key, None);
let embedding = result.keys.get(EMBEDDING_KEY).unwrap();
let embedding_vector_index = embedding
.float_list
.as_ref()
.unwrap()
.vector_index
.as_ref()
.unwrap();
assert_eq!(
embedding_vector_index.config.source_key,
Some(DOCUMENT_KEY.to_string())
);
let vector_index_config = defaults.config.clone();
assert!(vector_index_config.spann.is_some());
assert!(vector_index_config.hnsw.is_none());
assert_eq!(
embedding_vector_index.config.embedding_function,
Some(EmbeddingFunctionConfiguration::Known(
EmbeddingFunctionNewConfiguration {
name: "default".to_string(),
config: serde_json::json!({}),
},
))
);
assert_eq!(
defaults.config.embedding_function,
Some(EmbeddingFunctionConfiguration::Known(
EmbeddingFunctionNewConfiguration {
name: "default".to_string(),
config: serde_json::json!({}),
},
))
);
}
#[test]
fn test_reconcile_with_collection_config_both_non_default() {
let mut schema = Schema::new_default(KnnIndex::Hnsw);
schema.defaults.string = Some(StringValueType {
fts_index: Some(FtsIndexType {
enabled: true,
config: FtsIndexConfig {},
}),
string_inverted_index: None,
});
let mut collection_config = InternalCollectionConfiguration::default_hnsw();
if let VectorIndexConfiguration::Hnsw(ref mut hnsw_config) = collection_config.vector_index
{
hnsw_config.ef_construction = 500; }
let result = Schema::reconcile_schema_and_config(
Some(&schema),
Some(&collection_config),
KnnIndex::Spann,
);
assert!(result.is_err());
assert!(matches!(
result.unwrap_err(),
SchemaError::ConfigAndSchemaConflict
));
}
#[test]
fn test_reconcile_with_collection_config_hnsw_override() {
let schema = Schema::new_default(KnnIndex::Hnsw);
let collection_config = InternalCollectionConfiguration {
vector_index: VectorIndexConfiguration::Hnsw(InternalHnswConfiguration {
ef_construction: 300,
max_neighbors: 32,
ef_search: 50,
num_threads: 8,
batch_size: 200,
sync_threshold: 2000,
resize_factor: 1.5,
space: Space::L2,
}),
embedding_function: Some(EmbeddingFunctionConfiguration::Legacy),
};
let result =
Schema::reconcile_with_collection_config(&schema, &collection_config, KnnIndex::Hnsw)
.unwrap();
let embedding_override = result.keys.get(EMBEDDING_KEY).unwrap();
let vector_index = embedding_override
.float_list
.as_ref()
.unwrap()
.vector_index
.as_ref()
.unwrap();
assert!(vector_index.enabled);
assert_eq!(vector_index.config.space, Some(Space::L2));
assert_eq!(
vector_index.config.embedding_function,
Some(EmbeddingFunctionConfiguration::Legacy)
);
assert_eq!(
vector_index.config.source_key,
Some(DOCUMENT_KEY.to_string())
);
let hnsw_config = vector_index.config.hnsw.as_ref().unwrap();
assert_eq!(hnsw_config.ef_construction, Some(300));
assert_eq!(hnsw_config.max_neighbors, Some(32));
assert_eq!(hnsw_config.ef_search, Some(50));
assert_eq!(hnsw_config.num_threads, Some(8));
assert_eq!(hnsw_config.batch_size, Some(200));
assert_eq!(hnsw_config.sync_threshold, Some(2000));
assert_eq!(hnsw_config.resize_factor, Some(1.5));
assert!(vector_index.config.spann.is_none());
}
#[test]
fn test_reconcile_with_collection_config_spann_override() {
let schema = Schema::new_default(KnnIndex::Spann);
let collection_config = InternalCollectionConfiguration {
vector_index: VectorIndexConfiguration::Spann(InternalSpannConfiguration {
search_nprobe: 20,
search_rng_factor: 3.0,
search_rng_epsilon: 0.2,
nreplica_count: 5,
write_rng_factor: 2.0,
write_rng_epsilon: 0.1,
split_threshold: 2000,
num_samples_kmeans: 200,
initial_lambda: 0.8,
reassign_neighbor_count: 100,
merge_threshold: 800,
num_centers_to_merge_to: 20,
write_nprobe: 10,
ef_construction: 400,
ef_search: 60,
max_neighbors: 24,
space: Space::Cosine,
}),
embedding_function: None,
};
let result =
Schema::reconcile_with_collection_config(&schema, &collection_config, KnnIndex::Spann)
.unwrap();
let embedding_override = result.keys.get(EMBEDDING_KEY).unwrap();
let vector_index = embedding_override
.float_list
.as_ref()
.unwrap()
.vector_index
.as_ref()
.unwrap();
assert!(vector_index.enabled);
assert_eq!(vector_index.config.space, Some(Space::Cosine));
assert_eq!(vector_index.config.embedding_function, None);
assert_eq!(
vector_index.config.source_key,
Some(DOCUMENT_KEY.to_string())
);
assert!(vector_index.config.hnsw.is_none());
let spann_config = vector_index.config.spann.as_ref().unwrap();
assert_eq!(spann_config.search_nprobe, Some(20));
assert_eq!(spann_config.search_rng_factor, Some(3.0));
assert_eq!(spann_config.search_rng_epsilon, Some(0.2));
assert_eq!(spann_config.nreplica_count, Some(5));
assert_eq!(spann_config.write_rng_factor, Some(2.0));
assert_eq!(spann_config.write_rng_epsilon, Some(0.1));
assert_eq!(spann_config.split_threshold, Some(2000));
assert_eq!(spann_config.num_samples_kmeans, Some(200));
assert_eq!(spann_config.initial_lambda, Some(0.8));
assert_eq!(spann_config.reassign_neighbor_count, Some(100));
assert_eq!(spann_config.merge_threshold, Some(800));
assert_eq!(spann_config.num_centers_to_merge_to, Some(20));
assert_eq!(spann_config.write_nprobe, Some(10));
assert_eq!(spann_config.ef_construction, Some(400));
assert_eq!(spann_config.ef_search, Some(60));
assert_eq!(spann_config.max_neighbors, Some(24));
}
#[test]
fn test_reconcile_with_collection_config_updates_both_defaults_and_embedding() {
let schema = Schema::new_default(KnnIndex::Hnsw);
let collection_config = InternalCollectionConfiguration {
vector_index: VectorIndexConfiguration::Hnsw(InternalHnswConfiguration {
ef_construction: 300,
max_neighbors: 32,
ef_search: 50,
num_threads: 8,
batch_size: 200,
sync_threshold: 2000,
resize_factor: 1.5,
space: Space::L2,
}),
embedding_function: Some(EmbeddingFunctionConfiguration::Legacy),
};
let result =
Schema::reconcile_with_collection_config(&schema, &collection_config, KnnIndex::Hnsw)
.unwrap();
let defaults_vector_index = result
.defaults
.float_list
.as_ref()
.unwrap()
.vector_index
.as_ref()
.unwrap();
assert!(!defaults_vector_index.enabled);
assert_eq!(defaults_vector_index.config.space, Some(Space::L2));
assert_eq!(
defaults_vector_index.config.embedding_function,
Some(EmbeddingFunctionConfiguration::Legacy)
);
assert_eq!(defaults_vector_index.config.source_key, None);
let defaults_hnsw = defaults_vector_index.config.hnsw.as_ref().unwrap();
assert_eq!(defaults_hnsw.ef_construction, Some(300));
assert_eq!(defaults_hnsw.max_neighbors, Some(32));
let embedding_override = result.keys.get(EMBEDDING_KEY).unwrap();
let embedding_vector_index = embedding_override
.float_list
.as_ref()
.unwrap()
.vector_index
.as_ref()
.unwrap();
assert!(embedding_vector_index.enabled);
assert_eq!(embedding_vector_index.config.space, Some(Space::L2));
assert_eq!(
embedding_vector_index.config.embedding_function,
Some(EmbeddingFunctionConfiguration::Legacy)
);
assert_eq!(
embedding_vector_index.config.source_key,
Some(DOCUMENT_KEY.to_string())
);
let embedding_hnsw = embedding_vector_index.config.hnsw.as_ref().unwrap();
assert_eq!(embedding_hnsw.ef_construction, Some(300));
assert_eq!(embedding_hnsw.max_neighbors, Some(32));
}
#[test]
fn test_is_schema_default() {
let default_hnsw_schema = Schema::new_default(KnnIndex::Hnsw);
assert!(default_hnsw_schema.is_default());
let default_spann_schema = Schema::new_default(KnnIndex::Spann);
assert!(default_spann_schema.is_default());
let mut modified_schema = Schema::new_default(KnnIndex::Hnsw);
if let Some(ref mut string_type) = modified_schema.defaults.string {
if let Some(ref mut string_inverted) = string_type.string_inverted_index {
string_inverted.enabled = false; }
}
assert!(!modified_schema.is_default());
let mut schema_with_extra_overrides = Schema::new_default(KnnIndex::Hnsw);
schema_with_extra_overrides
.keys
.insert("custom_key".to_string(), ValueTypes::default());
assert!(!schema_with_extra_overrides.is_default());
}
#[test]
fn test_is_schema_default_with_space() {
let schema = Schema::new_default(KnnIndex::Hnsw);
assert!(schema.is_default());
let mut schema_with_space = Schema::new_default(KnnIndex::Hnsw);
if let Some(ref mut float_list) = schema_with_space.defaults.float_list {
if let Some(ref mut vector_index) = float_list.vector_index {
vector_index.config.space = Some(Space::Cosine);
}
}
assert!(!schema_with_space.is_default());
let mut schema_with_space_in_embedding_key = Schema::new_default(KnnIndex::Spann);
if let Some(ref mut embedding_key) = schema_with_space_in_embedding_key
.keys
.get_mut(EMBEDDING_KEY)
{
if let Some(ref mut float_list) = embedding_key.float_list {
if let Some(ref mut vector_index) = float_list.vector_index {
vector_index.config.space = Some(Space::Cosine);
}
}
}
assert!(!schema_with_space_in_embedding_key.is_default());
}
#[test]
fn test_is_schema_default_with_embedding_function() {
let schema = Schema::new_default(KnnIndex::Hnsw);
assert!(schema.is_default());
let mut schema_with_embedding_function = Schema::new_default(KnnIndex::Hnsw);
if let Some(ref mut float_list) = schema_with_embedding_function.defaults.float_list {
if let Some(ref mut vector_index) = float_list.vector_index {
vector_index.config.embedding_function =
Some(EmbeddingFunctionConfiguration::Legacy);
}
}
assert!(!schema_with_embedding_function.is_default());
let mut schema_with_embedding_function_in_embedding_key =
Schema::new_default(KnnIndex::Spann);
if let Some(ref mut embedding_key) = schema_with_embedding_function_in_embedding_key
.keys
.get_mut(EMBEDDING_KEY)
{
if let Some(ref mut float_list) = embedding_key.float_list {
if let Some(ref mut vector_index) = float_list.vector_index {
vector_index.config.embedding_function =
Some(EmbeddingFunctionConfiguration::Legacy);
}
}
}
assert!(!schema_with_embedding_function_in_embedding_key.is_default());
}
#[test]
fn test_add_merges_keys_by_value_type() {
let mut schema_a = Schema::new_default(KnnIndex::Hnsw);
let mut schema_b = Schema::new_default(KnnIndex::Hnsw);
let string_override = ValueTypes {
string: Some(StringValueType {
string_inverted_index: Some(StringInvertedIndexType {
enabled: true,
config: StringInvertedIndexConfig {},
}),
fts_index: None,
}),
..Default::default()
};
schema_a
.keys
.insert("custom_field".to_string(), string_override);
let float_override = ValueTypes {
float: Some(FloatValueType {
float_inverted_index: Some(FloatInvertedIndexType {
enabled: true,
config: FloatInvertedIndexConfig {},
}),
}),
..Default::default()
};
schema_b
.keys
.insert("custom_field".to_string(), float_override);
let merged = schema_a.merge(&schema_b).unwrap();
let merged_override = merged.keys.get("custom_field").unwrap();
assert!(merged_override.string.is_some());
assert!(merged_override.float.is_some());
assert!(
merged_override
.string
.as_ref()
.unwrap()
.string_inverted_index
.as_ref()
.unwrap()
.enabled
);
assert!(
merged_override
.float
.as_ref()
.unwrap()
.float_inverted_index
.as_ref()
.unwrap()
.enabled
);
}
#[test]
fn test_add_rejects_different_defaults() {
let schema_a = Schema::new_default(KnnIndex::Hnsw);
let mut schema_b = Schema::new_default(KnnIndex::Hnsw);
if let Some(string_type) = schema_b.defaults.string.as_mut() {
if let Some(string_index) = string_type.string_inverted_index.as_mut() {
string_index.enabled = false;
}
}
let err = schema_a.merge(&schema_b).unwrap_err();
assert!(matches!(err, SchemaError::DefaultsMismatch));
}
#[test]
fn test_add_detects_conflicting_value_type_configuration() {
let mut schema_a = Schema::new_default(KnnIndex::Hnsw);
let mut schema_b = Schema::new_default(KnnIndex::Hnsw);
let string_override_enabled = ValueTypes {
string: Some(StringValueType {
string_inverted_index: Some(StringInvertedIndexType {
enabled: true,
config: StringInvertedIndexConfig {},
}),
fts_index: None,
}),
..Default::default()
};
schema_a
.keys
.insert("custom_field".to_string(), string_override_enabled);
let string_override_disabled = ValueTypes {
string: Some(StringValueType {
string_inverted_index: Some(StringInvertedIndexType {
enabled: false,
config: StringInvertedIndexConfig {},
}),
fts_index: None,
}),
..Default::default()
};
schema_b
.keys
.insert("custom_field".to_string(), string_override_disabled);
let err = schema_a.merge(&schema_b).unwrap_err();
assert!(matches!(err, SchemaError::ConfigurationConflict { .. }));
}
#[test]
fn test_backward_compatibility_aliases() {
let old_format_json = r###"{
"defaults": {
"#string": {
"$fts_index": {
"enabled": true,
"config": {}
}
},
"#int": {
"$int_inverted_index": {
"enabled": true,
"config": {}
}
},
"#float_list": {
"$vector_index": {
"enabled": true,
"config": {
"spann": {
"search_nprobe": 10
}
}
}
}
},
"key_overrides": {
"#document": {
"#string": {
"$fts_index": {
"enabled": false,
"config": {}
}
}
}
}
}"###;
let schema_from_old: Schema = serde_json::from_str(old_format_json).unwrap();
let new_format_json = r###"{
"defaults": {
"string": {
"fts_index": {
"enabled": true,
"config": {}
}
},
"int": {
"int_inverted_index": {
"enabled": true,
"config": {}
}
},
"float_list": {
"vector_index": {
"enabled": true,
"config": {
"spann": {
"search_nprobe": 10
}
}
}
}
},
"keys": {
"#document": {
"string": {
"fts_index": {
"enabled": false,
"config": {}
}
}
}
}
}"###;
let schema_from_new: Schema = serde_json::from_str(new_format_json).unwrap();
assert_eq!(schema_from_old, schema_from_new);
assert!(schema_from_old.defaults.string.is_some());
assert!(schema_from_old
.defaults
.string
.as_ref()
.unwrap()
.fts_index
.is_some());
assert!(
schema_from_old
.defaults
.string
.as_ref()
.unwrap()
.fts_index
.as_ref()
.unwrap()
.enabled
);
assert!(schema_from_old.defaults.int.is_some());
assert!(schema_from_old
.defaults
.int
.as_ref()
.unwrap()
.int_inverted_index
.is_some());
assert!(schema_from_old.defaults.float_list.is_some());
assert!(schema_from_old
.defaults
.float_list
.as_ref()
.unwrap()
.vector_index
.is_some());
assert!(schema_from_old.keys.contains_key(DOCUMENT_KEY));
let doc_override = schema_from_old.keys.get(DOCUMENT_KEY).unwrap();
assert!(doc_override.string.is_some());
assert!(
!doc_override
.string
.as_ref()
.unwrap()
.fts_index
.as_ref()
.unwrap()
.enabled
);
let serialized = serde_json::to_string(&schema_from_old).unwrap();
assert!(serialized.contains(r#""keys":"#));
assert!(serialized.contains(r#""string":"#));
assert!(serialized.contains(r#""fts_index":"#));
assert!(serialized.contains(r#""int_inverted_index":"#));
assert!(serialized.contains(r#""vector_index":"#));
assert!(!serialized.contains(r#""key_overrides":"#));
assert!(!serialized.contains(r###""#string":"###));
assert!(!serialized.contains(r###""$fts_index":"###));
assert!(!serialized.contains(r###""$int_inverted_index":"###));
assert!(!serialized.contains(r###""$vector_index":"###));
}
#[test]
fn test_hnsw_index_config_validation() {
use validator::Validate;
let valid_config = HnswIndexConfig {
batch_size: Some(10),
sync_threshold: Some(100),
ef_construction: Some(100),
max_neighbors: Some(16),
..Default::default()
};
assert!(valid_config.validate().is_ok());
let invalid_batch_size = HnswIndexConfig {
batch_size: Some(1),
..Default::default()
};
assert!(invalid_batch_size.validate().is_err());
let invalid_sync_threshold = HnswIndexConfig {
sync_threshold: Some(1),
..Default::default()
};
assert!(invalid_sync_threshold.validate().is_err());
let boundary_config = HnswIndexConfig {
batch_size: Some(2),
sync_threshold: Some(2),
..Default::default()
};
assert!(boundary_config.validate().is_ok());
let all_none_config = HnswIndexConfig {
..Default::default()
};
assert!(all_none_config.validate().is_ok());
let other_fields_config = HnswIndexConfig {
ef_construction: Some(1),
max_neighbors: Some(1),
ef_search: Some(1),
num_threads: Some(1),
resize_factor: Some(0.1),
..Default::default()
};
assert!(other_fields_config.validate().is_ok());
}
#[test]
fn test_spann_index_config_validation() {
use validator::Validate;
let valid_config = SpannIndexConfig {
write_nprobe: Some(32),
nreplica_count: Some(4),
split_threshold: Some(100),
merge_threshold: Some(50),
reassign_neighbor_count: Some(32),
num_centers_to_merge_to: Some(4),
ef_construction: Some(100),
ef_search: Some(100),
max_neighbors: Some(32),
search_rng_factor: Some(1.0),
write_rng_factor: Some(1.0),
search_rng_epsilon: Some(7.5),
write_rng_epsilon: Some(7.5),
..Default::default()
};
assert!(valid_config.validate().is_ok());
let invalid_write_nprobe = SpannIndexConfig {
write_nprobe: Some(200),
..Default::default()
};
assert!(invalid_write_nprobe.validate().is_err());
let invalid_split_threshold = SpannIndexConfig {
split_threshold: Some(10),
..Default::default()
};
assert!(invalid_split_threshold.validate().is_err());
let invalid_split_threshold_high = SpannIndexConfig {
split_threshold: Some(250),
..Default::default()
};
assert!(invalid_split_threshold_high.validate().is_err());
let invalid_nreplica = SpannIndexConfig {
nreplica_count: Some(10),
..Default::default()
};
assert!(invalid_nreplica.validate().is_err());
let invalid_reassign = SpannIndexConfig {
reassign_neighbor_count: Some(100),
..Default::default()
};
assert!(invalid_reassign.validate().is_err());
let invalid_merge_threshold_low = SpannIndexConfig {
merge_threshold: Some(5),
..Default::default()
};
assert!(invalid_merge_threshold_low.validate().is_err());
let invalid_merge_threshold_high = SpannIndexConfig {
merge_threshold: Some(150),
..Default::default()
};
assert!(invalid_merge_threshold_high.validate().is_err());
let invalid_num_centers = SpannIndexConfig {
num_centers_to_merge_to: Some(10),
..Default::default()
};
assert!(invalid_num_centers.validate().is_err());
let invalid_ef_construction = SpannIndexConfig {
ef_construction: Some(300),
..Default::default()
};
assert!(invalid_ef_construction.validate().is_err());
let invalid_ef_search = SpannIndexConfig {
ef_search: Some(300),
..Default::default()
};
assert!(invalid_ef_search.validate().is_err());
let invalid_max_neighbors = SpannIndexConfig {
max_neighbors: Some(100),
..Default::default()
};
assert!(invalid_max_neighbors.validate().is_err());
let invalid_search_nprobe = SpannIndexConfig {
search_nprobe: Some(200),
..Default::default()
};
assert!(invalid_search_nprobe.validate().is_err());
let invalid_search_rng_factor_low = SpannIndexConfig {
search_rng_factor: Some(0.9),
..Default::default()
};
assert!(invalid_search_rng_factor_low.validate().is_err());
let invalid_search_rng_factor_high = SpannIndexConfig {
search_rng_factor: Some(1.1),
..Default::default()
};
assert!(invalid_search_rng_factor_high.validate().is_err());
let valid_search_rng_factor = SpannIndexConfig {
search_rng_factor: Some(1.0),
..Default::default()
};
assert!(valid_search_rng_factor.validate().is_ok());
let invalid_search_rng_epsilon_low = SpannIndexConfig {
search_rng_epsilon: Some(4.0),
..Default::default()
};
assert!(invalid_search_rng_epsilon_low.validate().is_err());
let invalid_search_rng_epsilon_high = SpannIndexConfig {
search_rng_epsilon: Some(11.0),
..Default::default()
};
assert!(invalid_search_rng_epsilon_high.validate().is_err());
let valid_search_rng_epsilon = SpannIndexConfig {
search_rng_epsilon: Some(7.5),
..Default::default()
};
assert!(valid_search_rng_epsilon.validate().is_ok());
let invalid_write_rng_factor_low = SpannIndexConfig {
write_rng_factor: Some(0.9),
..Default::default()
};
assert!(invalid_write_rng_factor_low.validate().is_err());
let invalid_write_rng_factor_high = SpannIndexConfig {
write_rng_factor: Some(1.1),
..Default::default()
};
assert!(invalid_write_rng_factor_high.validate().is_err());
let valid_write_rng_factor = SpannIndexConfig {
write_rng_factor: Some(1.0),
..Default::default()
};
assert!(valid_write_rng_factor.validate().is_ok());
let invalid_write_rng_epsilon_low = SpannIndexConfig {
write_rng_epsilon: Some(4.0),
..Default::default()
};
assert!(invalid_write_rng_epsilon_low.validate().is_err());
let invalid_write_rng_epsilon_high = SpannIndexConfig {
write_rng_epsilon: Some(11.0),
..Default::default()
};
assert!(invalid_write_rng_epsilon_high.validate().is_err());
let valid_write_rng_epsilon = SpannIndexConfig {
write_rng_epsilon: Some(7.5),
..Default::default()
};
assert!(valid_write_rng_epsilon.validate().is_ok());
let invalid_num_samples_kmeans = SpannIndexConfig {
num_samples_kmeans: Some(1500),
..Default::default()
};
assert!(invalid_num_samples_kmeans.validate().is_err());
let valid_num_samples_kmeans = SpannIndexConfig {
num_samples_kmeans: Some(500),
..Default::default()
};
assert!(valid_num_samples_kmeans.validate().is_ok());
let invalid_initial_lambda_high = SpannIndexConfig {
initial_lambda: Some(150.0),
..Default::default()
};
assert!(invalid_initial_lambda_high.validate().is_err());
let invalid_initial_lambda_low = SpannIndexConfig {
initial_lambda: Some(50.0),
..Default::default()
};
assert!(invalid_initial_lambda_low.validate().is_err());
let valid_initial_lambda = SpannIndexConfig {
initial_lambda: Some(100.0),
..Default::default()
};
assert!(valid_initial_lambda.validate().is_ok());
let all_none_config = SpannIndexConfig {
..Default::default()
};
assert!(all_none_config.validate().is_ok());
}
#[test]
fn test_builder_pattern_crud_workflow() {
let schema = Schema::new_default(KnnIndex::Hnsw)
.create_index(
None,
IndexConfig::Vector(VectorIndexConfig {
space: Some(Space::Cosine),
embedding_function: None,
source_key: None,
hnsw: Some(HnswIndexConfig {
ef_construction: Some(200),
max_neighbors: Some(32),
ef_search: Some(50),
num_threads: None,
batch_size: None,
sync_threshold: None,
resize_factor: None,
}),
spann: None,
}),
)
.expect("vector config should succeed")
.create_index(
Some("category"),
IndexConfig::StringInverted(StringInvertedIndexConfig {}),
)
.expect("string inverted on key should succeed")
.create_index(
Some("year"),
IndexConfig::IntInverted(IntInvertedIndexConfig {}),
)
.expect("int inverted on key should succeed")
.create_index(
Some("rating"),
IndexConfig::FloatInverted(FloatInvertedIndexConfig {}),
)
.expect("float inverted on key should succeed")
.create_index(
Some("is_active"),
IndexConfig::BoolInverted(BoolInvertedIndexConfig {}),
)
.expect("bool inverted on key should succeed");
assert!(schema.keys.contains_key(EMBEDDING_KEY));
let embedding = schema.keys.get(EMBEDDING_KEY).unwrap();
assert!(embedding.float_list.is_some());
let vector_index = embedding
.float_list
.as_ref()
.unwrap()
.vector_index
.as_ref()
.unwrap();
assert!(vector_index.enabled);
assert_eq!(vector_index.config.space, Some(Space::Cosine));
assert_eq!(
vector_index.config.hnsw.as_ref().unwrap().ef_construction,
Some(200)
);
assert!(schema.keys.contains_key("category"));
assert!(schema.keys.contains_key("year"));
assert!(schema.keys.contains_key("rating"));
assert!(schema.keys.contains_key("is_active"));
let category = schema.keys.get("category").unwrap();
assert!(category.string.is_some());
let string_idx = category
.string
.as_ref()
.unwrap()
.string_inverted_index
.as_ref()
.unwrap();
assert!(string_idx.enabled);
let year = schema.keys.get("year").unwrap();
assert!(year.int.is_some());
let int_idx = year
.int
.as_ref()
.unwrap()
.int_inverted_index
.as_ref()
.unwrap();
assert!(int_idx.enabled);
let schema = schema
.delete_index(
Some("category"),
IndexConfig::StringInverted(StringInvertedIndexConfig {}),
)
.expect("delete string inverted should succeed")
.delete_index(
Some("year"),
IndexConfig::IntInverted(IntInvertedIndexConfig {}),
)
.expect("delete int inverted should succeed");
let category = schema.keys.get("category").unwrap();
let string_idx = category
.string
.as_ref()
.unwrap()
.string_inverted_index
.as_ref()
.unwrap();
assert!(!string_idx.enabled);
let year = schema.keys.get("year").unwrap();
let int_idx = year
.int
.as_ref()
.unwrap()
.int_inverted_index
.as_ref()
.unwrap();
assert!(!int_idx.enabled);
let rating = schema.keys.get("rating").unwrap();
let float_idx = rating
.float
.as_ref()
.unwrap()
.float_inverted_index
.as_ref()
.unwrap();
assert!(float_idx.enabled);
let is_active = schema.keys.get("is_active").unwrap();
let bool_idx = is_active
.boolean
.as_ref()
.unwrap()
.bool_inverted_index
.as_ref()
.unwrap();
assert!(bool_idx.enabled); }
#[test]
fn test_builder_create_index_validation_errors() {
let result = Schema::new_default(KnnIndex::Hnsw).create_index(
Some("my_vectors"),
IndexConfig::Vector(VectorIndexConfig {
space: Some(Space::L2),
embedding_function: None,
source_key: None,
hnsw: None,
spann: None,
}),
);
assert!(result.is_err());
assert!(matches!(
result.unwrap_err(),
SchemaBuilderError::VectorIndexMustBeGlobal { key } if key == "my_vectors"
));
let result = Schema::new_default(KnnIndex::Hnsw)
.create_index(Some("my_text"), IndexConfig::Fts(FtsIndexConfig {}));
assert!(result.is_err());
assert!(matches!(
result.unwrap_err(),
SchemaBuilderError::FtsIndexOnlyOnDocument
));
let schema = Schema::new_default(KnnIndex::Hnsw)
.create_index(Some(DOCUMENT_KEY), IndexConfig::Fts(FtsIndexConfig {}))
.expect("FTS on #document should succeed");
assert!(schema.is_fts_enabled());
let result = Schema::new_default(KnnIndex::Hnsw).create_index(
Some(DOCUMENT_KEY),
IndexConfig::StringInverted(StringInvertedIndexConfig {}),
);
assert!(result.is_err());
assert!(matches!(
result.unwrap_err(),
SchemaBuilderError::SpecialKeyModificationNotAllowed { .. }
));
let result = Schema::new_default(KnnIndex::Hnsw).create_index(
Some(EMBEDDING_KEY),
IndexConfig::IntInverted(IntInvertedIndexConfig {}),
);
assert!(result.is_err());
assert!(matches!(
result.unwrap_err(),
SchemaBuilderError::SpecialKeyModificationNotAllowed { .. }
));
let result = Schema::new_default(KnnIndex::Hnsw).create_index(
None,
IndexConfig::SparseVector(SparseVectorIndexConfig {
embedding_function: None,
source_key: None,
bm25: None,
}),
);
assert!(result.is_err());
assert!(matches!(
result.unwrap_err(),
SchemaBuilderError::SparseVectorRequiresKey
));
let result = Schema::new_default(KnnIndex::Hnsw)
.create_index(
Some("sparse1"),
IndexConfig::SparseVector(SparseVectorIndexConfig {
embedding_function: None,
source_key: None,
bm25: None,
}),
)
.expect("first sparse should succeed")
.create_index(
Some("sparse2"),
IndexConfig::SparseVector(SparseVectorIndexConfig {
embedding_function: None,
source_key: None,
bm25: None,
}),
);
assert!(result.is_err());
assert!(matches!(
result.unwrap_err(),
SchemaBuilderError::MultipleSparseVectorIndexes { existing_key } if existing_key == "sparse1"
));
}
#[test]
fn test_builder_delete_index_validation_errors() {
let result = Schema::new_default(KnnIndex::Hnsw).delete_index(
Some(EMBEDDING_KEY),
IndexConfig::StringInverted(StringInvertedIndexConfig {}),
);
assert!(result.is_err());
assert!(matches!(
result.unwrap_err(),
SchemaBuilderError::SpecialKeyModificationNotAllowed { .. }
));
let result = Schema::new_default(KnnIndex::Hnsw).delete_index(
Some(DOCUMENT_KEY),
IndexConfig::IntInverted(IntInvertedIndexConfig {}),
);
assert!(result.is_err());
assert!(matches!(
result.unwrap_err(),
SchemaBuilderError::SpecialKeyModificationNotAllowed { .. }
));
let result = Schema::new_default(KnnIndex::Hnsw).delete_index(
None,
IndexConfig::Vector(VectorIndexConfig {
space: None,
embedding_function: None,
source_key: None,
hnsw: None,
spann: None,
}),
);
assert!(result.is_err());
assert!(matches!(
result.unwrap_err(),
SchemaBuilderError::VectorIndexDeletionNotSupported
));
let schema = Schema::new_default(KnnIndex::Hnsw)
.delete_index(Some(DOCUMENT_KEY), IndexConfig::Fts(FtsIndexConfig {}))
.expect("FTS deletion should succeed");
assert!(!schema.is_fts_enabled());
let result = Schema::new_default(KnnIndex::Hnsw)
.create_index(
Some("sparse"),
IndexConfig::SparseVector(SparseVectorIndexConfig {
embedding_function: None,
source_key: None,
bm25: None,
}),
)
.expect("create should succeed")
.delete_index(
Some("sparse"),
IndexConfig::SparseVector(SparseVectorIndexConfig {
embedding_function: None,
source_key: None,
bm25: None,
}),
);
assert!(result.is_err());
assert!(matches!(
result.unwrap_err(),
SchemaBuilderError::SparseVectorIndexDeletionNotSupported
));
}
#[test]
fn test_fts_create_global_without_key_rejected() {
let result = Schema::new_default(KnnIndex::Hnsw)
.create_index(None, IndexConfig::Fts(FtsIndexConfig {}));
assert!(result.is_err());
assert!(matches!(
result.unwrap_err(),
SchemaBuilderError::FtsIndexOnlyOnDocument
));
}
#[test]
fn test_fts_delete_global_without_key_rejected() {
let result = Schema::new_default(KnnIndex::Hnsw)
.delete_index(None, IndexConfig::Fts(FtsIndexConfig {}));
assert!(result.is_err());
assert!(matches!(
result.unwrap_err(),
SchemaBuilderError::FtsIndexDeletionOnlyOnDocument
));
}
#[test]
fn test_fts_delete_on_custom_key_rejected() {
let result = Schema::new_default(KnnIndex::Hnsw)
.delete_index(Some("my_text"), IndexConfig::Fts(FtsIndexConfig {}));
assert!(result.is_err());
assert!(matches!(
result.unwrap_err(),
SchemaBuilderError::FtsIndexDeletionOnlyOnDocument
));
}
#[test]
fn test_reserved_key_prefix_create_index() {
let result = Schema::new_default(KnnIndex::Hnsw).create_index(
Some("#custom_field"),
IndexConfig::StringInverted(StringInvertedIndexConfig {}),
);
assert!(result.is_err());
assert!(matches!(
result.unwrap_err(),
SchemaBuilderError::ReservedKeyPrefix { key } if key == "#custom_field"
));
}
#[test]
fn test_reserved_key_prefix_delete_index() {
let result = Schema::new_default(KnnIndex::Hnsw).delete_index(
Some("#custom_field"),
IndexConfig::StringInverted(StringInvertedIndexConfig {}),
);
assert!(result.is_err());
assert!(matches!(
result.unwrap_err(),
SchemaBuilderError::ReservedKeyPrefix { key } if key == "#custom_field"
));
}
#[test]
fn test_is_fts_enabled_backward_compatibility() {
let schema = Schema::new_default(KnnIndex::Hnsw);
assert!(schema.is_fts_enabled());
let empty_schema = Schema {
defaults: ValueTypes::default(),
keys: HashMap::new(),
cmek: None,
source_attached_function_id: None,
};
assert!(empty_schema.is_fts_enabled());
}
#[test]
fn test_is_fts_enabled_after_disable() {
let schema = Schema::new_default(KnnIndex::Hnsw)
.delete_index(Some(DOCUMENT_KEY), IndexConfig::Fts(FtsIndexConfig {}))
.expect("FTS deletion should succeed");
assert!(!schema.is_fts_enabled());
}
#[test]
fn test_is_fts_enabled_after_reenable() {
let schema = Schema::new_default(KnnIndex::Hnsw)
.delete_index(Some(DOCUMENT_KEY), IndexConfig::Fts(FtsIndexConfig {}))
.expect("FTS deletion should succeed")
.create_index(Some(DOCUMENT_KEY), IndexConfig::Fts(FtsIndexConfig {}))
.expect("FTS creation should succeed");
assert!(schema.is_fts_enabled());
}
#[test]
fn test_fts_disabled_blocks_where_document_validation() {
use crate::{DocumentExpression, DocumentOperator};
let schema = Schema::new_default(KnnIndex::Hnsw)
.delete_index(Some(DOCUMENT_KEY), IndexConfig::Fts(FtsIndexConfig {}))
.expect("FTS deletion should succeed");
let where_clause = Where::Document(DocumentExpression {
operator: DocumentOperator::Contains,
pattern: "test query".to_string(),
});
let result = schema.is_metadata_where_indexing_enabled(&where_clause);
assert!(result.is_err());
assert!(matches!(
result.unwrap_err(),
FilterValidationError::FtsDisabled
));
}
#[test]
fn test_fts_enabled_allows_where_document_validation() {
use crate::{DocumentExpression, DocumentOperator};
let schema = Schema::new_default(KnnIndex::Hnsw);
let where_clause = Where::Document(DocumentExpression {
operator: DocumentOperator::Contains,
pattern: "test query".to_string(),
});
let result = schema.is_metadata_where_indexing_enabled(&where_clause);
assert!(result.is_ok());
}
#[test]
fn test_builder_pattern_chaining() {
let schema = Schema::new_default(KnnIndex::Hnsw)
.create_index(Some("tag1"), StringInvertedIndexConfig {}.into())
.unwrap()
.create_index(Some("tag2"), StringInvertedIndexConfig {}.into())
.unwrap()
.create_index(Some("tag3"), StringInvertedIndexConfig {}.into())
.unwrap()
.create_index(Some("count"), IntInvertedIndexConfig {}.into())
.unwrap()
.delete_index(Some("tag2"), StringInvertedIndexConfig {}.into())
.unwrap()
.create_index(Some("score"), FloatInvertedIndexConfig {}.into())
.unwrap();
assert!(
schema
.keys
.get("tag1")
.unwrap()
.string
.as_ref()
.unwrap()
.string_inverted_index
.as_ref()
.unwrap()
.enabled
);
assert!(
!schema
.keys
.get("tag2")
.unwrap()
.string
.as_ref()
.unwrap()
.string_inverted_index
.as_ref()
.unwrap()
.enabled
);
assert!(
schema
.keys
.get("tag3")
.unwrap()
.string
.as_ref()
.unwrap()
.string_inverted_index
.as_ref()
.unwrap()
.enabled
);
assert!(
schema
.keys
.get("count")
.unwrap()
.int
.as_ref()
.unwrap()
.int_inverted_index
.as_ref()
.unwrap()
.enabled
);
assert!(
schema
.keys
.get("score")
.unwrap()
.float
.as_ref()
.unwrap()
.float_inverted_index
.as_ref()
.unwrap()
.enabled
);
}
#[test]
fn test_schema_default_matches_python() {
let schema = Schema::default();
assert!(schema.defaults.string.is_some());
let string = schema.defaults.string.as_ref().unwrap();
assert!(!string.fts_index.as_ref().unwrap().enabled);
assert!(string.string_inverted_index.as_ref().unwrap().enabled);
assert!(schema.defaults.float_list.is_some());
let float_list = schema.defaults.float_list.as_ref().unwrap();
assert!(!float_list.vector_index.as_ref().unwrap().enabled);
let vector_config = &float_list.vector_index.as_ref().unwrap().config;
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);
assert!(schema.defaults.sparse_vector.is_some());
let sparse = schema.defaults.sparse_vector.as_ref().unwrap();
assert!(!sparse.sparse_vector_index.as_ref().unwrap().enabled);
assert!(schema.defaults.int.is_some());
assert!(
schema
.defaults
.int
.as_ref()
.unwrap()
.int_inverted_index
.as_ref()
.unwrap()
.enabled
);
assert!(schema.defaults.float.is_some());
assert!(
schema
.defaults
.float
.as_ref()
.unwrap()
.float_inverted_index
.as_ref()
.unwrap()
.enabled
);
assert!(schema.defaults.boolean.is_some());
assert!(
schema
.defaults
.boolean
.as_ref()
.unwrap()
.bool_inverted_index
.as_ref()
.unwrap()
.enabled
);
assert!(schema.keys.contains_key(DOCUMENT_KEY));
let doc = schema.keys.get(DOCUMENT_KEY).unwrap();
assert!(doc.string.is_some());
assert!(
doc.string
.as_ref()
.unwrap()
.fts_index
.as_ref()
.unwrap()
.enabled
);
assert!(
!doc.string
.as_ref()
.unwrap()
.string_inverted_index
.as_ref()
.unwrap()
.enabled
);
assert!(schema.keys.contains_key(EMBEDDING_KEY));
let embedding = schema.keys.get(EMBEDDING_KEY).unwrap();
assert!(embedding.float_list.is_some());
let vec_idx = embedding
.float_list
.as_ref()
.unwrap()
.vector_index
.as_ref()
.unwrap();
assert!(vec_idx.enabled);
assert_eq!(vec_idx.config.source_key, Some(DOCUMENT_KEY.to_string()));
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);
}
#[test]
fn test_schema_default_works_with_builder() {
let schema = Schema::default()
.create_index(Some("category"), StringInvertedIndexConfig {}.into())
.expect("should succeed");
assert!(schema.keys.contains_key("category"));
assert!(schema.keys.contains_key(DOCUMENT_KEY));
assert!(schema.keys.contains_key(EMBEDDING_KEY));
assert_eq!(schema.keys.len(), 3);
}
#[cfg(feature = "testing")]
mod proptests {
use super::*;
use crate::strategies::{
embedding_function_strategy, internal_collection_configuration_strategy,
internal_hnsw_configuration_strategy, internal_spann_configuration_strategy,
knn_index_strategy, space_strategy, TEST_NAME_PATTERN,
};
use crate::{
HnswIndexConfig, SpannIndexConfig, VectorIndexConfig, DOCUMENT_KEY, EMBEDDING_KEY,
};
use proptest::prelude::*;
use proptest::strategy::BoxedStrategy;
use proptest::string::string_regex;
use serde_json::json;
fn default_embedding_function_strategy(
) -> impl Strategy<Value = Option<EmbeddingFunctionConfiguration>> {
proptest::option::of(prop_oneof![
Just(EmbeddingFunctionConfiguration::Unknown),
Just(EmbeddingFunctionConfiguration::Known(
EmbeddingFunctionNewConfiguration {
name: "default".to_string(),
config: json!({ "alpha": 1 }),
}
)),
])
}
fn sparse_embedding_function_strategy(
) -> impl Strategy<Value = Option<EmbeddingFunctionConfiguration>> {
let known_strategy = string_regex(TEST_NAME_PATTERN).unwrap().prop_map(|name| {
EmbeddingFunctionConfiguration::Known(EmbeddingFunctionNewConfiguration {
name,
config: json!({ "alpha": 1 }),
})
});
proptest::option::of(prop_oneof![
Just(EmbeddingFunctionConfiguration::Unknown),
known_strategy,
])
}
fn non_default_internal_collection_configuration_strategy(
) -> impl Strategy<Value = InternalCollectionConfiguration> {
internal_collection_configuration_strategy()
.prop_filter("non-default configuration", |config| !config.is_default())
}
fn partial_hnsw_index_config_strategy() -> impl Strategy<Value = HnswIndexConfig> {
(
proptest::option::of(1usize..=512),
proptest::option::of(1usize..=128),
proptest::option::of(1usize..=512),
proptest::option::of(1usize..=64),
proptest::option::of(2usize..=4096),
proptest::option::of(2usize..=4096),
proptest::option::of(prop_oneof![
Just(0.5f64),
Just(1.0f64),
Just(1.5f64),
Just(2.0f64)
]),
)
.prop_map(
|(
ef_construction,
max_neighbors,
ef_search,
num_threads,
batch_size,
sync_threshold,
resize_factor,
)| HnswIndexConfig {
ef_construction,
max_neighbors,
ef_search,
num_threads,
batch_size,
sync_threshold,
resize_factor,
},
)
}
fn partial_spann_index_config_strategy() -> impl Strategy<Value = SpannIndexConfig> {
let epsilon_strategy = prop_oneof![Just(5.0f32), Just(7.5f32), Just(10.0f32)];
(
(
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), ),
(
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), ),
)
.prop_map(
|(
(
search_nprobe,
search_rng_factor,
search_rng_epsilon,
nreplica_count,
write_rng_factor,
write_rng_epsilon,
split_threshold,
num_samples_kmeans,
),
(
initial_lambda,
reassign_neighbor_count,
merge_threshold,
num_centers_to_merge_to,
write_nprobe,
ef_construction,
ef_search,
max_neighbors,
),
)| SpannIndexConfig {
search_nprobe,
search_rng_factor,
search_rng_epsilon,
nreplica_count,
write_rng_factor,
write_rng_epsilon,
split_threshold,
num_samples_kmeans,
initial_lambda,
reassign_neighbor_count,
merge_threshold,
num_centers_to_merge_to,
write_nprobe,
ef_construction,
ef_search,
max_neighbors,
center_drift_threshold: None,
quantize: Quantization::None,
},
)
}
proptest! {
#[test]
fn merge_hnsw_configs_preserves_user_overrides(
base in partial_hnsw_index_config_strategy(),
user in partial_hnsw_index_config_strategy(),
) {
let merged = Schema::merge_hnsw_configs(Some(&base), Some(&user))
.expect("merge should return Some when both are Some");
if user.ef_construction.is_some() {
prop_assert_eq!(merged.ef_construction, user.ef_construction);
}
if user.max_neighbors.is_some() {
prop_assert_eq!(merged.max_neighbors, user.max_neighbors);
}
if user.ef_search.is_some() {
prop_assert_eq!(merged.ef_search, user.ef_search);
}
if user.num_threads.is_some() {
prop_assert_eq!(merged.num_threads, user.num_threads);
}
if user.batch_size.is_some() {
prop_assert_eq!(merged.batch_size, user.batch_size);
}
if user.sync_threshold.is_some() {
prop_assert_eq!(merged.sync_threshold, user.sync_threshold);
}
if user.resize_factor.is_some() {
prop_assert_eq!(merged.resize_factor, user.resize_factor);
}
}
#[test]
fn merge_hnsw_configs_falls_back_to_base_when_user_is_none(
base in partial_hnsw_index_config_strategy(),
) {
let merged = Schema::merge_hnsw_configs(Some(&base), None)
.expect("merge should return Some when base is Some");
prop_assert_eq!(merged, base);
}
#[test]
fn merge_hnsw_configs_returns_user_when_base_is_none(
user in partial_hnsw_index_config_strategy(),
) {
let merged = Schema::merge_hnsw_configs(None, Some(&user))
.expect("merge should return Some when user is Some");
prop_assert_eq!(merged, user);
}
#[test]
fn merge_spann_configs_preserves_user_overrides(
base in partial_spann_index_config_strategy(),
user in partial_spann_index_config_strategy(),
) {
let merged = Schema::merge_spann_configs(Some(&base), Some(&user))
.expect("merge should return Ok")
.expect("merge should return Some when both are Some");
if user.search_nprobe.is_some() {
prop_assert_eq!(merged.search_nprobe, user.search_nprobe);
}
if user.search_rng_epsilon.is_some() {
prop_assert_eq!(merged.search_rng_epsilon, user.search_rng_epsilon);
}
if user.split_threshold.is_some() {
prop_assert_eq!(merged.split_threshold, user.split_threshold);
}
if user.ef_construction.is_some() {
prop_assert_eq!(merged.ef_construction, user.ef_construction);
}
if user.ef_search.is_some() {
prop_assert_eq!(merged.ef_search, user.ef_search);
}
if user.max_neighbors.is_some() {
prop_assert_eq!(merged.max_neighbors, user.max_neighbors);
}
}
#[test]
fn merge_spann_configs_falls_back_to_base_when_user_is_none(
base in partial_spann_index_config_strategy(),
) {
let merged = Schema::merge_spann_configs(Some(&base), None)
.expect("merge should return Ok")
.expect("merge should return Some when base is Some");
prop_assert_eq!(merged, base);
}
#[test]
fn merge_vector_index_config_preserves_user_overrides(
base in vector_index_config_strategy(),
user in vector_index_config_strategy(),
knn in knn_index_strategy(),
) {
let merged = Schema::merge_vector_index_config(&base, &user, knn)
.expect("merge should succeed");
if user.space.is_some() {
prop_assert_eq!(merged.space, user.space);
}
if user.embedding_function.is_some() {
prop_assert_eq!(merged.embedding_function, user.embedding_function);
}
if user.source_key.is_some() {
prop_assert_eq!(merged.source_key, user.source_key);
}
match knn {
KnnIndex::Hnsw => {
if let (Some(_base_hnsw), Some(user_hnsw)) = (&base.hnsw, &user.hnsw) {
let merged_hnsw = merged.hnsw.as_ref().expect("hnsw should be Some");
if user_hnsw.ef_construction.is_some() {
prop_assert_eq!(merged_hnsw.ef_construction, user_hnsw.ef_construction);
}
}
}
KnnIndex::Spann => {
if let (Some(_base_spann), Some(user_spann)) = (&base.spann, &user.spann) {
let merged_spann = merged.spann.as_ref().expect("spann should be Some");
if user_spann.search_nprobe.is_some() {
prop_assert_eq!(merged_spann.search_nprobe, user_spann.search_nprobe);
}
}
}
}
}
}
fn expected_vector_index_config(
config: &InternalCollectionConfiguration,
) -> VectorIndexConfig {
match &config.vector_index {
VectorIndexConfiguration::Hnsw(hnsw_config) => VectorIndexConfig {
space: Some(hnsw_config.space.clone()),
embedding_function: config.embedding_function.clone(),
source_key: None,
hnsw: Some(HnswIndexConfig {
ef_construction: Some(hnsw_config.ef_construction),
max_neighbors: Some(hnsw_config.max_neighbors),
ef_search: Some(hnsw_config.ef_search),
num_threads: Some(hnsw_config.num_threads),
batch_size: Some(hnsw_config.batch_size),
sync_threshold: Some(hnsw_config.sync_threshold),
resize_factor: Some(hnsw_config.resize_factor),
}),
spann: None,
},
VectorIndexConfiguration::Spann(spann_config) => VectorIndexConfig {
space: Some(spann_config.space.clone()),
embedding_function: config.embedding_function.clone(),
source_key: None,
hnsw: None,
spann: Some(SpannIndexConfig {
search_nprobe: Some(spann_config.search_nprobe),
search_rng_factor: Some(spann_config.search_rng_factor),
search_rng_epsilon: Some(spann_config.search_rng_epsilon),
nreplica_count: Some(spann_config.nreplica_count),
write_rng_factor: Some(spann_config.write_rng_factor),
write_rng_epsilon: Some(spann_config.write_rng_epsilon),
split_threshold: Some(spann_config.split_threshold),
num_samples_kmeans: Some(spann_config.num_samples_kmeans),
initial_lambda: Some(spann_config.initial_lambda),
reassign_neighbor_count: Some(spann_config.reassign_neighbor_count),
merge_threshold: Some(spann_config.merge_threshold),
num_centers_to_merge_to: Some(spann_config.num_centers_to_merge_to),
write_nprobe: Some(spann_config.write_nprobe),
ef_construction: Some(spann_config.ef_construction),
ef_search: Some(spann_config.ef_search),
max_neighbors: Some(spann_config.max_neighbors),
center_drift_threshold: None,
quantize: Quantization::None,
}),
},
}
}
fn non_special_key_strategy() -> BoxedStrategy<String> {
string_regex(TEST_NAME_PATTERN)
.unwrap()
.prop_filter("exclude special keys", |key| {
key != DOCUMENT_KEY && key != EMBEDDING_KEY
})
.boxed()
}
fn source_key_strategy() -> BoxedStrategy<Option<String>> {
proptest::option::of(prop_oneof![
Just(DOCUMENT_KEY.to_string()),
string_regex(TEST_NAME_PATTERN).unwrap(),
])
.boxed()
}
fn fts_index_type_strategy() -> impl Strategy<Value = FtsIndexType> {
any::<bool>().prop_map(|enabled| FtsIndexType {
enabled,
config: FtsIndexConfig {},
})
}
fn string_inverted_index_type_strategy() -> impl Strategy<Value = StringInvertedIndexType> {
any::<bool>().prop_map(|enabled| StringInvertedIndexType {
enabled,
config: StringInvertedIndexConfig {},
})
}
fn string_value_type_strategy() -> BoxedStrategy<Option<StringValueType>> {
proptest::option::of(
(
proptest::option::of(string_inverted_index_type_strategy()),
proptest::option::of(fts_index_type_strategy()),
)
.prop_map(|(string_inverted_index, fts_index)| {
StringValueType {
string_inverted_index,
fts_index,
}
}),
)
.boxed()
}
fn float_inverted_index_type_strategy() -> impl Strategy<Value = FloatInvertedIndexType> {
any::<bool>().prop_map(|enabled| FloatInvertedIndexType {
enabled,
config: FloatInvertedIndexConfig {},
})
}
fn float_value_type_strategy() -> BoxedStrategy<Option<FloatValueType>> {
proptest::option::of(
proptest::option::of(float_inverted_index_type_strategy()).prop_map(
|float_inverted_index| FloatValueType {
float_inverted_index,
},
),
)
.boxed()
}
fn int_inverted_index_type_strategy() -> impl Strategy<Value = IntInvertedIndexType> {
any::<bool>().prop_map(|enabled| IntInvertedIndexType {
enabled,
config: IntInvertedIndexConfig {},
})
}
fn int_value_type_strategy() -> BoxedStrategy<Option<IntValueType>> {
proptest::option::of(
proptest::option::of(int_inverted_index_type_strategy())
.prop_map(|int_inverted_index| IntValueType { int_inverted_index }),
)
.boxed()
}
fn bool_inverted_index_type_strategy() -> impl Strategy<Value = BoolInvertedIndexType> {
any::<bool>().prop_map(|enabled| BoolInvertedIndexType {
enabled,
config: BoolInvertedIndexConfig {},
})
}
fn bool_value_type_strategy() -> BoxedStrategy<Option<BoolValueType>> {
proptest::option::of(
proptest::option::of(bool_inverted_index_type_strategy()).prop_map(
|bool_inverted_index| BoolValueType {
bool_inverted_index,
},
),
)
.boxed()
}
fn sparse_vector_index_config_strategy() -> impl Strategy<Value = SparseVectorIndexConfig> {
(
sparse_embedding_function_strategy(),
source_key_strategy(),
proptest::option::of(any::<bool>()),
)
.prop_map(|(embedding_function, source_key, bm25)| {
SparseVectorIndexConfig {
embedding_function,
source_key,
bm25,
}
})
}
fn sparse_vector_value_type_strategy() -> BoxedStrategy<Option<SparseVectorValueType>> {
proptest::option::of(
(
any::<bool>(),
proptest::option::of(sparse_vector_index_config_strategy()),
)
.prop_map(|(enabled, config)| SparseVectorValueType {
sparse_vector_index: config.map(|cfg| SparseVectorIndexType {
enabled,
config: cfg,
}),
}),
)
.boxed()
}
fn hnsw_index_config_strategy() -> impl Strategy<Value = HnswIndexConfig> {
internal_hnsw_configuration_strategy().prop_map(|config| HnswIndexConfig {
ef_construction: Some(config.ef_construction),
max_neighbors: Some(config.max_neighbors),
ef_search: Some(config.ef_search),
num_threads: Some(config.num_threads),
batch_size: Some(config.batch_size),
sync_threshold: Some(config.sync_threshold),
resize_factor: Some(config.resize_factor),
})
}
fn spann_index_config_strategy() -> impl Strategy<Value = SpannIndexConfig> {
internal_spann_configuration_strategy().prop_map(|config| SpannIndexConfig {
search_nprobe: Some(config.search_nprobe),
search_rng_factor: Some(config.search_rng_factor),
search_rng_epsilon: Some(config.search_rng_epsilon),
nreplica_count: Some(config.nreplica_count),
write_rng_factor: Some(config.write_rng_factor),
write_rng_epsilon: Some(config.write_rng_epsilon),
split_threshold: Some(config.split_threshold),
num_samples_kmeans: Some(config.num_samples_kmeans),
initial_lambda: Some(config.initial_lambda),
reassign_neighbor_count: Some(config.reassign_neighbor_count),
merge_threshold: Some(config.merge_threshold),
num_centers_to_merge_to: Some(config.num_centers_to_merge_to),
write_nprobe: Some(config.write_nprobe),
ef_construction: Some(config.ef_construction),
ef_search: Some(config.ef_search),
max_neighbors: Some(config.max_neighbors),
center_drift_threshold: None,
quantize: Quantization::None,
})
}
fn vector_index_config_strategy() -> impl Strategy<Value = VectorIndexConfig> {
(
proptest::option::of(space_strategy()),
embedding_function_strategy(),
source_key_strategy(),
proptest::option::of(hnsw_index_config_strategy()),
proptest::option::of(spann_index_config_strategy()),
)
.prop_map(|(space, embedding_function, source_key, hnsw, spann)| {
VectorIndexConfig {
space,
embedding_function,
source_key,
hnsw,
spann,
}
})
}
fn vector_index_type_strategy() -> impl Strategy<Value = VectorIndexType> {
(any::<bool>(), vector_index_config_strategy())
.prop_map(|(enabled, config)| VectorIndexType { enabled, config })
}
fn float_list_value_type_strategy() -> BoxedStrategy<Option<FloatListValueType>> {
proptest::option::of(
proptest::option::of(vector_index_type_strategy())
.prop_map(|vector_index| FloatListValueType { vector_index }),
)
.boxed()
}
fn value_types_strategy() -> BoxedStrategy<ValueTypes> {
(
string_value_type_strategy(),
float_list_value_type_strategy(),
sparse_vector_value_type_strategy(),
int_value_type_strategy(),
float_value_type_strategy(),
bool_value_type_strategy(),
)
.prop_map(
|(string, float_list, sparse_vector, int, float, boolean)| ValueTypes {
string,
float_list,
sparse_vector,
int,
float,
boolean,
},
)
.boxed()
}
fn schema_strategy() -> BoxedStrategy<Schema> {
(
value_types_strategy(),
proptest::collection::hash_map(
non_special_key_strategy(),
value_types_strategy(),
0..=3,
),
proptest::option::of(value_types_strategy()),
proptest::option::of(value_types_strategy()),
)
.prop_map(
|(defaults, mut extra_keys, document_override, embedding_override)| {
if let Some(doc) = document_override {
extra_keys.insert(DOCUMENT_KEY.to_string(), doc);
}
if let Some(embed) = embedding_override {
extra_keys.insert(EMBEDDING_KEY.to_string(), embed);
}
Schema {
defaults,
keys: extra_keys,
cmek: None,
source_attached_function_id: None,
}
},
)
.boxed()
}
fn force_non_default_schema(mut schema: Schema) -> Schema {
if schema.is_default() {
if let Some(string_value) = schema
.defaults
.string
.as_mut()
.and_then(|string_value| string_value.string_inverted_index.as_mut())
{
string_value.enabled = !string_value.enabled;
} else {
schema.defaults.string = Some(StringValueType {
string_inverted_index: Some(StringInvertedIndexType {
enabled: false,
config: StringInvertedIndexConfig {},
}),
fts_index: None,
});
}
}
schema
}
fn non_default_schema_strategy() -> BoxedStrategy<Schema> {
schema_strategy().prop_map(force_non_default_schema).boxed()
}
fn extract_vector_configs(schema: &Schema) -> (VectorIndexConfig, VectorIndexConfig) {
let defaults = schema
.defaults
.float_list
.as_ref()
.and_then(|fl| fl.vector_index.as_ref())
.map(|vi| vi.config.clone())
.expect("defaults vector index missing");
let embedding = schema
.keys
.get(EMBEDDING_KEY)
.and_then(|value_types| value_types.float_list.as_ref())
.and_then(|fl| fl.vector_index.as_ref())
.map(|vi| vi.config.clone())
.expect("#embedding vector index missing");
(defaults, embedding)
}
proptest! {
#[test]
fn reconcile_schema_and_config_matches_convert_for_config_only(
config in internal_collection_configuration_strategy(),
knn in knn_index_strategy(),
) {
let result = Schema::reconcile_schema_and_config(None, Some(&config), knn)
.expect("reconciliation should succeed");
let (defaults_vi, embedding_vi) = extract_vector_configs(&result);
let expected_config = expected_vector_index_config(&config);
prop_assert_eq!(defaults_vi, expected_config.clone());
let mut expected_embedding_config = expected_config;
expected_embedding_config.source_key = Some(DOCUMENT_KEY.to_string());
prop_assert_eq!(embedding_vi, expected_embedding_config);
prop_assert_eq!(result.keys.len(), 2);
}
}
proptest! {
#[test]
fn reconcile_schema_and_config_errors_when_both_non_default(
config in non_default_internal_collection_configuration_strategy(),
knn in knn_index_strategy(),
) {
let schema = Schema::try_from(&config)
.expect("conversion should succeed");
prop_assume!(!schema.is_default());
let result = Schema::reconcile_schema_and_config(Some(&schema), Some(&config), knn);
prop_assert!(matches!(result, Err(SchemaError::ConfigAndSchemaConflict)));
}
}
proptest! {
#[test]
fn reconcile_schema_and_config_matches_schema_only_path(
schema in schema_strategy(),
knn in knn_index_strategy(),
) {
let result = Schema::reconcile_schema_and_config(Some(&schema), None, knn)
.expect("reconciliation should succeed");
let (defaults_vi, embedding_vi) = extract_vector_configs(&result);
if let Some(schema_float_list) = schema.defaults.float_list.as_ref() {
if let Some(schema_vi) = schema_float_list.vector_index.as_ref() {
if let Some(schema_space) = &schema_vi.config.space {
prop_assert_eq!(defaults_vi.space, Some(schema_space.clone()));
}
if let Some(schema_ef) = &schema_vi.config.embedding_function {
prop_assert_eq!(defaults_vi.embedding_function, Some(schema_ef.clone()));
}
match knn {
KnnIndex::Hnsw => {
if let Some(schema_hnsw) = &schema_vi.config.hnsw {
if let Some(merged_hnsw) = &defaults_vi.hnsw {
if let Some(schema_ef_construction) = schema_hnsw.ef_construction {
prop_assert_eq!(merged_hnsw.ef_construction, Some(schema_ef_construction));
}
}
}
}
KnnIndex::Spann => {
if let Some(schema_spann) = &schema_vi.config.spann {
if let Some(merged_spann) = &defaults_vi.spann {
if let Some(schema_search_nprobe) = schema_spann.search_nprobe {
prop_assert_eq!(merged_spann.search_nprobe, Some(schema_search_nprobe));
}
}
}
}
}
}
}
if let Some(embedding_values) = schema.keys.get(EMBEDDING_KEY) {
if let Some(embedding_float_list) = embedding_values.float_list.as_ref() {
if let Some(embedding_vi_type) = embedding_float_list.vector_index.as_ref() {
if let Some(schema_space) = &embedding_vi_type.config.space {
prop_assert_eq!(embedding_vi.space, Some(schema_space.clone()));
}
}
}
}
}
}
proptest! {
#[test]
fn reconcile_schema_and_config_with_default_schema_and_default_config_applies_embedding_function(
embedding_function in default_embedding_function_strategy(),
knn in knn_index_strategy(),
) {
let schema = Schema::new_default(knn);
let mut config = match knn {
KnnIndex::Hnsw => InternalCollectionConfiguration::default_hnsw(),
KnnIndex::Spann => InternalCollectionConfiguration::default_spann(),
};
config.embedding_function = embedding_function.clone();
let result = Schema::reconcile_schema_and_config(
Some(&schema),
Some(&config),
knn,
)
.expect("reconciliation should succeed");
let (defaults_vi, embedding_vi) = extract_vector_configs(&result);
if let Some(ef) = embedding_function {
prop_assert_eq!(defaults_vi.embedding_function, Some(ef.clone()));
prop_assert_eq!(embedding_vi.embedding_function, Some(ef));
} else {
prop_assert_eq!(defaults_vi.embedding_function, None);
prop_assert_eq!(embedding_vi.embedding_function, None);
}
}
}
proptest! {
#[test]
fn reconcile_schema_and_config_with_default_config_keeps_non_default_schema(
schema in non_default_schema_strategy(),
knn in knn_index_strategy(),
) {
let default_config = match knn {
KnnIndex::Hnsw => InternalCollectionConfiguration::default_hnsw(),
KnnIndex::Spann => InternalCollectionConfiguration::default_spann(),
};
let result = Schema::reconcile_schema_and_config(
Some(&schema),
Some(&default_config),
knn,
)
.expect("reconciliation should succeed");
let (defaults_vi, embedding_vi) = extract_vector_configs(&result);
if let Some(schema_float_list) = schema.defaults.float_list.as_ref() {
if let Some(schema_vi) = schema_float_list.vector_index.as_ref() {
if let Some(schema_space) = &schema_vi.config.space {
prop_assert_eq!(defaults_vi.space, Some(schema_space.clone()));
}
if let Some(schema_ef) = &schema_vi.config.embedding_function {
prop_assert_eq!(defaults_vi.embedding_function, Some(schema_ef.clone()));
}
}
}
if let Some(embedding_values) = schema.keys.get(EMBEDDING_KEY) {
if let Some(embedding_float_list) = embedding_values.float_list.as_ref() {
if let Some(embedding_vi_type) = embedding_float_list.vector_index.as_ref() {
if let Some(schema_space) = &embedding_vi_type.config.space {
prop_assert_eq!(embedding_vi.space, Some(schema_space.clone()));
}
}
}
}
}
}
}
}