diff --git a/crates/lance-context-core/src/lib.rs b/crates/lance-context-core/src/lib.rs index 7ccfaa2..4435f3e 100644 --- a/crates/lance-context-core/src/lib.rs +++ b/crates/lance-context-core/src/lib.rs @@ -102,3 +102,5 @@ pub use lance::Error as LanceError; // Re-export the Lance session type so the server can build one shared, // capacity-bounded cache session across all resident rollout stores. pub use lance::session::Session; + +pub use store_base::PreparedCompaction; diff --git a/crates/lance-context-core/src/rollout_append.rs b/crates/lance-context-core/src/rollout_append.rs index 4c5c816..13b02ca 100644 --- a/crates/lance-context-core/src/rollout_append.rs +++ b/crates/lance-context-core/src/rollout_append.rs @@ -934,6 +934,166 @@ mod tests { MergeMemoryBudget::new(8 * 1024 * 1024) } + #[derive(Debug)] + struct NoPreparationCommit; + impl crate::merge_write_scope::CommitAuthorizer for NoPreparationCommit { + fn authorize<'a>( + &'a self, + _: &'a str, + _: u64, + ) -> std::pin::Pin> + Send + 'a>> + { + Box::pin(async { Err(Error::io("unexpected preparation manifest commit")) }) + } + } + + async fn prepared_compaction_fixture(uri: &str) -> (RolloutStore, AppendCoordinator) { + let a = writer(uri, "a").await; + let mut coordinator = AppendCoordinator::open(uri, None).await.unwrap(); + for i in 0..4 { + put(&a, &format!("seed-{i}"), 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(); + } + (a, coordinator) + } + + async fn prepare_without_commit(uri: &str) -> crate::PreparedCompaction { + crate::merge_write_scope::MergeWriteScope::with_pinned_authorizer(Arc::new( + NoPreparationCommit, + )) + .run(async { + let store = + RolloutStore::open_existing_with_options(uri, RolloutStoreOptions::default()) + .await + .unwrap(); + store + .prepare_compaction(crate::CompactionConfig { + target_rows_per_fragment: 1024, + max_source_fragments: Some(128), + ..Default::default() + }) + .await + .unwrap() + }) + .await + } + + #[tokio::test] + async fn prepared_compaction_preserves_append_and_watermarks_in_both_orders() { + for merge_first in [true, false] { + let dir = tempfile::tempdir().unwrap(); + let uri = dir.path().to_str().unwrap(); + let (a, mut coordinator) = prepared_compaction_fixture(uri).await; + let before = Dataset::open(uri).await.unwrap().version().version; + let prepared = prepare_without_commit(uri).await; + assert!(!prepared.is_empty()); + assert_eq!( + Dataset::open(uri).await.unwrap().version().version, + before, + "preparation must not publish even a reservation" + ); + put(&a, "during-compact", 8192).await; + let (plans, _) = coordinator + .plan(&["a".into()], 64, 1024 * 1024) + .await + .unwrap(); + let staged = stage(uri, plans[0].clone(), budget(), None).await.unwrap(); + let high = staged.plan.generations.last().unwrap().number; + let shard = staged.plan.shard; + let mut publisher = + RolloutStore::open_existing_with_options(uri, RolloutStoreOptions::default()) + .await + .unwrap(); + if merge_first { + coordinator.commit(vec![staged.clone()]).await.unwrap(); + } + let metrics = publisher + .commit_prepared_compaction(prepared) + .await + .unwrap(); + assert_eq!(metrics.fragments_removed, 4); + if !merge_first { + coordinator.commit(vec![staged.clone()]).await.unwrap(); + } + assert_eq!( + coordinator.commit(vec![staged]).await.unwrap(), + 0, + "replay after rewrite must not duplicate rows" + ); + assert_eq!(rows(uri).await, 5); + let current = Dataset::open(uri).await.unwrap(); + assert_eq!(watermarks(¤t).await.unwrap().get(&shard), Some(&high)); + let reader = + RolloutStore::open_existing_with_options(uri, RolloutStoreOptions::default()) + .await + .unwrap(); + for i in 0..4 { + assert_eq!( + reader.get_blob(&format!("seed-{i}")).await.unwrap(), + Some(vec![42; 4096]) + ); + } + assert_eq!( + reader.get_blob("during-compact").await.unwrap(), + Some(vec![42; 8192]) + ); + } + } + + #[tokio::test] + async fn prepared_compaction_cannot_publish_with_revoked_authority() { + let dir = tempfile::tempdir().unwrap(); + let uri = dir.path().to_str().unwrap(); + let (_writer, _) = prepared_compaction_fixture(uri).await; + let prepared = prepare_without_commit(uri).await; + let before = Dataset::open(uri).await.unwrap().version().version; + let result = crate::merge_write_scope::MergeWriteScope::with_pinned_authorizer(Arc::new( + NoPreparationCommit, + )) + .run(async { + let mut publisher = + RolloutStore::open_existing_with_options(uri, RolloutStoreOptions::default()) + .await + .unwrap(); + publisher.commit_prepared_compaction(prepared).await + }) + .await; + assert!(result.is_err()); + assert_eq!(Dataset::open(uri).await.unwrap().version().version, before); + assert_eq!(rows(uri).await, 4); + } + + #[tokio::test] + async fn prepared_compaction_rejects_deleted_sources_without_resurrection() { + let dir = tempfile::tempdir().unwrap(); + let uri = dir.path().to_str().unwrap(); + let (_a, _) = prepared_compaction_fixture(uri).await; + let prepared = prepare_without_commit(uri).await; + let mut latest = Dataset::open(uri).await.unwrap(); + latest.delete("id = 'seed-0'").await.unwrap(); + let version = latest.version().version; + let mut publisher = + RolloutStore::open_existing_with_options(uri, RolloutStoreOptions::default()) + .await + .unwrap(); + let error = publisher + .commit_prepared_compaction(prepared) + .await + .unwrap_err(); + assert!( + error.to_string().contains("source fragment changed"), + "{error}" + ); + assert_eq!(Dataset::open(uri).await.unwrap().version().version, version); + assert_eq!(rows(uri).await, 3); + assert!(publisher.get_blob("seed-0").await.unwrap().is_none()); + } + #[tokio::test] async fn parallel_files_publish_once_and_preserve_new_writes() { let dir = tempfile::tempdir().unwrap(); diff --git a/crates/lance-context-core/src/rollout_store.rs b/crates/lance-context-core/src/rollout_store.rs index 9d47a00..712a589 100644 --- a/crates/lance-context-core/src/rollout_store.rs +++ b/crates/lance-context-core/src/rollout_store.rs @@ -706,6 +706,22 @@ impl RolloutStore { self.base.compact(options).await } + /// Prepare immutable compaction output without a base-table commit. + pub async fn prepare_compaction( + &self, + options: CompactionConfig, + ) -> LanceResult { + self.base.prepare_compaction(options).await + } + + /// Publish prepared output using this freshly opened, fenced handle. + pub async fn commit_prepared_compaction( + &mut self, + prepared: crate::PreparedCompaction, + ) -> LanceResult { + self.base.commit_prepared_compaction(prepared).await + } + /// See `StorageBase::base_data_files`. pub fn base_data_files(&self) -> Vec { self.base.base_data_files() diff --git a/crates/lance-context-core/src/store_base.rs b/crates/lance-context-core/src/store_base.rs index 3f6fbd5..8bd2ffa 100644 --- a/crates/lance-context-core/src/store_base.rs +++ b/crates/lance-context-core/src/store_base.rs @@ -51,7 +51,7 @@ use lance::dataset::mem_wal::{ }; use lance::dataset::optimize::{ commit_compaction, compact_files, plan_compaction, CompactionMetrics, CompactionMode, - CompactionOptions, + CompactionOptions, RewriteResult, }; use lance::dataset::{ builder::DatasetBuilder, Dataset, MergeInsertBuilder, NewColumnTransform, WhenMatched, @@ -111,17 +111,38 @@ pub(crate) const DEFAULT_MERGE_MAX_BYTES: usize = 1024 * 1024 * 1024; /// read unaffordable: at ~1,500 per shard, reads exhausted a 32 GiB worker. pub(crate) const DEFAULT_PENDING_GENERATIONS_WARN: usize = 256; +/// Unpublished rewrite files. No dataset handle or commit authorization escapes +/// preparation. A fresh, fenced handle must publish this result. +#[derive(Debug)] +pub struct PreparedCompaction { + read_version: u64, + uri: String, + schema: Schema, + options: CompactionOptions, + tasks: Vec, +} + +impl PreparedCompaction { + pub fn read_version(&self) -> u64 { + self.read_version + } + + pub fn is_empty(&self) -> bool { + self.tasks.is_empty() + } +} + /// Execute only the first `max_source_fragments` from a Lance compaction plan. /// /// Lance's built-in `max_source_fragments` stops before a whole planned task /// that exceeds the budget. A long run of tiny fragments can therefore produce /// one oversized task and compact nothing. Truncating the public task data /// keeps the rewrite contiguous while guaranteeing incremental progress. -async fn compact_files_incremental( - dataset: &mut Dataset, +async fn prepare_compaction_files( + dataset: &Dataset, mut options: CompactionOptions, max_source_fragments: usize, -) -> LanceResult { +) -> LanceResult> { options.max_source_fragments = None; let plan = plan_compaction(dataset, &options).await?; let mut remaining = max_source_fragments; @@ -137,7 +158,7 @@ async fn compact_files_incremental( } } if tasks.is_empty() { - return Ok(CompactionMetrics::default()); + return Ok(Vec::new()); } let concurrency = options.num_threads.unwrap_or(1).max(1); @@ -151,6 +172,16 @@ async fn compact_files_incremental( .try_collect() .await?; + Ok(completed) +} + +async fn compact_files_incremental( + dataset: &mut Dataset, + options: CompactionOptions, + max_source_fragments: usize, +) -> LanceResult { + let completed = + prepare_compaction_files(dataset, options.clone(), max_source_fragments).await?; commit_compaction( dataset, completed, @@ -1527,6 +1558,97 @@ impl StorageBase { // ------------------------------------------------- compaction & indexing + fn compaction_options(config: &CompactionConfig) -> CompactionOptions { + 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, + max_bytes_per_file: config.max_bytes_per_file, + batch_size: config.batch_size, + max_source_fragments: config.max_source_fragments, + compaction_mode: config + .try_binary_copy + .then_some(CompactionMode::TryBinaryCopy), + // Every base table here carries a MemWAL index, which is fieldless + // (it tracks shard/generation bookkeeping, not a data column). + // Lance's inline index remap panics on a fieldless index ("An index + // existed with no fields"), so defer remapping: compaction records a + // fragment-reuse index and remaps lazily instead of touching the + // MemWAL index during the rewrite. + defer_index_remap: true, + ..Default::default() + } + } + + /// Read a fixed snapshot and write new immutable files, without publishing. + /// Callers must forbid manifest commits while this future runs. + pub async fn prepare_compaction( + &self, + config: CompactionConfig, + ) -> LanceResult { + self.ensure_writable()?; + let options = Self::compaction_options(&config); + let tasks = prepare_compaction_files( + &self.dataset, + options.clone(), + config.max_source_fragments.unwrap_or(usize::MAX).max(1), + ) + .await?; + Ok(PreparedCompaction { + read_version: self.dataset.version().version, + uri: self.dataset.uri().to_owned(), + schema: Schema::from(self.dataset.schema()), + options, + tasks, + }) + } + + /// Publish prepared files under a fresh owner. Preserve the original + /// RewriteResult read versions so Lance checks every intervening write. + pub async fn commit_prepared_compaction( + &mut self, + prepared: PreparedCompaction, + ) -> LanceResult { + self.ensure_writable()?; + self.dataset.checkout_latest().await?; + if prepared.uri != self.dataset.uri() + || prepared.schema != Schema::from(self.dataset.schema()) + { + return Err(LanceError::invalid_input( + "compaction dataset or schema changed; reprepare", + )); + } + // Reject a stale rewrite before any ReserveFragments/index commit. + // Append-only merges leave these source fragments unchanged. + let sources: HashMap<_, _> = self + .dataset + .manifest() + .fragments + .iter() + .map(|f| (f.id, f)) + .collect(); + for task in &prepared.tasks { + for old in &task.original_fragments { + if sources.get(&old.id).copied() != Some(old) { + return Err(LanceError::invalid_input( + "compaction source fragment changed; reprepare", + )); + } + } + } + let metrics = commit_compaction( + &mut self.dataset, + prepared.tasks, + Arc::new(DatasetIndexRemapperOptions::default()), + &prepared.options, + ) + .await?; + self.reload().await?; + Ok(metrics) + } + /// Compact the base table's small fragments into larger ones. /// /// Every WAL merge `append`s a new fragment to the base table, so a @@ -1548,27 +1670,7 @@ impl StorageBase { self.ensure_writable()?; let config = options.unwrap_or_default(); - 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, - max_bytes_per_file: config.max_bytes_per_file, - batch_size: config.batch_size, - max_source_fragments: config.max_source_fragments, - compaction_mode: config - .try_binary_copy - .then_some(CompactionMode::TryBinaryCopy), - // Every base table here carries a MemWAL index, which is fieldless - // (it tracks shard/generation bookkeeping, not a data column). - // Lance's inline index remap panics on a fieldless index ("An index - // existed with no fields"), so defer remapping: compaction records a - // fragment-reuse index and remaps lazily instead of touching the - // MemWAL index during the rewrite. - defer_index_remap: true, - ..Default::default() - }; + let lance_options = Self::compaction_options(&config); let result = match config.max_source_fragments { Some(max_source_fragments) => { diff --git a/crates/lance-context-master/src/config.rs b/crates/lance-context-master/src/config.rs index ad925bb..c1315a3 100644 --- a/crates/lance-context-master/src/config.rs +++ b/crates/lance-context-master/src/config.rs @@ -320,6 +320,15 @@ mod tests { /// Deadlines for local table mutations on explicitly owned targets. #[derive(Debug, Clone, clap::Args)] pub struct MaintenanceConfig { + /// Prepare rollout compaction files without holding the table write lock. + /// Explicit rollout targets or '*'; empty keeps legacy behavior. Updates to + /// source fragments invalidate preparation instead of being overwritten. + #[arg(long, env = "COMPACTION_PREPARE_TARGETS", value_delimiter = ',')] + pub compaction_prepare_targets: Vec, + /// Yield prepared output after this long waiting for another table writer. + /// This limits local capacity retention; it never cancels the current writer. + #[arg(long, env = "COMPACTION_COMMIT_WAIT_SECS", default_value_t = 120)] + pub compaction_commit_wait_secs: u64, /// Legacy wire field; owned work now uses the real-progress idle timeout. #[arg(long, env = "MAINTENANCE_TIMEOUT_SECS", default_value_t = 3600)] pub maintenance_timeout_secs: u64, @@ -333,6 +342,8 @@ pub struct MaintenanceConfig { impl Default for MaintenanceConfig { fn default() -> Self { Self { + compaction_prepare_targets: Vec::new(), + compaction_commit_wait_secs: 120, maintenance_timeout_secs: 3600, maintenance_idle_timeout_secs: 600, maintenance_drain_timeout_secs: 30, diff --git a/crates/lance-context-master/src/scheduler.rs b/crates/lance-context-master/src/scheduler.rs index f909f25..c0890ab 100644 --- a/crates/lance-context-master/src/scheduler.rs +++ b/crates/lance-context-master/src/scheduler.rs @@ -133,7 +133,7 @@ struct TaskClaimTiming { /// `timing` carries how long the dispatch loop spent claiming this task and /// waiting for a concurrency permit, so every phase of the task's life lands on /// one metric rather than only the work window. -async fn run_task(state: &Arc, claim: TaskClaim, timing: TaskClaimTiming) { +async fn run_task(state: &Arc, mut claim: TaskClaim, timing: TaskClaimTiming) { let task = claim.task.clone(); let kind = kind_label(task.kind); @@ -147,28 +147,32 @@ async fn run_task(state: &Arc, claim: TaskClaim, timing: TaskClaimT .record(timing.permit_wait.as_secs_f64()); let started = std::time::Instant::now(); - let outcome = match crate::maintenance_execution::reconcile_previous(state, &claim).await { - Err(error) => Err(error), - Ok(()) - if state.config.merge_rollout.draining(&task.target) - && task.kind != TaskKind::MergeWal => - { - Err("target draining legacy writers for merge protocol transition".into()) - } - Ok(()) => match task.kind { - TaskKind::MergeWal => crate::merge_execution::run_merge_wal(state, &claim).await, - _ => { - crate::maintenance_execution::run(state, &claim, async { - match task.kind { - TaskKind::Compact => run_compaction(state, &task).await, - TaskKind::IndexId => run_index_id(state, &task.target).await, - TaskKind::Repair => run_repair(state, &claim).await, - TaskKind::MergeWal => unreachable!(), - } - }) - .await + let outcome = if claim.preparing_compaction() { + run_prepared_compaction(state, &mut claim).await + } else { + match crate::maintenance_execution::reconcile_previous(state, &claim).await { + Err(error) => Err(error), + Ok(()) + if state.config.merge_rollout.draining(&task.target) + && task.kind != TaskKind::MergeWal => + { + Err("target draining legacy writers for merge protocol transition".into()) } - }, + Ok(()) => match task.kind { + TaskKind::MergeWal => crate::merge_execution::run_merge_wal(state, &claim).await, + _ => { + crate::maintenance_execution::run(state, &claim, async { + match task.kind { + TaskKind::Compact => run_compaction(state, &task).await, + TaskKind::IndexId => run_index_id(state, &task.target).await, + TaskKind::Repair => run_repair(state, &claim).await, + TaskKind::MergeWal => unreachable!(), + } + }) + .await + } + }, + } }; let work_elapsed = started.elapsed(); let result = if outcome.is_ok() { "success" } else { "failed" }; @@ -320,6 +324,150 @@ async fn run_repair(state: &Arc, claim: &TaskClaim) -> Result