Skip to content

Commit

Permalink
[CLN] Cleanup hnsw provider to not know about segments (#3044)
Browse files Browse the repository at this point in the history 
## Description of changes

*Summarize the changes made by this PR.*
 - Improvements & Bug fixes
	 - Leaky abstraction. The Hnsw provider knew about segments and was parsing it to extract stuff like index params. The distributed hnsw segment writer and reader should instead be doing that and provider should not know about it.

## Test plan
*How are these changes tested?*
- [x] Tests pass locally with `pytest` for python, `yarn test` for js, `cargo test` for rust

## Documentation Changes
None
  • Loading branch information
@sanketkedia
sanketkedia authored Nov 14, 2024
1 parent b55c26a commit 3f56aeb
Show file tree
Hide file tree
Showing 6 changed files with 284 additions and 250 deletions.
118 changes: 15 additions & 103 deletions rust/index/src/hnsw.rs
View file
Original file line number Diff line number Diff line change
@@ -1,16 +1,17 @@
use super::{Index, IndexConfig, IndexUuid, PersistentIndex};
use chroma_error::{ChromaError, ErrorCodes};
use chroma_types::{Metadata, MetadataValue, MetadataValueConversionError, Segment};
use chroma_types::MetadataValueConversionError;
use std::ffi::CString;
use std::ffi::{c_char, c_int};
use std::path::Path;
use std::str::Utf8Error;
use thiserror::Error;
use tracing::instrument;

const DEFAULT_MAX_ELEMENTS: usize = 10000;
const DEFAULT_HNSW_M: usize = 16;
const DEFAULT_HNSW_EF_CONSTRUCTION: usize = 100;
const DEFAULT_HNSW_EF_SEARCH: usize = 10;
pub const DEFAULT_MAX_ELEMENTS: usize = 10000;
pub const DEFAULT_HNSW_M: usize = 16;
pub const DEFAULT_HNSW_EF_CONSTRUCTION: usize = 100;
pub const DEFAULT_HNSW_EF_SEARCH: usize = 10;

// https://doc.rust-lang.org/nomicon/ffi.html#representing-opaque-structs
#[repr(C)]
Expand Down Expand Up @@ -50,10 +51,12 @@ impl ChromaError for HnswIndexFromSegmentError {
}

impl HnswIndexConfig {
pub fn from_segment(
segment: &Segment,
persist_path: &std::path::Path,
) -> Result<HnswIndexConfig, Box<HnswIndexFromSegmentError>> {
pub fn new(
m: usize,
ef_construction: usize,
ef_search: usize,
persist_path: &Path,
) -> Result<Self, Box<HnswIndexFromSegmentError>> {
let persist_path = match persist_path.to_str() {
Some(persist_path) => persist_path,
None => {
Expand All @@ -62,53 +65,11 @@ impl HnswIndexConfig {
)))
}
};
let metadata = match &segment.metadata {
Some(metadata) => metadata,
None => {
// TODO: This should error, but the configuration is not stored correctly
// after the configuration is refactored to be always stored and doesn't rely on defaults we can fix this
return Ok(HnswIndexConfig {
max_elements: DEFAULT_MAX_ELEMENTS,
m: DEFAULT_HNSW_M,
ef_construction: DEFAULT_HNSW_EF_CONSTRUCTION,
ef_search: DEFAULT_HNSW_EF_SEARCH,
random_seed: 0,
persist_path: persist_path.to_string(),
});
}
};

fn get_metadata_value_as<'a, T>(
metadata: &'a Metadata,
key: &str,
) -> Result<T, Box<HnswIndexFromSegmentError>>
where
T: TryFrom<&'a MetadataValue, Error = MetadataValueConversionError>,
{
let res = match metadata.get(key) {
Some(value) => T::try_from(value),
None => {
return Err(Box::new(HnswIndexFromSegmentError::MissingConfig(
key.to_string(),
)))
}
};
match res {
Ok(value) => Ok(value),
Err(e) => Err(Box::new(HnswIndexFromSegmentError::MetadataValueError(e))),
}
}

let m = get_metadata_value_as::<i64>(metadata, "hnsw:M").unwrap_or(DEFAULT_HNSW_M as i64);
let ef_construction = get_metadata_value_as::<i64>(metadata, "hnsw:construction_ef")
.unwrap_or(DEFAULT_HNSW_EF_CONSTRUCTION as i64);
let ef_search = get_metadata_value_as::<i64>(metadata, "hnsw:search_ef")
.unwrap_or(DEFAULT_HNSW_EF_SEARCH as i64);
Ok(HnswIndexConfig {
max_elements: DEFAULT_MAX_ELEMENTS,
m: m as usize,
ef_construction: ef_construction as usize,
ef_search: ef_search as usize,
m,
ef_construction,
ef_search,
random_seed: 0,
persist_path: persist_path.to_string(),
})
Expand Down Expand Up @@ -428,12 +389,9 @@ pub mod test {
use super::*;
use crate::utils;
use chroma_distance::DistanceFunction;
use chroma_types::CollectionUuid;
use chroma_types::SegmentUuid;
use rand::seq::IteratorRandom;
use rayon::prelude::*;
use rayon::ThreadPoolBuilder;
use std::collections::HashMap;
use tempfile::tempdir;
use uuid::Uuid;

Expand Down Expand Up @@ -826,52 +784,6 @@ pub mod test {
});
}

#[test]
fn parameter_defaults() {
let segment = Segment {
id: SegmentUuid::new(),
r#type: chroma_types::SegmentType::HnswDistributed,
scope: chroma_types::SegmentScope::VECTOR,
metadata: Some(HashMap::new()),
collection: CollectionUuid(Uuid::new_v4()),
file_path: HashMap::new(),
};

let persist_path = tempdir().unwrap().path().to_owned();
let config = HnswIndexConfig::from_segment(&segment, &persist_path)
.expect("Failed to create config from segment");

assert_eq!(config.max_elements, DEFAULT_MAX_ELEMENTS);
assert_eq!(config.m, DEFAULT_HNSW_M);
assert_eq!(config.ef_construction, DEFAULT_HNSW_EF_CONSTRUCTION);
assert_eq!(config.ef_search, DEFAULT_HNSW_EF_SEARCH);
assert_eq!(config.random_seed, 0);
assert_eq!(config.persist_path, persist_path.to_str().unwrap());

// Try partial metadata
let mut metadata = HashMap::new();
metadata.insert("hnsw:M".to_string(), MetadataValue::Int(10_i64));

let segment = Segment {
id: SegmentUuid::new(),
r#type: chroma_types::SegmentType::HnswDistributed,
scope: chroma_types::SegmentScope::VECTOR,
metadata: Some(metadata),
collection: CollectionUuid(Uuid::new_v4()),
file_path: HashMap::new(),
};

let config = HnswIndexConfig::from_segment(&segment, &persist_path)
.expect("Failed to create config from segment");

assert_eq!(config.max_elements, DEFAULT_MAX_ELEMENTS);
assert_eq!(config.m, 10);
assert_eq!(config.ef_construction, DEFAULT_HNSW_EF_CONSTRUCTION);
assert_eq!(config.ef_search, DEFAULT_HNSW_EF_SEARCH);
assert_eq!(config.random_seed, 0);
assert_eq!(config.persist_path, persist_path.to_str().unwrap());
}

#[test]
fn it_can_catch_error() {
let n = 10;
Expand Down
Loading

0 comments on commit 3f56aeb

Please sign in to comment.