Skip to main content

multivector/
collections.rs

1//! Collection namespaces share the engine contract, not document IDs or indexes.
2use crate::{Durability, IndexConfig, IndexError, MultiVectorIndex};
3use std::{
4    collections::BTreeMap,
5    fs,
6    path::{Path, PathBuf},
7    sync::{Arc, Mutex},
8};
9
10pub struct Collections {
11    root: PathBuf,
12    durability: Durability,
13    indexes: Mutex<BTreeMap<String, Arc<MultiVectorIndex>>>,
14}
15impl Collections {
16    pub fn open(root: impl AsRef<Path>, durability: Durability) -> Result<Self, IndexError> {
17        let root = root.as_ref().to_owned();
18        fs::create_dir_all(&root)?;
19        let mut indexes = BTreeMap::new();
20        for entry in fs::read_dir(&root)? {
21            let entry = entry?;
22            if !entry.file_type()?.is_dir() {
23                continue;
24            }
25            let name = entry
26                .file_name()
27                .into_string()
28                .map_err(|_| IndexError::Invalid("invalid collection name".into()))?;
29            validate_name(&name)?;
30            if entry.path().join("manifest.json").exists() {
31                indexes.insert(
32                    name,
33                    Arc::new(MultiVectorIndex::open_existing(entry.path(), durability)?),
34                );
35            }
36        }
37        Ok(Self {
38            root,
39            durability,
40            indexes: Mutex::new(indexes),
41        })
42    }
43    pub fn create(
44        &self,
45        name: &str,
46        config: IndexConfig,
47    ) -> Result<Arc<MultiVectorIndex>, IndexError> {
48        validate_name(name)?;
49        let mut indexes = self.indexes.lock().unwrap();
50        if indexes.contains_key(name) {
51            return Err(IndexError::Invalid("collection already exists".into()));
52        }
53        let index = Arc::new(MultiVectorIndex::open_with_durability(
54            self.root.join(name),
55            config,
56            self.durability,
57        )?);
58        index.initialize()?;
59        indexes.insert(name.into(), Arc::clone(&index));
60        Ok(index)
61    }
62    pub fn get(&self, name: &str) -> Option<Arc<MultiVectorIndex>> {
63        self.indexes.lock().unwrap().get(name).cloned()
64    }
65    pub fn names(&self) -> Vec<String> {
66        self.indexes.lock().unwrap().keys().cloned().collect()
67    }
68}
69fn validate_name(name: &str) -> Result<(), IndexError> {
70    if name.is_empty()
71        || name.len() > 64
72        || !name
73            .bytes()
74            .all(|c| c.is_ascii_alphanumeric() || c == b'_' || c == b'-')
75    {
76        return Err(IndexError::Invalid(
77            "collection names require 1..=64 ASCII letters, digits, '-' or '_'".into(),
78        ));
79    }
80    Ok(())
81}