From 3e7fd2661eef0f91ef534e4a59f67ad5801fe82b Mon Sep 17 00:00:00 2001 From: Beinan Date: Mon, 27 Jul 2026 22:37:47 +0000 Subject: [PATCH 1/2] refactor(core): migrate ContextStore onto StorageBase MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Second step of #214, after #215. `ContextStore` now delegates its storage layer to `StorageBase` instead of carrying its own forked implementation. Closes the gaps that made this store diverge from rollout: - **WAL merge, which this store never had.** Flushed generations accumulated under `_mem_wal/` forever and every read unioned all of them, so read cost grew without bound in the number of writes. It now gets the same merge as rollout, exposed as `maybe_merge_wal` (count trigger), `cleanup_wal` (time trigger) and `pending_wal_generations`. - **Compaction now defers index remap.** The base table carries a fieldless MemWAL index and Lance's inline remap panics on it. - **Compaction and index creation no longer drop `storage_options`.** Both reloaded via a bare `Dataset::open`, which silently discarded credentials and broke on any authenticated object store. Both now reload through `StorageBase::load_with_options`. - **Resident writer.** `add` used to open, `put` and `close` a `ShardWriter` on every call; it now reuses one resident writer with fence-retry, so an append no longer pays a cold DNS resolution + TCP/TLS handshake + epoch claim. - **`add` takes `&self`.** The writer lives behind the base's mutex, so the server takes a read lock and concurrent appends stop serializing on the store's `RwLock`. Sharding changes to match rollout: writes go to the shard owned by this writer instance (`ContextStoreOptions::shard_id`) rather than one derived per `(bot_id, session_id)`. One instance owns exactly one shard, so the epoch-fencing invariant holds by construction, and merge has a single well-defined shard to drain. `derive_region_id` is removed. `add` visibility is preserved, not silently weakened. `StorageBase` gains a `seal_on_put` switch: rollout keeps its deferred seal (durable on return, visible via the flush sweeper or `?flush=true`), while `ContextStore` defaults to `seal_on_add: true` — required for correctness, since its id/external-id uniqueness validation, upsert and tombstones all read the table back. Callers appending trusted batches can opt into rollout's throughput profile with `seal_on_add: false`. `ContextStore` is no longer `Clone`: it owns a `StorageBase`, which owns the resident writer and a `Drop` that seals it, so it must have exactly one owner. The two clone sites now open their own handle instead — the background compactor (which only rewrites base-table fragments and would otherwise contend with the write path) and Python's `Context.fork` (which branches the in-memory context over the same dataset). Three new tests pin the behavior: `add` is visible on return by default, deferred seal is durable-but-invisible until `flush`, and WAL generations merge into the base table and stay readable. Full workspace suite passes (180 core lib tests, up from 177), clippy clean. Co-Authored-By: Claude --- crates/lance-context-core/src/api_impl.rs | 5 +- crates/lance-context-core/src/eval.rs | 10 +- crates/lance-context-core/src/namespace.rs | 2 +- .../lance-context-core/src/rollout_store.rs | 5 + crates/lance-context-core/src/store.rs | 630 +++++++++++------- crates/lance-context-core/src/store_base.rs | 127 +++- .../src/routes/records.rs | 5 +- python/src/lib.rs | 17 +- 8 files changed, 524 insertions(+), 277 deletions(-) diff --git a/crates/lance-context-core/src/api_impl.rs b/crates/lance-context-core/src/api_impl.rs index 4d945a6..a33e2de 100644 --- a/crates/lance-context-core/src/api_impl.rs +++ b/crates/lance-context-core/src/api_impl.rs @@ -41,7 +41,10 @@ impl ContextStoreApi for ContextStore { } let count = core_records.len(); - let version = self.add(&core_records).await.map_err(to_ctx_err)?; + // Disambiguate: the inherent `ContextStore::add`, not this trait method. + let version = ContextStore::add(self, &core_records) + .await + .map_err(to_ctx_err)?; Ok(AddRecordsResponse { version, ids, diff --git a/crates/lance-context-core/src/eval.rs b/crates/lance-context-core/src/eval.rs index 7b51c92..e74bd91 100644 --- a/crates/lance-context-core/src/eval.rs +++ b/crates/lance-context-core/src/eval.rs @@ -529,7 +529,7 @@ mod tests { 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 store = ContextStore::open(&uri).await.unwrap(); let a = embedding(&store, &[1.0]); let b = embedding(&store, &[0.5]); let c = embedding(&store, &[0.0, 1.0]); @@ -579,7 +579,7 @@ mod tests { 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 store = ContextStore::open(&uri).await.unwrap(); let q = embedding(&store, &[1.0]); // the only relevant doc is retired -> hidden by default. let mut retired = record("doc-a", "alpha", q.clone()); @@ -624,7 +624,7 @@ mod tests { 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 store = ContextStore::open(&uri).await.unwrap(); let shared = embedding(&store, &[1.0]); let mut a = record("doc-a", "alpha", shared.clone()); a.tenant = Some("x".to_string()); @@ -662,7 +662,7 @@ mod tests { 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 store = ContextStore::open(&uri).await.unwrap(); let a = embedding(&store, &[1.0]); let b = embedding(&store, &[0.0, 1.0]); store @@ -701,7 +701,7 @@ mod tests { 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 store = ContextStore::open(&uri).await.unwrap(); let a = embedding(&store, &[1.0]); let b = embedding(&store, &[0.5]); store diff --git a/crates/lance-context-core/src/namespace.rs b/crates/lance-context-core/src/namespace.rs index c03ca12..7b501af 100644 --- a/crates/lance-context-core/src/namespace.rs +++ b/crates/lance-context-core/src/namespace.rs @@ -493,7 +493,7 @@ mod tests { let selector = selector(&[("tenant", "acme"), ("source", "memory")]); let partition = namespace.resolve_partition(&selector).unwrap(); - let mut store = namespace.context(&selector).await.unwrap(); + let store = namespace.context(&selector).await.unwrap(); store .add(&[record("rec-1", "acme", "memory")]) .await diff --git a/crates/lance-context-core/src/rollout_store.rs b/crates/lance-context-core/src/rollout_store.rs index e90b072..f5b0954 100644 --- a/crates/lance-context-core/src/rollout_store.rs +++ b/crates/lance-context-core/src/rollout_store.rs @@ -424,6 +424,11 @@ impl RolloutStore { // claim-check columns), so a WAL merge first evolves an older // base table to the current schema before appending. latest_schema: Some(Arc::new(rollout_schema())), + // Rollout defers the seal: high fan-in appends must not + // serialize behind a per-append seal, and rollout rows are + // immutable so nothing reads back before writing. The server's + // flush sweeper (and `?flush=true`) provide visibility. + seal_on_put: false, }, create_if_missing, ) diff --git a/crates/lance-context-core/src/store.rs b/crates/lance-context-core/src/store.rs index 5344ade..176e2d0 100644 --- a/crates/lance-context-core/src/store.rs +++ b/crates/lance-context-core/src/store.rs @@ -1,5 +1,7 @@ use std::cmp::Ordering; use std::collections::{HashMap, HashSet}; +use std::future::Future; +use std::pin::Pin; use std::sync::Arc; use std::time::Duration; @@ -11,22 +13,19 @@ use arrow_array::builder::{ use arrow_array::types::Int8Type; use arrow_array::{ Array, ArrayRef, DictionaryArray, FixedSizeListArray, Float32Array, Int32Array, Int64Array, - LargeBinaryArray, LargeStringArray, ListArray, RecordBatch, RecordBatchIterator, StringArray, - StructArray, TimestampMicrosecondArray, + LargeBinaryArray, LargeStringArray, ListArray, RecordBatch, StringArray, StructArray, + TimestampMicrosecondArray, }; use arrow_schema::{ArrowError, DataType, Field, FieldRef, Schema, TimeUnit}; use chrono::{DateTime, Timelike, Utc}; use futures::TryStreamExt; -use lance::dataset::mem_wal::{ - DatasetMemWalExt, LsmScanner, ShardManifestStore, ShardSnapshot, ShardWriterConfig, -}; -use lance::dataset::optimize::{compact_files, CompactionMetrics, CompactionOptions}; +use lance::dataset::mem_wal::LsmScanner; +use lance::dataset::optimize::CompactionMetrics; +use lance::dataset::Dataset; use lance::dataset::NewColumnTransform; -use lance::dataset::{builder::DatasetBuilder, Dataset, WriteMode, WriteParams}; use lance::index::DatasetIndexExt; use lance::io::{ObjectStore, ObjectStoreParams, ObjectStoreRegistry, StorageOptionsAccessor}; use lance::{Error as LanceError, Result as LanceResult}; -use lance_index::mem_wal::MEM_WAL_INDEX_NAME; use lance_index::scalar::ScalarIndexParams; use lance_index::IndexType; use tokio::sync::Mutex; @@ -39,11 +38,11 @@ use crate::record::{ SearchResult, StateMetadata, UpdateResult, UpsertResult, LIFECYCLE_ACTIVE, }; use crate::serde::CONTENT_TYPE_TOMBSTONE; +use crate::store_base::{StorageBase, StorageBaseOptions}; /// Embedding length used for the semantic index column. const DEFAULT_EMBEDDING_DIM: i32 = 1536; const DEFAULT_SEARCH_LIMIT: usize = 10; -const DEFAULT_MANIFEST_SCAN_BATCH_SIZE: usize = 16; const RRF_K: f32 = 60.0; const ID_INDEX_NAME: &str = "id_idx"; pub(crate) const RELATIONSHIPS_COLUMN: &str = "relationships"; @@ -186,24 +185,49 @@ struct CompactionState { /// Valid column names that may use blob encoding. const VALID_BLOB_COLUMNS: &[&str] = &["text_payload", "binary_payload"]; +/// Open a store handle for the background compactor, with the recursion broken. +/// +/// `ContextStore::open_inner` may spawn the background compaction task, and that +/// task needs to reopen the store — a mutual recursion between two `async fn`s. +/// Rust cannot infer `Send` for either of the resulting opaque future types, so +/// the task fails to compile even though the recursion is unreachable at runtime +/// (the reopen passes `spawn_background_compaction: false`). +/// +/// Boxing to `dyn Future` here erases the type and cuts the cycle. It lives +/// outside the `impl` so the erasure is what the spawned task names, rather than +/// the opaque type it is trying to break out of. +fn open_for_compaction( + uri: &str, + options: ContextStoreOptions, +) -> Pin> + Send + '_>> { + Box::pin(ContextStore::open_inner(uri, options, false, false)) +} + /// Persistent Lance-backed context store. -#[derive(Clone)] pub struct ContextStore { - dataset: Dataset, + /// Shared storage layer: dataset handle, resident MemWAL writer, durable + /// append, flush, WAL merge, compaction, indexing and the LSM scanner. + /// + /// Behind an `Arc` because `ContextStore` is `Clone` (the background + /// compactor and the Python `fork` both rely on it) while `StorageBase` + /// owns a resident `ShardWriter` and a `Drop` impl, so it must have exactly + /// one owner. Sharing the base is also what makes clones agree about the + /// resident writer instead of each opening their own. + base: StorageBase, compaction_state: Arc>, pub compaction_config: CompactionConfig, blob_columns: HashSet, id_index_type: IdIndexType, embedding_dim: i32, distance_metric: DistanceMetric, - /// Object-store configuration used to resolve external payload references - /// (see [`ContextStore::fetch_payload`]). Reuses the same options the - /// dataset was opened with so referenced media can live in the same bucket. - storage_options: Option>, } /// Additional configuration when opening a [`ContextStore`]. -#[derive(Debug, Clone, Default)] +/// +/// [`Default`] is hand-written rather than derived because `seal_on_add` +/// must default to `true` (preserving read-your-write), which `derive(Default)` +/// would silently turn into `false`. +#[derive(Debug, Clone)] pub struct ContextStoreOptions { pub storage_options: Option>, pub compaction: CompactionConfig, @@ -222,6 +246,53 @@ pub struct ContextStoreOptions { /// is used; passing a different metric here is an error. `None` defaults to /// the persisted metric (or `L2` for datasets created before persistence). pub distance_metric: Option, + /// Stable identity of this writer instance, mapped to a MemWAL shard. + /// + /// Each instance owns exactly one shard, so no two instances ever contend + /// and the epoch-fencing invariant holds by construction. `None` falls back + /// to a single `"default"` shard, which is correct for single-instance use + /// but must be set per-instance when running multiple writers. + pub shard_id: Option, + /// Fold this instance's flushed MemWAL generations into the base table once + /// it has accumulated this many. `None`/`0` disables the count trigger. + /// + /// Without a merge trigger, generations accumulate under `_mem_wal/` forever + /// and every read unions all of them — the read amplification that + /// previously had no bound at all on this store. + pub merge_after_generations: Option, + /// Whether [`ContextStore::add`] seals the memtable before returning, so the + /// rows it wrote are immediately readable. + /// + /// Defaults to `true`, which is this store's historical behavior and is + /// **required for correctness on the paths that read the table back**: + /// id/external-id uniqueness validation, `upsert_by_external_id`, and + /// tombstone-based deletes. + /// + /// Set `false` to get [`crate::RolloutStore`]'s throughput profile: `add` + /// becomes a durable WAL append only, concurrent appends stop serializing + /// behind a per-append seal, and one flushed generation per call is no + /// longer emitted. In exchange there is **no read-your-write guarantee** — + /// visibility then depends on [`ContextStore::flush`] being driven + /// periodically. Only appropriate for trusted bulk appends that do not rely + /// on the read-back paths above. + pub seal_on_add: bool, +} + +impl Default for ContextStoreOptions { + fn default() -> Self { + Self { + storage_options: None, + compaction: CompactionConfig::default(), + embedding_dim: None, + blob_columns: HashSet::new(), + id_index_type: IdIndexType::default(), + distance_metric: None, + shard_id: None, + merge_after_generations: None, + // Read-your-write by default; see the field docs. + seal_on_add: true, + } + } } impl ContextStoreOptions { @@ -339,7 +410,7 @@ impl ContextStore { /// Open a dataset with explicit object store configuration (e.g. S3 credentials). /// Creates the dataset if it does not exist. pub async fn open_with_options(uri: &str, options: ContextStoreOptions) -> LanceResult { - Self::open_inner(uri, options, true).await + Self::open_inner(uri, options, true, true).await } /// Open an **existing** context dataset. Unlike [`Self::open_with_options`], @@ -354,13 +425,14 @@ impl ContextStore { uri: &str, options: ContextStoreOptions, ) -> LanceResult { - Self::open_inner(uri, options, false).await + Self::open_inner(uri, options, false, true).await } async fn open_inner( uri: &str, options: ContextStoreOptions, create_if_missing: bool, + spawn_background_compaction: bool, ) -> LanceResult { // Validate blob_columns for col in &options.blob_columns { @@ -417,8 +489,31 @@ impl ContextStore { } } - let mut store = Self { + let base = StorageBase::from_dataset( dataset, + StorageBaseOptions { + storage_options, + shard_id: options.shard_id.clone(), + merge_after_generations: options.merge_after_generations, + session: None, + schema: Arc::new(arrow_schema.clone()), + key_column: "id".to_string(), + // The context schema is parameterized per dataset (blob columns, + // embedding width, optional feature columns), so there is no + // single "latest" schema to evolve an older table to. Schema + // migration stays explicit, via `migrate_relationships_column`. + latest_schema: None, + // Context writes read the table back — id/external-id uniqueness + // validation, upsert, and tombstones all need read-your-write — + // so seal by default. Callers appending trusted batches can opt + // into rollout's deferred-seal throughput instead. + seal_on_put: options.seal_on_add, + }, + ) + .await?; + + let mut store = Self { + base, compaction_state: Arc::new(Mutex::new(CompactionState { background_task: None, is_compacting: false, @@ -431,14 +526,20 @@ impl ContextStore { id_index_type: options.id_index_type, embedding_dim, distance_metric, - storage_options, }; // Ensure id index if configured store.ensure_id_index().await?; - // Start background compaction if enabled - store.start_background_compaction().await?; + // Start background compaction if enabled. + // + // Skipped when this open *is* the background compactor reopening the + // store: that path would otherwise re-enter here and spawn another + // compactor, making the recursion type-level (the future can never be + // proven `Send`) rather than merely infinite. + if spawn_background_compaction { + store.start_background_compaction().await?; + } Ok(store) } @@ -452,7 +553,7 @@ impl ContextStore { /// URI of the underlying Lance dataset. #[must_use] pub fn uri(&self) -> &str { - self.dataset.uri() + self.base.dataset.uri() } /// Distance metric this context ranks vector-search results with. @@ -462,65 +563,28 @@ impl ContextStore { } /// Append context records to the store and return the new dataset version. - pub async fn add(&mut self, entries: &[ContextRecord]) -> LanceResult { + pub async fn add(&self, entries: &[ContextRecord]) -> LanceResult { if entries.is_empty() { - return Ok(self.dataset.manifest.version); + return Ok(self.base.version()); } self.validate_unique_ids(entries).await?; self.write_entries(entries).await } - async fn write_entries(&mut self, entries: &[ContextRecord]) -> LanceResult { + /// Encode and append `entries` through the shared storage layer. + /// + /// Takes `&self`: the resident writer lives behind the base's mutex, so + /// appends no longer need exclusive access. Whether the rows are readable + /// when this returns is governed by + /// [`ContextStoreOptions::seal_on_add`] (default `true`). + async fn write_entries(&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> = - HashMap::new(); - for entry in entries { - let key = (entry.bot_id.clone(), entry.session_id.clone()); - groups.entry(key).or_default().push(entry.clone()); - } - - // Ensure MemWAL is initialized (once for the dataset) - { - let indices = self.dataset.load_indices().await?; - let has_mem_wal = indices.iter().any(|i| i.name == MEM_WAL_INDEX_NAME); - - if !has_mem_wal { - // ZoneMap indices are not supported by MemWAL; exclude them - let maintained_indexes: Vec = indices - .iter() - .filter(|i| { - !(self.id_index_type == IdIndexType::ZoneMap && i.name == ID_INDEX_NAME) - }) - .map(|i| i.name.clone()) - .collect(); - self.dataset - .initialize_mem_wal() - .unsharded() - .maintained_indexes(maintained_indexes) - .execute() - .await?; - } - } - - for ((bot_id, session_id), group_entries) in groups { - let region_id = Self::derive_region_id(&bot_id, &session_id); - let batch = self.records_to_batch(&group_entries)?; - let config = ShardWriterConfig { - shard_id: region_id, - ..Default::default() - }; - - let writer = self.dataset.mem_wal_writer(region_id, config).await?; - writer.put(vec![batch]).await?; - writer.close().await?; + return Ok(self.base.version()); } - - Ok(self.dataset.manifest.version) + let batch = self.records_to_batch(entries)?; + self.base.put(vec![batch]).await?; + Ok(self.base.version()) } /// Resolve a record's external payload reference to its bytes on demand. @@ -572,7 +636,7 @@ impl ContextStore { /// same credentials/endpoint apply when resolving external payload URIs. fn payload_store_params(&self) -> ObjectStoreParams { let mut params = ObjectStoreParams::default(); - if let Some(options) = &self.storage_options { + if let Some(options) = &self.base.storage_options { params.storage_options_accessor = Some(Arc::new( StorageOptionsAccessor::with_static_options(options.clone()), )); @@ -1197,23 +1261,9 @@ impl ContextStore { Ok((existing_ids, existing_external_ids)) } - fn derive_region_id(bot_id: &Option, session_id: &Option) -> Uuid { - let mut input = String::new(); - - if let Some(bid) = bot_id { - input.push_str(bid); - } - input.push('#'); - if let Some(sid) = session_id { - input.push_str(sid); - } - - // Use OID namespace as a base for our deterministic UUIDs - Uuid::new_v5(&Uuid::NAMESPACE_OID, input.as_bytes()) - } - fn has_relationships_column(&self) -> bool { - self.dataset + self.base + .dataset .schema() .field_paths() .iter() @@ -1221,7 +1271,8 @@ impl ContextStore { } fn has_external_id_column(&self) -> bool { - self.dataset + self.base + .dataset .schema() .field_paths() .iter() @@ -1230,7 +1281,7 @@ impl ContextStore { /// Current dataset version. pub fn version(&self) -> u64 { - self.dataset.manifest.version + self.base.version() } /// Add the relationships column to an older dataset if it is missing. @@ -1243,7 +1294,8 @@ impl ContextStore { } let schema = Arc::new(Schema::new(vec![relationship_field()])); - self.dataset + self.base + .dataset .add_columns(NewColumnTransform::AllNulls(schema), None, None) .await?; Ok(true) @@ -1251,15 +1303,15 @@ impl ContextStore { /// Checkout a specific dataset version. pub async fn checkout(&mut self, version_id: u64) -> LanceResult<()> { - let dataset = self.dataset.checkout_version(version_id).await?; - self.dataset = dataset; + let dataset = self.base.dataset.checkout_version(version_id).await?; + self.base.dataset = dataset; Ok(()) } /// Retrieve a single record by its unique ID. pub async fn get(&self, id: &str) -> LanceResult> { let escaped_id = id.replace('\'', "''"); - let mut scanner = self.dataset.scan(); + let mut scanner = self.base.dataset.scan(); scanner.filter(&format!("id = '{}'", escaped_id))?; scanner.limit(Some(1), None)?; @@ -1604,44 +1656,18 @@ impl ContextStore { Ok(results) } + /// LSM scanner over the base table unioned with every shard's flushed + /// MemWAL generations, deduplicating by `id`. async fn lsm_scanner(&self) -> LanceResult { - let object_store = self.dataset.object_store(None).await?; - let branch_location = self.dataset.branch_location(); - let shard_ids = self.dataset.list_mem_wal_latest_shard_ids().await?; - - let mut shard_snapshots = Vec::with_capacity(shard_ids.len()); - for shard_id in shard_ids { - let manifest_store = ShardManifestStore::new( - object_store.clone(), - &branch_location.path, - shard_id, - DEFAULT_MANIFEST_SCAN_BATCH_SIZE, - ); - let Some(manifest) = manifest_store.read_latest().await? else { - continue; - }; - - let mut snapshot = ShardSnapshot::new(shard_id) - .with_spec_id(manifest.shard_spec_id) - .with_current_generation(manifest.current_generation); - for flushed in manifest.flushed_generations { - snapshot = snapshot.with_flushed_generation(flushed.generation, flushed.path); - } - shard_snapshots.push(snapshot); - } - - Ok(LsmScanner::new( - Arc::new(self.dataset.clone()), - shard_snapshots, - vec!["id".to_string()], - )) + self.base.lsm_scanner().await } /// Top-level column names to read for a projection (drops the excluded /// payload columns; everything else is always loaded so lifecycle /// filtering and metadata stay correct). fn projected_columns(&self, projection: ReadProjection) -> Vec { - self.dataset + self.base + .dataset .schema() .fields .iter() @@ -1700,7 +1726,7 @@ impl ContextStore { info!( "Starting compaction: {} fragments", - self.dataset.count_fragments() + self.base.dataset.count_fragments() ); let start = std::time::Instant::now(); @@ -1716,18 +1742,11 @@ impl ContextStore { state.is_compacting = true; } - // Build Lance CompactionOptions - let lance_options = CompactionOptions { - target_rows_per_fragment: config.target_rows_per_fragment, - max_rows_per_group: config.max_rows_per_group, - materialize_deletions: config.materialize_deletions, - materialize_deletions_threshold: config.materialize_deletions_threshold, - num_threads: config.num_threads, - ..Default::default() - }; - - // Run compaction - let result = compact_files(&mut self.dataset, lance_options, None).await; + // Run compaction through the shared layer, which sets + // `defer_index_remap` (the base table carries a fieldless MemWAL index + // that Lance's inline remap panics on) and reloads the handle with this + // store's storage options and session. + let result = self.base.compact(Some(config.clone())).await; // Update state let mut state = self.compaction_state.lock().await; @@ -1749,8 +1768,10 @@ impl ContextStore { metrics.files_added ); - // Reload dataset to see new version - self.dataset = Dataset::open(self.dataset.uri()).await?; + // `StorageBase::compact` already reloaded the handle, with the + // storage options and session attached -- the bare + // `Dataset::open` that used to run here dropped both, which + // broke reopening on any credentialed object store. // Ensure id index exists after compaction // (handles first-time creation on previously empty dataset) @@ -1768,9 +1789,53 @@ impl ContextStore { } } + /// Seal the active memtable so previously added rows become readable on + /// every instance. + /// + /// A no-op when no writer is resident, and unnecessary under the default + /// [`ContextStoreOptions::seal_on_add`] (`true`), where `add` already seals. + /// Stores that opted into deferred sealing must drive this — periodically or + /// explicitly — or their writes stay durable but invisible. + pub async fn flush(&self) -> LanceResult<()> { + self.base.flush().await + } + + /// Gracefully close the resident MemWAL writer, draining its background + /// tasks and sealing whatever it still buffers. Idempotent. + pub async fn close(&mut self) -> LanceResult<()> { + self.base.close().await + } + + /// Fold this instance's flushed MemWAL generations into the base table when + /// it has accumulated [`ContextStoreOptions::merge_after_generations`] of + /// them. Returns the number reclaimed; a no-op when the trigger is unset. + /// + /// # This store previously never merged at all + /// + /// Flushed generations accumulated under `_mem_wal/` forever, and every read + /// unioned all of them, so read cost grew without bound in the number of + /// writes. Merging is what keeps that bounded — drive it from a sweeper, or + /// use [`Self::cleanup_wal`] for the time-based trigger. + pub async fn maybe_merge_wal(&mut self) -> LanceResult { + self.base.maybe_merge_own_shard().await + } + + /// Seal, then fold **every** pending flushed generation into the base table. + /// The time half of the "time OR count" trigger, so deliberately not gated + /// by the count threshold. Returns the number of generations reclaimed. + pub async fn cleanup_wal(&mut self) -> LanceResult { + self.base.cleanup_own_shard().await + } + + /// Number of flushed MemWAL generations pending merge into the base table, + /// across all shards. Read-only: never merges. + pub async fn pending_wal_generations(&self) -> LanceResult { + self.base.pending_wal_generations().await + } + /// Check if compaction should run based on configuration thresholds. pub async fn should_compact(&self) -> LanceResult { - let fragment_count = self.dataset.count_fragments(); + let fragment_count = self.base.dataset.count_fragments(); if fragment_count < self.compaction_config.min_fragments { return Ok(false); @@ -1797,7 +1862,7 @@ impl ContextStore { let state = self.compaction_state.lock().await; Ok(CompactionStats { - total_fragments: self.dataset.count_fragments(), + total_fragments: self.base.dataset.count_fragments(), is_compacting: state.is_compacting, last_compaction: state.last_compaction, last_error: state.last_error.clone(), @@ -1811,7 +1876,7 @@ impl ContextStore { return Ok(()); } - let indices = self.dataset.load_indices().await?; + let indices = self.base.dataset.load_indices().await?; if indices.iter().any(|i| i.name == ID_INDEX_NAME) { return Ok(()); } @@ -1831,16 +1896,17 @@ impl ContextStore { let params = ScalarIndexParams::default(); - self.dataset + self.base + .dataset .create_index_builder(&["id"], index_type, ¶ms) .name(ID_INDEX_NAME.to_string()) .replace(true) .await?; - // Reload dataset to pick up new index - self.dataset = Dataset::open(self.dataset.uri()).await?; - - Ok(()) + // Reload through the base so the new index is visible to subsequent + // reads, keeping the storage options and session (a bare + // `Dataset::open` here silently dropped them). + self.base.reload().await } /// Start background compaction task if enabled. @@ -1860,19 +1926,54 @@ impl ContextStore { self.compaction_config.check_interval_secs, self.compaction_config.min_fragments ); - let mut store_clone = self.clone(); + // The task opens its *own* store handle rather than cloning this one. + // `ContextStore` is no longer `Clone`: it owns a `StorageBase`, which + // owns the resident MemWAL writer and a `Drop` that seals it, so it must + // have exactly one owner. A second handle is the right model anyway -- + // compaction only rewrites base-table fragments and takes `&mut`, so + // sharing a handle with the write path would mean contending for it. + let uri = self.uri().to_string(); let interval_secs = self.compaction_config.check_interval_secs; + let options = ContextStoreOptions { + storage_options: self.base.storage_options.clone(), + // `enabled: false` is load-bearing, not tidiness: the reopened + // handle would otherwise spawn its own background compactor, which + // makes `start_background_compaction` infinitely recursive (and the + // resulting future un-`Send`, so it does not even compile). + compaction: CompactionConfig { + enabled: false, + ..self.compaction_config.clone() + }, + embedding_dim: Some(self.embedding_dim), + blob_columns: self.blob_columns.clone(), + id_index_type: self.id_index_type, + distance_metric: Some(self.distance_metric), + shard_id: None, + merge_after_generations: None, + // A compactor never appends, so the seal mode is irrelevant to it; + // deferring keeps it from ever emitting a generation. + seal_on_add: false, + }; let task = tokio::spawn(async move { + let compaction_options = options; let mut interval = tokio::time::interval(Duration::from_secs(interval_secs)); loop { interval.tick().await; - match store_clone.should_compact().await { + let mut store = match open_for_compaction(&uri, compaction_options.clone()).await { + Ok(store) => store, + Err(e) => { + error!("Background compaction could not open store: {}", e); + continue; + } + }; + + match store.should_compact().await { Ok(true) => { info!("Background compaction triggered"); - if let Err(e) = store_clone.compact(None).await { + if let Err(e) = store.compact(None).await { error!("Background compaction failed: {}", e); } } @@ -2050,14 +2151,7 @@ impl ContextStore { uri: &str, storage_options: Option>, ) -> LanceResult { - if let Some(options) = storage_options { - DatasetBuilder::from_uri(uri) - .with_storage_options(options) - .load() - .await - } else { - Dataset::open(uri).await - } + StorageBase::load_with_options(uri, storage_options, None).await } async fn create_with_options( @@ -2077,62 +2171,47 @@ impl ContextStore { embedding_dim, distance_metric, )); - let empty_batch = RecordBatch::new_empty(schema.clone()); - let batches = RecordBatchIterator::new( - vec![Ok::(empty_batch)].into_iter(), - schema.clone(), - ); - - let mut params = WriteParams { - mode: WriteMode::Create, - ..Default::default() - }; - - if let Some(options) = storage_options { - let store_params = ObjectStoreParams { - storage_options_accessor: Some(Arc::new( - StorageOptionsAccessor::with_static_options(options), - )), - ..Default::default() - }; - params.store_params = Some(store_params); - } - - Dataset::write(batches, uri, Some(params)).await + StorageBase::create_with_options(uri, schema, storage_options, None).await } fn records_to_batch(&self, entries: &[ContextRecord]) -> LanceResult { let include_external_id = self + .base .dataset .schema() .field_paths() .iter() .any(|path| path == "external_id"); let include_lifecycle = self + .base .dataset .schema() .field_paths() .iter() .any(|path| path == "expires_at"); let include_metadata = self + .base .dataset .schema() .field_paths() .iter() .any(|path| path == "metadata"); let include_tenant = self + .base .dataset .schema() .field_paths() .iter() .any(|path| path == "tenant"); let include_source = self + .base .dataset .schema() .field_paths() .iter() .any(|path| path == "source"); let include_external_reference = self + .base .dataset .schema() .field_paths() @@ -2458,7 +2537,7 @@ impl ContextStore { ]); } - let schema: Arc = Arc::new(self.dataset.schema().into()); + let schema: Arc = Arc::new(self.base.dataset.schema().into()); let arrays = schema .fields() .iter() @@ -3118,8 +3197,11 @@ fn sql_quoted_list(values: &[&str]) -> String { #[cfg(test)] mod tests { use super::*; + use crate::serde::CONTENT_TYPE_TEXT; + use arrow_array::RecordBatchIterator; use chrono::{Duration as ChronoDuration, Utc}; + use lance::dataset::{WriteMode, WriteParams}; use tempfile::TempDir; fn make_embedding_with_dim(dim: usize, pivot: f32) -> Vec { @@ -3184,7 +3266,7 @@ mod tests { let runtime = tokio::runtime::Runtime::new().unwrap(); runtime.block_on(async { - let mut store = ContextStore::open(&uri).await.unwrap(); + let store = ContextStore::open(&uri).await.unwrap(); // Offload bytes to the object store, then reference them by URI. let written = store.put_payload(&object_uri, &payload).await.unwrap(); @@ -3219,7 +3301,7 @@ mod tests { 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 store = ContextStore::open(&uri).await.unwrap(); // Unknown id resolves to None rather than erroring. assert_eq!(store.fetch_payload("does-not-exist").await.unwrap(), None); @@ -3238,7 +3320,7 @@ mod tests { 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 store = ContextStore::open(&uri).await.unwrap(); let first = text_record("a", 0.0); let second = text_record("b", 1.0); store.add(&[first.clone(), second.clone()]).await.unwrap(); @@ -3324,7 +3406,7 @@ mod tests { // Default (L2): `near` should rank first. let l2_dir = TempDir::new().unwrap(); - let mut l2_store = ContextStore::open(&l2_dir.path().to_string_lossy()) + let l2_store = ContextStore::open(&l2_dir.path().to_string_lossy()) .await .unwrap(); l2_store @@ -3343,7 +3425,7 @@ mod tests { distance_metric: Some(DistanceMetric::Cosine), ..Default::default() }; - let mut cos_store = + let cos_store = ContextStore::open_with_options(&cos_dir.path().to_string_lossy(), cos_opts) .await .unwrap(); @@ -3363,7 +3445,7 @@ mod tests { distance_metric: Some(DistanceMetric::Dot), ..Default::default() }; - let mut dot_store = + let dot_store = ContextStore::open_with_options(&dot_dir.path().to_string_lossy(), dot_opts) .await .unwrap(); @@ -3395,7 +3477,7 @@ mod tests { distance_metric: Some(DistanceMetric::Cosine), ..Default::default() }; - let mut store = ContextStore::open_with_options(&uri, opts).await.unwrap(); + let store = ContextStore::open_with_options(&uri, opts).await.unwrap(); store .add(&[ text_record_with("aligned", aligned.clone()), @@ -3465,7 +3547,7 @@ mod tests { 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 store = ContextStore::open(&uri).await.unwrap(); let mut semantic_near = text_record("semantic-near", 0.0); semantic_near.text_payload = Some("general rollout risk guidance".to_string()); let mut exact_policy = text_record("exact-policy", 1.0); @@ -3507,7 +3589,7 @@ mod tests { embedding_dim: Some(3), ..Default::default() }; - let mut store = ContextStore::open_with_options(&uri, options) + let store = ContextStore::open_with_options(&uri, options) .await .unwrap(); assert_eq!(store.embedding_dim(), 3); @@ -3544,7 +3626,7 @@ mod tests { 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 store = ContextStore::open(&uri).await.unwrap(); assert_eq!(store.embedding_dim(), DEFAULT_EMBEDDING_DIM); store.add(&[text_record("default-dim", 0.0)]).await.unwrap(); drop(store); @@ -3594,7 +3676,7 @@ mod tests { 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 store = ContextStore::open(&uri).await.unwrap(); let active = text_record("active", 0.0); let mut expired = text_record("expired", 0.0); expired.expires_at = Some(Utc::now() - ChronoDuration::minutes(1)); @@ -3642,7 +3724,7 @@ mod tests { 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 store = ContextStore::open(&uri).await.unwrap(); let old = text_record("old", 0.0); let mut replacement = text_record("new", 1.0); replacement.supersedes_id = Some(old.id.clone()); @@ -3671,7 +3753,7 @@ mod tests { 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 store = ContextStore::open(&uri).await.unwrap(); let active = text_record("active", 1.0); let mut expired_better_match = text_record("expired", 0.0); expired_better_match.expires_at = Some(Utc::now() - ChronoDuration::minutes(1)); @@ -3700,7 +3782,7 @@ mod tests { 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 store = ContextStore::open(&uri).await.unwrap(); let mut record = text_record("a", 0.0); record.external_id = Some("doc-123#chunk-1".to_string()); store.add(std::slice::from_ref(&record)).await.unwrap(); @@ -4150,7 +4232,7 @@ mod tests { 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 store = ContextStore::open(&uri).await.unwrap(); let mut related = text_record("related", 0.0); related.relationships = vec![ Relationship { @@ -4259,7 +4341,7 @@ mod tests { 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 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(); @@ -4281,7 +4363,7 @@ mod tests { 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 store = ContextStore::open(&uri).await.unwrap(); let mut record = text_record("a", 0.0); record.content_type = CONTENT_TYPE_TOMBSTONE.to_string(); @@ -4300,7 +4382,7 @@ mod tests { 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 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(); @@ -4318,7 +4400,7 @@ mod tests { 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 store = ContextStore::open(&uri).await.unwrap(); let err = store .add(&[text_record("same", 0.0), text_record("same", 1.0)]) .await @@ -4337,7 +4419,7 @@ mod tests { 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 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); @@ -4492,7 +4574,7 @@ mod tests { id_index_type: IdIndexType::BTree, ..Default::default() }; - let mut store = ContextStore::open_with_options(&uri, options) + let store = ContextStore::open_with_options(&uri, options) .await .unwrap(); @@ -4607,7 +4689,7 @@ mod tests { 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 store = ContextStore::open(&uri).await.unwrap(); let tricky = "o'brien#chunk-1"; let mut first = text_record("a", 0.0); @@ -4744,35 +4826,105 @@ mod tests { } #[test] - fn test_region_id_derivation_explicit() { - let bot_id = Some("bot-123".to_string()); - let session_id = Some("session-456".to_string()); + fn add_is_visible_on_return_by_default() { + // The default `seal_on_add: true` contract. `ContextStore`'s uniqueness + // validation, upsert and tombstones all read the table back, so an `add` + // that returned before sealing would silently break them. + 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 store = ContextStore::open(&uri).await.unwrap(); + store.add(&[text_record("r1", 0.0)]).await.unwrap(); - let region_id_1 = ContextStore::derive_region_id(&bot_id, &session_id); - let region_id_2 = ContextStore::derive_region_id(&bot_id, &session_id); + // No flush: the row must already be readable. + assert_eq!(store.list(None, None).await.unwrap().len(), 1); + assert!(store.get_by_id("r1").await.unwrap().is_some()); + }); + } - assert_eq!( - region_id_1, region_id_2, - "Region ID should be deterministic for same inputs" - ); + #[test] + fn deferred_seal_makes_add_durable_but_invisible_until_flush() { + // Opting into `RolloutStore`'s write profile: `add` is a durable WAL + // append only, and visibility waits for an explicit flush. Mirrors + // `rollout_store::tests::add_is_durable_but_not_visible_until_flush`. + 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 store = ContextStore::open_with_options( + &uri, + ContextStoreOptions { + seal_on_add: false, + ..Default::default() + }, + ) + .await + .unwrap(); - let other_session = Some("session-789".to_string()); - let region_id_3 = ContextStore::derive_region_id(&bot_id, &other_session); + store.add(&[text_record("r1", 0.0)]).await.unwrap(); + assert_eq!( + store.list(None, None).await.unwrap().len(), + 0, + "deferred seal must not publish the row" + ); - assert_ne!( - region_id_1, region_id_3, - "Region ID should differ for different inputs" - ); + store.flush().await.unwrap(); + assert_eq!(store.list(None, None).await.unwrap().len(), 1); + assert!(store.get_by_id("r1").await.unwrap().is_some()); + }); + } - // Test None/None case (now deterministic based on empty strings) - let region_id_none = ContextStore::derive_region_id(&None, &None); - let region_id_none_2 = ContextStore::derive_region_id(&None, &None); - assert_eq!( - region_id_none, region_id_none_2, - "Region ID for None/None should be deterministic" - ); + #[test] + fn wal_generations_merge_into_the_base_table() { + // This store previously never merged: flushed generations accumulated + // under `_mem_wal/` forever and every read unioned all of them. Each + // sealing `add` emits one generation, so after N adds the merge must + // reclaim them and leave every row readable from the base table. + 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(); + for i in 0..3 { + store + .add(&[text_record(&format!("r{i}"), i as f32)]) + .await + .unwrap(); + } + + assert!( + store.pending_wal_generations().await.unwrap() > 0, + "sealing adds should leave generations pending" + ); + + let reclaimed = store.cleanup_wal().await.unwrap(); + assert!(reclaimed > 0, "cleanup must merge the pending generations"); + assert_eq!(store.pending_wal_generations().await.unwrap(), 0); + + // Rows survive the merge and are still readable. + let ids: Vec = store + .list(None, None) + .await + .unwrap() + .iter() + .map(|r| r.id.clone()) + .collect(); + assert_eq!(ids.len(), 3); + for i in 0..3 { + assert!(ids.contains(&format!("r{i}"))); + } + }); } + // `test_region_id_derivation_explicit` was deleted with the per-conversation + // MemWAL sharding it covered. Writes now go to the shard owned by this + // writer instance (see `ContextStoreOptions::shard_id`), matching + // `RolloutStore` -- one instance owns exactly one shard, so the epoch-fencing + // invariant holds by construction. `test_add_multiple_regions` below still + // pins the user-visible behavior: records with differing bot/session ids + // round-trip through a single store. + #[test] fn test_add_multiple_regions() { let dir = TempDir::new().unwrap(); @@ -4780,7 +4932,7 @@ mod tests { let runtime = tokio::runtime::Runtime::new().unwrap(); runtime.block_on(async { - let mut store = ContextStore::open(&uri).await.unwrap(); + let store = ContextStore::open(&uri).await.unwrap(); // Create records for different regions let mut record1 = text_record("r1", 0.0); @@ -4821,7 +4973,7 @@ mod tests { blob_columns: HashSet::from(["binary_payload".to_string()]), ..Default::default() }; - let mut store = ContextStore::open_with_options(&uri, options) + let store = ContextStore::open_with_options(&uri, options) .await .unwrap(); @@ -4854,7 +5006,7 @@ mod tests { blob_columns: HashSet::from(["text_payload".to_string()]), ..Default::default() }; - let mut store = ContextStore::open_with_options(&uri, options) + let store = ContextStore::open_with_options(&uri, options) .await .unwrap(); @@ -4896,7 +5048,7 @@ mod tests { ]), ..Default::default() }; - let mut store = ContextStore::open_with_options(&uri, options) + let store = ContextStore::open_with_options(&uri, options) .await .unwrap(); @@ -5038,7 +5190,7 @@ mod tests { .unwrap(); // Index should be created eagerly on open - let indices = store.dataset.load_indices().await.unwrap(); + let indices = store.base.dataset.load_indices().await.unwrap(); assert!( indices.iter().any(|i| i.name == ID_INDEX_NAME), "btree index should be created on open" @@ -5054,7 +5206,7 @@ mod tests { store.compact(None).await.unwrap(); // Index should still exist after compaction - let indices = store.dataset.load_indices().await.unwrap(); + let indices = store.base.dataset.load_indices().await.unwrap(); assert!( indices.iter().any(|i| i.name == ID_INDEX_NAME), "btree index should persist after compaction" @@ -5078,7 +5230,7 @@ mod tests { .unwrap(); // Index should be created eagerly on open - let indices = store.dataset.load_indices().await.unwrap(); + let indices = store.base.dataset.load_indices().await.unwrap(); assert!( indices.iter().any(|i| i.name == ID_INDEX_NAME), "zonemap index should be created on open" @@ -5092,7 +5244,7 @@ mod tests { } store.compact(None).await.unwrap(); - let indices = store.dataset.load_indices().await.unwrap(); + let indices = store.base.dataset.load_indices().await.unwrap(); assert!( indices.iter().any(|i| i.name == ID_INDEX_NAME), "zonemap index should persist after compaction" @@ -5112,7 +5264,7 @@ mod tests { store.add(&[text_record("no-idx-1", 0.0)]).await.unwrap(); store.compact(None).await.unwrap(); - let indices = store.dataset.load_indices().await.unwrap(); + let indices = store.base.dataset.load_indices().await.unwrap(); assert!( !indices.iter().any(|i| i.name == ID_INDEX_NAME), "no id index should be created when IdIndexType::None" @@ -5157,7 +5309,7 @@ mod tests { 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 store = ContextStore::open(&uri).await.unwrap(); let mut record = text_record("img", 0.0); record.content_type = "image/png".to_string(); record.binary_payload = Some(vec![1, 2, 3, 4]); @@ -5213,7 +5365,7 @@ mod tests { 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 store = ContextStore::open(&uri).await.unwrap(); let mut a = text_record("a", 0.0); a.binary_payload = Some(vec![9, 9, 9]); let mut b = text_record("b", 1.0); diff --git a/crates/lance-context-core/src/store_base.rs b/crates/lance-context-core/src/store_base.rs index d2c57e9..eb9990c 100644 --- a/crates/lance-context-core/src/store_base.rs +++ b/crates/lance-context-core/src/store_base.rs @@ -154,6 +154,27 @@ pub(crate) struct StorageBaseOptions { /// older base table to it by adding missing nullable columns as all-nulls /// (see [`StorageBase::ensure_latest_schema`]). `None` disables evolution. pub latest_schema: Option>, + /// Whether [`StorageBase::put`] seals the memtable before returning, making + /// the rows immediately readable. + /// + /// This is the one deliberate behavioral difference between the stores, and + /// it is a **visibility/throughput trade, not a durability one** — a `put` + /// is durable either way. + /// + /// - `false` (rollout): `put` is a durable WAL append only. Concurrent + /// appends are not serialized behind a per-append seal, but there is **no + /// read-your-write guarantee** — visibility is bounded by whatever drives + /// [`StorageBase::flush`] (the server's flush sweeper, 30s by default), by + /// Lance's own memtable-size thresholds, or by an explicit per-request + /// flush. + /// - `true` (context): `put` seals before returning, so a subsequent read + /// sees the rows. Required by any store whose writes read the table back — + /// `ContextStore`'s id/external-id uniqueness validation, upsert, and + /// tombstones all depend on it. + /// + /// Sealing per write produces a flushed generation per call, which is + /// exactly the read amplification the high-fan-in rollout path avoids. + pub seal_on_put: bool, } /// Schema-agnostic Lance storage: dataset handle, MemWAL write path, WAL merge, @@ -175,6 +196,8 @@ pub(crate) struct StorageBase { pub key_column: String, /// Latest expected schema for merge-time evolution; see the option. latest_schema: Option>, + /// Whether `put` seals before returning; see [`StorageBaseOptions::seal_on_put`]. + seal_on_put: bool, /// Self-merge threshold; `0` disables it. merge_after_generations: usize, /// Timestamp of the last successful [`Self::compact`] on this handle. @@ -227,15 +250,9 @@ impl StorageBase { schema, key_column, latest_schema, + seal_on_put, } = options; - if schema.field_with_name(&key_column).is_err() { - return Err(ArrowError::SchemaError(format!( - "key column '{key_column}' is not present in the store schema" - )) - .into()); - } - let dataset = match Self::load_with_options(uri, storage_options.clone(), session.clone()).await { Ok(dataset) => dataset, @@ -251,6 +268,48 @@ impl StorageBase { Err(err) => return Err(err), }; + Self::from_dataset( + dataset, + StorageBaseOptions { + storage_options, + shard_id, + merge_after_generations, + session, + schema, + key_column, + latest_schema, + seal_on_put, + }, + ) + .await + } + + /// Wrap an already-open [`Dataset`]. + /// + /// For stores that must inspect the dataset before they can describe + /// themselves: `ContextStore` reads the embedding width and distance metric + /// out of the existing schema (and validates them against the caller's + /// request) before it knows what schema to hand the base. [`Self::open`] is + /// the simpler path when the schema is known up front. + pub async fn from_dataset(dataset: Dataset, options: StorageBaseOptions) -> LanceResult { + let StorageBaseOptions { + storage_options, + shard_id, + merge_after_generations, + session, + schema, + key_column, + latest_schema, + seal_on_put, + } = options; + + if schema.field_with_name(&key_column).is_err() { + return Err(ArrowError::SchemaError(format!( + "key column '{key_column}' is not present in the store schema" + )) + .into()); + } + let mut base = Self { dataset, write_shard: derive_shard_id(shard_id.as_deref()), @@ -258,6 +317,7 @@ impl StorageBase { session, key_column, latest_schema, + seal_on_put, merge_after_generations: merge_after_generations.unwrap_or(0), last_compaction: None, total_compactions: 0, @@ -302,26 +362,28 @@ impl StorageBase { /// Durably append `batches` through this instance's MemWAL shard. /// - /// # Durable on return, *not* visible on return + /// # Always durable on return; visible on return only if `seal_on_put` /// - /// The only per-append work is `put`, which returns once the WAL entry has - /// been PUT to object storage. The rows are then **durable** — they survive - /// a crash and are replayed on reopen — but they are **not yet readable**, - /// by this instance or any other. A row becomes visible only after its - /// memtable is sealed into a flushed generation and committed to the shard - /// manifest, which happens in [`Self::flush`] (also performed by - /// [`Self::close`], and by the merge path). + /// `put` returns once the WAL entry has been PUT to object storage. The rows + /// are then **durable** — they survive a crash and are replayed on reopen. + /// Whether they are also **readable** depends on + /// [`StorageBaseOptions::seal_on_put`]: /// - /// Callers therefore get **no read-your-write guarantee**; a caller that - /// needs the row readable immediately must `put(..).await` then - /// `flush().await`. + /// - `false` (rollout): the rows sit in the active memtable, invisible to + /// every reader including this one, until something seals them into a + /// flushed generation and commits it to the shard manifest — [`Self::flush`] + /// (driven by the server's flush sweeper), [`Self::close`], the merge path, + /// or Lance's own memtable-size thresholds. Callers get **no + /// read-your-write guarantee**. + /// - `true` (context): this call seals before returning, so a subsequent + /// read sees the rows. /// - /// This decoupling is deliberate: sealing on the append path serialized - /// concurrent appends behind one seal+drain. Keeping only the durable `put` - /// here lets appends run concurrently, and reuses a single resident - /// [`ShardWriter`] so the shard epoch is claimed once and the object-store - /// connection is pooled, rather than paying a cold DNS resolution + - /// TCP/TLS handshake + epoch claim per append. + /// Deferring the seal is what lets concurrent appends run without + /// serializing behind one seal+drain, and it avoids emitting a flushed + /// generation per call. Either way a single resident [`ShardWriter`] is + /// reused, so the shard epoch is claimed once and the object-store + /// connection is pooled, rather than paying a cold DNS resolution + TCP/TLS + /// handshake + epoch claim per append. /// /// Retried exactly once on a fence: a fence means a merge superseded our /// epoch, and invalidating + reopening re-claims the current one. Rows are @@ -332,7 +394,7 @@ impl StorageBase { return Ok(()); } let started = timer_start!(); - let result = self.put_inner(batches).await; + let result = self.put_sealing(batches).await; // Success-path latency only. A failed append's *duration* is not // actionable, but its *rate* is, so errors increment a flat counter // rather than doubling this histogram's series count. @@ -346,6 +408,19 @@ impl StorageBase { result } + /// [`Self::put_inner`], then seal if this store wants read-your-write. + /// + /// The seal is inside the timed region because for a `seal_on_put` store it + /// is part of what the caller is waiting for; excluding it would understate + /// the write latency such a store actually sees. + async fn put_sealing(&self, batches: Vec) -> LanceResult<()> { + self.put_inner(batches).await?; + if self.seal_on_put { + self.flush().await?; + } + Ok(()) + } + async fn put_inner(&self, batches: Vec) -> LanceResult<()> { let writer = self.resident_writer().await?; match writer.put(batches.clone()).await { @@ -995,7 +1070,7 @@ impl StorageBase { /// Reload the base dataset handle through [`Self::load_with_options`], so /// the shared session and storage options are never dropped. - async fn reload(&mut self) -> LanceResult<()> { + pub async fn reload(&mut self) -> LanceResult<()> { let uri = self.dataset.uri().to_string(); self.dataset = Self::load_with_options(&uri, self.storage_options.clone(), self.session.clone()) diff --git a/crates/lance-context-server/src/routes/records.rs b/crates/lance-context-server/src/routes/records.rs index 0a7395e..2365607 100644 --- a/crates/lance-context-server/src/routes/records.rs +++ b/crates/lance-context-server/src/routes/records.rs @@ -45,7 +45,10 @@ pub async fn add_records( } let count = core_records.len(); - let mut store = store_lock.write().await; + // Read lock: `add` is `&self` now that the resident writer lives behind the + // storage layer's own mutex, so concurrent appends no longer serialize on + // this store's RwLock. + let store = store_lock.read().await; let version = store .add(&core_records) .await diff --git a/python/src/lib.rs b/python/src/lib.rs index 9cc9283..79fff0a 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -220,6 +220,7 @@ fn context_options_from_py<'py>( blob_columns: blob_set, id_index_type: id_idx, distance_metric: metric, + ..Default::default() }) } @@ -773,13 +774,21 @@ impl Context { self.inner.snapshot(label) } - fn fork(&self, branch_name: &str) -> Self { - Self { + fn fork(&self, py: Python<'_>, branch_name: &str) -> PyResult { + // Opens a second handle on the same dataset rather than cloning this + // one: `ContextStore` owns a resident MemWAL writer (and a `Drop` that + // seals it), so it is deliberately not `Clone` -- two owners would mean + // two writers racing for one shard. A fork branches the in-memory + // `Context` and shares the underlying dataset, which a fresh handle + // gives it. + let uri = self.store.uri().to_string(); + let store = py.allow_threads(|| self.runtime.block_on(ContextStore::open(&uri))); + Ok(Self { inner: self.inner.fork(branch_name), - store: self.store.clone(), + store: store.map_err(to_py_err)?, runtime: Arc::clone(&self.runtime), run_id: new_run_id(), - } + }) } fn checkout(&mut self, py: Python<'_>, version_id: u64) -> PyResult<()> { From e98fc8281ca1e8b61d6e3f3a4db7fb87f3ce2237 Mon Sep 17 00:00:00 2001 From: Beinan Date: Mon, 27 Jul 2026 23:05:36 +0000 Subject: [PATCH 2/2] refactor(core): migrate DatagenStore onto StorageBase MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Third and last store of #214 Part 1, after #215 and #216. All three fixed-schema stores now share one storage implementation. Closes this store's gaps: - **Compaction, which it had none of.** Every WAL merge appends a fragment to the base table and datagen seals once per checkpoint batch, so fragments grew without bound and scans degraded with them. Now exposed as `compact` / `should_compact` / `compaction_stats`, with `defer_index_remap` set (the fieldless MemWAL index panics Lance's inline remap). - **Scalar index, which it also had none of.** `create_event_id_index` builds a ZoneMap on `event_id`; point lookups by event id previously always scanned. - **`load_with_options` dropped storage options.** It fell back to a bare `Dataset::open` whenever no options were passed, and never attached a session, so every generation read allocated Lance's 6 GiB default cache. Deletes the near-line-for-line copy of rollout's `merge_own_shard` (plus `merge_own_shard_if_ready`, the writer lifecycle, `ensure_mem_wal`, shard snapshot discovery, the `Drop` impl and a second `is_fenced_error`), and picks up the streaming-safe version in the base instead of buffering every flushed generation in memory. `seal_on_put: true`: each append is one complete checkpoint batch that must be durable *and* readable before the writer moves on, so a crash cannot expose a partially checkpointed step. This preserves the store's existing put -> seal -> drain behavior exactly. `key_column` is `event_id`, not `id` — event ids are derived deterministically (UUIDv5 over item/checkpoint/ordinal), which is what makes a retried append idempotent under LSM dedup. The base already took the key column as a parameter, so no change was needed there. New test `compaction_folds_merged_fragments_and_preserves_reads` pins the behavior that did not exist before: fragments accumulate across merges, compaction reduces them, and every event stays readable. Full suite passes — 181 core lib tests, 182 Python tests, clippy clean. Co-Authored-By: Claude --- .../lance-context-core/src/datagen_store.rs | 567 +++++------------- 1 file changed, 149 insertions(+), 418 deletions(-) diff --git a/crates/lance-context-core/src/datagen_store.rs b/crates/lance-context-core/src/datagen_store.rs index 11b8295..738d7a5 100644 --- a/crates/lance-context-core/src/datagen_store.rs +++ b/crates/lance-context-core/src/datagen_store.rs @@ -8,7 +8,7 @@ //! the base table with every flushed shard and de-duplicate by deterministic //! `event_id`, making a retried checkpoint batch idempotent. -use std::collections::{HashMap, HashSet}; +use std::collections::HashMap; use std::sync::Arc; use arrow_array::builder::{ @@ -17,33 +17,26 @@ use arrow_array::builder::{ }; use arrow_array::{ Array, ArrayRef, BooleanArray, Float64Array, Int32Array, Int64Array, LargeBinaryArray, - LargeStringArray, RecordBatch, RecordBatchIterator, StringArray, TimestampMicrosecondArray, - UInt64Array, + LargeStringArray, RecordBatch, StringArray, TimestampMicrosecondArray, UInt64Array, }; use arrow_schema::{ArrowError, DataType, Field, Schema, TimeUnit}; -use futures::{stream, StreamExt, TryStreamExt}; -use lance::dataset::mem_wal::{ - DatasetMemWalExt, LsmScanner, ShardManifestStore, ShardSnapshot, ShardWriter, ShardWriterConfig, -}; -use lance::dataset::{builder::DatasetBuilder, Dataset, WriteMode, WriteParams}; -use lance::index::DatasetIndexExt; -use lance::io::{ObjectStoreParams, StorageOptionsAccessor}; +use futures::TryStreamExt; +use lance::dataset::mem_wal::LsmScanner; +use lance::dataset::optimize::CompactionMetrics; +use lance::dataset::Dataset; use lance::{Error as LanceError, Result as LanceResult}; -use lance_index::mem_wal::{ShardManifest, MEM_WAL_INDEX_NAME}; use tokio::task::JoinHandle; use tracing::{info, warn}; -use uuid::Uuid; use crate::datagen::{ datagen_failures, datagen_trajectory, fold_datagen_events, DatagenBlobValue, DatagenEvent, DatagenEventType, DatagenFailure, DatagenItemLookup, DatagenItemStatus, DatagenRootItemStatuses, DatagenStepCursor, DatagenStepKind, DatagenValue, }; -use crate::store::{column_as, column_as_optional, timestamp_from_micros}; -use crate::store_base::{align_batch_to_schema, derive_shard_id, is_not_found_error}; - -const DEFAULT_MANIFEST_SCAN_BATCH_SIZE: usize = 16; -const DEFAULT_SHARD_SCAN_CONCURRENCY: usize = 16; +use crate::store::{ + column_as, column_as_optional, timestamp_from_micros, CompactionConfig, CompactionStats, +}; +use crate::store_base::{is_not_found_error, StorageBase, StorageBaseOptions}; /// Configuration for opening a [`DatagenStore`]. #[derive(Debug, Clone, Default)] @@ -61,13 +54,13 @@ pub struct DatagenStoreOptions { } /// A single append-only Lance dataset for datagen checkpoint events. +/// +/// Storage mechanics live in [`StorageBase`]; what remains here is the datagen +/// schema ([`datagen_log_schema`]), event encode/decode, and the typed +/// fold/trajectory read APIs. pub struct DatagenStore { - dataset: Dataset, - write_shard: Uuid, - storage_options: Option>, - merge_after_generations: usize, + base: StorageBase, cleanup_interval_secs: u64, - write_writer: Option, } impl DatagenStore { @@ -91,33 +84,43 @@ impl DatagenStore { options: DatagenStoreOptions, create_if_missing: bool, ) -> LanceResult { - let storage_options = options.storage_options.clone(); - let dataset = match Self::load_with_options(uri, storage_options.clone()).await { - Ok(dataset) => dataset, - Err(LanceError::DatasetNotFound { .. }) if create_if_missing => { - Self::create_with_options(uri, storage_options.clone()).await? - } - Err(error) => return Err(error), - }; + let base = StorageBase::open( + uri, + StorageBaseOptions { + storage_options: options.storage_options.clone(), + shard_id: options.shard_id.clone(), + merge_after_generations: options.merge_after_generations, + session: None, + schema: Arc::new(datagen_log_schema()), + // Datagen keys on `event_id`, not `id`: event ids are derived + // deterministically (UUIDv5 over item/checkpoint/ordinal), which + // is what makes a retried append idempotent under LSM dedup. + key_column: "event_id".to_string(), + latest_schema: None, + // Each append is one complete checkpoint batch and must be + // durable *and* readable before the writer moves on, so a crash + // cannot expose a partially checkpointed step. This preserves + // the store's existing put -> seal -> drain behavior. + seal_on_put: true, + }, + create_if_missing, + ) + .await?; Ok(Self { - dataset, - write_shard: derive_shard_id(options.shard_id.as_deref()), - storage_options, - merge_after_generations: options.merge_after_generations.unwrap_or(0), + base, cleanup_interval_secs: options.cleanup_interval_secs.unwrap_or(0), - write_writer: None, }) } #[must_use] pub fn uri(&self) -> &str { - self.dataset.uri() + self.base.dataset.uri() } #[must_use] pub fn version(&self) -> u64 { - self.dataset.manifest.version + self.base.version() } /// Append one or more complete checkpoint batches. @@ -127,18 +130,16 @@ impl DatagenStore { /// the same call so a crash cannot expose a partially checkpointed step. pub async fn append(&mut self, events: &[DatagenEvent]) -> LanceResult { if events.is_empty() { - return Ok(self.dataset.manifest.version); + return Ok(self.base.version()); } validate_write_batch(events)?; let batch = events_to_batch(events)?; - self.ensure_mem_wal().await?; - self.write_with_resident_writer(&batch).await?; - if self.merge_after_generations > 0 { - self.merge_own_shard_if_ready(self.merge_after_generations) - .await?; - } - Ok(self.dataset.manifest.version) + // Encoding and validation are schema-specific and stay here; the + // resident writer, fence retry, seal+drain and metrics are the base's. + self.base.put(vec![batch]).await?; + self.base.maybe_merge_own_shard().await?; + Ok(self.base.version()) } /// Append one completed step boundary atomically. @@ -151,51 +152,9 @@ impl DatagenStore { self.append(events).await } - async fn write_with_resident_writer(&mut self, batch: &RecordBatch) -> LanceResult<()> { - match self.put_seal_drain(batch).await { - Ok(()) => Ok(()), - Err(error) if is_fenced_error(&error) => { - self.write_writer = None; - self.put_seal_drain(batch).await - } - Err(error) => Err(error), - } - } - - async fn put_seal_drain(&mut self, batch: &RecordBatch) -> LanceResult<()> { - self.ensure_write_writer().await?; - let writer = self - .write_writer - .as_ref() - .expect("ensure_write_writer set the writer"); - writer.put(vec![batch.clone()]).await?; - writer.force_seal_active().await?; - writer.wait_for_flush_drain().await?; - Ok(()) - } - - async fn ensure_write_writer(&mut self) -> LanceResult<()> { - if self.write_writer.is_some() { - return Ok(()); - } - let config = ShardWriterConfig { - shard_id: self.write_shard, - ..Default::default() - }; - self.write_writer = Some( - self.dataset - .mem_wal_writer(self.write_shard, config) - .await?, - ); - Ok(()) - } - /// Gracefully stop this store's resident MemWAL writer. pub async fn close(&mut self) -> LanceResult<()> { - if let Some(writer) = self.write_writer.take() { - writer.close().await?; - } - Ok(()) + self.base.close().await } /// Read one item's event history without materializing blob bytes. @@ -273,14 +232,15 @@ impl DatagenStore { /// matching `_rowid` is then passed to `take_rows` for an O(single blob) /// payload read. pub async fn get_blob(&self, event_id: &str) -> LanceResult>> { - let snapshots = self.wal_shard_snapshots().await?; + let snapshots = self.base.wal_shard_snapshots().await?; let mut generations: Vec<(u64, String)> = snapshots .iter() .flat_map(|snapshot| { snapshot.flushed_generations.iter().map(|generation| { ( generation.generation, - self.flushed_generation_uri(snapshot.shard_id, &generation.path), + self.base + .flushed_generation_uri(snapshot.shard_id, &generation.path), ) }) }) @@ -293,7 +253,7 @@ impl DatagenStore { // base table (checked below), so skip it rather than failing the // whole lookup. Now that merges no longer block appends, this window // is genuinely reachable. (Mirrors the rollout store.) - let dataset = match self.open_flushed_dataset(&uri).await { + let dataset = match self.base.open_flushed_dataset(&uri).await { Ok(dataset) => dataset, Err(err) if is_not_found_error(&err) => continue, Err(err) => return Err(err), @@ -303,25 +263,63 @@ impl DatagenStore { } } - Ok(Self::get_blob_from_dataset(&self.dataset, event_id) + Ok(Self::get_blob_from_dataset(&self.base.dataset, event_id) .await? .flatten()) } /// Number of flushed generations waiting across all writer shards. pub async fn pending_wal_generations(&self) -> LanceResult { - Ok(self - .wal_shard_snapshots() - .await? - .iter() - .map(|snapshot| snapshot.flushed_generations.len()) - .sum()) + self.base.pending_wal_generations().await } /// Merge every currently flushed generation owned by this writer into the /// base table. pub async fn cleanup_own_shard(&mut self) -> LanceResult { - self.merge_own_shard_if_ready(1).await + self.base.cleanup_own_shard().await + } + + /// Compact the base table's small fragments into larger ones. + /// + /// # This store previously had no compaction at all + /// + /// Every WAL merge appends a fragment to the base table, and datagen seals + /// once per checkpoint batch, so fragments accumulated without bound and + /// scans degraded accordingly. Like the other stores this must be driven by + /// a *single* external trigger, not a per-worker timer: compaction rewrites + /// the shared base table and Lance treats two concurrent `Rewrite` commits + /// as a conflict. + pub async fn compact( + &mut self, + options: Option, + ) -> LanceResult { + self.base.compact(options).await + } + + /// Whether the base table has enough fragments to be worth compacting, + /// honoring the config's quiet hours. + #[must_use] + pub fn should_compact(&self, config: &CompactionConfig) -> bool { + self.base.should_compact(config) + } + + /// Current compaction statistics for the base table. + #[must_use] + pub fn compaction_stats(&self) -> CompactionStats { + self.base.compaction_stats() + } + + /// Build a ZoneMap scalar index on `event_id`, the table's key column. + /// Idempotent. Datagen previously had no scalar index, so every point + /// lookup by event id scanned. + pub async fn create_event_id_index(&mut self) -> LanceResult<()> { + self.base.create_key_zonemap_index().await + } + + /// Seal the active memtable. A no-op here in normal operation, since + /// `append` already seals each checkpoint batch before returning. + pub async fn flush(&self) -> LanceResult<()> { + self.base.flush().await } /// Start periodic cleanup after the configured interval. The task keeps @@ -349,17 +347,17 @@ impl DatagenStore { match tokio::time::timeout(pass_timeout, guard.cleanup_own_shard()).await { Ok(Ok(0)) => {} Ok(Ok(reclaimed)) => info!( - shard = %guard.write_shard, + shard = %guard.base.write_shard, reclaimed, "datagen WAL cleanup merged flushed generations" ), Ok(Err(error)) => warn!( - shard = %guard.write_shard, + shard = %guard.base.write_shard, error = %error, "datagen WAL cleanup failed" ), Err(_) => warn!( - shard = %guard.write_shard, + shard = %guard.base.write_shard, timeout_secs = pass_timeout.as_secs(), "datagen WAL cleanup timed out" ), @@ -385,158 +383,6 @@ impl DatagenStore { }); Ok(events) } - - async fn merge_own_shard_if_ready(&mut self, threshold: usize) -> LanceResult { - let object_store = self.dataset.object_store(None).await?; - let branch_location = self.dataset.branch_location(); - let manifest_store = ShardManifestStore::new( - object_store, - &branch_location.path, - self.write_shard, - DEFAULT_MANIFEST_SCAN_BATCH_SIZE, - ); - let Some(manifest) = manifest_store.read_latest().await? else { - return Ok(0); - }; - let pending = manifest.flushed_generations.len(); - if pending == 0 || pending < threshold.max(1) { - return Ok(0); - } - self.merge_own_shard(&manifest_store, &manifest).await?; - Ok(pending) - } - - /// Fold this instance's flushed MemWAL generations into the base table and - /// drain them from the shard manifest. - /// - /// # Concurrency: does not stop the writer - /// - /// A merge only touches *sealed* generations — history — while appends write - /// the active memtable at the WAL tail. They operate on disjoint data, so a - /// merge must not block the write path. - /// - /// That requires *not* calling `claim_epoch`. Claiming bumps `writer_epoch`, - /// which fences every live writer of the shard including our own, forcing a - /// `close()` first and an exclusive lock to hide the window. The epoch is an - /// *ownership* token, not a per-commit token: `commit_update` only rejects a - /// writer whose epoch is **older** than the stored one, and Lance's own - /// flush path reuses one epoch for every manifest commit a writer makes. So - /// this reuses the shard's current epoch and leaves the writer untouched. - /// - /// Lance sanctions concurrent draining explicitly: recovery keys off - /// `replay_after_wal_entry_position` and deliberately does not consult - /// `flushed_generations`, "since an external compactor may legitimately - /// drain that vector back to empty". - /// - /// # Surgical drain, not blanket clear - /// - /// The drain retains everything this call did not merge. `commit_update` - /// re-runs the closure against a freshly-read manifest on every CAS retry, - /// so a *relative* edit composes with a concurrent flush's append, while an - /// absolute `flushed_generations = []` would silently discard a generation - /// that was never merged. `..current.clone()` preserves - /// `current_generation`, `replay_after_wal_entry_position` and - /// `wal_entry_position_last_seen`, which a concurrent flush advances. - async fn merge_own_shard( - &mut self, - manifest_store: &ShardManifestStore, - manifest: &ShardManifest, - ) -> LanceResult<()> { - if manifest.flushed_generations.is_empty() { - return Ok(()); - } - - // NOTE: deliberately no `self.close()` here. Closing the resident writer - // was only needed because the drain below used to `claim_epoch`, which - // fenced our own writer. Reusing the current epoch removes that need, so - // appends continue uninterrupted for the whole merge. - - // Align to the *live* dataset schema rather than the compile-time - // `datagen_log_schema()`: a base table written by an older binary can - // lack columns the current schema has, and merging a batch built against - // the compile-time schema into it fails. This mirrors the rollout store, - // which hit exactly this and now aligns against the live schema. - let merge_schema: Arc = Arc::new(self.dataset.schema().into()); - - let base_uri = self.dataset.uri().trim_end_matches('/').to_string(); - let mut merged_generations = HashSet::new(); - let mut merged_paths = Vec::new(); - let mut batches = Vec::new(); - for flushed in &manifest.flushed_generations { - let generation_uri = format!( - "{}/_mem_wal/{}/{}", - base_uri, self.write_shard, flushed.path - ); - let generation = - Self::load_with_options(&generation_uri, self.storage_options.clone()).await?; - let mut stream = generation.scan().try_into_stream().await?; - while let Some(batch) = stream.try_next().await? { - if batch.num_rows() > 0 { - batches.push(align_batch_to_schema(batch, merge_schema.clone())?); - } - } - merged_generations.insert(flushed.generation); - merged_paths.push(flushed.path.clone()); - } - - if !batches.is_empty() { - let reader = RecordBatchIterator::new( - batches.into_iter().map(Ok::), - merge_schema, - ); - let mut params = WriteParams { - mode: WriteMode::Append, - ..Default::default() - }; - if let Some(options) = &self.storage_options { - params.store_params = Some(ObjectStoreParams { - storage_options_accessor: Some(Arc::new( - StorageOptionsAccessor::with_static_options(options.clone()), - )), - ..Default::default() - }); - } - self.dataset.append(reader, Some(params)).await?; - } - - // Reuse the shard's *current* epoch rather than claiming a new one: - // claiming would fence our own live writer. `commit_update` still fails - // cleanly if a genuinely new writer has claimed the shard meanwhile — - // its epoch would exceed ours — which is the correct outcome. - let epoch = manifest.writer_epoch; - manifest_store - .commit_update(epoch, |current| ShardManifest { - version: current.version + 1, - flushed_generations: current - .flushed_generations - .iter() - .filter(|generation| !merged_generations.contains(&generation.generation)) - .cloned() - .collect(), - ..current.clone() - }) - .await?; - - let object_store = self.dataset.object_store(None).await?; - let branch_path = self.dataset.branch_location().path.clone(); - for path in merged_paths { - let generation_path = branch_path - .clone() - .join("_mem_wal") - .join(self.write_shard.to_string().as_str()) - .join(path.as_str()); - if let Err(error) = object_store.remove_dir_all(generation_path).await { - warn!( - shard = %self.write_shard, - generation_path = %path, - error = %error, - "failed to delete merged datagen WAL generation" - ); - } - } - Ok(()) - } - async fn get_blob_from_dataset( dataset: &Dataset, event_id: &str, @@ -571,7 +417,8 @@ impl DatagenStore { } fn non_blob_columns(&self) -> Vec { - self.dataset + self.base + .dataset .schema() .fields .iter() @@ -580,132 +427,10 @@ impl DatagenStore { .collect() } + /// LSM scanner over the base table unioned with every shard's flushed + /// generations, deduplicating by `event_id`. async fn lsm_scanner(&self) -> LanceResult { - Ok(LsmScanner::new( - Arc::new(self.dataset.clone()), - self.wal_shard_snapshots().await?, - vec!["event_id".to_string()], - )) - } - - async fn wal_shard_snapshots(&self) -> LanceResult> { - let object_store = self.dataset.object_store(None).await?; - let branch_path = self.dataset.branch_location().path.clone(); - let shard_ids = self.dataset.list_mem_wal_latest_shard_ids().await?; - - let snapshots: Vec> = stream::iter(shard_ids) - .map(|shard_id| { - let object_store = object_store.clone(); - let branch_path = branch_path.clone(); - async move { - let manifest_store = ShardManifestStore::new( - object_store, - &branch_path, - shard_id, - DEFAULT_MANIFEST_SCAN_BATCH_SIZE, - ); - let Some(manifest) = manifest_store.read_latest().await? else { - return Ok(None); - }; - let mut snapshot = ShardSnapshot::new(shard_id) - .with_spec_id(manifest.shard_spec_id) - .with_current_generation(manifest.current_generation); - for flushed in manifest.flushed_generations { - snapshot = - snapshot.with_flushed_generation(flushed.generation, flushed.path); - } - Ok::<_, LanceError>(Some(snapshot)) - } - }) - .buffer_unordered(DEFAULT_SHARD_SCAN_CONCURRENCY) - .try_collect() - .await?; - Ok(snapshots.into_iter().flatten().collect()) - } - - async fn ensure_mem_wal(&mut self) -> LanceResult<()> { - if self.mem_wal_index_present().await? { - return Ok(()); - } - match self - .dataset - .initialize_mem_wal() - .unsharded() - .execute() - .await - { - Ok(()) => Ok(()), - Err(error) => { - let uri = self.dataset.uri().to_string(); - self.dataset = Self::load_with_options(&uri, self.storage_options.clone()).await?; - if self.mem_wal_index_present().await? { - Ok(()) - } else { - Err(error) - } - } - } - } - - async fn mem_wal_index_present(&self) -> LanceResult { - let indices = self.dataset.load_indices().await?; - Ok(indices.iter().any(|index| index.name == MEM_WAL_INDEX_NAME)) - } - - fn flushed_generation_uri(&self, shard_id: Uuid, path: &str) -> String { - format!( - "{}/_mem_wal/{shard_id}/{path}", - self.dataset.uri().trim_end_matches('/') - ) - } - - async fn open_flushed_dataset(&self, uri: &str) -> LanceResult { - let mut builder = DatasetBuilder::from_uri(uri).with_session(self.dataset.session()); - if let Some(options) = self.storage_options.clone() { - builder = builder.with_storage_options(options); - } - builder.load().await - } - - async fn load_with_options( - uri: &str, - storage_options: Option>, - ) -> LanceResult { - if let Some(options) = storage_options { - DatasetBuilder::from_uri(uri) - .with_storage_options(options) - .load() - .await - } else { - Dataset::open(uri).await - } - } - - async fn create_with_options( - uri: &str, - storage_options: Option>, - ) -> LanceResult { - let schema = Arc::new(datagen_log_schema()); - let batches = RecordBatchIterator::new( - vec![Ok::(RecordBatch::new_empty( - schema.clone(), - ))] - .into_iter(), - schema, - ); - let mut params = WriteParams { - mode: WriteMode::Create, - ..Default::default() - }; - if let Some(options) = storage_options { - params.store_params = Some(ObjectStoreParams { - storage_options_accessor: Some(Arc::new( - StorageOptionsAccessor::with_static_options(options), - )), - ..Default::default() - }); - } - Dataset::write(batches, uri, Some(params)).await + self.base.lsm_scanner().await } } @@ -1192,44 +917,6 @@ fn invalid_input(message: impl Into) -> LanceError { LanceError::from(ArrowError::InvalidArgumentError(message.into())) } -impl Drop for DatagenStore { - /// Best-effort drain of a still-resident writer's background tasks. - /// - /// [`ShardWriter`] has no `Drop`, so dropping it without `close().await` - /// leaks its background tasks. When a Tokio runtime is available we move the - /// writer into a detached task that closes it; otherwise we can only drop it. - fn drop(&mut self) { - if let Some(writer) = self.write_writer.take() { - if let Ok(handle) = tokio::runtime::Handle::try_current() { - handle.spawn(async move { - // Log rather than discard: a failing close leaks the - // writer's background tasks and, if the memtable was still - // buffered, strands rows that are durable in the WAL but - // never sealed. Silently swallowing that leaves no signal - // anywhere. (Mirrors the rollout store.) - if let Err(error) = writer.close().await { - tracing::warn!( - error = %error, - "failed to close datagen MemWAL writer on drop; \ - its background tasks may leak" - ); - } - }); - } else { - tracing::debug!( - "datagen store dropped outside a Tokio runtime; \ - MemWAL writer background tasks cannot be drained" - ); - } - } - } -} - -fn is_fenced_error(error: &LanceError) -> bool { - let text = error.to_string(); - text.contains("fenced") || text.contains("Fenced") -} - #[cfg(test)] mod tests { use super::*; @@ -1588,6 +1275,50 @@ mod tests { /// reopening lazily. Reusing the shard's current epoch removes the need, so /// the resident writer stays valid across a merge and appends immediately /// after one must still succeed and be readable. + #[test] + fn compaction_folds_merged_fragments_and_preserves_reads() { + // `DatagenStore` had no compaction at all before it moved onto + // `StorageBase`: every WAL merge appends a fragment and datagen seals + // once per checkpoint batch, so fragments grew without bound. + 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 = DatagenStore::open(&uri).await.unwrap(); + + // One cleanup pass appends one fragment, so merge after each append + // to accumulate several -- this is exactly the growth pattern that + // had no bound before compaction existed here. + for seq in 0..4 { + store + .append(&[completed_step(seq, seq as u32)]) + .await + .unwrap(); + store.cleanup_own_shard().await.unwrap(); + } + + let before = store.compaction_stats().total_fragments; + assert!(before > 1, "each merge appends a fragment; got {before}"); + + store + .compact(Some(CompactionConfig { + min_fragments: 2, + ..Default::default() + })) + .await + .unwrap(); + + assert!( + store.compaction_stats().total_fragments < before, + "compaction must reduce the fragment count" + ); + + // Every event survives and is still readable by item. + let events = store.events_for_item("item-1").await.unwrap(); + assert_eq!(events.len(), 4); + }); + } + #[test] fn append_after_merge_reuses_the_writer_without_being_fenced() { let directory = TempDir::new().unwrap();