use crate::hnsw::distance_metric::DistanceMetric;
#[derive(Debug, Clone, PartialEq)]
pub struct HnswConfig {
pub dimension: usize,
pub m: usize,
pub ef_construction: usize,
pub ef_search: usize,
pub ml: u8,
pub distance_metric: DistanceMetric,
pub enable_multilayer: bool,
pub multilayer_level_distribution_base: Option<usize>,
pub multilayer_deterministic_seed: Option<u64>,
}
impl HnswConfig {
pub fn new(
dimension: usize,
m: usize,
ef_construction: usize,
distance_metric: DistanceMetric,
) -> Self {
HnswConfig {
dimension,
m,
ef_construction,
ef_search: ef_construction, ml: 16, distance_metric,
enable_multilayer: false, multilayer_level_distribution_base: None,
multilayer_deterministic_seed: None,
}
}
}
impl Default for HnswConfig {
fn default() -> Self {
HnswConfig {
dimension: 768, m: 16, ef_construction: 200, ef_search: 50, ml: 16, distance_metric: DistanceMetric::Cosine,
enable_multilayer: false, multilayer_level_distribution_base: None, multilayer_deterministic_seed: None, }
}
}
pub fn hnsw_config() -> crate::hnsw::builder::HnswConfigBuilder {
crate::hnsw::builder::HnswConfigBuilder::new()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_default_config() {
let config = HnswConfig::default();
assert_eq!(config.dimension, 768);
assert_eq!(config.m, 16);
assert_eq!(config.ef_construction, 200);
assert_eq!(config.ef_search, 50);
assert_eq!(config.ml, 16);
assert_eq!(config.distance_metric, DistanceMetric::Cosine);
}
#[test]
fn test_config_clone() {
let config1 = HnswConfig {
dimension: 256,
m: 12,
ef_construction: 150,
ef_search: 40,
ml: 12,
distance_metric: DistanceMetric::Manhattan,
enable_multilayer: true,
multilayer_level_distribution_base: Some(12),
multilayer_deterministic_seed: Some(123),
};
let config2 = config1.clone();
assert_eq!(config1, config2);
}
#[test]
fn test_high_precision_config() {
let config = HnswConfig {
dimension: 1536,
m: 32,
ef_construction: 400,
ef_search: 100,
ml: 24,
distance_metric: DistanceMetric::Cosine,
enable_multilayer: false,
multilayer_level_distribution_base: None,
multilayer_deterministic_seed: None,
};
assert_eq!(config.dimension, 1536);
assert_eq!(config.m, 32);
assert!(config.ef_construction >= config.m);
assert!(!config.enable_multilayer);
}
#[test]
fn test_fast_construction_config() {
let config = HnswConfig {
dimension: 384,
m: 8,
ef_construction: 100,
ef_search: 20,
ml: 12,
distance_metric: DistanceMetric::Euclidean,
enable_multilayer: false,
multilayer_level_distribution_base: None,
multilayer_deterministic_seed: None,
};
assert_eq!(config.dimension, 384);
assert_eq!(config.m, 8);
assert_eq!(config.ef_construction, 100);
assert_eq!(config.ef_search, 20);
assert!(!config.enable_multilayer);
}
#[test]
fn test_hnsw_config_function() {
let config = hnsw_config()
.dimension(512)
.m_connections(24)
.distance_metric(DistanceMetric::Euclidean)
.build()
.unwrap();
assert_eq!(config.dimension, 512);
assert_eq!(config.m, 24);
assert_eq!(config.distance_metric, DistanceMetric::Euclidean);
assert!(!config.enable_multilayer); }
#[test]
fn test_multilayer_config_defaults() {
let config = HnswConfig::default();
assert!(!config.enable_multilayer);
assert_eq!(config.multilayer_level_distribution_base, None);
assert_eq!(config.multilayer_deterministic_seed, None);
}
#[test]
fn test_multilayer_config_enabled() {
let config = HnswConfig {
dimension: 256,
m: 16,
ef_construction: 200,
ef_search: 50,
ml: 16,
distance_metric: DistanceMetric::Cosine,
enable_multilayer: true,
multilayer_level_distribution_base: Some(16),
multilayer_deterministic_seed: Some(42),
};
assert!(config.enable_multilayer);
assert_eq!(config.multilayer_level_distribution_base, Some(16));
assert_eq!(config.multilayer_deterministic_seed, Some(42));
}
#[test]
fn test_multilayer_config_defaults_derivation() {
let config = HnswConfig {
dimension: 512,
m: 20,
ef_construction: 200,
ef_search: 50,
ml: 16,
distance_metric: DistanceMetric::Euclidean,
enable_multilayer: true,
multilayer_level_distribution_base: None, multilayer_deterministic_seed: None, };
assert!(config.enable_multilayer);
assert_eq!(config.multilayer_level_distribution_base, None);
assert_eq!(config.multilayer_deterministic_seed, None);
}
#[test]
fn test_multilayer_config_validation() {
let config = HnswConfig {
dimension: 768,
m: 16,
ef_construction: 200,
ef_search: 50,
ml: 16,
distance_metric: DistanceMetric::Cosine,
enable_multilayer: true,
multilayer_level_distribution_base: Some(32), multilayer_deterministic_seed: Some(12345),
};
assert_eq!(config.dimension, 768);
assert_eq!(config.m, 16);
assert_eq!(config.multilayer_level_distribution_base, Some(32));
assert_eq!(config.multilayer_deterministic_seed, Some(12345));
}
#[test]
fn test_single_layer_vs_multi_layer_config() {
let single_layer = HnswConfig {
dimension: 512,
m: 16,
ef_construction: 200,
ef_search: 50,
ml: 16,
distance_metric: DistanceMetric::Cosine,
enable_multilayer: false,
multilayer_level_distribution_base: None,
multilayer_deterministic_seed: None,
};
let multi_layer = HnswConfig {
dimension: 512,
m: 16,
ef_construction: 200,
ef_search: 50,
ml: 16,
distance_metric: DistanceMetric::Cosine,
enable_multilayer: true,
multilayer_level_distribution_base: Some(16),
multilayer_deterministic_seed: Some(42),
};
assert_ne!(single_layer, multi_layer);
assert!(!single_layer.enable_multilayer);
assert!(multi_layer.enable_multilayer);
}
}