chroma_types/
segment.rs

1use super::{
2    CollectionUuid, Metadata, MetadataValueConversionError, SegmentScope,
3    SegmentScopeConversionError,
4};
5use crate::collection_schema::Schema;
6use crate::{chroma_proto, DatabaseUuid};
7use chroma_error::{ChromaError, ErrorCodes};
8use chroma_proto::CollectionVersionInfo;
9use std::{collections::HashMap, str::FromStr};
10use thiserror::Error;
11use tonic::Status;
12use uuid::Uuid;
13
14pub const USER_ID_TO_OFFSET_ID: &str = "user_id_to_offset_id";
15pub const OFFSET_ID_TO_USER_ID: &str = "offset_id_to_user_id";
16pub const OFFSET_ID_TO_DATA: &str = "offset_id_to_data";
17pub const MAX_OFFSET_ID: &str = "max_offset_id";
18pub const USER_ID_BLOOM_FILTER: &str = "user_id_bloom_filter";
19
20pub const FULL_TEXT_PLS: &str = "full_text_pls";
21pub const FULL_TEXT_TOKEN: &str = "full_text_token";
22pub const STRING_METADATA: &str = "string_metadata";
23pub const BOOL_METADATA: &str = "bool_metadata";
24pub const F32_METADATA: &str = "f32_metadata";
25pub const U32_METADATA: &str = "u32_metadata";
26
27pub const SPARSE_MAX: &str = "sparse_max";
28pub const SPARSE_OFFSET_VALUE: &str = "sparse_offset_value";
29pub const SPARSE_POSTING: &str = "sparse_posting";
30
31/// Delimiter separating the sparse-index file_path family (one of the
32/// `SPARSE_*` constants) from the metadata key it indexes, e.g.
33/// `sparse_posting::my_field`. The `file_path` map key is internal to Chroma
34/// (never used as a storage path component), so this delimiter only needs to
35/// be unambiguous when parsing the family back out.
36pub const SPARSE_FILE_PATH_DELIMITER: &str = "::";
37
38/// Which sparse-index blockfile family a `file_path` map key refers to.
39#[derive(Clone, Copy, Debug, PartialEq, Eq)]
40pub enum SparseFilePathKind {
41    /// MaxScore posting blockfile (`SPARSE_POSTING`).
42    Posting,
43    /// WAND block-max blockfile (`SPARSE_MAX`).
44    Max,
45    /// WAND offset-value blockfile (`SPARSE_OFFSET_VALUE`).
46    OffsetValue,
47}
48
49/// Per-key `file_path` map key for the MaxScore posting blockfile.
50pub fn sparse_posting_key(metadata_key: &str) -> String {
51    format!("{SPARSE_POSTING}{SPARSE_FILE_PATH_DELIMITER}{metadata_key}")
52}
53
54/// Per-key `file_path` map key for the WAND block-max blockfile.
55pub fn sparse_max_key(metadata_key: &str) -> String {
56    format!("{SPARSE_MAX}{SPARSE_FILE_PATH_DELIMITER}{metadata_key}")
57}
58
59/// Per-key `file_path` map key for the WAND offset-value blockfile.
60pub fn sparse_offset_value_key(metadata_key: &str) -> String {
61    format!("{SPARSE_OFFSET_VALUE}{SPARSE_FILE_PATH_DELIMITER}{metadata_key}")
62}
63
64/// Classify a `file_path` map key as a sparse-index entry, returning its
65/// family and the metadata key it indexes. Returns `None` for non-sparse
66/// keys. The metadata key is `None` for the legacy/global anonymous layout
67/// (bare `SPARSE_*` constants written before per-key indexing existed) and
68/// `Some(key)` for the per-key layout (`SPARSE_*::key`).
69pub fn parse_sparse_file_path_key(name: &str) -> Option<(SparseFilePathKind, Option<String>)> {
70    match name {
71        SPARSE_POSTING => return Some((SparseFilePathKind::Posting, None)),
72        SPARSE_MAX => return Some((SparseFilePathKind::Max, None)),
73        SPARSE_OFFSET_VALUE => return Some((SparseFilePathKind::OffsetValue, None)),
74        _ => {}
75    }
76    // Order is irrelevant: none of the `SPARSE_*` constants (with the
77    // trailing delimiter) is a prefix of another.
78    for (family, kind) in [
79        (SPARSE_POSTING, SparseFilePathKind::Posting),
80        (SPARSE_OFFSET_VALUE, SparseFilePathKind::OffsetValue),
81        (SPARSE_MAX, SparseFilePathKind::Max),
82    ] {
83        let with_delimiter = format!("{family}{SPARSE_FILE_PATH_DELIMITER}");
84        if let Some(metadata_key) = name.strip_prefix(&with_delimiter) {
85            return Some((kind, Some(metadata_key.to_string())));
86        }
87    }
88    None
89}
90
91pub const HNSW_PATH: &str = "hnsw_path";
92pub const VERSION_MAP_PATH: &str = "version_map_path";
93pub const POSTING_LIST_PATH: &str = "posting_list_path";
94pub const MAX_HEAD_ID_BF_PATH: &str = "max_head_id_path";
95
96pub const QUANTIZED_SPANN_CLUSTER: &str = "quantized_spann_cluster";
97pub const QUANTIZED_SPANN_SCALAR_METADATA: &str = "quantized_spann_scalar_metadata";
98pub const QUANTIZED_SPANN_EMBEDDING_METADATA: &str = "quantized_spann_embedding_metadata";
99pub const QUANTIZED_SPANN_RAW_CENTROID: &str = "quantized_spann_raw_centroid";
100pub const QUANTIZED_SPANN_QUANTIZED_CENTROID: &str = "quantized_spann_quantized_centroid";
101
102/// SegmentUuid is a wrapper around Uuid to provide a type for the segment id.
103#[derive(Copy, Clone, Debug, Default, Eq, PartialEq, Ord, PartialOrd, Hash)]
104pub struct SegmentUuid(pub Uuid);
105
106impl SegmentUuid {
107    pub fn new() -> Self {
108        SegmentUuid(Uuid::new_v4())
109    }
110}
111
112impl FromStr for SegmentUuid {
113    type Err = SegmentConversionError;
114
115    fn from_str(s: &str) -> Result<Self, SegmentConversionError> {
116        match Uuid::parse_str(s) {
117            Ok(uuid) => Ok(SegmentUuid(uuid)),
118            Err(_) => Err(SegmentConversionError::InvalidUuid),
119        }
120    }
121}
122
123impl std::fmt::Display for SegmentUuid {
124    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
125        write!(f, "{}", self.0)
126    }
127}
128
129#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
130pub enum SegmentType {
131    BlockfileMetadata,
132    BlockfileRecord,
133    HnswDistributed,
134    HnswLocalMemory,
135    HnswLocalPersisted,
136    Sqlite,
137    Spann,
138    QuantizedSpann,
139}
140
141impl From<SegmentType> for String {
142    fn from(segment_type: SegmentType) -> String {
143        match segment_type {
144            SegmentType::BlockfileMetadata => "urn:chroma:segment/metadata/blockfile".to_string(),
145            SegmentType::BlockfileRecord => "urn:chroma:segment/record/blockfile".to_string(),
146            SegmentType::HnswDistributed => {
147                "urn:chroma:segment/vector/hnsw-distributed".to_string()
148            }
149            SegmentType::HnswLocalMemory => {
150                "urn:chroma:segment/vector/hnsw-local-memory".to_string()
151            }
152            SegmentType::HnswLocalPersisted => {
153                "urn:chroma:segment/vector/hnsw-local-persisted".to_string()
154            }
155            SegmentType::Spann => "urn:chroma:segment/vector/spann".to_string(),
156            SegmentType::QuantizedSpann => "urn:chroma:segment/vector/quantized-spann".to_string(),
157            SegmentType::Sqlite => "urn:chroma:segment/metadata/sqlite".to_string(),
158        }
159    }
160}
161
162impl TryFrom<&str> for SegmentType {
163    type Error = SegmentConversionError;
164
165    fn try_from(segment_type: &str) -> Result<Self, Self::Error> {
166        match segment_type {
167            "urn:chroma:segment/metadata/blockfile" => Ok(SegmentType::BlockfileMetadata),
168            "urn:chroma:segment/record/blockfile" => Ok(SegmentType::BlockfileRecord),
169            "urn:chroma:segment/vector/hnsw-distributed" => Ok(SegmentType::HnswDistributed),
170            "urn:chroma:segment/vector/hnsw-local-memory" => Ok(SegmentType::HnswLocalMemory),
171            "urn:chroma:segment/vector/hnsw-local-persisted" => Ok(Self::HnswLocalPersisted),
172            "urn:chroma:segment/vector/spann" => Ok(SegmentType::Spann),
173            "urn:chroma:segment/vector/quantized-spann" => Ok(SegmentType::QuantizedSpann),
174            "urn:chroma:segment/metadata/sqlite" => Ok(SegmentType::Sqlite),
175            _ => Err(SegmentConversionError::InvalidSegmentType),
176        }
177    }
178}
179
180#[derive(Clone, Debug, PartialEq)]
181pub struct Segment {
182    pub id: SegmentUuid,
183    pub r#type: SegmentType,
184    pub scope: SegmentScope,
185    pub collection: CollectionUuid,
186    pub metadata: Option<Metadata>,
187    pub file_path: HashMap<String, Vec<String>>,
188}
189
190impl Segment {
191    /// Returns this segment with its identity/metadata preserved but without
192    /// any on-disk file paths.
193    pub fn empty_segment(&self) -> Self {
194        let mut segment = self.clone();
195        segment.file_path.clear();
196        segment
197    }
198
199    pub fn historical_segment_for_version(
200        &self,
201        version: &CollectionVersionInfo,
202        segment_id: SegmentUuid,
203    ) -> Result<Self, String> {
204        let mut historical_segment = self.clone();
205        let all_segments_info = version.segment_info.as_ref().ok_or_else(|| {
206            format!(
207                "Invariant violation: collection version {} is missing segment info",
208                version.version
209            )
210        })?;
211
212        let segment_info = all_segments_info
213            .segment_compaction_info
214            .iter()
215            .find(|segment_info| segment_info.segment_id == segment_id.to_string())
216            .ok_or_else(|| {
217                format!(
218                    "Invariant violation: collection version {} is missing segment {}",
219                    version.version, segment_id
220                )
221            })?;
222
223        historical_segment.file_path = segment_info
224            .file_paths
225            .iter()
226            .map(|(key, value)| (key.clone(), value.paths.clone()))
227            .collect::<HashMap<_, _>>();
228        Ok(historical_segment)
229    }
230
231    // INVARIANT: THIS ALWAYS RETURNS AT LEAST ONE SHARD
232    pub fn get_shards(&self) -> Result<Vec<SegmentShard>, SegmentShardError> {
233        let num_shards = self.num_shards()?;
234
235        // Create a SegmentShard for each shard index, propagating any errors
236        let shards: Result<Vec<SegmentShard>, SegmentShardError> = (0..num_shards)
237            .map(|shard_index| SegmentShard::try_from((self, shard_index as u32)))
238            .collect();
239
240        shards
241    }
242
243    pub fn prefetch_supported(&self) -> bool {
244        matches!(
245            self.r#type,
246            SegmentType::BlockfileMetadata
247                | SegmentType::BlockfileRecord
248                | SegmentType::QuantizedSpann
249                | SegmentType::Spann
250        )
251    }
252
253    /// Returns the file paths that should be prefetched for this segment.
254    /// If shard_index is None, returns the active shard's file paths. If shard_index is Some, returns
255    /// only the file paths for that shard.
256    pub fn filepaths_to_prefetch(&self, shard_index: Option<u32>) -> Vec<String> {
257        let mut res = Vec::new();
258        match self.r#type {
259            SegmentType::QuantizedSpann => {
260                for key in [
261                    QUANTIZED_SPANN_CLUSTER,
262                    QUANTIZED_SPANN_EMBEDDING_METADATA,
263                    QUANTIZED_SPANN_SCALAR_METADATA,
264                ] {
265                    if let Some(paths) = self.file_path.get(key) {
266                        if let Some(path) = match shard_index {
267                            Some(index) => paths.get(index as usize),
268                            None => paths.last(),
269                        } {
270                            res.push(path.clone());
271                        }
272                    }
273                }
274            }
275            SegmentType::Spann => {
276                if let Some(pl_path) = self.file_path.get(POSTING_LIST_PATH) {
277                    if let Some(path) = match shard_index {
278                        Some(index) => pl_path.get(index as usize),
279                        None => pl_path.last(),
280                    } {
281                        res.push(path.clone());
282                    }
283                }
284            }
285            SegmentType::BlockfileMetadata | SegmentType::BlockfileRecord => {
286                for (key, paths) in &self.file_path {
287                    if key == USER_ID_BLOOM_FILTER {
288                        continue;
289                    }
290                    if let Some(path) = match shard_index {
291                        Some(index) => paths.get(index as usize),
292                        None => paths.last(),
293                    } {
294                        res.push(path.clone());
295                    }
296                }
297            }
298            _ => {}
299        }
300        res
301    }
302
303    /// Returns the number of shards for this segment.
304    /// Derives from the length of file_path Vecs (all must be equal per the
305    /// shard count invariant). Returns 1 if no file paths are present.
306    pub fn num_shards(&self) -> Result<usize, SegmentShardError> {
307        let mut values = self.file_path.values();
308        let num_shards = match values.next() {
309            Some(paths) => paths.len().max(1),
310            None => return Ok(1),
311        };
312        for (key, paths) in &self.file_path {
313            if paths.len() != num_shards {
314                return Err(SegmentShardError::MismatchedShardCounts {
315                    key: key.clone(),
316                    actual: paths.len(),
317                    expected: num_shards,
318                });
319            }
320        }
321        Ok(num_shards)
322    }
323
324    /// Clears file paths for a specific shard index.
325    /// This is useful during shard-specific rebuilds where we want to regenerate
326    /// only one shard's files while preserving others.
327    pub fn clear_shard_file_paths(&mut self, shard_index: u32) {
328        let shard_idx = shard_index as usize;
329        for (_, paths) in self.file_path.iter_mut() {
330            if paths.len() > shard_idx {
331                paths[shard_idx].clear();
332            }
333        }
334        if self
335            .file_path
336            .values()
337            .all(|paths| paths.iter().all(|p| p.is_empty()))
338        {
339            self.file_path.clear();
340        }
341    }
342
343    /// Check that the on-disk file_path shape matches what the collection schema
344    /// implies. Returns `Ok(())` when they agree or the segment is uninitialized
345    /// (empty file_path). Returns `Err(SchemaMismatchError)` on the first
346    /// mismatch detected.
347    ///
348    /// Validates that all file_path keys have the same shard count and that the
349    /// set of keys present is consistent with the schema.
350    pub fn matches_schema(&self, schema: &Schema) -> Result<(), SchemaMismatchError> {
351        if self.file_path.is_empty() {
352            return Ok(());
353        }
354        // Validate all keys have the same number of shards.
355        self.num_shards().map_err(SchemaMismatchError::ShardError)?;
356
357        match self.scope {
358            SegmentScope::VECTOR => self.check_vector_consistency(schema),
359            SegmentScope::METADATA => self.check_metadata_consistency(schema),
360            SegmentScope::RECORD => self.check_record_consistency(),
361            SegmentScope::SQLITE => Ok(()),
362        }
363    }
364
365    fn mismatch(&self, detail: String) -> SchemaMismatchError {
366        SchemaMismatchError::Mismatch {
367            segment_id: self.id,
368            scope: self.scope.clone(),
369            detail,
370        }
371    }
372
373    /// Verify that all `required` keys exist in file_path and no `denied` keys
374    /// exist. Uninitialized segments (no relevant keys at all) pass.
375    fn require_keys(&self, required: &[&str], denied: &[&str]) -> Result<(), SchemaMismatchError> {
376        let has_required: Vec<&str> = required
377            .iter()
378            .filter(|k| self.file_path.contains_key(**k))
379            .copied()
380            .collect();
381        let has_denied: Vec<&str> = denied
382            .iter()
383            .filter(|k| self.file_path.contains_key(**k))
384            .copied()
385            .collect();
386
387        // No relevant keys at all — uninitialized, OK.
388        if has_required.is_empty() && has_denied.is_empty() {
389            return Ok(());
390        }
391
392        // Denied keys must not be present.
393        if !has_denied.is_empty() {
394            return Err(self.mismatch(format!(
395                "unexpected keys present: {}",
396                has_denied.join(", ")
397            )));
398        }
399
400        // All required keys must be present — no partial state.
401        if has_required.len() != required.len() {
402            let missing: Vec<&str> = required
403                .iter()
404                .filter(|k| !self.file_path.contains_key(**k))
405                .copied()
406                .collect();
407            return Err(self.mismatch(format!("missing required keys: {}", missing.join(", "))));
408        }
409
410        Ok(())
411    }
412
413    fn check_vector_consistency(&self, schema: &Schema) -> Result<(), SchemaMismatchError> {
414        if schema.get_spann_config().is_none() {
415            return Ok(());
416        }
417
418        let spann_keys: &[&str] = &[
419            HNSW_PATH,
420            VERSION_MAP_PATH,
421            POSTING_LIST_PATH,
422            MAX_HEAD_ID_BF_PATH,
423        ];
424        let qspann_keys: &[&str] = &[
425            QUANTIZED_SPANN_CLUSTER,
426            QUANTIZED_SPANN_SCALAR_METADATA,
427            QUANTIZED_SPANN_EMBEDDING_METADATA,
428            QUANTIZED_SPANN_RAW_CENTROID,
429            QUANTIZED_SPANN_QUANTIZED_CENTROID,
430        ];
431
432        if schema.is_quantization_enabled() {
433            self.require_keys(qspann_keys, spann_keys)
434        } else {
435            self.require_keys(spann_keys, qspann_keys)
436        }
437    }
438
439    fn check_metadata_consistency(&self, schema: &Schema) -> Result<(), SchemaMismatchError> {
440        let base_keys: &[&str] = &[STRING_METADATA, BOOL_METADATA, F32_METADATA, U32_METADATA];
441        self.require_keys(base_keys, &[])?;
442
443        if schema.is_fts_enabled() {
444            if schema.is_token_bitmap_fts_enabled() {
445                self.require_keys(&[FULL_TEXT_TOKEN], &[FULL_TEXT_PLS])?;
446            } else {
447                self.require_keys(&[FULL_TEXT_PLS], &[FULL_TEXT_TOKEN])?;
448            }
449        }
450
451        self.check_sparse_consistency(schema)
452    }
453
454    /// Validate the sparse-index `file_path` layout against the schema.
455    ///
456    /// Two layouts are accepted:
457    /// - Legacy/global (`SPARSE_*` constants): a single anonymous index written
458    ///   before per-key indexing existed. Its on-disk algorithm must match the
459    ///   schema (mismatch triggers a rebuild), preserving prior behavior.
460    /// - Per-key (`SPARSE_*::key`): one index per enabled sparse metadata key.
461    ///   Each present group must form a complete WAND pair or a MaxScore
462    ///   posting, matching that key's algorithm.
463    ///
464    /// Segments with no sparse entries (uninitialized or not yet compacted)
465    /// pass, so a legacy collection can lazily rewrite to per-key layout on its
466    /// next normal compaction without an explicit rebuild.
467    fn check_sparse_consistency(&self, schema: &Schema) -> Result<(), SchemaMismatchError> {
468        if schema.is_sparse_index_enabled() {
469            let has_legacy = self.file_path.contains_key(SPARSE_MAX)
470                || self.file_path.contains_key(SPARSE_OFFSET_VALUE)
471                || self.file_path.contains_key(SPARSE_POSTING);
472            if has_legacy {
473                let wand_keys: &[&str] = &[SPARSE_MAX, SPARSE_OFFSET_VALUE];
474                let maxscore_keys: &[&str] = &[SPARSE_POSTING];
475                if schema.is_maxscore_enabled() {
476                    self.require_keys(maxscore_keys, wand_keys)?;
477                } else {
478                    self.require_keys(wand_keys, maxscore_keys)?;
479                }
480            }
481        }
482
483        for key in schema.enabled_sparse_keys() {
484            let max_key = sparse_max_key(&key);
485            let offset_value_key = sparse_offset_value_key(&key);
486            let posting_key = sparse_posting_key(&key);
487
488            let has_per_key = self.file_path.contains_key(&max_key)
489                || self.file_path.contains_key(&offset_value_key)
490                || self.file_path.contains_key(&posting_key);
491            if !has_per_key {
492                continue;
493            }
494
495            let wand_keys: &[&str] = &[max_key.as_str(), offset_value_key.as_str()];
496            let maxscore_keys: &[&str] = &[posting_key.as_str()];
497            if schema.is_key_maxscore_enabled(&key) {
498                self.require_keys(maxscore_keys, wand_keys)?;
499            } else {
500                self.require_keys(wand_keys, maxscore_keys)?;
501            }
502        }
503
504        Ok(())
505    }
506
507    fn check_record_consistency(&self) -> Result<(), SchemaMismatchError> {
508        let required: &[&str] = &[
509            USER_ID_TO_OFFSET_ID,
510            OFFSET_ID_TO_USER_ID,
511            OFFSET_ID_TO_DATA,
512            MAX_OFFSET_ID,
513        ];
514        self.require_keys(required, &[])
515    }
516
517    pub fn extract_prefix_and_id(path: &str) -> Result<(&str, uuid::Uuid), uuid::Error> {
518        let (prefix, id) = match path.rfind('/') {
519            Some(pos) => (&path[..pos], &path[pos + 1..]),
520            None => ("", path),
521        };
522        match Uuid::try_parse(id) {
523            Ok(uid) => Ok((prefix, uid)),
524            Err(e) => Err(e),
525        }
526    }
527
528    pub fn construct_prefix_path(&self, tenant: &str, database_id: &DatabaseUuid) -> String {
529        Self::construct_prefix_path_impl(tenant, database_id, &self.collection, &self.id)
530    }
531
532    fn construct_prefix_path_impl(
533        tenant: &str,
534        database_id: &DatabaseUuid,
535        collection: &CollectionUuid,
536        segment_id: &SegmentUuid,
537    ) -> String {
538        format!(
539            "tenant/{}/database/{}/collection/{}/segment/{}",
540            tenant, database_id, collection, segment_id
541        )
542    }
543}
544
545#[derive(Clone, Debug, PartialEq)]
546pub struct SegmentShard {
547    pub id: SegmentUuid,
548    pub r#type: SegmentType,
549    pub scope: SegmentScope,
550    pub collection: CollectionUuid,
551    pub metadata: Option<Metadata>,
552    pub file_path: HashMap<String, String>,
553}
554
555impl SegmentShard {
556    pub fn construct_prefix_path(&self, tenant: &str, database_id: &DatabaseUuid) -> String {
557        Segment::construct_prefix_path_impl(tenant, database_id, &self.collection, &self.id)
558    }
559}
560
561#[derive(Error, Debug)]
562pub enum SegmentShardError {
563    #[error("Empty path vector for key '{0}'")]
564    EmptyPathVector(String),
565    #[error("Empty path string for key '{0}'")]
566    EmptyPathString(String),
567    #[error("Shard index {index} out of bounds for key '{key}' (len {len})")]
568    ShardIndexOutOfBounds { key: String, index: u32, len: usize },
569    #[error("Mismatched shard counts: key '{key}' has {actual} entries, expected {expected}")]
570    MismatchedShardCounts {
571        key: String,
572        actual: usize,
573        expected: usize,
574    },
575    #[error("No shards found")]
576    EmptyShards,
577}
578
579impl ChromaError for SegmentShardError {
580    fn code(&self) -> ErrorCodes {
581        ErrorCodes::Internal
582    }
583}
584
585/// Error returned by [`Segment::matches_schema`] when the on-disk file_path
586/// shape does not match what the collection schema implies.
587#[derive(Error, Debug)]
588pub enum SchemaMismatchError {
589    #[error("Schema/file_path mismatch on segment {segment_id} (scope {scope:?}): {detail}")]
590    Mismatch {
591        segment_id: SegmentUuid,
592        scope: SegmentScope,
593        detail: String,
594    },
595    #[error("Error inspecting shards: {0}")]
596    ShardError(SegmentShardError),
597}
598
599impl ChromaError for SchemaMismatchError {
600    fn code(&self) -> ErrorCodes {
601        match self {
602            Self::Mismatch { .. } => ErrorCodes::FailedPrecondition,
603            Self::ShardError(e) => e.code(),
604        }
605    }
606}
607
608impl TryFrom<(&Segment, u32)> for SegmentShard {
609    type Error = SegmentShardError;
610
611    fn try_from((segment, shard_index): (&Segment, u32)) -> Result<Self, Self::Error> {
612        let mut file_path = HashMap::new();
613        // If there are no shards in the filepaths this for loop won't
614        // run and this function will return an empty SegmentShard.
615        for (key, paths) in &segment.file_path {
616            match paths.get(shard_index as usize) {
617                Some(path) => {
618                    // During rebuild, clear_shard_file_paths sets paths to empty strings
619                    // to signal "create new". Skip these so writers see missing keys instead.
620                    if !path.is_empty() {
621                        file_path.insert(key.clone(), path.clone());
622                    }
623                }
624                None if paths.is_empty() => {
625                    return Err(SegmentShardError::EmptyPathVector(key.clone()));
626                }
627                None => {
628                    return Err(SegmentShardError::ShardIndexOutOfBounds {
629                        key: key.clone(),
630                        index: shard_index,
631                        len: paths.len(),
632                    });
633                }
634            }
635        }
636
637        Ok(SegmentShard {
638            id: segment.id,
639            r#type: segment.r#type,
640            scope: segment.scope.clone(),
641            collection: segment.collection,
642            metadata: segment.metadata.clone(),
643            file_path,
644        })
645    }
646}
647
648impl Segment {
649    pub fn new_shard(&self) -> SegmentShard {
650        SegmentShard {
651            id: self.id,
652            r#type: self.r#type,
653            scope: self.scope.clone(),
654            collection: self.collection,
655            metadata: self.metadata.clone(),
656            file_path: HashMap::new(),
657        }
658    }
659}
660
661#[derive(Error, Debug)]
662pub enum SegmentConversionError {
663    #[error("Invalid UUID")]
664    InvalidUuid,
665    #[error(transparent)]
666    MetadataValueConversionError(#[from] MetadataValueConversionError),
667    #[error(transparent)]
668    SegmentScopeConversionError(#[from] SegmentScopeConversionError),
669    #[error("Invalid segment type")]
670    InvalidSegmentType,
671}
672
673impl ChromaError for SegmentConversionError {
674    fn code(&self) -> ErrorCodes {
675        match self {
676            SegmentConversionError::InvalidUuid => ErrorCodes::InvalidArgument,
677            SegmentConversionError::InvalidSegmentType => ErrorCodes::InvalidArgument,
678            SegmentConversionError::SegmentScopeConversionError(e) => e.code(),
679            SegmentConversionError::MetadataValueConversionError(e) => e.code(),
680        }
681    }
682}
683
684impl From<SegmentConversionError> for Status {
685    fn from(value: SegmentConversionError) -> Self {
686        Status::invalid_argument(value.to_string())
687    }
688}
689
690impl TryFrom<chroma_proto::Segment> for Segment {
691    type Error = SegmentConversionError;
692
693    fn try_from(proto_segment: chroma_proto::Segment) -> Result<Self, Self::Error> {
694        let mut proto_segment = proto_segment;
695
696        let segment_uuid = match SegmentUuid::from_str(&proto_segment.id) {
697            Ok(uuid) => uuid,
698            Err(_) => return Err(SegmentConversionError::InvalidUuid),
699        };
700        let collection_uuid = match Uuid::try_parse(&proto_segment.collection) {
701            Ok(uuid) => uuid,
702            Err(_) => return Err(SegmentConversionError::InvalidUuid),
703        };
704        let collection_uuid = CollectionUuid(collection_uuid);
705        let segment_metadata: Option<Metadata> = match proto_segment.metadata {
706            Some(proto_metadata) => match proto_metadata.try_into() {
707                Ok(metadata) => Some(metadata),
708                Err(e) => return Err(SegmentConversionError::MetadataValueConversionError(e)),
709            },
710            None => None,
711        };
712        let scope: SegmentScope = match proto_segment.scope.try_into() {
713            Ok(scope) => scope,
714            Err(e) => return Err(SegmentConversionError::SegmentScopeConversionError(e)),
715        };
716
717        let segment_type: SegmentType = proto_segment.r#type.as_str().try_into()?;
718
719        let mut file_paths = HashMap::new();
720        let drain = proto_segment.file_paths.drain();
721        for (key, value) in drain {
722            file_paths.insert(key, value.paths);
723        }
724
725        Ok(Segment {
726            id: segment_uuid,
727            r#type: segment_type,
728            scope,
729            collection: collection_uuid,
730            metadata: segment_metadata,
731            file_path: file_paths,
732        })
733    }
734}
735
736impl From<Segment> for chroma_proto::Segment {
737    fn from(value: Segment) -> Self {
738        Self {
739            id: value.id.0.to_string(),
740            r#type: value.r#type.into(),
741            scope: chroma_proto::SegmentScope::from(value.scope) as i32,
742            collection: value.collection.0.to_string(),
743            metadata: value.metadata.map(Into::into),
744            file_paths: value
745                .file_path
746                .into_iter()
747                .map(|(name, paths)| (name, chroma_proto::FilePaths { paths }))
748                .collect(),
749        }
750    }
751}
752
753pub fn test_segment(collection_uuid: CollectionUuid, scope: SegmentScope) -> Segment {
754    let r#type = match scope {
755        SegmentScope::METADATA => SegmentType::BlockfileMetadata,
756        SegmentScope::RECORD => SegmentType::BlockfileRecord,
757        SegmentScope::VECTOR => SegmentType::HnswDistributed,
758        SegmentScope::SQLITE => unimplemented!("Sqlite segment is not implemented"),
759    };
760    Segment {
761        id: SegmentUuid::new(),
762        r#type,
763        scope,
764        collection: collection_uuid,
765        metadata: None,
766        file_path: HashMap::new(),
767    }
768}
769
770#[cfg(test)]
771mod tests {
772
773    use super::*;
774    use crate::MetadataValue;
775
776    #[test]
777    fn test_segment_try_from() {
778        let mut metadata = chroma_proto::UpdateMetadata {
779            metadata: HashMap::new(),
780        };
781        metadata.metadata.insert(
782            "foo".to_string(),
783            chroma_proto::UpdateMetadataValue {
784                value: Some(chroma_proto::update_metadata_value::Value::IntValue(42)),
785            },
786        );
787        let proto_segment = chroma_proto::Segment {
788            id: "00000000-0000-0000-0000-000000000000".to_string(),
789            r#type: "urn:chroma:segment/vector/hnsw-distributed".to_string(),
790            scope: chroma_proto::SegmentScope::Vector as i32,
791            collection: "00000000-0000-0000-0000-000000000000".to_string(),
792            metadata: Some(metadata),
793            file_paths: HashMap::new(),
794        };
795        let converted_segment: Segment = proto_segment.try_into().unwrap();
796        assert_eq!(converted_segment.id, SegmentUuid(Uuid::nil()));
797        assert_eq!(converted_segment.r#type, SegmentType::HnswDistributed);
798        assert_eq!(converted_segment.scope, SegmentScope::VECTOR);
799        assert_eq!(converted_segment.collection, CollectionUuid(Uuid::nil()));
800        let metadata = converted_segment.metadata.unwrap();
801        assert_eq!(metadata.len(), 1);
802        assert_eq!(metadata.get("foo").unwrap(), &MetadataValue::Int(42));
803    }
804
805    #[test]
806    fn test_segment_construct_prefix_path() {
807        let segment = Segment {
808            id: SegmentUuid(Uuid::nil()),
809            r#type: SegmentType::BlockfileMetadata,
810            scope: SegmentScope::METADATA,
811            collection: CollectionUuid(Uuid::nil()),
812            metadata: None,
813            file_path: HashMap::new(),
814        };
815        let tenant = "test_tenant";
816        let database_id = &DatabaseUuid(Uuid::nil());
817        let prefix_path = segment.construct_prefix_path(tenant, database_id);
818        assert_eq!(
819            prefix_path,
820            "tenant/test_tenant/database/00000000-0000-0000-0000-000000000000/collection/00000000-0000-0000-0000-000000000000/segment/00000000-0000-0000-0000-000000000000"
821        );
822    }
823
824    #[test]
825    fn test_segment_extract_prefix_and_id() {
826        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";
827        let (prefix, id) =
828            Segment::extract_prefix_and_id(path).expect("Failed to extract prefix and id");
829        assert_eq!(
830            prefix,
831            "tenant/test_tenant/database/00000000-0000-0000-0000-000000000000/collection/00000000-0000-0000-0000-000000000000/segment/00000000-0000-0000-0000-000000000000"
832        );
833        assert_eq!(
834            id,
835            Uuid::from_str("00000000-0000-0000-0000-000000000001").expect("Cannot happen")
836        );
837    }
838
839    #[test]
840    fn test_segment_extract_prefix_and_id_legacy() {
841        let path = "00000000-0000-0000-0000-000000000001";
842        let (prefix, id) =
843            Segment::extract_prefix_and_id(path).expect("Failed to extract prefix and id");
844        assert_eq!(prefix, "");
845        assert_eq!(
846            id,
847            Uuid::from_str("00000000-0000-0000-0000-000000000001").expect("Cannot happen")
848        );
849    }
850
851    #[test]
852    fn test_num_shards_empty_file_path() {
853        let segment = Segment {
854            id: SegmentUuid(Uuid::nil()),
855            r#type: SegmentType::BlockfileRecord,
856            scope: SegmentScope::RECORD,
857            collection: CollectionUuid(Uuid::nil()),
858            metadata: None,
859            file_path: HashMap::new(),
860        };
861        assert_eq!(segment.num_shards().unwrap(), 1);
862    }
863
864    #[test]
865    fn test_num_shards_single_shard() {
866        let mut file_path = HashMap::new();
867        file_path.insert("key_a".to_string(), vec!["path_a_0".to_string()]);
868        file_path.insert("key_b".to_string(), vec!["path_b_0".to_string()]);
869        let segment = Segment {
870            id: SegmentUuid(Uuid::nil()),
871            r#type: SegmentType::BlockfileRecord,
872            scope: SegmentScope::RECORD,
873            collection: CollectionUuid(Uuid::nil()),
874            metadata: None,
875            file_path,
876        };
877        assert_eq!(segment.num_shards().unwrap(), 1);
878    }
879
880    #[test]
881    fn test_num_shards_multi_shard() {
882        let mut file_path = HashMap::new();
883        file_path.insert(
884            "key_a".to_string(),
885            vec!["a0".to_string(), "a1".to_string(), "a2".to_string()],
886        );
887        file_path.insert(
888            "key_b".to_string(),
889            vec!["b0".to_string(), "b1".to_string(), "b2".to_string()],
890        );
891        let segment = Segment {
892            id: SegmentUuid(Uuid::nil()),
893            r#type: SegmentType::BlockfileRecord,
894            scope: SegmentScope::RECORD,
895            collection: CollectionUuid(Uuid::nil()),
896            metadata: None,
897            file_path,
898        };
899        assert_eq!(segment.num_shards().unwrap(), 3);
900    }
901
902    #[test]
903    fn test_num_shards_mismatched_lengths() {
904        let mut file_path = HashMap::new();
905        file_path.insert(
906            "key_a".to_string(),
907            vec!["a0".to_string(), "a1".to_string()],
908        );
909        file_path.insert(
910            "key_b".to_string(),
911            vec!["b0".to_string(), "b1".to_string(), "b2".to_string()],
912        );
913        let segment = Segment {
914            id: SegmentUuid(Uuid::nil()),
915            r#type: SegmentType::BlockfileRecord,
916            scope: SegmentScope::RECORD,
917            collection: CollectionUuid(Uuid::nil()),
918            metadata: None,
919            file_path,
920        };
921        let err = segment.num_shards().unwrap_err();
922        assert!(matches!(
923            err,
924            SegmentShardError::MismatchedShardCounts { .. }
925        ));
926    }
927
928    #[test]
929    fn test_clear_shard_file_paths() {
930        // Create a segment with 3 shards
931        let mut file_path = HashMap::new();
932        file_path.insert(
933            "metadata".to_string(),
934            vec![
935                "/path/to/shard0/metadata".to_string(),
936                "/path/to/shard1/metadata".to_string(),
937                "/path/to/shard2/metadata".to_string(),
938            ],
939        );
940        file_path.insert(
941            "index".to_string(),
942            vec![
943                "/path/to/shard0/index".to_string(),
944                "/path/to/shard1/index".to_string(),
945                "/path/to/shard2/index".to_string(),
946            ],
947        );
948        file_path.insert(
949            "data".to_string(),
950            vec![
951                "/path/to/shard0/data".to_string(),
952                "/path/to/shard1/data".to_string(),
953                "/path/to/shard2/data".to_string(),
954            ],
955        );
956
957        let mut segment = Segment {
958            id: SegmentUuid(Uuid::nil()),
959            r#type: SegmentType::BlockfileRecord,
960            scope: SegmentScope::RECORD,
961            collection: CollectionUuid(Uuid::nil()),
962            metadata: None,
963            file_path,
964        };
965
966        // Clear only shard 1's file paths
967        segment.clear_shard_file_paths(1);
968
969        // Verify shard 0 and shard 2 are untouched
970        assert_eq!(
971            segment.file_path.get("metadata").unwrap()[0],
972            "/path/to/shard0/metadata"
973        );
974        assert_eq!(
975            segment.file_path.get("metadata").unwrap()[2],
976            "/path/to/shard2/metadata"
977        );
978        assert_eq!(
979            segment.file_path.get("index").unwrap()[0],
980            "/path/to/shard0/index"
981        );
982        assert_eq!(
983            segment.file_path.get("index").unwrap()[2],
984            "/path/to/shard2/index"
985        );
986        assert_eq!(
987            segment.file_path.get("data").unwrap()[0],
988            "/path/to/shard0/data"
989        );
990        assert_eq!(
991            segment.file_path.get("data").unwrap()[2],
992            "/path/to/shard2/data"
993        );
994
995        // Verify shard 1 is cleared (empty string)
996        assert_eq!(segment.file_path.get("metadata").unwrap()[1], "");
997        assert_eq!(segment.file_path.get("index").unwrap()[1], "");
998        assert_eq!(segment.file_path.get("data").unwrap()[1], "");
999
1000        // Test clearing a shard index that's out of bounds (should not panic)
1001        segment.clear_shard_file_paths(5);
1002
1003        // Verify nothing changed for the out-of-bounds case
1004        assert_eq!(
1005            segment.file_path.get("metadata").unwrap()[0],
1006            "/path/to/shard0/metadata"
1007        );
1008    }
1009
1010    // ── matches_schema tests ────────────────────────────────────────────────
1011
1012    use crate::collection_configuration::KnnIndex;
1013    use crate::collection_schema::{
1014        Quantization, Schema, SparseIndexAlgorithm, SparseVectorIndexConfig, SparseVectorIndexType,
1015        SparseVectorValueType,
1016    };
1017
1018    /// Helper: create a Spann schema (full precision, no quantization).
1019    fn spann_schema() -> Schema {
1020        Schema::new_default(KnnIndex::Spann)
1021    }
1022
1023    /// Helper: create a Spann schema with quantization enabled.
1024    fn quantized_spann_schema() -> Schema {
1025        let mut schema = Schema::new_default(KnnIndex::Spann);
1026        if let Some(spann_config) = schema.get_spann_config_mut() {
1027            spann_config.quantize = Quantization::FourBitRabitQWithUSearch;
1028        }
1029        schema
1030    }
1031
1032    /// Helper: create an HNSW schema.
1033    fn hnsw_schema() -> Schema {
1034        Schema::new_default(KnnIndex::Hnsw)
1035    }
1036
1037    /// Helper: create a schema with sparse index enabled (Wand).
1038    fn sparse_wand_schema() -> Schema {
1039        let mut schema = Schema::new_default(KnnIndex::Hnsw);
1040        schema
1041            .keys
1042            .entry("sparse_key".to_string())
1043            .or_default()
1044            .sparse_vector = Some(SparseVectorValueType {
1045            sparse_vector_index: Some(SparseVectorIndexType {
1046                enabled: true,
1047                config: SparseVectorIndexConfig {
1048                    embedding_function: None,
1049                    source_key: None,
1050                    bm25: None,
1051                    algorithm: SparseIndexAlgorithm::Wand,
1052                },
1053            }),
1054        });
1055        schema
1056    }
1057
1058    /// Helper: create a schema with sparse index enabled (MaxScore).
1059    fn sparse_maxscore_schema() -> Schema {
1060        let mut schema = sparse_wand_schema();
1061        schema.set_sparse_algorithm(SparseIndexAlgorithm::MaxScore);
1062        schema
1063    }
1064
1065    fn make_vector_segment(file_path: HashMap<String, Vec<String>>) -> Segment {
1066        Segment {
1067            id: SegmentUuid(Uuid::nil()),
1068            r#type: SegmentType::Spann,
1069            scope: SegmentScope::VECTOR,
1070            collection: CollectionUuid(Uuid::nil()),
1071            metadata: None,
1072            file_path,
1073        }
1074    }
1075
1076    fn make_metadata_segment(file_path: HashMap<String, Vec<String>>) -> Segment {
1077        Segment {
1078            id: SegmentUuid(Uuid::nil()),
1079            r#type: SegmentType::BlockfileMetadata,
1080            scope: SegmentScope::METADATA,
1081            collection: CollectionUuid(Uuid::nil()),
1082            metadata: None,
1083            file_path,
1084        }
1085    }
1086
1087    #[test]
1088    fn test_matches_schema_empty_file_path() {
1089        let seg = make_vector_segment(HashMap::new());
1090        assert!(seg.matches_schema(&spann_schema()).is_ok());
1091        assert!(seg.matches_schema(&quantized_spann_schema()).is_ok());
1092        assert!(seg.matches_schema(&hnsw_schema()).is_ok());
1093    }
1094
1095    #[test]
1096    fn test_matches_schema_spann_matches_spann() {
1097        let mut fp = HashMap::new();
1098        fp.insert(HNSW_PATH.to_string(), vec!["path/a".to_string()]);
1099        fp.insert(VERSION_MAP_PATH.to_string(), vec!["path/b".to_string()]);
1100        fp.insert(POSTING_LIST_PATH.to_string(), vec!["path/c".to_string()]);
1101        fp.insert(MAX_HEAD_ID_BF_PATH.to_string(), vec!["path/d".to_string()]);
1102        let seg = make_vector_segment(fp);
1103        assert!(seg.matches_schema(&spann_schema()).is_ok());
1104    }
1105
1106    #[test]
1107    fn test_matches_schema_qspann_matches_qspann() {
1108        let mut fp = HashMap::new();
1109        fp.insert(
1110            QUANTIZED_SPANN_CLUSTER.to_string(),
1111            vec!["path/a".to_string()],
1112        );
1113        fp.insert(
1114            QUANTIZED_SPANN_SCALAR_METADATA.to_string(),
1115            vec!["path/b".to_string()],
1116        );
1117        fp.insert(
1118            QUANTIZED_SPANN_EMBEDDING_METADATA.to_string(),
1119            vec!["path/c".to_string()],
1120        );
1121        fp.insert(
1122            QUANTIZED_SPANN_RAW_CENTROID.to_string(),
1123            vec!["path/d".to_string()],
1124        );
1125        fp.insert(
1126            QUANTIZED_SPANN_QUANTIZED_CENTROID.to_string(),
1127            vec!["path/e".to_string()],
1128        );
1129        let seg = make_vector_segment(fp);
1130        assert!(seg.matches_schema(&quantized_spann_schema()).is_ok());
1131    }
1132
1133    #[test]
1134    fn test_matches_schema_spann_on_disk_qspann_in_schema() {
1135        let mut fp = HashMap::new();
1136        fp.insert(HNSW_PATH.to_string(), vec!["path/a".to_string()]);
1137        fp.insert(VERSION_MAP_PATH.to_string(), vec!["path/b".to_string()]);
1138        fp.insert(POSTING_LIST_PATH.to_string(), vec!["path/c".to_string()]);
1139        fp.insert(MAX_HEAD_ID_BF_PATH.to_string(), vec!["path/d".to_string()]);
1140        let seg = make_vector_segment(fp);
1141        let err = seg.matches_schema(&quantized_spann_schema()).unwrap_err();
1142        assert!(matches!(err, SchemaMismatchError::Mismatch { .. }));
1143    }
1144
1145    #[test]
1146    fn test_matches_schema_qspann_on_disk_spann_in_schema() {
1147        let mut fp = HashMap::new();
1148        fp.insert(
1149            QUANTIZED_SPANN_CLUSTER.to_string(),
1150            vec!["path/a".to_string()],
1151        );
1152        fp.insert(
1153            QUANTIZED_SPANN_SCALAR_METADATA.to_string(),
1154            vec!["path/b".to_string()],
1155        );
1156        fp.insert(
1157            QUANTIZED_SPANN_EMBEDDING_METADATA.to_string(),
1158            vec!["path/c".to_string()],
1159        );
1160        fp.insert(
1161            QUANTIZED_SPANN_RAW_CENTROID.to_string(),
1162            vec!["path/d".to_string()],
1163        );
1164        fp.insert(
1165            QUANTIZED_SPANN_QUANTIZED_CENTROID.to_string(),
1166            vec!["path/e".to_string()],
1167        );
1168        let seg = make_vector_segment(fp);
1169        let err = seg.matches_schema(&spann_schema()).unwrap_err();
1170        assert!(matches!(err, SchemaMismatchError::Mismatch { .. }));
1171    }
1172
1173    #[test]
1174    fn test_matches_schema_hnsw_schema_always_ok() {
1175        // HNSW schema does not check file_path — no migration variants.
1176        let mut fp = HashMap::new();
1177        fp.insert(HNSW_PATH.to_string(), vec!["path/a".to_string()]);
1178        let seg = make_vector_segment(fp);
1179        assert!(seg.matches_schema(&hnsw_schema()).is_ok());
1180    }
1181
1182    fn make_record_segment(file_path: HashMap<String, Vec<String>>) -> Segment {
1183        Segment {
1184            id: SegmentUuid(Uuid::nil()),
1185            r#type: SegmentType::BlockfileRecord,
1186            scope: SegmentScope::RECORD,
1187            collection: CollectionUuid(Uuid::nil()),
1188            metadata: None,
1189            file_path,
1190        }
1191    }
1192
1193    #[test]
1194    fn test_matches_schema_record_empty() {
1195        let seg = make_record_segment(HashMap::new());
1196        assert!(seg.matches_schema(&spann_schema()).is_ok());
1197    }
1198
1199    #[test]
1200    fn test_matches_schema_record_complete() {
1201        let mut fp = HashMap::new();
1202        fp.insert(USER_ID_TO_OFFSET_ID.to_string(), vec!["a".to_string()]);
1203        fp.insert(OFFSET_ID_TO_USER_ID.to_string(), vec!["b".to_string()]);
1204        fp.insert(OFFSET_ID_TO_DATA.to_string(), vec!["c".to_string()]);
1205        fp.insert(MAX_OFFSET_ID.to_string(), vec!["d".to_string()]);
1206        let seg = make_record_segment(fp);
1207        assert!(seg.matches_schema(&spann_schema()).is_ok());
1208    }
1209
1210    #[test]
1211    fn test_matches_schema_record_missing_key() {
1212        let mut fp = HashMap::new();
1213        fp.insert(USER_ID_TO_OFFSET_ID.to_string(), vec!["a".to_string()]);
1214        fp.insert(OFFSET_ID_TO_USER_ID.to_string(), vec!["b".to_string()]);
1215        fp.insert(OFFSET_ID_TO_DATA.to_string(), vec!["c".to_string()]);
1216        // MAX_OFFSET_ID intentionally omitted
1217        let seg = make_record_segment(fp);
1218        let err = seg.matches_schema(&spann_schema()).unwrap_err();
1219        assert!(matches!(err, SchemaMismatchError::Mismatch { .. }));
1220    }
1221
1222    #[test]
1223    fn test_matches_schema_sparse_wand_matches_wand() {
1224        let mut fp = HashMap::new();
1225        fp.insert(SPARSE_MAX.to_string(), vec!["path/a".to_string()]);
1226        fp.insert(SPARSE_OFFSET_VALUE.to_string(), vec!["path/b".to_string()]);
1227        let seg = make_metadata_segment(fp);
1228        assert!(seg.matches_schema(&sparse_wand_schema()).is_ok());
1229    }
1230
1231    #[test]
1232    fn test_matches_schema_sparse_maxscore_matches_maxscore() {
1233        let mut fp = HashMap::new();
1234        fp.insert(SPARSE_POSTING.to_string(), vec!["path/a".to_string()]);
1235        let seg = make_metadata_segment(fp);
1236        assert!(seg.matches_schema(&sparse_maxscore_schema()).is_ok());
1237    }
1238
1239    #[test]
1240    fn test_matches_schema_sparse_wand_on_disk_maxscore_in_schema() {
1241        let mut fp = HashMap::new();
1242        fp.insert(SPARSE_MAX.to_string(), vec!["path/a".to_string()]);
1243        fp.insert(SPARSE_OFFSET_VALUE.to_string(), vec!["path/b".to_string()]);
1244        let seg = make_metadata_segment(fp);
1245        let err = seg.matches_schema(&sparse_maxscore_schema()).unwrap_err();
1246        assert!(matches!(err, SchemaMismatchError::Mismatch { .. }));
1247    }
1248
1249    #[test]
1250    fn test_matches_schema_sparse_maxscore_on_disk_wand_in_schema() {
1251        let mut fp = HashMap::new();
1252        fp.insert(SPARSE_POSTING.to_string(), vec!["path/a".to_string()]);
1253        let seg = make_metadata_segment(fp);
1254        let err = seg.matches_schema(&sparse_wand_schema()).unwrap_err();
1255        assert!(matches!(err, SchemaMismatchError::Mismatch { .. }));
1256    }
1257
1258    #[test]
1259    fn test_matches_schema_sparse_empty_file_path() {
1260        let seg = make_metadata_segment(HashMap::new());
1261        assert!(seg.matches_schema(&sparse_maxscore_schema()).is_ok());
1262        assert!(seg.matches_schema(&sparse_wand_schema()).is_ok());
1263    }
1264
1265    #[test]
1266    fn test_matches_schema_sparse_disabled_index_no_check() {
1267        // Schema with no sparse index enabled — should not check.
1268        let schema = hnsw_schema();
1269        assert!(!schema.is_sparse_index_enabled());
1270        let mut fp = HashMap::new();
1271        fp.insert(SPARSE_POSTING.to_string(), vec!["path/a".to_string()]);
1272        let seg = make_metadata_segment(fp);
1273        assert!(seg.matches_schema(&schema).is_ok());
1274    }
1275
1276    #[test]
1277    fn test_matches_schema_sparse_both_keys_present() {
1278        // Both Wand and MaxScore keys present — invalid state.
1279        let mut fp = HashMap::new();
1280        fp.insert(SPARSE_POSTING.to_string(), vec!["path/a".to_string()]);
1281        fp.insert(SPARSE_MAX.to_string(), vec!["path/b".to_string()]);
1282        fp.insert(SPARSE_OFFSET_VALUE.to_string(), vec!["path/c".to_string()]);
1283        let seg = make_metadata_segment(fp);
1284        let err = seg.matches_schema(&sparse_maxscore_schema()).unwrap_err();
1285        assert!(matches!(err, SchemaMismatchError::Mismatch { .. }));
1286    }
1287
1288    #[test]
1289    fn test_matches_schema_sparse_per_key_wand_matches_wand() {
1290        let mut fp = HashMap::new();
1291        fp.insert(sparse_max_key("sparse_key"), vec!["path/a".to_string()]);
1292        fp.insert(
1293            sparse_offset_value_key("sparse_key"),
1294            vec!["path/b".to_string()],
1295        );
1296        let seg = make_metadata_segment(fp);
1297        assert!(seg.matches_schema(&sparse_wand_schema()).is_ok());
1298    }
1299
1300    #[test]
1301    fn test_matches_schema_sparse_per_key_maxscore_matches_maxscore() {
1302        let mut fp = HashMap::new();
1303        fp.insert(sparse_posting_key("sparse_key"), vec!["path/a".to_string()]);
1304        let seg = make_metadata_segment(fp);
1305        assert!(seg.matches_schema(&sparse_maxscore_schema()).is_ok());
1306    }
1307
1308    #[test]
1309    fn test_matches_schema_sparse_per_key_algorithm_mismatch() {
1310        // On-disk per-key WAND pair but schema declares MaxScore for the key.
1311        let mut fp = HashMap::new();
1312        fp.insert(sparse_max_key("sparse_key"), vec!["path/a".to_string()]);
1313        fp.insert(
1314            sparse_offset_value_key("sparse_key"),
1315            vec!["path/b".to_string()],
1316        );
1317        let seg = make_metadata_segment(fp);
1318        let err = seg.matches_schema(&sparse_maxscore_schema()).unwrap_err();
1319        assert!(matches!(err, SchemaMismatchError::Mismatch { .. }));
1320    }
1321
1322    #[test]
1323    fn test_matches_schema_sparse_per_key_incomplete_wand_pair() {
1324        // Only one half of a WAND pair present — incomplete, must mismatch.
1325        let mut fp = HashMap::new();
1326        fp.insert(sparse_max_key("sparse_key"), vec!["path/a".to_string()]);
1327        let seg = make_metadata_segment(fp);
1328        let err = seg.matches_schema(&sparse_wand_schema()).unwrap_err();
1329        assert!(matches!(err, SchemaMismatchError::Mismatch { .. }));
1330    }
1331}
chroma_types/segment.rs

chroma_types/
segment.rs
