multivector/
collections.rs1use 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}