diff --git a/.github/workflows/rust-test.yml b/.github/workflows/rust-test.yml index f55f5d04..5ff351a0 100644 --- a/.github/workflows/rust-test.yml +++ b/.github/workflows/rust-test.yml @@ -140,8 +140,27 @@ jobs: -p lance-context-core --lib -- registry --ignored - name: Merge coverage reports run: cargo llvm-cov report --lcov --output-path lcov.info + - name: Save coverage report + uses: actions/upload-artifact@v4 + with: + name: coverage-lcov + path: lcov.info + if-no-files-found: error + retention-days: 7 + + coverage-upload: + needs: coverage + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - uses: actions/checkout@v4 + - uses: actions/download-artifact@v4 + with: + name: coverage-lcov - name: Upload coverage to Codecov - uses: codecov/codecov-action@v4 + # v5 uses the maintained shell uploader, including signature checks. + # Keep network upload independently retryable after tests have passed. + uses: codecov/codecov-action@0fb7174895f61a3b6b78fc075e0cd60383518dac # v5 with: files: lcov.info fail_ci_if_error: false diff --git a/crates/lance-context-core/src/generic_store.rs b/crates/lance-context-core/src/generic_store.rs index c1aa98be..340d0b69 100644 --- a/crates/lance-context-core/src/generic_store.rs +++ b/crates/lance-context-core/src/generic_store.rs @@ -19,14 +19,14 @@ //! back on open, so callers pass it once. Reopening with a *different* spec is //! an error rather than a silent reinterpretation of existing data. -use std::collections::{HashMap, HashSet}; +use std::collections::{BTreeSet, HashMap, HashSet}; use std::sync::Arc; -use arrow_array::RecordBatch; +use arrow_array::{RecordBatch, StringArray}; use arrow_schema::{ArrowError, Schema}; use datafusion::prelude::{col, lit, Expr}; -use futures::TryStreamExt; -use lance::dataset::mem_wal::ShardManifestStore; +use futures::{FutureExt, TryStreamExt}; +use lance::dataset::mem_wal::{scanner::ShardSnapshot, ShardManifestStore}; use lance::dataset::optimize::CompactionMetrics; use lance::session::Session; use lance::{Error as LanceError, Result as LanceResult}; @@ -384,11 +384,7 @@ impl GenericStore { } .await; match read { - Err(error) - if error.to_string().contains( - "RowAddrTreeMap::from_sorted_iter called with non-sorted input", - ) => - { + Err(error) if is_wal_pk_batch_error(&error) => { // Lance 9's batched WAL membership probe can reject // a PK sidecar that its point lookup can still read. // Discard partial results and retry all keys through @@ -445,22 +441,131 @@ impl GenericStore { offset: Option, columns: &[String], ) -> LanceResult> { - let refs: Vec<&str> = columns.iter().map(String::as_str).collect(); - let mut scanner = self.base.lsm_scanner().await?.project(&refs)?; if let Some(filter) = filter { - scanner = scanner.filter(filter)?; - } - if limit.is_some() || offset.is_some() { - scanner = scanner.limit( - map_i64_bound("limit", limit)?, - map_i64_bound("offset", offset)?, - )?; + // Keep point/filtered reads on this handle's view. The server's + // refresh-on-miss path refreshes the cached handle itself; reading + // a newer clone here would bypass that cache update. + let refs: Vec<&str> = columns.iter().map(String::as_str).collect(); + let mut scanner = self + .base + .lsm_scanner() + .await? + .project(&refs)? + .filter(filter)?; + if limit.is_some() || offset.is_some() { + scanner = scanner.limit( + map_i64_bound("limit", limit)?, + map_i64_bound("offset", offset)?, + )?; + } + let mut stream = scanner.try_into_stream().await?; + let mut rows = Vec::new(); + while let Some(batch) = stream.try_next().await? { + rows.extend(batch_to_rows(&self.spec, &batch)?); + } + return Ok(rows); } - let mut stream = scanner.try_into_stream().await?; + // Only unfiltered lists use this compatibility recovery. Keep retries + // on one captured base/WAL view, including the fallback. + self.base + .read_consistent(ListSource::All, |dataset, snapshots| async move { + let refs: Vec<&str> = columns.iter().map(String::as_str).collect(); + let mut scanner = StorageBase::lsm_scanner_for_dataset( + &dataset, ID_COLUMN, ListSource::All, snapshots.clone(), + ).project(&refs)?; + if limit.is_some() || offset.is_some() { + scanner = scanner.limit( + map_i64_bound("limit", limit)?, map_i64_bound("offset", offset)?, + )?; + } + let read = async { + let mut stream = scanner.try_into_stream().await?; + let mut rows = Vec::new(); + while let Some(batch) = stream.try_next().await? { + rows.extend(batch_to_rows(&self.spec, &batch)?); + } + Ok::<_, LanceError>(rows) + }.await; + match read { + Err(error) if is_wal_pk_batch_error(&error) => { + tracing::warn!("unfiltered list retrying bounded native point lookups after a WAL PK index error"); + // Never turn an index failure into an unbounded ID cache + // or a long-running scan of a production table. + tokio::time::timeout( + std::time::Duration::from_secs(10), + self.scan_via_points(&dataset, &snapshots, limit, offset, columns), + ).await.map_err(|_| LanceError::io( + "WAL PK index fallback exceeded its 10 second budget" + ))? + } + result => result, + } + }) + // Erase this recovery future so callers do not inherit the full + // Lance scan + retry type (master tests otherwise exceed rustc + // layout query depth before they can run). + .boxed() + .await + } + + /// Narrow compatibility path for Lance 9's broken batched PK sidecars. + /// Physical scans supply only candidate keys; the native point planner + /// still decides visibility/newest-wins, including replay and tombstones. + /// Filtering is deliberately not pushed into these older physical sources. + async fn scan_via_points( + &self, + dataset: &lance::Dataset, + snapshots: &[ShardSnapshot], + limit: Option, + offset: Option, + columns: &[String], + ) -> LanceResult> { + let mut ids = BTreeSet::new(); + let mut id_bytes = 0usize; + let mut scanned = 0usize; + // One dataset at a time; do not allocate a path list proportional to + // pending WAL generations before the key/work budgets can reject it. + collect_scan_keys(dataset, &mut ids, &mut id_bytes, &mut scanned).await?; + for snapshot in snapshots { + for generation in &snapshot.flushed_generations { + let path = self + .base + .flushed_generation_uri(snapshot.shard_id, &generation.path); + let source = StorageBase::load_with_options( + &path, + self.base.storage_options.clone(), + Some(dataset.session()), + ) + .await?; + collect_scan_keys(&source, &mut ids, &mut id_bytes, &mut scanned).await?; + } + } + let refs: Vec<&str> = columns.iter().map(String::as_str).collect(); let mut rows = Vec::new(); - while let Some(batch) = stream.try_next().await? { - rows.extend(batch_to_rows(&self.spec, &batch)?); + let mut skip = offset.unwrap_or(0); + for id in ids { + if limit.is_some_and(|n| rows.len() >= n) { + break; + } + let scanner = StorageBase::lsm_scanner_for_dataset( + dataset, + ID_COLUMN, + ListSource::All, + snapshots.to_vec(), + ) + .project(&refs)? + .filter_expr(col(ID_COLUMN).eq(lit(id))); + let mut stream = scanner.try_into_stream().await?; + while let Some(batch) = stream.try_next().await? { + for row in batch_to_rows(&self.spec, &batch)? { + if skip > 0 { + skip -= 1; + } else { + rows.push(row); + } + } + } } Ok(rows) } @@ -679,6 +784,48 @@ impl GenericStore { } } +/// Collect only candidate keys; native point reads resolve their visibility. +async fn collect_scan_keys( + source: &lance::Dataset, + ids: &mut BTreeSet, + id_bytes: &mut usize, + scanned: &mut usize, +) -> LanceResult<()> { + let mut scanner = source.scan(); + scanner.project(&[ID_COLUMN])?; + scanner.batch_size(256); + let mut stream = scanner.try_into_stream().await?; + while let Some(batch) = stream.try_next().await? { + *scanned += batch.num_rows(); + if *scanned > 16 * MAX_BATCH_GET_IDS { + return Err(LanceError::io("WAL PK index fallback scan budget exceeded")); + } + let keys = batch + .column(0) + .as_any() + .downcast_ref::() + .ok_or_else(|| LanceError::io("generic ID must be Utf8"))?; + for key in keys.iter().flatten() { + if !ids.contains(key) { + if ids.len() >= MAX_BATCH_GET_IDS + || id_bytes.saturating_add(key.len()) > 1024 * 1024 + { + return Err(LanceError::io("WAL PK index fallback key budget exceeded")); + } + *id_bytes += key.len(); + ids.insert(key.to_owned()); + } + } + } + Ok(()) +} + +fn is_wal_pk_batch_error(error: &LanceError) -> bool { + error + .to_string() + .contains("RowAddrTreeMap::from_sorted_iter called with non-sorted input") +} + #[cfg(test)] // `GenericStore` owns a `Dataset` and is not `Debug`, so `expect_err` (which // formats the Ok value) is unavailable; `.err().expect(..)` is the alternative. @@ -720,6 +867,83 @@ mod tests { } } + #[tokio::test] + async fn point_scan_fallback_preserves_newest_rows_projection_and_paging() { + let dir = TempDir::new().unwrap(); + let mut store = GenericStore::open(dir.path().to_str().unwrap(), spec(), sealing()) + .await + .unwrap(); + store + .add(&[ + row(json!({"id":"a", "score":1.0})), + row(json!({"id":"b", "score":2.0})), + row(json!({"id":"c", "score":3.0})), + ]) + .await + .unwrap(); + store.cleanup_wal().await.unwrap(); + store + .add(&[row(json!({"id":"b", "score":20.0}))]) + .await + .unwrap(); + store + .add(&[ + row(json!({"id":"b", "score":21.0})), + row(json!({"id":"d", "score":4.0})), + ]) + .await + .unwrap(); + let dataset = store.base.dataset.clone(); + let snapshots = store.base.wal_shard_snapshots().await.unwrap(); + let columns = vec!["id".to_string(), "score".to_string()]; + let rows = store + .scan_via_points(&dataset, &snapshots, None, None, &columns) + .await + .unwrap(); + assert_eq!( + rows.iter() + .map(|r| r["id"].as_str().unwrap()) + .collect::>(), + ["a", "b", "c", "d"] + ); + assert_eq!(rows[1]["score"], json!(21.0)); + assert!(rows.iter().all(|r| r.len() == 2)); + let page = store + .scan_via_points(&dataset, &snapshots, Some(2), Some(1), &columns) + .await + .unwrap(); + assert_eq!(page, rows[1..3]); + assert!(store + .scan_via_points(&dataset, &snapshots, Some(0), None, &columns) + .await + .unwrap() + .is_empty()); + store.close().await.unwrap(); + } + + #[tokio::test] + async fn point_scan_fallback_refuses_an_unbounded_key_cache() { + let dir = TempDir::new().unwrap(); + let mut store = GenericStore::open(dir.path().to_str().unwrap(), spec(), sealing()) + .await + .unwrap(); + let rows: Vec<_> = (0..=MAX_BATCH_GET_IDS) + .map(|i| row(json!({"id":format!("key-{i}"),"score":1.0}))) + .collect(); + store.add(&rows).await.unwrap(); + let dataset = store.base.dataset.clone(); + let snapshots = store.base.wal_shard_snapshots().await.unwrap(); + let error = store + .scan_via_points(&dataset, &snapshots, Some(1), None, &["id".into()]) + .await + .unwrap_err(); + assert!( + error.to_string().contains("fallback key budget exceeded"), + "{error}" + ); + store.close().await.unwrap(); + } + #[tokio::test] async fn batch_get_matches_visible_rows_across_generations_and_merge() { let dir = TempDir::new().unwrap(); diff --git a/crates/lance-context-core/src/rollout_append.rs b/crates/lance-context-core/src/rollout_append.rs index ada9c34d..4c5c8162 100644 --- a/crates/lance-context-core/src/rollout_append.rs +++ b/crates/lance-context-core/src/rollout_append.rs @@ -18,6 +18,7 @@ use lance::dataset::{ use lance::index::DatasetIndexExt; use lance_index::mem_wal::{MemWalIndexDetails, MergedGeneration, MEM_WAL_INDEX_NAME}; use lance_table::format::{pb, Fragment}; +use lance_table::system_index::mem_wal::ShardManifest; use serde::{Deserialize, Serialize}; use uuid::Uuid; @@ -27,6 +28,8 @@ use lance::Result; const CUTOVER_PREFIX: &str = "lance-context.rollout-append.cutover."; const WRITER_EPOCH_PREFIX: &str = "lance-context.rollout-append.writer-epoch."; +const REPLAY_RANGES_PREFIX: &str = "lance-context.rollout-append.replay-ranges."; + pub const MAX_PLAN_GENERATIONS: usize = 256; pub const MAX_STAGE_BYTES: usize = 256 * 1024 * 1024; @@ -36,6 +39,190 @@ pub struct Generation { pub path: String, } +/// Precise replay generations, with the old broad cutoff retained for rolling +/// upgrades. Old coordinators/workers may still use that conservative prefix. +/// Only trust this state while its epoch AND cutoff match the durable old keys. +#[derive(Clone, Debug, Serialize, Deserialize)] +struct ReplayRanges { + epoch: u64, + through: u64, + manifest_version: u64, + ranges: Vec<(u64, u64)>, +} + +impl ReplayRanges { + fn load(dataset: &Dataset, shard: Uuid) -> Result> { + let Some(json) = dataset + .metadata() + .get(&format!("{REPLAY_RANGES_PREFIX}{shard}")) + else { + return Ok(None); + }; + let state: Self = + serde_json::from_str(json).map_err(|_| Error::io("invalid rollout replay ranges"))?; + if state.manifest_version == 0 + || state.ranges.len() > MAX_PLAN_GENERATIONS + || state + .ranges + .iter() + .any(|(lo, hi)| lo > hi || *hi > state.through) + || state.ranges.windows(2).any(|r| r[0].1 >= r[1].0) + { + return Err(Error::io("invalid rollout replay ranges")); + } + if dataset + .metadata() + .get(&format!("{WRITER_EPOCH_PREFIX}{shard}")) + != Some(&state.epoch.to_string()) + || dataset.metadata().get(&format!("{CUTOVER_PREFIX}{shard}")) + != Some(&state.through.to_string()) + { + // An older coordinator advanced the boundary. Import its entire + // prefix rather than trusting an optimization it did not update. + return Ok(None); + } + Ok(Some(state)) + } + + fn contains(&self, generation: u64) -> bool { + self.ranges + .iter() + .any(|(lo, hi)| (*lo..=*hi).contains(&generation)) + } + + fn normalize(&mut self, watermark: Option) { + self.ranges + .retain(|(_, hi)| watermark.is_none_or(|w| *hi > w)); + self.ranges.sort_unstable(); + let mut merged: Vec<(u64, u64)> = Vec::new(); + for (lo, hi) in self.ranges.drain(..) { + if let Some(last) = merged.last_mut() { + if lo <= last.1.saturating_add(1) { + last.1 = last.1.max(hi); + continue; + } + } + merged.push((lo, hi)); + } + // Bound metadata even if a writer repeatedly restarts without draining. + // Coalescing adds checks, never removes a necessary check. + if merged.len() > MAX_PLAN_GENERATIONS { + let excess = merged.len() - MAX_PLAN_GENERATIONS; + merged[excess].0 = merged[0].0; + merged.drain(..excess); + } + self.ranges = merged; + } +} + +/// Find the first manifest of each unseen writer epoch. Claiming an epoch +/// preserves current_generation; replay fills exactly that first memtable. +/// Immutable generations older than the claim cannot acquire replayed rows. +/// Binary search bounds this to 64 metadata GETs per transition, with a safe +/// broad-prefix fallback when history is missing or too many epochs elapsed. +async fn replay_generations( + store: &ShardManifestStore, + previous: &ReplayRanges, + current: &ShardManifest, +) -> Result> { + if previous.manifest_version >= current.version || previous.epoch >= current.writer_epoch { + return Err(Error::io("rollout replay history moved backwards")); + } + let mut version = previous.manifest_version; + let mut epoch = previous.epoch; + let mut reads = 0; + let mut ranges = Vec::new(); + while epoch < current.writer_epoch { + let mut low = version + 1; + let mut high = current.version; + let mut first = current.clone(); + while low < high { + if reads >= 64 { + return Err(Error::io("rollout replay history read budget exceeded")); + } + let mid = low + (high - low) / 2; + let manifest = store.read_version(mid).await?; + reads += 1; + crate::merge_write_scope::checkpoint(); + if manifest.writer_epoch > epoch { + high = mid; + first = manifest; + } else { + low = mid + 1; + } + } + if first.version != low + || first.writer_epoch <= epoch + || first.writer_epoch > current.writer_epoch + || first.current_generation > current.current_generation + { + return Err(Error::io("inconsistent rollout replay history")); + } + ranges.push((first.current_generation, first.current_generation)); + version = first.version; + epoch = first.writer_epoch; + } + Ok(ranges) +} + +/// Upgrade the old ever-growing cutoff using its first epoch-aware boundary. +/// This is a one-time metadata-only search: historical data files are never read. +/// Missing/pruned history leaves the caller on the conservative prefix. +async fn import_replay_ranges( + dataset: &Dataset, + shard: Uuid, + current: &ShardManifest, + through: u64, +) -> Result { + let epoch_key = format!("{WRITER_EPOCH_PREFIX}{shard}"); + let cutoff_key = format!("{CUTOVER_PREFIX}{shard}"); + let mut low = 1; + let mut high = dataset.version().version; + let mut first = dataset.clone(); + let mut reads = 0; + while low < high { + if reads >= 64 { + return Err(Error::io("rollout cutover history read budget exceeded")); + } + let mid = low + (high - low) / 2; + let snapshot = dataset.checkout_version(mid).await?; + reads += 1; + crate::merge_write_scope::checkpoint(); + if snapshot.metadata().contains_key(&epoch_key) { + high = mid; + first = snapshot; + } else { + low = mid + 1; + } + } + let parse = |key: &str| -> Result { + first + .metadata() + .get(key) + .and_then(|s| s.parse().ok()) + .ok_or_else(|| Error::io("missing initial rollout replay boundary")) + }; + let epoch = parse(&epoch_key)?; + let initial_cutoff = parse(&cutoff_key)?; + if first.version().version != low || epoch > current.writer_epoch || initial_cutoff > through { + return Err(Error::io("inconsistent rollout cutover history")); + } + let mut state = ReplayRanges { + epoch, + through, + // Start at the beginning of the immutable shard history. Versions in + // the original epoch are skipped by replay_generations' binary search. + manifest_version: 1, + ranges: vec![(0, initial_cutoff)], + }; + if epoch < current.writer_epoch { + state.ranges.extend( + replay_generations(&shard_store(dataset, shard).await?, &state, current).await?, + ); + } + Ok(state) +} + /// A bounded, immutable prefix of one shard. URIs/credentials are never supplied /// by an RPC caller: each process resolves the rollout name in its own registry. #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] @@ -224,6 +411,7 @@ impl AppendCoordinator { let mut pending = Vec::new(); let mut cutovers = HashMap::new(); let mut reclaimed = 0; + let marks = watermarks(&self.dataset).await?; for shard in shard_ids { reclaimed += reconcile_shard(&self.dataset, shard).await?; let Some(manifest) = shard_store(&self.dataset, shard) @@ -262,12 +450,9 @@ impl AppendCoordinator { if previous_epoch.is_some_and(|epoch| epoch > manifest.writer_epoch) { return Err(Error::io("rollout writer epoch moved backwards")); } - let cutoff = if previous_epoch != Some(manifest.writer_epoch) { - // A restarted/fenced writer can replay already-published WAL - // bytes into a NEW generation when its flush cursor lagged. - // Generation watermarks alone cannot identify those rows. - // Protect the active generation too: planning may observe an - // epoch claim before the recovered memtable is flushed. + let prior_ranges = ReplayRanges::load(&self.dataset, shard)?; + let cutoff = if previous_epoch != Some(manifest.writer_epoch) || prior_ranges.is_none() + { let high = manifest .flushed_generations .iter() @@ -276,6 +461,57 @@ impl AppendCoordinator { .unwrap_or(0) .max(manifest.current_generation) .max(previous_cutoff.unwrap_or(0)); + let mut state = if let Some(mut state) = prior_ranges { + match replay_generations( + &shard_store(&self.dataset, shard).await?, + &state, + &manifest, + ) + .await + { + Ok(ranges) => state.ranges.extend(ranges), + Err(error) => { + tracing::warn!(%shard, %error, "using conservative rollout replay boundary"); + state.ranges.push((0, high)); + } + } + state + } else { + let conservative = ReplayRanges { + epoch: manifest.writer_epoch, + through: high, + manifest_version: manifest.version, + ranges: vec![( + 0, + if previous_epoch == Some(manifest.writer_epoch) { + previous_cutoff.unwrap_or(high) + } else { + high + }, + )], + }; + if previous_epoch.is_some() { + match import_replay_ranges(&self.dataset, shard, &manifest, high).await { + Ok(state) => state, + Err(error) => { + tracing::warn!(%shard, %error, "preserving legacy rollout replay prefix"); + conservative + } + } + } else { + // Initial cutover: the old commit/drain ambiguity must + // be checked once. Future restarts add only replayed generations. + conservative + } + }; + state.epoch = manifest.writer_epoch; + state.through = high; + state.manifest_version = manifest.version; + state.normalize(marks.get(&shard).copied()); + cutovers.insert( + format!("{REPLAY_RANGES_PREFIX}{shard}"), + serde_json::to_string(&state)?, + ); cutovers.insert(key, high.to_string()); cutovers.insert(epoch_key, manifest.writer_epoch.to_string()); high @@ -477,6 +713,7 @@ pub async fn stage( "staging plan does not match durable cutover", )); } + let replay_ranges = ReplayRanges::load(&dataset, plan.shard)?; let mut reservation = budget.reserve(plan.max_bytes.min(budget.limit())).await; let mut batches = Vec::new(); let mut bytes = 0usize; @@ -505,7 +742,12 @@ pub async fn stage( current.push(batch); crate::merge_write_scope::checkpoint(); } - batches.extend(current); + let needs_lookup = replay_ranges + .as_ref() + .map_or(generation.number <= plan.legacy_through, |ranges| { + ranges.contains(generation.number) + }); + batches.extend(current.into_iter().map(|batch| (batch, needs_lookup))); completed += 1; if bytes >= plan.max_bytes { break; @@ -514,9 +756,13 @@ pub async fn stage( // Rollout IDs are immutable. Dedup identical retries within this batch; // during migration also exclude IDs already published by a legacy merge. let mut seen = HashSet::new(); - if plan.generations[0].number <= plan.legacy_through { + let lookup_started = std::time::Instant::now(); + let mut lookup_ids = 0; + if batches.iter().any(|(_, needs_lookup)| *needs_lookup) + && !dataset.manifest().fragments.is_empty() + { let mut ids = HashSet::new(); - for batch in &batches { + for (batch, _) in batches.iter().filter(|(_, needs_lookup)| *needs_lookup) { let column = batch .column_by_name("id") .unwrap() @@ -526,6 +772,7 @@ pub async fn stage( ids.extend(column.iter().flatten().map(str::to_owned)); } let ids: Vec<_> = ids.into_iter().collect(); + lookup_ids = ids.len(); for chunk in ids.chunks(1024) { let mut scanner = dataset.scan(); scanner.project(&["id"])?; @@ -544,8 +791,15 @@ pub async fn stage( } } } + tracing::info!( + shard = %plan.shard, + completed_generations = completed, + legacy_lookup_ids = lookup_ids, + legacy_lookup_seconds = lookup_started.elapsed().as_secs_f64(), + "rollout append legacy lookup completed" + ); let mut output = Vec::new(); - for batch in batches { + for (batch, _) in batches { let ids = batch .column_by_name("id") .unwrap() @@ -826,6 +1080,398 @@ mod tests { restarted.close().await.unwrap(); } + #[tokio::test] + async fn restart_does_not_make_clean_backlog_scan_base_ids() { + let dir = tempfile::tempdir().unwrap(); + let uri = dir.path().to_str().unwrap(); + let mut a = writer(uri, "a").await; + let mut coordinator = AppendCoordinator::open(uri, None).await.unwrap(); + // Drain the one-time legacy boundary, including its active generation. + for id in ["seed", "tail"] { + put(&a, id, 4096).await; + let (plans, _) = coordinator + .plan(&["a".into()], 64, 1024 * 1024) + .await + .unwrap(); + let part = stage(uri, plans[0].clone(), budget(), None).await.unwrap(); + coordinator.commit(vec![part]).await.unwrap(); + } + for i in 0..8 { + put(&a, &format!("clean-{i}"), 4096).await; + } + a.close().await.unwrap(); + let shard = derive_shard_id(Some("a")); + let store = shard_store(&coordinator.dataset, shard).await.unwrap(); + let before = store.read_latest().await.unwrap().unwrap(); + // Real replay: lag the durable cursor, so reopening replays "tail" and + // the eight still-pending rows into the first generation of the epoch. + store + .commit_update(before.writer_epoch, |m| { + let mut next = m.clone(); + next.version += 1; + next.replay_after_wal_entry_position = 1; + next + }) + .await + .unwrap(); + let mut restarted = writer(uri, "a").await; + restarted + .add(&[artifact_record("new", &[42; 4096])]) + .await + .unwrap(); + // Observe the new epoch BEFORE replay's first flush. + let (plans, _) = coordinator + .plan(&["a".into()], 8, 1024 * 1024) + .await + .unwrap(); + assert_eq!(plans[0].generations.len(), 8); + let state = ReplayRanges::load(&coordinator.dataset, shard) + .unwrap() + .unwrap(); + assert!(plans[0] + .generations + .iter() + .all(|g| !state.contains(g.number))); + assert!(state.contains(before.current_generation)); + + // Make base payload files unavailable. Staging the clean backlog must + // still work: metadata/watermarks suffice, not even an ID scan is needed. + let mut hidden = Vec::new(); + for fragment in coordinator.dataset.manifest().fragments.iter() { + for file in &fragment.files { + let path = dir.path().join("data").join(&file.path); + let saved = path.with_extension("saved"); + std::fs::rename(&path, &saved).unwrap(); + hidden.push((path, saved)); + } + } + let part = stage(uri, plans[0].clone(), budget(), None).await.unwrap(); + assert_eq!(part.rows, 8); + for (path, saved) in hidden { + std::fs::rename(saved, path).unwrap(); + } + coordinator.commit(vec![part]).await.unwrap(); + restarted.flush().await.unwrap(); + put(&restarted, "steady", 4096).await; + let (plans, _) = coordinator + .plan(&["a".into()], 64, 1024 * 1024) + .await + .unwrap(); + assert_eq!(plans[0].generations.len(), 2); + let state = ReplayRanges::load(&coordinator.dataset, shard) + .unwrap() + .unwrap(); + assert!(state.contains(plans[0].generations[0].number)); + assert!(!state.contains(plans[0].generations[1].number)); + let part = stage(uri, plans[0].clone(), budget(), None).await.unwrap(); + assert_eq!(part.rows, 2, "exclude committed tail AND replayed backlog"); + coordinator.commit(vec![part.clone()]).await.unwrap(); + assert_eq!(coordinator.commit(vec![part]).await.unwrap(), 0); + assert_eq!(rows(uri).await, 12); + let reader = RolloutStore::open_existing_with_options(uri, RolloutStoreOptions::default()) + .await + .unwrap(); + for id in ["seed", "tail", "clean-0", "clean-7", "new", "steady"] { + assert_eq!(reader.get_blob(id).await.unwrap(), Some(vec![42; 4096])); + } + restarted.close().await.unwrap(); + } + + #[tokio::test] + async fn multiple_unobserved_epochs_protect_only_their_replay_generations() { + let dir = tempfile::tempdir().unwrap(); + let uri = dir.path().to_str().unwrap(); + let mut a = writer(uri, "a").await; + a.add(&[artifact_record("seed", &[42; 4096])]) + .await + .unwrap(); + let mut coordinator = AppendCoordinator::open(uri, None).await.unwrap(); + coordinator + .plan(&["a".into()], 64, 1024 * 1024) + .await + .unwrap(); + let shard = derive_shard_id(Some("a")); + let store = shard_store(&coordinator.dataset, shard).await.unwrap(); + let mut expected = vec![1]; + let mut clean = Vec::new(); + for epoch in 0..3 { + put(&a, &format!("first-{epoch}"), 4096).await; + let current = store.read_latest().await.unwrap().unwrap(); + clean.push(current.current_generation); + put(&a, &format!("steady-{epoch}"), 4096).await; + a.close().await.unwrap(); + expected.push( + store + .read_latest() + .await + .unwrap() + .unwrap() + .current_generation, + ); + a = writer(uri, "a").await; + a.add(&[artifact_record(&format!("opened-{epoch}"), &[42; 4096])]) + .await + .unwrap(); + } + coordinator + .plan(&["a".into()], 64, 1024 * 1024) + .await + .unwrap(); + let state = ReplayRanges::load(&coordinator.dataset, shard) + .unwrap() + .unwrap(); + assert!(expected.into_iter().all(|g| state.contains(g))); + assert!(clean.into_iter().all(|g| !state.contains(g))); + a.close().await.unwrap(); + } + + #[tokio::test] + async fn old_coordinator_boundary_invalidates_precise_ranges() { + let dir = tempfile::tempdir().unwrap(); + let uri = dir.path().to_str().unwrap(); + let mut a = writer(uri, "a").await; + a.add(&[artifact_record("seed", &[42; 4096])]) + .await + .unwrap(); + let mut coordinator = AppendCoordinator::open(uri, None).await.unwrap(); + coordinator + .plan(&["a".into()], 64, 1024 * 1024) + .await + .unwrap(); + let shard = derive_shard_id(Some("a")); + assert!(ReplayRanges::load(&coordinator.dataset, shard) + .unwrap() + .is_some()); + a.close().await.unwrap(); + a = writer(uri, "a").await; + put(&a, "new", 4096).await; + let current = shard_store(&coordinator.dataset, shard) + .await + .unwrap() + .read_latest() + .await + .unwrap() + .unwrap(); + // Simulate an older coordinator, which updates only the legacy keys. + coordinator + .dataset + .update_metadata([ + ( + format!("{WRITER_EPOCH_PREFIX}{shard}"), + current.writer_epoch.to_string(), + ), + ( + format!("{CUTOVER_PREFIX}{shard}"), + current.current_generation.to_string(), + ), + ]) + .await + .unwrap(); + assert!(ReplayRanges::load(&coordinator.dataset, shard) + .unwrap() + .is_none()); + coordinator + .plan(&["a".into()], 64, 1024 * 1024) + .await + .unwrap(); + let state = ReplayRanges::load(&coordinator.dataset, shard) + .unwrap() + .unwrap(); + assert!(state.contains(1)); + assert!(state.contains(current.flushed_generations[0].generation)); + assert!(!state.contains(current.current_generation)); + a.close().await.unwrap(); + } + + #[tokio::test] + async fn missing_epoch_history_preserves_conservative_checks() { + let dir = tempfile::tempdir().unwrap(); + let uri = dir.path().to_str().unwrap(); + let mut a = writer(uri, "a").await; + a.add(&[artifact_record("seed", &[42; 4096])]) + .await + .unwrap(); + let mut coordinator = AppendCoordinator::open(uri, None).await.unwrap(); + coordinator + .plan(&["a".into()], 64, 1024 * 1024) + .await + .unwrap(); + for i in 0..4 { + put(&a, &format!("row-{i}"), 4096).await; + } + a.close().await.unwrap(); + let mut restarted = writer(uri, "a").await; + restarted + .add(&[artifact_record("next", &[42; 4096])]) + .await + .unwrap(); + let shard = derive_shard_id(Some("a")); + let store = shard_store(&coordinator.dataset, shard).await.unwrap(); + let current = store.read_latest().await.unwrap().unwrap(); + // Remove only old metadata, retaining the exact latest version so that + // normal WAL discovery and publication are still possible. + let manifest_dir = dir + .path() + .join("_mem_wal") + .join(shard.to_string()) + .join("manifest"); + let latest = format!("{:064b}.binpb", current.version.reverse_bits()); + for entry in std::fs::read_dir(manifest_dir).unwrap() { + let entry = entry.unwrap(); + if entry.file_name().to_string_lossy().ends_with(".binpb") + && entry.file_name() != latest.as_str() + { + std::fs::remove_file(entry.path()).unwrap(); + } + } + let (plans, _) = coordinator + .plan(&["a".into()], 64, 1024 * 1024) + .await + .unwrap(); + let state = ReplayRanges::load(&coordinator.dataset, shard) + .unwrap() + .unwrap(); + assert!(plans[0] + .generations + .iter() + .all(|g| state.contains(g.number))); + let part = stage(uri, plans[0].clone(), budget(), None).await.unwrap(); + coordinator.commit(vec![part]).await.unwrap(); + assert_eq!(rows(uri).await, 5); + restarted.close().await.unwrap(); + } + + #[tokio::test] + #[ignore = "opt-in populated-base replay lookup benchmark"] + async fn benchmark_replay_lookup_elision() { + let dir = tempfile::tempdir().unwrap(); + let uri = dir.path().to_str().unwrap(); + let mut a = writer(uri, "a").await; + let mut coordinator = AppendCoordinator::open(uri, None).await.unwrap(); + // Many base fragments, then a clean backlog and a writer restart. + for i in 0..64 { + let records: Vec<_> = (0..16) + .map(|row| artifact_record(&format!("base-{i}-{row}"), &[42; 4096])) + .collect(); + a.add(&records).await.unwrap(); + a.flush().await.unwrap(); + let (plans, _) = coordinator + .plan(&["a".into()], 64, 1024 * 1024) + .await + .unwrap(); + let part = stage(uri, plans[0].clone(), budget(), None).await.unwrap(); + coordinator.commit(vec![part]).await.unwrap(); + } + for i in 0..32 { + put(&a, &format!("pending-{i}"), 4096).await; + } + a.close().await.unwrap(); + let mut restarted = writer(uri, "a").await; + restarted + .add(&[artifact_record("next", &[42; 4096])]) + .await + .unwrap(); + let shard = derive_shard_id(Some("a")); + let current = shard_store(&coordinator.dataset, shard) + .await + .unwrap() + .read_latest() + .await + .unwrap() + .unwrap(); + // The old coordinator expands its cutoff over the entire clean backlog. + coordinator + .dataset + .update_metadata([ + ( + format!("{WRITER_EPOCH_PREFIX}{shard}"), + current.writer_epoch.to_string(), + ), + ( + format!("{CUTOVER_PREFIX}{shard}"), + current.current_generation.to_string(), + ), + ]) + .await + .unwrap(); + assert!(ReplayRanges::load(&coordinator.dataset, shard) + .unwrap() + .is_none()); + let old_plan = AppendPlan { + base_version: coordinator.version(), + shard, + generations: current + .flushed_generations + .iter() + .map(|g| Generation { + number: g.generation, + path: g.path.clone(), + }) + .collect(), + max_bytes: 1024 * 1024, + legacy_through: current.current_generation, + }; + let (plans, _) = coordinator + .plan(&["a".into()], 64, 1024 * 1024) + .await + .unwrap(); + let new_plan = plans[0].clone(); + let state = ReplayRanges::load(&coordinator.dataset, shard) + .unwrap() + .unwrap(); + assert!(new_plan + .generations + .iter() + .all(|g| !state.contains(g.number))); + for old in [true, false, false, true] { + let start = std::time::Instant::now(); + let part = stage( + uri, + if old { + old_plan.clone() + } else { + new_plan.clone() + }, + budget(), + None, + ) + .await + .unwrap(); + assert_eq!(part.rows, 32); + assert_eq!(part.completed, 32); + println!("replay_lookup_benchmark mode={} base_fragments=64 generations=32 rows={} seconds={:.6}", if old { "legacy" } else { "precise" }, part.rows, start.elapsed().as_secs_f64()); + } + let part = stage(uri, new_plan, budget(), None).await.unwrap(); + coordinator.commit(vec![part.clone()]).await.unwrap(); + assert_eq!(coordinator.commit(vec![part]).await.unwrap(), 0); + let dataset = Dataset::open(uri).await.unwrap(); + let mut scanner = dataset.scan(); + scanner.project(&["id"]).unwrap(); + let mut stream = scanner.try_into_stream().await.unwrap(); + let mut ids = HashSet::new(); + let mut count = 0; + while let Some(batch) = stream.try_next().await.unwrap() { + let column = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + count += column.len(); + ids.extend(column.iter().flatten().map(str::to_owned)); + } + assert_eq!(count, 1056); + assert_eq!(ids.len(), count); + for i in 0..64 { + for row in 0..16 { + assert!(ids.contains(&format!("base-{i}-{row}"))); + } + } + for i in 0..32 { + assert!(ids.contains(&format!("pending-{i}"))); + } + restarted.close().await.unwrap(); + } + #[derive(Debug)] struct FailDrain(AtomicBool); impl CommitAuthorizer for FailDrain { diff --git a/docs/rollout-parallel-append.md b/docs/rollout-parallel-append.md index aff35baf..691735b2 100644 --- a/docs/rollout-parallel-append.md +++ b/docs/rollout-parallel-append.md @@ -85,15 +85,30 @@ While staging this prefix, probe the base for IDs in chunks of 1024, projecting until the historical prefix is consumed; it does not read old payloads or delete base rows. Generations beyond the cutover do not require this migration lookup. -The coordinator also persists the last observed writer epoch per shard. On a -writer restart or fence, it extends the ID-check boundary through the current -active generation as well as every flushed generation already visible. A WAL -replay cursor can lag a flush: already-published rows can then be replayed into a -new generation, so generation watermarks alone are insufficient. Planning before -the recovered memtable flushes must protect that active generation too. The -boundary stays fixed while the writer epoch is unchanged; subsequent generations -return to append without base ID probes. Upgrading metadata that lacks a writer -epoch performs this bounded check once before using the fast path. +The coordinator persists the last observed writer epoch and precise replay ranges +per shard. Lance 9.0.0 replays WAL into the first memtable opened by a new writer; +the epoch-claim manifest identifies that generation. A bounded binary search of +immutable shard manifests locates each unseen epoch's first generation. Only +those generations need replay ID checks. Previously flushed, immutable backlog +and later generations in the same epoch keep the lookup-free append path, even +when they share a staging batch with replayed data. The first replay generation +is protected even if planning happens before it flushes. + +The original legacy cutover prefix still needs ID checks until consumed. Upgrading +an old, repeatedly expanded boundary performs a one-time binary search of base +metadata for the first epoch-aware cutover, then reconstructs subsequent replay +generations from shard metadata. Both searches are limited to 64 metadata reads; +missing/pruned history falls back to the conservative prefix and does not prevent +merge. Ranges are bounded to 256 per shard, coalescing conservatively if needed. +No full-table ID cache, new index, or payload scan is introduced by this planning. + +The old cutoff and epoch keys remain present for rolling upgrades. Old workers +may perform extra checks; new workers use precise ranges only when both old keys +match their recorded values. If an old coordinator advances the keys without +updating the ranges, staging falls back to the old prefix until a new coordinator +reconstructs it. The append RPC and atomic data/watermark publication are unchanged. +The `rollout append legacy lookup completed` log reports `legacy_lookup_ids` and +`legacy_lookup_seconds`; clean batches report zero candidate IDs. Merge retries are idempotent through the atomic generation watermark. This is not an arbitrary upsert or ingest-deduplication API: replaying the same ID into different @@ -121,7 +136,10 @@ maintenance runs; point queries must retain their existing uncovered-fragment sc Tests exercise parallel file preparation and a single combined commit, live writes arriving between stage and publication, duplicate results, partial stale prefixes, byte bounds, legacy append-without-drain migration, restart after base publication, -and fallback through the existing merge implementation. Master HTTP tests require +and fallback through the existing merge implementation. Replay tests also cover +multiple unobserved epochs, old-coordinator boundary updates, missing shard history, +and a clean backlog that stages successfully with every base data file unavailable. +Master HTTP tests require two workers to reach a barrier together, catching accidental serial fan-out. A local ARM64 debug-build benchmark (Lance 9.0.0, one process, local filesystem) @@ -149,3 +167,13 @@ ROLLOUT_APPEND_BENCH=1 cargo test -p lance-context-core --lib \ `ROLLOUT_APPEND_BENCH_ROOT` optionally selects a scratch object-store prefix. Each fixture creates a new UUID subdirectory; it never scans an existing table. + +A separate local populated-base comparison can be run with: + +```bash +cargo test -p lance-context-core --lib benchmark_replay_lookup_elision -- --ignored --nocapture +``` + +It stages the same 32 pending generations against 64 populated base fragments in +legacy/precise/precise/legacy order, then verifies all 1,056 final IDs and duplicate +commit idempotence. Fixture data is confined to a new local temporary directory.