diff --git a/crates/lance-context-core/src/rollout_store.rs b/crates/lance-context-core/src/rollout_store.rs index c724161..ffb519f 100644 --- a/crates/lance-context-core/src/rollout_store.rs +++ b/crates/lance-context-core/src/rollout_store.rs @@ -3055,8 +3055,23 @@ mod tests { ..CompactionConfig::default() })); - let metrics = store.compact(None).await.unwrap(); + let metrics = store + .compact(Some(CompactionConfig { + min_fragments: 2, + num_threads: Some(1), + max_bytes_per_file: Some(1024 * 1024), + batch_size: Some(1), + max_source_fragments: Some(4), + try_binary_copy: true, + ..Default::default() + })) + .await + .unwrap(); assert!(metrics.fragments_removed > 0); + assert!( + metrics.fragments_removed <= 4, + "one incremental pass must honor max_source_fragments" + ); let after = store.base.dataset.count_fragments(); assert!( diff --git a/crates/lance-context-core/src/store.rs b/crates/lance-context-core/src/store.rs index 176e2d0..00229fb 100644 --- a/crates/lance-context-core/src/store.rs +++ b/crates/lance-context-core/src/store.rs @@ -67,6 +67,14 @@ pub struct CompactionConfig { pub materialize_deletions_threshold: f32, /// Number of threads for compaction (None = auto). pub num_threads: Option, + /// Maximum bytes per output file (None = Lance default). + pub max_bytes_per_file: Option, + /// Rows per input scan batch (None = Lance default). + pub batch_size: Option, + /// Maximum source fragments rewritten by one compaction run. + pub max_source_fragments: Option, + /// Try page-level binary copy before falling back to row decoding. + pub try_binary_copy: bool, /// Interval in seconds between compaction checks. pub check_interval_secs: u64, /// Quiet hours during which compaction is skipped [(start_hour, end_hour)]. @@ -83,6 +91,10 @@ impl Default for CompactionConfig { materialize_deletions: true, materialize_deletions_threshold: 0.1, num_threads: None, + max_bytes_per_file: None, + batch_size: None, + max_source_fragments: None, + try_binary_copy: false, check_interval_secs: 300, quiet_hours: vec![], } diff --git a/crates/lance-context-core/src/store_base.rs b/crates/lance-context-core/src/store_base.rs index eb9990c..5c8bcdd 100644 --- a/crates/lance-context-core/src/store_base.rs +++ b/crates/lance-context-core/src/store_base.rs @@ -42,10 +42,14 @@ use arrow_array::{new_null_array, RecordBatch, RecordBatchIterator}; use arrow_schema::{ArrowError, Schema}; use chrono::{DateTime, Utc}; use futures::{stream, StreamExt, TryStreamExt}; +use lance::dataset::index::DatasetIndexRemapperOptions; use lance::dataset::mem_wal::{ DatasetMemWalExt, LsmScanner, ShardManifestStore, ShardSnapshot, ShardWriter, ShardWriterConfig, }; -use lance::dataset::optimize::{compact_files, CompactionMetrics, CompactionOptions}; +use lance::dataset::optimize::{ + commit_compaction, compact_files, plan_compaction, CompactionMetrics, CompactionMode, + CompactionOptions, +}; use lance::dataset::{ builder::DatasetBuilder, Dataset, NewColumnTransform, WriteMode, WriteParams, }; @@ -70,6 +74,55 @@ pub(crate) const DEFAULT_MANIFEST_SCAN_BATCH_SIZE: usize = 16; /// concurrently while collecting observability metrics. pub(crate) const DEFAULT_OBSERVE_CONCURRENCY: usize = 16; +/// 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, + mut options: CompactionOptions, + max_source_fragments: usize, +) -> LanceResult { + options.max_source_fragments = None; + let plan = plan_compaction(dataset, &options).await?; + let mut remaining = max_source_fragments; + let mut tasks = Vec::new(); + for mut task in plan.compaction_tasks() { + if remaining == 0 { + break; + } + task.task.fragments.truncate(remaining); + remaining -= task.task.fragments.len(); + if !task.task.fragments.is_empty() { + tasks.push(task); + } + } + if tasks.is_empty() { + return Ok(CompactionMetrics::default()); + } + + let concurrency = options.num_threads.unwrap_or(1).max(1); + let dataset_snapshot = dataset.clone(); + let completed = stream::iter(tasks) + .map(|task| { + let dataset_snapshot = dataset_snapshot.clone(); + async move { task.execute(&dataset_snapshot).await } + }) + .buffer_unordered(concurrency) + .try_collect() + .await?; + + commit_compaction( + dataset, + completed, + Arc::new(DatasetIndexRemapperOptions::default()), + &options, + ) + .await +} + /// Name of the scalar index on the base table's key column. One name across /// every store, so all tables index their primary key identically. pub(crate) const ID_INDEX_NAME: &str = "id_idx"; @@ -967,6 +1020,12 @@ impl StorageBase { 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 @@ -977,7 +1036,19 @@ impl StorageBase { ..Default::default() }; - match compact_files(&mut self.dataset, lance_options, None).await { + let result = match config.max_source_fragments { + Some(max_source_fragments) => { + compact_files_incremental( + &mut self.dataset, + lance_options, + max_source_fragments.max(1), + ) + .await + } + None => compact_files(&mut self.dataset, lance_options, None).await, + }; + + match result { Ok(metrics) => { // Reload the handle so the caller (and subsequent reads on this // instance) observe the compacted version. diff --git a/crates/lance-context-master/src/config.rs b/crates/lance-context-master/src/config.rs index 36264af..07d7e74 100644 --- a/crates/lance-context-master/src/config.rs +++ b/crates/lance-context-master/src/config.rs @@ -82,6 +82,31 @@ pub struct MasterConfig { #[arg(long, env = "TARGET_ROWS_PER_FRAGMENT", default_value_t = 1_048_576)] pub target_rows_per_fragment: usize, + /// Maximum number of compactions executing in this master process. + #[arg(long, env = "COMPACTION_CONCURRENCY", default_value_t = 1)] + pub compaction_concurrency: usize, + + /// Lance rewrite tasks executing inside one compaction. + #[arg(long, env = "COMPACTION_THREADS", default_value_t = 1)] + pub compaction_threads: usize, + + /// Rows per input batch when compaction must decode and re-encode data. + #[arg(long, env = "COMPACTION_BATCH_SIZE", default_value_t = 8)] + pub compaction_batch_size: usize, + + /// Maximum source fragments rewritten by one compaction task. `0` disables + /// the incremental limit. + #[arg(long, env = "COMPACTION_MAX_SOURCE_FRAGMENTS", default_value_t = 32)] + pub compaction_max_source_fragments: usize, + + /// Maximum bytes per compacted output file. `0` uses Lance's default. + #[arg( + long, + env = "COMPACTION_MAX_BYTES_PER_FILE", + default_value_t = 1_073_741_824 + )] + pub compaction_max_bytes_per_file: usize, + /// Interval, in seconds, between automatic WAL-merge sweeps. Each sweep /// enqueues a `MergeWal` task for every experiment whose pending MemWAL /// generation count crosses `merge_wal_min_generations`; the task fans out @@ -211,4 +236,14 @@ mod tests { .unwrap(); assert_eq!(disabled.rollout_cache_bytes, 0); } + + #[test] + fn compaction_defaults_bound_parallelism_and_rewrite_size() { + let config = MasterConfig::try_parse_from(["lance-context-master"]).unwrap(); + assert_eq!(config.compaction_concurrency, 1); + assert_eq!(config.compaction_threads, 1); + assert_eq!(config.compaction_batch_size, 8); + assert_eq!(config.compaction_max_source_fragments, 32); + assert_eq!(config.compaction_max_bytes_per_file, 1024 * 1024 * 1024); + } } diff --git a/crates/lance-context-master/src/routes.rs b/crates/lance-context-master/src/routes.rs index ac68251..6d23ac7 100644 --- a/crates/lance-context-master/src/routes.rs +++ b/crates/lance-context-master/src/routes.rs @@ -612,6 +612,11 @@ mod tests { compaction_interval_secs: 0, min_fragments: 16, target_rows_per_fragment: 1_048_576, + compaction_concurrency: 1, + compaction_threads: 1, + compaction_batch_size: 8, + compaction_max_source_fragments: 32, + compaction_max_bytes_per_file: 1024 * 1024 * 1024, merge_wal_interval_secs: 0, merge_wal_min_generations: 8, worker_endpoints: vec![], diff --git a/crates/lance-context-master/src/scanner.rs b/crates/lance-context-master/src/scanner.rs index 2a04c57..ad68268 100644 --- a/crates/lance-context-master/src/scanner.rs +++ b/crates/lance-context-master/src/scanner.rs @@ -13,7 +13,8 @@ use std::time::Duration; use chrono::Utc; use futures::stream::{self, StreamExt}; -use lance_context_core::{RolloutStore, RolloutStoreOptions}; +use lance_context_core::{CompactionConfig, RolloutStore, RolloutStoreOptions}; +use tokio::sync::Semaphore; use tokio::task::JoinHandle; use crate::state::MasterState; @@ -263,6 +264,8 @@ async fn scan_once_inner(state: &Arc) -> lance::Result { retire_after, Utc::now().timestamp_millis(), state.rollout_store_options(), + state.compaction_config(), + state.compaction_permits.clone(), ) .await; if !retired.is_empty() { @@ -436,6 +439,8 @@ async fn retire_cold_experiments( retire_after: Duration, now_ms: i64, options: RolloutStoreOptions, + compaction_config: CompactionConfig, + compaction_permits: Arc, ) -> HashSet { if retire_after.is_zero() { return HashSet::new(); @@ -447,7 +452,15 @@ async fn retire_cold_experiments( if row.last_updated > cutoff_ms { continue; } - match prepare_for_retirement(&row.name, &row.uri, options.clone()).await { + match prepare_for_retirement( + &row.name, + &row.uri, + options.clone(), + compaction_config.clone(), + compaction_permits.clone(), + ) + .await + { Ok(true) => { retired.insert(row.name.clone()); metrics::counter!("master_stats_experiments_retired_total").increment(1); @@ -485,6 +498,8 @@ async fn prepare_for_retirement( name: &str, uri: &str, options: RolloutStoreOptions, + compaction_config: CompactionConfig, + compaction_permits: Arc, ) -> lance::Result { let mut store = match tokio::time::timeout( OBSERVE_TIMEOUT, @@ -509,7 +524,11 @@ async fn prepare_for_retirement( // 2. Compact, so the retired table is not left as many small fragments that // nothing will ever come back to tidy. - tokio::time::timeout(RETIRE_TIMEOUT, store.compact(None)) + let _permit = compaction_permits + .acquire_owned() + .await + .map_err(|_| lance::Error::io("compaction semaphore closed"))?; + tokio::time::timeout(RETIRE_TIMEOUT, store.compact(Some(compaction_config))) .await .map_err(|_| lance::Error::io(format!("compaction timed out retiring '{name}'")))??; @@ -904,6 +923,20 @@ mod retirement_tests { } } + fn compaction_limits() -> (CompactionConfig, Arc) { + ( + CompactionConfig { + num_threads: Some(1), + batch_size: Some(8), + max_source_fragments: Some(32), + max_bytes_per_file: Some(1024 * 1024 * 1024), + try_binary_copy: true, + ..Default::default() + }, + Arc::new(Semaphore::new(1)), + ) + } + /// Retirement must drain the WAL before dropping the row. /// /// The sweeps read the stats table and nothing else, so a retired @@ -932,11 +965,14 @@ mod retirement_tests { let now = Utc::now().timestamp_millis(); let old = now - Duration::from_secs(30 * 86_400).as_millis() as i64; + let (compaction_config, compaction_permits) = compaction_limits(); let retired = retire_cold_experiments( &[row("e", &uri, old)], Duration::from_secs(7 * 86_400), now, RolloutStoreOptions::default(), + compaction_config, + compaction_permits, ) .await; @@ -969,11 +1005,14 @@ mod retirement_tests { } let now = Utc::now().timestamp_millis(); + let (compaction_config, compaction_permits) = compaction_limits(); let retired = retire_cold_experiments( &[row("e", &uri, now)], Duration::from_secs(7 * 86_400), now, RolloutStoreOptions::default(), + compaction_config, + compaction_permits, ) .await; assert!( @@ -986,11 +1025,14 @@ mod retirement_tests { #[tokio::test] async fn zero_window_disables_retirement() { let now = Utc::now().timestamp_millis(); + let (compaction_config, compaction_permits) = compaction_limits(); let retired = retire_cold_experiments( &[row("e", "/nonexistent", 0)], Duration::from_secs(0), now, RolloutStoreOptions::default(), + compaction_config, + compaction_permits, ) .await; assert!(retired.is_empty()); @@ -1002,11 +1044,14 @@ mod retirement_tests { async fn unpreparable_experiment_is_kept() { let now = Utc::now().timestamp_millis(); let old = now - Duration::from_secs(30 * 86_400).as_millis() as i64; + let (compaction_config, compaction_permits) = compaction_limits(); let retired = retire_cold_experiments( &[row("gone", "/no/such/dataset.lance", old)], Duration::from_secs(7 * 86_400), now, RolloutStoreOptions::default(), + compaction_config, + compaction_permits, ) .await; assert!( diff --git a/crates/lance-context-master/src/scheduler.rs b/crates/lance-context-master/src/scheduler.rs index 07e7b67..07c7d39 100644 --- a/crates/lance-context-master/src/scheduler.rs +++ b/crates/lance-context-master/src/scheduler.rs @@ -44,16 +44,6 @@ use crate::task_store::TaskClaim; /// anything still over the threshold is picked up by the next tick. const MAX_SWEEP_ENQUEUE: usize = 256; -/// Build the [`CompactionConfig`] the scheduler applies, from master config. -pub fn compaction_config(state: &MasterState) -> CompactionConfig { - CompactionConfig { - enabled: true, - min_fragments: state.config.min_fragments, - target_rows_per_fragment: state.config.target_rows_per_fragment, - ..Default::default() - } -} - fn kind_label(kind: TaskKind) -> &'static str { match kind { TaskKind::Compact => "compact", @@ -177,7 +167,13 @@ async fn index_id_inner(state: &Arc, name: &str) -> Result, name: &str) -> Result { let uri = state.rollout_uri(name); - let config = compaction_config(state); + let config = state.compaction_config(); + let _permit = state + .compaction_permits + .clone() + .acquire_owned() + .await + .map_err(|_| "compaction semaphore closed".to_string())?; let mut store = RolloutStore::open_existing_with_options(&uri, state.rollout_store_options()) .await @@ -377,7 +373,7 @@ pub async fn sweep_candidates(state: &Arc) -> lance::Result } async fn sweep_candidates_inner(state: &Arc) -> lance::Result { - let config = compaction_config(state); + let config = state.compaction_config(); // Quiet-hours gate applies to the whole sweep. if in_quiet_hours(&config) { return Ok(0); @@ -559,6 +555,11 @@ mod tests { // Low threshold so a handful of appends crosses it. min_fragments: 2, target_rows_per_fragment: 1_048_576, + compaction_concurrency: 1, + compaction_threads: 1, + compaction_batch_size: 8, + compaction_max_source_fragments: 32, + compaction_max_bytes_per_file: 1024 * 1024 * 1024, merge_wal_interval_secs: 0, merge_wal_min_generations: 2, worker_endpoints: vec![], diff --git a/crates/lance-context-master/src/state.rs b/crates/lance-context-master/src/state.rs index d88171a..94928ce 100644 --- a/crates/lance-context-master/src/state.rs +++ b/crates/lance-context-master/src/state.rs @@ -3,9 +3,11 @@ use std::num::NonZeroUsize; use std::sync::Arc; -use lance_context_core::{join_uri, RolloutRegistry, RolloutStore, RolloutStoreOptions, Session}; +use lance_context_core::{ + join_uri, CompactionConfig, RolloutRegistry, RolloutStore, RolloutStoreOptions, Session, +}; use lru::LruCache; -use tokio::sync::{Mutex, RwLock}; +use tokio::sync::{Mutex, RwLock, Semaphore}; use crate::config::MasterConfig; use crate::discovery; @@ -32,6 +34,22 @@ fn build_rollout_session(cache_bytes: usize) -> Option> { Some(RolloutStore::build_session(index_bytes, metadata_bytes)) } +fn build_compaction_config(config: &MasterConfig) -> CompactionConfig { + CompactionConfig { + enabled: true, + min_fragments: config.min_fragments, + target_rows_per_fragment: config.target_rows_per_fragment, + num_threads: Some(config.compaction_threads.max(1)), + max_bytes_per_file: (config.compaction_max_bytes_per_file > 0) + .then_some(config.compaction_max_bytes_per_file), + batch_size: Some(config.compaction_batch_size.max(1)), + max_source_fragments: (config.compaction_max_source_fragments > 0) + .then_some(config.compaction_max_source_fragments), + try_binary_copy: true, + ..Default::default() + } +} + /// Shared state for the master process. /// /// The data-plane owns steady-state registry writes (store create/delete); the @@ -59,6 +77,8 @@ pub struct MasterState { pub task_store: TaskStore, /// Shared HTTP client for fanning `MergeWal` tasks out to worker endpoints. pub http: reqwest::Client, + /// Process-wide compaction permits shared by scheduler and retirement work. + pub(crate) compaction_permits: Arc, /// Whether this process has already run one `_stats` maintenance pass. /// /// The first pass runs without a timeout so a deployment carrying a version @@ -93,6 +113,7 @@ impl MasterState { }; let base_uri = config.data_dir.clone(); let rollout_session = build_rollout_session(config.rollout_cache_bytes); + let compaction_concurrency = config.compaction_concurrency.max(1); let registry_uri = join_uri(&base_uri, "_registry.rollout.lance"); let stats_uri = join_uri(&base_uri, "_stats.rollout.lance"); let mut registry = RolloutRegistry::open_or_create(®istry_uri, None).await?; @@ -116,6 +137,7 @@ impl MasterState { config, task_store, http: reqwest::Client::new(), + compaction_permits: Arc::new(Semaphore::new(compaction_concurrency)), stats_maintenance_done: std::sync::atomic::AtomicBool::new(false), stats_maintenance_failures: std::sync::atomic::AtomicU64::new(0), stats_last_reclaimed_version: std::sync::atomic::AtomicU64::new(0), @@ -137,6 +159,11 @@ impl MasterState { } } + /// Memory-bounded compaction settings for every master-initiated rewrite. + pub(crate) fn compaction_config(&self) -> CompactionConfig { + build_compaction_config(&self.config) + } + /// Return a cached rollout handle for the master records browser, opening /// it without holding the cache lock on a miss. pub async fn get_or_open_record_store( @@ -185,6 +212,11 @@ mod tests { compaction_interval_secs: 0, min_fragments: 16, target_rows_per_fragment: 1_048_576, + compaction_concurrency: 1, + compaction_threads: 1, + compaction_batch_size: 8, + compaction_max_source_fragments: 32, + compaction_max_bytes_per_file: 1024 * 1024 * 1024, merge_wal_interval_secs: 0, merge_wal_min_generations: 8, worker_endpoints: vec![], @@ -211,6 +243,23 @@ mod tests { assert!(build_rollout_session(7).is_some()); } + #[test] + fn compaction_config_maps_bounds_and_zero_disables_optional_limits() { + let dir = TempDir::new().unwrap(); + let mut config = test_config(&dir); + config.compaction_threads = 0; + config.compaction_batch_size = 0; + config.compaction_max_source_fragments = 0; + config.compaction_max_bytes_per_file = 0; + + let compaction = build_compaction_config(&config); + assert_eq!(compaction.num_threads, Some(1)); + assert_eq!(compaction.batch_size, Some(1)); + assert_eq!(compaction.max_source_fragments, None); + assert_eq!(compaction.max_bytes_per_file, None); + assert!(compaction.try_binary_copy); + } + #[tokio::test] #[ignore = "requires ETCD_TEST_ENDPOINTS"] async fn startup_backfills_legacy_rollout_datasets() { diff --git a/crates/lance-context-master/src/task_store.rs b/crates/lance-context-master/src/task_store.rs index 2ec190c..5bc7090 100644 --- a/crates/lance-context-master/src/task_store.rs +++ b/crates/lance-context-master/src/task_store.rs @@ -932,6 +932,11 @@ mod tests { compaction_interval_secs: 0, min_fragments: 16, target_rows_per_fragment: 1_048_576, + compaction_concurrency: 1, + compaction_threads: 1, + compaction_batch_size: 8, + compaction_max_source_fragments: 32, + compaction_max_bytes_per_file: 1024 * 1024 * 1024, merge_wal_interval_secs: 0, merge_wal_min_generations: 8, worker_endpoints: vec![], diff --git a/deploy/kubernetes/README.md b/deploy/kubernetes/README.md index 416772c..7177f53 100644 --- a/deploy/kubernetes/README.md +++ b/deploy/kubernetes/README.md @@ -36,7 +36,13 @@ master replica sees current stats. Two periodic sweeps run on the master and feed the shared scheduler queue: - **Compaction** (`COMPACTION_INTERVAL_SECS`, `MIN_FRAGMENTS`) rewrites an - experiment's base-table fragments locally on the master. + experiment's base-table fragments locally on the master. Large inline blobs + make row decoding expensive, so master compaction first attempts Lance + binary-copy and otherwise uses bounded input batches. Keep + `COMPACTION_CONCURRENCY=1` and `COMPACTION_THREADS=1` unless the pod memory + limit is sized for parallel rewrites. `COMPACTION_MAX_SOURCE_FRAGMENTS` + makes large stores converge incrementally across sweeps, while + `COMPACTION_MAX_BYTES_PER_FILE` prevents giant output files. - **WAL merge** (`MERGE_WAL_INTERVAL_SECS`, `MERGE_WAL_MIN_GENERATIONS`) enqueues a `MergeWal` task for every experiment whose pending MemWAL generation count (from the periodically-scanned stats table) crosses the threshold. The task diff --git a/deploy/kubernetes/master.yaml b/deploy/kubernetes/master.yaml index 9ea5175..8f74a44 100644 --- a/deploy/kubernetes/master.yaml +++ b/deploy/kubernetes/master.yaml @@ -44,6 +44,19 @@ spec: # below the pod memory limit to leave room for query working sets. - name: ROLLOUT_CACHE_BYTES value: "536870912" + # Compaction is process-wide serialized and incrementally rewrites + # small batches of source fragments. These defaults keep large + # inline payloads from multiplying into unbounded rewrite buffers. + - name: COMPACTION_CONCURRENCY + value: "1" + - name: COMPACTION_THREADS + value: "1" + - name: COMPACTION_BATCH_SIZE + value: "8" + - name: COMPACTION_MAX_SOURCE_FRAGMENTS + value: "32" + - name: COMPACTION_MAX_BYTES_PER_FILE + value: "1073741824" - name: ETCD_USERNAME valueFrom: secretKeyRef: