From f46fbb5ba4f7dad08fe79dc13fbbd793b45abab2 Mon Sep 17 00:00:00 2001 From: Allen Cheng Date: Wed, 24 Jun 2026 15:43:03 -0700 Subject: [PATCH 1/3] perf: validate uniqueness via projected indexed scans (#100) Uniqueness validation previously listed and deserialized the entire dataset (all columns, all history including expired/retired/superseded/ tombstoned rows) into a Vec on every add / add_many / upsert / update, then linear-scanned for collisions. That made each write O(N) in store size and incremental/streaming appends O(N^2). Replace the full scan with find_existing_keys(), which issues projected, filtered LsmScanner scans over only the candidate keys: - `id IN (...)` projecting just `id` + `content_type` - `external_id IN (...)` projecting just `external_id` + `content_type` Filters are pushed down to each data source, so the id lookup is index-accelerated when an id scalar index is configured, and neither scan deserializes full records or materializes the whole dataset. Behavior is preserved exactly: within-batch duplicate detection and the existing error messages are unchanged; tombstones remain excluded from collisions (a key whose only surviving row is a tombstone is reusable) while superseded / retired / expired non-tombstone rows still reserve their key. The external_id scan is guarded by a schema check so datasets without that column are unaffected. Caller-supplied external_ids are escaped before going into the IN filter. Tests: id/external_id collisions, within-batch dups, reuse-after-delete, reject-after-supersede, the indexed-scan path, large-store correctness, single-quote escaping, and an ignored append-cost benchmark (store grew ~20x while per-call append cost grew ~4x, not ~20x). Closes #100 --- crates/lance-context-core/src/store.rs | 479 ++++++++++++++++++++++++- 1 file changed, 460 insertions(+), 19 deletions(-) diff --git a/crates/lance-context-core/src/store.rs b/crates/lance-context-core/src/store.rs index 19556fd..9febee0 100644 --- a/crates/lance-context-core/src/store.rs +++ b/crates/lance-context-core/src/store.rs @@ -696,19 +696,22 @@ impl ContextStore { } } - for record in self - .list_with_options(None, None, LifecycleQueryOptions::new(true, true)) - .await? - { - if ids.contains(record.id.as_str()) { + let id_list: Vec<&str> = ids.iter().copied().collect(); + let external_id_list: Vec<&str> = external_ids.iter().copied().collect(); + let (existing_ids, existing_external_ids) = + self.find_existing_keys(&id_list, &external_id_list).await?; + + // Report collisions in input order for deterministic, intuitive errors. + for entry in entries { + if existing_ids.contains(entry.id.as_str()) { return Err(ArrowError::InvalidArgumentError(format!( "id '{}' already exists", - record.id + entry.id )) .into()); } - if let Some(external_id) = record.external_id { - if external_ids.contains(external_id.as_str()) { + if let Some(external_id) = &entry.external_id { + if existing_external_ids.contains(external_id.as_str()) { return Err(ArrowError::InvalidArgumentError(format!( "external_id '{}' already exists", external_id @@ -722,21 +725,99 @@ impl ContextStore { } async fn validate_new_record_id(&self, entry: &ContextRecord) -> LanceResult<()> { - for record in self - .list_with_options(None, None, LifecycleQueryOptions::new(true, true)) - .await? - { - if record.id == entry.id { - return Err(ArrowError::InvalidArgumentError(format!( - "id '{}' already exists", - entry.id - )) - .into()); - } + let id = entry.id.as_str(); + let (existing_ids, _) = self.find_existing_keys(&[id], &[]).await?; + if existing_ids.contains(id) { + return Err(ArrowError::InvalidArgumentError(format!( + "id '{}' already exists", + entry.id + )) + .into()); } Ok(()) } + /// Resolve which of the candidate `id` / `external_id` values already exist + /// among non-tombstone records. + /// + /// Unlike a full `list` + deserialize, this issues projected, filtered + /// scans that read only the `id` / `external_id` / `content_type` columns + /// for rows matching the candidate keys. When an `id` scalar index is + /// configured the `id` lookup is index-accelerated, so uniqueness + /// validation cost is bounded by the candidate batch (and index + /// selectivity) rather than the total number of stored records — history, + /// retired, expired, and superseded rows included. + /// + /// Tombstones are skipped to preserve existing behavior: a key whose only + /// surviving row is a tombstone is free for reuse, while a + /// superseded/retired/expired (non-tombstone) row still reserves its key. + async fn find_existing_keys( + &self, + ids: &[&str], + external_ids: &[&str], + ) -> LanceResult<(HashSet, HashSet)> { + let mut existing_ids = HashSet::new(); + let mut existing_external_ids = HashSet::new(); + + let candidate_ids: HashSet<&str> = ids.iter().copied().collect(); + let candidate_external_ids: HashSet<&str> = external_ids.iter().copied().collect(); + + if !candidate_ids.is_empty() { + let filter = format!("id IN ({})", sql_quoted_list(ids)); + let scanner = self + .lsm_scanner() + .await? + .project(&["id", "content_type"]) + .filter(&filter)?; + let mut stream = scanner.try_into_stream().await?; + while let Some(batch) = stream.try_next().await? { + let id_array = column_as::(&batch, "id")?; + let content_type_array = column_as::(&batch, "content_type")?; + for row in 0..batch.num_rows() { + if content_type_array.value(row) == CONTENT_TYPE_TOMBSTONE { + continue; + } + let id = id_array.value(row); + if candidate_ids.contains(id) { + existing_ids.insert(id.to_string()); + } + } + } + } + + if !candidate_external_ids.is_empty() && self.has_external_id_column() { + let filter = format!("external_id IN ({})", sql_quoted_list(external_ids)); + let scanner = self + .lsm_scanner() + .await? + .project(&["external_id", "content_type"]) + .filter(&filter)?; + let mut stream = scanner.try_into_stream().await?; + while let Some(batch) = stream.try_next().await? { + let content_type_array = column_as::(&batch, "content_type")?; + let Some(external_id_array) = + column_as_optional::(&batch, "external_id") + else { + continue; + }; + for row in 0..batch.num_rows() { + if content_type_array.value(row) == CONTENT_TYPE_TOMBSTONE { + continue; + } + if external_id_array.is_null(row) { + continue; + } + let external_id = external_id_array.value(row); + if candidate_external_ids.contains(external_id) { + existing_external_ids.insert(external_id.to_string()); + } + } + } + } + + Ok((existing_ids, existing_external_ids)) + } + fn derive_region_id(bot_id: &Option, session_id: &Option) -> Uuid { let mut input = String::new(); @@ -760,6 +841,14 @@ impl ContextStore { .any(|path| path == RELATIONSHIPS_COLUMN) } + fn has_external_id_column(&self) -> bool { + self.dataset + .schema() + .field_paths() + .iter() + .any(|path| path == "external_id") + } + /// Current dataset version. pub fn version(&self) -> u64 { self.dataset.manifest.version @@ -2478,6 +2567,17 @@ where .and_then(|col| col.as_ref().as_any().downcast_ref::()) } +/// Render a SQL `IN (...)` value list as comma-separated quoted string +/// literals, escaping any embedded single quotes. Callers ensure the input is +/// non-empty before building the surrounding `IN ()` clause. +fn sql_quoted_list(values: &[&str]) -> String { + values + .iter() + .map(|value| format!("'{}'", value.replace('\'', "''"))) + .collect::>() + .join(",") +} + #[cfg(test)] mod tests { use super::*; @@ -3342,6 +3442,347 @@ mod tests { }); } + #[test] + fn add_rejects_duplicate_id_against_existing() { + 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(); + store.add(&[text_record("dup", 0.0)]).await.unwrap(); + + let err = store.add(&[text_record("dup", 1.0)]).await.unwrap_err(); + let message = err.to_string(); + assert!( + message.contains("id 'dup'") && message.contains("already exists"), + "unexpected error message: {message}" + ); + }); + } + + #[test] + fn add_rejects_duplicate_id_within_batch() { + 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 err = store + .add(&[text_record("same", 0.0), text_record("same", 1.0)]) + .await + .unwrap_err(); + let message = err.to_string(); + assert!( + message.contains("duplicate id 'same' in batch"), + "unexpected error message: {message}" + ); + }); + } + + #[test] + fn add_rejects_duplicate_external_id_within_batch() { + 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("ext".to_string()); + let mut second = text_record("b", 1.0); + second.external_id = Some("ext".to_string()); + + let err = store.add(&[first, second]).await.unwrap_err(); + let message = err.to_string(); + assert!( + message.contains("duplicate external_id 'ext' in batch"), + "unexpected error message: {message}" + ); + }); + } + + /// A record removed via tombstone frees its `external_id` for reuse: + /// validation skips tombstones, so a later insert with the same + /// `external_id` succeeds. + #[test] + fn add_allows_external_id_reuse_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("ext".to_string()); + store.add(std::slice::from_ref(&first)).await.unwrap(); + assert!(store.delete_by_external_id("ext").await.unwrap()); + + let mut reused = text_record("b", 1.0); + reused.external_id = Some("ext".to_string()); + store + .add(std::slice::from_ref(&reused)) + .await + .expect("external_id should be reusable after delete"); + + let visible = store.get_by_external_id("ext").await.unwrap().unwrap(); + assert_eq!(visible.id, reused.id); + }); + } + + /// A record removed via tombstone frees its `id` for reuse. + #[test] + fn add_allows_id_reuse_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 first = text_record("dup", 0.0); + store.add(std::slice::from_ref(&first)).await.unwrap(); + assert!(store.delete_by_id("dup").await.unwrap()); + + store + .add(&[text_record("dup", 1.0)]) + .await + .expect("id should be reusable after delete"); + + let visible = store.get_by_id("dup").await.unwrap().unwrap(); + assert_eq!(visible.id, "dup"); + }); + } + + /// A superseded (non-tombstone) record still reserves its `external_id`, + /// so a plain `add` with that `external_id` is rejected. + #[test] + fn add_rejects_external_id_after_supersede() { + 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("ext".to_string()); + store.upsert_by_external_id(first).await.unwrap(); + + let mut successor = text_record("b", 1.0); + successor.external_id = Some("ext".to_string()); + store.upsert_by_external_id(successor).await.unwrap(); + + // The original is now superseded (non-tombstone) and still present, + // so its external_id is taken. + let mut conflict = text_record("c", 2.0); + conflict.external_id = Some("ext".to_string()); + let err = store.add(&[conflict]).await.unwrap_err(); + let message = err.to_string(); + assert!( + message.contains("external_id 'ext'") && message.contains("already exists"), + "unexpected error message: {message}" + ); + }); + } + + /// Uniqueness validation must behave identically when an `id` scalar index + /// is configured (the indexed-lookup path), covering both a rejected + /// collision and a permitted reuse-after-delete. + #[test] + fn validate_uniqueness_with_btree_index() { + 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 options = ContextStoreOptions { + id_index_type: IdIndexType::BTree, + ..Default::default() + }; + let mut store = ContextStore::open_with_options(&uri, options) + .await + .unwrap(); + + let mut first = text_record("idx-a", 0.0); + first.external_id = Some("ext".to_string()); + store.add(std::slice::from_ref(&first)).await.unwrap(); + + // Duplicate id rejected via the indexed lookup. + let dup_id = store.add(&[text_record("idx-a", 1.0)]).await.unwrap_err(); + assert!( + dup_id.to_string().contains("id 'idx-a'") + && dup_id.to_string().contains("already exists") + ); + + // Duplicate external_id rejected. + let mut dup_ext = text_record("idx-b", 1.0); + dup_ext.external_id = Some("ext".to_string()); + let dup_ext_err = store.add(&[dup_ext]).await.unwrap_err(); + assert!( + dup_ext_err.to_string().contains("external_id 'ext'") + && dup_ext_err.to_string().contains("already exists") + ); + + // After delete, both keys become reusable. + assert!(store.delete_by_id("idx-a").await.unwrap()); + let mut reused = text_record("idx-a", 2.0); + reused.external_id = Some("ext".to_string()); + store + .add(std::slice::from_ref(&reused)) + .await + .expect("keys should be reusable after delete with index configured"); + }); + } + + /// Validation stays correct as the store grows: a duplicate is rejected and + /// a fresh key accepted against a store with many historical records, + /// exercising the projected/filtered scan path over a larger dataset. + #[test] + fn validate_uniqueness_against_large_store() { + 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 options = ContextStoreOptions { + id_index_type: IdIndexType::BTree, + ..Default::default() + }; + let mut store = ContextStore::open_with_options(&uri, options) + .await + .unwrap(); + + for i in 0..300 { + let mut record = text_record(&format!("rec-{i}"), i as f32); + record.external_id = Some(format!("ext-{i}")); + store.add(std::slice::from_ref(&record)).await.unwrap(); + } + + // Existing id and external_id both rejected. + let mut dup = text_record("rec-150", 0.0); + dup.external_id = Some("ext-999".to_string()); + assert!(store + .add(&[dup]) + .await + .unwrap_err() + .to_string() + .contains("id 'rec-150'")); + + let mut dup_ext = text_record("rec-new", 0.0); + dup_ext.external_id = Some("ext-42".to_string()); + assert!(store + .add(&[dup_ext]) + .await + .unwrap_err() + .to_string() + .contains("external_id 'ext-42'")); + + // A fresh record still inserts cleanly. + let mut fresh = text_record("rec-300", 0.0); + fresh.external_id = Some("ext-300".to_string()); + store.add(std::slice::from_ref(&fresh)).await.unwrap(); + assert!(store.get_by_id("rec-300").await.unwrap().is_some()); + }); + } + + /// Benchmark guarding against the O(N) regression: with an `id` index, + /// per-append validation cost should not grow proportionally to store size. + /// Ignored by default (timing-sensitive); run with + /// `cargo test -p lance-context-core -- --ignored append_cost`. + #[test] + #[ignore = "timing-sensitive benchmark; run explicitly with --ignored"] + fn append_cost_does_not_grow_linearly() { + use std::time::Instant; + + 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 options = ContextStoreOptions { + id_index_type: IdIndexType::BTree, + ..Default::default() + }; + let mut store = ContextStore::open_with_options(&uri, options) + .await + .unwrap(); + + // Time a window of single-record appends at a given store size, + // compacting first so the cost reflects validation against the base + // table rather than accumulated MemWAL generations. + async fn time_window(store: &mut ContextStore, tag: &str, window: usize) -> f64 { + store.compact(None).await.unwrap(); + let start = Instant::now(); + for i in 0..window { + let id = format!("{tag}-probe-{i}"); + store.add(&[text_record(&id, i as f32)]).await.unwrap(); + } + start.elapsed().as_secs_f64() / window as f64 + } + + // Grow the store to `count` rows using batched appends (few commits) + // so the benchmark isolates per-call validation cost, not raw write + // throughput. + async fn seed(store: &mut ContextStore, tag: &str, count: usize) { + let chunk = 100; + let mut i = 0; + while i < count { + let batch: Vec = (i..(i + chunk).min(count)) + .map(|j| text_record(&format!("{tag}-seed-{j}"), j as f32)) + .collect(); + store.add(&batch).await.unwrap(); + i += chunk; + } + store.compact(None).await.unwrap(); + } + + let window = 30; + seed(&mut store, "small", 100).await; + let small = time_window(&mut store, "small", window).await; + + seed(&mut store, "big", 2000).await; + let large = time_window(&mut store, "big", window).await; + + let ratio = large / small.max(f64::EPSILON); + eprintln!( + "append per-call: small={small:.6}s large={large:.6}s ratio={ratio:.2} (store grew ~20x)" + ); + assert!( + ratio < 8.0, + "append cost appears to scale with store size (ratio {ratio:.2}); \ + expected roughly constant per-call validation" + ); + }); + } + + /// `external_id` values are caller-supplied and flow into a SQL `IN (...)` + /// filter, so embedded single quotes must be escaped rather than break the + /// scan. Both insertion and duplicate detection must handle them. + #[test] + fn validation_handles_external_id_with_single_quote() { + 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 tricky = "o'brien#chunk-1"; + + let mut first = text_record("a", 0.0); + first.external_id = Some(tricky.to_string()); + store.add(std::slice::from_ref(&first)).await.unwrap(); + + // Re-using the same quoted external_id is still detected as a dup. + let mut dup = text_record("b", 1.0); + dup.external_id = Some(tricky.to_string()); + let err = store.add(&[dup]).await.unwrap_err(); + assert!( + err.to_string().contains("already exists"), + "unexpected error message: {err}" + ); + + // A different quoted external_id inserts cleanly. + let mut other = text_record("c", 2.0); + other.external_id = Some("d'angelo#chunk-2".to_string()); + store.add(std::slice::from_ref(&other)).await.unwrap(); + assert!(store + .get_by_external_id("d'angelo#chunk-2") + .await + .unwrap() + .is_some()); + }); + } + #[test] fn delete_by_external_id_hides_record_from_default_reads() { let dir = TempDir::new().unwrap(); From 668606820030ae90977ebbb622b643ea97d9afc2 Mon Sep 17 00:00:00 2001 From: Allen Cheng Date: Wed, 24 Jun 2026 16:52:35 -0700 Subject: [PATCH 2/3] feat: add bulk upsert by external_id (#102) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds an insert-or-replace-by-`external_id` batch API across every layer. Previously every upsert path was single-record and `add_many` rejected existing `external_id`s, so there was no way to idempotently upsert a batch in one operation. Core: `ContextStore::upsert_many_by_external_id(records)` applies insert-or-replace per `external_id` in one logical operation and returns one `UpsertResult` per input record (inserted vs. replaced, resulting id), all carrying the final version. Semantics (parity with the single-record `upsert_by_external_id`): - every record must carry a non-empty `external_id`; - duplicate `id`s / `external_id`s within the batch are rejected (parity with `add_many`); - validation is all-or-nothing — nothing is written if any record is invalid; - caller-supplied supersession fields are ignored (replacement is managed by `external_id`); - an insert whose `external_id` already exists on a non-tombstone but hidden row is rejected, exactly as a single insert is. Composes with the #100 indexed validation: a batch resolves existing `external_id`s in a single scan and validates `id` uniqueness via the indexed path, rather than full-scanning per record. The whole batch is written in one pass (single version bump where records share a shard). Plumbed through every surface for parity: - api: `ContextStoreApi::upsert_many` + `UpsertRecordsRequest` / `UpsertRecordsResponse` / `UpsertResultDto`; - core dispatch (`api_impl`), unified store, and remote client; - REST: `PUT /contexts/{name}/records/batch`; - Python: `Context.upsert_many(records, key="external_id")` (+ binding). Tests: core (insert / replace / mixed / idempotent / within-batch dup / missing external_id / id collision / empty / indexed path), REST (insert+replace, empty, missing external_id), and Python (`test_upsert_many.py`). README gains a bulk-upsert example. Closes #102 --- README.md | 19 + crates/lance-context-api/src/lib.rs | 27 + crates/lance-context-client/src/lib.rs | 27 + crates/lance-context-core/src/api_impl.rs | 53 +- crates/lance-context-core/src/store.rs | 469 ++++++++++++++++++ crates/lance-context-server/src/routes/mod.rs | 4 + .../src/routes/records.rs | 174 ++++++- crates/lance-context/src/unified.rs | 9 +- python/python/lance_context/api.py | 100 ++++ python/src/lib.rs | 62 +++ python/tests/test_upsert_many.py | 105 ++++ 11 files changed, 1046 insertions(+), 3 deletions(-) create mode 100644 python/tests/test_upsert_many.py diff --git a/README.md b/README.md index d693166..53a1c06 100644 --- a/README.md +++ b/README.md @@ -172,6 +172,25 @@ ctx.add_many([ }, ]) +# Bulk insert-or-replace by external_id (idempotent re-ingestion). New +# external_ids are inserted; existing ones are replaced (the previous row is +# superseded), all in one storage operation. Returns one result per record. +results = ctx.upsert_many([ + { + "role": "source", + "content": "Chunk 1, revised", + "content_type": "text/markdown", + "external_id": "doc-77#chunk-1", + }, + { + "role": "source", + "content": "Brand new chunk", + "content_type": "text/markdown", + "external_id": "doc-77#chunk-2", + }, +]) +print([(r["inserted"], r["replaced_id"]) for r in results]) + # Deferred embeddings: raw-first capture, enrich later. # # Bulk ingestion often needs to persist source chunks immediately and compute diff --git a/crates/lance-context-api/src/lib.rs b/crates/lance-context-api/src/lib.rs index 246c587..e0a0328 100644 --- a/crates/lance-context-api/src/lib.rs +++ b/crates/lance-context-api/src/lib.rs @@ -39,6 +39,11 @@ pub trait ContextStoreApi { request: &UpsertRecordRequest, ) -> impl Future> + Send; + fn upsert_many( + &mut self, + request: &UpsertRecordsRequest, + ) -> impl Future> + Send; + fn update( &mut self, request: &UpdateRecordRequest, @@ -225,6 +230,28 @@ pub struct UpsertRecordResponse { pub record: RecordDto, } +#[derive(Debug, Serialize, Deserialize)] +pub struct UpsertRecordsRequest { + pub records: Vec, + #[serde(default = "default_upsert_key")] + pub key: String, +} + +/// Per-record outcome of a batch upsert, in input order. +#[derive(Debug, Serialize, Deserialize)] +pub struct UpsertResultDto { + pub inserted: bool, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub replaced_id: Option, + pub record: RecordDto, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct UpsertRecordsResponse { + pub version: u64, + pub results: Vec, +} + #[derive(Debug, Clone, Default, Serialize, Deserialize)] pub struct RecordPatchDto { #[serde(default, skip_serializing_if = "Option::is_none")] diff --git a/crates/lance-context-client/src/lib.rs b/crates/lance-context-client/src/lib.rs index 2a75496..02b6a7c 100644 --- a/crates/lance-context-client/src/lib.rs +++ b/crates/lance-context-client/src/lib.rs @@ -71,6 +71,19 @@ impl ContextStoreApi for RemoteContextStore { Ok(resp) } + async fn upsert_many( + &mut self, + request: &UpsertRecordsRequest, + ) -> ContextResult { + let resp = self + .client + .upsert_records(&self.context_name, request) + .await + .map_err(to_ctx_err)?; + self.cached_version = resp.version; + Ok(resp) + } + async fn update( &mut self, request: &UpdateRecordRequest, @@ -339,6 +352,20 @@ impl ContextClient { Self::handle_response(resp).await } + pub async fn upsert_records( + &self, + name: &str, + req: &UpsertRecordsRequest, + ) -> Result { + let resp = self + .http + .put(self.url(&format!("/contexts/{}/records/batch", name))) + .json(req) + .send() + .await?; + Self::handle_response(resp).await + } + pub async fn update_record( &self, name: &str, diff --git a/crates/lance-context-core/src/api_impl.rs b/crates/lance-context-core/src/api_impl.rs index f287484..32a1317 100644 --- a/crates/lance-context-core/src/api_impl.rs +++ b/crates/lance-context-core/src/api_impl.rs @@ -7,7 +7,7 @@ use lance_context_api::{ ContextError, ContextResult, ContextStoreApi, DeleteRecordResponse, RecordDto, RecordPatchDto, RelationshipDto, RetrieveRequest, RetrieveResultDto, SearchRequest, SearchResultDto, StateMetadataDto, UpdateRecordRequest, UpdateRecordResponse, UpsertRecordRequest, - UpsertRecordResponse, + UpsertRecordResponse, UpsertRecordsRequest, UpsertRecordsResponse, UpsertResultDto, }; use crate::record::{ @@ -74,6 +74,57 @@ impl ContextStoreApi for ContextStore { }) } + async fn upsert_many( + &mut self, + request: &UpsertRecordsRequest, + ) -> ContextResult { + if request.key != "external_id" { + return Err(ContextError::InvalidRequest(format!( + "upsert key '{}' is not supported; use 'external_id'", + request.key + ))); + } + if request.records.is_empty() { + return Err(ContextError::InvalidRequest( + "upsert_many requires at least one record".to_string(), + )); + } + for (index, record) in request.records.iter().enumerate() { + if record.external_id.as_deref().is_none_or(str::is_empty) { + return Err(ContextError::InvalidRequest(format!( + "upsert_many requires record.external_id (records[{index}])" + ))); + } + } + + let core_records: Vec = request + .records + .iter() + .map(|r| { + record_from_add_request(r, Uuid::new_v4().to_string(), Uuid::new_v4().to_string()) + }) + .collect(); + + let results = ContextStore::upsert_many_by_external_id(self, core_records) + .await + .map_err(to_ctx_err)?; + let version = results + .last() + .map(|r| r.version) + .unwrap_or_else(|| ContextStore::version(self)); + Ok(UpsertRecordsResponse { + version, + results: results + .into_iter() + .map(|r| UpsertResultDto { + inserted: r.inserted, + replaced_id: r.replaced_id, + record: record_to_dto(r.record), + }) + .collect(), + }) + } + async fn update( &mut self, request: &UpdateRecordRequest, diff --git a/crates/lance-context-core/src/store.rs b/crates/lance-context-core/src/store.rs index 9febee0..9f1609a 100644 --- a/crates/lance-context-core/src/store.rs +++ b/crates/lance-context-core/src/store.rs @@ -266,6 +266,17 @@ fn relationship_struct_builder() -> StructBuilder { ) } +/// Per-`external_id` resolution computed in a single scan for batch upsert. +#[derive(Default)] +struct ExternalIdState { + /// Ids of currently-visible records carrying this external_id (default + /// lifecycle: not tombstoned/expired/retired/superseded). + visible_ids: Vec, + /// Whether any non-tombstone row (visible or hidden) carries it. Mirrors + /// the uniqueness check `add` performs on the single-upsert insert path. + has_non_tombstone: bool, +} + impl ContextStore { /// Open an existing context dataset or create a new one with the project schema. pub async fn open(uri: &str) -> LanceResult { @@ -512,6 +523,215 @@ impl ContextStore { } } + /// Insert-or-replace a batch of records keyed by `external_id`, in one + /// logical operation. + /// + /// For each record: if a currently-visible record with the same + /// `external_id` exists, it is replaced append-only (the successor gets + /// `supersedes_id` set to the existing record id and the original is hidden + /// from default reads); otherwise the record is inserted. All rows are + /// written in a single pass, so records sharing a shard land in a single + /// version bump. + /// + /// Semantics (parity with [`Self::upsert_by_external_id`] and `add_many`): + /// - every record must carry a non-empty `external_id`; + /// - duplicate `id`s or `external_id`s *within* the batch are rejected; + /// - validation is all-or-nothing — if any record is invalid, nothing is + /// written; + /// - caller-supplied supersession fields are ignored (replacement is + /// managed by `external_id`); + /// - an insert whose `external_id` already exists on a non-tombstone but + /// hidden record is rejected, exactly as a single insert would be. + /// + /// Existing-key resolution and `id` uniqueness validation are each done in + /// a single scan for the whole batch (not per record), composing with the + /// indexed `id` validation so a batch does not full-scan per record. + /// + /// Returns one [`UpsertResult`] per input record, in input order, all + /// carrying the final dataset version. + pub async fn upsert_many_by_external_id( + &mut self, + mut records: Vec, + ) -> LanceResult> { + if records.is_empty() { + return Ok(Vec::new()); + } + + // 1. Per-record validation + within-batch duplicate detection. + let mut seen_ids: HashSet<&str> = HashSet::with_capacity(records.len()); + let mut seen_external_ids: HashSet<&str> = HashSet::with_capacity(records.len()); + for record in &records { + let Some(external_id) = record.external_id.as_deref() else { + return Err(ArrowError::InvalidArgumentError( + "upsert_many_by_external_id requires external_id on every record".to_string(), + ) + .into()); + }; + if external_id.is_empty() { + return Err(ArrowError::InvalidArgumentError( + "upsert_many_by_external_id requires a non-empty external_id".to_string(), + ) + .into()); + } + if record.is_tombstone() { + return Err(ArrowError::InvalidArgumentError(format!( + "content_type '{}' is reserved for internal tombstones", + CONTENT_TYPE_TOMBSTONE + )) + .into()); + } + if !seen_ids.insert(record.id.as_str()) { + return Err(ArrowError::InvalidArgumentError(format!( + "duplicate id '{}' in batch", + record.id + )) + .into()); + } + if !seen_external_ids.insert(external_id) { + return Err(ArrowError::InvalidArgumentError(format!( + "duplicate external_id '{}' in batch", + external_id + )) + .into()); + } + } + + // 2. Replacement is managed here; ignore caller-supplied supersession. + for record in &mut records { + record.supersedes_id = None; + record.superseded_by_id = None; + } + + // 3. id uniqueness against the store (indexed). external_id is NOT + // rejected here: an existing external_id means "replace". + let id_list: Vec<&str> = records.iter().map(|r| r.id.as_str()).collect(); + let (existing_ids, _) = self.find_existing_keys(&id_list, &[]).await?; + if let Some(record) = records + .iter() + .find(|r| existing_ids.contains(r.id.as_str())) + { + return Err(ArrowError::InvalidArgumentError(format!( + "id '{}' already exists", + record.id + )) + .into()); + } + + // 4. Resolve every external_id to its visible record (if any) in one scan. + let external_id_list: Vec<&str> = records + .iter() + .map(|r| r.external_id.as_deref().unwrap_or_default()) + .collect(); + let states = self.external_id_states(&external_id_list).await?; + + // 5. Wire supersession + per-record outcomes, mirroring the single path. + let mut outcomes: Vec<(bool, Option)> = Vec::with_capacity(records.len()); + for record in &mut records { + let external_id = record.external_id.as_deref().unwrap_or_default(); + match states.get(external_id) { + Some(state) if state.visible_ids.len() > 1 => { + return Err(ArrowError::InvalidArgumentError(format!( + "external_id '{}' matches multiple visible records", + external_id + )) + .into()); + } + Some(state) if state.visible_ids.len() == 1 => { + let existing_id = state.visible_ids[0].clone(); + record.supersedes_id = Some(existing_id.clone()); + outcomes.push((false, Some(existing_id))); + } + Some(state) if state.has_non_tombstone => { + // No visible record, but a hidden non-tombstone row already + // holds this external_id — an insert would collide, exactly + // as the single-record insert path (via `add`) rejects it. + return Err(ArrowError::InvalidArgumentError(format!( + "external_id '{}' already exists", + external_id + )) + .into()); + } + _ => outcomes.push((true, None)), + } + } + + // 6. Single write for the whole batch. + let version = self.write_entries(&records).await?; + + Ok(records + .into_iter() + .zip(outcomes) + .map(|(record, (inserted, replaced_id))| UpsertResult { + record, + inserted, + replaced_id, + version, + }) + .collect()) + } + + /// Resolve, for each candidate `external_id`, the set of currently-visible + /// record ids and whether any non-tombstone row carries it — in a single + /// projected, filtered scan rather than a full dataset list. + /// + /// A record that supersedes another keeps the same `external_id`, so the + /// supersession relation is resolved correctly within the filtered set for + /// every flow that creates supersession through the public API. + async fn external_id_states( + &self, + external_ids: &[&str], + ) -> LanceResult> { + let mut states: HashMap = HashMap::new(); + let candidates: HashSet<&str> = external_ids + .iter() + .copied() + .filter(|value| !value.is_empty()) + .collect(); + if candidates.is_empty() { + return Ok(states); + } + + let filter_values: Vec<&str> = candidates.iter().copied().collect(); + let filter = format!("external_id IN ({})", sql_quoted_list(&filter_values)); + let scanner = self.lsm_scanner().await?.filter(&filter)?; + let mut stream = scanner.try_into_stream().await?; + let mut rows: Vec = Vec::new(); + while let Some(batch) = stream.try_next().await? { + rows.extend(batch_to_records(&batch)?); + } + + let superseded_ids: HashSet = rows + .iter() + .filter_map(|record| { + let supersedes_id = record.supersedes_id.as_ref()?; + if supersedes_id == &record.id { + None + } else { + Some(supersedes_id.clone()) + } + }) + .collect(); + + let options = LifecycleQueryOptions::default(); + for record in rows { + let Some(external_id) = record.external_id.as_deref() else { + continue; + }; + if !candidates.contains(external_id) { + continue; + } + let entry = states.entry(external_id.to_string()).or_default(); + if !record.is_tombstone() { + entry.has_non_tombstone = true; + } + if options.is_visible(&record) && !superseded_ids.contains(&record.id) { + entry.visible_ids.push(record.id); + } + } + + Ok(states) + } + /// Partially update mutable fields on a visible record by internal id. /// /// The update is append-only: it writes a replacement record that @@ -3168,6 +3388,255 @@ mod tests { }); } + fn upsert_record(id: &str, external_id: &str, pivot: f32) -> ContextRecord { + let mut record = text_record(id, pivot); + record.external_id = Some(external_id.to_string()); + record + } + + #[test] + fn upsert_many_inserts_new_records() { + 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 batch = vec![ + upsert_record("a", "ext-a", 0.0), + upsert_record("b", "ext-b", 1.0), + ]; + let results = store.upsert_many_by_external_id(batch).await.unwrap(); + + assert_eq!(results.len(), 2); + assert!(results.iter().all(|r| r.inserted)); + assert!(results.iter().all(|r| r.replaced_id.is_none())); + assert_eq!(results[0].version, results[1].version); + + let visible = store.list(None, None).await.unwrap(); + assert_eq!(visible.len(), 2); + }); + } + + #[test] + fn upsert_many_replaces_existing_and_is_idempotent() { + 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 first = vec![ + upsert_record("a1", "ext-a", 0.0), + upsert_record("b1", "ext-b", 1.0), + ]; + store.upsert_many_by_external_id(first).await.unwrap(); + + // Re-apply with new ids: both should replace (supersede) the + // originals, and only the successors remain visible. + let second = vec![ + upsert_record("a2", "ext-a", 2.0), + upsert_record("b2", "ext-b", 3.0), + ]; + let results = store.upsert_many_by_external_id(second).await.unwrap(); + + assert!(results.iter().all(|r| !r.inserted)); + assert_eq!(results[0].replaced_id.as_deref(), Some("a1")); + assert_eq!(results[1].replaced_id.as_deref(), Some("b1")); + assert_eq!(results[0].record.supersedes_id.as_deref(), Some("a1")); + + let visible = store.list(None, None).await.unwrap(); + assert_eq!(visible.len(), 2); + let visible_ids: HashSet<&str> = visible.iter().map(|r| r.id.as_str()).collect(); + assert_eq!( + visible_ids, + HashSet::from(["a2", "b2"]), + "only the successors should be visible" + ); + + // Idempotent re-application again still leaves one visible per key. + let third = vec![ + upsert_record("a3", "ext-a", 4.0), + upsert_record("b3", "ext-b", 5.0), + ]; + store.upsert_many_by_external_id(third).await.unwrap(); + assert_eq!(store.list(None, None).await.unwrap().len(), 2); + }); + } + + #[test] + fn upsert_many_handles_mixed_insert_and_replace() { + 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(); + store + .upsert_many_by_external_id(vec![upsert_record("a1", "ext-a", 0.0)]) + .await + .unwrap(); + + let batch = vec![ + upsert_record("a2", "ext-a", 1.0), // replace + upsert_record("c1", "ext-c", 2.0), // insert + ]; + let results = store.upsert_many_by_external_id(batch).await.unwrap(); + + assert_eq!(results.len(), 2); + assert!(!results[0].inserted); + assert_eq!(results[0].replaced_id.as_deref(), Some("a1")); + assert!(results[1].inserted); + assert!(results[1].replaced_id.is_none()); + + let visible_ids: HashSet = store + .list(None, None) + .await + .unwrap() + .into_iter() + .map(|r| r.id) + .collect(); + assert_eq!( + visible_ids, + HashSet::from(["a2".to_string(), "c1".to_string()]) + ); + }); + } + + #[test] + fn upsert_many_rejects_within_batch_duplicate_external_id() { + 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 batch = vec![ + upsert_record("a", "dup", 0.0), + upsert_record("b", "dup", 1.0), + ]; + let err = store.upsert_many_by_external_id(batch).await.unwrap_err(); + assert!( + err.to_string() + .contains("duplicate external_id 'dup' in batch"), + "unexpected error: {err}" + ); + // Nothing was written (all-or-nothing). + assert_eq!(store.list(None, None).await.unwrap().len(), 0); + }); + } + + #[test] + fn upsert_many_rejects_within_batch_duplicate_id() { + 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 batch = vec![ + upsert_record("same", "ext-a", 0.0), + upsert_record("same", "ext-b", 1.0), + ]; + let err = store.upsert_many_by_external_id(batch).await.unwrap_err(); + assert!( + err.to_string().contains("duplicate id 'same' in batch"), + "unexpected error: {err}" + ); + }); + } + + #[test] + fn upsert_many_rejects_missing_external_id() { + 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 no_ext = vec![text_record("a", 0.0)]; + let err = store.upsert_many_by_external_id(no_ext).await.unwrap_err(); + assert!(err.to_string().contains("external_id"), "unexpected: {err}"); + + let mut empty = text_record("b", 0.0); + empty.external_id = Some(String::new()); + let err = store + .upsert_many_by_external_id(vec![empty]) + .await + .unwrap_err(); + assert!( + err.to_string().contains("non-empty external_id"), + "unexpected: {err}" + ); + }); + } + + #[test] + fn upsert_many_rejects_id_collision_with_store() { + 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(); + store.add(&[text_record("taken", 0.0)]).await.unwrap(); + + let batch = vec![upsert_record("taken", "ext-a", 1.0)]; + let err = store.upsert_many_by_external_id(batch).await.unwrap_err(); + assert!( + err.to_string().contains("id 'taken'") + && err.to_string().contains("already exists"), + "unexpected error: {err}" + ); + }); + } + + #[test] + fn upsert_many_empty_batch_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(); + let results = store.upsert_many_by_external_id(Vec::new()).await.unwrap(); + assert!(results.is_empty()); + }); + } + + #[test] + fn upsert_many_matches_single_upsert_with_btree_index() { + 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 options = ContextStoreOptions { + id_index_type: IdIndexType::BTree, + ..Default::default() + }; + let mut store = ContextStore::open_with_options(&uri, options) + .await + .unwrap(); + + // Seed one record via the single-record path. + store + .upsert_by_external_id(upsert_record("a1", "ext-a", 0.0)) + .await + .unwrap(); + + // Batch replaces ext-a and inserts ext-b through the indexed path. + let results = store + .upsert_many_by_external_id(vec![ + upsert_record("a2", "ext-a", 1.0), + upsert_record("b1", "ext-b", 2.0), + ]) + .await + .unwrap(); + assert_eq!(results[0].replaced_id.as_deref(), Some("a1")); + assert!(results[1].inserted); + + assert_eq!( + store.get_by_external_id("ext-a").await.unwrap().unwrap().id, + "a2" + ); + }); + } + #[test] fn update_by_external_id_patches_mutable_fields_and_preserves_payload() { let dir = TempDir::new().unwrap(); diff --git a/crates/lance-context-server/src/routes/mod.rs b/crates/lance-context-server/src/routes/mod.rs index f7182c6..2e33035 100644 --- a/crates/lance-context-server/src/routes/mod.rs +++ b/crates/lance-context-server/src/routes/mod.rs @@ -27,6 +27,10 @@ pub fn router() -> Router> { "/api/v1/contexts/{name}/records", put(records::upsert_record), ) + .route( + "/api/v1/contexts/{name}/records/batch", + put(records::upsert_records), + ) .route( "/api/v1/contexts/{name}/records", patch(records::update_record), diff --git a/crates/lance-context-server/src/routes/records.rs b/crates/lance-context-server/src/routes/records.rs index e1a3ff6..24c52b2 100644 --- a/crates/lance-context-server/src/routes/records.rs +++ b/crates/lance-context-server/src/routes/records.rs @@ -7,6 +7,7 @@ use lance_context_api::{ AddRecordsRequest, AddRecordsResponse, DeleteRecordResponse, GetRecordResponse, ListRecordsResponse, RecordDto, RecordPatchDto, RelationshipDto, StateMetadataDto, UpdateRecordRequest, UpdateRecordResponse, UpsertRecordRequest, UpsertRecordResponse, + UpsertRecordsRequest, UpsertRecordsResponse, UpsertResultDto, }; use lance_context_core::{ ContextRecord, LifecycleQueryOptions, RecordFilters, RecordPatch, Relationship, StateMetadata, @@ -113,6 +114,69 @@ pub async fn upsert_record( )) } +pub async fn upsert_records( + State(state): State>, + Path(name): Path, + Json(req): Json, +) -> Result<(axum::http::StatusCode, Json), AppError> { + if req.key != "external_id" { + return Err(AppError::InvalidRequest(format!( + "upsert key '{}' is not supported; use 'external_id'", + req.key + ))); + } + if req.records.is_empty() { + return Err(AppError::InvalidRequest( + "records array must not be empty".to_string(), + )); + } + for (index, record) in req.records.iter().enumerate() { + if record.external_id.as_deref().is_none_or(str::is_empty) { + return Err(AppError::InvalidRequest(format!( + "upsert requires record.external_id (records[{index}])" + ))); + } + } + + let stores = state.stores.read().await; + let store_lock = stores + .get(&name) + .ok_or_else(|| AppError::NotFound(format!("Context '{}' does not exist", name)))? + .clone(); + drop(stores); + + let core_records: Vec = req + .records + .iter() + .map(|r| record_from_add_request(r, Uuid::new_v4().to_string(), Uuid::new_v4().to_string())) + .collect(); + + let mut store = store_lock.write().await; + let results = store + .upsert_many_by_external_id(core_records) + .await + .map_err(AppError::from_lance)?; + let version = results + .last() + .map(|r| r.version) + .unwrap_or_else(|| store.version()); + + Ok(( + axum::http::StatusCode::OK, + Json(UpsertRecordsResponse { + version, + results: results + .into_iter() + .map(|r| UpsertResultDto { + inserted: r.inserted, + replaced_id: r.replaced_id, + record: record_to_dto(r.record), + }) + .collect(), + }), + )) +} + pub async fn update_record( State(state): State>, Path(name): Path, @@ -480,7 +544,7 @@ mod tests { use chrono::{Duration, Utc}; use lance_context_api::{ AddRecordRequest, AddRecordsRequest, RecordPatchDto, UpdateRecordRequest, - UpsertRecordRequest, + UpsertRecordRequest, UpsertRecordsRequest, }; use lance_context_core::ContextStore; use tempfile::TempDir; @@ -665,6 +729,114 @@ mod tests { ); } + #[tokio::test] + async fn upsert_records_batch_inserts_and_replaces() { + let context_name = "ctx"; + let (state, _dir) = test_state(context_name).await; + + let mut a = text_record("a-old"); + a.external_id = Some("ext-a".to_string()); + let mut b = text_record("b-value"); + b.external_id = Some("ext-b".to_string()); + + // First batch: two inserts. + let (status, Json(first)) = upsert_records( + State(state.clone()), + Path(context_name.to_string()), + Json(UpsertRecordsRequest { + records: vec![a, b], + key: "external_id".to_string(), + }), + ) + .await + .unwrap(); + assert_eq!(status, axum::http::StatusCode::OK); + assert_eq!(first.results.len(), 2); + assert!(first.results.iter().all(|r| r.inserted)); + let a_id = first.results[0].record.id.clone(); + + // Second batch: replace ext-a, insert ext-c. + let mut a_new = text_record("a-new"); + a_new.external_id = Some("ext-a".to_string()); + let mut c = text_record("c-value"); + c.external_id = Some("ext-c".to_string()); + let (_, Json(second)) = upsert_records( + State(state.clone()), + Path(context_name.to_string()), + Json(UpsertRecordsRequest { + records: vec![a_new, c], + key: "external_id".to_string(), + }), + ) + .await + .unwrap(); + assert!(!second.results[0].inserted); + assert_eq!( + second.results[0].replaced_id.as_deref(), + Some(a_id.as_str()) + ); + assert!(second.results[1].inserted); + + // ext-a now resolves to the replacement; three external_ids visible. + let Json(after) = get_record_by_external_id( + State(state.clone()), + Path(context_name.to_string()), + Query(ExternalIdParams { + external_id: "ext-a".to_string(), + }), + ) + .await + .unwrap(); + assert_eq!(after.record.unwrap().text_payload.as_deref(), Some("a-new")); + + let Json(listed) = list_records( + State(state), + Path(context_name.to_string()), + Query(ListParams { + limit: None, + offset: None, + ..Default::default() + }), + ) + .await + .unwrap(); + assert_eq!(listed.records.len(), 3); + } + + #[tokio::test] + async fn upsert_records_batch_rejects_empty() { + let context_name = "ctx"; + let (state, _dir) = test_state(context_name).await; + let err = upsert_records( + State(state), + Path(context_name.to_string()), + Json(UpsertRecordsRequest { + records: vec![], + key: "external_id".to_string(), + }), + ) + .await + .unwrap_err(); + assert!(matches!(err, AppError::InvalidRequest(_))); + } + + #[tokio::test] + async fn upsert_records_batch_requires_external_id() { + let context_name = "ctx"; + let (state, _dir) = test_state(context_name).await; + let err = upsert_records( + State(state), + Path(context_name.to_string()), + Json(UpsertRecordsRequest { + records: vec![text_record("no external id")], + key: "external_id".to_string(), + }), + ) + .await + .unwrap_err(); + assert!(matches!(err, AppError::InvalidRequest(_))); + } + #[tokio::test] async fn update_by_external_id_patches_visible_record() { let context_name = "ctx"; diff --git a/crates/lance-context/src/unified.rs b/crates/lance-context/src/unified.rs index 43af986..7847c04 100644 --- a/crates/lance-context/src/unified.rs +++ b/crates/lance-context/src/unified.rs @@ -4,7 +4,7 @@ use lance_context_api::{ AddRecordRequest, AddRecordsResponse, CompactRequest, CompactResponse, CompactStatsResponse, ContextError, ContextResult, ContextStoreApi, DeleteRecordResponse, RecordDto, RetrieveRequest, RetrieveResultDto, SearchRequest, SearchResultDto, UpdateRecordRequest, UpdateRecordResponse, - UpsertRecordRequest, UpsertRecordResponse, + UpsertRecordRequest, UpsertRecordResponse, UpsertRecordsRequest, UpsertRecordsResponse, }; use lance_context_core::{ ContextStore as LocalStore, ContextStoreOptions, DistanceMetric, IdIndexType, @@ -129,6 +129,13 @@ impl ContextStoreApi for ContextStore { dispatch_mut!(self, upsert, request) } + async fn upsert_many( + &mut self, + request: &UpsertRecordsRequest, + ) -> ContextResult { + dispatch_mut!(self, upsert_many, request) + } + async fn update( &mut self, request: &UpdateRecordRequest, diff --git a/python/python/lance_context/api.py b/python/python/lance_context/api.py index 8bfd43d..3ee495e 100644 --- a/python/python/lance_context/api.py +++ b/python/python/lance_context/api.py @@ -863,6 +863,106 @@ def add_many(self, records: Iterable[Mapping[str, Any]]) -> None: self._auto_embed_batch(normalized) self._inner.add_many(normalized) + def upsert_many( + self, + records: Iterable[Mapping[str, Any]], + *, + key: str = "external_id", + ) -> list[dict[str, Any]]: + """Insert-or-replace multiple records by ``external_id`` in one call. + + Each record accepts the same fields as :meth:`add` and must include a + non-empty ``external_id``. For each record, the currently-visible record + with the same ``external_id`` is replaced (superseded); otherwise the + record is inserted. Duplicate ``id``/``external_id`` values within the + batch are rejected, and the whole batch is validated before anything is + written. + + Returns one result mapping per input record, in order, each with + ``inserted``, ``replaced_id``, ``version`` and the resulting ``record``. + """ + if key != "external_id": + raise ValueError("Only key='external_id' is currently supported") + + normalized: list[dict[str, Any]] = [] + for index, record in enumerate(records): + if not isinstance(record, Mapping): + raise TypeError(f"records[{index}] must be a mapping") + if "role" not in record: + raise ValueError(f"records[{index}] is missing required key 'role'") + if "content" not in record: + raise ValueError(f"records[{index}] is missing required key 'content'") + if not record.get("external_id"): + raise ValueError(f"records[{index}] is missing required key 'external_id'") + + content_type = record.get("content_type") + data_type = record.get("data_type") + if content_type is not None and data_type is not None: + raise ValueError( + f"records[{index}] specifies both content_type and data_type" + ) + if content_type is None: + content_type = data_type + + payload, resolved_type = _normalize_content(record["content"], content_type) + normalized.append( + { + "role": record["role"], + "content": payload, + "data_type": resolved_type, + "embedding": record.get("embedding"), + "bot_id": record.get( + "bot_id", getattr(self, "_default_fields", {}).get("bot_id") + ), + "session_id": record.get( + "session_id", + getattr(self, "_default_fields", {}).get("session_id"), + ), + "tenant": record.get( + "tenant", getattr(self, "_default_fields", {}).get("tenant") + ), + "source": record.get( + "source", getattr(self, "_default_fields", {}).get("source") + ), + "external_id": record.get("external_id"), + "run_id": record.get("run_id"), + "created_at": _coerce_timestamp( + record.get("created_at"), + field_name=f"records[{index}].created_at", + ), + "state_metadata": _normalize_state_metadata( + record.get("state_metadata") + ), + "metadata_json": _json_dumps(record.get("metadata"), "metadata"), + "relationships_json": _json_dumps( + record.get("relationships"), "relationships" + ), + "expires_at": _coerce_timestamp( + record.get("expires_at"), + field_name=f"records[{index}].expires_at", + ), + "retention_policy": record.get("retention_policy"), + "lifecycle_status": record.get("lifecycle_status"), + "retired_at": _coerce_timestamp( + record.get("retired_at"), + field_name=f"records[{index}].retired_at", + ), + "retired_reason": record.get("retired_reason"), + } + ) + + self._auto_embed_batch(normalized) + results = self._inner.upsert_many(normalized, key) + return [ + { + "inserted": bool(result["inserted"]), + "replaced_id": result.get("replaced_id"), + "version": result["version"], + "record": _normalize_record(result["record"]), + } + for result in results + ] + def ingest_records( self, records: Iterable[Mapping[str, Any]], diff --git a/python/src/lib.rs b/python/src/lib.rs index c7acaae..4d48dc3 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -552,6 +552,68 @@ impl Context { Ok(()) } + #[pyo3(signature = (records, key = "external_id"))] + fn upsert_many( + &mut self, + py: Python<'_>, + records: &Bound<'_, PyAny>, + key: &str, + ) -> PyResult> { + if key != "external_id" { + return Err(PyRuntimeError::new_err(format!( + "upsert key '{key}' is not supported; use 'external_id'" + ))); + } + + let mut prepared = Vec::new(); + for (index, item) in records.try_iter()?.enumerate() { + let item = item?; + let dict = item + .downcast::() + .map_err(|_| PyTypeError::new_err(format!("records[{index}] must be a dict")))?; + let record = self.prepare_record_from_dict(dict, index)?; + if record + .record + .external_id + .as_deref() + .is_none_or(str::is_empty) + { + return Err(PyRuntimeError::new_err(format!( + "upsert_many requires external_id (records[{index}])" + ))); + } + prepared.push(record); + } + + if prepared.is_empty() { + return Ok(Vec::new()); + } + + let context_records: Vec = + prepared.iter().map(|item| item.record.clone()).collect(); + let results = py.allow_threads(|| { + self.runtime + .block_on(self.store.upsert_many_by_external_id(context_records)) + }); + let results = results.map_err(to_py_err)?; + + for item in &prepared { + self.inner + .add(&item.role, &item.inner_content, item.data_type.as_deref()); + } + + let mut out = Vec::with_capacity(results.len()); + for result in results { + let dict = PyDict::new(py); + dict.set_item("inserted", result.inserted)?; + dict.set_item("replaced_id", result.replaced_id)?; + dict.set_item("version", result.version)?; + dict.set_item("record", record_to_py(py, result.record)?)?; + out.push(dict.into_pyobject(py)?.unbind().into()); + } + Ok(out) + } + #[pyo3(signature = (label = None))] fn snapshot(&mut self, label: Option<&str>) -> String { self.inner.snapshot(label) diff --git a/python/tests/test_upsert_many.py b/python/tests/test_upsert_many.py new file mode 100644 index 0000000..8cc2c2a --- /dev/null +++ b/python/tests/test_upsert_many.py @@ -0,0 +1,105 @@ +"""Tests for batch insert-or-replace (upsert_many).""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pytest + +if TYPE_CHECKING: + from pathlib import Path + +from lance_context.api import Context + + +def test_upsert_many_inserts_new_records(tmp_path: Path) -> None: + uri = str(tmp_path / "context.lance") + ctx = Context.create(uri) + + results = ctx.upsert_many( + [ + {"role": "user", "content": "a", "external_id": "ext-a"}, + {"role": "user", "content": "b", "external_id": "ext-b"}, + ] + ) + + assert len(results) == 2 + assert all(r["inserted"] for r in results) + assert all(r["replaced_id"] is None for r in results) + assert len(ctx.list()) == 2 + + +def test_upsert_many_replaces_and_is_idempotent(tmp_path: Path) -> None: + uri = str(tmp_path / "context.lance") + ctx = Context.create(uri) + + ctx.upsert_many( + [ + {"role": "user", "content": "a-old", "external_id": "ext-a"}, + {"role": "user", "content": "b-old", "external_id": "ext-b"}, + ] + ) + results = ctx.upsert_many( + [ + {"role": "user", "content": "a-new", "external_id": "ext-a"}, + {"role": "user", "content": "b-new", "external_id": "ext-b"}, + ] + ) + + assert all(not r["inserted"] for r in results) + assert all(r["replaced_id"] is not None for r in results) + + # Only the successors remain visible, one per external_id. + visible = ctx.list() + assert len(visible) == 2 + assert ctx.get(external_id="ext-a")["text"] == "a-new" + assert ctx.get(external_id="ext-b")["text"] == "b-new" + + +def test_upsert_many_mixed_insert_and_replace(tmp_path: Path) -> None: + uri = str(tmp_path / "context.lance") + ctx = Context.create(uri) + + ctx.upsert_many([{"role": "user", "content": "a", "external_id": "ext-a"}]) + results = ctx.upsert_many( + [ + {"role": "user", "content": "a2", "external_id": "ext-a"}, + {"role": "user", "content": "c", "external_id": "ext-c"}, + ] + ) + + assert not results[0]["inserted"] + assert results[0]["replaced_id"] is not None + assert results[1]["inserted"] + assert len(ctx.list()) == 2 + + +def test_upsert_many_rejects_within_batch_duplicate_external_id(tmp_path: Path) -> None: + uri = str(tmp_path / "context.lance") + ctx = Context.create(uri) + + with pytest.raises(RuntimeError, match="duplicate external_id"): + ctx.upsert_many( + [ + {"role": "user", "content": "a", "external_id": "dup"}, + {"role": "user", "content": "b", "external_id": "dup"}, + ] + ) + + # All-or-nothing: nothing was written. + assert ctx.list() == [] + + +def test_upsert_many_requires_external_id(tmp_path: Path) -> None: + uri = str(tmp_path / "context.lance") + ctx = Context.create(uri) + + with pytest.raises(ValueError, match="external_id"): + ctx.upsert_many([{"role": "user", "content": "a"}]) + + +def test_upsert_many_empty_batch_returns_empty(tmp_path: Path) -> None: + uri = str(tmp_path / "context.lance") + ctx = Context.create(uri) + + assert ctx.upsert_many([]) == [] From f8935480795d1c66fb7beb5c8b2a3137218d73a3 Mon Sep 17 00:00:00 2001 From: Allen Cheng Date: Wed, 24 Jun 2026 16:56:26 -0700 Subject: [PATCH 3/3] style: wrap long line in upsert_many to satisfy ruff (#102) --- python/python/lance_context/api.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/python/lance_context/api.py b/python/python/lance_context/api.py index 3ee495e..b771634 100644 --- a/python/python/lance_context/api.py +++ b/python/python/lance_context/api.py @@ -893,7 +893,9 @@ def upsert_many( if "content" not in record: raise ValueError(f"records[{index}] is missing required key 'content'") if not record.get("external_id"): - raise ValueError(f"records[{index}] is missing required key 'external_id'") + raise ValueError( + f"records[{index}] is missing required key 'external_id'" + ) content_type = record.get("content_type") data_type = record.get("data_type")