From 9f3d0cd9362fe2b0685aa77545c82ac679010306 Mon Sep 17 00:00:00 2001 From: Allen Cheng Date: Mon, 8 Jun 2026 23:33:45 -0700 Subject: [PATCH 1/2] feat: add logical delete API --- README.md | 9 ++ crates/lance-context-core/src/record.rs | 9 ++ crates/lance-context-core/src/serde.rs | 1 + crates/lance-context-core/src/store.rs | 155 +++++++++++++++++++++++- python/python/lance_context/api.py | 17 +++ python/python/tests/test_context.py | 13 ++ python/src/lib.rs | 27 +++++ python/tests/test_delete.py | 64 ++++++++++ python/tests/test_search.py | 51 ++++++++ 9 files changed, 345 insertions(+), 1 deletion(-) create mode 100644 python/tests/test_delete.py diff --git a/README.md b/README.md index d6261d6..905ada8 100644 --- a/README.md +++ b/README.md @@ -60,6 +60,8 @@ ctx.add( external_id="conversation-2026-03-01#turn-1", ) print(ctx.get(external_id="conversation-2026-03-01#turn-1")) +ctx.delete(external_id="conversation-2026-03-01#turn-1") +assert ctx.get(external_id="conversation-2026-03-01#turn-1") is None from PIL import Image image = Image.new("RGB", (2, 2), color="teal") @@ -133,6 +135,13 @@ metrics = ctx.compact() print(f"Compaction removed {metrics['fragments_removed']} fragments") ``` +`delete()` and its alias `forget()` write a versioned tombstone for the target +record and return `False` if the id is already absent. Default `list()`, +`get()`, and `search()` calls hide tombstoned records, but this is logical +forgetting rather than guaranteed physical erasure: older dataset versions and +underlying files may still contain the original payload until retention or +physical cleanup policies remove them. + ### Rust ```rust diff --git a/crates/lance-context-core/src/record.rs b/crates/lance-context-core/src/record.rs index 95cf27b..3f88f18 100644 --- a/crates/lance-context-core/src/record.rs +++ b/crates/lance-context-core/src/record.rs @@ -1,5 +1,7 @@ use chrono::{DateTime, Utc}; +use crate::serde::CONTENT_TYPE_TOMBSTONE; + /// Structured metadata captured alongside each context entry. #[derive(Debug, Clone, Default)] pub struct StateMetadata { @@ -26,6 +28,13 @@ pub struct ContextRecord { pub embedding: Option>, } +impl ContextRecord { + #[must_use] + pub fn is_tombstone(&self) -> bool { + self.content_type == CONTENT_TYPE_TOMBSTONE + } +} + /// Result returned from a vector similarity search. #[derive(Debug, Clone)] pub struct SearchResult { diff --git a/crates/lance-context-core/src/serde.rs b/crates/lance-context-core/src/serde.rs index d0c97b1..2af01c3 100644 --- a/crates/lance-context-core/src/serde.rs +++ b/crates/lance-context-core/src/serde.rs @@ -5,6 +5,7 @@ use serde::{Deserialize, Serialize}; pub const CONTENT_TYPE_TEXT: &str = "text/plain"; pub const CONTENT_TYPE_ARROW_STREAM: &str = "application/vnd.apache.arrow.stream"; +pub const CONTENT_TYPE_TOMBSTONE: &str = "application/vnd.lance-context.tombstone"; #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct SerializedContent { diff --git a/crates/lance-context-core/src/store.rs b/crates/lance-context-core/src/store.rs index 5f1f6a1..aa58366 100644 --- a/crates/lance-context-core/src/store.rs +++ b/crates/lance-context-core/src/store.rs @@ -32,6 +32,7 @@ use tracing::{error, info, warn}; use uuid::Uuid; use crate::record::{ContextRecord, SearchResult, StateMetadata}; +use crate::serde::CONTENT_TYPE_TOMBSTONE; /// Embedding length used for the semantic index column. const DEFAULT_EMBEDDING_DIM: i32 = 1536; @@ -204,6 +205,13 @@ impl ContextStore { } self.validate_unique_ids(entries).await?; + self.write_entries(entries).await + } + + async fn write_entries(&mut self, entries: &[ContextRecord]) -> LanceResult { + if entries.is_empty() { + return Ok(self.dataset.manifest.version); + } // Group entries by (bot_id, session_id) let mut groups: HashMap<(Option, Option), Vec> = @@ -252,10 +260,56 @@ impl ContextStore { Ok(self.dataset.manifest.version) } + /// Logically forget a record by internal storage id. + /// + /// This writes a tombstone with the same primary key, preserving prior + /// dataset versions while hiding the record from default reads. + pub async fn delete_by_id(&mut self, id: &str) -> LanceResult { + let Some(record) = self.get_by_id(id).await? else { + return Ok(false); + }; + self.write_tombstone_for(record).await?; + Ok(true) + } + + /// Logically forget a record by caller-supplied external id. + pub async fn delete_by_external_id(&mut self, external_id: &str) -> LanceResult { + let Some(record) = self.get_by_external_id(external_id).await? else { + return Ok(false); + }; + self.write_tombstone_for(record).await?; + Ok(true) + } + + async fn write_tombstone_for(&mut self, record: ContextRecord) -> LanceResult { + let tombstone = ContextRecord { + id: record.id, + external_id: record.external_id, + run_id: record.run_id, + bot_id: record.bot_id, + session_id: record.session_id, + created_at: Utc::now(), + role: record.role, + state_metadata: None, + content_type: CONTENT_TYPE_TOMBSTONE.to_string(), + text_payload: None, + binary_payload: None, + embedding: None, + }; + self.write_entries(std::slice::from_ref(&tombstone)).await + } + async fn validate_unique_ids(&self, entries: &[ContextRecord]) -> LanceResult<()> { let mut ids = HashSet::new(); let mut external_ids = HashSet::new(); for entry in entries { + if entry.is_tombstone() { + return Err(ArrowError::InvalidArgumentError(format!( + "content_type '{}' is reserved for internal tombstones", + CONTENT_TYPE_TOMBSTONE + )) + .into()); + } if !ids.insert(entry.id.as_str()) { return Err(ArrowError::InvalidArgumentError(format!( "duplicate id '{}' in batch", @@ -333,7 +387,11 @@ impl ContextStore { let mut stream = scanner.try_into_stream().await?; let mut results = Vec::new(); while let Some(batch) = stream.try_next().await? { - results.extend(batch_to_records(&batch)?); + results.extend( + batch_to_records(&batch)? + .into_iter() + .filter(|record| !record.is_tombstone()), + ); } if let Some(offset) = offset { @@ -1347,6 +1405,101 @@ mod tests { }); } + #[test] + fn add_rejects_reserved_tombstone_content_type() { + let dir = TempDir::new().unwrap(); + let uri = dir.path().to_string_lossy().to_string(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + runtime.block_on(async { + let mut store = ContextStore::open(&uri).await.unwrap(); + let mut record = text_record("a", 0.0); + record.content_type = CONTENT_TYPE_TOMBSTONE.to_string(); + + let err = store.add(&[record]).await.unwrap_err(); + let message = err.to_string(); + assert!( + message.contains("reserved") && message.contains("tombstone"), + "unexpected error message: {message}" + ); + }); + } + + #[test] + fn delete_by_external_id_hides_record_from_default_reads() { + let dir = TempDir::new().unwrap(); + let uri = dir.path().to_string_lossy().to_string(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + runtime.block_on(async { + let mut store = ContextStore::open(&uri).await.unwrap(); + let mut first = text_record("a", 0.0); + first.external_id = Some("doc-123#chunk-1".to_string()); + let second = text_record("b", 2.0); + store.add(&[first.clone(), second.clone()]).await.unwrap(); + + assert!(store + .delete_by_external_id("doc-123#chunk-1") + .await + .unwrap()); + + assert!(store + .get_by_external_id("doc-123#chunk-1") + .await + .unwrap() + .is_none()); + assert!(store.get_by_id(&first.id).await.unwrap().is_none()); + + let records = store.list(None, None).await.unwrap(); + assert_eq!(records.len(), 1); + assert_eq!(records[0].id, second.id); + + let query = make_embedding(0.0); + let hits = store.search(&query, Some(10)).await.unwrap(); + assert_eq!(hits.len(), 1); + assert_eq!(hits[0].record.id, second.id); + }); + } + + #[test] + fn delete_missing_id_is_noop() { + let dir = TempDir::new().unwrap(); + let uri = dir.path().to_string_lossy().to_string(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + runtime.block_on(async { + let mut store = ContextStore::open(&uri).await.unwrap(); + assert!(!store.delete_by_id("missing").await.unwrap()); + assert!(!store.delete_by_external_id("missing").await.unwrap()); + }); + } + + #[test] + fn external_id_can_be_reused_after_delete() { + let dir = TempDir::new().unwrap(); + let uri = dir.path().to_string_lossy().to_string(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + runtime.block_on(async { + let mut store = ContextStore::open(&uri).await.unwrap(); + let mut first = text_record("a", 0.0); + first.external_id = Some("doc-123#chunk-1".to_string()); + store.add(std::slice::from_ref(&first)).await.unwrap(); + assert!(store + .delete_by_external_id("doc-123#chunk-1") + .await + .unwrap()); + + let mut replacement = text_record("b", 1.0); + replacement.external_id = first.external_id.clone(); + store.add(std::slice::from_ref(&replacement)).await.unwrap(); + + let by_external_id = store + .get_by_external_id("doc-123#chunk-1") + .await + .unwrap() + .unwrap(); + assert_eq!(by_external_id.id, replacement.id); + assert_eq!(store.list(None, None).await.unwrap().len(), 1); + }); + } + #[test] fn test_region_id_derivation_explicit() { let bot_id = Some("bot-123".to_string()); diff --git a/python/python/lance_context/api.py b/python/python/lance_context/api.py index d431def..a6cbab9 100644 --- a/python/python/lance_context/api.py +++ b/python/python/lance_context/api.py @@ -410,6 +410,23 @@ def get( return None return _normalize_record(result) + def delete(self, *, id: str | None = None, external_id: str | None = None) -> bool: + """Logically forget one entry by internal id or caller-supplied external id. + + Returns True when an entry was found and tombstoned, False when the + identifier is already absent. This is a versioned logical delete: + default reads hide the entry, but older dataset versions and underlying + storage files may still contain the original payload until retention or + physical cleanup policies remove them. + """ + if (id is None) == (external_id is None): + raise ValueError("Specify exactly one of id or external_id") + return bool(self._inner.delete(id, external_id)) + + def forget(self, *, id: str | None = None, external_id: str | None = None) -> bool: + """Alias for :meth:`delete`.""" + return self.delete(id=id, external_id=external_id) + def compact( self, *, diff --git a/python/python/tests/test_context.py b/python/python/tests/test_context.py index a286123..aca5d80 100644 --- a/python/python/tests/test_context.py +++ b/python/python/tests/test_context.py @@ -29,3 +29,16 @@ def test_context_snapshot_and_fork(): assert fork.branch() == "branch-a" assert fork.entries() == ctx.entries() assert fork.version() == ctx.version() + + +def test_context_delete_and_forget_by_external_id(tmp_path): + ctx = lc.Context.create(str(tmp_path / "context.lance")) + ctx.add("user", "stale", external_id="doc-1#chunk-1") + + entry = ctx.get(external_id="doc-1#chunk-1") + assert entry is not None + + assert ctx.delete(external_id="doc-1#chunk-1") is True + assert ctx.get(external_id="doc-1#chunk-1") is None + assert ctx.get(id=entry["id"]) is None + assert ctx.forget(external_id="doc-1#chunk-1") is False diff --git a/python/src/lib.rs b/python/src/lib.rs index f40a460..8ba997b 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -316,6 +316,33 @@ impl Context { record.map(|record| record_to_py(py, record)).transpose() } + #[pyo3(signature = (id = None, external_id = None))] + fn delete( + &mut self, + py: Python<'_>, + id: Option, + external_id: Option, + ) -> PyResult { + match (id, external_id) { + (Some(id), None) => py.allow_threads(|| { + self.runtime + .block_on(self.store.delete_by_id(&id)) + .map_err(to_py_err) + }), + (None, Some(external_id)) => py.allow_threads(|| { + self.runtime + .block_on(self.store.delete_by_external_id(&external_id)) + .map_err(to_py_err) + }), + (None, None) => Err(PyRuntimeError::new_err( + "delete() requires either id or external_id", + )), + (Some(_), Some(_)) => Err(PyRuntimeError::new_err( + "delete() accepts only one of id or external_id", + )), + } + } + #[pyo3(signature = (target_rows_per_fragment=None, materialize_deletions=None))] fn compact( &mut self, diff --git a/python/tests/test_delete.py b/python/tests/test_delete.py new file mode 100644 index 0000000..058c975 --- /dev/null +++ b/python/tests/test_delete.py @@ -0,0 +1,64 @@ +"""Tests for logical deletion / forgetting.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pathlib import Path + +from lance_context.api import Context + + +def _embedding(value: float) -> list[float]: + vector = [0.0] * 1536 + vector[0] = value + return vector + + +def test_delete_by_external_id_hides_record_from_default_reads(tmp_path: Path) -> None: + uri = str(tmp_path / "context.lance") + ctx = Context.create(uri) + + ctx.add( + "user", + "stale memory", + embedding=_embedding(0.0), + external_id="doc-123#chunk-4", + ) + ctx.add( + "user", + "fresh memory", + embedding=_embedding(1.0), + external_id="doc-456#chunk-1", + ) + + stale = ctx.get(external_id="doc-123#chunk-4") + assert stale is not None + stale_id = stale["id"] + + assert ctx.delete(external_id="doc-123#chunk-4") is True + assert ctx.delete(external_id="doc-123#chunk-4") is False + assert ctx.get(external_id="doc-123#chunk-4") is None + assert ctx.get(id=stale_id) is None + + entries = ctx.list() + assert [entry["external_id"] for entry in entries] == ["doc-456#chunk-1"] + + hits = ctx.search(_embedding(0.0), limit=10) + assert [hit["external_id"] for hit in hits] == ["doc-456#chunk-1"] + + +def test_external_id_can_be_reused_after_delete(tmp_path: Path) -> None: + uri = str(tmp_path / "context.lance") + ctx = Context.create(uri) + + ctx.add("user", "stale memory", external_id="doc-123#chunk-4") + assert ctx.forget(external_id="doc-123#chunk-4") is True + + ctx.add("user", "replacement memory", external_id="doc-123#chunk-4") + + entry = ctx.get(external_id="doc-123#chunk-4") + assert entry is not None + assert entry["text"] == "replacement memory" + assert [item["text"] for item in ctx.list()] == ["replacement memory"] diff --git a/python/tests/test_search.py b/python/tests/test_search.py index 51dbe1e..e89ab70 100644 --- a/python/tests/test_search.py +++ b/python/tests/test_search.py @@ -15,6 +15,7 @@ def __init__(self) -> None: self.search_calls: list[tuple[list[float], int | None]] = [] self.list_calls: list[tuple[int | None, int | None]] = [] self.get_calls: list[tuple[str | None, str | None]] = [] + self.delete_calls: list[tuple[str | None, str | None]] = [] self.add_calls: list[ tuple[ str, @@ -47,6 +48,10 @@ def get(self, id: str | None, external_id: str | None): return self.list(None, None)[0] return None + def delete(self, id: str | None, external_id: str | None): + self.delete_calls.append((id, external_id)) + return id == "rec-1" or external_id == "source-1" + def search(self, vector: list[float], limit: int | None): self.search_calls.append((vector, limit)) return [ @@ -230,6 +235,52 @@ def test_context_get_requires_exactly_one_identifier(): ctx.get(id="rec-1", external_id="source-1") +def test_context_delete_by_external_id(): + ctx = Context.__new__(Context) + dummy = DummyInner() + ctx._inner = dummy # type: ignore[attr-defined] + + assert ctx.delete(external_id="source-1") is True + assert dummy.delete_calls == [(None, "source-1")] + + +def test_context_delete_by_id(): + ctx = Context.__new__(Context) + dummy = DummyInner() + ctx._inner = dummy # type: ignore[attr-defined] + + assert ctx.delete(id="rec-1") is True + assert dummy.delete_calls == [("rec-1", None)] + + +def test_context_delete_missing_returns_false(): + ctx = Context.__new__(Context) + dummy = DummyInner() + ctx._inner = dummy # type: ignore[attr-defined] + + assert ctx.delete(external_id="missing") is False + + +def test_context_forget_aliases_delete(): + ctx = Context.__new__(Context) + dummy = DummyInner() + ctx._inner = dummy # type: ignore[attr-defined] + + assert ctx.forget(external_id="source-1") is True + assert dummy.delete_calls == [(None, "source-1")] + + +def test_context_delete_requires_exactly_one_identifier(): + ctx = Context.__new__(Context) + dummy = DummyInner() + ctx._inner = dummy # type: ignore[attr-defined] + + with pytest.raises(ValueError, match="exactly one"): + ctx.delete() + with pytest.raises(ValueError, match="exactly one"): + ctx.delete(id="rec-1", external_id="source-1") + + def test_context_list_default_args(): ctx = Context.__new__(Context) dummy = DummyInner() From d0160e2152c55337f6fc91e287937df9b84a09f2 Mon Sep 17 00:00:00 2001 From: Allen Cheng Date: Tue, 9 Jun 2026 00:04:34 -0700 Subject: [PATCH 2/2] test: cover delete by record id --- crates/lance-context-core/src/store.rs | 32 +++++++++++++++++++++++++ python/tests/test_delete.py | 33 ++++++++++++++++++++++++++ 2 files changed, 65 insertions(+) diff --git a/crates/lance-context-core/src/store.rs b/crates/lance-context-core/src/store.rs index aa58366..3c550ab 100644 --- a/crates/lance-context-core/src/store.rs +++ b/crates/lance-context-core/src/store.rs @@ -1459,6 +1459,38 @@ mod tests { }); } + #[test] + fn delete_by_id_hides_record_from_default_reads() { + let dir = TempDir::new().unwrap(); + let uri = dir.path().to_string_lossy().to_string(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + runtime.block_on(async { + let mut store = ContextStore::open(&uri).await.unwrap(); + let mut first = text_record("a", 0.0); + first.external_id = Some("doc-123#chunk-1".to_string()); + let second = text_record("b", 2.0); + store.add(&[first.clone(), second.clone()]).await.unwrap(); + + assert!(store.delete_by_id(&first.id).await.unwrap()); + + assert!(store.get_by_id(&first.id).await.unwrap().is_none()); + assert!(store + .get_by_external_id("doc-123#chunk-1") + .await + .unwrap() + .is_none()); + + let records = store.list(None, None).await.unwrap(); + assert_eq!(records.len(), 1); + assert_eq!(records[0].id, second.id); + + let query = make_embedding(0.0); + let hits = store.search(&query, Some(10)).await.unwrap(); + assert_eq!(hits.len(), 1); + assert_eq!(hits[0].record.id, second.id); + }); + } + #[test] fn delete_missing_id_is_noop() { let dir = TempDir::new().unwrap(); diff --git a/python/tests/test_delete.py b/python/tests/test_delete.py index 058c975..74f2586 100644 --- a/python/tests/test_delete.py +++ b/python/tests/test_delete.py @@ -49,6 +49,39 @@ def test_delete_by_external_id_hides_record_from_default_reads(tmp_path: Path) - assert [hit["external_id"] for hit in hits] == ["doc-456#chunk-1"] +def test_delete_by_id_hides_record_from_default_reads(tmp_path: Path) -> None: + uri = str(tmp_path / "context.lance") + ctx = Context.create(uri) + + ctx.add( + "user", + "stale memory", + embedding=_embedding(0.0), + external_id="doc-123#chunk-4", + ) + ctx.add( + "user", + "fresh memory", + embedding=_embedding(1.0), + external_id="doc-456#chunk-1", + ) + + stale = ctx.get(external_id="doc-123#chunk-4") + assert stale is not None + stale_id = stale["id"] + + assert ctx.delete(id=stale_id) is True + assert ctx.delete(id=stale_id) is False + assert ctx.get(id=stale_id) is None + assert ctx.get(external_id="doc-123#chunk-4") is None + + entries = ctx.list() + assert [entry["external_id"] for entry in entries] == ["doc-456#chunk-1"] + + hits = ctx.search(_embedding(0.0), limit=10) + assert [hit["external_id"] for hit in hits] == ["doc-456#chunk-1"] + + def test_external_id_can_be_reused_after_delete(tmp_path: Path) -> None: uri = str(tmp_path / "context.lance") ctx = Context.create(uri)