@@ -8,7 +8,9 @@ use chroma_index::hnsw_provider::{
88 HnswIndexProviderOpenError , HnswIndexRef ,
99} ;
1010use chroma_index:: { Index , IndexUuid } ;
11- use chroma_types:: { Collection , HnswParametersFromSegmentError , SegmentUuid } ;
11+ use chroma_types:: {
12+ Collection , HnswParametersFromSegmentError , KnnIndex , Schema , SchemaError , SegmentUuid ,
13+ } ;
1214use chroma_types:: { MaterializedLogOperation , Segment } ;
1315use std:: collections:: HashMap ;
1416use std:: fmt:: Debug ;
@@ -55,6 +57,8 @@ pub enum DistributedHNSWSegmentFromSegmentError {
5557 MissingHnswConfiguration ,
5658 #[ error( "Could not parse HNSW configuration: {0}" ) ]
5759 InvalidHnswConfiguration ( #[ from] HnswParametersFromSegmentError ) ,
60+ #[ error( "Invalid schema: {0}" ) ]
61+ InvalidSchema ( #[ source] SchemaError ) ,
5862}
5963
6064impl ChromaError for DistributedHNSWSegmentFromSegmentError {
@@ -72,6 +76,7 @@ impl ChromaError for DistributedHNSWSegmentFromSegmentError {
7276 DistributedHNSWSegmentFromSegmentError :: InvalidHnswConfiguration ( _) => {
7377 ErrorCodes :: Internal
7478 }
79+ DistributedHNSWSegmentFromSegmentError :: InvalidSchema ( e) => e. code ( ) ,
7580 }
7681 }
7782}
@@ -96,9 +101,14 @@ impl DistributedHNSWSegmentWriter {
96101 hnsw_index_provider : HnswIndexProvider ,
97102 ) -> Result < Box < DistributedHNSWSegmentWriter > , Box < DistributedHNSWSegmentFromSegmentError > >
98103 {
99- let hnsw_configuration = collection
100- . config
101- . get_hnsw_config_with_legacy_fallback ( segment)
104+ let schema = if let Some ( schema) = & collection. schema {
105+ schema
106+ } else {
107+ & Schema :: convert_collection_config_to_schema ( & collection. config , KnnIndex :: Hnsw )
108+ . map_err ( DistributedHNSWSegmentFromSegmentError :: InvalidSchema ) ?
109+ } ;
110+ let hnsw_configuration = schema
111+ . get_internal_hnsw_config_with_legacy_fallback ( segment)
102112 . map_err ( DistributedHNSWSegmentFromSegmentError :: InvalidHnswConfiguration ) ?
103113 . ok_or ( DistributedHNSWSegmentFromSegmentError :: MissingHnswConfiguration ) ?;
104114
@@ -313,9 +323,13 @@ impl DistributedHNSWSegmentReader {
313323 hnsw_index_provider : HnswIndexProvider ,
314324 ) -> Result < Box < DistributedHNSWSegmentReader > , Box < DistributedHNSWSegmentFromSegmentError > >
315325 {
316- let hnsw_configuration = collection
317- . config
318- . get_hnsw_config_with_legacy_fallback ( segment)
326+ let schema = collection. schema . as_ref ( ) . ok_or_else ( || {
327+ DistributedHNSWSegmentFromSegmentError :: InvalidSchema ( SchemaError :: InvalidSchema {
328+ reason : "Schema is None" . to_string ( ) ,
329+ } )
330+ } ) ?;
331+ let hnsw_configuration = schema
332+ . get_internal_hnsw_config_with_legacy_fallback ( segment)
319333 . map_err ( DistributedHNSWSegmentFromSegmentError :: InvalidHnswConfiguration ) ?
320334 . ok_or ( DistributedHNSWSegmentFromSegmentError :: MissingHnswConfiguration ) ?;
321335
@@ -394,7 +408,7 @@ pub mod test {
394408 use chroma_index:: { HnswIndexConfig , DEFAULT_MAX_ELEMENTS } ;
395409 use chroma_types:: {
396410 Collection , CollectionUuid , InternalCollectionConfiguration , InternalHnswConfiguration ,
397- Segment , SegmentUuid ,
411+ KnnIndex , Schema , Segment , SegmentUuid ,
398412 } ;
399413 use tempfile:: tempdir;
400414 use uuid:: Uuid ;
@@ -423,18 +437,20 @@ pub mod test {
423437 config. persist_path,
424438 Some ( persist_path. to_str( ) . unwrap( ) . to_string( ) )
425439 ) ;
440+ let config = InternalCollectionConfiguration {
441+ vector_index : chroma_types:: VectorIndexConfiguration :: Hnsw ( InternalHnswConfiguration {
442+ max_neighbors : 10 ,
443+ ..Default :: default ( )
444+ } ) ,
445+ embedding_function : None ,
446+ } ;
426447
427448 // Try partial override
428449 let collection = Collection {
429- config : InternalCollectionConfiguration {
430- vector_index : chroma_types:: VectorIndexConfiguration :: Hnsw (
431- InternalHnswConfiguration {
432- max_neighbors : 10 ,
433- ..Default :: default ( )
434- } ,
435- ) ,
436- embedding_function : None ,
437- } ,
450+ config : config. clone ( ) ,
451+ schema : Some (
452+ Schema :: convert_collection_config_to_schema ( & config, KnnIndex :: Hnsw ) . unwrap ( ) ,
453+ ) ,
438454 ..Default :: default ( )
439455 } ;
440456
@@ -448,9 +464,12 @@ pub mod test {
448464 } ;
449465
450466 let hnsw_params = collection
451- . config
452- . get_hnsw_config_with_legacy_fallback ( & segment)
467+ . schema
468+ . as_ref ( )
469+ . map ( |schema| schema. get_internal_hnsw_config_with_legacy_fallback ( & segment) )
470+ . transpose ( )
453471 . unwrap ( )
472+ . flatten ( )
454473 . unwrap ( ) ;
455474 let config = HnswIndexConfig :: new_persistent (
456475 hnsw_params. max_neighbors ,
0 commit comments