completr 0.2.0__tar.gz → 0.2.2__tar.gz

{completr-0.2.0 → completr-0.2.2}/Cargo.lock

@@ -735,7 +735,7 @@ dependencies = [
 
 [[package]]
 name = "completr"
-version = "0.2.0"
+version = "0.2.2"
 dependencies = [
  "arc-swap",
  "async-trait",
@@ -765,7 +765,7 @@ dependencies = [
 
 [[package]]
 name = "completr-cli"
-version = "0.2.0"
+version = "0.2.2"
 dependencies = [
  "clap",
  "completr",
@@ -776,7 +776,7 @@ dependencies = [
 
 [[package]]
 name = "completr-py"
-version = "0.2.0"
+version = "0.2.2"
 dependencies = [
  "completr",
  "mimalloc",

{completr-0.2.0 → completr-0.2.2}/Cargo.toml

@@ -3,7 +3,7 @@ resolver = "2"
 members = ["crates/completr", "crates/completr-py"]
 
 [workspace.package]
-version = "0.2.0"
+version = "0.2.2"
 edition = "2021"
 rust-version = "1.94.1"
 license = "MIT"

{completr-0.2.0 → completr-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: completr
-Version: 0.2.0
+Version: 0.2.2
 Classifier: Development Status :: 4 - Beta
 Classifier: Intended Audience :: Developers
 Classifier: Operating System :: MacOS

{completr-0.2.0 → completr-0.2.2}/crates/completr/Cargo.toml

@@ -53,6 +53,7 @@ aws-credentials = ["aws", "dep:aws-config", "dep:aws-credential-types", "dep:asy
 [dev-dependencies]
 turbovec = "1.0.0"
 tempfile = "3"
+async-trait = "0.1"
 xxhash-rust = { version = "0.8", features = ["xxh3"] }
 serde_json = { version = "1", features = ["float_roundtrip"] }
 tokio = { version = "1", features = ["macros", "rt-multi-thread"] }

{completr-0.2.0 → completr-0.2.2}/crates/completr/src/client.rs

@@ -28,6 +28,8 @@ pub struct ConnectOptions {
     pub cache_dir: Option<PathBuf>,
     /// Storage options, such as credentials or a region, passed to the object store.
     pub storage: Vec<(String, String)>,
+    /// Tags on every object written, e.g. to select bucket lifecycle rules.
+    pub tags: Vec<(String, String)>,
     pub build: BuildOptions,
     pub index: IndexOptions,
 }
@@ -38,6 +40,7 @@ impl Default for ConnectOptions {
             sync_every: Some(Duration::from_secs(5)),
             cache_dir: None,
             storage: Vec::new(),
+            tags: Vec::new(),
             build: BuildOptions::default(),
             index: IndexOptions::default(),
         }
@@ -48,6 +51,7 @@ crate::setters!(ConnectOptions {
     sync_every: Option<Duration>,
     cache_dir: Option<PathBuf>,
     storage: Vec<(String, String)>,
+    tags: Vec<(String, String)>,
     build: BuildOptions,
     index: IndexOptions,
 });
@@ -218,7 +222,8 @@ const LEGACY_SETTINGS: &str = "collection/";
 pub async fn connect(url: &str, options: ConnectOptions) -> Result<Client, Error> {
     let mut database = Database::open(url, options.storage)
         .await?
-        .with_build_options(options.build);
+        .with_build_options(options.build)
+        .with_tags(options.tags);
     if let Some(dir) = &options.cache_dir {
         database = database.with_cache_dir(dir)?;
     }

{completr-0.2.0 → completr-0.2.2}/crates/completr/src/database.rs

@@ -196,6 +196,17 @@ impl Database {
         Ok(Self::new(Store::open_with(url, options).await?))
     }
 
+    /// Tags every object this database writes, e.g. to select bucket lifecycle rules.
+    pub fn with_tags<K: AsRef<str>, V: AsRef<str>>(
+        self,
+        tags: impl IntoIterator<Item = (K, V)>,
+    ) -> Self {
+        Self {
+            store: self.store.with_tags(tags),
+            ..self
+        }
+    }
+
     /// Caches downloaded segments in `dir` and memory-maps them.
     pub fn with_cache_dir(self, dir: impl AsRef<std::path::Path>) -> Result<Self, Error> {
         Ok(Self {

{completr-0.2.0 → completr-0.2.2}/crates/completr/src/store.rs

@@ -7,7 +7,7 @@ use futures::{StreamExt, TryStreamExt};
 use object_store::local::LocalFileSystem;
 use object_store::memory::InMemory;
 use object_store::path::Path;
-use object_store::{ObjectStore, ObjectStoreExt, PutMode, PutOptions, PutPayload};
+use object_store::{ObjectStore, ObjectStoreExt, PutMode, PutOptions, PutPayload, TagSet};
 use url::Url;
 
 use crate::{Error, Segment};
@@ -29,6 +29,8 @@ pub struct Store {
     /// Where remote segments are downloaded to and memory-mapped from, per store URL.
     cache: Option<PathBuf>,
     namespace: String,
+    /// Tags on every object written, e.g. to select bucket lifecycle rules.
+    tags: TagSet,
 }
 
 impl Store {
@@ -80,6 +82,7 @@ impl Store {
             local_root: None,
             cache: None,
             namespace,
+            tags: TagSet::default(),
         })
     }
 
@@ -93,6 +96,7 @@ impl Store {
             local_root: Some(root),
             cache: None,
             namespace: String::new(),
+            tags: TagSet::default(),
         })
     }
 
@@ -115,9 +119,31 @@ impl Store {
             .to_vec())
     }
 
+    /// Tags every object this store writes from now on; S3 and Azure keep them, other stores ignore them.
+    pub fn with_tags<K: AsRef<str>, V: AsRef<str>>(
+        mut self,
+        tags: impl IntoIterator<Item = (K, V)>,
+    ) -> Self {
+        let mut set = TagSet::default();
+        for (key, value) in tags {
+            set.push(key.as_ref(), value.as_ref());
+        }
+        self.tags = set;
+        self
+    }
+
+    /// The tags written with every object, URL-encoded as S3 takes them.
+    pub fn tags(&self) -> &str {
+        self.tags.encoded()
+    }
+
     pub async fn put(&self, key: &str, data: Vec<u8>) -> Result<(), Error> {
+        let options = PutOptions {
+            tags: self.tags.clone(),
+            ..PutOptions::default()
+        };
         self.inner
-            .put(&self.path(key), PutPayload::from(data))
+            .put_opts(&self.path(key), PutPayload::from(data), options)
             .await?;
         Ok(())
     }
@@ -127,6 +153,7 @@ impl Store {
     pub async fn put_if_absent(&self, key: &str, data: Vec<u8>) -> Result<bool, Error> {
         let options = PutOptions {
             mode: PutMode::Create,
+            tags: self.tags.clone(),
             ..PutOptions::default()
         };
         match self
@@ -431,6 +458,13 @@ impl BlockingStore {
         self.0.with_cache_dir(dir).map(Self)
     }
 
+    pub fn with_tags<K: AsRef<str>, V: AsRef<str>>(
+        self,
+        tags: impl IntoIterator<Item = (K, V)>,
+    ) -> Self {
+        Self(self.0.with_tags(tags))
+    }
+
     pub fn prune_cache<'a>(&self, keep: impl IntoIterator<Item = &'a str>) -> Result<usize, Error> {
         self.0.prune_cache(keep)
     }
@@ -557,6 +591,123 @@ mod tests {
     use super::*;
     use crate::Document;
 
+    /// An in-memory store that records the tags each write carries.
+    #[derive(Debug, Default)]
+    struct Recording {
+        inner: InMemory,
+        tags: std::sync::Mutex<Vec<String>>,
+    }
+
+    impl std::fmt::Display for Recording {
+        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+            write!(f, "Recording")
+        }
+    }
+
+    #[async_trait::async_trait]
+    impl ObjectStore for Recording {
+        async fn put_opts(
+            &self,
+            location: &Path,
+            payload: PutPayload,
+            opts: PutOptions,
+        ) -> object_store::Result<object_store::PutResult> {
+            self.tags
+                .lock()
+                .unwrap()
+                .push(opts.tags.encoded().to_owned());
+            self.inner.put_opts(location, payload, opts).await
+        }
+
+        async fn put_multipart_opts(
+            &self,
+            location: &Path,
+            opts: object_store::PutMultipartOptions,
+        ) -> object_store::Result<Box<dyn object_store::MultipartUpload>> {
+            self.inner.put_multipart_opts(location, opts).await
+        }
+
+        async fn get_opts(
+            &self,
+            location: &Path,
+            options: object_store::GetOptions,
+        ) -> object_store::Result<object_store::GetResult> {
+            self.inner.get_opts(location, options).await
+        }
+
+        fn delete_stream(
+            &self,
+            locations: futures::stream::BoxStream<'static, object_store::Result<Path>>,
+        ) -> futures::stream::BoxStream<'static, object_store::Result<Path>> {
+            self.inner.delete_stream(locations)
+        }
+
+        fn list(
+            &self,
+            prefix: Option<&Path>,
+        ) -> futures::stream::BoxStream<'static, object_store::Result<object_store::ObjectMeta>>
+        {
+            self.inner.list(prefix)
+        }
+
+        async fn list_with_delimiter(
+            &self,
+            prefix: Option<&Path>,
+        ) -> object_store::Result<object_store::ListResult> {
+            self.inner.list_with_delimiter(prefix).await
+        }
+
+        async fn copy_opts(
+            &self,
+            from: &Path,
+            to: &Path,
+            options: object_store::CopyOptions,
+        ) -> object_store::Result<()> {
+            self.inner.copy_opts(from, to, options).await
+        }
+    }
+
+    #[tokio::test]
+    async fn every_write_carries_the_tags() {
+        let recording = Arc::new(Recording::default());
+        let store = Store {
+            inner: recording.clone(),
+            prefix: Path::default(),
+            local_root: None,
+            cache: None,
+            namespace: "recording".into(),
+            tags: TagSet::default(),
+        }
+        .with_tags([
+            ("LifecycleRule", "KeepForever"),
+            ("team", "search & ranking"),
+        ]);
+        assert_eq!(
+            store.tags(),
+            "LifecycleRule=KeepForever&team=search+%26+ranking"
+        );
+
+        let database = crate::Database::new(store);
+        let mut txn = database.begin().await.unwrap();
+        txn.append_documents("songs", [Document::new(1, "Rust", 0.5)], [])
+            .unwrap();
+        txn.commit().await.unwrap();
+        let mut changes = crate::ChangeSet::new();
+        changes.upsert("songs", [Document::new(2, "Go", 0.5)]);
+        database.submit(changes).await.unwrap();
+        database
+            .acquire_lease("ingestor", "me", std::time::Duration::from_secs(60))
+            .await
+            .unwrap();
+
+        let written = recording.tags.lock().unwrap().clone();
+        // A segment, a version, a change set and a lease.
+        assert!(written.len() >= 4, "{written:?}");
+        assert!(written
+            .iter()
+            .all(|t| t == "LifecycleRule=KeepForever&team=search+%26+ranking"));
+    }
+
     #[test]
     fn local_and_memory_round_trip() {
         let dir = tempfile::tempdir().unwrap();

{completr-0.2.0 → completr-0.2.2}/crates/completr-py/Cargo.toml

@@ -18,4 +18,7 @@ pyo3 = { version = "0.29", features = ["extension-module", "abi3-py311"] }
 pyo3-log = "0.13"
 tracing = { version = "0.1", default-features = false, features = ["log"] }
 completr_rs = { package = "completr", path = "../completr", features = ["aws-credentials", "gcp", "azure"] }
+
+# Not on macOS, where every mimalloc in a process shares fixed thread slots and pyarrow ships one.
+[target.'cfg(target_os = "linux")'.dependencies]
 mimalloc = { version = "0.1", default-features = false }

{completr-0.2.0 → completr-0.2.2}/crates/completr-py/python/completr/__init__.pyi

@@ -21,6 +21,7 @@ def connect(
     options: Mapping[str, str] | None = None,
     cache_dir: str | PathLike[str] | None = None,
     *,
+    tags: Mapping[str, str] | None = None,
     min_word_chars: int = 3,
     max_edit_distance: int = 2,
     fuzzy_prefix_chars: int = 7,

{completr-0.2.0 → completr-0.2.2}/crates/completr-py/python/completr/completr.pyi

@@ -311,7 +311,12 @@ class Engine:
 
 class Store:
     def __init__(
-        self, url: str, options: Mapping[str, str] | None = None, cache_dir: str | PathLike[str] | None = None
+        self,
+        url: str,
+        options: Mapping[str, str] | None = None,
+        cache_dir: str | PathLike[str] | None = None,
+        *,
+        tags: Mapping[str, str] | None = None,
     ) -> None: ...
     def prune_cache(self, keep: Sequence[str]) -> int: ...
     def get(self, key: str) -> bytes: ...
@@ -394,6 +399,7 @@ class Database:
         options: Mapping[str, str] | None = None,
         cache_dir: str | PathLike[str] | None = None,
         *,
+        tags: Mapping[str, str] | None = None,
         min_word_chars: int = 3,
         max_edit_distance: int = 2,
         fuzzy_prefix_chars: int = 7,
@@ -564,6 +570,7 @@ class Client:
         sync_every: float | None = 5.0,
         cache_dir: str | PathLike[str] | None = None,
         options: Mapping[str, str] | None = None,
+        tags: Mapping[str, str] | None = None,
         min_word_chars: int = 3,
         max_edit_distance: int = 2,
         fuzzy_prefix_chars: int = 7,

{completr-0.2.0 → completr-0.2.2}/crates/completr-py/src/lib.rs

@@ -10,7 +10,8 @@ use pyo3::prelude::*;
 use pyo3::sync::PyOnceLock;
 use pyo3::types::{PyBytes, PyDict, PyList, PyString, PyType};
 
-// Returns freed memory to the system, which the platform allocator keeps after a build.
+// Returns freed memory to the system, which glibc keeps after a build.
+#[cfg(target_os = "linux")]
 #[global_allocator]
 static ALLOCATOR: mimalloc::MiMalloc = mimalloc::MiMalloc;
 
@@ -1542,17 +1543,19 @@ struct Store(completr_rs::BlockingStore);
 #[pymethods]
 impl Store {
     #[new]
-    #[pyo3(signature = (url, options = None, cache_dir = None))]
+    #[pyo3(signature = (url, options = None, cache_dir = None, *, tags = None))]
     fn new(
         py: Python<'_>,
         url: &str,
         options: Option<HashMap<String, String>>,
         cache_dir: Option<PathBuf>,
+        tags: Option<HashMap<String, String>>,
     ) -> PyResult<Self> {
         let options = options.unwrap_or_default();
         let store = py
             .detach(|| completr_rs::BlockingStore::open_with(url, options))
-            .map_err(to_py_err)?;
+            .map_err(to_py_err)?
+            .with_tags(tags.unwrap_or_default());
         let store = match cache_dir {
             Some(dir) => store.with_cache_dir(dir).map_err(to_py_err)?,
             None => store,
@@ -1636,11 +1639,11 @@ impl Database {
 
 #[pymethods]
 impl Database {
-    /// `cache_dir` keeps downloaded segments on local disk, memory-mapped. The keyword settings
-    /// apply to segments this database builds from documents.
+    /// `cache_dir` keeps downloaded segments on local disk, memory-mapped; `tags` go on every object
+    /// written. The other keyword settings apply to segments this database builds from documents.
     #[new]
     #[pyo3(signature = (
-        url, options = None, cache_dir = None, *,
+        url, options = None, cache_dir = None, *, tags = None,
         min_word_chars = 3, max_edit_distance = 2, fuzzy_prefix_chars = 7, vector_bits = 4, compact_keys = false, build_threads = 1,
     ))]
     #[allow(clippy::too_many_arguments)]
@@ -1649,6 +1652,7 @@ impl Database {
         url: &str,
         options: Option<HashMap<String, String>>,
         cache_dir: Option<PathBuf>,
+        tags: Option<HashMap<String, String>>,
         min_word_chars: u8,
         max_edit_distance: u8,
         fuzzy_prefix_chars: u8,
@@ -1659,7 +1663,8 @@ impl Database {
         let options = options.unwrap_or_default();
         let database = py
             .detach(|| completr_rs::block_on(completr_rs::Database::open(url, options)))
-            .map_err(to_py_err)?;
+            .map_err(to_py_err)?
+            .with_tags(tags.unwrap_or_default());
         let database = match cache_dir {
             Some(dir) => database.with_cache_dir(dir).map_err(to_py_err)?,
             None => database,
@@ -2518,7 +2523,7 @@ impl Client {
     /// `None` syncs only on `sync()`. The keyword settings apply to segments this client builds.
     #[new]
     #[pyo3(signature = (
-        url = "memory://", *, sync_every = Some(5.0), cache_dir = None, options = None,
+        url = "memory://", *, sync_every = Some(5.0), cache_dir = None, options = None, tags = None,
         min_word_chars = 3, max_edit_distance = 2, fuzzy_prefix_chars = 7, vector_bits = 4, compact_keys = false, build_threads = 1,
     ))]
     #[allow(clippy::too_many_arguments)]
@@ -2528,6 +2533,7 @@ impl Client {
         sync_every: Option<f64>,
         cache_dir: Option<PathBuf>,
         options: Option<HashMap<String, String>>,
+        tags: Option<HashMap<String, String>>,
         min_word_chars: u8,
         max_edit_distance: u8,
         fuzzy_prefix_chars: u8,
@@ -2547,6 +2553,7 @@ impl Client {
             .sync_every(sync_every.map(|s| Duration::from_secs_f64(s.max(0.001))))
             .cache_dir(cache_dir)
             .storage(options.unwrap_or_default().into_iter().collect())
+            .tags(tags.unwrap_or_default().into_iter().collect())
             .build(build);
         let client = py
             .detach(|| completr_rs::block_on(completr_rs::connect(url, connect)))

completr-0.2.2/crates/completr-py/tests/test_allocators.py

@@ -0,0 +1,20 @@
+import pytest
+
+import completr
+
+pyarrow = pytest.importorskip("pyarrow")
+
+
+def test_completes_alongside_pyarrow_allocations(tmp_path):
+    pool = pyarrow.default_memory_pool()
+    table = pyarrow.table({"text": [f"skill {i}" for i in range(10_000)]})
+    assert pool.bytes_allocated() > 0
+    db = completr.connect(str(tmp_path / "db"))
+    with db.namespace("en-US").begin() as txn:
+        txn.overwrite("owner", [{"id": f"s{i}", "text": f"Skill {i}", "popularity": 0.5} for i in range(2_000)])
+    namespace = completr.connect(str(tmp_path / "db")).engine(sync_every=None).namespace("en-US")
+    for _ in range(200):
+        hits = namespace.complete("skill 1", ["owner"])
+        assert [str(h.id) for h in hits]
+        table = table.slice(1)
+    assert len(table) > 0

{completr-0.2.0 → completr-0.2.2}/crates/completr-py/tests/test_database.py

@@ -3,6 +3,8 @@ import time
 
 import pytest
 
+import completr
+
 from completr import ConflictError
 from completr import ChangeSet, Database, Engine, Replica, Store, Ingestor
 
@@ -184,3 +186,14 @@ def test_threaded_queries_during_updates(database):
     assert len({id(index) for index, _, _ in samples}) > 1
     for index, query, result in samples:
         assert [(h.id, h.score, h.kind) for h in index.complete(query, 10)] == result
+
+
+def test_tags_are_accepted_everywhere_and_ignored_by_local_stores(tmp_path):
+    tags = {"LifecycleRule": "KeepForever"}
+    db = completr.connect(str(tmp_path / "db"), tags=tags)
+    with db.begin() as txn:
+        txn.append("songs", DOCS[:2])
+    assert db.index_names() == ["songs"]
+    completr.Store(str(tmp_path / "store"), tags=tags).put("a", b"x")
+    client = completr.Client(str(tmp_path / "client"), sync_every=None, tags=tags)
+    client.create_collection("songs").add(DOCS[:2])