chroma_types/
collection.rs

1use std::str::FromStr;
2
3use super::{Metadata, MetadataValueConversionError};
4use crate::{
5    chroma_proto, test_segment, CollectionConfiguration, InternalCollectionConfiguration, Schema,
6    SchemaError, Segment, SegmentScope, UpdateCollectionConfiguration, UpdateMetadata,
7};
8use chroma_error::{ChromaError, ErrorCodes};
9use serde::{Deserialize, Serialize};
10use std::time::{Duration, SystemTime};
11use thiserror::Error;
12use uuid::Uuid;
13
14#[cfg(feature = "pyo3")]
15use pyo3::types::PyAnyMethods;
16
17/// CollectionUuid is a wrapper around Uuid to provide a type for the collection id.
18#[derive(
19    Copy, Clone, Debug, Default, Deserialize, Eq, PartialEq, Ord, PartialOrd, Hash, Serialize,
20)]
21#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
22pub struct CollectionUuid(pub Uuid);
23
24/// DatabaseUuid is a wrapper around Uuid to provide a type for the database id.
25#[derive(
26    Copy, Clone, Debug, Default, Deserialize, Eq, PartialEq, Ord, PartialOrd, Hash, Serialize,
27)]
28#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
29pub struct DatabaseUuid(pub Uuid);
30
31impl DatabaseUuid {
32    pub fn new() -> Self {
33        DatabaseUuid(Uuid::new_v4())
34    }
35}
36
37impl CollectionUuid {
38    pub fn new() -> Self {
39        CollectionUuid(Uuid::new_v4())
40    }
41
42    pub fn storage_prefix_for_log(&self) -> String {
43        format!("logs/{}", self)
44    }
45}
46
47impl std::str::FromStr for CollectionUuid {
48    type Err = uuid::Error;
49
50    fn from_str(s: &str) -> Result<Self, Self::Err> {
51        match Uuid::parse_str(s) {
52            Ok(uuid) => Ok(CollectionUuid(uuid)),
53            Err(err) => Err(err),
54        }
55    }
56}
57
58impl std::str::FromStr for DatabaseUuid {
59    type Err = uuid::Error;
60
61    fn from_str(s: &str) -> Result<Self, Self::Err> {
62        match Uuid::parse_str(s) {
63            Ok(uuid) => Ok(DatabaseUuid(uuid)),
64            Err(err) => Err(err),
65        }
66    }
67}
68
69impl std::fmt::Display for DatabaseUuid {
70    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
71        write!(f, "{}", self.0)
72    }
73}
74
75impl std::fmt::Display for CollectionUuid {
76    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
77        write!(f, "{}", self.0)
78    }
79}
80
81fn serialize_internal_collection_configuration<S: serde::Serializer>(
82    config: &InternalCollectionConfiguration,
83    serializer: S,
84) -> Result<S::Ok, S::Error> {
85    let collection_config: CollectionConfiguration = config.clone().into();
86    collection_config.serialize(serializer)
87}
88
89fn deserialize_internal_collection_configuration<'de, D: serde::Deserializer<'de>>(
90    deserializer: D,
91) -> Result<InternalCollectionConfiguration, D::Error> {
92    let collection_config = CollectionConfiguration::deserialize(deserializer)?;
93    collection_config
94        .try_into()
95        .map_err(serde::de::Error::custom)
96}
97
98#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
99#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
100#[cfg_attr(feature = "pyo3", pyo3::pyclass)]
101pub struct Collection {
102    #[serde(rename = "id")]
103    pub collection_id: CollectionUuid,
104    pub name: String,
105    #[serde(
106        serialize_with = "serialize_internal_collection_configuration",
107        deserialize_with = "deserialize_internal_collection_configuration",
108        rename = "configuration_json"
109    )]
110    #[cfg_attr(feature = "utoipa", schema(value_type = CollectionConfiguration))]
111    pub config: InternalCollectionConfiguration,
112    pub schema: Option<Schema>,
113    pub metadata: Option<Metadata>,
114    pub dimension: Option<i32>,
115    pub tenant: String,
116    pub database: String,
117    pub log_position: i64,
118    pub version: i32,
119    #[serde(skip)]
120    pub total_records_post_compaction: u64,
121    #[serde(skip)]
122    pub size_bytes_post_compaction: u64,
123    #[serde(skip)]
124    pub last_compaction_time_secs: u64,
125    #[serde(skip)]
126    pub version_file_path: Option<String>,
127    #[serde(skip)]
128    pub root_collection_id: Option<CollectionUuid>,
129    #[serde(skip)]
130    pub lineage_file_path: Option<String>,
131    #[serde(skip, default = "SystemTime::now")]
132    pub updated_at: SystemTime,
133    #[serde(skip)]
134    pub database_id: DatabaseUuid,
135}
136
137impl Default for Collection {
138    fn default() -> Self {
139        Self {
140            collection_id: CollectionUuid::new(),
141            name: "".to_string(),
142            config: InternalCollectionConfiguration::default_hnsw(),
143            schema: None,
144            metadata: None,
145            dimension: None,
146            tenant: "".to_string(),
147            database: "".to_string(),
148            log_position: 0,
149            version: 0,
150            total_records_post_compaction: 0,
151            size_bytes_post_compaction: 0,
152            last_compaction_time_secs: 0,
153            version_file_path: None,
154            root_collection_id: None,
155            lineage_file_path: None,
156            updated_at: SystemTime::now(),
157            database_id: DatabaseUuid::new(),
158        }
159    }
160}
161
162#[cfg(feature = "pyo3")]
163#[pyo3::pymethods]
164impl Collection {
165    #[getter]
166    fn id<'py>(&self, py: pyo3::Python<'py>) -> pyo3::PyResult<pyo3::Bound<'py, pyo3::PyAny>> {
167        let res = pyo3::prelude::PyModule::import(py, "uuid")?
168            .getattr("UUID")?
169            .call1((self.collection_id.to_string(),))?;
170        Ok(res)
171    }
172
173    #[getter]
174    fn configuration<'py>(
175        &self,
176        py: pyo3::Python<'py>,
177    ) -> pyo3::PyResult<pyo3::Bound<'py, pyo3::PyAny>> {
178        let config: crate::CollectionConfiguration = self.config.clone().into();
179        let config_json_str = serde_json::to_string(&config).unwrap();
180        let res = pyo3::prelude::PyModule::import(py, "json")?
181            .getattr("loads")?
182            .call1((config_json_str,))?;
183        Ok(res)
184    }
185
186    #[getter]
187    pub fn name(&self) -> &str {
188        &self.name
189    }
190
191    #[getter]
192    pub fn metadata(&self) -> Option<Metadata> {
193        self.metadata.clone()
194    }
195
196    #[getter]
197    pub fn dimension(&self) -> Option<i32> {
198        self.dimension
199    }
200
201    #[getter]
202    pub fn tenant(&self) -> &str {
203        &self.tenant
204    }
205
206    #[getter]
207    pub fn database(&self) -> &str {
208        &self.database
209    }
210}
211
212impl Collection {
213    /// Reconcile the collection schema and configuration, ensuring both are consistent.
214    pub fn reconcile_schema_with_config(&mut self) -> Result<(), SchemaError> {
215        let reconciled_schema =
216            Schema::reconcile_schema_and_config(self.schema.clone(), Some(self.config.clone()))
217                .map_err(|reason| SchemaError::InvalidSchema { reason })?;
218
219        self.config = InternalCollectionConfiguration::try_from(&reconciled_schema)
220            .map_err(|reason| SchemaError::InvalidSchema { reason })?;
221        self.schema = Some(reconciled_schema);
222
223        Ok(())
224    }
225
226    pub fn test_collection(dim: i32) -> Self {
227        Collection {
228            name: "test_collection".to_string(),
229            dimension: Some(dim),
230            tenant: "default_tenant".to_string(),
231            database: "default_database".to_string(),
232            database_id: DatabaseUuid::new(),
233            ..Default::default()
234        }
235    }
236}
237
238#[derive(Error, Debug)]
239pub enum CollectionConversionError {
240    #[error("Invalid config: {0}")]
241    InvalidConfig(#[from] serde_json::Error),
242    #[error("Invalid UUID")]
243    InvalidUuid,
244    #[error(transparent)]
245    MetadataValueConversionError(#[from] MetadataValueConversionError),
246    #[error("Missing Database Id")]
247    MissingDatabaseId,
248}
249
250impl ChromaError for CollectionConversionError {
251    fn code(&self) -> ErrorCodes {
252        match self {
253            CollectionConversionError::InvalidConfig(_) => ErrorCodes::InvalidArgument,
254            CollectionConversionError::InvalidUuid => ErrorCodes::InvalidArgument,
255            CollectionConversionError::MetadataValueConversionError(e) => e.code(),
256            CollectionConversionError::MissingDatabaseId => ErrorCodes::Internal,
257        }
258    }
259}
260
261impl TryFrom<chroma_proto::Collection> for Collection {
262    type Error = CollectionConversionError;
263
264    fn try_from(proto_collection: chroma_proto::Collection) -> Result<Self, Self::Error> {
265        let collection_id = CollectionUuid::from_str(&proto_collection.id)
266            .map_err(|_| CollectionConversionError::InvalidUuid)?;
267        let collection_metadata: Option<Metadata> = match proto_collection.metadata {
268            Some(proto_metadata) => match proto_metadata.try_into() {
269                Ok(metadata) => Some(metadata),
270                Err(e) => return Err(CollectionConversionError::MetadataValueConversionError(e)),
271            },
272            None => None,
273        };
274        // TODO(@codetheweb): this be updated to error with "missing field" once all SysDb deployments are up-to-date
275        let updated_at = match proto_collection.updated_at {
276            Some(updated_at) => {
277                SystemTime::UNIX_EPOCH
278                    + Duration::new(updated_at.seconds as u64, updated_at.nanos as u32)
279            }
280            None => SystemTime::now(),
281        };
282        let database_id = match proto_collection.database_id {
283            Some(db_id) => DatabaseUuid::from_str(&db_id)
284                .map_err(|_| CollectionConversionError::InvalidUuid)?,
285            None => {
286                return Err(CollectionConversionError::MissingDatabaseId);
287            }
288        };
289        let schema = match proto_collection.schema_str {
290            Some(schema_str) if !schema_str.is_empty() => Some(serde_json::from_str(&schema_str)?),
291            _ => None,
292        };
293
294        Ok(Collection {
295            collection_id,
296            name: proto_collection.name,
297            config: serde_json::from_str(&proto_collection.configuration_json_str)?,
298            schema,
299            metadata: collection_metadata,
300            dimension: proto_collection.dimension,
301            tenant: proto_collection.tenant,
302            database: proto_collection.database,
303            log_position: proto_collection.log_position,
304            version: proto_collection.version,
305            total_records_post_compaction: proto_collection.total_records_post_compaction,
306            size_bytes_post_compaction: proto_collection.size_bytes_post_compaction,
307            last_compaction_time_secs: proto_collection.last_compaction_time_secs,
308            version_file_path: proto_collection.version_file_path,
309            root_collection_id: proto_collection
310                .root_collection_id
311                .map(|uuid| CollectionUuid(Uuid::try_parse(&uuid).unwrap())),
312            lineage_file_path: proto_collection.lineage_file_path,
313            updated_at,
314            database_id,
315        })
316    }
317}
318
319#[derive(Error, Debug)]
320pub enum CollectionToProtoError {
321    #[error("Could not serialize config: {0}")]
322    ConfigSerialization(#[from] serde_json::Error),
323}
324
325impl ChromaError for CollectionToProtoError {
326    fn code(&self) -> ErrorCodes {
327        match self {
328            CollectionToProtoError::ConfigSerialization(_) => ErrorCodes::Internal,
329        }
330    }
331}
332
333impl TryFrom<Collection> for chroma_proto::Collection {
334    type Error = CollectionToProtoError;
335
336    fn try_from(value: Collection) -> Result<Self, Self::Error> {
337        Ok(Self {
338            id: value.collection_id.0.to_string(),
339            name: value.name,
340            configuration_json_str: serde_json::to_string(&value.config)?,
341            schema_str: value
342                .schema
343                .map(|s| serde_json::to_string(&s))
344                .transpose()?,
345            metadata: value.metadata.map(Into::into),
346            dimension: value.dimension,
347            tenant: value.tenant,
348            database: value.database,
349            log_position: value.log_position,
350            version: value.version,
351            total_records_post_compaction: value.total_records_post_compaction,
352            size_bytes_post_compaction: value.size_bytes_post_compaction,
353            last_compaction_time_secs: value.last_compaction_time_secs,
354            version_file_path: value.version_file_path,
355            root_collection_id: value.root_collection_id.map(|uuid| uuid.0.to_string()),
356            lineage_file_path: value.lineage_file_path,
357            updated_at: Some(value.updated_at.into()),
358            database_id: Some(value.database_id.0.to_string()),
359        })
360    }
361}
362
363#[derive(Clone, Debug)]
364pub struct CollectionAndSegments {
365    pub collection: Collection,
366    pub metadata_segment: Segment,
367    pub record_segment: Segment,
368    pub vector_segment: Segment,
369}
370
371impl CollectionAndSegments {
372    // If dimension is not set and vector segment has no files,
373    // we assume this is an uninitialized collection
374    pub fn is_uninitialized(&self) -> bool {
375        self.collection.dimension.is_none() && self.vector_segment.file_path.is_empty()
376    }
377
378    pub fn test(dim: i32) -> Self {
379        let collection = Collection::test_collection(dim);
380        let collection_uuid = collection.collection_id;
381        Self {
382            collection,
383            metadata_segment: test_segment(collection_uuid, SegmentScope::METADATA),
384            record_segment: test_segment(collection_uuid, SegmentScope::RECORD),
385            vector_segment: test_segment(collection_uuid, SegmentScope::VECTOR),
386        }
387    }
388}
389
390#[derive(Deserialize, Serialize, Debug, Clone)]
391#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
392pub struct CreateCollectionPayload {
393    pub name: String,
394    pub schema: Option<Schema>,
395    pub configuration: Option<CollectionConfiguration>,
396    pub metadata: Option<Metadata>,
397    #[serde(default)]
398    pub get_or_create: bool,
399}
400
401#[derive(Deserialize, Serialize, Debug, Clone)]
402#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
403pub struct UpdateCollectionPayload {
404    pub new_name: Option<String>,
405    pub new_metadata: Option<UpdateMetadata>,
406    pub new_configuration: Option<UpdateCollectionConfiguration>,
407}
408
409#[cfg(test)]
410mod test {
411    use super::*;
412
413    #[test]
414    fn test_collection_try_from() {
415        // Create a valid Schema and serialize it
416        let schema = Schema::new_default(crate::KnnIndex::Spann);
417        let schema_str = serde_json::to_string(&schema).unwrap();
418
419        let proto_collection = chroma_proto::Collection {
420            id: "00000000-0000-0000-0000-000000000000".to_string(),
421            name: "foo".to_string(),
422            configuration_json_str: "{\"a\": \"param\", \"b\": \"param2\", \"3\": true}"
423                .to_string(),
424            schema_str: Some(schema_str),
425            metadata: None,
426            dimension: None,
427            tenant: "baz".to_string(),
428            database: "qux".to_string(),
429            log_position: 0,
430            version: 0,
431            total_records_post_compaction: 0,
432            size_bytes_post_compaction: 0,
433            last_compaction_time_secs: 0,
434            version_file_path: Some("version_file_path".to_string()),
435            root_collection_id: Some("00000000-0000-0000-0000-000000000000".to_string()),
436            lineage_file_path: Some("lineage_file_path".to_string()),
437            updated_at: Some(prost_types::Timestamp {
438                seconds: 1,
439                nanos: 1,
440            }),
441            database_id: Some("00000000-0000-0000-0000-000000000000".to_string()),
442        };
443        let converted_collection: Collection = proto_collection.try_into().unwrap();
444        assert_eq!(
445            converted_collection.collection_id,
446            CollectionUuid(Uuid::nil())
447        );
448        assert_eq!(converted_collection.name, "foo".to_string());
449        assert_eq!(converted_collection.metadata, None);
450        assert_eq!(converted_collection.dimension, None);
451        assert_eq!(converted_collection.tenant, "baz".to_string());
452        assert_eq!(converted_collection.database, "qux".to_string());
453        assert_eq!(converted_collection.total_records_post_compaction, 0);
454        assert_eq!(converted_collection.size_bytes_post_compaction, 0);
455        assert_eq!(converted_collection.last_compaction_time_secs, 0);
456        assert_eq!(
457            converted_collection.version_file_path,
458            Some("version_file_path".to_string())
459        );
460        assert_eq!(
461            converted_collection.root_collection_id,
462            Some(CollectionUuid(Uuid::nil()))
463        );
464        assert_eq!(
465            converted_collection.lineage_file_path,
466            Some("lineage_file_path".to_string())
467        );
468        assert_eq!(
469            converted_collection.updated_at,
470            SystemTime::UNIX_EPOCH + Duration::new(1, 1)
471        );
472        assert_eq!(converted_collection.database_id, DatabaseUuid(Uuid::nil()));
473    }
474
475    #[test]
476    fn storage_prefix_for_log_format() {
477        let collection_id = Uuid::parse_str("34e72052-5e60-47cb-be88-19a9715b7026")
478            .map(CollectionUuid)
479            .unwrap();
480        let prefix = collection_id.storage_prefix_for_log();
481        assert_eq!("logs/34e72052-5e60-47cb-be88-19a9715b7026", prefix);
482    }
483}
chroma_types/collection.rs

chroma_types/
collection.rs
