chroma_types/
segment.rs

1use super::{
2    CollectionUuid, Metadata, MetadataValueConversionError, SegmentScope,
3    SegmentScopeConversionError,
4};
5use crate::{chroma_proto, DatabaseUuid};
6use chroma_error::{ChromaError, ErrorCodes};
7use std::{collections::HashMap, str::FromStr};
8use thiserror::Error;
9use tonic::Status;
10use uuid::Uuid;
11
12pub const USER_ID_TO_OFFSET_ID: &str = "user_id_to_offset_id";
13pub const OFFSET_ID_TO_USER_ID: &str = "offset_id_to_user_id";
14pub const OFFSET_ID_TO_DATA: &str = "offset_id_to_data";
15pub const MAX_OFFSET_ID: &str = "max_offset_id";
16pub const USER_ID_BLOOM_FILTER: &str = "user_id_bloom_filter";
17
18pub const FULL_TEXT_PLS: &str = "full_text_pls";
19pub const STRING_METADATA: &str = "string_metadata";
20pub const BOOL_METADATA: &str = "bool_metadata";
21pub const F32_METADATA: &str = "f32_metadata";
22pub const U32_METADATA: &str = "u32_metadata";
23
24pub const SPARSE_MAX: &str = "sparse_max";
25pub const SPARSE_OFFSET_VALUE: &str = "sparse_offset_value";
26
27pub const HNSW_PATH: &str = "hnsw_path";
28pub const VERSION_MAP_PATH: &str = "version_map_path";
29pub const POSTING_LIST_PATH: &str = "posting_list_path";
30pub const MAX_HEAD_ID_BF_PATH: &str = "max_head_id_path";
31
32pub const QUANTIZED_SPANN_CLUSTER: &str = "quantized_spann_cluster";
33pub const QUANTIZED_SPANN_SCALAR_METADATA: &str = "quantized_spann_scalar_metadata";
34pub const QUANTIZED_SPANN_EMBEDDING_METADATA: &str = "quantized_spann_embedding_metadata";
35pub const QUANTIZED_SPANN_RAW_CENTROID: &str = "quantized_spann_raw_centroid";
36pub const QUANTIZED_SPANN_QUANTIZED_CENTROID: &str = "quantized_spann_quantized_centroid";
37
38/// SegmentUuid is a wrapper around Uuid to provide a type for the segment id.
39#[derive(Copy, Clone, Debug, Default, Eq, PartialEq, Ord, PartialOrd, Hash)]
40pub struct SegmentUuid(pub Uuid);
41
42impl SegmentUuid {
43    pub fn new() -> Self {
44        SegmentUuid(Uuid::new_v4())
45    }
46}
47
48impl FromStr for SegmentUuid {
49    type Err = SegmentConversionError;
50
51    fn from_str(s: &str) -> Result<Self, SegmentConversionError> {
52        match Uuid::parse_str(s) {
53            Ok(uuid) => Ok(SegmentUuid(uuid)),
54            Err(_) => Err(SegmentConversionError::InvalidUuid),
55        }
56    }
57}
58
59impl std::fmt::Display for SegmentUuid {
60    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
61        write!(f, "{}", self.0)
62    }
63}
64
65#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
66pub enum SegmentType {
67    BlockfileMetadata,
68    BlockfileRecord,
69    HnswDistributed,
70    HnswLocalMemory,
71    HnswLocalPersisted,
72    Sqlite,
73    Spann,
74    QuantizedSpann,
75}
76
77impl From<SegmentType> for String {
78    fn from(segment_type: SegmentType) -> String {
79        match segment_type {
80            SegmentType::BlockfileMetadata => "urn:chroma:segment/metadata/blockfile".to_string(),
81            SegmentType::BlockfileRecord => "urn:chroma:segment/record/blockfile".to_string(),
82            SegmentType::HnswDistributed => {
83                "urn:chroma:segment/vector/hnsw-distributed".to_string()
84            }
85            SegmentType::HnswLocalMemory => {
86                "urn:chroma:segment/vector/hnsw-local-memory".to_string()
87            }
88            SegmentType::HnswLocalPersisted => {
89                "urn:chroma:segment/vector/hnsw-local-persisted".to_string()
90            }
91            SegmentType::Spann => "urn:chroma:segment/vector/spann".to_string(),
92            SegmentType::QuantizedSpann => "urn:chroma:segment/vector/quantized-spann".to_string(),
93            SegmentType::Sqlite => "urn:chroma:segment/metadata/sqlite".to_string(),
94        }
95    }
96}
97
98impl TryFrom<&str> for SegmentType {
99    type Error = SegmentConversionError;
100
101    fn try_from(segment_type: &str) -> Result<Self, Self::Error> {
102        match segment_type {
103            "urn:chroma:segment/metadata/blockfile" => Ok(SegmentType::BlockfileMetadata),
104            "urn:chroma:segment/record/blockfile" => Ok(SegmentType::BlockfileRecord),
105            "urn:chroma:segment/vector/hnsw-distributed" => Ok(SegmentType::HnswDistributed),
106            "urn:chroma:segment/vector/hnsw-local-memory" => Ok(SegmentType::HnswLocalMemory),
107            "urn:chroma:segment/vector/hnsw-local-persisted" => Ok(Self::HnswLocalPersisted),
108            "urn:chroma:segment/vector/spann" => Ok(SegmentType::Spann),
109            "urn:chroma:segment/vector/quantized-spann" => Ok(SegmentType::QuantizedSpann),
110            "urn:chroma:segment/metadata/sqlite" => Ok(SegmentType::Sqlite),
111            _ => Err(SegmentConversionError::InvalidSegmentType),
112        }
113    }
114}
115
116#[derive(Clone, Debug, PartialEq)]
117pub struct Segment {
118    pub id: SegmentUuid,
119    pub r#type: SegmentType,
120    pub scope: SegmentScope,
121    pub collection: CollectionUuid,
122    pub metadata: Option<Metadata>,
123    pub file_path: HashMap<String, Vec<String>>,
124}
125
126impl Segment {
127    // INVARIANT: THIS ALWAYS RETURNS AT LEAST ONE SHARD
128    pub fn get_shards(&self) -> Result<Vec<SegmentShard>, SegmentShardError> {
129        // If there are no file paths, return empty vector
130        if self.file_path.is_empty() {
131            let vec = vec![self.new_shard()];
132            return Ok(vec);
133        }
134
135        // Check that all file-path keys have the same number of shards
136        let counts: std::collections::HashSet<usize> =
137            self.file_path.values().map(|v| v.len()).collect();
138        if counts.len() > 1 {
139            let first_key = self.file_path.keys().next().unwrap();
140            let expected = self.file_path[first_key].len();
141            return Err(SegmentShardError::MismatchedShardCounts {
142                key: format!("segment {}", self.id),
143                actual: counts.into_iter().next().unwrap(),
144                expected,
145            });
146        }
147
148        // All paths should have the same number of shards
149        let num_shards = self
150            .file_path
151            .values()
152            .next()
153            .map(|paths| paths.len())
154            .unwrap_or(0);
155
156        // Create a SegmentShard for each shard index, propagating any errors
157        let shards: Result<Vec<SegmentShard>, SegmentShardError> = (0..num_shards)
158            .map(|shard_index| SegmentShard::try_from((self, shard_index as u32)))
159            .collect();
160
161        shards
162    }
163
164    pub fn prefetch_supported(&self) -> bool {
165        matches!(
166            self.r#type,
167            SegmentType::BlockfileMetadata
168                | SegmentType::BlockfileRecord
169                | SegmentType::QuantizedSpann
170                | SegmentType::Spann
171        )
172    }
173
174    /// Returns the file paths that should be prefetched for this segment.
175    /// If shard_index is None, returns the active shard's file paths. If shard_index is Some, returns
176    /// only the file paths for that shard.
177    pub fn filepaths_to_prefetch(&self, shard_index: Option<u32>) -> Vec<String> {
178        let mut res = Vec::new();
179        match self.r#type {
180            SegmentType::QuantizedSpann => {
181                for key in [
182                    QUANTIZED_SPANN_CLUSTER,
183                    QUANTIZED_SPANN_EMBEDDING_METADATA,
184                    QUANTIZED_SPANN_SCALAR_METADATA,
185                ] {
186                    if let Some(paths) = self.file_path.get(key) {
187                        if let Some(path) = match shard_index {
188                            Some(index) => paths.get(index as usize),
189                            None => paths.last(),
190                        } {
191                            res.push(path.clone());
192                        }
193                    }
194                }
195            }
196            SegmentType::Spann => {
197                if let Some(pl_path) = self.file_path.get(POSTING_LIST_PATH) {
198                    if let Some(path) = match shard_index {
199                        Some(index) => pl_path.get(index as usize),
200                        None => pl_path.last(),
201                    } {
202                        res.push(path.clone());
203                    }
204                }
205            }
206            SegmentType::BlockfileMetadata | SegmentType::BlockfileRecord => {
207                for (key, paths) in &self.file_path {
208                    if key == USER_ID_BLOOM_FILTER {
209                        continue;
210                    }
211                    if let Some(path) = match shard_index {
212                        Some(index) => paths.get(index as usize),
213                        None => paths.last(),
214                    } {
215                        res.push(path.clone());
216                    }
217                }
218            }
219            _ => {}
220        }
221        res
222    }
223
224    /// Returns the number of shards for this segment.
225    /// Derives from the length of file_path Vecs (all must be equal per the
226    /// shard count invariant). Returns 1 if no file paths are present.
227    pub fn num_shards(&self) -> Result<usize, SegmentShardError> {
228        let mut values = self.file_path.values();
229        let num_shards = match values.next() {
230            Some(paths) => paths.len().max(1),
231            None => return Ok(1),
232        };
233        for (key, paths) in &self.file_path {
234            if paths.len() != num_shards {
235                return Err(SegmentShardError::MismatchedShardCounts {
236                    key: key.clone(),
237                    actual: paths.len(),
238                    expected: num_shards,
239                });
240            }
241        }
242        Ok(num_shards)
243    }
244
245    pub fn extract_prefix_and_id(path: &str) -> Result<(&str, uuid::Uuid), uuid::Error> {
246        let (prefix, id) = match path.rfind('/') {
247            Some(pos) => (&path[..pos], &path[pos + 1..]),
248            None => ("", path),
249        };
250        match Uuid::try_parse(id) {
251            Ok(uid) => Ok((prefix, uid)),
252            Err(e) => Err(e),
253        }
254    }
255
256    pub fn construct_prefix_path(&self, tenant: &str, database_id: &DatabaseUuid) -> String {
257        Self::construct_prefix_path_impl(tenant, database_id, &self.collection, &self.id)
258    }
259
260    fn construct_prefix_path_impl(
261        tenant: &str,
262        database_id: &DatabaseUuid,
263        collection: &CollectionUuid,
264        segment_id: &SegmentUuid,
265    ) -> String {
266        format!(
267            "tenant/{}/database/{}/collection/{}/segment/{}",
268            tenant, database_id, collection, segment_id
269        )
270    }
271}
272
273#[derive(Clone, Debug, PartialEq)]
274pub struct SegmentShard {
275    pub id: SegmentUuid,
276    pub r#type: SegmentType,
277    pub scope: SegmentScope,
278    pub collection: CollectionUuid,
279    pub metadata: Option<Metadata>,
280    pub file_path: HashMap<String, String>,
281}
282
283impl SegmentShard {
284    pub fn construct_prefix_path(&self, tenant: &str, database_id: &DatabaseUuid) -> String {
285        Segment::construct_prefix_path_impl(tenant, database_id, &self.collection, &self.id)
286    }
287}
288
289#[derive(Error, Debug)]
290pub enum SegmentShardError {
291    #[error("Empty path vector for key '{0}'")]
292    EmptyPathVector(String),
293    #[error("Empty path string for key '{0}'")]
294    EmptyPathString(String),
295    #[error("Shard index {index} out of bounds for key '{key}' (len {len})")]
296    ShardIndexOutOfBounds { key: String, index: u32, len: usize },
297    #[error("Mismatched shard counts: key '{key}' has {actual} entries, expected {expected}")]
298    MismatchedShardCounts {
299        key: String,
300        actual: usize,
301        expected: usize,
302    },
303    #[error("No shards found")]
304    EmptyShards,
305}
306
307impl ChromaError for SegmentShardError {
308    fn code(&self) -> ErrorCodes {
309        ErrorCodes::Internal
310    }
311}
312
313impl TryFrom<(&Segment, u32)> for SegmentShard {
314    type Error = SegmentShardError;
315
316    fn try_from((segment, shard_index): (&Segment, u32)) -> Result<Self, Self::Error> {
317        let mut file_path = HashMap::new();
318        for (key, paths) in &segment.file_path {
319            match paths.get(shard_index as usize) {
320                Some(path) if path.is_empty() => {
321                    return Err(SegmentShardError::EmptyPathString(key.clone()));
322                }
323                Some(path) => {
324                    file_path.insert(key.clone(), path.clone());
325                }
326                None if paths.is_empty() => {
327                    return Err(SegmentShardError::EmptyPathVector(key.clone()));
328                }
329                None => {
330                    return Err(SegmentShardError::ShardIndexOutOfBounds {
331                        key: key.clone(),
332                        index: shard_index,
333                        len: paths.len(),
334                    });
335                }
336            }
337        }
338
339        Ok(SegmentShard {
340            id: segment.id,
341            r#type: segment.r#type,
342            scope: segment.scope.clone(),
343            collection: segment.collection,
344            metadata: segment.metadata.clone(),
345            file_path,
346        })
347    }
348}
349
350impl Segment {
351    pub fn new_shard(&self) -> SegmentShard {
352        SegmentShard {
353            id: self.id,
354            r#type: self.r#type,
355            scope: self.scope.clone(),
356            collection: self.collection,
357            metadata: self.metadata.clone(),
358            file_path: HashMap::new(),
359        }
360    }
361}
362
363#[derive(Error, Debug)]
364pub enum SegmentConversionError {
365    #[error("Invalid UUID")]
366    InvalidUuid,
367    #[error(transparent)]
368    MetadataValueConversionError(#[from] MetadataValueConversionError),
369    #[error(transparent)]
370    SegmentScopeConversionError(#[from] SegmentScopeConversionError),
371    #[error("Invalid segment type")]
372    InvalidSegmentType,
373}
374
375impl ChromaError for SegmentConversionError {
376    fn code(&self) -> ErrorCodes {
377        match self {
378            SegmentConversionError::InvalidUuid => ErrorCodes::InvalidArgument,
379            SegmentConversionError::InvalidSegmentType => ErrorCodes::InvalidArgument,
380            SegmentConversionError::SegmentScopeConversionError(e) => e.code(),
381            SegmentConversionError::MetadataValueConversionError(e) => e.code(),
382        }
383    }
384}
385
386impl From<SegmentConversionError> for Status {
387    fn from(value: SegmentConversionError) -> Self {
388        Status::invalid_argument(value.to_string())
389    }
390}
391
392impl TryFrom<chroma_proto::Segment> for Segment {
393    type Error = SegmentConversionError;
394
395    fn try_from(proto_segment: chroma_proto::Segment) -> Result<Self, Self::Error> {
396        let mut proto_segment = proto_segment;
397
398        let segment_uuid = match SegmentUuid::from_str(&proto_segment.id) {
399            Ok(uuid) => uuid,
400            Err(_) => return Err(SegmentConversionError::InvalidUuid),
401        };
402        let collection_uuid = match Uuid::try_parse(&proto_segment.collection) {
403            Ok(uuid) => uuid,
404            Err(_) => return Err(SegmentConversionError::InvalidUuid),
405        };
406        let collection_uuid = CollectionUuid(collection_uuid);
407        let segment_metadata: Option<Metadata> = match proto_segment.metadata {
408            Some(proto_metadata) => match proto_metadata.try_into() {
409                Ok(metadata) => Some(metadata),
410                Err(e) => return Err(SegmentConversionError::MetadataValueConversionError(e)),
411            },
412            None => None,
413        };
414        let scope: SegmentScope = match proto_segment.scope.try_into() {
415            Ok(scope) => scope,
416            Err(e) => return Err(SegmentConversionError::SegmentScopeConversionError(e)),
417        };
418
419        let segment_type: SegmentType = proto_segment.r#type.as_str().try_into()?;
420
421        let mut file_paths = HashMap::new();
422        let drain = proto_segment.file_paths.drain();
423        for (key, value) in drain {
424            file_paths.insert(key, value.paths);
425        }
426
427        Ok(Segment {
428            id: segment_uuid,
429            r#type: segment_type,
430            scope,
431            collection: collection_uuid,
432            metadata: segment_metadata,
433            file_path: file_paths,
434        })
435    }
436}
437
438impl From<Segment> for chroma_proto::Segment {
439    fn from(value: Segment) -> Self {
440        Self {
441            id: value.id.0.to_string(),
442            r#type: value.r#type.into(),
443            scope: chroma_proto::SegmentScope::from(value.scope) as i32,
444            collection: value.collection.0.to_string(),
445            metadata: value.metadata.map(Into::into),
446            file_paths: value
447                .file_path
448                .into_iter()
449                .map(|(name, paths)| (name, chroma_proto::FilePaths { paths }))
450                .collect(),
451        }
452    }
453}
454
455pub fn test_segment(collection_uuid: CollectionUuid, scope: SegmentScope) -> Segment {
456    let r#type = match scope {
457        SegmentScope::METADATA => SegmentType::BlockfileMetadata,
458        SegmentScope::RECORD => SegmentType::BlockfileRecord,
459        SegmentScope::VECTOR => SegmentType::HnswDistributed,
460        SegmentScope::SQLITE => unimplemented!("Sqlite segment is not implemented"),
461    };
462    Segment {
463        id: SegmentUuid::new(),
464        r#type,
465        scope,
466        collection: collection_uuid,
467        metadata: None,
468        file_path: HashMap::new(),
469    }
470}
471
472#[cfg(test)]
473mod tests {
474
475    use super::*;
476    use crate::MetadataValue;
477
478    #[test]
479    fn test_segment_try_from() {
480        let mut metadata = chroma_proto::UpdateMetadata {
481            metadata: HashMap::new(),
482        };
483        metadata.metadata.insert(
484            "foo".to_string(),
485            chroma_proto::UpdateMetadataValue {
486                value: Some(chroma_proto::update_metadata_value::Value::IntValue(42)),
487            },
488        );
489        let proto_segment = chroma_proto::Segment {
490            id: "00000000-0000-0000-0000-000000000000".to_string(),
491            r#type: "urn:chroma:segment/vector/hnsw-distributed".to_string(),
492            scope: chroma_proto::SegmentScope::Vector as i32,
493            collection: "00000000-0000-0000-0000-000000000000".to_string(),
494            metadata: Some(metadata),
495            file_paths: HashMap::new(),
496        };
497        let converted_segment: Segment = proto_segment.try_into().unwrap();
498        assert_eq!(converted_segment.id, SegmentUuid(Uuid::nil()));
499        assert_eq!(converted_segment.r#type, SegmentType::HnswDistributed);
500        assert_eq!(converted_segment.scope, SegmentScope::VECTOR);
501        assert_eq!(converted_segment.collection, CollectionUuid(Uuid::nil()));
502        let metadata = converted_segment.metadata.unwrap();
503        assert_eq!(metadata.len(), 1);
504        assert_eq!(metadata.get("foo").unwrap(), &MetadataValue::Int(42));
505    }
506
507    #[test]
508    fn test_segment_construct_prefix_path() {
509        let segment = Segment {
510            id: SegmentUuid(Uuid::nil()),
511            r#type: SegmentType::BlockfileMetadata,
512            scope: SegmentScope::METADATA,
513            collection: CollectionUuid(Uuid::nil()),
514            metadata: None,
515            file_path: HashMap::new(),
516        };
517        let tenant = "test_tenant";
518        let database_id = &DatabaseUuid(Uuid::nil());
519        let prefix_path = segment.construct_prefix_path(tenant, database_id);
520        assert_eq!(
521            prefix_path,
522            "tenant/test_tenant/database/00000000-0000-0000-0000-000000000000/collection/00000000-0000-0000-0000-000000000000/segment/00000000-0000-0000-0000-000000000000"
523        );
524    }
525
526    #[test]
527    fn test_segment_extract_prefix_and_id() {
528        let path = "tenant/test_tenant/database/00000000-0000-0000-0000-000000000000/collection/00000000-0000-0000-0000-000000000000/segment/00000000-0000-0000-0000-000000000000/00000000-0000-0000-0000-000000000001";
529        let (prefix, id) =
530            Segment::extract_prefix_and_id(path).expect("Failed to extract prefix and id");
531        assert_eq!(
532            prefix,
533            "tenant/test_tenant/database/00000000-0000-0000-0000-000000000000/collection/00000000-0000-0000-0000-000000000000/segment/00000000-0000-0000-0000-000000000000"
534        );
535        assert_eq!(
536            id,
537            Uuid::from_str("00000000-0000-0000-0000-000000000001").expect("Cannot happen")
538        );
539    }
540
541    #[test]
542    fn test_segment_extract_prefix_and_id_legacy() {
543        let path = "00000000-0000-0000-0000-000000000001";
544        let (prefix, id) =
545            Segment::extract_prefix_and_id(path).expect("Failed to extract prefix and id");
546        assert_eq!(prefix, "");
547        assert_eq!(
548            id,
549            Uuid::from_str("00000000-0000-0000-0000-000000000001").expect("Cannot happen")
550        );
551    }
552
553    #[test]
554    fn test_num_shards_empty_file_path() {
555        let segment = Segment {
556            id: SegmentUuid(Uuid::nil()),
557            r#type: SegmentType::BlockfileRecord,
558            scope: SegmentScope::RECORD,
559            collection: CollectionUuid(Uuid::nil()),
560            metadata: None,
561            file_path: HashMap::new(),
562        };
563        assert_eq!(segment.num_shards().unwrap(), 1);
564    }
565
566    #[test]
567    fn test_num_shards_single_shard() {
568        let mut file_path = HashMap::new();
569        file_path.insert("key_a".to_string(), vec!["path_a_0".to_string()]);
570        file_path.insert("key_b".to_string(), vec!["path_b_0".to_string()]);
571        let segment = Segment {
572            id: SegmentUuid(Uuid::nil()),
573            r#type: SegmentType::BlockfileRecord,
574            scope: SegmentScope::RECORD,
575            collection: CollectionUuid(Uuid::nil()),
576            metadata: None,
577            file_path,
578        };
579        assert_eq!(segment.num_shards().unwrap(), 1);
580    }
581
582    #[test]
583    fn test_num_shards_multi_shard() {
584        let mut file_path = HashMap::new();
585        file_path.insert(
586            "key_a".to_string(),
587            vec!["a0".to_string(), "a1".to_string(), "a2".to_string()],
588        );
589        file_path.insert(
590            "key_b".to_string(),
591            vec!["b0".to_string(), "b1".to_string(), "b2".to_string()],
592        );
593        let segment = Segment {
594            id: SegmentUuid(Uuid::nil()),
595            r#type: SegmentType::BlockfileRecord,
596            scope: SegmentScope::RECORD,
597            collection: CollectionUuid(Uuid::nil()),
598            metadata: None,
599            file_path,
600        };
601        assert_eq!(segment.num_shards().unwrap(), 3);
602    }
603
604    #[test]
605    fn test_num_shards_mismatched_lengths() {
606        let mut file_path = HashMap::new();
607        file_path.insert(
608            "key_a".to_string(),
609            vec!["a0".to_string(), "a1".to_string()],
610        );
611        file_path.insert(
612            "key_b".to_string(),
613            vec!["b0".to_string(), "b1".to_string(), "b2".to_string()],
614        );
615        let segment = Segment {
616            id: SegmentUuid(Uuid::nil()),
617            r#type: SegmentType::BlockfileRecord,
618            scope: SegmentScope::RECORD,
619            collection: CollectionUuid(Uuid::nil()),
620            metadata: None,
621            file_path,
622        };
623        let err = segment.num_shards().unwrap_err();
624        assert!(matches!(
625            err,
626            SegmentShardError::MismatchedShardCounts { .. }
627        ));
628    }
629}
chroma_types/segment.rs

chroma_types/
segment.rs
