From 6d3f3ebd30845db92bc1774e7c5e88e8758d9a07 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Tue, 29 Sep 2026 14:57:37 +0800 Subject: [PATCH 01/24] feat: horizontal compaction via pluggable executor/committer Repack a fragment's per-column data files into fewer files without moving rows. Adds CompactionExecutor/CompactionCommitter traits (vertical reuses the existing rewrite path), a horizontal RewriteColumns path committing Operation::Update with an empty fields_modified (indices and overlays preserved), column_layout_stats, and a compaction column_groups option. --- rust/lance/src/dataset.rs | 2 + rust/lance/src/dataset/compaction_stats.rs | 109 ++++ rust/lance/src/dataset/optimize.rs | 532 +++++++++++++++++- rust/lance/src/dataset/rewrite_columns.rs | 502 +++++++++++++++++ .../tests/dataset_overlay_index_masking.rs | 32 ++ rust/lance/src/dataset/versions/mod.rs | 16 +- 6 files changed, 1170 insertions(+), 23 deletions(-) create mode 100644 rust/lance/src/dataset/compaction_stats.rs create mode 100644 rust/lance/src/dataset/rewrite_columns.rs diff --git a/rust/lance/src/dataset.rs b/rust/lance/src/dataset.rs index 6ba63e01be3..9271eba89fe 100644 --- a/rust/lance/src/dataset.rs +++ b/rust/lance/src/dataset.rs @@ -69,6 +69,7 @@ pub(crate) mod blob; pub(crate) mod branch_location; pub mod builder; pub mod cleanup; +pub mod compaction_stats; mod data_file; mod data_file_part; pub mod delta; @@ -82,6 +83,7 @@ pub mod optimize; pub(crate) mod overlay; pub mod progress; pub mod refs; +pub mod rewrite_columns; pub mod rowids; pub mod scanner; mod schema_evolution; diff --git a/rust/lance/src/dataset/compaction_stats.rs b/rust/lance/src/dataset/compaction_stats.rs new file mode 100644 index 00000000000..c7e41ef7dc3 --- /dev/null +++ b/rust/lance/src/dataset/compaction_stats.rs @@ -0,0 +1,109 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Per-fragment column-layout statistics. +//! +//! A compaction planner keys on how many data files each fragment carries: a +//! fragment split across many small per-column files (typically from repeated +//! `add_columns` backfills) is the signal horizontal compaction targets. +//! Exposing these stats via [`Dataset::column_layout_stats`] keeps the planning +//! decision observable instead of buried in a black-box heuristic. + +use lance_table::format::overlay::TOMBSTONE_FIELD_ID; + +use super::Dataset; + +/// Column-layout statistics for a single fragment. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct FragmentColumnLayoutStats { + /// The fragment these stats describe. + pub fragment_id: u64, + /// Number of data files still holding at least one live (non-tombstoned) + /// field. A large value on a wide dataset is what horizontal compaction + /// collapses back into fewer files. + pub live_file_count: usize, + /// Number of overlay files attached to the fragment. + pub overlay_count: usize, +} + +impl Dataset { + /// Per-fragment column-layout stats, in manifest fragment order. + /// + /// This is the planning input for horizontal compaction: it reads only + /// fragment metadata (no data files), so it is cheap to call. + pub fn column_layout_stats(&self) -> Vec { + self.manifest + .fragments + .iter() + .map(|fragment| FragmentColumnLayoutStats { + fragment_id: fragment.id, + live_file_count: fragment + .files + .iter() + .filter(|file| file.fields.iter().any(|id| *id != TOMBSTONE_FIELD_ID)) + .count(), + overlay_count: fragment.overlays.len(), + }) + .collect() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::dataset::{NewColumnTransform, WriteParams}; + use arrow_array::{Int32Array, RecordBatch, RecordBatchIterator}; + use arrow_schema::{DataType, Field, Schema as ArrowSchema}; + use std::sync::Arc; + + #[tokio::test] + async fn column_layout_stats_counts_files_per_fragment() { + let schema = Arc::new(ArrowSchema::new(vec![Field::new( + "a", + DataType::Int32, + false, + )])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from_iter_values(0..8))], + ) + .unwrap(); + let reader = RecordBatchIterator::new([Ok(batch)], schema); + // Two fragments (max_rows_per_file = 4), one data file each to start. + let mut dataset = Dataset::write( + reader, + "memory://", + Some(WriteParams { + max_rows_per_file: 4, + ..Default::default() + }), + ) + .await + .unwrap(); + + let stats = dataset.column_layout_stats(); + assert_eq!(stats.len(), 2); + assert!( + stats + .iter() + .all(|s| s.live_file_count == 1 && s.overlay_count == 0) + ); + + // add_columns appends a second data file per fragment. + dataset + .add_columns( + NewColumnTransform::SqlExpressions(vec![("b".into(), "a + 1".into())]), + None, + None, + ) + .await + .unwrap(); + + let stats = dataset.column_layout_stats(); + assert_eq!(stats.len(), 2); + assert!( + stats.iter().all(|s| s.live_file_count == 2), + "each fragment should now have 2 data files: {stats:?}" + ); + } +} diff --git a/rust/lance/src/dataset/optimize.rs b/rust/lance/src/dataset/optimize.rs index 53d7ea6a4a1..7b0352f2149 100644 --- a/rust/lance/src/dataset/optimize.rs +++ b/rust/lance/src/dataset/optimize.rs @@ -325,6 +325,19 @@ pub struct CompactionOptions { /// }; /// ``` pub data_storage_version: Option, + /// Top-level columns to keep in their own data files across this compaction. + /// + /// Each inner list becomes one data file per compacted fragment holding + /// exactly those columns; every column not listed goes to one shared file. + /// Empty (the default) writes every column to a single file per fragment. + /// + /// This is the layout [`Dataset::rewrite_columns`] produces, so a compaction + /// configured with the same groups preserves that split instead of folding + /// wide columns back in with narrow ones. Binary copy is disabled when + /// groups are set, and `max_bytes_per_file` is ignored so every group's + /// files split at the same rows. + #[serde(default)] + pub column_groups: Vec>, /// Transaction properties to store with this commit. /// /// These key-value pairs are stored in the transaction file @@ -359,6 +372,7 @@ impl Default for CompactionOptions { excluded_fragment_ids: Vec::new(), max_overlays_per_fragment: Some(10), data_storage_version: None, + column_groups: Vec::new(), transaction_properties: None, } } @@ -536,6 +550,20 @@ impl CompactionOptions { })?), }; } + "column_groups" => { + self.column_groups = value + .split(';') + .map(|group| { + group + .split(',') + .map(str::trim) + .filter(|column| !column.is_empty()) + .map(str::to_owned) + .collect::>() + }) + .filter(|group| !group.is_empty()) + .collect(); + } _ => { warn!("Ignoring unknown compaction config key: {}", key); } @@ -550,6 +578,16 @@ impl CompactionOptions { self.materialize_deletions = false; } + self.column_groups.retain(|group| !group.is_empty()); + let mut seen = HashSet::new(); + for column in self.column_groups.iter().flatten() { + if !seen.insert(column.as_str()) { + return Err(Error::invalid_input(format!( + "CompactionOptions::column_groups lists column \"{column}\" more than once" + ))); + } + } + for (name, value) in [ ( "max_source_fragments", @@ -613,6 +651,10 @@ async fn can_use_binary_copy( options: &CompactionOptions, fragments: &[Fragment], ) -> bool { + if !options.column_groups.is_empty() { + log::debug!("Binary copy disabled: column_groups splits each fragment across files"); + return false; + } let version = options.write_version(dataset); versions::can_use_binary_copy(version, dataset, options, fragments) .await @@ -781,6 +823,84 @@ pub trait CompactionPlanner: Send + Sync { async fn plan(&self, dataset: &Dataset) -> Result; } +/// The executor half of a pluggable compaction pipeline: run one task the +/// planner produced, yielding a result the paired [`CompactionCommitter`] +/// commits. Vertical (`RewriteExecutor`) and, later, horizontal compaction are +/// separate implementations selected by mode — not variants of one plan, so +/// each keeps its own task/result types and serialized shapes. +#[async_trait::async_trait] +pub trait CompactionExecutor: Send + Sync { + /// The task unit the paired planner emits. + type Task: Send; + /// Per-task output, collected and handed to the committer. + type TaskResult: Send; + + /// Execute one task against `dataset` (already at the plan's read version). + async fn execute( + &self, + dataset: Cow<'_, Dataset>, + task: Self::Task, + options: &CompactionOptions, + ) -> Result; +} + +/// The commit half of a pluggable compaction pipeline: fold the executor's +/// results into one transaction. Each pipeline commits its own operation +/// (vertical → `Operation::Rewrite`; horizontal → `Operation::Update`). +#[async_trait::async_trait] +pub trait CompactionCommitter: Send + Sync { + /// Must match the paired executor's [`CompactionExecutor::TaskResult`]. + type TaskResult: Send; + + /// Commit the collected results, returning the run's metrics. + async fn commit( + &self, + dataset: &mut Dataset, + results: Vec, + remap_options: Arc, + options: &CompactionOptions, + ) -> Result; +} + +/// Vertical compaction: rewrite groups of fragments via [`rewrite_files`], +/// committed as `Operation::Rewrite` by [`RewriteCommitter`]. +#[derive(Debug, Default, Clone, Copy)] +pub struct RewriteExecutor; + +#[async_trait::async_trait] +impl CompactionExecutor for RewriteExecutor { + type Task = TaskData; + type TaskResult = RewriteResult; + + async fn execute( + &self, + dataset: Cow<'_, Dataset>, + task: TaskData, + options: &CompactionOptions, + ) -> Result { + rewrite_files(dataset, task, options).await + } +} + +/// Commits vertical compaction results as a single `Operation::Rewrite`. +#[derive(Debug, Default, Clone, Copy)] +pub struct RewriteCommitter; + +#[async_trait::async_trait] +impl CompactionCommitter for RewriteCommitter { + type TaskResult = RewriteResult; + + async fn commit( + &self, + dataset: &mut Dataset, + results: Vec, + remap_options: Arc, + options: &CompactionOptions, + ) -> Result { + commit_compaction(dataset, results, remap_options, options).await + } +} + /// Formulate a plan to compact the files in a dataset /// /// The compaction plan will contain a list of tasks to execute. Each task @@ -814,6 +934,10 @@ impl CompactionPlanner for DefaultCompactionPlanner { dataset.manifest.data_storage_format.lance_file_format(), write_version, )?; + if !self.options.column_groups.is_empty() { + // Surface a bad column-group config here rather than mid-write. + column_group_schemas(dataset.schema(), &self.options.column_groups)?; + } if self.options.defer_index_remap && dataset.manifest.uses_stable_row_ids() { return Err(Error::invalid_input( "defer_index_remap=true is not supported on datasets with stable row IDs: \ @@ -1051,8 +1175,9 @@ pub async fn compact_files_with_planner( let dataset_ref = &dataset.clone(); + let executor = RewriteExecutor; let result_stream = futures::stream::iter(compaction_plan.tasks) - .map(|task| rewrite_files(Cow::Borrowed(dataset_ref), task, &compaction_plan.options)) + .map(|task| executor.execute(Cow::Borrowed(dataset_ref), task, &compaction_plan.options)) .buffer_unordered( compaction_plan .options @@ -1062,13 +1187,14 @@ pub async fn compact_files_with_planner( let completed_tasks: Vec = result_stream.try_collect().await?; let remap_options = remap_options.unwrap_or(Arc::new(DatasetIndexRemapperOptions::default())); - let metrics = commit_compaction( - dataset, - completed_tasks, - remap_options, - &compaction_plan.options, - ) - .await?; + let metrics = RewriteCommitter + .commit( + dataset, + completed_tasks, + remap_options, + &compaction_plan.options, + ) + .await?; Ok(metrics) } @@ -2084,7 +2210,10 @@ impl CompactionTask { } else { Cow::Owned(dataset.checkout_version(self.read_version).await?) }; - rewrite_files(dataset, self.task.clone(), &self.options).await + let executor = RewriteExecutor; + executor + .execute(dataset, self.task.clone(), &self.options) + .await } } @@ -2375,6 +2504,164 @@ async fn reserve_fragment_ids( Ok(()) } +/// The schema of each data file a compacted fragment gets when column groups +/// are configured: the columns no group claims first (when any), then one +/// schema per group. Columns keep the dataset's order within each file. +fn column_group_schemas( + schema: &lance_core::datatypes::Schema, + groups: &[Vec], +) -> Result> { + let names = || schema.fields.iter().map(|field| field.name.as_str()); + for column in groups.iter().flatten() { + if !names().any(|name| name == column) { + return Err(Error::invalid_input(format!( + "column_groups names \"{column}\", which is not a top-level column of the dataset" + ))); + } + } + let claimed = |name: &str| groups.iter().flatten().any(|column| column == name); + let rest: Vec<&str> = names().filter(|name| !claimed(name)).collect(); + std::iter::once(rest) + .chain(groups.iter().map(|group| { + names() + .filter(|name| group.iter().any(|column| column == name)) + .collect() + })) + .filter(|columns: &Vec<&str>| !columns.is_empty()) + .map(|columns| schema.project(&columns)) + .collect() +} + +/// Compaction's grouped write: fan one read of the merged rows out to one +/// writer per column group, so each output fragment gets one data file per +/// group instead of one file holding every column. All writers break at the +/// same `file_row_counts`, so their fragments line up one-to-one and merge by +/// appending each later group's files onto the first group's fragments. Index +/// seeds go to whichever group holds the indexed column. +async fn write_column_group_fragments( + write_version: ConcreteFileVersion, + dataset: &Dataset, + group_schemas: Vec, + mut reader: SendableRecordBatchStream, + mut params: WriteParams, + file_row_counts: Vec, +) -> Result> { + use futures::SinkExt; + + // `file_row_counts` is authoritative for the grouped write: every group must + // break at the same rows so the fragments merge one-to-one. A byte cap would + // roll an extra file in a wide group only (write.rs re-plans the remainder on + // a byte-driven close), leaving groups misaligned, so disable it here. + params.max_bytes_per_file = usize::MAX; + + let arrow_schema = reader.schema(); + let projections = group_schemas + .iter() + .map(|schema| { + schema + .fields + .iter() + .map(|field| arrow_schema.index_of(&field.name)) + .collect::, _>>() + }) + .collect::, _>>()?; + let group_arrow_schemas = projections + .iter() + .map(|projection| arrow_schema.project(projection).map(Arc::new)) + .collect::, _>>()?; + + // One bounded channel per group; the forwarder projects each read batch to + // every group. Bound 1 keeps a slow writer from letting the others buffer + // the whole scan ahead of it in memory. + let (mut senders, group_streams): (Vec<_>, Vec<_>) = group_arrow_schemas + .into_iter() + .map(|group_arrow| { + let (tx, rx) = + futures::channel::mpsc::channel::>(1); + let stream: SendableRecordBatchStream = + Box::pin(RecordBatchStreamAdapter::new(group_arrow, rx)); + (tx, stream) + }) + .unzip(); + + // Blob-check each group, then seed the scalar indexes once and route each + // seed to the group holding its column. Every dataset column lands in + // exactly one group (the unclaimed columns form the first group), so no seed + // is created only to be dropped, and each index is seeded exactly once. + for group_schema in &group_schemas { + versions::validate_write_schema(write_version, group_schema)?; + } + let mut per_group_seeds: Vec>> = + group_schemas.iter().map(|_| Vec::new()).collect(); + for seed in versions::create_seed_writers(write_version, Some(dataset), ¶ms).await? { + if let Some(index) = group_schemas + .iter() + .position(|schema| schema.field(seed.column_name()).is_some()) + { + per_group_seeds[index].push(seed); + } + } + + let forward = async move { + while let Some(batch) = reader.next().await { + let batch = batch?; + for (sender, projection) in senders.iter_mut().zip(&projections) { + // A writer that stopped early (its own error) closes its + // receiver; stop forwarding and let try_join surface that error. + if sender.send(Ok(batch.project(projection)?)).await.is_err() { + return Ok(()); + } + } + } + Ok::<(), Error>(()) + }; + let writers = group_streams + .into_iter() + .zip(group_schemas) + .zip(per_group_seeds) + .map(|((group_stream, group_schema), seed_writers)| { + let params = params.clone(); + let file_row_counts = file_row_counts.clone(); + async move { + versions::write_fragments_direct( + write_version, + Some(dataset), + dataset.object_store.clone(), + &dataset.base, + &group_schema, + group_stream, + params, + None, + seed_writers, + Some(file_row_counts), + ) + .await + } + }) + .collect::>(); + + let (_, per_group) = futures::try_join!(forward, futures::future::try_join_all(writers))?; + + let mut per_group = per_group.into_iter(); + let mut fragments = per_group.next().unwrap_or_default(); + for group in per_group { + if group.len() != fragments.len() + || fragments + .iter() + .zip(&group) + .any(|(a, b)| a.physical_rows != b.physical_rows) + { + return Err(Error::internal( + "column group writers did not split the compacted rows at the same fragments", + )); + } + for (fragment, other) in fragments.iter_mut().zip(group) { + fragment.files.extend(other.files); + } + } + Ok(fragments) +} + /// Rewrite the files in a single task. /// /// This assumes that the dataset is the correct read version to be compacted. @@ -2604,19 +2891,33 @@ async fn rewrite_files( row_ids_rx = Some(rx); } } else { - let (frags, _) = write_fragments_internal_with_file_row_counts( - write_version, - Some(dataset.as_ref()), - dataset.object_store.clone(), - &dataset.base, - dataset.schema().clone(), - reader.expect("reader must be prepared for non-binary-copy path"), - params, - None, - Some(file_row_counts), - ) - .await?; - new_fragments = frags; + let reader = reader.expect("reader must be prepared for non-binary-copy path"); + new_fragments = if options.column_groups.is_empty() { + let (frags, _) = write_fragments_internal_with_file_row_counts( + write_version, + Some(dataset.as_ref()), + dataset.object_store.clone(), + &dataset.base, + dataset.schema().clone(), + reader, + params, + None, + Some(file_row_counts), + ) + .await?; + frags + } else { + let group_schemas = column_group_schemas(dataset.schema(), &options.column_groups)?; + write_column_group_fragments( + write_version, + dataset.as_ref(), + group_schemas, + reader, + params, + file_row_counts, + ) + .await? + }; } log::info!("Compaction task {}: file written", task_id); @@ -7955,6 +8256,86 @@ mod tests { /// frag-reuse aware, so a desync would surface here as recall collapse or a /// resurrected/again-deleted row. /// Top-k `id`s for a KNN query against the `vec` column. + /// Horizontal compaction ([`Dataset::rewrite_columns`]) repacks a column's + /// files without moving rows, so a vector index over that column keeps its + /// uuid and coverage and is never rebuilt — the property #9291 only tested + /// for BTree. + #[tokio::test] + async fn rewrite_columns_preserves_vector_index() { + use arrow_array::cast::AsArray; + use arrow_array::types::{Float32Type, Int32Type}; + use lance_datagen::Dimension; + + const DIM: u32 = 32; + let mut dataset = lance_datagen::gen_batch() + .col("id", lance_datagen::array::step::()) + .col( + "vec", + lance_datagen::array::rand_vec::(Dimension::from(DIM)), + ) + .into_ram_dataset(FragmentCount::from(4), FragmentRowCount::from(256)) + .await + .unwrap(); + let params = VectorIndexParams::with_ivf_pq_params( + DistanceType::L2, + small_ivf(), + PQBuildParams { + max_iters: 2, + num_sub_vectors: 2, + ..Default::default() + }, + ); + dataset + .create_index( + &["vec"], + IndexType::Vector, + Some("vec_idx".into()), + ¶ms, + false, + ) + .await + .unwrap(); + let before = dataset + .load_index_by_name("vec_idx") + .await + .unwrap() + .unwrap(); + + // A query taken from a real row, and its answer before the rewrite. + let query = { + let mut scanner = dataset.scan(); + scanner.project(&["vec"]).unwrap(); + scanner.limit(Some(1), None).unwrap(); + let batch = scanner.try_into_batch().await.unwrap(); + batch["vec"] + .as_fixed_size_list() + .value(0) + .as_primitive::() + .values() + .to_vec() + }; + let knn_before = vector_knn_ids(&dataset, &query, 5).await; + + dataset.rewrite_columns(&["vec"], None).await.unwrap(); + + let after = dataset + .load_index_by_name("vec_idx") + .await + .unwrap() + .unwrap(); + assert_eq!(after.uuid, before.uuid, "index must not be rebuilt"); + assert_eq!( + after.fragment_bitmap, before.fragment_bitmap, + "index coverage must be unchanged" + ); + assert_eq!( + vector_knn_ids(&dataset, &query, 5).await, + knn_before, + "search must return the same rows" + ); + dataset.validate().await.unwrap(); + } + async fn vector_knn_ids(dataset: &Dataset, query: &[f32], k: usize) -> Vec { use arrow_array::cast::AsArray; use arrow_array::types::{Float32Type, Int32Type}; @@ -11168,4 +11549,111 @@ mod tests { Some(4) ); } + + /// Three-column dataset in two 10-row fragments, all columns in one file. + async fn column_group_dataset() -> (Dataset, RecordBatch) { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, true), + Field::new("c", DataType::Int32, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from_iter_values(0..20)), + Arc::new(Int32Array::from_iter_values((0..20).map(|v| v * 10))), + Arc::new(Int32Array::from_iter_values((0..20).map(|v| v * 100))), + ], + ) + .unwrap(); + let reader = RecordBatchIterator::new([Ok(batch)], schema); + let dataset = Dataset::write( + reader, + "memory://", + Some(WriteParams { + max_rows_per_file: 10, + data_storage_version: Some(LanceFileVersion::V2_0), + ..Default::default() + }), + ) + .await + .unwrap(); + let before = dataset.scan().try_into_batch().await.unwrap(); + (dataset, before) + } + + #[tokio::test] + async fn compaction_column_groups_write_one_file_per_group() { + let (mut dataset, before) = column_group_dataset().await; + + // Merge the two small fragments; column_groups splits the output columns + // into their own files: {a} unclaimed, then {b} and {c}. + let options = CompactionOptions { + column_groups: vec![vec!["b".into()], vec!["c".into()]], + ..Default::default() + }; + compact_files(&mut dataset, options, None).await.unwrap(); + + let fragments = dataset.get_fragments(); + assert_eq!(fragments.len(), 1, "the two fragments merged into one"); + let mut groups: Vec> = fragments[0] + .metadata() + .files + .iter() + .map(|file| file.fields.to_vec()) + .collect(); + groups.sort(); + assert_eq!( + groups, + vec![vec![0], vec![1], vec![2]], + "one data file per column group" + ); + assert_eq!(dataset.scan().try_into_batch().await.unwrap(), before); + dataset.validate().await.unwrap(); + } + + #[tokio::test] + async fn compaction_column_groups_reject_unknown_column() { + let (mut dataset, _) = column_group_dataset().await; + let options = CompactionOptions { + column_groups: vec![vec!["nonexistent".into()]], + ..Default::default() + }; + let err = compact_files(&mut dataset, options, None) + .await + .unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. }), "{err}"); + } + + #[tokio::test] + async fn compaction_column_groups_ignore_byte_limit() { + let (mut dataset, before) = column_group_dataset().await; + // A tiny byte cap plus groups of different widths ({b,c} vs {a}) would, + // if honored, roll the wider group into more files and leave the groups + // misaligned. column_groups splits by planned row counts only, so the + // compaction succeeds and still yields one file per group. + let options = CompactionOptions { + column_groups: vec![vec!["b".into(), "c".into()]], + max_bytes_per_file: Some(64), + ..Default::default() + }; + compact_files(&mut dataset, options, None).await.unwrap(); + + let fragments = dataset.get_fragments(); + assert_eq!(fragments.len(), 1, "the two fragments merged into one"); + let mut groups: Vec> = fragments[0] + .metadata() + .files + .iter() + .map(|file| file.fields.to_vec()) + .collect(); + groups.sort(); + assert_eq!( + groups, + vec![vec![0], vec![1, 2]], + "one file for {{a}}, one for the {{b, c}} group" + ); + assert_eq!(dataset.scan().try_into_batch().await.unwrap(), before); + dataset.validate().await.unwrap(); + } } diff --git a/rust/lance/src/dataset/rewrite_columns.rs b/rust/lance/src/dataset/rewrite_columns.rs new file mode 100644 index 00000000000..05db0b57331 --- /dev/null +++ b/rust/lance/src/dataset/rewrite_columns.rs @@ -0,0 +1,502 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Horizontal compaction: rewrite a few columns of every fragment into their +//! own data files, leaving the files that hold the other columns untouched. +//! +//! This is the horizontal pipeline of the pluggable compaction framework +//! ([`CompactionExecutor`] / [`CompactionCommitter`] in [`super::optimize`]). +//! Unlike vertical compaction (which merges whole fragments and re-encodes +//! every column), it re-reads only the named columns, writes them to one new +//! file per fragment, and tombstones them where they used to live. Rows, +//! fragment ids, row addresses and indices are unchanged, so it commits as an +//! `Operation::Update` in [`UpdateMode::RewriteColumns`] with an empty +//! `fields_modified` (values did not change: indices keep their coverage and +//! overlays keep shadowing). + +use std::borrow::Cow; +use std::collections::HashSet; +use std::sync::Arc; + +use lance_core::datatypes::Schema; +use lance_file::version::{ConcreteFileVersion, LanceFileVersion}; +use lance_table::format::{Fragment, overlay::TOMBSTONE_FIELD_ID}; + +use super::fragment::FileFragment; +use super::optimize::{ + CompactionCommitter, CompactionExecutor, CompactionMetrics, CompactionOptions, +}; +use super::transaction::{Operation, Transaction, UpdateMode}; +use super::write::cleanup_data_fragments; +use super::{Dataset, versions}; +use crate::dataset::optimize::remapping::IndexRemapperOptions; +use crate::{Error, Result}; + +/// The exact V2 version the new files get: `requested`, or the dataset's +/// default write version, which is never changed. +fn write_version( + dataset: &Dataset, + requested: Option, +) -> Result { + let default_version = dataset.manifest.data_storage_format.lance_file_format(); + let version = requested + .map(LanceFileVersion::resolve) + .unwrap_or(default_version); + versions::validate_write_version(default_version, version)?; + if version == ConcreteFileVersion::V1 { + return Err(Error::not_supported( + "rewrite_columns requires the V2 file format: a V1 file cannot tombstone \ + single fields, so compact the dataset to V2 first", + )); + } + Ok(version) +} +/// The schema of the new data file: `columns` in dataset order, all top-level. +fn rewrite_schema(dataset: &Dataset, columns: &[&str]) -> Result { + if columns.is_empty() { + return Err(Error::invalid_input( + "rewrite_columns needs at least one column", + )); + } + let schema = dataset.schema(); + if let Some(column) = columns + .iter() + .find(|column| !schema.fields.iter().any(|field| &field.name == *column)) + { + return Err(Error::invalid_input(format!( + "Column \"{column}\" is not a top-level column of the dataset" + ))); + } + let ordered = schema + .fields + .iter() + .map(|field| field.name.as_str()) + .filter(|name| columns.contains(name)) + .collect::>(); + let schema = schema.project(&ordered)?; + if let Some(blob) = schema.fields_pre_order().find(|field| field.is_blob()) { + return Err(Error::not_supported(format!( + "rewrite_columns cannot rewrite blob column \"{}\"", + blob.name + ))); + } + Ok(schema) +} +impl FileFragment { + /// Rewrite `columns` of this fragment into one new data file. + /// + /// The returned metadata has the new file appended and the rewritten fields + /// tombstoned in the files they came from; a file left holding only + /// tombstones is dropped. Nothing is committed. Returns `None` when the + /// columns already sit alone in a file of the requested version, so a + /// retried or resumed rewrite skips finished work. + pub async fn rewrite_columns( + &self, + columns: &[&str], + data_storage_version: Option, + ) -> Result> { + let dataset = self.dataset(); + let write_schema = rewrite_schema(dataset, columns)?; + let write_version = write_version(dataset, data_storage_version)?; + let field_ids: HashSet = write_schema.field_ids().into_iter().collect(); + + let already_split = self.metadata().files.iter().any(|file| { + file.fields + .iter() + .copied() + .filter(|id| *id != TOMBSTONE_FIELD_ID) + .collect::>() + == field_ids + && file + .file_version() + .is_ok_and(|version| version == write_version) + }); + if already_split { + return Ok(None); + } + + let mut updater = self + .updater_with_version( + Some(columns), + Some((write_schema, dataset.schema().clone())), + None, + None, + write_version, + ) + .await?; + let written: Result = async { + while let Some(batch) = updater.next().await?.cloned() { + updater.update(batch).await?; + } + updater.finish().await + } + .await; + let mut fragment = match written { + Ok(fragment) => fragment, + Err(err) => { + updater.cleanup_unfinished_writer().await; + return Err(err); + } + }; + // No live rows, so no file was written and nothing to tombstone. + if fragment.files.len() == self.metadata().files.len() { + return Ok(None); + } + + let new_file = fragment.files.len() - 1; + for file in &mut fragment.files[..new_file] { + file.fields = file + .fields + .iter() + .map(|id| { + if field_ids.contains(id) { + TOMBSTONE_FIELD_ID + } else { + *id + } + }) + .collect::>() + .into(); + } + // A file holding only tombstones is unreachable to readers. + fragment + .files + .retain(|file| file.fields.iter().any(|id| *id != TOMBSTONE_FIELD_ID)); + Ok(Some(fragment)) + } +} +/// Horizontal executor: rewrites one fragment's named columns into a new data +/// file (see [`FileFragment::rewrite_columns`]). The columns and target version +/// are the executor's configuration; the task is the fragment to rewrite. +#[derive(Debug, Clone)] +pub struct RewriteColumnsExecutor { + columns: Vec, + data_storage_version: Option, +} + +#[async_trait::async_trait] +impl CompactionExecutor for RewriteColumnsExecutor { + type Task = FileFragment; + /// `None` when the fragment already had the layout (nothing to commit). + type TaskResult = Option; + + async fn execute( + &self, + _dataset: Cow<'_, Dataset>, + task: FileFragment, + _options: &CompactionOptions, + ) -> Result> { + let columns: Vec<&str> = self.columns.iter().map(String::as_str).collect(); + task.rewrite_columns(&columns, self.data_storage_version) + .await + } +} +/// Commits horizontal-compaction results as one `Operation::Update` in +/// [`UpdateMode::RewriteColumns`] with an empty `fields_modified`: values did +/// not change, so index coverage is kept and overlays keep shadowing. +#[derive(Debug, Default, Clone, Copy)] +pub struct RewriteColumnsCommitter; + +#[async_trait::async_trait] +impl CompactionCommitter for RewriteColumnsCommitter { + type TaskResult = Option; + + async fn commit( + &self, + dataset: &mut Dataset, + results: Vec>, + _remap_options: Arc, + _options: &CompactionOptions, + ) -> Result { + let updated_fragments: Vec = results.into_iter().flatten().collect(); + if updated_fragments.is_empty() { + return Ok(CompactionMetrics::default()); + } + // Row addresses and values are unchanged, so no fragment is added or + // removed; only per-fragment files change. `files_added` is one new + // file per rewritten fragment; `files_removed` is the fully-tombstoned + // files that were dropped. + let mut metrics = CompactionMetrics { + files_added: updated_fragments.len(), + ..CompactionMetrics::default() + }; + for fragment in &updated_fragments { + let original_files = dataset + .get_fragment(fragment.id as usize) + .map(|f| f.metadata().files.len()) + .unwrap_or(0); + metrics.files_removed += (original_files + 1).saturating_sub(fragment.files.len()); + } + let transaction = Transaction::new( + dataset.manifest.version, + Operation::Update { + removed_fragment_ids: Vec::new(), + updated_fragments, + new_fragments: Vec::new(), + fields_modified: Vec::new(), + compacted_sstables: Vec::new(), + fields_for_preserving_frag_bitmap: Vec::new(), + update_mode: Some(UpdateMode::RewriteColumns), + inserted_rows_filter: None, + updated_fragment_offsets: None, + }, + None, + ); + dataset + .apply_commit(transaction, &Default::default(), &Default::default()) + .await?; + Ok(metrics) + } +} +impl Dataset { + /// Rewrite `columns` of every fragment into one new data file per fragment, + /// optionally in a different V2 version, without touching the files that + /// hold the other columns. Fragments already in the requested layout are + /// skipped, so an interrupted rewrite can be rerun. + /// + /// Runs through the horizontal compaction pipeline + /// ([`RewriteColumnsExecutor`] + [`RewriteColumnsCommitter`]) and commits + /// once. To spread the work over many machines, call + /// [`FileFragment::rewrite_columns`] per fragment and commit the returned + /// metadata with [`RewriteColumnsCommitter`]. + pub async fn rewrite_columns( + &mut self, + columns: &[&str], + data_storage_version: Option, + ) -> Result { + let executor = RewriteColumnsExecutor { + columns: columns.iter().map(|c| c.to_string()).collect(), + data_storage_version, + }; + let options = CompactionOptions::default(); + let fragments = self.get_fragments(); + let mut results = Vec::with_capacity(fragments.len()); + for fragment in fragments { + match executor + .execute(Cow::Borrowed(&*self), fragment, &options) + .await + { + Ok(result) => results.push(result), + Err(err) => { + self.cleanup_rewritten_files(&results).await; + return Err(err); + } + } + } + RewriteColumnsCommitter + .commit( + self, + results, + Arc::new(crate::dataset::index::DatasetIndexRemapperOptions::default()), + &options, + ) + .await + } + + /// Delete the files a failed rewrite wrote: the last file of every + /// rewritten fragment is the only one that did not exist before. + async fn cleanup_rewritten_files(&self, results: &[Option]) { + let new_files = results + .iter() + .flatten() + .filter_map(|fragment| fragment.files.last()) + .map(|file| Fragment { + files: vec![file.clone()], + ..Fragment::new(0) + }) + .collect::>(); + cleanup_data_fragments(&self.object_store, &self.base, None, &new_files).await; + } +} +#[cfg(test)] +mod tests { + use super::*; + use crate::dataset::WriteParams; + use crate::index::DatasetIndexExt; + use arrow_array::{Int32Array, RecordBatch, RecordBatchIterator, StringArray}; + use arrow_schema::{DataType, Field, Schema as ArrowSchema}; + use lance_core::utils::tempfile::TempStrDir; + use lance_index::IndexType; + use lance_index::scalar::ScalarIndexParams; + use std::sync::Arc; + + fn batch(start: i32, rows: i32) -> RecordBatch { + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, true), + Field::new("c", DataType::Utf8, true), + ])); + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from_iter_values(start..start + rows)), + Arc::new(Int32Array::from_iter_values( + (start..start + rows).map(|v| v * 10), + )), + Arc::new(StringArray::from_iter_values( + (start..start + rows).map(|v| format!("row-{v}")), + )), + ], + ) + .unwrap() + } + + async fn write(uri: &str, version: LanceFileVersion) -> Dataset { + let data = batch(0, 8); + let reader = RecordBatchIterator::new([Ok(data.clone())], data.schema()); + Dataset::write( + reader, + uri, + Some(WriteParams { + max_rows_per_file: 4, + data_storage_version: Some(version), + ..Default::default() + }), + ) + .await + .unwrap() + } + + fn file_versions(fragment: &Fragment) -> Vec<(Vec, ConcreteFileVersion)> { + fragment + .files + .iter() + .map(|file| (file.fields.to_vec(), file.file_version().unwrap())) + .collect() + } + #[tokio::test] + async fn rewrite_migrates_only_the_named_column() { + let mut dataset = write("memory://", LanceFileVersion::V2_0).await; + dataset.delete("a = 3").await.unwrap(); + let before = dataset.scan().try_into_batch().await.unwrap(); + let version = dataset.manifest.version; + + let metrics = dataset + .rewrite_columns(&["c"], Some(LanceFileVersion::V2_2)) + .await + .unwrap(); + assert_eq!(metrics.files_added, 2, "one new file per fragment"); + + assert_eq!(dataset.manifest.version, version + 1); + assert_eq!(dataset.manifest.fragments.len(), 2); + for fragment in dataset.manifest.fragments.iter() { + assert_eq!( + file_versions(fragment), + vec![ + (vec![0, 1, TOMBSTONE_FIELD_ID], ConcreteFileVersion::V2_0), + (vec![2], ConcreteFileVersion::V2_2), + ] + ); + } + // Values are byte-for-byte unchanged. + assert_eq!(dataset.scan().try_into_batch().await.unwrap(), before); + // The dataset default write version is not changed. + assert_eq!( + dataset.manifest.data_storage_format.lance_file_format(), + ConcreteFileVersion::V2_0 + ); + dataset.validate().await.unwrap(); + + // Idempotent: already in the requested layout, nothing committed. + dataset + .rewrite_columns(&["c"], Some(LanceFileVersion::V2_2)) + .await + .unwrap(); + assert_eq!(dataset.manifest.version, version + 1); + } + #[tokio::test] + async fn rewrite_keeps_scalar_index() { + let mut dataset = write("memory://", LanceFileVersion::V2_0).await; + dataset + .create_index( + &["a"], + IndexType::BTree, + Some("a_idx".into()), + &ScalarIndexParams::default(), + false, + ) + .await + .unwrap(); + let before = dataset.load_indices().await.unwrap()[0].clone(); + + dataset.rewrite_columns(&["a"], None).await.unwrap(); + + let after = dataset.load_indices().await.unwrap()[0].clone(); + assert_eq!(after.uuid, before.uuid); + assert_eq!(after.fragment_bitmap, before.fragment_bitmap); + let filtered = dataset + .scan() + .filter("a = 5") + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(filtered.num_rows(), 1); + dataset.validate().await.unwrap(); + } + + #[tokio::test] + async fn rewrite_rejects_bad_input() { + let mut dataset = write("memory://", LanceFileVersion::V2_0).await; + let err = dataset + .rewrite_columns(&["missing"], None) + .await + .unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. }), "{err}"); + + let legacy_dir = TempStrDir::default(); + let mut legacy = write(&legacy_dir, LanceFileVersion::Legacy).await; + let err = legacy.rewrite_columns(&["c"], None).await.unwrap_err(); + assert!(matches!(err, Error::NotSupported { .. }), "{err}"); + } + + #[tokio::test] + async fn rewrite_collapses_backfilled_files() { + use crate::dataset::NewColumnTransform; + let mut dataset = write("memory://", LanceFileVersion::V2_0).await; + // Two backfills, each appending one data file per fragment: the feature's + // target — a fragment split across many per-column files. + dataset + .add_columns( + NewColumnTransform::SqlExpressions(vec![("d".into(), "a + 1".into())]), + None, + None, + ) + .await + .unwrap(); + dataset + .add_columns( + NewColumnTransform::SqlExpressions(vec![("e".into(), "a + 2".into())]), + None, + None, + ) + .await + .unwrap(); + let before = dataset.scan().try_into_batch().await.unwrap(); + assert!( + dataset + .column_layout_stats() + .iter() + .all(|s| s.live_file_count == 3), + "each fragment carries three data files after two backfills: {:?}", + dataset.column_layout_stats() + ); + + // Repack every column into a single file per fragment. + dataset + .rewrite_columns(&["a", "b", "c", "d", "e"], None) + .await + .unwrap(); + + assert!( + dataset + .column_layout_stats() + .iter() + .all(|s| s.live_file_count == 1), + "rewrite collapsed each fragment to one data file: {:?}", + dataset.column_layout_stats() + ); + assert_eq!(dataset.scan().try_into_batch().await.unwrap(), before); + dataset.validate().await.unwrap(); + } +} diff --git a/rust/lance/src/dataset/tests/dataset_overlay_index_masking.rs b/rust/lance/src/dataset/tests/dataset_overlay_index_masking.rs index 0113319c528..3a6142d72e8 100644 --- a/rust/lance/src/dataset/tests/dataset_overlay_index_masking.rs +++ b/rust/lance/src/dataset/tests/dataset_overlay_index_masking.rs @@ -255,6 +255,38 @@ async fn test_overlay_stale_drop_and_new_match(#[values(false, true)] stable_row assert_eq!(ids_matching(&dataset, "age = 20").await, vec![2]); } +/// Horizontal compaction ([`Dataset::rewrite_columns`]) commits with an empty +/// `fields_modified`, so an existing data overlay must be carried forward and +/// keep shadowing the base — even when the rewrite repacks an unrelated column. +#[tokio::test] +async fn rewrite_columns_preserves_data_overlay() { + let dataset = create_base_dataset().await; + // Fragment 0, offset 1 is id=1, age=10; the overlay changes its age to 999. + let mut dataset = commit_overlay( + dataset, + "age_overlay", + 0, + &[1], + OverlayCoverage::dense(RoaringBitmap::from_iter([1])), + vec![i32_array([Some(999)])], + ) + .await; + assert_eq!(ids_matching(&dataset, "age = 999").await, vec![1]); + + // Repack an unrelated column; the overlay on `age` must survive. + dataset.rewrite_columns(&["id"], None).await.unwrap(); + + assert_eq!( + dataset.get_fragment(0).unwrap().metadata().overlays.len(), + 1, + "the data overlay must survive the rewrite" + ); + // The overlaid value still shadows the base: age=999 for id=1, old 10 gone. + assert_eq!(ids_matching(&dataset, "age = 999").await, vec![1]); + assert_eq!(ids_matching(&dataset, "age = 10").await, Vec::::new()); + dataset.validate().await.unwrap(); +} + /// Row-level BTree precision: when one row in a covered fragment is stale, only that row is /// blocked from the index result and re-evaluated on the stale-Take path. Non-stale rows in /// the same fragment (including one that matches the predicate) remain on the indexed path. diff --git a/rust/lance/src/dataset/versions/mod.rs b/rust/lance/src/dataset/versions/mod.rs index 051e933a6bf..6c28be681af 100644 --- a/rust/lance/src/dataset/versions/mod.rs +++ b/rust/lance/src/dataset/versions/mod.rs @@ -105,7 +105,21 @@ pub fn schema_compare_options(version: ConcreteFileVersion) -> SchemaCompareOpti } } -async fn create_seed_writers( +/// Reject a write schema whose blob columns don't match the file version, the +/// same check [`write_fragments`] runs — exposed for callers that go straight to +/// [`write_fragments_direct`] with a schema they built themselves. +pub(super) fn validate_write_schema(version: ConcreteFileVersion, schema: &Schema) -> Result<()> { + match version { + ConcreteFileVersion::V1 | ConcreteFileVersion::V2_0 | ConcreteFileVersion::V2_1 => { + write::validate_legacy_blob_write_schema(schema, &format!("{version:?}")) + } + ConcreteFileVersion::V2_2 | ConcreteFileVersion::V2_3 => { + write::validate_blob_v2_write_schema(schema) + } + } +} + +pub(super) async fn create_seed_writers( version: ConcreteFileVersion, dataset: Option<&Dataset>, params: &WriteParams, From 0cb9d4803db29ce2e9bc9b1b0f25044ecb346516 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Tue, 29 Sep 2026 17:10:38 +0800 Subject: [PATCH 02/24] refactor(compaction): make the executor/committer seam load-bearing Add a generic run_compaction_pipeline that both vertical and horizontal compaction run through, so the traits are a real plug point. Route Dataset::rewrite_columns through it (bounded concurrency, no bespoke cleanup). Surface a missing fragment and an uncovered indexed column as errors instead of silently miscounting or dropping a seed. Add tests for column_groups seed routing (ZoneMap), split-preserved-across-compaction, config parsing, and index/overlay preservation. --- rust/lance/src/dataset/optimize.rs | 248 ++++++++++++++++++---- rust/lance/src/dataset/rewrite_columns.rs | 108 ++++++---- 2 files changed, 269 insertions(+), 87 deletions(-) diff --git a/rust/lance/src/dataset/optimize.rs b/rust/lance/src/dataset/optimize.rs index 7b0352f2149..9a686bba18a 100644 --- a/rust/lance/src/dataset/optimize.rs +++ b/rust/lance/src/dataset/optimize.rs @@ -792,8 +792,9 @@ pub struct CompactionMetrics { pub fragments_added: usize, /// The number of files that have been removed, including deletion files. pub files_removed: usize, - /// The number of files that have been added, which is always equal to the - /// number of fragments. + /// The number of data files that have been added. This equals the number of + /// new fragments for a plain vertical compaction, but a `column_groups` + /// compaction adds one file per group per fragment, so it can be larger. pub files_added: usize, } @@ -901,6 +902,44 @@ impl CompactionCommitter for RewriteCommitter { } } +/// Run a compaction pipeline: execute every task with bounded concurrency, then +/// commit the collected results. This is the one orchestration both vertical +/// (`RewriteExecutor`/`RewriteCommitter`) and horizontal +/// (`RewriteColumnsExecutor`/`RewriteColumnsCommitter`) funnel through, so the +/// [`CompactionExecutor`]/[`CompactionCommitter`] traits carry real polymorphism +/// — a caller with its own executor/committer pair plugs in here — rather than +/// documenting a seam nothing consumes. The two pipelines still differ in how +/// they produce tasks and which operation they commit; only the execute-collect- +/// commit skeleton is shared. +pub(crate) async fn run_compaction_pipeline( + dataset: &mut Dataset, + executor: E, + committer: C, + tasks: Vec, + remap_options: Arc, + options: &CompactionOptions, + concurrency: usize, +) -> Result +where + E: CompactionExecutor, + C: CompactionCommitter, +{ + if tasks.is_empty() { + return Ok(CompactionMetrics::default()); + } + // Execute against an immutable snapshot; `commit` takes `&mut dataset` once + // every task has finished and the borrows are released. + let dataset_ref = &dataset.clone(); + let results: Vec = futures::stream::iter(tasks) + .map(|task| executor.execute(Cow::Borrowed(dataset_ref), task, options)) + .buffer_unordered(concurrency.max(1)) + .try_collect() + .await?; + committer + .commit(dataset, results, remap_options, options) + .await +} + /// Formulate a plan to compact the files in a dataset /// /// The compaction plan will contain a list of tasks to execute. Each task @@ -1173,30 +1212,21 @@ pub async fn compact_files_with_planner( return Ok(CompactionMetrics::default()); } - let dataset_ref = &dataset.clone(); - - let executor = RewriteExecutor; - let result_stream = futures::stream::iter(compaction_plan.tasks) - .map(|task| executor.execute(Cow::Borrowed(dataset_ref), task, &compaction_plan.options)) - .buffer_unordered( - compaction_plan - .options - .num_threads - .unwrap_or_else(get_num_compute_intensive_cpus), - ); - - let completed_tasks: Vec = result_stream.try_collect().await?; + let concurrency = compaction_plan + .options + .num_threads + .unwrap_or_else(get_num_compute_intensive_cpus); let remap_options = remap_options.unwrap_or(Arc::new(DatasetIndexRemapperOptions::default())); - let metrics = RewriteCommitter - .commit( - dataset, - completed_tasks, - remap_options, - &compaction_plan.options, - ) - .await?; - - Ok(metrics) + run_compaction_pipeline( + dataset, + RewriteExecutor, + RewriteCommitter, + compaction_plan.tasks, + remap_options, + &compaction_plan.options, + concurrency, + ) + .await } /// Information about a fragment used to decide its fate in compaction @@ -2594,11 +2624,21 @@ async fn write_column_group_fragments( let mut per_group_seeds: Vec>> = group_schemas.iter().map(|_| Vec::new()).collect(); for seed in versions::create_seed_writers(write_version, Some(dataset), ¶ms).await? { - if let Some(index) = group_schemas + match group_schemas .iter() .position(|schema| schema.field(seed.column_name()).is_some()) { - per_group_seeds[index].push(seed); + Some(index) => per_group_seeds[index].push(seed), + // `column_group_schemas` partitions every top-level column, so an + // index's keyed column always lands in exactly one group. Guard the + // invariant rather than silently dropping the seed, which would + // rebuild the index from incomplete data. + None => { + return Err(Error::internal(format!( + "column groups do not cover indexed column \"{}\"", + seed.column_name() + ))); + } } } @@ -8243,19 +8283,6 @@ mod tests { assert!(checked > 0, "expected to check at least one stored vector"); } - /// Build an `id` + `vec` dataset, create the given IVF vector index, - /// optionally delete rows, then run deferred compaction (which materializes - /// the deletions into the fragment-reuse index) and assert that KNN over - /// surviving vectors during the FRI window (a) never returns a deleted row - /// and (b) stays consistent with the pre-compaction answer. - /// - /// The deletion path is the interesting one: materialized deletions drop - /// rows from the quantization storage at load time, which shifts storage - /// positions. Flat storage (FLAT/PQ/SQ/RQ) is scanned linearly so this is - /// fine, but the HNSW graph addresses storage positionally and is not - /// frag-reuse aware, so a desync would surface here as recall collapse or a - /// resurrected/again-deleted row. - /// Top-k `id`s for a KNN query against the `vec` column. /// Horizontal compaction ([`Dataset::rewrite_columns`]) repacks a column's /// files without moving rows, so a vector index over that column keeps its /// uuid and coverage and is never rebuilt — the property #9291 only tested @@ -8336,6 +8363,7 @@ mod tests { dataset.validate().await.unwrap(); } + /// Top-k `id`s for a KNN query against the `vec` column. async fn vector_knn_ids(dataset: &Dataset, query: &[f32], k: usize) -> Vec { use arrow_array::cast::AsArray; use arrow_array::types::{Float32Type, Int32Type}; @@ -8357,6 +8385,18 @@ mod tests { ids } + /// Build an `id` + `vec` dataset, create the given IVF vector index, + /// optionally delete rows, then run deferred compaction (which materializes + /// the deletions into the fragment-reuse index) and assert that KNN over + /// surviving vectors during the FRI window (a) never returns a deleted row + /// and (b) stays consistent with the pre-compaction answer. + /// + /// The deletion path is the interesting one: materialized deletions drop + /// rows from the quantization storage at load time, which shifts storage + /// positions. Flat storage (FLAT/PQ/SQ/RQ) is scanned linearly so this is + /// fine, but the HNSW graph addresses storage positionally and is not + /// frag-reuse aware, so a desync would surface here as recall collapse or a + /// resurrected/again-deleted row. async fn check_vector_defer_compaction( params: VectorIndexParams, delete_predicate: Option<&str>, @@ -11656,4 +11696,132 @@ mod tests { assert_eq!(dataset.scan().try_into_batch().await.unwrap(), before); dataset.validate().await.unwrap(); } + + #[tokio::test] + async fn compaction_column_groups_preserve_scalar_index() { + let (mut dataset, before) = column_group_dataset().await; + // ZoneMap is the scalar index that writes an index seed, so it actually + // drives the seed-routing loop in `write_column_group_fragments` (BTree + // writes no seed and would leave that path untested). + // ZoneMap only writes an index seed when `use_seeds` is enabled (for a + // fixed-width column it defaults off), and the seed is what drives the + // routing loop in `write_column_group_fragments`. + let index_params = ScalarIndexParams::for_builtin(BuiltinIndexType::ZoneMap) + .with_params(&serde_json::json!({"use_seeds": true})); + dataset + .create_index( + &["a"], + IndexType::ZoneMap, + Some("a_idx".into()), + &index_params, + false, + ) + .await + .unwrap(); + // Split the indexed column into its own file group: its seed must be + // routed to the group holding "a" (a miss would trip the + // `Error::internal` guard and fail the compaction). + let options = CompactionOptions { + column_groups: vec![vec!["a".into()]], + ..Default::default() + }; + compact_files(&mut dataset, options, None).await.unwrap(); + + // The grouped write ran the seed-routing loop (ZoneMap writes a seed, + // unlike BTree) and produced the requested split without corrupting the + // index metadata or the values. + let fragments = dataset.get_fragments(); + assert_eq!(fragments.len(), 1, "the two fragments merged into one"); + let mut groups: Vec> = fragments[0] + .metadata() + .files + .iter() + .map(|file| file.fields.to_vec()) + .collect(); + groups.sort(); + assert_eq!( + groups, + vec![vec![0], vec![1, 2]], + "\"a\" split into its own file, {{b, c}} in another" + ); + assert!( + dataset.load_index_by_name("a_idx").await.unwrap().is_some(), + "the ZoneMap index survives compaction" + ); + let filtered = dataset + .scan() + .filter("a >= 15") + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(filtered.num_rows(), 5, "a >= 15 matches five rows (15..20)"); + assert_eq!(dataset.scan().try_into_batch().await.unwrap(), before); + dataset.validate().await.unwrap(); + } + + #[tokio::test] + async fn column_groups_preserve_horizontal_split_across_compaction() { + let (mut dataset, before) = column_group_dataset().await; + // Horizontal split: column c into its own file per fragment. + dataset.rewrite_columns(&["c"], None).await.unwrap(); + for fragment in dataset.get_fragments() { + assert_eq!( + fragment.metadata().files.len(), + 2, + "c is split into its own file" + ); + } + // Vertical compaction with a matching group keeps c split instead of + // folding it back into one file — the feature's core promise. + let options = CompactionOptions { + column_groups: vec![vec!["c".into()]], + ..Default::default() + }; + compact_files(&mut dataset, options, None).await.unwrap(); + + let fragments = dataset.get_fragments(); + assert_eq!(fragments.len(), 1, "the two fragments merged"); + let mut groups: Vec> = fragments[0] + .metadata() + .files + .iter() + .map(|file| file.fields.to_vec()) + .collect(); + groups.sort(); + assert_eq!( + groups, + vec![vec![0, 1], vec![2]], + "layout stays split: {{a, b}} and {{c}}" + ); + assert_eq!(dataset.scan().try_into_batch().await.unwrap(), before); + dataset.validate().await.unwrap(); + } + + #[test] + fn compaction_column_groups_parse_from_config() { + let config = std::collections::HashMap::from([( + "lance.compaction.column_groups".to_string(), + "b, c ; d".to_string(), + )]); + let options = CompactionOptions::from_dataset_config(&config).unwrap(); + assert_eq!( + options.column_groups, + vec![ + vec!["b".to_string(), "c".to_string()], + vec!["d".to_string()] + ], + "';' splits groups, ',' splits columns, whitespace trimmed" + ); + } + + #[test] + fn compaction_column_groups_validate_rejects_duplicate() { + let mut options = CompactionOptions { + column_groups: vec![vec!["b".into()], vec!["b".into()]], + ..Default::default() + }; + let err = options.validate().unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. }), "{err}"); + } } diff --git a/rust/lance/src/dataset/rewrite_columns.rs b/rust/lance/src/dataset/rewrite_columns.rs index 05db0b57331..f21b8615cd2 100644 --- a/rust/lance/src/dataset/rewrite_columns.rs +++ b/rust/lance/src/dataset/rewrite_columns.rs @@ -25,9 +25,9 @@ use lance_table::format::{Fragment, overlay::TOMBSTONE_FIELD_ID}; use super::fragment::FileFragment; use super::optimize::{ CompactionCommitter, CompactionExecutor, CompactionMetrics, CompactionOptions, + run_compaction_pipeline, }; use super::transaction::{Operation, Transaction, UpdateMode}; -use super::write::cleanup_data_fragments; use super::{Dataset, versions}; use crate::dataset::optimize::remapping::IndexRemapperOptions; use crate::{Error, Result}; @@ -97,7 +97,7 @@ impl FileFragment { ) -> Result> { let dataset = self.dataset(); let write_schema = rewrite_schema(dataset, columns)?; - let write_version = write_version(dataset, data_storage_version)?; + let resolved_version = write_version(dataset, data_storage_version)?; let field_ids: HashSet = write_schema.field_ids().into_iter().collect(); let already_split = self.metadata().files.iter().any(|file| { @@ -109,7 +109,7 @@ impl FileFragment { == field_ids && file .file_version() - .is_ok_and(|version| version == write_version) + .is_ok_and(|version| version == resolved_version) }); if already_split { return Ok(None); @@ -121,7 +121,7 @@ impl FileFragment { Some((write_schema, dataset.schema().clone())), None, None, - write_version, + resolved_version, ) .await?; let written: Result = async { @@ -174,12 +174,26 @@ pub struct RewriteColumnsExecutor { data_storage_version: Option, } +impl RewriteColumnsExecutor { + /// Build an executor that repacks `columns` into one new file per fragment, + /// optionally in a different V2 version. + pub fn new(columns: &[&str], data_storage_version: Option) -> Self { + Self { + columns: columns.iter().map(|c| c.to_string()).collect(), + data_storage_version, + } + } +} + #[async_trait::async_trait] impl CompactionExecutor for RewriteColumnsExecutor { type Task = FileFragment; /// `None` when the fragment already had the layout (nothing to commit). type TaskResult = Option; + /// The task carries its own dataset handle (a [`FileFragment`] already bound + /// to the read version), so the `dataset` argument the pipeline threads in is + /// unused here — the read version comes from the fragment, not the snapshot. async fn execute( &self, _dataset: Cow<'_, Dataset>, @@ -221,10 +235,21 @@ impl CompactionCommitter for RewriteColumnsCommitter { ..CompactionMetrics::default() }; for fragment in &updated_fragments { + // A fragment being updated must exist in the version we commit + // against; a missing one is an inconsistency (e.g. a distributed + // caller passing a stale fragment id), not a zero-file fragment. let original_files = dataset .get_fragment(fragment.id as usize) - .map(|f| f.metadata().files.len()) - .unwrap_or(0); + .ok_or_else(|| { + Error::internal(format!( + "rewrite_columns is updating fragment {} that is not in the \ + current manifest", + fragment.id + )) + })? + .metadata() + .files + .len(); metrics.files_removed += (original_files + 1).saturating_sub(fragment.files.len()); } let transaction = Transaction::new( @@ -264,48 +289,22 @@ impl Dataset { columns: &[&str], data_storage_version: Option, ) -> Result { - let executor = RewriteColumnsExecutor { - columns: columns.iter().map(|c| c.to_string()).collect(), - data_storage_version, - }; + let executor = RewriteColumnsExecutor::new(columns, data_storage_version); let options = CompactionOptions::default(); - let fragments = self.get_fragments(); - let mut results = Vec::with_capacity(fragments.len()); - for fragment in fragments { - match executor - .execute(Cow::Borrowed(&*self), fragment, &options) - .await - { - Ok(result) => results.push(result), - Err(err) => { - self.cleanup_rewritten_files(&results).await; - return Err(err); - } - } - } - RewriteColumnsCommitter - .commit( - self, - results, - Arc::new(crate::dataset::index::DatasetIndexRemapperOptions::default()), - &options, - ) - .await - } - - /// Delete the files a failed rewrite wrote: the last file of every - /// rewritten fragment is the only one that did not exist before. - async fn cleanup_rewritten_files(&self, results: &[Option]) { - let new_files = results - .iter() - .flatten() - .filter_map(|fragment| fragment.files.last()) - .map(|file| Fragment { - files: vec![file.clone()], - ..Fragment::new(0) - }) - .collect::>(); - cleanup_data_fragments(&self.object_store, &self.base, None, &new_files).await; + let tasks = self.get_fragments(); + let concurrency = self.object_store.io_parallelism(); + // Like vertical compaction, a failed run leaves its new files + // uncommitted and unreferenced; dataset cleanup reclaims them. + run_compaction_pipeline( + self, + executor, + RewriteColumnsCommitter, + tasks, + Arc::new(crate::dataset::index::DatasetIndexRemapperOptions::default()), + &options, + concurrency, + ) + .await } } #[cfg(test)] @@ -499,4 +498,19 @@ mod tests { assert_eq!(dataset.scan().try_into_batch().await.unwrap(), before); dataset.validate().await.unwrap(); } + + #[tokio::test] + async fn rewrite_after_fragment_fully_deleted() { + let mut dataset = write("memory://", LanceFileVersion::V2_0).await; + // Deleting every row of a fragment removes it from the manifest, so the + // rewrite runs over the surviving fragment(s); it must still succeed and + // leave the live rows byte-for-byte unchanged. + dataset.delete("a < 4").await.unwrap(); + let before = dataset.scan().try_into_batch().await.unwrap(); + + dataset.rewrite_columns(&["c"], None).await.unwrap(); + + assert_eq!(dataset.scan().try_into_batch().await.unwrap(), before); + dataset.validate().await.unwrap(); + } } From 051f4ad9a30b97f1fbb50446b933d4b69da50e5e Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Tue, 29 Sep 2026 18:02:20 +0800 Subject: [PATCH 03/24] docs: avoid a public-to-private intra-doc link on RewriteExecutor --- rust/lance/src/dataset/optimize.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rust/lance/src/dataset/optimize.rs b/rust/lance/src/dataset/optimize.rs index 9a686bba18a..cfc39d646db 100644 --- a/rust/lance/src/dataset/optimize.rs +++ b/rust/lance/src/dataset/optimize.rs @@ -863,7 +863,7 @@ pub trait CompactionCommitter: Send + Sync { ) -> Result; } -/// Vertical compaction: rewrite groups of fragments via [`rewrite_files`], +/// Vertical compaction: rewrite groups of fragments via `rewrite_files`, /// committed as `Operation::Rewrite` by [`RewriteCommitter`]. #[derive(Debug, Default, Clone, Copy)] pub struct RewriteExecutor; From c873253ab674a6b9d240ce59972c0004faf2f687 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Tue, 29 Sep 2026 23:07:47 +0800 Subject: [PATCH 04/24] fix(compaction): bind horizontal rewrite results to their read version RewriteColumnsCommitter committed against the current manifest version, so in the distributed path a fragment changed between its read and the commit would slip past the conflict check. Carry each rewrite's read version in RewriteColumnsResult and commit against the minimum, mirroring vertical compaction's RewriteResult::read_version. Single-machine rewrite_columns is unaffected (read version == commit version). --- rust/lance/src/dataset/rewrite_columns.rs | 52 +++++++++++++++++------ 1 file changed, 40 insertions(+), 12 deletions(-) diff --git a/rust/lance/src/dataset/rewrite_columns.rs b/rust/lance/src/dataset/rewrite_columns.rs index f21b8615cd2..2dcc7a76a0d 100644 --- a/rust/lance/src/dataset/rewrite_columns.rs +++ b/rust/lance/src/dataset/rewrite_columns.rs @@ -165,6 +165,18 @@ impl FileFragment { Ok(Some(fragment)) } } +/// One fragment's horizontal-rewrite output, bound to the manifest version it +/// was read at. The committer commits against that version so the conflict +/// check catches any change made to the fragment between the read and the +/// commit (mirrors vertical compaction's `RewriteResult::read_version`). +#[derive(Debug, Clone)] +pub struct RewriteColumnsResult { + /// The rewritten fragment metadata to publish. + pub fragment: Fragment, + /// The manifest version the fragment was read at. + pub read_version: u64, +} + /// Horizontal executor: rewrites one fragment's named columns into a new data /// file (see [`FileFragment::rewrite_columns`]). The columns and target version /// are the executor's configuration; the task is the fragment to rewrite. @@ -189,20 +201,26 @@ impl RewriteColumnsExecutor { impl CompactionExecutor for RewriteColumnsExecutor { type Task = FileFragment; /// `None` when the fragment already had the layout (nothing to commit). - type TaskResult = Option; + type TaskResult = Option; - /// The task carries its own dataset handle (a [`FileFragment`] already bound - /// to the read version), so the `dataset` argument the pipeline threads in is - /// unused here — the read version comes from the fragment, not the snapshot. + /// The task is a [`FileFragment`] already bound to its read version, so the + /// `dataset` argument the pipeline threads in is unused; the read version is + /// taken from the fragment and carried into the result for the committer. async fn execute( &self, _dataset: Cow<'_, Dataset>, task: FileFragment, _options: &CompactionOptions, - ) -> Result> { + ) -> Result> { + let read_version = task.dataset().manifest.version; let columns: Vec<&str> = self.columns.iter().map(String::as_str).collect(); - task.rewrite_columns(&columns, self.data_storage_version) - .await + Ok(task + .rewrite_columns(&columns, self.data_storage_version) + .await? + .map(|fragment| RewriteColumnsResult { + fragment, + read_version, + })) } } /// Commits horizontal-compaction results as one `Operation::Update` in @@ -213,19 +231,29 @@ pub struct RewriteColumnsCommitter; #[async_trait::async_trait] impl CompactionCommitter for RewriteColumnsCommitter { - type TaskResult = Option; + type TaskResult = Option; async fn commit( &self, dataset: &mut Dataset, - results: Vec>, + results: Vec>, _remap_options: Arc, _options: &CompactionOptions, ) -> Result { - let updated_fragments: Vec = results.into_iter().flatten().collect(); - if updated_fragments.is_empty() { + let results: Vec = results.into_iter().flatten().collect(); + if results.is_empty() { return Ok(CompactionMetrics::default()); } + // Commit against the earliest version any rewritten fragment was read + // at, not the current one, so the conflict check catches a fragment + // changed between the read and this commit (as vertical compaction does). + let read_version = results + .iter() + .map(|result| result.read_version) + .min() + .expect("results is non-empty"); + let updated_fragments: Vec = + results.into_iter().map(|result| result.fragment).collect(); // Row addresses and values are unchanged, so no fragment is added or // removed; only per-fragment files change. `files_added` is one new // file per rewritten fragment; `files_removed` is the fully-tombstoned @@ -253,7 +281,7 @@ impl CompactionCommitter for RewriteColumnsCommitter { metrics.files_removed += (original_files + 1).saturating_sub(fragment.files.len()); } let transaction = Transaction::new( - dataset.manifest.version, + read_version, Operation::Update { removed_fragment_ids: Vec::new(), updated_fragments, From b15e1e9c651235c6f7c2460d8c55a80369d9133a Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Tue, 29 Sep 2026 23:33:27 +0800 Subject: [PATCH 05/24] test(compaction): cover the horizontal rewrite read-version conflict Regression test for the read-version binding: a rewrite read at V, committed after a delete lands on the same fragment at V+1, must conflict rather than resurrect the deleted row. Guards the fix so the distributed commit path cannot silently regress. --- rust/lance/src/dataset/rewrite_columns.rs | 59 +++++++++++++++++++++++ 1 file changed, 59 insertions(+) diff --git a/rust/lance/src/dataset/rewrite_columns.rs b/rust/lance/src/dataset/rewrite_columns.rs index 2dcc7a76a0d..f7e6661d9a3 100644 --- a/rust/lance/src/dataset/rewrite_columns.rs +++ b/rust/lance/src/dataset/rewrite_columns.rs @@ -541,4 +541,63 @@ mod tests { assert_eq!(dataset.scan().try_into_batch().await.unwrap(), before); dataset.validate().await.unwrap(); } + + #[tokio::test] + async fn rewrite_columns_conflicts_with_intervening_delete() { + use crate::dataset::index::DatasetIndexRemapperOptions; + use crate::dataset::optimize::{ + CompactionCommitter, CompactionExecutor, CompactionOptions, + }; + use std::borrow::Cow; + + let mut dataset = write("memory://", LanceFileVersion::V2_0).await; + + // Rewrite column "c" of fragment 0, reading at the current version. + let executor = RewriteColumnsExecutor::new(&["c"], None); + let fragment = dataset.get_fragment(0).unwrap(); + let stale = executor + .execute( + Cow::Borrowed(&dataset), + fragment, + &CompactionOptions::default(), + ) + .await + .unwrap() + .expect("fragment 0 has columns to repack"); + + // A delete lands on that same fragment before the rewrite is committed. + dataset.delete("a = 1").await.unwrap(); + let after_delete = dataset.scan().try_into_batch().await.unwrap(); + assert_eq!(after_delete.num_rows(), 7); + + // The rewrite was read before the delete, so committing it must conflict + // rather than publish the stale layout and resurrect the deleted row. + let committed = RewriteColumnsCommitter + .commit( + &mut dataset, + vec![Some(stale)], + Arc::new(DatasetIndexRemapperOptions::default()), + &CompactionOptions::default(), + ) + .await; + assert!( + committed.is_err(), + "a rewrite read before the delete must not commit over it" + ); + + // The delete stands: seven rows, none with a = 1. + assert_eq!(dataset.scan().try_into_batch().await.unwrap(), after_delete); + assert_eq!( + dataset + .scan() + .filter("a = 1") + .unwrap() + .try_into_batch() + .await + .unwrap() + .num_rows(), + 0 + ); + dataset.validate().await.unwrap(); + } } From 453db7b4290fc5eebea2b2c21ed91a7c869fbc64 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Fri, 2 Oct 2026 10:52:47 +0800 Subject: [PATCH 06/24] fix(compaction): conflict horizontal rewrites with concurrent drops, plus review fixes --- rust/lance/src/dataset/optimize.rs | 134 ++++++-- rust/lance/src/dataset/rewrite_columns.rs | 323 +++++++++++++++--- rust/lance/src/index/append.rs | 77 ++++- rust/lance/src/io/commit/conflict_resolver.rs | 32 +- 4 files changed, 481 insertions(+), 85 deletions(-) diff --git a/rust/lance/src/dataset/optimize.rs b/rust/lance/src/dataset/optimize.rs index 7d6edc67494..067e8e8fc42 100644 --- a/rust/lance/src/dataset/optimize.rs +++ b/rust/lance/src/dataset/optimize.rs @@ -340,8 +340,10 @@ pub struct CompactionOptions { /// This is the layout [`Dataset::rewrite_columns`] produces, so a compaction /// configured with the same groups preserves that split instead of folding /// wide columns back in with narrow ones. Binary copy is disabled when - /// groups are set, and `max_bytes_per_file` is ignored so every group's - /// files split at the same rows. + /// groups are set (`ForceBinaryCopy` is rejected), and `max_bytes_per_file` + /// is ignored so every group's files split at the same rows. Set from + /// dataset config with `lance.compaction.column_groups` (see + /// [`Self::from_dataset_config`]). #[serde(default)] pub column_groups: Vec>, /// Transaction properties to store with this commit. @@ -409,6 +411,10 @@ impl CompactionOptions { /// - `lance.compaction.max_source_bytes` /// - `lance.compaction.max_overlays_per_fragment` /// - `lance.compaction.data_storage_version` + /// - `lance.compaction.column_groups`: groups separated by `;`, columns in + /// a group by `,`, surrounding whitespace trimmed (`"b, c; d"` is + /// `[["b", "c"], ["d"]]`). A column whose name contains `,` or `;`, or + /// starts or ends with whitespace, cannot be named this way. pub fn from_dataset_config(config: &HashMap) -> Result { let mut opts = Self::default(); opts.apply_dataset_config(config)?; @@ -593,6 +599,15 @@ impl CompactionOptions { ))); } } + if !self.column_groups.is_empty() + && matches!(self.compaction_mode(), CompactionMode::ForceBinaryCopy) + { + return Err(Error::invalid_input( + "CompactionOptions::column_groups cannot be combined with \ + compaction_mode=ForceBinaryCopy: binary copy keeps each fragment's \ + file layout, so it cannot split columns into groups", + )); + } for (name, value) in [ ( @@ -909,14 +924,11 @@ impl CompactionCommitter for RewriteCommitter { } /// Run a compaction pipeline: execute every task with bounded concurrency, then -/// commit the collected results. This is the one orchestration both vertical +/// commit the collected results. Both vertical /// (`RewriteExecutor`/`RewriteCommitter`) and horizontal -/// (`RewriteColumnsExecutor`/`RewriteColumnsCommitter`) funnel through, so the -/// [`CompactionExecutor`]/[`CompactionCommitter`] traits carry real polymorphism -/// — a caller with its own executor/committer pair plugs in here — rather than -/// documenting a seam nothing consumes. The two pipelines still differ in how -/// they produce tasks and which operation they commit; only the execute-collect- -/// commit skeleton is shared. +/// (`RewriteColumnsExecutor`/`RewriteColumnsCommitter`) compaction run through +/// it. The two pipelines still differ in how they produce tasks and which +/// operation they commit; only the execute-collect-commit skeleton is shared. pub(crate) async fn run_compaction_pipeline( dataset: &mut Dataset, executor: E, @@ -2604,12 +2616,21 @@ fn column_group_schemas( groups: &[Vec], ) -> Result> { let names = || schema.fields.iter().map(|field| field.name.as_str()); + let mut seen = HashSet::new(); for column in groups.iter().flatten() { if !names().any(|name| name == column) { return Err(Error::invalid_input(format!( "column_groups names \"{column}\", which is not a top-level column of the dataset" ))); } + // Checked here as well as in `validate()`: a custom planner or a + // hand-built `CompactionTask` reaches `rewrite_files` without it, and a + // column in two groups would be written to two files. + if !seen.insert(column.as_str()) { + return Err(Error::invalid_input(format!( + "column_groups lists column \"{column}\" more than once" + ))); + } } let claimed = |name: &str| groups.iter().flatten().any(|column| column == name); let rest: Vec<&str> = names().filter(|name| !claimed(name)).collect(); @@ -2704,18 +2725,18 @@ async fn write_column_group_fragments( } } + // Returns which writer stopped reading first, if one did: a writer that + // fails closes its receiver, and its error is the one to report. let forward = async move { while let Some(batch) = reader.next().await { let batch = batch?; - for (sender, projection) in senders.iter_mut().zip(&projections) { - // A writer that stopped early (its own error) closes its - // receiver; stop forwarding and let try_join surface that error. + for (index, (sender, projection)) in senders.iter_mut().zip(&projections).enumerate() { if sender.send(Ok(batch.project(projection)?)).await.is_err() { - return Ok(()); + return Ok(Some(index)); } } } - Ok::<(), Error>(()) + Ok::, Error>(None) }; let writers = group_streams .into_iter() @@ -2742,21 +2763,57 @@ async fn write_column_group_fragments( }) .collect::>(); - let (_, per_group) = futures::try_join!(forward, futures::future::try_join_all(writers))?; - - let mut per_group = per_group.into_iter(); - let mut fragments = per_group.next().unwrap_or_default(); - for group in per_group { - if group.len() != fragments.len() - || fragments + // Run every writer to completion rather than dropping the others on the + // first error: a writer cleans up its own files only when it fails, so a + // group that finished must be cleaned up here if the write as a whole fails. + let (forwarded, per_group) = futures::join!(forward, futures::future::join_all(writers)); + let mut written = Vec::with_capacity(per_group.len()); + let mut errors = Vec::new(); + for (index, result) in per_group.into_iter().enumerate() { + match result { + Ok(fragments) => written.push(fragments), + Err(err) => errors.push((index, err)), + } + } + let failure = match forwarded { + // The read failed; the writers' errors follow from it. + Err(err) => Some(err), + Ok(Some(stopped)) => Some( + match errors.into_iter().find(|(index, _)| *index == stopped) { + Some((_, err)) => err, + // Not expected: a writer only finishes once its input ends, + // which needs the forwarder to have dropped every sender. + None => Error::internal(format!( + "column group writer {stopped} stopped reading before the compacted rows ended" + )), + }, + ), + Ok(None) => errors.into_iter().next().map(|(_, err)| err), + }; + let misaligned = written.iter().skip(1).any(|group| { + group.len() != written[0].len() + || written[0] .iter() - .zip(&group) + .zip(group) .any(|(a, b)| a.physical_rows != b.physical_rows) - { - return Err(Error::internal( + }); + let failure = failure.or_else(|| { + misaligned.then(|| { + Error::internal( "column group writers did not split the compacted rows at the same fragments", - )); + ) + }) + }); + if let Some(err) = failure { + for fragments in &written { + cleanup_data_fragments(&dataset.object_store, &dataset.base, None, fragments).await; } + return Err(err); + } + + let mut written = written.into_iter(); + let mut fragments = written.next().unwrap_or_default(); + for group in written { for (fragment, other) in fragments.iter_mut().zip(group) { fragment.files.extend(other.files); } @@ -12519,4 +12576,31 @@ mod tests { let err = options.validate().unwrap_err(); assert!(matches!(err, Error::InvalidInput { .. }), "{err}"); } + + /// `rewrite_files` builds its group schemas without `validate()` (a custom + /// planner or a hand-built `CompactionTask`), so the duplicate check has to + /// hold there too. + #[test] + fn column_group_schemas_reject_duplicate() { + let schema = lance_core::datatypes::Schema::try_from(&Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, true), + ])) + .unwrap(); + let err = column_group_schemas(&schema, &[vec!["b".into()], vec!["b".into()]]).unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. }), "{err}"); + assert!(err.to_string().contains("more than once"), "{err}"); + } + + #[test] + fn compaction_column_groups_validate_rejects_force_binary_copy() { + let mut options = CompactionOptions { + column_groups: vec![vec!["b".into()]], + compaction_mode: Some(CompactionMode::ForceBinaryCopy), + ..Default::default() + }; + let err = options.validate().unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. }), "{err}"); + assert!(err.to_string().contains("ForceBinaryCopy"), "{err}"); + } } diff --git a/rust/lance/src/dataset/rewrite_columns.rs b/rust/lance/src/dataset/rewrite_columns.rs index f7e6661d9a3..b258fd426c4 100644 --- a/rust/lance/src/dataset/rewrite_columns.rs +++ b/rust/lance/src/dataset/rewrite_columns.rs @@ -86,25 +86,35 @@ impl FileFragment { /// Rewrite `columns` of this fragment into one new data file. /// /// The returned metadata has the new file appended and the rewritten fields - /// tombstoned in the files they came from; a file left holding only - /// tombstones is dropped. Nothing is committed. Returns `None` when the - /// columns already sit alone in a file of the requested version, so a - /// retried or resumed rewrite skips finished work. + /// tombstoned in the files they came from; a file left holding no column + /// of the schema is dropped, unless it carries the fragment's spilled row + /// lineage. Nothing is committed: commit the result with + /// [`RewriteColumnsCommitter`]. Returns `None` when the columns already sit + /// alone in a file of the requested version, so a retried or resumed + /// rewrite skips finished work. pub async fn rewrite_columns( &self, columns: &[&str], data_storage_version: Option, - ) -> Result> { + ) -> Result> { let dataset = self.dataset(); let write_schema = rewrite_schema(dataset, columns)?; let resolved_version = write_version(dataset, data_storage_version)?; let field_ids: HashSet = write_schema.field_ids().into_iter().collect(); + // Ids of columns `drop_columns` removed stay in their files, and a + // spilled row lineage sequence sits in a data file under a reserved id, + // so a file's ids are only compared and kept by what the schema reads. + let schema_ids: HashSet = dataset + .schema() + .fields_pre_order() + .map(|field| field.id) + .collect(); let already_split = self.metadata().files.iter().any(|file| { file.fields .iter() .copied() - .filter(|id| *id != TOMBSTONE_FIELD_ID) + .filter(|id| schema_ids.contains(id)) .collect::>() == field_ids && file @@ -158,11 +168,22 @@ impl FileFragment { .collect::>() .into(); } - // A file holding only tombstones is unreachable to readers. - fragment - .files - .retain(|file| file.fields.iter().any(|id| *id != TOMBSTONE_FIELD_ID)); - Ok(Some(fragment)) + // Drop a file left with nothing the schema reads, the rule + // `drop_columns` applies, but keep one carrying the fragment's spilled + // row lineage: it is the only copy. + let spilled = fragment.spilled_row_lineage_field_ids(); + fragment.files.retain(|file| { + file.fields + .iter() + .any(|id| schema_ids.contains(id) || spilled.contains(id)) + }); + // The new file is the last one and is always kept. + let files_removed = self.metadata().files.len() + 1 - fragment.files.len(); + Ok(Some(RewriteColumnsResult { + fragment, + read_version: dataset.manifest.version, + files_removed, + })) } } /// One fragment's horizontal-rewrite output, bound to the manifest version it @@ -175,6 +196,8 @@ pub struct RewriteColumnsResult { pub fragment: Fragment, /// The manifest version the fragment was read at. pub read_version: u64, + /// How many of the fragment's data files the rewrite dropped. + pub files_removed: usize, } /// Horizontal executor: rewrites one fragment's named columns into a new data @@ -204,23 +227,17 @@ impl CompactionExecutor for RewriteColumnsExecutor { type TaskResult = Option; /// The task is a [`FileFragment`] already bound to its read version, so the - /// `dataset` argument the pipeline threads in is unused; the read version is - /// taken from the fragment and carried into the result for the committer. + /// `dataset` argument the pipeline threads in is unused; the result carries + /// the fragment's read version for the committer. async fn execute( &self, _dataset: Cow<'_, Dataset>, task: FileFragment, _options: &CompactionOptions, ) -> Result> { - let read_version = task.dataset().manifest.version; let columns: Vec<&str> = self.columns.iter().map(String::as_str).collect(); - Ok(task - .rewrite_columns(&columns, self.data_storage_version) - .await? - .map(|fragment| RewriteColumnsResult { - fragment, - read_version, - })) + task.rewrite_columns(&columns, self.data_storage_version) + .await } } /// Commits horizontal-compaction results as one `Operation::Update` in @@ -252,34 +269,17 @@ impl CompactionCommitter for RewriteColumnsCommitter { .map(|result| result.read_version) .min() .expect("results is non-empty"); - let updated_fragments: Vec = - results.into_iter().map(|result| result.fragment).collect(); // Row addresses and values are unchanged, so no fragment is added or - // removed; only per-fragment files change. `files_added` is one new - // file per rewritten fragment; `files_removed` is the fully-tombstoned - // files that were dropped. - let mut metrics = CompactionMetrics { - files_added: updated_fragments.len(), + // removed; only per-fragment files change: one new file per rewritten + // fragment, and the files each rewrite dropped. A fragment changed or + // removed since its read is left to the commit's conflict check. + let metrics = CompactionMetrics { + files_added: results.len(), + files_removed: results.iter().map(|result| result.files_removed).sum(), ..CompactionMetrics::default() }; - for fragment in &updated_fragments { - // A fragment being updated must exist in the version we commit - // against; a missing one is an inconsistency (e.g. a distributed - // caller passing a stale fragment id), not a zero-file fragment. - let original_files = dataset - .get_fragment(fragment.id as usize) - .ok_or_else(|| { - Error::internal(format!( - "rewrite_columns is updating fragment {} that is not in the \ - current manifest", - fragment.id - )) - })? - .metadata() - .files - .len(); - metrics.files_removed += (original_files + 1).saturating_sub(fragment.files.len()); - } + let updated_fragments: Vec = + results.into_iter().map(|result| result.fragment).collect(); let transaction = Transaction::new( read_version, Operation::Update { @@ -311,7 +311,12 @@ impl Dataset { /// ([`RewriteColumnsExecutor`] + [`RewriteColumnsCommitter`]) and commits /// once. To spread the work over many machines, call /// [`FileFragment::rewrite_columns`] per fragment and commit the returned - /// metadata with [`RewriteColumnsCommitter`]. + /// results with [`RewriteColumnsCommitter`]. + /// + /// The new files carry only the rewritten columns. On a table with stable + /// row ids whose row lineage was spilled into a data file, that file is + /// kept for the lineage even after every column has moved out of it, until + /// a vertical compaction rewrites the fragment. pub async fn rewrite_columns( &mut self, columns: &[&str], @@ -542,8 +547,17 @@ mod tests { dataset.validate().await.unwrap(); } + /// A delete that lands on the fragment between the rewrite's read and its + /// commit, whether it removes one row or the whole fragment, must surface + /// as a retryable conflict. + #[rstest::rstest] + #[case::one_row("a = 1", 7)] + #[case::whole_fragment("a < 4", 4)] #[tokio::test] - async fn rewrite_columns_conflicts_with_intervening_delete() { + async fn rewrite_columns_conflicts_with_intervening_delete( + #[case] predicate: &str, + #[case] rows_left: usize, + ) { use crate::dataset::index::DatasetIndexRemapperOptions; use crate::dataset::optimize::{ CompactionCommitter, CompactionExecutor, CompactionOptions, @@ -566,31 +580,32 @@ mod tests { .expect("fragment 0 has columns to repack"); // A delete lands on that same fragment before the rewrite is committed. - dataset.delete("a = 1").await.unwrap(); + dataset.delete(predicate).await.unwrap(); let after_delete = dataset.scan().try_into_batch().await.unwrap(); - assert_eq!(after_delete.num_rows(), 7); + assert_eq!(after_delete.num_rows(), rows_left); // The rewrite was read before the delete, so committing it must conflict // rather than publish the stale layout and resurrect the deleted row. - let committed = RewriteColumnsCommitter + let err = RewriteColumnsCommitter .commit( &mut dataset, vec![Some(stale)], Arc::new(DatasetIndexRemapperOptions::default()), &CompactionOptions::default(), ) - .await; + .await + .unwrap_err(); assert!( - committed.is_err(), - "a rewrite read before the delete must not commit over it" + matches!(err, Error::RetryableCommitConflict { .. }), + "a rewrite read before the delete must not commit over it: {err}" ); - // The delete stands: seven rows, none with a = 1. + // The delete stands. assert_eq!(dataset.scan().try_into_batch().await.unwrap(), after_delete); assert_eq!( dataset .scan() - .filter("a = 1") + .filter(predicate) .unwrap() .try_into_batch() .await @@ -600,4 +615,200 @@ mod tests { ); dataset.validate().await.unwrap(); } + + /// `drop_columns` leaves a dropped column's id in the file that held it. + /// Rewriting the rest of that file's columns leaves it with nothing the + /// schema reads, so the rewrite must drop it, as `drop_columns` would. + #[tokio::test] + async fn rewrite_drops_file_left_with_only_dropped_columns() { + use crate::dataset::NewColumnTransform; + let mut dataset = write("memory://", LanceFileVersion::V2_0).await; + dataset + .add_columns( + NewColumnTransform::SqlExpressions(vec![("d".into(), "a + 1".into())]), + None, + None, + ) + .await + .unwrap(); + dataset.drop_columns(&["c"]).await.unwrap(); + let before = dataset.scan().try_into_batch().await.unwrap(); + + dataset + .rewrite_columns(&["a", "b", "d"], None) + .await + .unwrap(); + + for fragment in dataset.get_fragments() { + assert_eq!( + fragment.metadata().files.len(), + 1, + "only the new file is left: {:?}", + fragment.metadata().files + ); + } + assert_eq!(dataset.scan().try_into_batch().await.unwrap(), before); + dataset.validate().await.unwrap(); + } + + /// A file whose only other column was dropped already holds exactly the + /// requested columns, so there is nothing to rewrite. + #[tokio::test] + async fn rewrite_skips_file_whose_other_columns_were_dropped() { + let mut dataset = write("memory://", LanceFileVersion::V2_0).await; + dataset.drop_columns(&["c"]).await.unwrap(); + let version = dataset.manifest.version; + + dataset.rewrite_columns(&["a", "b"], None).await.unwrap(); + + assert_eq!(dataset.manifest.version, version, "nothing was committed"); + dataset.validate().await.unwrap(); + } + + /// A `drop_columns` that commits between a rewrite's read and its commit + /// would leave a file holding only the dropped column if the rewrite went + /// through, whether that is the rewritten column's new file or a file the + /// drop removed, so the commit must conflict. When the rewritten column + /// survived the drop, rerunning the rewrite on the new version succeeds. + #[rstest::rstest] + #[case::rewritten_column("c")] + #[case::other_column("d")] + #[tokio::test] + async fn rewrite_columns_conflicts_with_concurrent_drop(#[case] dropped: &str) { + use crate::dataset::NewColumnTransform; + use crate::dataset::index::DatasetIndexRemapperOptions; + use crate::dataset::optimize::{ + CompactionCommitter, CompactionExecutor, CompactionOptions, + }; + use std::borrow::Cow; + + let mut dataset = write("memory://", LanceFileVersion::V2_0).await; + // `d` gets a file of its own in every fragment. + dataset + .add_columns( + NewColumnTransform::SqlExpressions(vec![("d".into(), "a + 1".into())]), + None, + None, + ) + .await + .unwrap(); + + let executor = RewriteColumnsExecutor::new(&["c"], None); + let mut stale = Vec::new(); + for fragment in dataset.get_fragments() { + stale.push( + executor + .execute( + Cow::Borrowed(&dataset), + fragment, + &CompactionOptions::default(), + ) + .await + .unwrap(), + ); + } + + dataset.drop_columns(&[dropped]).await.unwrap(); + let version = dataset.manifest.version; + let expected = dataset.scan().try_into_batch().await.unwrap(); + + let err = RewriteColumnsCommitter + .commit( + &mut dataset, + stale, + Arc::new(DatasetIndexRemapperOptions::default()), + &CompactionOptions::default(), + ) + .await + .unwrap_err(); + assert!( + matches!(err, Error::RetryableCommitConflict { .. }), + "{err}" + ); + dataset.checkout_latest().await.unwrap(); + assert_eq!(dataset.manifest.version, version, "nothing was committed"); + dataset.validate().await.unwrap(); + + if dropped != "c" { + dataset.rewrite_columns(&["c"], None).await.unwrap(); + assert_eq!(dataset.scan().try_into_batch().await.unwrap(), expected); + dataset.validate().await.unwrap(); + } + } + + /// The concurrent-drop check must not trip on a file holding only + /// tombstones, which an in-place merge_insert leaves in its post-image for + /// the commit to drop: a rename racing such a commit has to go through. + #[tokio::test] + async fn rewrite_columns_commit_ignores_tombstoned_file_on_rename() { + use crate::dataset::index::DatasetIndexRemapperOptions; + use crate::dataset::optimize::{ + CompactionCommitter, CompactionExecutor, CompactionOptions, + }; + use crate::dataset::{ColumnAlteration, NewColumnTransform}; + use std::borrow::Cow; + + let mut dataset = write("memory://", LanceFileVersion::V2_0).await; + // `d` gets a file of its own in every fragment. + dataset + .add_columns( + NewColumnTransform::SqlExpressions(vec![("d".into(), "a + 1".into())]), + None, + None, + ) + .await + .unwrap(); + + // `d` already sits alone, so repack it together with `c`. + let executor = RewriteColumnsExecutor::new(&["c", "d"], None); + let mut results = Vec::new(); + for fragment in dataset.get_fragments() { + // The file `d` lived in, fully tombstoned, as merge_insert leaves it. + let mut tombstoned = fragment.metadata().files[1].clone(); + tombstoned.fields = vec![TOMBSTONE_FIELD_ID].into(); + let mut result = executor + .execute( + Cow::Borrowed(&dataset), + fragment, + &CompactionOptions::default(), + ) + .await + .unwrap() + .expect("c and d have files to repack"); + result.fragment.files.insert(1, tombstoned); + results.push(Some(result)); + } + + dataset + .alter_columns(&[ColumnAlteration::new("b".into()).rename("b2".into())]) + .await + .unwrap(); + let expected = dataset.scan().try_into_batch().await.unwrap(); + + RewriteColumnsCommitter + .commit( + &mut dataset, + results, + Arc::new(DatasetIndexRemapperOptions::default()), + &CompactionOptions::default(), + ) + .await + .unwrap(); + + for fragment in dataset.get_fragments() { + let layout: Vec> = fragment + .metadata() + .files + .iter() + .map(|file| file.fields.to_vec()) + .collect(); + assert_eq!( + layout, + vec![vec![0, 1, TOMBSTONE_FIELD_ID], vec![2, 3]], + "the rewrite landed and the commit dropped the tombstoned file" + ); + } + assert_eq!(dataset.scan().try_into_batch().await.unwrap(), expected); + dataset.validate().await.unwrap(); + } } diff --git a/rust/lance/src/index/append.rs b/rust/lance/src/index/append.rs index 7fe4b7feac5..934ef6ca2dc 100644 --- a/rust/lance/src/index/append.rs +++ b/rust/lance/src/index/append.rs @@ -243,6 +243,10 @@ pub async fn build_per_segment_filters( /// Attempt to read seed buffers for `column_name` from `fragments`' data files. /// +/// A seed is written into the data file that holds the column, which after a +/// `column_groups` compaction need not be the fragment's first file, so each +/// fragment's file is found by `field_id`. +/// /// Returns `Some(vec)` only if every fragment has a seed entry; returns `None` /// if any fragment is missing a seed or its data file cannot be opened. /// Index-type-specific validation (e.g. `rows_per_zone` checks) is left to the @@ -251,6 +255,7 @@ async fn try_harvest_seeds( dataset: &Dataset, fragments: &[Fragment], column_name: &str, + field_id: i32, ) -> Result>> { if fragments.is_empty() { return Ok(Some(Vec::new())); @@ -265,7 +270,11 @@ async fn try_harvest_seeds( ); for fragment in fragments { - let Some(data_file) = fragment.files.first() else { + let Some(data_file) = fragment + .files + .iter() + .find(|file| file.fields.contains(&field_id)) + else { return Ok(None); }; @@ -534,7 +543,13 @@ async fn merge_scalar_indices<'a>( // Only open data files looking for seeds when the plugin confirms this // index type and configuration can actually produce them. let maybe_created = if plugin.might_use_seeds(&index_details) { - if let Some(seeds) = try_harvest_seeds(dataset.as_ref(), unindexed, column_name).await? + if let Some(seeds) = try_harvest_seeds( + dataset.as_ref(), + unindexed, + column_name, + reference_idx.fields[0], + ) + .await? { plugin .update_from_seeds(seeds, reference_index.clone(), &index_details, &new_store) @@ -4941,4 +4956,62 @@ mod tests { .num_rows(); assert_eq!(total, 150, "no rows may be lost across compaction + merge"); } + + /// A `column_groups` compaction writes an indexed column's seed into that + /// column's group file, which need not be the fragment's first file; the + /// harvest has to find it there. + #[tokio::test] + async fn test_harvest_seeds_from_column_group_file() { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from_iter_values(0..20)), + Arc::new(Int32Array::from_iter_values(0..20)), + ], + ) + .unwrap(); + let reader = RecordBatchIterator::new([Ok(batch)], schema); + let mut dataset = Dataset::write( + reader, + "memory://", + Some(WriteParams { + max_rows_per_file: 10, + ..Default::default() + }), + ) + .await + .unwrap(); + let params = ScalarIndexParams::for_builtin(BuiltinIndexType::ZoneMap) + .with_params(&serde_json::json!({"use_seeds": true})); + dataset + .create_index(&["a"], IndexType::ZoneMap, None, ¶ms, false) + .await + .unwrap(); + // {b} is unclaimed and becomes the first file; {a} is the second. + let options = CompactionOptions { + column_groups: vec![vec!["a".into()]], + ..Default::default() + }; + compact_files(&mut dataset, options, None).await.unwrap(); + + let fragments: Vec = dataset + .get_fragments() + .iter() + .map(|fragment| fragment.metadata().clone()) + .collect(); + let a_id = dataset.schema().field("a").unwrap().id; + assert!( + !fragments[0].files[0].fields.contains(&a_id), + "the indexed column is not in the first file: {:?}", + fragments[0].files + ); + let seeds = try_harvest_seeds(&dataset, &fragments, "a", a_id) + .await + .unwrap(); + assert!(seeds.is_some(), "the seed in the group file is found"); + } } diff --git a/rust/lance/src/io/commit/conflict_resolver.rs b/rust/lance/src/io/commit/conflict_resolver.rs index 08f978f6a9d..8a26c538946 100644 --- a/rust/lance/src/io/commit/conflict_resolver.rs +++ b/rust/lance/src/io/commit/conflict_resolver.rs @@ -18,7 +18,7 @@ use lance_index::frag_reuse::FRAG_REUSE_INDEX_NAME; use lance_index::mem_wal::{CompactedSsTable, MEM_WAL_INDEX_NAME}; use lance_select::{RowAddrTreeMap, RowSetOps}; use lance_table::format::IndexMetadata; -use lance_table::format::overlay::OverlayCoverage; +use lance_table::format::overlay::{OverlayCoverage, TOMBSTONE_FIELD_ID}; use lance_table::format::pb::fragment_reuse_index_details::{InlineContent, Transition}; use lance_table::system_index::frag_reuse::lineage::TaggedLineage; use lance_table::system_index::frag_reuse::metadata::is_tagged; @@ -648,6 +648,34 @@ impl<'a> TransactionRebase<'a> { ); } } + // An in-place column rewrite publishes each fragment's file + // list as it built it at the read version, so a file the + // projection removed (it held only dropped fields) would come + // back, and so would a new file holding only a dropped field. + // `fields_modified` can't catch this: a rewrite that moves + // values without changing them (horizontal compaction) leaves + // it empty. Apply the projection's own keep rule to our files. + // A file holding only tombstones is skipped: an in-place + // merge_insert leaves one behind and the commit drops it. + if matches!(self_update_mode, Some(UpdateMode::RewriteColumns)) { + let remaining: HashSet = + schema.fields_pre_order().map(|field| field.id).collect(); + let drops_a_file = self_updated_fragments.iter().any(|fragment| { + let spilled = fragment.spilled_row_lineage_field_ids(); + fragment.files.iter().any(|file| { + file.fields.iter().any(|id| *id != TOMBSTONE_FIELD_ID) + && !file + .fields + .iter() + .any(|id| remaining.contains(id) || spilled.contains(id)) + }) + }); + if drops_a_file { + return Err( + self.retryable_conflict_err(other_transaction, other_version) + ); + } + } Ok(()) } Operation::DataOverlay { groups } => { @@ -1027,7 +1055,7 @@ impl<'a> TransactionRebase<'a> { fragment.files.iter().any(|file| { file.fields.iter().any(|field| { *field - == lance_table::format::overlay::TOMBSTONE_FIELD_ID + == TOMBSTONE_FIELD_ID }) }) }) From f8b46baade62e9d9f7331b8f23f6d5b86e5e4f5e Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Fri, 2 Oct 2026 11:16:35 +0800 Subject: [PATCH 07/24] fix(compaction): count only files holding schema columns in column_layout_stats --- rust/lance/src/dataset/compaction_stats.rs | 69 ++++++++++++++++++++-- 1 file changed, 64 insertions(+), 5 deletions(-) diff --git a/rust/lance/src/dataset/compaction_stats.rs b/rust/lance/src/dataset/compaction_stats.rs index c7e41ef7dc3..462d49a1bd8 100644 --- a/rust/lance/src/dataset/compaction_stats.rs +++ b/rust/lance/src/dataset/compaction_stats.rs @@ -9,7 +9,7 @@ //! Exposing these stats via [`Dataset::column_layout_stats`] keeps the planning //! decision observable instead of buried in a black-box heuristic. -use lance_table::format::overlay::TOMBSTONE_FIELD_ID; +use std::collections::HashSet; use super::Dataset; @@ -18,9 +18,11 @@ use super::Dataset; pub struct FragmentColumnLayoutStats { /// The fragment these stats describe. pub fragment_id: u64, - /// Number of data files still holding at least one live (non-tombstoned) - /// field. A large value on a wide dataset is what horizontal compaction - /// collapses back into fewer files. + /// Number of data files holding at least one column of the dataset schema. + /// A large value on a wide dataset is what horizontal compaction collapses + /// back into fewer files. A file left holding only tombstones or dropped + /// columns, or kept only for the fragment's spilled row lineage, holds no + /// column a read touches and is not counted. pub live_file_count: usize, /// Number of overlay files attached to the fragment. pub overlay_count: usize, @@ -32,6 +34,11 @@ impl Dataset { /// This is the planning input for horizontal compaction: it reads only /// fragment metadata (no data files), so it is cheap to call. pub fn column_layout_stats(&self) -> Vec { + let schema_ids: HashSet = self + .schema() + .fields_pre_order() + .map(|field| field.id) + .collect(); self.manifest .fragments .iter() @@ -40,7 +47,7 @@ impl Dataset { live_file_count: fragment .files .iter() - .filter(|file| file.fields.iter().any(|id| *id != TOMBSTONE_FIELD_ID)) + .filter(|file| file.fields.iter().any(|id| schema_ids.contains(id))) .count(), overlay_count: fragment.overlays.len(), }) @@ -106,4 +113,56 @@ mod tests { "each fragment should now have 2 data files: {stats:?}" ); } + + /// Only files holding a schema column count. The extra files are added to + /// the manifest by hand: the commit paths that leave such files behind (an + /// in-place column update after a drop, a full horizontal rewrite on a + /// table whose lineage was spilled into a data file) are not what this + /// test is about. + #[tokio::test] + async fn column_layout_stats_skips_files_without_schema_columns() { + use lance_file::version::ConcreteFileVersion; + use lance_table::format::{ + DataFile, ROW_ID_FIELD_ID, RowIdMeta, overlay::TOMBSTONE_FIELD_ID, + }; + + let schema = Arc::new(ArrowSchema::new(vec![Field::new( + "a", + DataType::Int32, + false, + )])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from_iter_values(0..4))], + ) + .unwrap(); + let reader = RecordBatchIterator::new([Ok(batch)], schema); + let mut dataset = Dataset::write(reader, "memory://", None).await.unwrap(); + + let dropped_id = dataset.schema().max_field_id().unwrap() + 1; + let mut manifest = dataset.manifest.as_ref().clone(); + let mut fragments = manifest.fragments.as_ref().clone(); + let file = |fields: Vec| { + let indices = (0..fields.len() as i32).collect(); + DataFile::new( + "extra.lance", + fields, + indices, + ConcreteFileVersion::V2_0, + None, + None, + ) + }; + fragments[0].files.push(file(vec![TOMBSTONE_FIELD_ID])); + fragments[0].files.push(file(vec![dropped_id])); + fragments[0] + .files + .push(file(vec![TOMBSTONE_FIELD_ID, ROW_ID_FIELD_ID])); + fragments[0].row_id_meta = Some(RowIdMeta::Column); + manifest.fragments = Arc::new(fragments); + dataset.manifest = Arc::new(manifest); + + let stats = dataset.column_layout_stats(); + assert_eq!(stats[0].live_file_count, 1, "{stats:?}"); + } } From 30888e86339af3ec58bc42e12667e5aba2a704dd Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Fri, 2 Oct 2026 11:51:40 +0800 Subject: [PATCH 08/24] fix(compaction): pass the new write_fragments_direct argument in the grouped write --- rust/lance/src/dataset/optimize.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/rust/lance/src/dataset/optimize.rs b/rust/lance/src/dataset/optimize.rs index ca0d36c7dc0..14146ea4548 100644 --- a/rust/lance/src/dataset/optimize.rs +++ b/rust/lance/src/dataset/optimize.rs @@ -2757,6 +2757,7 @@ async fn write_column_group_fragments( None, seed_writers, Some(file_row_counts), + None, ) .await } From 668d99955353c9b322fc381a2f0f4e72f9307a50 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Fri, 2 Oct 2026 20:24:19 +0800 Subject: [PATCH 09/24] feat(transaction): add data_change to DataReplacement and allow several groups per fragment --- java/lance-jni/src/transaction.rs | 20 +- .../org/lance/operation/DataReplacement.java | 55 +++- protos/transaction.proto | 6 + python/python/lance/dataset.py | 6 + python/src/transaction.rs | 13 +- .../src/system_index/frag_reuse/gate.rs | 1 + rust/lance-table/src/transaction/conflicts.rs | 12 +- .../src/transaction/index_maintenance.rs | 11 +- .../src/transaction/manifest_build.rs | 236 +++++++++++++----- rust/lance-table/src/transaction/operation.rs | 31 +-- rust/lance-table/src/transaction/proto.rs | 66 ++++- rust/lance/src/dataset/blob.rs | 2 + rust/lance/src/dataset/rowids/spill.rs | 1 + rust/lance/src/dataset/schema_evolution.rs | 1 + .../lance/src/dataset/tests/data_file_part.rs | 1 + rust/lance/src/dataset/tests/dataset_index.rs | 1 + .../src/dataset/tests/dataset_merge_update.rs | 12 + .../dataset/tests/fragment_write_columns.rs | 5 +- rust/lance/src/io/commit.rs | 3 +- rust/lance/src/io/commit/conflict_resolver.rs | 181 +++++++++++++- 20 files changed, 544 insertions(+), 120 deletions(-) diff --git a/java/lance-jni/src/transaction.rs b/java/lance-jni/src/transaction.rs index eb06d599f2d..40f0df3522a 100644 --- a/java/lance-jni/src/transaction.rs +++ b/java/lance-jni/src/transaction.rs @@ -1125,13 +1125,19 @@ fn convert_to_java_operation_inner<'local>( )?; Ok(java_operation) } - Operation::DataReplacement { replacements } => { + Operation::DataReplacement { + replacements, + data_change, + } => { let java_replacements = export_vec(env, &replacements)?; Ok(env.new_object( "org/lance/operation/DataReplacement", - "(Ljava/util/List;)V", - &[JValue::Object(&java_replacements)], + "(Ljava/util/List;Z)V", + &[ + JValue::Object(&java_replacements), + JValue::Bool(data_change as u8), + ], )?) } Operation::DataOverlay { groups } => { @@ -1998,7 +2004,13 @@ fn convert_to_rust_operation( import_vec_from_method(env, java_operation, "replacements", |env, replacement| { replacement.extract_object(env) })?; - Operation::DataReplacement { replacements } + let data_change = env + .call_method(java_operation, "dataChange", "()Z", &[])? + .z()?; + Operation::DataReplacement { + replacements, + data_change, + } } "DataOverlay" => { let groups = import_vec_from_method(env, java_operation, "getGroups", |env, group| { diff --git a/java/src/main/java/org/lance/operation/DataReplacement.java b/java/src/main/java/org/lance/operation/DataReplacement.java index 42f0c12615f..484149d2191 100644 --- a/java/src/main/java/org/lance/operation/DataReplacement.java +++ b/java/src/main/java/org/lance/operation/DataReplacement.java @@ -21,21 +21,22 @@ import java.util.Objects; /** - * Replace data in a column in the dataset with new data. This is used for null column population - * where we replace an entirely null column with a new column that has data. + * Replace the data files backing some fields of existing fragments with new files, without moving + * rows. Each group names a fragment and one new data file for it. At commit, a file holding exactly + * the new file's fields is swapped for it; otherwise the new file's fields are tombstoned where they + * live and the new file is appended. A fragment can take several groups, one per new file. Used for + * null column population, and by compaction to repack columns into fewer files. * - *

This operation will only allow replacing files that contain the same schema e.g. if the - * original files contain columns A, B, C and the new files contain only columns A, B then the - * operation is not allowed. - * - *

Corollary to the above: the operation will also not allow replacing files unless the affected - * columns all have the same datafile layout across the fragments being replaced. + *

{@code dataChange == false} declares that the new files hold the same values as the files they + * replace: indices keep their coverage, overlays keep shadowing, and no row is reported as updated. */ public class DataReplacement implements Operation { private final List replacements; + private final boolean dataChange; - private DataReplacement(List replacements) { + private DataReplacement(List replacements, boolean dataChange) { this.replacements = replacements; + this.dataChange = dataChange; } /** @@ -47,6 +48,15 @@ public List replacements() { return replacements; } + /** + * Whether the new files change any value. + * + * @return false if the values were only moved to new files + */ + public boolean dataChange() { + return dataChange; + } + @Override public String name() { return "DataReplacement"; @@ -54,7 +64,10 @@ public String name() { @Override public String toString() { - return MoreObjects.toStringHelper(this).add("replacements", replacements).toString(); + return MoreObjects.toStringHelper(this) + .add("replacements", replacements) + .add("dataChange", dataChange) + .toString(); } @Override @@ -62,7 +75,12 @@ public boolean equals(Object o) { if (this == o) return true; if (o == null || getClass() != o.getClass()) return false; DataReplacement that = (DataReplacement) o; - return Objects.equals(replacements, that.replacements); + return dataChange == that.dataChange && Objects.equals(replacements, that.replacements); + } + + @Override + public int hashCode() { + return Objects.hash(replacements, dataChange); } /** @@ -77,6 +95,7 @@ public static Builder builder() { /** Builder for DataReplacement. */ public static class Builder { private List replacements; + private boolean dataChange = true; public Builder() {} @@ -91,13 +110,25 @@ public Builder replacements(List replacements) { return this; } + /** + * Set whether the new files change any value. Defaults to true; pass false only when the new + * files hold the same values as the files they replace. + * + * @param dataChange whether the new files change any value + * @return this builder + */ + public Builder dataChange(boolean dataChange) { + this.dataChange = dataChange; + return this; + } + /** * Build a new DataReplacement. * * @return a new DataReplacement */ public DataReplacement build() { - return new DataReplacement(replacements); + return new DataReplacement(replacements, dataChange); } } diff --git a/protos/transaction.proto b/protos/transaction.proto index 6e01831b4af..e3e626f62be 100644 --- a/protos/transaction.proto +++ b/protos/transaction.proto @@ -365,6 +365,12 @@ message Transaction { // An operation that replaces the data in a region of the table with new data. message DataReplacement { repeated DataReplacementGroup replacements = 1; + /* Whether the new files change any value. Absent means true, which is what + * a writer that predates this field meant. False means the values were only + * moved to new files: index coverage, overlays, and row update versions are + * left as they are. + */ + optional bool data_change = 2; } /* Overlay files to append to a single fragment, in order (the last entry is diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index 13af2bbe19f..d068ba6d9d9 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -6603,9 +6603,15 @@ class DataReplacementGroup: class DataReplacement(BaseOperation): """ Operation that replaces existing datafiles in the dataset. + + ``data_change=False`` declares that the new files hold the same values + as the files they replace (compaction moved them). Indices then keep + their coverage, overlays keep shadowing, and no row is reported as + updated. """ replacements: List[LanceOperation.DataReplacementGroup] + data_change: bool = True @dataclass class DataOverlayFile: diff --git a/python/src/transaction.rs b/python/src/transaction.rs index b3f2311b3f4..267353d0f0d 100644 --- a/python/src/transaction.rs +++ b/python/src/transaction.rs @@ -583,8 +583,12 @@ impl FromPyObject<'_, '_> for PyLance { } "DataReplacement" => { let replacements = extract_vec(&ob.getattr("replacements")?)?; + let data_change = ob.getattr("data_change")?.extract::()?; - let op = Operation::DataReplacement { replacements }; + let op = Operation::DataReplacement { + replacements, + data_change, + }; Ok(Self(op)) } @@ -751,12 +755,15 @@ impl<'py> IntoPyObject<'py> for PyLance<&Operation> { updated_fragment_offsets, )) } - Operation::DataReplacement { replacements } => { + Operation::DataReplacement { + replacements, + data_change, + } => { let replacements = export_vec(py, replacements.as_slice())?; let cls = namespace .getattr("DataReplacement") .expect("Failed to get DataReplacement class"); - cls.call1((replacements,)) + cls.call1((replacements, *data_change)) } Operation::DataOverlay { groups } => { let groups = export_vec(py, groups.as_slice())?; diff --git a/rust/lance-table/src/system_index/frag_reuse/gate.rs b/rust/lance-table/src/system_index/frag_reuse/gate.rs index 42d953532c7..04ddbbc3b81 100644 --- a/rust/lance-table/src/system_index/frag_reuse/gate.rs +++ b/rust/lance-table/src/system_index/frag_reuse/gate.rs @@ -373,6 +373,7 @@ mod tests { None, ), )], + data_change: true, }, "data_overlay" => Operation::DataOverlay { groups: vec![DataOverlayGroup { diff --git a/rust/lance-table/src/transaction/conflicts.rs b/rust/lance-table/src/transaction/conflicts.rs index cba1b5fb547..47a8c868621 100644 --- a/rust/lance-table/src/transaction/conflicts.rs +++ b/rust/lance-table/src/transaction/conflicts.rs @@ -198,9 +198,15 @@ impl PartialEq for Operation { && a_field == b_field } ( - Self::DataReplacement { replacements: a }, - Self::DataReplacement { replacements: b }, - ) => a.len() == b.len() && a.iter().all(|r| b.contains(r)), + Self::DataReplacement { + replacements: a, + data_change: a_change, + }, + Self::DataReplacement { + replacements: b, + data_change: b_change, + }, + ) => a_change == b_change && a.len() == b.len() && a.iter().all(|r| b.contains(r)), // Handle all remaining combinations. // We spell out all combinations explicitly to prevent // us accidentally handling a new case in the wrong way. diff --git a/rust/lance-table/src/transaction/index_maintenance.rs b/rust/lance-table/src/transaction/index_maintenance.rs index 6c43eb58e51..d1cd681338a 100644 --- a/rust/lance-table/src/transaction/index_maintenance.rs +++ b/rust/lance-table/src/transaction/index_maintenance.rs @@ -98,8 +98,9 @@ impl Transaction { /// id, unexpanded: an update with `fields_modified` rewrites those /// fields in every updated fragment; a merge rewrites, in each fragment /// present in `previous_fragments`, the fields whose backing data file - /// changed (`merge_rewritten_fields`); a data replacement rewrites the - /// fields its new files carry, read through `schema`. Any other + /// changed (`merge_rewritten_fields`); a data replacement that changes + /// data rewrites the fields its new files carry, read through `schema`. + /// Any other /// operation rewrites nothing. `previous_fragments` is the caller's /// "before" list: the current manifest's for a commit, the read /// version's for a rebase. @@ -120,7 +121,11 @@ impl Transaction { Operation::Merge { fragments, .. } => { Self::merge_rewritten_fields(previous_fragments, fragments) } - Operation::DataReplacement { replacements } => replacements + // Values that only moved to new files invalidate no index. + Operation::DataReplacement { + replacements, + data_change: true, + } => replacements .iter() .map(|DataReplacementGroup(fragment_id, new_file)| { let mut fields: Vec = new_file diff --git a/rust/lance-table/src/transaction/manifest_build.rs b/rust/lance-table/src/transaction/manifest_build.rs index ee16618fb1f..976703f55e4 100644 --- a/rust/lance-table/src/transaction/manifest_build.rs +++ b/rust/lance-table/src/transaction/manifest_build.rs @@ -1247,7 +1247,10 @@ impl Transaction { Operation::Restore { .. } => { unreachable!() } - Operation::DataReplacement { replacements } => { + Operation::DataReplacement { + replacements, + data_change, + } => { log::warn!( "Building manifest with DataReplacement operation. This operation is not stable yet, please use with caution." ); @@ -1257,56 +1260,24 @@ impl Transaction { .map(|DataReplacementGroup(fragment_id, new_file)| (fragment_id, new_file)) .unzip(); - // 1. make sure the new files all have the same fields / or empty - // NOTE: arguably this requirement could be relaxed in the future - // for the sake of simplicity, we require the new files to have the same fields - if new_datafiles - .iter() - .map(|f| f.fields.clone()) - .collect::>() - .len() - > 1 - { - let field_info = new_datafiles - .iter() - .enumerate() - .map(|(id, f)| (id, f.fields.clone())) - .fold("".to_string(), |acc, (id, fields)| { - format!("{}File {}: {:?}\n", acc, id, fields) - }); - - return Err(Error::invalid_input(format!( - "All new data files must have the same fields, but found different fields:\n{field_info}" - ))); - } - let existing_fragments = maybe_existing_fragments?; - // Collect replaced field IDs before consuming new_datafiles - let replaced_fields: Vec = new_datafiles - .first() - .map(|f| { - f.schema(&schema) - .field_ids() - .iter() - .chain(f.fields.iter()) - .filter(|&&id| id >= 0) - .map(|&id| id as u32) - .collect() - }) - .unwrap_or_default(); - - // 2. check that the fragments being modified have isomorphic layouts along the columns being replaced - // 3. add modified fragments to final_fragments + // Apply the groups in order. Several groups may name the same + // fragment (one new file each), so each applies to the + // fragment as the previous ones left it. + let mut replaced: HashMap = HashMap::new(); for (frag_id, new_file) in old_fragment_ids.iter().zip(new_datafiles) { - let frag = existing_fragments - .iter() - .find(|f| f.id == **frag_id) - .ok_or_else(|| { - Error::invalid_input( - "Fragment being replaced not found in existing fragments", - ) - })?; + let frag = match replaced.get(*frag_id) { + Some(frag) => frag, + None => existing_fragments + .iter() + .find(|f| f.id == **frag_id) + .ok_or_else(|| { + Error::invalid_input( + "Fragment being replaced not found in existing fragments", + ) + })?, + }; let mut new_frag = frag.clone(); // Physical mappings differ across V2 encodings (a nested @@ -1425,19 +1396,31 @@ impl Transaction { // value though -- the conflict resolver rebases these two // precisely because the overlay wins -- so it stays, and // being newer it stays last, preserving the ordering. - let (mut superseded, newer): (Vec<_>, Vec<_>) = new_frag - .overlays - .drain(..) - .partition(|overlay| overlay.committed_version <= self.read_version); - crate::format::overlay::tombstone_overlay_fields( - &mut superseded, - &replaced_fields, - ); - superseded.extend(newer); - new_frag.overlays = superseded; + // Values that only moved supersede nothing. + if *data_change { + let replaced_fields: Vec = new_file + .schema(&schema) + .field_ids() + .iter() + .chain(new_file.fields.iter()) + .filter(|&&id| id >= 0) + .map(|&id| id as u32) + .collect(); + let (mut superseded, newer): (Vec<_>, Vec<_>) = new_frag + .overlays + .drain(..) + .partition(|overlay| overlay.committed_version <= self.read_version); + crate::format::overlay::tombstone_overlay_fields( + &mut superseded, + &replaced_fields, + ); + superseded.extend(newer); + new_frag.overlays = superseded; + } - final_fragments.push(new_frag); + replaced.insert(**frag_id, new_frag); } + final_fragments.extend(replaced.into_values()); let fragments_changed = old_fragment_ids .iter() @@ -1457,7 +1440,7 @@ impl Transaction { // A replacement changes what its rows read as, so stamp them // updated. Without this, get_updated_rows never reports them and // an incremental consumer skips them for good. - if next_row_id.is_some() { + if next_row_id.is_some() && *data_change { let new_version = current_manifest.map_or(1, |m| m.version + 1); for fragment in final_fragments .iter_mut() @@ -2160,6 +2143,7 @@ mod tests { None, ), )], + data_change: true, }, _ => unreachable!(), } @@ -2867,6 +2851,33 @@ mod tests { ); } + /// A replacement whose values only moved to new files keeps every index's + /// coverage of the replaced fragment. + #[test] + fn prepare_indices_keeps_coverage_for_moved_values() { + let mut manifest = manifest_with_file("a.lance"); + manifest.reader_feature_flags &= !FLAG_FRAGMENT_REUSE_INDEX; + manifest.writer_feature_flags &= !FLAG_FRAGMENT_REUSE_INDEX; + let mut segment = sample_index_metadata("id_idx"); + segment.fragment_bitmap = Some([0u32, 3].into_iter().collect()); + let mut operation = in_place_rewrite("data_replacement", &manifest, vec![0]); + let Operation::DataReplacement { data_change, .. } = &mut operation else { + unreachable!() + }; + *data_change = false; + let transaction = Transaction::new(manifest.version, operation, None); + let prepared = transaction + .prepare_indices( + Some(&manifest), + vec![segment.clone()], + &default_build_config(), + None, + FragReuseUpdate::None, + ) + .unwrap(); + assert_eq!(prepared.prepared(), &[segment]); + } + /// Operations that rewrite no column in place, `Restore` included, /// prepare the list unchanged. #[test] @@ -4347,6 +4358,7 @@ mod tests { 0, DataFile::new_legacy_from_fields("f5-new.lance", vec![5], None), )], + data_change: true, }, None, ); @@ -4364,6 +4376,109 @@ mod tests { assert_eq!(frag.overlays[0].data_file.fields.as_ref(), &[3, -2]); } + #[test] + fn test_data_replacement_without_data_change_keeps_overlays() { + // Values that only moved to a new file supersede no overlay: both + // overlays keep shadowing the replaced field. + let mut fragment = Fragment::new(0); + fragment.files = vec![ + DataFile::new_legacy_from_fields("f3.lance", vec![3], None), + DataFile::new_legacy_from_fields("f5.lance", vec![5], None), + ]; + let overlay = |path: &str, fields: Vec| DataOverlayFile { + data_file: DataFile::new_legacy_from_fields(path, fields, None), + coverage: OverlayCoverage::dense(roaring::RoaringBitmap::from_iter([0u32])), + committed_version: 1, + }; + fragment.overlays = vec![overlay("o3.lance", vec![3]), overlay("o5.lance", vec![5])]; + let overlays = fragment.overlays.clone(); + + let schema = ArrowSchema::new(vec![ArrowField::new("id", DataType::Int32, false)]); + let manifest = Manifest::new( + LanceSchema::try_from(&schema).unwrap(), + Arc::new(vec![fragment]), + crate::format::DataStorageFormat::new(ConcreteFileVersion::V2_0), + HashMap::new(), + ); + let txn = Transaction::new( + manifest.version, + Operation::DataReplacement { + replacements: vec![DataReplacementGroup( + 0, + DataFile::new_legacy_from_fields("f5-new.lance", vec![5], None), + )], + data_change: false, + }, + None, + ); + let (result, _) = txn + .build_manifest(Some(&manifest), vec![], "txn", &default_build_config()) + .unwrap(); + + let frag = &result.fragments[0]; + assert!(frag.files.iter().any(|f| f.path == "f5-new.lance")); + assert_eq!(frag.overlays, overlays); + } + + #[test] + fn test_data_replacement_applies_several_groups_to_one_fragment() { + // Two groups repack four single-field files of one fragment into two + // files. Each group applies to the fragment as the previous one left + // it, and the files left holding no schema field are dropped. + let schema = ArrowSchema::new(vec![ + ArrowField::new("x", DataType::Int32, true), + ArrowField::new("a", DataType::Int32, true), + ArrowField::new("v", DataType::Int32, true), + ArrowField::new("y", DataType::Int32, true), + ]); + let lance_schema = LanceSchema::try_from(&schema).unwrap(); + let ids: Vec = lance_schema.fields.iter().map(|field| field.id).collect(); + let file = |path: &str, fields: Vec| { + let indices = (0..fields.len() as i32).collect(); + DataFile::new(path, fields, indices, ConcreteFileVersion::V2_0, None, None) + }; + let mut fragment = Fragment::new(0); + fragment.files = ids + .iter() + .map(|id| file(&format!("f{id}.lance"), vec![*id])) + .collect(); + let manifest = Manifest::new( + lance_schema, + Arc::new(vec![fragment, Fragment::new(1)]), + crate::format::DataStorageFormat::new(ConcreteFileVersion::V2_0), + HashMap::new(), + ); + + let txn = Transaction::new( + manifest.version, + Operation::DataReplacement { + replacements: vec![ + DataReplacementGroup(0, file("g0.lance", vec![ids[0], ids[1]])), + DataReplacementGroup(0, file("g1.lance", vec![ids[2], ids[3]])), + ], + data_change: false, + }, + None, + ); + let (result, _) = txn + .build_manifest(Some(&manifest), vec![], "txn", &default_build_config()) + .unwrap(); + + assert_eq!(result.fragments.len(), 2, "the untouched fragment stays"); + let layout: Vec<(&str, Vec)> = result.fragments[0] + .files + .iter() + .map(|file| (file.path.as_str(), file.fields.to_vec())) + .collect(); + assert_eq!( + layout, + vec![ + ("g0.lance", vec![ids[0], ids[1]]), + ("g1.lance", vec![ids[2], ids[3]]), + ] + ); + } + /// Replace `fields` in `fragment` at `read_version`, against a manifest /// at `manifest_version` whose schema declares field ids 3 ("x"), 4 ("a"), /// 5 ("v") and 6 ("y"). @@ -4407,6 +4522,7 @@ mod tests { None, ), )], + data_change: true, }, None, ); diff --git a/rust/lance-table/src/transaction/operation.rs b/rust/lance-table/src/transaction/operation.rs index 09f73cf6155..b3b75d79dee 100644 --- a/rust/lance-table/src/transaction/operation.rs +++ b/rust/lance-table/src/transaction/operation.rs @@ -110,24 +110,25 @@ pub enum Operation { /// `groups`, and the manifest is the durable record of the history. frag_reuse_index: Option, }, - /// Replace data in a column in the dataset with new data. This is used for - /// null column population where we replace an entirely null column with a - /// new column that has data. + /// Replace the data files backing some fields of existing fragments with + /// new files, without moving rows. Each group names a fragment and one new + /// data file for it. At commit, against the fragment as it stands then, a + /// file holding exactly the new file's fields (in the same file version) is + /// swapped for it; otherwise the new file's fields are tombstoned where + /// they live and the new file is appended, and a file left holding no + /// field of the schema is dropped. A fragment can take several groups, one + /// per new file, applied in order. Used for null column population, and by + /// horizontal compaction to repack columns into fewer files. /// - /// This operation will only allow replacing files that contain the same schema - /// e.g. if the original files contain columns A, B, C and the new files contain - /// only columns A, B then the operation is not allowed. As we would need to split - /// the original files into two files, one with column A, B and the other with column C. - /// - /// Corollary to the above: the operation will also not allow replacing files unless the - /// affected columns all have the same datafile layout across the fragments being replaced. - /// - /// e.g. if fragments being replaced contain files with different schema layouts on - /// the column being replaced, the operation is not allowed. - /// say `frag_1: [A] [B, C]` and `frag_2: [A, B] [C]` and we are trying to replace column A - /// with a new column A, the operation is not allowed. + /// The fields of a legacy (V1) data file cannot be tombstoned one by one, + /// so a field held by a V1 file can only be replaced by an exact match. DataReplacement { replacements: Vec, + /// Whether the new files change any value. `false` means the values + /// were only moved to new files (horizontal compaction): indices keep + /// their coverage of the replaced fields, overlays keep shadowing, and + /// no row is stamped as updated. + data_change: bool, }, /// Attach overlay files to fragments, supplying new values for a subset of /// `(physical offset, field)` cells without rewriting the fragments' base diff --git a/rust/lance-table/src/transaction/proto.rs b/rust/lance-table/src/transaction/proto.rs index 0e9e64853fe..6b3138f93aa 100644 --- a/rust/lance-table/src/transaction/proto.rs +++ b/rust/lance-table/src/transaction/proto.rs @@ -373,12 +373,16 @@ impl TryFrom for Transaction { } } Some(pb::transaction::Operation::DataReplacement( - pb::transaction::DataReplacement { replacements }, + pb::transaction::DataReplacement { + replacements, + data_change, + }, )) => Operation::DataReplacement { replacements: replacements .into_iter() .map(DataReplacementGroup::try_from) .collect::>>()?, + data_change: data_change.unwrap_or(true), }, Some(pb::transaction::Operation::UpdateMemWalState( pb::transaction::UpdateMemWalState { compacted_sstables }, @@ -673,14 +677,18 @@ impl From<&Transaction> for pb::Transaction { schema_metadata: Default::default(), field_metadata: Default::default(), }), - Operation::DataReplacement { replacements } => { - pb::transaction::Operation::DataReplacement(pb::transaction::DataReplacement { - replacements: replacements - .iter() - .map(pb::transaction::DataReplacementGroup::from) - .collect(), - }) - } + Operation::DataReplacement { + replacements, + data_change, + } => pb::transaction::Operation::DataReplacement(pb::transaction::DataReplacement { + replacements: replacements + .iter() + .map(pb::transaction::DataReplacementGroup::from) + .collect(), + // Written only when false, so a transaction file from a data + // changing replacement stays byte-identical to an older writer's. + data_change: (!data_change).then_some(false), + }), Operation::DataOverlay { groups } => { pb::transaction::Operation::DataOverlay(pb::transaction::DataOverlay { groups: groups @@ -882,4 +890,44 @@ mod tests { other => panic!("expected DataOverlay, got {other:?}"), } } + + /// `data_change` survives the round trip, and a transaction written + /// before the field existed reads back as a data change. + #[rstest::rstest] + #[case::moved(false)] + #[case::changed(true)] + fn test_data_replacement_data_change_roundtrips(#[case] data_change: bool) { + let transaction = Transaction::new( + 1, + Operation::DataReplacement { + replacements: vec![DataReplacementGroup( + 0, + DataFile::new_legacy_from_fields("new.lance", vec![3], None), + )], + data_change, + }, + None, + ); + let mut message = pb::Transaction::from(&transaction); + let decoded = Transaction::try_from(message.clone()).unwrap(); + assert!(matches!( + decoded.operation, + Operation::DataReplacement { data_change: d, .. } if d == data_change + )); + + if let Some(pb::transaction::Operation::DataReplacement(replacement)) = + message.operation.as_mut() + { + assert_eq!(replacement.data_change, (!data_change).then_some(false)); + replacement.data_change = None; + } + let legacy = Transaction::try_from(message).unwrap(); + assert!(matches!( + legacy.operation, + Operation::DataReplacement { + data_change: true, + .. + } + )); + } } diff --git a/rust/lance/src/dataset/blob.rs b/rust/lance/src/dataset/blob.rs index 823adaa3fd8..da34547687d 100644 --- a/rust/lance/src/dataset/blob.rs +++ b/rust/lance/src/dataset/blob.rs @@ -6819,6 +6819,7 @@ mod tests { uuid: Uuid::new_v4().hyphenated().to_string(), operation: Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, data_file)], + data_change: true, }, tag: None, transaction_properties: None, @@ -7002,6 +7003,7 @@ mod tests { uuid: Uuid::new_v4().hyphenated().to_string(), operation: Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, data_file)], + data_change: true, }, tag: None, transaction_properties: None, diff --git a/rust/lance/src/dataset/rowids/spill.rs b/rust/lance/src/dataset/rowids/spill.rs index e9e5fff33a2..72f04f28720 100644 --- a/rust/lance/src/dataset/rowids/spill.rs +++ b/rust/lance/src/dataset/rowids/spill.rs @@ -2149,6 +2149,7 @@ mod tests { let read_version = dataset.manifest.version; let operation = Operation::DataReplacement { replacements: vec![replacement], + data_change: true, }; dataset = Dataset::commit( Arc::new(dataset), diff --git a/rust/lance/src/dataset/schema_evolution.rs b/rust/lance/src/dataset/schema_evolution.rs index b57f12c80b4..d177326394f 100644 --- a/rust/lance/src/dataset/schema_evolution.rs +++ b/rust/lance/src/dataset/schema_evolution.rs @@ -3600,6 +3600,7 @@ mod test { dataset.version().version, Operation::DataReplacement { replacements: vec![replacement], + data_change: true, }, None, ); diff --git a/rust/lance/src/dataset/tests/data_file_part.rs b/rust/lance/src/dataset/tests/data_file_part.rs index 35ae388c7c4..fda33abde5d 100644 --- a/rust/lance/src/dataset/tests/data_file_part.rs +++ b/rust/lance/src/dataset/tests/data_file_part.rs @@ -93,6 +93,7 @@ async fn commit(dataset: &Dataset, replacement: DataReplacementGroup) -> Result< WriteDestination::Dataset(Arc::new(dataset.clone())), Operation::DataReplacement { replacements: vec![replacement], + data_change: true, }, Some(dataset.version_id()), None, diff --git a/rust/lance/src/dataset/tests/dataset_index.rs b/rust/lance/src/dataset/tests/dataset_index.rs index cc44fced886..6dbcf5f2dbc 100644 --- a/rust/lance/src/dataset/tests/dataset_index.rs +++ b/rust/lance/src/dataset/tests/dataset_index.rs @@ -3043,6 +3043,7 @@ async fn test_partial_compound_hybrid_prunes_same_path_different_base_rewrite() WriteDestination::Dataset(Arc::new(dataset)), Operation::DataReplacement { replacements: vec![DataReplacementGroup(1, replacement_file)], + data_change: true, }, Some(read_version), None, diff --git a/rust/lance/src/dataset/tests/dataset_merge_update.rs b/rust/lance/src/dataset/tests/dataset_merge_update.rs index 2a5a148fe79..fbf082a4276 100644 --- a/rust/lance/src/dataset/tests/dataset_merge_update.rs +++ b/rust/lance/src/dataset/tests/dataset_merge_update.rs @@ -690,6 +690,7 @@ async fn test_datafile_replacement() { WriteDestination::Dataset(dataset.clone()), Operation::DataReplacement { replacements: vec![], + data_change: true, }, Some(1), None, @@ -719,6 +720,7 @@ async fn test_datafile_replacement() { WriteDestination::Dataset(Arc::new(dataset)), Operation::DataReplacement { replacements: vec![], + data_change: true, }, Some(3), None, @@ -773,6 +775,7 @@ async fn test_datafile_replacement() { WriteDestination::Dataset(Arc::new(dataset)), Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, new_data_file)], + data_change: true, }, Some(4), None, @@ -890,6 +893,7 @@ async fn test_datafile_partial_replacement() { WriteDestination::Dataset(Arc::new(dataset)), Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, new_data_file)], + data_change: true, }, Some(3), None, @@ -951,6 +955,7 @@ async fn test_datafile_partial_replacement() { WriteDestination::Dataset(Arc::new(dataset)), Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, new_data_file)], + data_change: true, }, Some(4), None, @@ -1056,6 +1061,7 @@ async fn test_datafile_replacement_error() { WriteDestination::Dataset(Arc::new(dataset.clone())), Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, new_data_file)], + data_change: true, }, // read at the current version (after the Merge above) Some(dataset.manifest.version), @@ -2621,6 +2627,7 @@ async fn test_data_replacement_advances_row_lineage() { WriteDestination::Dataset(Arc::new(dataset)), Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, new_data_file)], + data_change: true, }, Some(read_version), None, @@ -2730,6 +2737,7 @@ async fn test_data_replacement_invalidates_index_bitmap() { WriteDestination::Dataset(Arc::new(dataset)), Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, new_data_file)], + data_change: true, }, Some(read_version), None, @@ -3014,6 +3022,7 @@ async fn test_data_replacement_populates_invalidated_bitmap() { WriteDestination::Dataset(Arc::new(dataset)), Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, new_data_file)], + data_change: true, }, Some(read_version), None, @@ -3133,6 +3142,7 @@ async fn test_fts_stale_entries_after_data_replacement() { WriteDestination::Dataset(Arc::new(dataset)), Operation::DataReplacement { replacements: vec![DataReplacementGroup(1, new_data_file)], + data_change: true, }, Some(read_version), None, @@ -3273,6 +3283,7 @@ async fn test_cross_column_fast_search_blocks_column_local_stale_postings() { WriteDestination::Dataset(Arc::new(dataset)), Operation::DataReplacement { replacements: vec![DataReplacementGroup(1, replacement_file)], + data_change: true, }, Some(read_version), None, @@ -3444,6 +3455,7 @@ async fn test_vector_index_after_data_replacement() { WriteDestination::Dataset(Arc::new(dataset)), Operation::DataReplacement { replacements: vec![DataReplacementGroup(frag1_id, new_data_file)], + data_change: true, }, Some(read_version), None, diff --git a/rust/lance/src/dataset/tests/fragment_write_columns.rs b/rust/lance/src/dataset/tests/fragment_write_columns.rs index ad0e184d721..4d054b53d02 100644 --- a/rust/lance/src/dataset/tests/fragment_write_columns.rs +++ b/rust/lance/src/dataset/tests/fragment_write_columns.rs @@ -97,7 +97,10 @@ async fn commit(dataset: &Dataset, replacements: Vec) -> R let read_version = dataset.manifest.version; Dataset::commit( WriteDestination::Dataset(Arc::new(dataset.clone())), - Operation::DataReplacement { replacements }, + Operation::DataReplacement { + replacements, + data_change: true, + }, Some(read_version), None, None, diff --git a/rust/lance/src/io/commit.rs b/rust/lance/src/io/commit.rs index cbec4753876..ce0be3efc5d 100644 --- a/rust/lance/src/io/commit.rs +++ b/rust/lance/src/io/commit.rs @@ -1487,7 +1487,8 @@ async fn prepare_attempt( Operation::Update { fields_modified, .. } => !fields_modified.is_empty(), - Operation::Merge { .. } | Operation::DataReplacement { .. } => true, + Operation::Merge { .. } => true, + Operation::DataReplacement { data_change, .. } => *data_change, _ => false, }; let ledger = if may_rewrite_in_place { diff --git a/rust/lance/src/io/commit/conflict_resolver.rs b/rust/lance/src/io/commit/conflict_resolver.rs index 6faf60ccf69..b5808c27357 100644 --- a/rust/lance/src/io/commit/conflict_resolver.rs +++ b/rust/lance/src/io/commit/conflict_resolver.rs @@ -107,14 +107,18 @@ fn may_alter_nullability(operation: &Operation) -> bool { /// Whether `operation` can commit rows that falsify such a change: by writing /// a null into a scanned field, or by omitting a required column entirely (a /// stale append's fragments read as null for columns they predate). `Delete` -/// only removes rows and `Rewrite` preserves the values it moves; `Project` -/// already conflicts with projections and merges elsewhere. +/// only removes rows, and `Rewrite` and a `DataReplacement` without +/// `data_change` preserve the values they move; `Project` already conflicts +/// with projections and merges elsewhere. fn supplies_values(operation: &Operation) -> bool { matches!( operation, Operation::Append { .. } | Operation::Update { .. } - | Operation::DataReplacement { .. } + | Operation::DataReplacement { + data_change: true, + .. + } | Operation::DataOverlay { .. } ) } @@ -256,7 +260,7 @@ impl<'a> TransactionRebase<'a> { reuse, }) } - Operation::DataReplacement { replacements } => { + Operation::DataReplacement { replacements, .. } => { let modified_fragment_ids = replacements.iter().map(|r| r.0).collect::>(); let initial_fragments = @@ -348,6 +352,19 @@ impl<'a> TransactionRebase<'a> { ) } + /// Whether this transaction is a replacement whose values only moved + /// (compaction). Losing its target to a concurrent drop or delete costs + /// nothing a fresh plan cannot redo, so it retries rather than fails. + fn only_moves_values(&self) -> bool { + matches!( + self.transaction.operation, + Operation::DataReplacement { + data_change: false, + .. + } + ) + } + #[track_caller] fn data_replacement_target_removed_err( &self, @@ -355,6 +372,9 @@ impl<'a> TransactionRebase<'a> { other_transaction: &Transaction, other_version: u64, ) -> Error { + if self.only_moves_values() { + return self.retryable_conflict_err(other_transaction, other_version); + } Error::incompatible_transaction_source( format!( "DataReplacement target fragment {} was removed by concurrent {} at version {}.", @@ -372,6 +392,9 @@ impl<'a> TransactionRebase<'a> { other_transaction: &Transaction, other_version: u64, ) -> Error { + if self.only_moves_values() { + return self.retryable_conflict_err(other_transaction, other_version); + } Error::incompatible_transaction_source( format!( "DataReplacement target field {} in fragment {} was dropped by concurrent {} at version {}.", @@ -1205,7 +1228,12 @@ impl<'a> TransactionRebase<'a> { } } Operation::UpdateConfig { .. } => Ok(()), - Operation::DataReplacement { replacements } => { + // Values that only moved to new files leave what the index + // was built from unchanged. + Operation::DataReplacement { + data_change: false, .. + } => Ok(()), + Operation::DataReplacement { replacements, .. } => { // A data replacement only conflicts if it is updating a field the // index depends on -- whether keyed on or merely carried, since // `fields` lists both (see `IndexMetadata::covering_fields`). @@ -1416,7 +1444,7 @@ impl<'a> TransactionRebase<'a> { Ok(()) } } - Operation::DataReplacement { replacements } => { + Operation::DataReplacement { replacements, .. } => { // These conflict if the rewrite touches any of the fragments being replaced. for replacement in replacements { for group in groups { @@ -1644,7 +1672,7 @@ impl<'a> TransactionRebase<'a> { other_transaction: &Transaction, other_version: u64, ) -> Result<()> { - if let Operation::DataReplacement { replacements } = &self.transaction.operation { + if let Operation::DataReplacement { replacements, .. } = &self.transaction.operation { match &other_transaction.operation { Operation::Append { .. } | Operation::Clone { .. } @@ -1743,6 +1771,12 @@ impl<'a> TransactionRebase<'a> { // the index's fragment bitmap, which would lead to fewer conflicts. However // this would introduce fragment bitmaps with holes which may not be well tested // yet. For now, we don't allow this case. + // + // Values that only moved still match what the index was + // built from, so those never conflict. + if self.only_moves_values() { + return Ok(()); + } let newly_depended_fields = new_indices .iter() .flat_map(|idx| idx.fields.iter()) @@ -1775,6 +1809,7 @@ impl<'a> TransactionRebase<'a> { } Operation::DataReplacement { replacements: other_replacements, + .. } => { // These conflict if there is overlap in fragment id && fields. for replacement in replacements { @@ -4323,6 +4358,7 @@ mod tests { 1, DataFile::new_legacy_from_fields("r.lance", vec![0], None), )], + data_change: true, }, Compatible, ), @@ -4721,6 +4757,7 @@ mod tests { "replacement", Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, file())], + data_change: true, }, ), ( @@ -5845,13 +5882,40 @@ mod tests { .map(|f| f.id) .chain(removed_fragment_ids.iter().copied()), ), - Operation::DataReplacement { replacements } => { + Operation::DataReplacement { replacements, .. } => { Box::new(replacements.iter().map(|r| r.0)) } Operation::DataOverlay { groups } => Box::new(groups.iter().map(|g| g.fragment_id)), } } + /// A schema with fields 0, 2, and 3: field 1 was dropped. + fn project_schema_without_field_1() -> lance_core::datatypes::Schema { + let arrow_schema = Schema::new( + ["a", "b", "c", "d"] + .map(|name| Field::new(name, DataType::Int32, true)) + .to_vec(), + ); + let schema = lance_core::datatypes::Schema::try_from(&arrow_schema).unwrap(); + schema.project_by_ids(&[0, 2, 3], true) + } + + fn index_on_fields(fields: Vec) -> IndexMetadata { + IndexMetadata { + uuid: Uuid::new_v4(), + name: "idx".to_string(), + fields, + covering_fields: vec![], + dataset_version: 1, + fragment_bitmap: Some(RoaringBitmap::from_iter([0u32])), + index_details: None, + index_version: 0, + created_at: None, + base_id: None, + files: None, + } + } + #[tokio::test] async fn test_conflicts_data_replacement() { use io::commit::conflict_resolver::tests::{ConflictResult::*, modified_fragment_ids}; @@ -5871,9 +5935,11 @@ mod tests { "Different fragments", Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01.clone())], + data_change: true, }, Operation::DataReplacement { replacements: vec![DataReplacementGroup(1, data_file_frag1_fields01)], + data_change: true, }, Compatible, ), @@ -5881,9 +5947,11 @@ mod tests { "Same fragment, different fields", Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01.clone())], + data_change: true, }, Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, data_file_frag0_fields23)], + data_change: true, }, Compatible, ), @@ -5891,9 +5959,11 @@ mod tests { "Same fragment, same fields", Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01.clone())], + data_change: true, }, Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01.clone())], + data_change: true, }, Retryable, ), @@ -5901,12 +5971,14 @@ mod tests { "Same fragment, overlapping fields", Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01.clone())], + data_change: true, }, Operation::DataReplacement { replacements: vec![DataReplacementGroup( 0, DataFile::new_legacy_from_fields("path0_12", vec![1, 2], None), )], + data_change: true, }, Retryable, ), @@ -5914,6 +5986,7 @@ mod tests { "DataReplacement vs Rewrite on same fragment", Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01.clone())], + data_change: true, }, Operation::Rewrite { groups: vec![RewriteGroup { @@ -5929,6 +6002,7 @@ mod tests { "DataReplacement vs Rewrite on different fragment", Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01.clone())], + data_change: true, }, Operation::Rewrite { groups: vec![RewriteGroup { @@ -5947,6 +6021,7 @@ mod tests { "DataReplacement vs Update (RewriteColumns) on a different field", Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01.clone())], + data_change: true, }, Operation::Update { updated_fragments: vec![Fragment::new(0)], @@ -5966,6 +6041,7 @@ mod tests { "DataReplacement vs Update (RewriteColumns) with inserts on a different field", Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01.clone())], + data_change: true, }, Operation::Update { updated_fragments: vec![Fragment::new(0)], @@ -5984,6 +6060,7 @@ mod tests { "DataReplacement vs Update (RewriteColumns) that rewrote one of our fields", Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01.clone())], + data_change: true, }, Operation::Update { updated_fragments: vec![Fragment::new(0)], @@ -6002,6 +6079,7 @@ mod tests { "DataReplacement vs Update (RewriteRows) that moved our rows", Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01.clone())], + data_change: true, }, Operation::Update { updated_fragments: vec![Fragment::new(0)], @@ -6020,6 +6098,7 @@ mod tests { "DataReplacement vs Update that removed our fragment", Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01.clone())], + data_change: true, }, Operation::Update { updated_fragments: vec![], @@ -6038,6 +6117,7 @@ mod tests { "DataReplacement vs Update (RewriteRows) that moved a different fragment's rows", Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01.clone())], + data_change: true, }, Operation::Update { updated_fragments: vec![Fragment::new(1)], @@ -6056,6 +6136,7 @@ mod tests { "DataReplacement vs Delete (deletion-vector only) on same fragment", Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01.clone())], + data_change: true, }, Operation::Delete { deleted_fragment_ids: vec![], @@ -6068,6 +6149,7 @@ mod tests { "DataReplacement vs Delete that removes the fragment", Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01.clone())], + data_change: true, }, Operation::Delete { deleted_fragment_ids: vec![0], @@ -6080,7 +6162,8 @@ mod tests { ( "DataReplacement vs Merge", Operation::DataReplacement { - replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01)], + replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01.clone())], + data_change: true, }, Operation::Merge { fragments: vec![Fragment::new(0)], @@ -6121,6 +6204,86 @@ mod tests { 0, DataFile::new_legacy_from_fields("path0_3", vec![3], None), )], + data_change: true, + }, + Retryable, + ), + // A replacement whose values only moved (compaction) retries where + // a data-changing one fails, and never conflicts with an index. + ( + "Moved DataReplacement vs Delete that removes the fragment", + Operation::DataReplacement { + replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01.clone())], + data_change: false, + }, + Operation::Delete { + deleted_fragment_ids: vec![0], + updated_fragments: vec![], + predicate: "a > 0".to_string(), + }, + Retryable, + ), + ( + "Moved DataReplacement vs Project that dropped one of its fields", + Operation::DataReplacement { + replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01.clone())], + data_change: false, + }, + Operation::Project { + schema: project_schema_without_field_1(), + preserves_nullability: true, + }, + Retryable, + ), + ( + "DataReplacement vs Project that dropped one of its fields", + Operation::DataReplacement { + replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01.clone())], + data_change: true, + }, + Operation::Project { + schema: project_schema_without_field_1(), + preserves_nullability: true, + }, + NotCompatible, + ), + ( + "Moved DataReplacement vs CreateIndex on one of its fields", + Operation::DataReplacement { + replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01.clone())], + data_change: false, + }, + Operation::CreateIndex { + new_indices: vec![index_on_fields(vec![1])], + removed_indices: vec![], + }, + Compatible, + ), + ( + // op1 is CreateIndex for the same reason as the case above. + "CreateIndex on a field vs moved DataReplacement of that field", + Operation::CreateIndex { + new_indices: vec![index_on_fields(vec![1])], + removed_indices: vec![], + }, + Operation::DataReplacement { + replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01.clone())], + data_change: false, + }, + Compatible, + ), + ( + "Moved DataReplacement vs DataReplacement of an overlapping field", + Operation::DataReplacement { + replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01.clone())], + data_change: false, + }, + Operation::DataReplacement { + replacements: vec![DataReplacementGroup( + 0, + DataFile::new_legacy_from_fields("path0_1", vec![1], None), + )], + data_change: true, }, Retryable, ), From 6c157265638cdc682496cbc21dcbd84e42d17854 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Fri, 2 Oct 2026 23:44:26 +0800 Subject: [PATCH 10/24] feat(compaction): repack columns as a compact_files task committed via DataReplacement --- java/lance-jni/src/blocking_dataset.rs | 33 +- java/lance-jni/src/optimize.rs | 139 ++- java/lance-jni/src/utils.rs | 30 +- .../java/org/lance/compaction/Compaction.java | 10 +- .../lance/compaction/CompactionOptions.java | 106 ++- .../org/lance/compaction/CompactionScope.java | 34 + .../org/lance/compaction/CompactionTask.java | 10 +- .../org/lance/compaction/RewriteResult.java | 42 + .../java/org/lance/compaction/TaskData.java | 32 +- .../org/lance/operation/DataReplacement.java | 6 +- .../test/java/org/lance/CompactionTest.java | 51 ++ python/python/lance/dataset.py | 29 + python/python/lance/lance/optimize.pyi | 3 +- python/python/lance/optimize.py | 22 + python/python/tests/test_optimize.py | 50 ++ python/src/dataset/optimize.rs | 41 +- rust/lance-table/src/transaction/operation.rs | 4 +- rust/lance/src/dataset.rs | 1 - rust/lance/src/dataset/compaction_stats.rs | 25 +- rust/lance/src/dataset/optimize.rs | 515 +++++++---- rust/lance/src/dataset/optimize/repack.rs | 374 ++++++++ .../src/dataset/optimize/tests/repack.rs | 613 +++++++++++++ rust/lance/src/dataset/rewrite_columns.rs | 814 ------------------ .../tests/dataset_overlay_index_masking.rs | 19 +- rust/lance/src/index.rs | 7 +- rust/lance/src/index/frag_reuse.rs | 1 + rust/lance/src/io/commit/conflict_resolver.rs | 34 +- 27 files changed, 1982 insertions(+), 1063 deletions(-) create mode 100644 java/src/main/java/org/lance/compaction/CompactionScope.java create mode 100644 rust/lance/src/dataset/optimize/repack.rs create mode 100644 rust/lance/src/dataset/optimize/tests/repack.rs delete mode 100644 rust/lance/src/dataset/rewrite_columns.rs diff --git a/java/lance-jni/src/blocking_dataset.rs b/java/lance-jni/src/blocking_dataset.rs index 320475def0a..91e4af7b1f6 100644 --- a/java/lance-jni/src/blocking_dataset.rs +++ b/java/lance-jni/src/blocking_dataset.rs @@ -10,9 +10,9 @@ use crate::namespace::{ use crate::session::{handle_from_session, session_from_handle}; use crate::traits::{FromJObjectWithEnv, FromJString, export_vec, import_vec, import_vec_to_rust}; use crate::utils::{ - build_compaction_options, extract_base_store_params, extract_storage_options, - extract_write_params, get_scalar_index_params, get_vector_index_params, to_java_map, - to_rust_map, + apply_repack_options, build_compaction_options, extract_base_store_params, + extract_storage_options, extract_write_params, get_scalar_index_params, + get_vector_index_params, to_java_map, to_rust_map, }; use crate::{block_on, traits::IntoJava}; use arrow::array::RecordBatchReader; @@ -3616,7 +3616,22 @@ fn convert_java_compaction_options_to_rust( )? .l()?; - build_compaction_options( + let max_data_files_per_fragment = env + .call_method( + &java_options, + "getMaxDataFilesPerFragment", + "()Ljava/util/Optional;", + &[], + )? + .l()?; + let column_groups = env + .call_method(&java_options, "getColumnGroups", "()Ljava/util/List;", &[])? + .l()?; + let scope = env + .call_method(&java_options, "getScope", "()Ljava/util/Optional;", &[])? + .l()?; + + let mut options = build_compaction_options( env, &target_rows_per_fragment, &max_rows_per_group, @@ -3634,7 +3649,15 @@ fn convert_java_compaction_options_to_rust( &excluded_fragment_ids, &data_storage_version, config, - ) + )?; + apply_repack_options( + env, + &mut options, + &max_data_files_per_fragment, + &column_groups, + &scope, + )?; + Ok(options) } #[unsafe(no_mangle)] diff --git a/java/lance-jni/src/optimize.rs b/java/lance-jni/src/optimize.rs index ce0151d202a..edc630cc3a3 100644 --- a/java/lance-jni/src/optimize.rs +++ b/java/lance-jni/src/optimize.rs @@ -11,8 +11,9 @@ use jni::{ use lance::dataset::{ index::DatasetIndexRemapperOptions, optimize::{ - CompactionMetrics, CompactionMode, CompactionOptions, CompactionPlan, CompactionTask, - IndexRemapperOptions, RewriteResult, TaskData, commit_compaction, plan_compaction, + CompactionMetrics, CompactionMode, CompactionOptions, CompactionPlan, CompactionScope, + CompactionTask, CompactionTaskKind, IndexRemapperOptions, RepackedFiles, RewriteResult, + TaskData, commit_compaction, plan_compaction, }, }; @@ -23,12 +24,13 @@ use crate::{ FromJObjectWithEnv, IntoJava, export_vec, import_vec_from_method, import_vec_to_rust, }, utils::{ - build_compaction_options, to_java_boolean_obj, to_java_float_obj, to_java_list, - to_java_long_obj, to_java_optional, + apply_repack_options, build_compaction_options, to_java_boolean_obj, to_java_float_obj, + to_java_list, to_java_long_obj, to_java_optional, }, }; use crate::error::Result; +use crate::ffi::JNIEnvExt; #[unsafe(no_mangle)] pub extern "system" fn Java_org_lance_compaction_Compaction_nativePlanCompaction<'local>( @@ -50,6 +52,9 @@ pub extern "system" fn Java_org_lance_compaction_Compaction_nativePlanCompaction max_source_bytes: JObject, // Optional excluded_fragment_ids: JObject, // List data_storage_version: JObject, // Optional + max_data_files_per_fragment: JObject, // Optional + column_groups: JObject, // List> + scope: JObject, // Optional ) -> JObject<'local> { ok_or_throw_with_return!( env, @@ -70,7 +75,10 @@ pub extern "system" fn Java_org_lance_compaction_Compaction_nativePlanCompaction max_source_rows, max_source_bytes, excluded_fragment_ids, - data_storage_version + data_storage_version, + max_data_files_per_fragment, + column_groups, + scope ), JObject::null() ) @@ -95,13 +103,16 @@ fn inner_plan_compaction<'local>( max_source_bytes: JObject, // Optional excluded_fragment_ids: JObject, // List data_storage_version: JObject, // Optional + max_data_files_per_fragment: JObject, // Optional + column_groups: JObject, // List> + scope: JObject, // Optional ) -> Result> { let config = { let dataset = unsafe { env.get_rust_field::<_, _, BlockingDataset>(&java_dataset, NATIVE_DATASET) }?; dataset.inner.manifest.config.clone() }; - let compaction_options = build_compaction_options( + let mut compaction_options = build_compaction_options( env, &target_rows_per_fragment, &max_rows_per_group, @@ -120,6 +131,13 @@ fn inner_plan_compaction<'local>( &data_storage_version, &config, )?; + apply_repack_options( + env, + &mut compaction_options, + &max_data_files_per_fragment, + &column_groups, + &scope, + )?; let plan = { let dataset = @@ -258,6 +276,9 @@ pub extern "system" fn Java_org_lance_compaction_CompactionTask_nativeExecute<'l max_source_bytes: JObject, // Optional excluded_fragment_ids: JObject, // List data_storage_version: JObject, // Optional + max_data_files_per_fragment: JObject, // Optional + column_groups: JObject, // List> + scope: JObject, // Optional ) -> JObject<'local> { ok_or_throw_with_return!( env, @@ -280,7 +301,10 @@ pub extern "system" fn Java_org_lance_compaction_CompactionTask_nativeExecute<'l max_source_rows, max_source_bytes, excluded_fragment_ids, - data_storage_version + data_storage_version, + max_data_files_per_fragment, + column_groups, + scope ), JObject::null() ) @@ -307,6 +331,9 @@ fn inner_execute_task<'local>( max_source_bytes: JObject, // Optional excluded_fragment_ids: JObject, // List data_storage_version: JObject, // Optional + max_data_files_per_fragment: JObject, // Optional + column_groups: JObject, // List> + scope: JObject, // Optional ) -> Result> { let task_data: TaskData = task_data.extract_object(env)?; let config = { @@ -314,7 +341,7 @@ fn inner_execute_task<'local>( unsafe { env.get_rust_field::<_, _, BlockingDataset>(&java_dataset, NATIVE_DATASET) }?; dataset.inner.manifest.config.clone() }; - let compaction_options = build_compaction_options( + let mut compaction_options = build_compaction_options( env, &target_rows_per_fragment, &max_rows_per_group, @@ -333,6 +360,13 @@ fn inner_execute_task<'local>( &data_storage_version, &config, )?; + apply_repack_options( + env, + &mut compaction_options, + &max_data_files_per_fragment, + &column_groups, + &scope, + )?; let compaction_task = CompactionTask { task: task_data, read_version: read_version as u64, @@ -347,26 +381,48 @@ fn inner_execute_task<'local>( } const TASK_DATA_CLASS: &str = "org/lance/compaction/TaskData"; -const TASK_DATA_CONSTRUCTOR_SIG: &str = "(Ljava/util/List;)V"; +const TASK_DATA_CONSTRUCTOR_SIG: &str = "(Ljava/util/List;Ljava/util/List;)V"; const COMPACTION_METRICS_CLASS: &str = "org/lance/compaction/CompactionMetrics"; const COMPACTION_METRICS_CONSTRUCTOR_SIG: &str = "(JJJJ)V"; const COMPACTION_PLAN_CLASS: &str = "org/lance/compaction/CompactionPlan"; const COMPACTION_PLAN_CONSTRUCTOR_SIG: &str = "(Ljava/util/List;JLorg/lance/compaction/CompactionOptions;)V"; const REWRITE_RESULT_CLASS: &str = "org/lance/compaction/RewriteResult"; -const REWRITE_RESULT_CONSTRUCTOR_SIG: &str = - "(Lorg/lance/compaction/CompactionMetrics;Ljava/util/List;Ljava/util/List;J[B)V"; +const REWRITE_RESULT_CONSTRUCTOR_SIG: &str = "(Lorg/lance/compaction/CompactionMetrics;Ljava/util/List;Ljava/util/List;J[BLjava/lang/Long;Ljava/util/List;)V"; const COMPACTION_OPTIONS_CLASS: &str = "org/lance/compaction/CompactionOptions"; const COMPACTION_MODE_CLASS: &str = "org/lance/compaction/CompactionMode"; -const COMPACTION_OPTIONS_CONSTRUCTOR_SIG: &str = "(Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/List;Ljava/util/Optional;)V"; +const COMPACTION_OPTIONS_CONSTRUCTOR_SIG: &str = "(Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/List;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/List;Ljava/util/Optional;)V"; impl IntoJava for &TaskData { fn into_java<'a>(self, env: &mut JNIEnv<'a>) -> Result> { let fragments = export_vec(env, &self.fragments)?; + let repack_files = match &self.kind { + CompactionTaskKind::RewriteFragments => JObject::null(), + CompactionTaskKind::RepackColumns { files } => { + let mut lists = Vec::with_capacity(files.len()); + for file in files { + let ids = file + .iter() + .map(|id| { + Ok(env.new_object( + "java/lang/Integer", + "(I)V", + &[JValueGen::Int(*id)], + )?) + }) + .collect::>>()?; + lists.push(to_java_list(env, &ids)?); + } + to_java_list(env, &lists)? + } + }; Ok(env.new_object( TASK_DATA_CLASS, TASK_DATA_CONSTRUCTOR_SIG, - &[JValueGen::Object(&fragments)], + &[ + JValueGen::Object(&fragments), + JValueGen::Object(&repack_files), + ], )?) } } @@ -445,6 +501,25 @@ impl IntoJava for &CompactionOptions { None => JObject::null(), }; let data_storage_version_opt = to_java_optional(env, data_storage_version)?; + let max_data_files_per_fragment = + to_java_long_obj(env, self.max_data_files_per_fragment.map(|v| v as i64))?; + let max_data_files_per_fragment_opt = to_java_optional(env, max_data_files_per_fragment)?; + let mut column_groups = Vec::with_capacity(self.column_groups.len()); + for group in &self.column_groups { + let names = group + .iter() + .map(|name| Ok(JObject::from(env.new_string(name)?))) + .collect::>>()?; + column_groups.push(to_java_list(env, &names)?); + } + let column_groups = to_java_list(env, &column_groups)?; + let scope = match self.scope { + CompactionScope::All => "all", + CompactionScope::RewriteFragments => "rewrite_fragments", + CompactionScope::RepackColumns => "repack_columns", + }; + let scope: JObject = env.new_string(scope)?.into(); + let scope_opt = to_java_optional(env, scope)?; Ok(env.new_object( COMPACTION_OPTIONS_CLASS, @@ -465,6 +540,9 @@ impl IntoJava for &CompactionOptions { JValueGen::Object(&max_source_bytes_opt), JValueGen::Object(&excluded_fragment_ids), JValueGen::Object(&data_storage_version_opt), + JValueGen::Object(&max_data_files_per_fragment_opt), + JValueGen::Object(&column_groups), + JValueGen::Object(&scope_opt), ], )?) } @@ -496,6 +574,13 @@ impl IntoJava for &RewriteResult { } else { JObject::null() }; + let (repacked_fragment_id, repacked_files) = match &self.repacked_files { + Some(repacked) => ( + to_java_long_obj(env, Some(repacked.fragment_id as i64))?, + export_vec(env, &repacked.files)?, + ), + None => (JObject::null(), JObject::null()), + }; Ok(env.new_object( REWRITE_RESULT_CLASS, REWRITE_RESULT_CONSTRUCTOR_SIG, @@ -505,6 +590,8 @@ impl IntoJava for &RewriteResult { JValueGen::Object(&original_fragments), JValueGen::Long(self.read_version as i64), JValueGen::Object(&row_addrs), + JValueGen::Object(&repacked_fragment_id), + JValueGen::Object(&repacked_files), ], )?) } @@ -534,8 +621,19 @@ impl FromJObjectWithEnv for JObject<'_> { let task_data = import_vec_from_method(env, self, "getFragments", |env, fragment| { fragment.extract_object(env) })?; + let repack_files = env + .call_method(self, "getRepackFiles", "()Ljava/util/List;", &[])? + .l()?; + let kind = if repack_files.is_null() { + CompactionTaskKind::RewriteFragments + } else { + let files = + import_vec_to_rust(env, &repack_files, |env, file| env.get_integers(&file))?; + CompactionTaskKind::RepackColumns { files } + }; Ok(TaskData { fragments: task_data, + kind, }) } } @@ -569,12 +667,27 @@ impl FromJObjectWithEnv for JObject<'_> { } else { Some(env.convert_byte_array(row_addrs_obj)?) }; + let repacked_fragment_id = env + .call_method(self, "getRepackedFragmentId", "()Ljava/lang/Long;", &[])? + .l()?; + let repacked_files = if repacked_fragment_id.is_null() { + None + } else { + let fragment_id = env + .call_method(&repacked_fragment_id, "longValue", "()J", &[])? + .j()? as u64; + let files = import_vec_from_method(env, self, "getRepackedFiles", |env, file| { + file.extract_object(env) + })?; + Some(RepackedFiles { fragment_id, files }) + }; Ok(RewriteResult { metrics, new_fragments, read_version, original_fragments, row_addrs, + repacked_files, }) } } diff --git a/java/lance-jni/src/utils.rs b/java/lance-jni/src/utils.rs index 5a3f795880c..e9d950c7aa1 100644 --- a/java/lance-jni/src/utils.rs +++ b/java/lance-jni/src/utils.rs @@ -8,7 +8,7 @@ use arrow_schema::{DataType, Field}; use jni::JNIEnv; use jni::objects::{JFloatArray, JMap, JObject, JString, JValue, JValueGen}; use jni::sys::{jboolean, jfloat, jlong}; -use lance::dataset::optimize::{CompactionMode, CompactionOptions}; +use lance::dataset::optimize::{CompactionMode, CompactionOptions, CompactionScope}; use lance::dataset::{WriteMode, WriteParams}; use lance::index::vector::{IndexFileVersion, StageParams, VectorIndexParams}; use lance::io::ObjectStoreParams; @@ -286,6 +286,34 @@ pub fn build_compaction_options( Ok(compaction_options) } +/// Apply the column repack options a Java `CompactionOptions` carries. An empty +/// `column_groups` keeps the table config's groups. +pub fn apply_repack_options( + env: &mut JNIEnv, + options: &mut CompactionOptions, + max_data_files_per_fragment: &JObject, // Optional + column_groups: &JObject, // List> + scope: &JObject, // Optional +) -> Result<()> { + if let Some(max) = env.get_long_opt(max_data_files_per_fragment)? { + options.max_data_files_per_fragment = Some(usize::try_from(max).map_err(|_| { + Error::input_error(format!( + "max_data_files_per_fragment must be positive, got {max}" + )) + })?); + } + let groups = crate::traits::import_vec_to_rust(env, column_groups, |env, group| { + env.get_strings(&group) + })?; + if !groups.is_empty() { + options.column_groups = groups; + } + if let Some(scope) = env.get_string_opt(scope)? { + options.scope = CompactionScope::try_from(scope.as_str())?; + } + Ok(()) +} + // Convert from Java Optional to Rust Option pub fn get_query(env: &mut JNIEnv, query_obj: JObject) -> Result> { let query = env.get_optional(&query_obj, |env, java_obj| { diff --git a/java/src/main/java/org/lance/compaction/Compaction.java b/java/src/main/java/org/lance/compaction/Compaction.java index e739d09dac8..7bd561f2446 100644 --- a/java/src/main/java/org/lance/compaction/Compaction.java +++ b/java/src/main/java/org/lance/compaction/Compaction.java @@ -50,7 +50,10 @@ public static CompactionPlan planCompaction( compactionOptions.getMaxSourceRows(), compactionOptions.getMaxSourceBytes(), compactionOptions.getExcludedFragmentIds(), - compactionOptions.getDataStorageVersion()); + compactionOptions.getDataStorageVersion(), + compactionOptions.getMaxDataFilesPerFragment(), + compactionOptions.getColumnGroups(), + compactionOptions.getScope()); } } @@ -155,5 +158,8 @@ private static native CompactionPlan nativePlanCompaction( Optional maxSourceRows, Optional maxSourceBytes, List excludedFragmentIds, - Optional dataStorageVersion); + Optional dataStorageVersion, + Optional maxDataFilesPerFragment, + List> columnGroups, + Optional scope); } diff --git a/java/src/main/java/org/lance/compaction/CompactionOptions.java b/java/src/main/java/org/lance/compaction/CompactionOptions.java index d83e618da5b..3dd7ed2750b 100644 --- a/java/src/main/java/org/lance/compaction/CompactionOptions.java +++ b/java/src/main/java/org/lance/compaction/CompactionOptions.java @@ -22,6 +22,7 @@ import java.io.ObjectOutputStream; import java.io.OptionalDataException; import java.io.Serializable; +import java.util.ArrayList; import java.util.Collections; import java.util.List; import java.util.Objects; @@ -54,7 +55,11 @@ public class CompactionOptions implements Serializable { private Optional maxSourceBytes; private List excludedFragmentIds; private Optional dataStorageVersion; + private Optional maxDataFilesPerFragment; + private List> columnGroups; + private Optional scope; + // Also called from the native layer. private CompactionOptions( Optional targetRowsPerFragment, Optional maxRowsPerGroup, @@ -70,7 +75,10 @@ private CompactionOptions( Optional maxSourceRows, Optional maxSourceBytes, List excludedFragmentIds, - Optional dataStorageVersion) { + Optional dataStorageVersion, + Optional maxDataFilesPerFragment, + List> columnGroups, + Optional scope) { this.targetRowsPerFragment = targetRowsPerFragment; this.maxRowsPerGroup = maxRowsPerGroup; this.maxBytesPerFile = maxBytesPerFile; @@ -86,6 +94,30 @@ private CompactionOptions( this.maxSourceBytes = maxSourceBytes; this.excludedFragmentIds = List.copyOf(excludedFragmentIds); this.dataStorageVersion = dataStorageVersion.map(DataStorageVersion::fromRustString); + this.maxDataFilesPerFragment = maxDataFilesPerFragment; + this.columnGroups = copyGroups(columnGroups); + this.scope = scope.map(CompactionOptions::scopeFromValue); + } + + private static List> copyGroups(List> groups) { + List> copy = new ArrayList<>(groups.size()); + for (List group : groups) { + copy.add(List.copyOf(group)); + } + return Collections.unmodifiableList(copy); + } + + public Optional getMaxDataFilesPerFragment() { + return maxDataFilesPerFragment; + } + + public List> getColumnGroups() { + return columnGroups; + } + + /** Returns the scope as its string value for the native layer. */ + public Optional getScope() { + return scope.map(CompactionScope::getValue); } public Optional getDeferIndexRemap() { @@ -172,6 +204,9 @@ public String toString() { .add("maxSourceBytes", maxSourceBytes.orElse(null)) .add("excludedFragmentIds", excludedFragmentIds) .add("dataStorageVersion", dataStorageVersion.orElse(null)) + .add("maxDataFilesPerFragment", maxDataFilesPerFragment.orElse(null)) + .add("columnGroups", columnGroups) + .add("scope", scope.orElse(null)) .toString(); } @@ -191,6 +226,9 @@ private void writeObject(ObjectOutputStream output) throws IOException { output.writeObject(maxSourceBytes.orElse(null)); output.writeObject(excludedFragmentIds); output.writeObject(getDataStorageVersion().orElse(null)); + output.writeObject(maxDataFilesPerFragment.orElse(null)); + output.writeObject(new ArrayList<>(columnGroups)); + output.writeObject(getScope().orElse(null)); } private void readObject(ObjectInputStream input) throws IOException, ClassNotFoundException { @@ -218,6 +256,32 @@ private void readObject(ObjectInputStream input) throws IOException, ClassNotFou this.maxSourceBytes = readTrailingLong(input); this.excludedFragmentIds = readTrailingLongList(input); this.dataStorageVersion = readTrailingString(input).map(DataStorageVersion::fromRustString); + this.maxDataFilesPerFragment = readTrailingLong(input); + this.columnGroups = readTrailingGroups(input); + this.scope = readTrailingString(input).map(CompactionOptions::scopeFromValue); + } + + private static CompactionScope scopeFromValue(String value) { + for (CompactionScope scope : CompactionScope.values()) { + if (scope.getValue().equals(value)) { + return scope; + } + } + throw new IllegalArgumentException("Unknown compaction scope: " + value); + } + + @SuppressWarnings("unchecked") + private static List> readTrailingGroups(ObjectInputStream input) + throws IOException, ClassNotFoundException { + try { + List> groups = (List>) input.readObject(); + return groups == null ? Collections.emptyList() : copyGroups(groups); + } catch (OptionalDataException e) { + if (!e.eof) { + throw e; + } + return Collections.emptyList(); + } } /** @@ -280,6 +344,9 @@ public static class Builder { private Optional maxSourceBytes = Optional.empty(); private List excludedFragmentIds = Collections.emptyList(); private Optional dataStorageVersion = Optional.empty(); + private Optional maxDataFilesPerFragment = Optional.empty(); + private List> columnGroups = Collections.emptyList(); + private Optional scope = Optional.empty(); private Builder() {} @@ -402,6 +469,38 @@ public Builder withExcludedFragmentIds(List excludedFragmentIds) { return this; } + /** + * Maximum number of data files a fragment may hold columns in before its columns are repacked + * into fewer files. Each {@code addColumns} backfill adds one file per fragment. A repack + * rewrites only the columns that move and keeps rows, fragment ids and indices as they are. + * Unset means no file-count trigger. + * + * @throws IllegalArgumentException if {@code maxDataFilesPerFragment} is not positive + */ + public Builder withMaxDataFilesPerFragment(long maxDataFilesPerFragment) { + this.maxDataFilesPerFragment = + Optional.of(positiveBudget("maxDataFilesPerFragment", maxDataFilesPerFragment)); + return this; + } + + /** + * Top-level columns to keep in their own data files. Each inner list becomes one data file per + * fragment; the columns no group names share one file. Fragments the compaction rewrites are + * written this way, and the others have their columns repacked to match. Names that are not + * top-level columns are ignored. An empty list keeps the {@code lance.compaction.column_groups} + * table config, if any. + */ + public Builder withColumnGroups(List> columnGroups) { + this.columnGroups = copyGroups(Objects.requireNonNull(columnGroups, "columnGroups")); + return this; + } + + /** Which kinds of task to plan. Defaults to {@link CompactionScope#ALL}. */ + public Builder withScope(CompactionScope scope) { + this.scope = Optional.of(Objects.requireNonNull(scope, "scope")); + return this; + } + /** * A max source budget of zero admits no work and a negative value would wrap around to an * effectively unlimited budget on the Rust side, so both are rejected here. Leave the option @@ -431,7 +530,10 @@ public CompactionOptions build() { maxSourceRows, maxSourceBytes, excludedFragmentIds, - dataStorageVersion.map(DataStorageVersion::toRustString)); + dataStorageVersion.map(DataStorageVersion::toRustString), + maxDataFilesPerFragment, + columnGroups, + scope.map(CompactionScope::getValue)); } } } diff --git a/java/src/main/java/org/lance/compaction/CompactionScope.java b/java/src/main/java/org/lance/compaction/CompactionScope.java new file mode 100644 index 00000000000..c5d509fbbb6 --- /dev/null +++ b/java/src/main/java/org/lance/compaction/CompactionScope.java @@ -0,0 +1,34 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.lance.compaction; + +/** Which kinds of task a compaction plans. */ +public enum CompactionScope { + /** Rewrite the fragments that need it and repack the columns of the others (default). */ + ALL("all"), + /** Only rewrite fragments. */ + REWRITE_FRAGMENTS("rewrite_fragments"), + /** Only repack columns. */ + REPACK_COLUMNS("repack_columns"); + + private final String value; + + CompactionScope(String value) { + this.value = value; + } + + public String getValue() { + return value; + } +} diff --git a/java/src/main/java/org/lance/compaction/CompactionTask.java b/java/src/main/java/org/lance/compaction/CompactionTask.java index ab4f9c7c7a2..e5cedbb6587 100644 --- a/java/src/main/java/org/lance/compaction/CompactionTask.java +++ b/java/src/main/java/org/lance/compaction/CompactionTask.java @@ -63,7 +63,10 @@ public RewriteResult execute(Dataset dataset) { compactionOptions.getMaxSourceRows(), compactionOptions.getMaxSourceBytes(), compactionOptions.getExcludedFragmentIds(), - compactionOptions.getDataStorageVersion()); + compactionOptions.getDataStorageVersion(), + compactionOptions.getMaxDataFilesPerFragment(), + compactionOptions.getColumnGroups(), + compactionOptions.getScope()); } } @@ -85,7 +88,10 @@ private native RewriteResult nativeExecute( Optional maxSourceRows, Optional maxSourceBytes, List excludedFragmentIds, - Optional dataStorageVersion); + Optional dataStorageVersion, + Optional maxDataFilesPerFragment, + List> columnGroups, + Optional scope); public CompactionOptions getCompactionOptions() { return compactionOptions; diff --git a/java/src/main/java/org/lance/compaction/RewriteResult.java b/java/src/main/java/org/lance/compaction/RewriteResult.java index c4c5d816c2c..f33c52a4f64 100644 --- a/java/src/main/java/org/lance/compaction/RewriteResult.java +++ b/java/src/main/java/org/lance/compaction/RewriteResult.java @@ -14,6 +14,7 @@ package org.lance.compaction; import org.lance.FragmentMetadata; +import org.lance.fragment.DataFile; import javax.annotation.Nullable; @@ -25,6 +26,10 @@ * committed later. */ public class RewriteResult implements Serializable { + // Pinned to the UID generated before the repack fields were added, so that results produced by + // older workers still deserialize during a rolling upgrade. + private static final long serialVersionUID = 4501818269828675274L; + private final CompactionMetrics metrics; private final List newFragments; private final List originalFragments; @@ -34,17 +39,54 @@ public class RewriteResult implements Serializable { // null for stable row IDs. @Nullable private final byte[] rowAddrs; + // Set only for a column repack: the fragment it repacked and the data files it wrote. + @Nullable private final Long repackedFragmentId; + @Nullable private final List repackedFiles; + public RewriteResult( CompactionMetrics metrics, List newFragments, List originalFragments, long readVersion, byte[] rowAddrs) { + this(metrics, newFragments, originalFragments, readVersion, rowAddrs, null, null); + } + + public RewriteResult( + CompactionMetrics metrics, + List newFragments, + List originalFragments, + long readVersion, + byte[] rowAddrs, + @Nullable Long repackedFragmentId, + @Nullable List repackedFiles) { this.metrics = metrics; this.newFragments = newFragments; this.originalFragments = originalFragments; this.readVersion = readVersion; this.rowAddrs = rowAddrs; + this.repackedFragmentId = repackedFragmentId; + this.repackedFiles = repackedFiles; + } + + /** + * The fragment a column repack wrote new data files for. + * + * @return null for the result of a fragment rewrite + */ + @Nullable + public Long getRepackedFragmentId() { + return repackedFragmentId; + } + + /** + * The data files a column repack wrote, empty when the fragment had no rows to write. + * + * @return null for the result of a fragment rewrite + */ + @Nullable + public List getRepackedFiles() { + return repackedFiles; } public long getReadVersion() { diff --git a/java/src/main/java/org/lance/compaction/TaskData.java b/java/src/main/java/org/lance/compaction/TaskData.java index 4ec5958215c..2d9cfa95f23 100644 --- a/java/src/main/java/org/lance/compaction/TaskData.java +++ b/java/src/main/java/org/lance/compaction/TaskData.java @@ -15,18 +15,48 @@ import org.lance.FragmentMetadata; +import javax.annotation.Nullable; + import java.io.Serializable; import java.util.List; -/** Data of compaction task. */ +/** + * Data of compaction task. + * + *

A task either rewrites its fragments into new fragments, or, when {@link #getRepackFiles()} is + * non-null, rewrites some columns of its one fragment into new data files and leaves the fragment + * id, rows, deletions, overlays and index coverage as they are. + */ public class TaskData implements Serializable { + // Pinned to the UID generated before repackFiles was added, so that tasks queued by older workers + // still deserialize (as fragment rewrites) during a rolling upgrade. + private static final long serialVersionUID = -4884632518342713596L; + private final List fragments; + // One entry per new data file: the field ids of the top-level columns it holds. + @Nullable private final List> repackFiles; + public TaskData(List fragments) { + this(fragments, null); + } + + public TaskData(List fragments, @Nullable List> repackFiles) { this.fragments = fragments; + this.repackFiles = repackFiles; } public List getFragments() { return fragments; } + + /** + * The new data files of a column repack, each the field ids of the top-level columns it holds. + * + * @return null for a task that rewrites its fragments + */ + @Nullable + public List> getRepackFiles() { + return repackFiles; + } } diff --git a/java/src/main/java/org/lance/operation/DataReplacement.java b/java/src/main/java/org/lance/operation/DataReplacement.java index 484149d2191..b46e123f28a 100644 --- a/java/src/main/java/org/lance/operation/DataReplacement.java +++ b/java/src/main/java/org/lance/operation/DataReplacement.java @@ -23,9 +23,9 @@ /** * Replace the data files backing some fields of existing fragments with new files, without moving * rows. Each group names a fragment and one new data file for it. At commit, a file holding exactly - * the new file's fields is swapped for it; otherwise the new file's fields are tombstoned where they - * live and the new file is appended. A fragment can take several groups, one per new file. Used for - * null column population, and by compaction to repack columns into fewer files. + * the new file's fields is swapped for it; otherwise the new file's fields are tombstoned where + * they live and the new file is appended. A fragment can take several groups, one per new file. + * Used for null column population, and by compaction to repack columns into fewer files. * *

{@code dataChange == false} declares that the new files hold the same values as the files they * replace: indices keep their coverage, overlays keep shadowing, and no row is reported as updated. diff --git a/java/src/test/java/org/lance/CompactionTest.java b/java/src/test/java/org/lance/CompactionTest.java index 4c874c4c878..d37875a5c27 100644 --- a/java/src/test/java/org/lance/CompactionTest.java +++ b/java/src/test/java/org/lance/CompactionTest.java @@ -18,8 +18,10 @@ import org.lance.compaction.CompactionMode; import org.lance.compaction.CompactionOptions; import org.lance.compaction.CompactionPlan; +import org.lance.compaction.CompactionScope; import org.lance.compaction.CompactionTask; import org.lance.compaction.RewriteResult; +import org.lance.schema.SqlExpressions; import org.apache.arrow.memory.RootAllocator; import org.junit.jupiter.api.Test; @@ -33,9 +35,11 @@ import java.io.ObjectInputStream; import java.io.ObjectOutputStream; import java.nio.file.Path; +import java.util.ArrayList; import java.util.Arrays; import java.util.Base64; import java.util.Collections; +import java.util.List; import java.util.Optional; import static org.junit.jupiter.api.Assertions.assertEquals; @@ -268,6 +272,53 @@ public void testCompactionModeRoundTrip(CompactionMode mode, @TempDir Path tempD } } + @Test + public void testRepackColumns(@TempDir Path tempDir) throws Exception { + String datasetPath = tempDir.resolve("test_repack_columns").toString(); + try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { + TestUtils.SimpleTestDataset testDataset = + new TestUtils.SimpleTestDataset(allocator, datasetPath); + testDataset.createEmptyDataset().close(); + testDataset.write(1, 10).close(); + try (Dataset dataset = testDataset.write(2, 10)) { + dataset.addColumns( + new SqlExpressions.Builder().withExpression("double_id", "id * 2").build(), + Optional.empty()); + assertEquals(2, dataset.getFragments().get(0).metadata().getFiles().size()); + + CompactionOptions options = + CompactionOptions.builder() + .withMaxDataFilesPerFragment(1) + .withScope(CompactionScope.REPACK_COLUMNS) + .build(); + CompactionPlan plan = Compaction.planCompaction(dataset, options); + assertEquals(Optional.of(1L), plan.getCompactionOptions().getMaxDataFilesPerFragment()); + assertEquals( + Optional.of(CompactionScope.REPACK_COLUMNS.getValue()), + plan.getCompactionOptions().getScope()); + assertEquals(2, plan.getCompactionTasks().size()); + + List results = new ArrayList<>(); + for (CompactionTask task : plan.getCompactionTasks()) { + task = serializeAndDeserialize(task); + assertEquals( + Collections.singletonList(Arrays.asList(0, 1, 2)), + task.getTaskData().getRepackFiles()); + RewriteResult result = serializeAndDeserialize(task.execute(dataset)); + assertEquals(1, result.getRepackedFiles().size()); + results.add(result); + } + Compaction.commitCompaction(dataset, results, plan.getCompactionOptions()); + + dataset.checkoutLatest(); + assertEquals(2, dataset.getFragments().size()); + for (Fragment fragment : dataset.getFragments()) { + assertEquals(1, fragment.metadata().getFiles().size()); + } + } + } + } + /** * A serialized CompactionOptions produced by the class as it existed before maxSourceRows and * maxSourceBytes were added (no declared serialVersionUID, stream ends after maxSourceFragments), diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index d068ba6d9d9..8a552edc031 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -7613,6 +7613,9 @@ def compact_files( max_source_bytes: Optional[int] = None, excluded_fragment_ids: Optional[list[int]] = None, data_storage_version: Optional[str] = None, + max_data_files_per_fragment: Optional[int] = None, + column_groups: Optional[list[list[str]]] = None, + scope: Optional[Literal["all", "rewrite_fragments", "repack_columns"]] = None, ) -> CompactionMetrics: """Compacts small files in the dataset, reducing total number of files. @@ -7620,6 +7623,8 @@ def compact_files( * Removes deleted rows from fragments * Removes dropped columns from fragments * Merges small fragments into larger ones + * Repacks a fragment's columns into fewer data files, when + ``max_data_files_per_fragment`` or ``column_groups`` asks for it This method preserves the insertion order of the dataset. This may mean it leaves small fragments in the dataset if they are not adjacent to @@ -7647,6 +7652,9 @@ def compact_files( ``lance.compaction.max_source_fragments``, ``lance.compaction.max_source_rows``, ``lance.compaction.max_source_bytes``, + ``lance.compaction.max_data_files_per_fragment``, + ``lance.compaction.column_groups``, + ``lance.compaction.scope``, ``lance.compaction.data_storage_version``. Parameters @@ -7728,6 +7736,24 @@ def compact_files( Uses the compaction config target when set, otherwise the dataset's default write version. Does not change that default or the versions of unselected files. V1/V2 cross-family targets are rejected. + max_data_files_per_fragment: int, optional + Maximum number of data files a fragment may hold columns in before + its columns are repacked into fewer files. Each ``add_columns`` + backfill adds one file per fragment. A repack rewrites only the + columns that move and keeps rows, fragment ids and indices as they + are. If not specified, uses the manifest config value, or no limit. + column_groups: list[list[str]], optional + Top-level columns to keep in their own data files. Each inner list + becomes one data file per fragment; the columns no group names share + one file. Fragments the compaction rewrites are written this way, + and the others have their columns repacked to match. Names that are + not top-level columns are ignored. If not specified, uses the + manifest config value (``"b, c; d"`` is ``[["b", "c"], ["d"]]``). + scope: str, optional + Which tasks to plan: ``"rewrite_fragments"`` only rewrites + fragments, ``"repack_columns"`` only repacks columns, and + ``"all"`` (the default) does both. A run that does both commits + two versions. Returns ------- @@ -7756,6 +7782,9 @@ def compact_files( max_source_bytes=max_source_bytes, excluded_fragment_ids=excluded_fragment_ids, data_storage_version=data_storage_version, + max_data_files_per_fragment=max_data_files_per_fragment, + column_groups=column_groups, + scope=scope, ).items() if v is not None } diff --git a/python/python/lance/lance/optimize.pyi b/python/python/lance/lance/optimize.pyi index c4b6b6546e6..f4949a5fe0a 100644 --- a/python/python/lance/lance/optimize.pyi +++ b/python/python/lance/lance/optimize.pyi @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import List, Optional +from typing import List, Literal, Optional from lance import LanceDataset from lance.fragment import FragmentMetadata @@ -32,6 +32,7 @@ class RewriteResult: class CompactionTask: read_version: int + kind: Literal["rewrite_fragments", "repack_columns"] fragments: List["FragmentMetadata"] def execute(self, dataset: "LanceDataset") -> RewriteResult: ... diff --git a/python/python/lance/optimize.py b/python/python/lance/optimize.py index ab18255c853..cff80ed61a7 100644 --- a/python/python/lance/optimize.py +++ b/python/python/lance/optimize.py @@ -121,6 +121,28 @@ class CompactionOptions(TypedDict, total=False): are not combined into the same task. Duplicate and unknown IDs are ignored. (default: None) """ + max_data_files_per_fragment: Optional[int] + """ + Maximum number of data files a fragment may hold columns in before its + columns are repacked into fewer files. Each ``add_columns`` backfill adds + one file per fragment. A repack rewrites only the columns that move and + keeps rows, fragment ids and indices as they are. (default: None, no + file-count trigger) + """ + column_groups: Optional[list[list[str]]] + """ + Top-level columns to keep in their own data files. Each inner list becomes + one data file per fragment; the columns no group names share one file. + Fragments the compaction rewrites are written this way, and the others + have their columns repacked to match. Names that are not top-level + columns are ignored. (default: None, one file per fragment) + """ + scope: Optional[Literal["all", "rewrite_fragments", "repack_columns"]] + """ + Which tasks to plan: ``"rewrite_fragments"`` only rewrites fragments, + ``"repack_columns"`` only repacks columns, ``"all"`` does both. + (default: "all") + """ data_storage_version: Optional[str] """ Output data file version, such as "2.2", "stable", or "next". If omitted, diff --git a/python/python/tests/test_optimize.py b/python/python/tests/test_optimize.py index 34199956c1e..fd2fccc3806 100644 --- a/python/python/tests/test_optimize.py +++ b/python/python/tests/test_optimize.py @@ -157,6 +157,56 @@ def test_compact_files_max_source_fragments(tmp_path: Path): assert len(dataset.get_fragments()) == 7 +def _backfilled(tmp_path: Path): + dataset = lance.write_dataset( + pa.table({"a": range(8), "b": range(8)}), + tmp_path / "dataset", + max_rows_per_file=4, + ) + dataset.add_columns({"c": "a + 1"}) + dataset.add_columns({"d": "a + 2"}) + return dataset + + +def _files_per_fragment(dataset): + return [len(fragment.metadata.files) for fragment in dataset.get_fragments()] + + +def test_compact_files_repacks_columns(tmp_path: Path): + dataset = _backfilled(tmp_path) + expected = dataset.to_table() + assert _files_per_fragment(dataset) == [3, 3] + + metrics = dataset.optimize.compact_files( + max_data_files_per_fragment=1, scope="repack_columns" + ) + + assert metrics.files_added == 2 + assert metrics.fragments_added == 0 + assert _files_per_fragment(dataset) == [1, 1] + assert dataset.to_table() == expected + assert [f.fragment_id for f in dataset.get_fragments()] == [0, 1] + + +def test_distributed_repack(tmp_path: Path): + dataset = _backfilled(tmp_path) + expected = dataset.to_table() + + plan = Compaction.plan( + dataset, options=dict(column_groups=[["d"]], scope="repack_columns") + ) + assert [task.kind for task in plan.tasks] == ["repack_columns"] * 2 + results = [ + pickle.loads(pickle.dumps(pickle.loads(pickle.dumps(task)).execute(dataset))) + for task in plan.tasks + ] + Compaction.commit(dataset, results) + + dataset = lance.dataset(dataset.uri) + assert _files_per_fragment(dataset) == [2, 2] + assert dataset.to_table() == expected + + def test_blob_compaction(tmp_path: Path): base_dir = tmp_path / "blob_dataset" blob_field = pa.field( diff --git a/python/src/dataset/optimize.rs b/python/src/dataset/optimize.rs index e17023b7450..2fb2e2578c4 100644 --- a/python/src/dataset/optimize.rs +++ b/python/src/dataset/optimize.rs @@ -15,8 +15,9 @@ use lance::dataset::{ index::DatasetIndexRemapperOptions, optimize::{ - CompactionMetrics, CompactionMode, CompactionOptions, CompactionPlan, CompactionTask, - RewriteResult, commit_compaction, compact_files, plan_compaction, + CompactionMetrics, CompactionMode, CompactionOptions, CompactionPlan, CompactionScope, + CompactionTask, CompactionTaskKind, RewriteResult, commit_compaction, compact_files, + plan_compaction, }, }; use pyo3::{exceptions::PyNotImplementedError, pyclass::CompareOp, types::PyTuple}; @@ -92,6 +93,21 @@ fn parse_compaction_options( opts.data_storage_version = Some(version.parse().infer_error()?); } } + "max_data_files_per_fragment" => { + opts.max_data_files_per_fragment = value.extract()?; + } + "column_groups" => { + opts.column_groups = value + .extract::>>>()? + .unwrap_or_default(); + } + "scope" => { + let scope: Option = value.extract()?; + if let Some(scope) = scope { + opts.scope = CompactionScope::try_from(scope.as_str()) + .map_err(|e| PyValueError::new_err(e.to_string()))?; + } + } _ => { return Err(PyValueError::new_err(format!( "Invalid compaction option: {}", @@ -129,8 +145,9 @@ pub struct PyCompactionMetrics { /// int : The number of files that have been removed, including deletion files. #[pyo3(get)] pub files_removed: usize, - /// int : The number of files that have been added, which is always equal to the - /// number of fragments. + /// int : The number of data files that have been added. A rewrite adds one + /// per new fragment (one per group with ``column_groups``), a column repack + /// one per new file. #[pyo3(get)] pub files_added: usize, } @@ -263,8 +280,10 @@ impl PyCompactionTask { .collect::>>()? .join(", "); Ok(format!( - "CompactionTask(read_version={}, fragments=[{}])", - self.0.read_version, fragment_reprs + "CompactionTask(read_version={}, kind={}, fragments=[{}])", + self.0.read_version, + self.kind(), + fragment_reprs )) } @@ -274,6 +293,16 @@ impl PyCompactionTask { self.0.read_version } + /// str : ``"rewrite_fragments"`` for a task that rewrites its fragments, + /// ``"repack_columns"`` for one that repacks its fragment's columns. + #[getter] + pub fn kind(&self) -> &'static str { + match self.0.task.kind { + CompactionTaskKind::RewriteFragments => "rewrite_fragments", + CompactionTaskKind::RepackColumns { .. } => "repack_columns", + } + } + /// List[lance.fragment.FragmentMetadata] : The fragments that will be compacted. #[getter] pub fn fragments<'py>(&self, py: Python<'py>) -> PyResult>> { diff --git a/rust/lance-table/src/transaction/operation.rs b/rust/lance-table/src/transaction/operation.rs index b3b75d79dee..f7121cc3aea 100644 --- a/rust/lance-table/src/transaction/operation.rs +++ b/rust/lance-table/src/transaction/operation.rs @@ -118,14 +118,14 @@ pub enum Operation { /// they live and the new file is appended, and a file left holding no /// field of the schema is dropped. A fragment can take several groups, one /// per new file, applied in order. Used for null column population, and by - /// horizontal compaction to repack columns into fewer files. + /// compaction to repack columns into fewer files. /// /// The fields of a legacy (V1) data file cannot be tombstoned one by one, /// so a field held by a V1 file can only be replaced by an exact match. DataReplacement { replacements: Vec, /// Whether the new files change any value. `false` means the values - /// were only moved to new files (horizontal compaction): indices keep + /// were only moved to new files (a compaction repack): indices keep /// their coverage of the replaced fields, overlays keep shadowing, and /// no row is stamped as updated. data_change: bool, diff --git a/rust/lance/src/dataset.rs b/rust/lance/src/dataset.rs index 3059abb79f2..0e1b5850f0f 100644 --- a/rust/lance/src/dataset.rs +++ b/rust/lance/src/dataset.rs @@ -83,7 +83,6 @@ pub mod optimize; pub(crate) mod overlay; pub mod progress; pub mod refs; -pub mod rewrite_columns; pub mod rowids; pub mod scanner; mod schema_evolution; diff --git a/rust/lance/src/dataset/compaction_stats.rs b/rust/lance/src/dataset/compaction_stats.rs index 462d49a1bd8..14b6ebd3f24 100644 --- a/rust/lance/src/dataset/compaction_stats.rs +++ b/rust/lance/src/dataset/compaction_stats.rs @@ -3,11 +3,11 @@ //! Per-fragment column-layout statistics. //! -//! A compaction planner keys on how many data files each fragment carries: a -//! fragment split across many small per-column files (typically from repeated -//! `add_columns` backfills) is the signal horizontal compaction targets. -//! Exposing these stats via [`Dataset::column_layout_stats`] keeps the planning -//! decision observable instead of buried in a black-box heuristic. +//! The default compaction planner repacks a fragment's columns when it holds +//! them in more than `max_data_files_per_fragment` files, counted here. +//! Repeated `add_columns` backfills are the usual cause. Exposing the counts +//! via [`Dataset::column_layout_stats`] lets a caller see why a fragment was, +//! or was not, picked, and lets a custom planner make its own call. use std::collections::HashSet; @@ -18,9 +18,9 @@ use super::Dataset; pub struct FragmentColumnLayoutStats { /// The fragment these stats describe. pub fragment_id: u64, - /// Number of data files holding at least one column of the dataset schema. - /// A large value on a wide dataset is what horizontal compaction collapses - /// back into fewer files. A file left holding only tombstones or dropped + /// Number of data files holding at least one column of the dataset schema, + /// the count `max_data_files_per_fragment` is compared with. A file left + /// holding only tombstones or dropped /// columns, or kept only for the fragment's spilled row lineage, holds no /// column a read touches and is not counted. pub live_file_count: usize, @@ -31,8 +31,8 @@ pub struct FragmentColumnLayoutStats { impl Dataset { /// Per-fragment column-layout stats, in manifest fragment order. /// - /// This is the planning input for horizontal compaction: it reads only - /// fragment metadata (no data files), so it is cheap to call. + /// The default planner's input for column repacks. It reads only fragment + /// metadata (no data files), so it is cheap to call. pub fn column_layout_stats(&self) -> Vec { let schema_ids: HashSet = self .schema() @@ -116,9 +116,8 @@ mod tests { /// Only files holding a schema column count. The extra files are added to /// the manifest by hand: the commit paths that leave such files behind (an - /// in-place column update after a drop, a full horizontal rewrite on a - /// table whose lineage was spilled into a data file) are not what this - /// test is about. + /// in-place column update after a drop, a repack on a table whose lineage + /// was spilled into a data file) are not what this test is about. #[tokio::test] async fn column_layout_stats_skips_files_without_schema_columns() { use lance_file::version::ConcreteFileVersion; diff --git a/rust/lance/src/dataset/optimize.rs b/rust/lance/src/dataset/optimize.rs index 14146ea4548..5f14cb8d6b8 100644 --- a/rust/lance/src/dataset/optimize.rs +++ b/rust/lance/src/dataset/optimize.rs @@ -146,6 +146,7 @@ use tracing::{info, warn}; pub(super) mod binary_copy; pub mod remapping; +mod repack; use crate::index::frag_reuse::build_new_frag_reuse_index; use crate::io::deletion::read_dataset_deletion_file; @@ -316,6 +317,22 @@ pub struct CompactionOptions { /// carries any overlay, or `None` to disable the overlay-count trigger /// entirely. pub max_overlays_per_fragment: Option, + /// Maximum number of data files a fragment may hold columns in before its + /// columns are repacked into fewer files. Each `add_columns` backfill adds + /// one file per fragment; a large, deletion-free fragment never qualifies + /// for a rewrite, so this is the trigger that collapses those files. The + /// repack rewrites only the columns that move and keeps rows, fragment ids + /// and indices as they are (see [`CompactionTaskKind::RepackColumns`]). + /// The count is [`FragmentColumnLayoutStats::live_file_count`]. + /// + /// Defaults to `None` (no file-count trigger). Must be at least 1. + /// + /// [`FragmentColumnLayoutStats::live_file_count`]: crate::dataset::compaction_stats::FragmentColumnLayoutStats::live_file_count + #[serde(default)] + pub max_data_files_per_fragment: Option, + /// Which kinds of task this run plans. Defaults to [`CompactionScope::All`]. + #[serde(default)] + pub scope: CompactionScope, /// Exact data file version for compacted output. /// /// If omitted, use the dataset's default write version without changing it. @@ -331,19 +348,22 @@ pub struct CompactionOptions { /// }; /// ``` pub data_storage_version: Option, - /// Top-level columns to keep in their own data files across this compaction. + /// Top-level columns to keep in their own data files. /// - /// Each inner list becomes one data file per compacted fragment holding - /// exactly those columns; every column not listed goes to one shared file. - /// Empty (the default) writes every column to a single file per fragment. + /// Each inner list becomes one data file per fragment holding exactly + /// those columns; every column not listed goes to one shared file. Empty + /// (the default) puts every column in a single file per fragment. /// - /// This is the layout [`Dataset::rewrite_columns`] produces, so a compaction - /// configured with the same groups preserves that split instead of folding - /// wide columns back in with narrow ones. Binary copy is disabled when - /// groups are set (`ForceBinaryCopy` is rejected), and `max_bytes_per_file` - /// is ignored so every group's files split at the same rows. Set from - /// dataset config with `lance.compaction.column_groups` (see - /// [`Self::from_dataset_config`]). + /// Fragments a compaction rewrites are written in this layout. A fragment + /// it leaves alone whose files do not match the layout has its columns + /// repacked to match (see [`CompactionTaskKind::RepackColumns`]), unless + /// `scope` is [`CompactionScope::RewriteFragments`]. A name that is not a + /// top-level column of the dataset (say, one dropped or renamed since the + /// groups were configured) is ignored with a warning. Binary copy is + /// disabled when groups are set (`ForceBinaryCopy` is rejected), and + /// `max_bytes_per_file` is ignored so every group's files split at the + /// same rows. Set from dataset config with `lance.compaction.column_groups` + /// (see [`Self::from_dataset_config`]). #[serde(default)] pub column_groups: Vec>, /// Transaction properties to store with this commit. @@ -379,6 +399,8 @@ impl Default for CompactionOptions { max_source_bytes: None, excluded_fragment_ids: Vec::new(), max_overlays_per_fragment: Some(10), + max_data_files_per_fragment: None, + scope: CompactionScope::All, data_storage_version: None, column_groups: Vec::new(), transaction_properties: None, @@ -410,6 +432,8 @@ impl CompactionOptions { /// - `lance.compaction.max_source_rows` /// - `lance.compaction.max_source_bytes` /// - `lance.compaction.max_overlays_per_fragment` + /// - `lance.compaction.max_data_files_per_fragment` + /// - `lance.compaction.scope`: `all`, `rewrite_fragments` or `repack_columns` /// - `lance.compaction.data_storage_version` /// - `lance.compaction.column_groups`: groups separated by `;`, columns in /// a group by `,`, surrounding whitespace trimmed (`"b, c; d"` is @@ -562,6 +586,17 @@ impl CompactionOptions { })?), }; } + "scope" => { + self.scope = CompactionScope::try_from(value.as_str())?; + } + "max_data_files_per_fragment" => { + self.max_data_files_per_fragment = Some(value.parse().map_err(|_| { + Error::invalid_input(format!( + "Invalid value for {}: '{}' (expected a positive integer)", + key, value + )) + })?); + } "column_groups" => { self.column_groups = value .split(';') @@ -609,6 +644,13 @@ impl CompactionOptions { )); } + if self.max_data_files_per_fragment == Some(0) { + return Err(Error::invalid_input( + "CompactionOptions::max_data_files_per_fragment must be at least 1 \ + (use None for no limit)", + )); + } + for (name, value) in [ ( "max_source_fragments", @@ -813,9 +855,9 @@ pub struct CompactionMetrics { pub fragments_added: usize, /// The number of files that have been removed, including deletion files. pub files_removed: usize, - /// The number of data files that have been added. This equals the number of - /// new fragments for a plain vertical compaction, but a `column_groups` - /// compaction adds one file per group per fragment, so it can be larger. + /// The number of data files that have been added: one per new fragment for + /// a fragment rewrite (one per group with `column_groups`), and one per new + /// file for a column repack. pub files_added: usize, } @@ -845,119 +887,49 @@ pub trait CompactionPlanner: Send + Sync { async fn plan(&self, dataset: &Dataset) -> Result; } -/// The executor half of a pluggable compaction pipeline: run one task the -/// planner produced, yielding a result the paired [`CompactionCommitter`] -/// commits. Vertical (`RewriteExecutor`) and, later, horizontal compaction are -/// separate implementations selected by mode — not variants of one plan, so -/// each keeps its own task/result types and serialized shapes. +/// Runs one task of a [`CompactionPlan`]. Pass an implementation to +/// [`compact_files_with_executor`] to run the tasks somewhere else (say, on a +/// cluster) or to wrap the default one; [`CompactionTask::execute`] runs the +/// default one for a task executed by hand. +/// +/// The result goes to [`commit_compaction`], which commits it by its kind. #[async_trait::async_trait] pub trait CompactionExecutor: Send + Sync { - /// The task unit the paired planner emits. - type Task: Send; - /// Per-task output, collected and handed to the committer. - type TaskResult: Send; - - /// Execute one task against `dataset` (already at the plan's read version). + /// Execute `task` against `dataset`, which is at the plan's read version. async fn execute( &self, - dataset: Cow<'_, Dataset>, - task: Self::Task, - options: &CompactionOptions, - ) -> Result; -} - -/// The commit half of a pluggable compaction pipeline: fold the executor's -/// results into one transaction. Each pipeline commits its own operation -/// (vertical → `Operation::Rewrite`; horizontal → `Operation::Update`). -#[async_trait::async_trait] -pub trait CompactionCommitter: Send + Sync { - /// Must match the paired executor's [`CompactionExecutor::TaskResult`]. - type TaskResult: Send; - - /// Commit the collected results, returning the run's metrics. - async fn commit( - &self, - dataset: &mut Dataset, - results: Vec, - remap_options: Arc, + dataset: &Dataset, + task: TaskData, options: &CompactionOptions, - ) -> Result; + ) -> Result; } -/// Vertical compaction: rewrite groups of fragments via `rewrite_files`, -/// committed as `Operation::Rewrite` by [`RewriteCommitter`]. +/// Runs each kind of task the way this crate does: a +/// [`CompactionTaskKind::RewriteFragments`] task rewrites its fragments into +/// new ones, a [`CompactionTaskKind::RepackColumns`] task writes new data +/// files for its fragment. #[derive(Debug, Default, Clone, Copy)] -pub struct RewriteExecutor; +pub struct DefaultCompactionExecutor; #[async_trait::async_trait] -impl CompactionExecutor for RewriteExecutor { - type Task = TaskData; - type TaskResult = RewriteResult; - +impl CompactionExecutor for DefaultCompactionExecutor { async fn execute( &self, - dataset: Cow<'_, Dataset>, + dataset: &Dataset, task: TaskData, options: &CompactionOptions, ) -> Result { - rewrite_files(dataset, task, options).await - } -} - -/// Commits vertical compaction results as a single `Operation::Rewrite`. -#[derive(Debug, Default, Clone, Copy)] -pub struct RewriteCommitter; - -#[async_trait::async_trait] -impl CompactionCommitter for RewriteCommitter { - type TaskResult = RewriteResult; - - async fn commit( - &self, - dataset: &mut Dataset, - results: Vec, - remap_options: Arc, - options: &CompactionOptions, - ) -> Result { - commit_compaction(dataset, results, remap_options, options).await + match task.kind { + CompactionTaskKind::RewriteFragments => { + rewrite_files(Cow::Borrowed(dataset), task, options).await + } + CompactionTaskKind::RepackColumns { .. } => { + repack::execute_repack(dataset, task, options).await + } + } } } -/// Run a compaction pipeline: execute every task with bounded concurrency, then -/// commit the collected results. Both vertical -/// (`RewriteExecutor`/`RewriteCommitter`) and horizontal -/// (`RewriteColumnsExecutor`/`RewriteColumnsCommitter`) compaction run through -/// it. The two pipelines still differ in how they produce tasks and which -/// operation they commit; only the execute-collect-commit skeleton is shared. -pub(crate) async fn run_compaction_pipeline( - dataset: &mut Dataset, - executor: E, - committer: C, - tasks: Vec, - remap_options: Arc, - options: &CompactionOptions, - concurrency: usize, -) -> Result -where - E: CompactionExecutor, - C: CompactionCommitter, -{ - if tasks.is_empty() { - return Ok(CompactionMetrics::default()); - } - // Execute against an immutable snapshot; `commit` takes `&mut dataset` once - // every task has finished and the borrows are released. - let dataset_ref = &dataset.clone(); - let results: Vec = futures::stream::iter(tasks) - .map(|task| executor.execute(Cow::Borrowed(dataset_ref), task, options)) - .buffer_unordered(concurrency.max(1)) - .try_collect() - .await?; - committer - .commit(dataset, results, remap_options, options) - .await -} - /// Formulate a plan to compact the files in a dataset /// /// The compaction plan will contain a list of tasks to execute. Each task @@ -966,6 +938,9 @@ where /// tasks may contain a single fragment when that fragment has deletions that /// are being materialized and doesn't have any neighbors that need to be /// compacted. +/// +/// Fragments no rewrite task takes are then checked for a column repack +/// against [`Dataset::column_layout_stats`], one task per fragment. #[derive(Debug, Clone, Default)] pub struct DefaultCompactionPlanner { options: CompactionOptions, @@ -1004,6 +979,33 @@ impl CompactionPlanner for DefaultCompactionPlanner { )); } + let mut all_tasks = if self.options.scope == CompactionScope::RepackColumns { + Vec::new() + } else { + self.plan_rewrites(dataset).await? + }; + if self.options.scope != CompactionScope::RewriteFragments { + let rewritten: HashSet = all_tasks + .iter() + .flat_map(|(task, _)| task.fragments.iter().map(|fragment| fragment.id)) + .collect(); + all_tasks.extend(self.plan_repacks(dataset, &rewritten)?); + } + + let tasks = limit_tasks_to_source_budget(&self.options, dataset.schema(), all_tasks)?; + + let mut options = self.options.clone(); + options.data_storage_version = Some(write_version.to_selector()); + let mut compaction_plan = CompactionPlan::new(dataset.manifest.version, options); + compaction_plan.extend_tasks(tasks); + + Ok(compaction_plan) + } +} + +impl DefaultCompactionPlanner { + /// The tasks that rewrite fragments, each with its live row count. + async fn plan_rewrites(&self, dataset: &Dataset) -> Result> { // get_fragments should be returning fragments in sorted order (by id) // and fragment ids should be unique let fragments = dataset.get_fragments(); @@ -1176,23 +1178,64 @@ impl CompactionPlanner for DefaultCompactionPlanner { .flat_map(|bin| bin.split_for_size(self.options.target_rows_per_fragment)) .map(|bin| { let live_rows = bin.row_counts.iter().sum(); - ( + (TaskData::rewrite_fragments(bin.fragments), live_rows) + }) + .collect(); + Ok(all_tasks) + } + + /// The tasks that repack the columns of a fragment not in `rewritten`, + /// each with its live row count. A fragment is repacked when its files + /// do not match `column_groups`, or when it holds columns in more than + /// `max_data_files_per_fragment` files, as + /// [`Dataset::column_layout_stats`] counts them. + fn plan_repacks( + &self, + dataset: &Dataset, + rewritten: &HashSet, + ) -> Result> { + let groups = if self.options.column_groups.is_empty() { + None + } else { + Some(column_group_ids( + dataset.schema(), + &self.options.column_groups, + )?) + }; + let max_files = self.options.max_data_files_per_fragment; + if groups.is_none() && max_files.is_none() { + return Ok(Vec::new()); + } + let stats = dataset.column_layout_stats(); + let tasks = dataset + .manifest + .fragments + .iter() + .zip(stats) + .filter(|(fragment, _)| { + !rewritten.contains(&fragment.id) + && !u32::try_from(fragment.id) + .is_ok_and(|id| self.excluded_fragment_ids.contains(id)) + }) + .filter_map(|(fragment, stats)| { + let files = repack::plan_fragment_repack( + dataset.schema(), + fragment, + stats.live_file_count, + groups.as_deref(), + max_files, + )?; + let live_rows = fragment.num_rows().unwrap_or_default(); + Some(( TaskData { - fragments: bin.fragments, + fragments: vec![fragment.clone()], + kind: CompactionTaskKind::RepackColumns { files }, }, live_rows, - ) + )) }) .collect(); - - let tasks = limit_tasks_to_source_budget(&self.options, dataset.schema(), all_tasks)?; - - let mut options = self.options.clone(); - options.data_storage_version = Some(write_version.to_selector()); - let mut compaction_plan = CompactionPlan::new(dataset.manifest.version, options); - compaction_plan.extend_tasks(tasks); - - Ok(compaction_plan) + Ok(tasks) } } @@ -1203,6 +1246,14 @@ impl CompactionPlanner for DefaultCompactionPlanner { /// * Removes dropped columns from fragments. /// * Merges fragments that are too small. /// +/// With [`CompactionOptions::max_data_files_per_fragment`] or +/// [`CompactionOptions::column_groups`] set, it also repacks the columns of a +/// fragment spread over too many data files, or laid out other than the +/// groups ask, into new files without moving its rows. A run that both +/// rewrites fragments and repacks columns commits two versions (see +/// [`commit_compaction`]). [`CompactionOptions::scope`] limits a run to one +/// kind of task. +/// /// This method tries to preserve the insertion order of rows in the dataset. /// /// If no compaction is needed, this method will not make a new version of the table. @@ -1220,6 +1271,17 @@ pub async fn compact_files_with_planner( dataset: &mut Dataset, remap_options: Option>, // These will be deprecated later planner: &dyn CompactionPlanner, +) -> Result { + compact_files_with_executor(dataset, remap_options, planner, &DefaultCompactionExecutor).await +} + +/// Plan with `planner`, run every task with `executor`, and commit the +/// results with [`commit_compaction`]. +pub async fn compact_files_with_executor( + dataset: &mut Dataset, + remap_options: Option>, // These will be deprecated later + planner: &dyn CompactionPlanner, + executor: &dyn CompactionExecutor, ) -> Result { let compaction_plan: CompactionPlan = planner.plan(dataset).await?; @@ -1227,8 +1289,14 @@ pub async fn compact_files_with_planner( // tagged entry, which only the deferred-remap commit path does; eager // remapping would rewrite provenance the tagged reader depends on. // Checked before any file is rewritten (commit_compaction re-checks for - // callers that commit externally planned results). - if !compaction_plan.options.defer_index_remap + // callers that commit externally planned results). A repack moves no row, + // so only a plan that rewrites fragments is refused. + let rewrites_fragments = compaction_plan + .tasks + .iter() + .any(|task| task.kind == CompactionTaskKind::RewriteFragments); + if rewrites_fragments + && !compaction_plan.options.defer_index_remap && dataset.manifest.writer_feature_flags & lance_table::feature_flags::FLAG_FRAGMENT_REUSE_INDEX != 0 @@ -1254,16 +1322,16 @@ pub async fn compact_files_with_planner( .num_threads .unwrap_or_else(get_num_compute_intensive_cpus); let remap_options = remap_options.unwrap_or(Arc::new(DatasetIndexRemapperOptions::default())); - run_compaction_pipeline( - dataset, - RewriteExecutor, - RewriteCommitter, - compaction_plan.tasks, - remap_options, - &compaction_plan.options, - concurrency, - ) - .await + let options = &compaction_plan.options; + // Execute against an immutable snapshot; the commit takes `&mut dataset` + // once every task has finished. + let snapshot = &dataset.clone(); + let results: Vec = futures::stream::iter(compaction_plan.tasks.clone()) + .map(|task| executor.execute(snapshot, task, options)) + .buffer_unordered(concurrency.max(1)) + .try_collect() + .await?; + commit_compaction(dataset, results, remap_options, options).await } /// Information about a fragment used to decide its fate in compaction @@ -2288,6 +2356,65 @@ async fn prepare_reader( pub struct TaskData { /// The fragments to compact. pub fragments: Vec, + /// What the task does with them. Absent in a task serialized before + /// repacking existed, which rewrote its fragments. + #[serde(default)] + pub kind: CompactionTaskKind, +} + +impl TaskData { + /// A task that rewrites `fragments` into new fragments. + pub fn rewrite_fragments(fragments: Vec) -> Self { + Self { + fragments, + kind: CompactionTaskKind::RewriteFragments, + } + } +} + +/// What a compaction task does. +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] +pub enum CompactionTaskKind { + /// Rewrite the task's fragments into new fragments, merging small ones, + /// dropping deleted rows and materializing overlays. Moves rows, so + /// indices covering the fragments are remapped. + #[default] + RewriteFragments, + /// Rewrite some columns of the task's one fragment into new data files, + /// leaving the fragment id, rows, deletions, overlays and index coverage + /// as they are. Each entry of `files` is one new file, given as the field + /// ids of the top-level columns it holds. The files the columns leave + /// keep their other columns; a file left with none is dropped. + RepackColumns { files: Vec> }, +} + +/// Which kinds of task a compaction plans. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +pub enum CompactionScope { + /// Rewrite the fragments that need it and repack the columns of the + /// others. + #[default] + All, + /// Only rewrite fragments. + RewriteFragments, + /// Only repack columns. + RepackColumns, +} + +impl TryFrom<&str> for CompactionScope { + type Error = Error; + + fn try_from(value: &str) -> std::result::Result { + match value.to_lowercase().as_str() { + "all" => Ok(Self::All), + "rewrite_fragments" => Ok(Self::RewriteFragments), + "repack_columns" => Ok(Self::RepackColumns), + _ => Err(Error::invalid_input(format!( + "Invalid compaction scope \"{}\". Valid values: \"all\", \"rewrite_fragments\", \"repack_columns\"", + value + ))), + } + } } /// A standalone task that can be serialized and sent to another machine for @@ -2314,9 +2441,8 @@ impl CompactionTask { } else { Cow::Owned(dataset.checkout_version(self.read_version).await?) }; - let executor = RewriteExecutor; - executor - .execute(dataset, self.task.clone(), &self.options) + DefaultCompactionExecutor + .execute(dataset.as_ref(), self.task.clone(), &self.options) .await } } @@ -2570,6 +2696,19 @@ pub struct RewriteResult { /// deferred index remap post-processing, or (2) used with reserved /// fragment IDs to build old-to-new mappings. pub row_addrs: Option>, + /// The new data files of a [`CompactionTaskKind::RepackColumns`] task. + /// `new_fragments` and `original_fragments` are then empty. + #[serde(default)] + pub repacked_files: Option, +} + +/// The data files a column repack wrote for one fragment. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub struct RepackedFiles { + /// The fragment the files belong to. + pub fragment_id: u64, + /// The new files, empty when the fragment had no rows to write. + pub files: Vec, } async fn reserve_fragment_ids( @@ -2611,6 +2750,10 @@ async fn reserve_fragment_ids( /// The schema of each data file a compacted fragment gets when column groups /// are configured: the columns no group claims first (when any), then one /// schema per group. Columns keep the dataset's order within each file. +/// +/// A name that is not a top-level column is skipped with a warning: the groups +/// usually come from table config, which `drop_columns` and a rename leave as +/// they are, and a stale name must not stop every later compaction. fn column_group_schemas( schema: &lance_core::datatypes::Schema, groups: &[Vec], @@ -2618,11 +2761,6 @@ fn column_group_schemas( let names = || schema.fields.iter().map(|field| field.name.as_str()); let mut seen = HashSet::new(); for column in groups.iter().flatten() { - if !names().any(|name| name == column) { - return Err(Error::invalid_input(format!( - "column_groups names \"{column}\", which is not a top-level column of the dataset" - ))); - } // Checked here as well as in `validate()`: a custom planner or a // hand-built `CompactionTask` reaches `rewrite_files` without it, and a // column in two groups would be written to two files. @@ -2631,6 +2769,12 @@ fn column_group_schemas( "column_groups lists column \"{column}\" more than once" ))); } + if !names().any(|name| name == column) { + warn!( + "Ignoring column_groups entry \"{column}\": it is not a top-level column \ + of the dataset" + ); + } } let claimed = |name: &str| groups.iter().flatten().any(|column| column == name); let rest: Vec<&str> = names().filter(|name| !claimed(name)).collect(); @@ -2645,6 +2789,17 @@ fn column_group_schemas( .collect() } +/// [`column_group_schemas`] as the field ids of each file's top-level columns. +fn column_group_ids( + schema: &lance_core::datatypes::Schema, + groups: &[Vec], +) -> Result>> { + Ok(column_group_schemas(schema, groups)? + .iter() + .map(|group| group.fields.iter().map(|field| field.id).collect()) + .collect()) +} + /// Compaction's grouped write: fan one read of the merged rows out to one /// writer per column group, so each output fragment gets one data file per /// group instead of one file holding every column. All writers break at the @@ -2844,6 +2999,7 @@ async fn rewrite_files( read_version: dataset.manifest.version, original_fragments: task.fragments, row_addrs: None, + repacked_files: None, }); } @@ -3191,6 +3347,7 @@ async fn rewrite_files( read_version: dataset.manifest.version, original_fragments: fragments, row_addrs, + repacked_files: None, }) } @@ -3555,11 +3712,32 @@ fn append_row_lineage_columns( /// they can be omitted and the successful tasks can be committed. However, once /// some of the tasks have been committed, the remainder of the tasks will not /// be able to be committed and should be considered cancelled. +/// +/// Fragment rewrites commit as one `Operation::Rewrite`; column repacks then +/// commit as one `Operation::DataReplacement` that moves values without +/// changing them, so a run with both kinds of task makes two versions. A +/// planner never puts one fragment in both kinds of task. If the second +/// commit fails, the first stays committed and the error is returned; the +/// repacks can be planned again. pub async fn commit_compaction( dataset: &mut Dataset, completed_tasks: Vec, remap_options: Arc, options: &CompactionOptions, +) -> Result { + let (repacks, rewrites): (Vec<_>, Vec<_>) = completed_tasks + .into_iter() + .partition(|task| task.repacked_files.is_some()); + let mut metrics = commit_rewrites(dataset, rewrites, remap_options, options).await?; + metrics += repack::commit_repacked_files(dataset, repacks, options).await?; + Ok(metrics) +} + +async fn commit_rewrites( + dataset: &mut Dataset, + completed_tasks: Vec, + remap_options: Arc, + options: &CompactionOptions, ) -> Result { if completed_tasks.is_empty() { return Ok(CompactionMetrics::default()); @@ -4063,6 +4241,7 @@ async fn cleanup_compaction_files_after_reservation_failure( mod tests { mod binary_copy; + mod repack; use self::remapping::RemappedIndex; use super::*; use crate::dataset::WriteDestination; @@ -9112,12 +9291,10 @@ mod tests { assert!(checked > 0, "expected to check at least one stored vector"); } - /// Horizontal compaction ([`Dataset::rewrite_columns`]) repacks a column's - /// files without moving rows, so a vector index over that column keeps its - /// uuid and coverage and is never rebuilt — the property #9291 only tested - /// for BTree. + /// A column repack moves no row, so a vector index over the repacked + /// column keeps its uuid and coverage and answers the same. #[tokio::test] - async fn rewrite_columns_preserves_vector_index() { + async fn repack_columns_preserves_vector_index() { use arrow_array::cast::AsArray; use arrow_array::types::{Float32Type, Int32Type}; use lance_datagen::Dimension; @@ -9172,7 +9349,10 @@ mod tests { }; let knn_before = vector_knn_ids(&dataset, &query, 5).await; - dataset.rewrite_columns(&["vec"], None).await.unwrap(); + repack_columns(&mut dataset, vec![vec!["vec".into()]]).await; + for fragment in dataset.get_fragments() { + assert_eq!(fragment.metadata().files.len(), 2, "vec is split out"); + } let after = dataset .load_index_by_name("vec_idx") @@ -12482,16 +12662,27 @@ mod tests { } #[tokio::test] - async fn compaction_column_groups_reject_unknown_column() { - let (mut dataset, _) = column_group_dataset().await; + async fn compaction_column_groups_ignore_unknown_column() { + // A configured group can outlive the column it names (drop_columns and + // renames leave the config as it is), so the name is skipped and the + // rest of the groups still apply. + let (mut dataset, before) = column_group_dataset().await; let options = CompactionOptions { - column_groups: vec![vec!["nonexistent".into()]], + column_groups: vec![vec!["nonexistent".into(), "c".into()]], ..Default::default() }; - let err = compact_files(&mut dataset, options, None) - .await - .unwrap_err(); - assert!(matches!(err, Error::InvalidInput { .. }), "{err}"); + compact_files(&mut dataset, options, None).await.unwrap(); + let fragments = dataset.get_fragments(); + assert_eq!(fragments.len(), 1); + let mut groups: Vec> = fragments[0] + .metadata() + .files + .iter() + .map(|file| file.fields.to_vec()) + .collect(); + groups.sort(); + assert_eq!(groups, vec![vec![0, 1], vec![2]]); + assert_eq!(dataset.scan().try_into_batch().await.unwrap(), before); } #[tokio::test] @@ -12589,11 +12780,21 @@ mod tests { dataset.validate().await.unwrap(); } + /// Run only the column repacks `column_groups` calls for. + async fn repack_columns(dataset: &mut Dataset, column_groups: Vec>) { + let options = CompactionOptions { + column_groups, + scope: CompactionScope::RepackColumns, + ..Default::default() + }; + compact_files(dataset, options, None).await.unwrap(); + } + #[tokio::test] - async fn column_groups_preserve_horizontal_split_across_compaction() { + async fn column_groups_preserve_repacked_split_across_compaction() { let (mut dataset, before) = column_group_dataset().await; - // Horizontal split: column c into its own file per fragment. - dataset.rewrite_columns(&["c"], None).await.unwrap(); + // Split column c into its own file per fragment. + repack_columns(&mut dataset, vec![vec!["c".into()]]).await; for fragment in dataset.get_fragments() { assert_eq!( fragment.metadata().files.len(), diff --git a/rust/lance/src/dataset/optimize/repack.rs b/rust/lance/src/dataset/optimize/repack.rs new file mode 100644 index 00000000000..fc7a667cc0d --- /dev/null +++ b/rust/lance/src/dataset/optimize/repack.rs @@ -0,0 +1,374 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Column repacking, the compaction task that rewrites one fragment's columns +//! into fewer data files without moving rows. +//! +//! Every `add_columns` backfill gives each fragment one more data file, and a +//! large fragment with few deletions never qualifies for a rewrite, so its +//! files pile up. A [`CompactionTaskKind::RepackColumns`] task reads some of +//! the fragment's columns, writes them to new files, and commits the files as +//! an `Operation::DataReplacement` with `data_change: false`: the fragment id, +//! row addresses, overlays and index coverage are left as they are. + +use std::collections::{BTreeSet, HashSet}; +use std::sync::Arc; + +use lance_core::datatypes::{Field as LanceField, Schema}; +use lance_file::version::ConcreteFileVersion; +use lance_table::format::{DataFile, Fragment}; + +use super::{ + CompactionMetrics, CompactionOptions, CompactionTaskKind, RepackedFiles, RewriteResult, + TaskData, +}; +use crate::dataset::fragment::FileFragment; +use crate::dataset::transaction::{DataReplacementGroup, Operation, TransactionBuilder}; +use crate::dataset::{Dataset, cleanup_data_fragments, versions}; +use crate::{Error, Result}; + +/// The field ids of `field` and everything beneath it. +fn subtree_ids(field: &LanceField, ids: &mut HashSet) { + ids.insert(field.id); + for child in &field.children { + subtree_ids(child, ids); + } +} + +fn holds_blob(field: &LanceField) -> bool { + field.is_blob() || field.children.iter().any(holds_blob) +} + +/// Every field id a data file answers for, the way a `DataReplacement` +/// commit counts it. +fn file_coverage(file: &DataFile, schema: &Schema) -> HashSet { + file.schema(schema) + .field_ids() + .into_iter() + .chain(file.fields.iter().copied()) + .filter(|id| *id >= 0) + .collect() +} + +/// The new files a repack of `fragment` should write, each the top-level +/// field ids it holds, or `None` when no repack is due. +/// +/// `groups` partitions the top-level columns into the files wanted, as +/// `column_group_schemas` orders them. Without groups, the columns of one +/// live file stay together (the file holding the fragment's spilled row +/// lineage, else the one holding the most bytes) and every other column goes +/// to one new file; when `max_files` is 1, every column goes to one file. A +/// group is written unless one live file holds exactly its columns. A column +/// the fragment has no data for (added as all nulls) is left out of its +/// group, and a group holding a blob column, or a column whose fields are +/// only partly in the fragment's files, is left as it is. +/// +/// The new files carry no row lineage, so a repack that would move every +/// column out of the file holding the fragment's spilled lineage is not +/// planned: that file would stay for the lineage alone, and a later rewrite +/// of the fragment is what reclaims it. +/// +/// A group out of place is a reason to repack; so is `live_files` (the +/// fragment's [`FragmentColumnLayoutStats::live_file_count`]) above +/// `max_files`. +/// +/// [`FragmentColumnLayoutStats::live_file_count`]: crate::dataset::compaction_stats::FragmentColumnLayoutStats::live_file_count +pub(super) fn plan_fragment_repack( + schema: &Schema, + fragment: &Fragment, + live_files: usize, + groups: Option<&[Vec]>, + max_files: Option, +) -> Option>> { + let over_file_limit = max_files.is_some_and(|max| live_files > max); + if groups.is_none() && !over_file_limit { + return None; + } + // A V1 file cannot tombstone single fields. + if fragment.files.iter().any(|file| { + !file + .file_version() + .is_ok_and(|version| version != ConcreteFileVersion::V1) + }) { + return None; + } + let schema_ids: HashSet = schema.fields_pre_order().map(|field| field.id).collect(); + let spilled = fragment.spilled_row_lineage_field_ids(); + let live: Vec<(&DataFile, HashSet)> = fragment + .files + .iter() + .filter(|file| file.fields.iter().any(|id| schema_ids.contains(id))) + .map(|file| (file, file_coverage(file, schema))) + .collect(); + let covered: HashSet = live.iter().flat_map(|(_, ids)| ids).copied().collect(); + let column_ids: Vec<(i32, HashSet)> = schema + .fields + .iter() + .map(|field| { + let mut ids = HashSet::new(); + subtree_ids(field, &mut ids); + (field.id, ids) + }) + .collect(); + let present: BTreeSet = column_ids + .iter() + .filter(|(_, ids)| !ids.is_disjoint(&covered)) + .map(|(column, _)| *column) + .collect(); + let movable: HashSet = schema + .fields + .iter() + .zip(&column_ids) + .filter(|(field, (_, ids))| !holds_blob(field) && ids.is_subset(&covered)) + .map(|(field, _)| field.id) + .collect(); + // The top-level columns each live file holds any field of. + let file_columns: Vec> = live + .iter() + .map(|(_, coverage)| { + column_ids + .iter() + .filter(|(_, ids)| !ids.is_disjoint(coverage)) + .map(|(column, _)| *column) + .collect() + }) + .collect(); + + let wanted: Vec> = match groups { + Some(groups) => groups + .iter() + .map(|group| group.iter().copied().collect()) + .collect(), + None if max_files == Some(1) => vec![present.clone()], + None => { + let kept = live + .iter() + .zip(&file_columns) + .max_by_key(|((file, _), columns)| { + ( + file.fields.iter().any(|id| spilled.contains(id)), + file.file_size_bytes.get().map(|size| size.get()), + columns.len(), + ) + }) + .map(|(_, columns)| columns.clone()) + .unwrap_or_default(); + let rest = present + .iter() + .copied() + .filter(|column| !kept.contains(column) && movable.contains(column)) + .collect(); + vec![kept, rest] + } + }; + + let out_of_place: Vec> = wanted + .into_iter() + .filter_map(|group| { + let group: BTreeSet = group.intersection(&present).copied().collect(); + if group.is_empty() || !group.iter().all(|column| movable.contains(column)) { + return None; + } + let holders: Vec<&BTreeSet> = file_columns + .iter() + .filter(|columns| !columns.is_disjoint(&group)) + .collect(); + let in_place = holders.len() == 1 && *holders[0] == group; + (!in_place).then(|| group.into_iter().collect()) + }) + .collect(); + let moved: HashSet = out_of_place.iter().flatten().copied().collect(); + let strands_lineage = live.iter().zip(&file_columns).any(|((file, _), columns)| { + file.fields.iter().any(|id| spilled.contains(id)) + && columns.iter().all(|column| moved.contains(column)) + }); + (!out_of_place.is_empty() && !strands_lineage).then_some(out_of_place) +} + +/// How many of `fragment`'s data files a commit of `written` would drop: the +/// files left holding no field of the schema and none of the fragment's +/// spilled row lineage, the rule the `DataReplacement` commit applies. +fn files_dropped(fragment: &Fragment, written: &[DataFile], schema: &Schema) -> usize { + let moved: HashSet = written + .iter() + .flat_map(|file| file_coverage(file, schema)) + .collect(); + let schema_ids: HashSet = schema.fields_pre_order().map(|field| field.id).collect(); + let spilled = fragment.spilled_row_lineage_field_ids(); + fragment + .files + .iter() + .filter(|file| { + !file + .fields + .iter() + .any(|id| (schema_ids.contains(id) && !moved.contains(id)) || spilled.contains(id)) + }) + .count() +} + +/// Write the new files of a [`CompactionTaskKind::RepackColumns`] task +/// against `dataset`, which must be at the plan's read version. +pub(super) async fn execute_repack( + dataset: &Dataset, + task: TaskData, + options: &CompactionOptions, +) -> Result { + let CompactionTaskKind::RepackColumns { files } = &task.kind else { + return Err(Error::internal("execute_repack called with a rewrite task")); + }; + let [fragment] = task.fragments.as_slice() else { + return Err(Error::invalid_input(format!( + "a RepackColumns task names exactly one fragment, got {}", + task.fragments.len() + ))); + }; + let write_version = options.write_version(dataset); + versions::validate_write_version( + dataset.manifest.data_storage_format.lance_file_format(), + write_version, + )?; + if write_version == ConcreteFileVersion::V1 { + return Err(Error::not_supported( + "repacking columns needs the V2 file format: a V1 file cannot tombstone \ + single fields", + )); + } + + // Read the base values only: the overlays stay on the fragment and keep + // shadowing them, so the new files hold exactly what the old ones did. + let mut base = fragment.clone(); + base.overlays.clear(); + let source = FileFragment::new(Arc::new(dataset.clone()), base); + let schema = dataset.schema(); + let batch_size = options.batch_size.map(|size| size as u32); + + let mut written: Vec = Vec::with_capacity(files.len()); + for column_ids in files { + let names: Vec<&str> = schema + .fields + .iter() + .filter(|field| column_ids.contains(&field.id)) + .map(|field| field.name.as_str()) + .collect(); + match write_one_file(&source, schema, &names, batch_size, write_version).await { + Ok(Some(file)) => written.push(file), + Ok(None) => {} + Err(err) => { + let partial = Fragment { + files: written, + ..Fragment::new(fragment.id) + }; + cleanup_data_fragments(&dataset.object_store, &dataset.base, None, &[partial]) + .await; + return Err(err); + } + } + } + + let metrics = CompactionMetrics { + files_added: written.len(), + files_removed: files_dropped(fragment, &written, schema), + ..Default::default() + }; + Ok(RewriteResult { + metrics, + new_fragments: Vec::new(), + read_version: dataset.manifest.version, + original_fragments: Vec::new(), + row_addrs: None, + repacked_files: Some(RepackedFiles { + fragment_id: fragment.id, + files: written, + }), + }) +} + +/// Write `columns` of `source` to one new data file. `None` when the +/// fragment has no rows to write. +async fn write_one_file( + source: &FileFragment, + schema: &Schema, + columns: &[&str], + batch_size: Option, + write_version: ConcreteFileVersion, +) -> Result> { + let write_schema = schema.project(columns)?; + let mut updater = source + .updater_with_version( + Some(columns), + Some((write_schema, schema.clone())), + batch_size, + None, + write_version, + ) + .await?; + let finished: Result = async { + while let Some(batch) = updater.next().await?.cloned() { + updater.update(batch).await?; + } + updater.finish().await + } + .await; + let fragment = match finished { + Ok(fragment) => fragment, + Err(err) => { + updater.cleanup_unfinished_writer().await; + return Err(err); + } + }; + // The updater appends the file it wrote to the fragment's files. + Ok((fragment.files.len() > source.metadata().files.len()) + .then(|| fragment.files.last().cloned()) + .flatten()) +} + +/// Commit the files of repack results as one `DataReplacement` that moves +/// values without changing them. Each result applies to its fragment as it +/// stands at commit, so a concurrent change to other columns of the fragment +/// is kept; the conflict check retries one that touches the moved columns. +pub(super) async fn commit_repacked_files( + dataset: &mut Dataset, + results: Vec, + options: &CompactionOptions, +) -> Result { + let mut metrics = CompactionMetrics::default(); + let mut read_version = u64::MAX; + let mut replacements = Vec::new(); + for result in results { + let Some(repacked) = result.repacked_files else { + continue; + }; + if repacked.files.is_empty() { + continue; + } + metrics += result.metrics; + read_version = read_version.min(result.read_version); + replacements.extend( + repacked + .files + .into_iter() + .map(|file| DataReplacementGroup(repacked.fragment_id, file)), + ); + } + if replacements.is_empty() { + return Ok(CompactionMetrics::default()); + } + let transaction = TransactionBuilder::new( + // The earliest version a repack read, so the conflict check covers + // every write since (the same reason vertical compaction uses it). + read_version, + Operation::DataReplacement { + replacements, + data_change: false, + }, + ) + .transaction_properties(options.transaction_properties.clone()) + .build(); + // Like a rewrite's, a repack result can be retried after an ambiguous + // success, so its files are left for cleanup rather than deleted here. + dataset + .apply_commit(transaction, &Default::default(), &Default::default()) + .await?; + Ok(metrics) +} diff --git a/rust/lance/src/dataset/optimize/tests/repack.rs b/rust/lance/src/dataset/optimize/tests/repack.rs new file mode 100644 index 00000000000..595e2b9ef4c --- /dev/null +++ b/rust/lance/src/dataset/optimize/tests/repack.rs @@ -0,0 +1,613 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Column repacking through `compact_files`. + +use super::*; +use crate::dataset::NewColumnTransform; +use arrow_array::StringArray; +use lance_index::IndexType; +use lance_index::scalar::ScalarIndexParams; + +fn abc_batch(start: i32, rows: i32) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, true), + Field::new("c", DataType::Utf8, true), + ])); + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from_iter_values(start..start + rows)), + Arc::new(Int32Array::from_iter_values( + (start..start + rows).map(|v| v * 10), + )), + Arc::new(StringArray::from_iter_values( + (start..start + rows).map(|v| format!("row-{v}")), + )), + ], + ) + .unwrap() +} + +/// Two fragments of four rows, each `a, b, c` in one file. +async fn write(uri: &str, version: LanceFileVersion) -> Dataset { + let data = abc_batch(0, 8); + let reader = RecordBatchIterator::new([Ok(data.clone())], data.schema()); + Dataset::write( + reader, + uri, + Some(WriteParams { + max_rows_per_file: 4, + data_storage_version: Some(version), + ..Default::default() + }), + ) + .await + .unwrap() +} + +/// `write`, then two backfills: every fragment holds its columns in three +/// files, `[a, b, c]`, `[d]` and `[e]`. +async fn write_backfilled() -> Dataset { + let mut dataset = write("memory://", LanceFileVersion::V2_0).await; + for (name, expr) in [("d", "a + 1"), ("e", "a + 2")] { + dataset + .add_columns( + NewColumnTransform::SqlExpressions(vec![(name.into(), expr.into())]), + None, + None, + ) + .await + .unwrap(); + } + dataset +} + +fn repack_options(max_files: Option, groups: Vec>) -> CompactionOptions { + CompactionOptions { + max_data_files_per_fragment: max_files, + column_groups: groups + .into_iter() + .map(|group| group.into_iter().map(String::from).collect()) + .collect(), + scope: CompactionScope::RepackColumns, + ..Default::default() + } +} + +fn layout(dataset: &Dataset) -> Vec>> { + dataset + .get_fragments() + .iter() + .map(|fragment| { + fragment + .metadata() + .files + .iter() + .map(|file| file.fields.to_vec()) + .collect() + }) + .collect() +} + +fn live_file_counts(dataset: &Dataset) -> Vec { + dataset + .column_layout_stats() + .iter() + .map(|stats| stats.live_file_count) + .collect() +} + +#[tokio::test] +async fn repack_collapses_backfilled_files() { + let mut dataset = write_backfilled().await; + let before = dataset.scan().try_into_batch().await.unwrap(); + let version = dataset.manifest.version; + assert_eq!(live_file_counts(&dataset), vec![3, 3]); + + let metrics = compact_files(&mut dataset, repack_options(Some(1), vec![]), None) + .await + .unwrap(); + + assert_eq!(dataset.manifest.version, version + 1, "one commit"); + assert_eq!(live_file_counts(&dataset), vec![1, 1]); + assert_eq!(layout(&dataset), vec![vec![vec![0, 1, 2, 3, 4]]; 2]); + assert_eq!(metrics.files_added, 2); + assert_eq!(metrics.files_removed, 6); + assert_eq!(metrics.fragments_added, 0); + assert_eq!(dataset.scan().try_into_batch().await.unwrap(), before); + dataset.validate().await.unwrap(); + + // Nothing is left to do, so nothing is committed. + compact_files(&mut dataset, repack_options(Some(1), vec![]), None) + .await + .unwrap(); + assert_eq!(dataset.manifest.version, version + 1); +} + +/// Above the limit, the columns of the biggest file stay where they are and +/// only the others are rewritten. +#[tokio::test] +async fn repack_keeps_the_biggest_file() { + let mut dataset = write_backfilled().await; + let before = dataset.scan().try_into_batch().await.unwrap(); + + compact_files(&mut dataset, repack_options(Some(2), vec![]), None) + .await + .unwrap(); + + for files in layout(&dataset) { + assert_eq!(files.len(), 2, "{files:?}"); + assert_eq!(files[0], vec![0, 1, 2], "the base file is untouched"); + assert_eq!(files[1], vec![3, 4]); + } + assert_eq!(dataset.scan().try_into_batch().await.unwrap(), before); + dataset.validate().await.unwrap(); +} + +#[tokio::test] +async fn repack_follows_column_groups() { + let mut dataset = write_backfilled().await; + let before = dataset.scan().try_into_batch().await.unwrap(); + + compact_files( + &mut dataset, + repack_options(None, vec![vec!["c", "e"]]), + None, + ) + .await + .unwrap(); + + for files in layout(&dataset) { + let mut files = files + .into_iter() + .map(|fields| { + fields + .into_iter() + .filter(|id| *id != TOMBSTONE_FIELD_ID) + .collect::>() + }) + .collect::>(); + files.sort(); + // The columns no group claims share one file. + assert_eq!(files, vec![vec![0, 1, 3], vec![2, 4]], "{files:?}"); + } + assert_eq!(dataset.scan().try_into_batch().await.unwrap(), before); + dataset.validate().await.unwrap(); + + // The layout matches the groups now. + let version = dataset.manifest.version; + compact_files( + &mut dataset, + repack_options(None, vec![vec!["c", "e"]]), + None, + ) + .await + .unwrap(); + assert_eq!(dataset.manifest.version, version); +} + +/// Deleted rows stay deleted and the new files line up with the old ones. +#[tokio::test] +async fn repack_keeps_deletions() { + let mut dataset = write_backfilled().await; + dataset.delete("a = 1 OR a = 6").await.unwrap(); + let before = dataset.scan().try_into_batch().await.unwrap(); + + let options = CompactionOptions { + // Neither small nor deleted enough to be rewritten. + target_rows_per_fragment: 4, + materialize_deletions_threshold: 0.5, + max_data_files_per_fragment: Some(1), + ..Default::default() + }; + compact_files(&mut dataset, options, None).await.unwrap(); + + assert_eq!(live_file_counts(&dataset), vec![1, 1]); + assert!( + dataset + .get_fragments() + .iter() + .all(|fragment| fragment.metadata().deletion_file.is_some()) + ); + assert_eq!(dataset.scan().try_into_batch().await.unwrap(), before); + dataset.validate().await.unwrap(); +} + +#[tokio::test] +async fn repack_keeps_scalar_index() { + let mut dataset = write_backfilled().await; + dataset + .create_index( + &["d"], + IndexType::BTree, + Some("d_idx".into()), + &ScalarIndexParams::default(), + false, + ) + .await + .unwrap(); + let before = dataset.load_indices().await.unwrap()[0].clone(); + + compact_files(&mut dataset, repack_options(Some(1), vec![]), None) + .await + .unwrap(); + + let after = dataset.load_indices().await.unwrap()[0].clone(); + assert_eq!(after.uuid, before.uuid, "the index is not rebuilt"); + assert_eq!(after.fragment_bitmap, before.fragment_bitmap); + let filtered = dataset + .scan() + .filter("d = 6") + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(filtered.num_rows(), 1); + dataset.validate().await.unwrap(); +} + +/// `drop_columns` leaves a dropped column's id in its file. A repack that +/// moves the other columns out drops that file, as `drop_columns` would. +#[tokio::test] +async fn repack_drops_file_left_with_only_dropped_columns() { + let mut dataset = write_backfilled().await; + dataset.drop_columns(&["c"]).await.unwrap(); + let before = dataset.scan().try_into_batch().await.unwrap(); + + compact_files(&mut dataset, repack_options(Some(1), vec![]), None) + .await + .unwrap(); + + assert_eq!(layout(&dataset), vec![vec![vec![0, 1, 3, 4]]; 2]); + assert_eq!(dataset.scan().try_into_batch().await.unwrap(), before); + dataset.validate().await.unwrap(); +} + +/// A run that rewrites some fragments and repacks others commits the +/// rewrite first and the repacks second, and never touches one fragment both +/// ways. +#[tokio::test] +async fn compaction_rewrites_and_repacks_in_one_run() { + let mut dataset = write_backfilled().await; + // Fragment 0 qualifies for a rewrite, fragment 1 only for a repack. + dataset.delete("a = 0 OR a = 1").await.unwrap(); + let before = dataset.scan().try_into_batch().await.unwrap(); + let version = dataset.manifest.version; + + let options = CompactionOptions { + target_rows_per_fragment: 3, + max_data_files_per_fragment: Some(1), + ..Default::default() + }; + let plan = plan_compaction(&dataset, &options).await.unwrap(); + assert_eq!(plan.tasks.len(), 2, "{plan:?}"); + assert_eq!(plan.tasks[0].kind, CompactionTaskKind::RewriteFragments); + assert_eq!(plan.tasks[0].fragments[0].id, 0); + assert!(matches!( + plan.tasks[1].kind, + CompactionTaskKind::RepackColumns { .. } + )); + assert_eq!(plan.tasks[1].fragments[0].id, 1); + + let metrics = compact_files(&mut dataset, options, None).await.unwrap(); + let last = dataset.manifest.version; + assert!(last >= version + 2); + for (version, operation) in [(last - 1, "Rewrite"), (last, "DataReplacement")] { + let transaction = dataset + .checkout_version(version) + .await + .unwrap() + .read_transaction() + .await + .unwrap() + .unwrap(); + assert_eq!(transaction.operation.to_string(), operation); + } + assert_eq!(metrics.fragments_removed, 1); + assert_eq!(metrics.fragments_added, 1); + assert_eq!(live_file_counts(&dataset), vec![1, 1]); + // The rewritten fragment gets a new id, so it now scans last; the repack + // changes no value. + let rewritten = dataset.checkout_version(last - 1).await.unwrap(); + let after_rewrite = rewritten.scan().try_into_batch().await.unwrap(); + assert_eq!(after_rewrite.num_rows(), before.num_rows()); + assert_eq!( + dataset.scan().try_into_batch().await.unwrap(), + after_rewrite + ); + dataset.validate().await.unwrap(); +} + +#[tokio::test] +async fn compaction_scope_selects_task_kinds() { + let mut dataset = write_backfilled().await; + dataset.delete("a = 0 OR a = 1").await.unwrap(); + let options = |scope| CompactionOptions { + target_rows_per_fragment: 2, + max_data_files_per_fragment: Some(1), + scope, + ..Default::default() + }; + + let rewrites = plan_compaction(&dataset, &options(CompactionScope::RewriteFragments)) + .await + .unwrap(); + assert_eq!(rewrites.tasks.len(), 1); + assert_eq!(rewrites.tasks[0].kind, CompactionTaskKind::RewriteFragments); + + let repacks = plan_compaction(&dataset, &options(CompactionScope::RepackColumns)) + .await + .unwrap(); + assert_eq!(repacks.tasks.len(), 2, "both fragments are repacked"); + assert!( + repacks + .tasks + .iter() + .all(|task| matches!(task.kind, CompactionTaskKind::RepackColumns { .. })) + ); +} + +/// Without `max_data_files_per_fragment` or `column_groups`, nothing is +/// repacked. +#[tokio::test] +async fn compaction_plans_no_repack_by_default() { + let dataset = write_backfilled().await; + let plan = plan_compaction(&dataset, &repack_options(None, vec![])) + .await + .unwrap(); + assert!(plan.tasks.is_empty(), "{plan:?}"); +} + +/// Repack tasks run on another machine: they are serialized, executed on +/// their own, and their results committed together. +#[tokio::test] +async fn repack_tasks_run_distributed() { + let mut dataset = write_backfilled().await; + let before = dataset.scan().try_into_batch().await.unwrap(); + let plan = plan_compaction(&dataset, &repack_options(Some(1), vec![])) + .await + .unwrap(); + assert_eq!(plan.tasks.len(), 2); + + let mut results = Vec::new(); + for task in plan.compaction_tasks() { + let task: CompactionTask = + serde_json::from_str(&serde_json::to_string(&task).unwrap()).unwrap(); + let result = task.execute(&dataset).await.unwrap(); + let result: RewriteResult = + serde_json::from_str(&serde_json::to_string(&result).unwrap()).unwrap(); + results.push(result); + } + commit_compaction( + &mut dataset, + results, + Arc::new(DatasetIndexRemapperOptions::default()), + plan.options(), + ) + .await + .unwrap(); + + assert_eq!(live_file_counts(&dataset), vec![1, 1]); + assert_eq!(dataset.scan().try_into_batch().await.unwrap(), before); + dataset.validate().await.unwrap(); +} + +/// A task serialized before `kind` existed rewrites its fragments. +#[test] +fn task_data_without_kind_rewrites_fragments() { + let task: TaskData = serde_json::from_str(r#"{"fragments": []}"#).unwrap(); + assert_eq!(task.kind, CompactionTaskKind::RewriteFragments); +} + +/// Run a repack of every fragment to the point of commit. +async fn stale_repack_results( + dataset: &Dataset, + options: &CompactionOptions, +) -> Vec { + let plan = plan_compaction(dataset, options).await.unwrap(); + let mut results = Vec::new(); + for task in plan.compaction_tasks() { + results.push(task.execute(dataset).await.unwrap()); + } + results +} + +async fn commit_results( + dataset: &mut Dataset, + results: Vec, +) -> Result { + commit_compaction( + dataset, + results, + Arc::new(DatasetIndexRemapperOptions::default()), + &CompactionOptions::default(), + ) + .await +} + +/// A delete that removes rows of a repacked fragment leaves the new files +/// aligned (a deletion only marks rows), so the repack still commits; one +/// that removes the whole fragment makes it retry. +#[rstest] +#[case::one_row("a = 1", true)] +#[case::whole_fragment("a < 4", false)] +#[tokio::test] +async fn repack_against_concurrent_delete(#[case] predicate: &str, #[case] commits: bool) { + let mut dataset = write_backfilled().await; + let stale = stale_repack_results(&dataset, &repack_options(Some(1), vec![])).await; + + dataset.delete(predicate).await.unwrap(); + let expected = dataset.scan().try_into_batch().await.unwrap(); + let version = dataset.manifest.version; + + let result = commit_results(&mut dataset, stale).await; + if commits { + result.unwrap(); + assert_eq!(live_file_counts(&dataset), vec![1, 1]); + } else { + let err = result.unwrap_err(); + assert!( + matches!(err, Error::RetryableCommitConflict { .. }), + "{err}" + ); + dataset.checkout_latest().await.unwrap(); + assert_eq!(dataset.manifest.version, version); + } + assert_eq!(dataset.scan().try_into_batch().await.unwrap(), expected); + dataset.validate().await.unwrap(); +} + +/// A concurrent `drop_columns` of a column the repack moves makes it retry. +/// One that drops a column the repack leaves alone does not: the new file is +/// applied to the fragment as the drop left it. +#[rstest] +#[case::moved_column("d", false)] +#[case::other_column("b", true)] +#[tokio::test] +async fn repack_against_concurrent_drop(#[case] dropped: &str, #[case] commits: bool) { + let mut dataset = write_backfilled().await; + // Moves d and e into one file and leaves the base file alone. + let stale = stale_repack_results(&dataset, &repack_options(Some(2), vec![])).await; + + dataset.drop_columns(&[dropped]).await.unwrap(); + let expected = dataset.scan().try_into_batch().await.unwrap(); + let version = dataset.manifest.version; + + let result = commit_results(&mut dataset, stale).await; + if commits { + result.unwrap(); + assert_eq!(live_file_counts(&dataset), vec![2, 2]); + } else { + let err = result.unwrap_err(); + assert!( + matches!(err, Error::RetryableCommitConflict { .. }), + "{err}" + ); + dataset.checkout_latest().await.unwrap(); + assert_eq!(dataset.manifest.version, version); + } + assert_eq!(dataset.scan().try_into_batch().await.unwrap(), expected); + dataset.validate().await.unwrap(); +} + +/// A concurrent update that writes new values into a moved column makes the +/// repack retry rather than publish the old values over the new ones. +#[tokio::test] +async fn repack_against_concurrent_update_of_moved_column() { + let mut dataset = write_backfilled().await; + let stale = stale_repack_results(&dataset, &repack_options(Some(1), vec![])).await; + + crate::dataset::UpdateBuilder::new(Arc::new(dataset.clone())) + .update_where("a = 2") + .unwrap() + .set("d", "100") + .unwrap() + .build() + .unwrap() + .execute() + .await + .unwrap(); + dataset.checkout_latest().await.unwrap(); + let expected = dataset.scan().try_into_batch().await.unwrap(); + + let err = commit_results(&mut dataset, stale).await.unwrap_err(); + assert!( + matches!(err, Error::RetryableCommitConflict { .. }), + "{err}" + ); + dataset.checkout_latest().await.unwrap(); + assert_eq!(dataset.scan().try_into_batch().await.unwrap(), expected); + dataset.validate().await.unwrap(); +} + +#[tokio::test] +async fn repack_skips_legacy_files() { + let dir = TempStrDir::default(); + let dataset = write(&dir, LanceFileVersion::Legacy).await; + let plan = plan_compaction(&dataset, &repack_options(None, vec![vec!["c"]])) + .await + .unwrap(); + assert!(plan.tasks.is_empty(), "{plan:?}"); +} + +#[test] +fn max_data_files_per_fragment_must_be_positive() { + let mut options = CompactionOptions { + max_data_files_per_fragment: Some(0), + ..Default::default() + }; + assert!(options.validate().is_err()); +} + +#[test] +fn repack_options_parse_from_config() { + let config = HashMap::from([ + ( + "lance.compaction.max_data_files_per_fragment".to_string(), + "4".to_string(), + ), + ( + "lance.compaction.scope".to_string(), + "repack_columns".to_string(), + ), + ]); + let options = CompactionOptions::from_dataset_config(&config).unwrap(); + assert_eq!(options.max_data_files_per_fragment, Some(4)); + assert_eq!(options.scope, CompactionScope::RepackColumns); +} + +/// The file holding a fragment's spilled row lineage keeps its columns: the +/// new files carry no lineage, so emptying that file would strand it. +#[test] +fn repack_keeps_the_spilled_lineage_file() { + use lance_table::format::{ROW_ID_FIELD_ID, RowIdMeta}; + let schema = lance_core::datatypes::Schema::try_from(&Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + Field::new("c", DataType::Int32, true), + ])) + .unwrap(); + let mut fragment = Fragment::new(0); + fragment.add_file( + "lineage.lance", + vec![0, ROW_ID_FIELD_ID], + vec![0, 1], + lance_file::version::ConcreteFileVersion::V2_0, + None, + ); + for (path, field) in [("b.lance", 1), ("c.lance", 2)] { + fragment.add_file( + path, + vec![field], + vec![0], + lance_file::version::ConcreteFileVersion::V2_0, + None, + ); + } + fragment.row_id_meta = Some(RowIdMeta::Column); + + // Above the limit, the lineage file's column stays and the others move. + assert_eq!( + crate::dataset::optimize::repack::plan_fragment_repack( + &schema, + &fragment, + 3, + None, + Some(2) + ), + Some(vec![vec![1, 2]]) + ); + // Moving every column would leave the lineage file behind, so no repack. + assert_eq!( + crate::dataset::optimize::repack::plan_fragment_repack( + &schema, + &fragment, + 3, + None, + Some(1) + ), + None + ); +} diff --git a/rust/lance/src/dataset/rewrite_columns.rs b/rust/lance/src/dataset/rewrite_columns.rs deleted file mode 100644 index b258fd426c4..00000000000 --- a/rust/lance/src/dataset/rewrite_columns.rs +++ /dev/null @@ -1,814 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright The Lance Authors - -//! Horizontal compaction: rewrite a few columns of every fragment into their -//! own data files, leaving the files that hold the other columns untouched. -//! -//! This is the horizontal pipeline of the pluggable compaction framework -//! ([`CompactionExecutor`] / [`CompactionCommitter`] in [`super::optimize`]). -//! Unlike vertical compaction (which merges whole fragments and re-encodes -//! every column), it re-reads only the named columns, writes them to one new -//! file per fragment, and tombstones them where they used to live. Rows, -//! fragment ids, row addresses and indices are unchanged, so it commits as an -//! `Operation::Update` in [`UpdateMode::RewriteColumns`] with an empty -//! `fields_modified` (values did not change: indices keep their coverage and -//! overlays keep shadowing). - -use std::borrow::Cow; -use std::collections::HashSet; -use std::sync::Arc; - -use lance_core::datatypes::Schema; -use lance_file::version::{ConcreteFileVersion, LanceFileVersion}; -use lance_table::format::{Fragment, overlay::TOMBSTONE_FIELD_ID}; - -use super::fragment::FileFragment; -use super::optimize::{ - CompactionCommitter, CompactionExecutor, CompactionMetrics, CompactionOptions, - run_compaction_pipeline, -}; -use super::transaction::{Operation, Transaction, UpdateMode}; -use super::{Dataset, versions}; -use crate::dataset::optimize::remapping::IndexRemapperOptions; -use crate::{Error, Result}; - -/// The exact V2 version the new files get: `requested`, or the dataset's -/// default write version, which is never changed. -fn write_version( - dataset: &Dataset, - requested: Option, -) -> Result { - let default_version = dataset.manifest.data_storage_format.lance_file_format(); - let version = requested - .map(LanceFileVersion::resolve) - .unwrap_or(default_version); - versions::validate_write_version(default_version, version)?; - if version == ConcreteFileVersion::V1 { - return Err(Error::not_supported( - "rewrite_columns requires the V2 file format: a V1 file cannot tombstone \ - single fields, so compact the dataset to V2 first", - )); - } - Ok(version) -} -/// The schema of the new data file: `columns` in dataset order, all top-level. -fn rewrite_schema(dataset: &Dataset, columns: &[&str]) -> Result { - if columns.is_empty() { - return Err(Error::invalid_input( - "rewrite_columns needs at least one column", - )); - } - let schema = dataset.schema(); - if let Some(column) = columns - .iter() - .find(|column| !schema.fields.iter().any(|field| &field.name == *column)) - { - return Err(Error::invalid_input(format!( - "Column \"{column}\" is not a top-level column of the dataset" - ))); - } - let ordered = schema - .fields - .iter() - .map(|field| field.name.as_str()) - .filter(|name| columns.contains(name)) - .collect::>(); - let schema = schema.project(&ordered)?; - if let Some(blob) = schema.fields_pre_order().find(|field| field.is_blob()) { - return Err(Error::not_supported(format!( - "rewrite_columns cannot rewrite blob column \"{}\"", - blob.name - ))); - } - Ok(schema) -} -impl FileFragment { - /// Rewrite `columns` of this fragment into one new data file. - /// - /// The returned metadata has the new file appended and the rewritten fields - /// tombstoned in the files they came from; a file left holding no column - /// of the schema is dropped, unless it carries the fragment's spilled row - /// lineage. Nothing is committed: commit the result with - /// [`RewriteColumnsCommitter`]. Returns `None` when the columns already sit - /// alone in a file of the requested version, so a retried or resumed - /// rewrite skips finished work. - pub async fn rewrite_columns( - &self, - columns: &[&str], - data_storage_version: Option, - ) -> Result> { - let dataset = self.dataset(); - let write_schema = rewrite_schema(dataset, columns)?; - let resolved_version = write_version(dataset, data_storage_version)?; - let field_ids: HashSet = write_schema.field_ids().into_iter().collect(); - // Ids of columns `drop_columns` removed stay in their files, and a - // spilled row lineage sequence sits in a data file under a reserved id, - // so a file's ids are only compared and kept by what the schema reads. - let schema_ids: HashSet = dataset - .schema() - .fields_pre_order() - .map(|field| field.id) - .collect(); - - let already_split = self.metadata().files.iter().any(|file| { - file.fields - .iter() - .copied() - .filter(|id| schema_ids.contains(id)) - .collect::>() - == field_ids - && file - .file_version() - .is_ok_and(|version| version == resolved_version) - }); - if already_split { - return Ok(None); - } - - let mut updater = self - .updater_with_version( - Some(columns), - Some((write_schema, dataset.schema().clone())), - None, - None, - resolved_version, - ) - .await?; - let written: Result = async { - while let Some(batch) = updater.next().await?.cloned() { - updater.update(batch).await?; - } - updater.finish().await - } - .await; - let mut fragment = match written { - Ok(fragment) => fragment, - Err(err) => { - updater.cleanup_unfinished_writer().await; - return Err(err); - } - }; - // No live rows, so no file was written and nothing to tombstone. - if fragment.files.len() == self.metadata().files.len() { - return Ok(None); - } - - let new_file = fragment.files.len() - 1; - for file in &mut fragment.files[..new_file] { - file.fields = file - .fields - .iter() - .map(|id| { - if field_ids.contains(id) { - TOMBSTONE_FIELD_ID - } else { - *id - } - }) - .collect::>() - .into(); - } - // Drop a file left with nothing the schema reads, the rule - // `drop_columns` applies, but keep one carrying the fragment's spilled - // row lineage: it is the only copy. - let spilled = fragment.spilled_row_lineage_field_ids(); - fragment.files.retain(|file| { - file.fields - .iter() - .any(|id| schema_ids.contains(id) || spilled.contains(id)) - }); - // The new file is the last one and is always kept. - let files_removed = self.metadata().files.len() + 1 - fragment.files.len(); - Ok(Some(RewriteColumnsResult { - fragment, - read_version: dataset.manifest.version, - files_removed, - })) - } -} -/// One fragment's horizontal-rewrite output, bound to the manifest version it -/// was read at. The committer commits against that version so the conflict -/// check catches any change made to the fragment between the read and the -/// commit (mirrors vertical compaction's `RewriteResult::read_version`). -#[derive(Debug, Clone)] -pub struct RewriteColumnsResult { - /// The rewritten fragment metadata to publish. - pub fragment: Fragment, - /// The manifest version the fragment was read at. - pub read_version: u64, - /// How many of the fragment's data files the rewrite dropped. - pub files_removed: usize, -} - -/// Horizontal executor: rewrites one fragment's named columns into a new data -/// file (see [`FileFragment::rewrite_columns`]). The columns and target version -/// are the executor's configuration; the task is the fragment to rewrite. -#[derive(Debug, Clone)] -pub struct RewriteColumnsExecutor { - columns: Vec, - data_storage_version: Option, -} - -impl RewriteColumnsExecutor { - /// Build an executor that repacks `columns` into one new file per fragment, - /// optionally in a different V2 version. - pub fn new(columns: &[&str], data_storage_version: Option) -> Self { - Self { - columns: columns.iter().map(|c| c.to_string()).collect(), - data_storage_version, - } - } -} - -#[async_trait::async_trait] -impl CompactionExecutor for RewriteColumnsExecutor { - type Task = FileFragment; - /// `None` when the fragment already had the layout (nothing to commit). - type TaskResult = Option; - - /// The task is a [`FileFragment`] already bound to its read version, so the - /// `dataset` argument the pipeline threads in is unused; the result carries - /// the fragment's read version for the committer. - async fn execute( - &self, - _dataset: Cow<'_, Dataset>, - task: FileFragment, - _options: &CompactionOptions, - ) -> Result> { - let columns: Vec<&str> = self.columns.iter().map(String::as_str).collect(); - task.rewrite_columns(&columns, self.data_storage_version) - .await - } -} -/// Commits horizontal-compaction results as one `Operation::Update` in -/// [`UpdateMode::RewriteColumns`] with an empty `fields_modified`: values did -/// not change, so index coverage is kept and overlays keep shadowing. -#[derive(Debug, Default, Clone, Copy)] -pub struct RewriteColumnsCommitter; - -#[async_trait::async_trait] -impl CompactionCommitter for RewriteColumnsCommitter { - type TaskResult = Option; - - async fn commit( - &self, - dataset: &mut Dataset, - results: Vec>, - _remap_options: Arc, - _options: &CompactionOptions, - ) -> Result { - let results: Vec = results.into_iter().flatten().collect(); - if results.is_empty() { - return Ok(CompactionMetrics::default()); - } - // Commit against the earliest version any rewritten fragment was read - // at, not the current one, so the conflict check catches a fragment - // changed between the read and this commit (as vertical compaction does). - let read_version = results - .iter() - .map(|result| result.read_version) - .min() - .expect("results is non-empty"); - // Row addresses and values are unchanged, so no fragment is added or - // removed; only per-fragment files change: one new file per rewritten - // fragment, and the files each rewrite dropped. A fragment changed or - // removed since its read is left to the commit's conflict check. - let metrics = CompactionMetrics { - files_added: results.len(), - files_removed: results.iter().map(|result| result.files_removed).sum(), - ..CompactionMetrics::default() - }; - let updated_fragments: Vec = - results.into_iter().map(|result| result.fragment).collect(); - let transaction = Transaction::new( - read_version, - Operation::Update { - removed_fragment_ids: Vec::new(), - updated_fragments, - new_fragments: Vec::new(), - fields_modified: Vec::new(), - compacted_sstables: Vec::new(), - fields_for_preserving_frag_bitmap: Vec::new(), - update_mode: Some(UpdateMode::RewriteColumns), - inserted_rows_filter: None, - updated_fragment_offsets: None, - }, - None, - ); - dataset - .apply_commit(transaction, &Default::default(), &Default::default()) - .await?; - Ok(metrics) - } -} -impl Dataset { - /// Rewrite `columns` of every fragment into one new data file per fragment, - /// optionally in a different V2 version, without touching the files that - /// hold the other columns. Fragments already in the requested layout are - /// skipped, so an interrupted rewrite can be rerun. - /// - /// Runs through the horizontal compaction pipeline - /// ([`RewriteColumnsExecutor`] + [`RewriteColumnsCommitter`]) and commits - /// once. To spread the work over many machines, call - /// [`FileFragment::rewrite_columns`] per fragment and commit the returned - /// results with [`RewriteColumnsCommitter`]. - /// - /// The new files carry only the rewritten columns. On a table with stable - /// row ids whose row lineage was spilled into a data file, that file is - /// kept for the lineage even after every column has moved out of it, until - /// a vertical compaction rewrites the fragment. - pub async fn rewrite_columns( - &mut self, - columns: &[&str], - data_storage_version: Option, - ) -> Result { - let executor = RewriteColumnsExecutor::new(columns, data_storage_version); - let options = CompactionOptions::default(); - let tasks = self.get_fragments(); - let concurrency = self.object_store.io_parallelism(); - // Like vertical compaction, a failed run leaves its new files - // uncommitted and unreferenced; dataset cleanup reclaims them. - run_compaction_pipeline( - self, - executor, - RewriteColumnsCommitter, - tasks, - Arc::new(crate::dataset::index::DatasetIndexRemapperOptions::default()), - &options, - concurrency, - ) - .await - } -} -#[cfg(test)] -mod tests { - use super::*; - use crate::dataset::WriteParams; - use crate::index::DatasetIndexExt; - use arrow_array::{Int32Array, RecordBatch, RecordBatchIterator, StringArray}; - use arrow_schema::{DataType, Field, Schema as ArrowSchema}; - use lance_core::utils::tempfile::TempStrDir; - use lance_index::IndexType; - use lance_index::scalar::ScalarIndexParams; - use std::sync::Arc; - - fn batch(start: i32, rows: i32) -> RecordBatch { - let schema = Arc::new(ArrowSchema::new(vec![ - Field::new("a", DataType::Int32, false), - Field::new("b", DataType::Int32, true), - Field::new("c", DataType::Utf8, true), - ])); - RecordBatch::try_new( - schema, - vec![ - Arc::new(Int32Array::from_iter_values(start..start + rows)), - Arc::new(Int32Array::from_iter_values( - (start..start + rows).map(|v| v * 10), - )), - Arc::new(StringArray::from_iter_values( - (start..start + rows).map(|v| format!("row-{v}")), - )), - ], - ) - .unwrap() - } - - async fn write(uri: &str, version: LanceFileVersion) -> Dataset { - let data = batch(0, 8); - let reader = RecordBatchIterator::new([Ok(data.clone())], data.schema()); - Dataset::write( - reader, - uri, - Some(WriteParams { - max_rows_per_file: 4, - data_storage_version: Some(version), - ..Default::default() - }), - ) - .await - .unwrap() - } - - fn file_versions(fragment: &Fragment) -> Vec<(Vec, ConcreteFileVersion)> { - fragment - .files - .iter() - .map(|file| (file.fields.to_vec(), file.file_version().unwrap())) - .collect() - } - #[tokio::test] - async fn rewrite_migrates_only_the_named_column() { - let mut dataset = write("memory://", LanceFileVersion::V2_0).await; - dataset.delete("a = 3").await.unwrap(); - let before = dataset.scan().try_into_batch().await.unwrap(); - let version = dataset.manifest.version; - - let metrics = dataset - .rewrite_columns(&["c"], Some(LanceFileVersion::V2_2)) - .await - .unwrap(); - assert_eq!(metrics.files_added, 2, "one new file per fragment"); - - assert_eq!(dataset.manifest.version, version + 1); - assert_eq!(dataset.manifest.fragments.len(), 2); - for fragment in dataset.manifest.fragments.iter() { - assert_eq!( - file_versions(fragment), - vec![ - (vec![0, 1, TOMBSTONE_FIELD_ID], ConcreteFileVersion::V2_0), - (vec![2], ConcreteFileVersion::V2_2), - ] - ); - } - // Values are byte-for-byte unchanged. - assert_eq!(dataset.scan().try_into_batch().await.unwrap(), before); - // The dataset default write version is not changed. - assert_eq!( - dataset.manifest.data_storage_format.lance_file_format(), - ConcreteFileVersion::V2_0 - ); - dataset.validate().await.unwrap(); - - // Idempotent: already in the requested layout, nothing committed. - dataset - .rewrite_columns(&["c"], Some(LanceFileVersion::V2_2)) - .await - .unwrap(); - assert_eq!(dataset.manifest.version, version + 1); - } - #[tokio::test] - async fn rewrite_keeps_scalar_index() { - let mut dataset = write("memory://", LanceFileVersion::V2_0).await; - dataset - .create_index( - &["a"], - IndexType::BTree, - Some("a_idx".into()), - &ScalarIndexParams::default(), - false, - ) - .await - .unwrap(); - let before = dataset.load_indices().await.unwrap()[0].clone(); - - dataset.rewrite_columns(&["a"], None).await.unwrap(); - - let after = dataset.load_indices().await.unwrap()[0].clone(); - assert_eq!(after.uuid, before.uuid); - assert_eq!(after.fragment_bitmap, before.fragment_bitmap); - let filtered = dataset - .scan() - .filter("a = 5") - .unwrap() - .try_into_batch() - .await - .unwrap(); - assert_eq!(filtered.num_rows(), 1); - dataset.validate().await.unwrap(); - } - - #[tokio::test] - async fn rewrite_rejects_bad_input() { - let mut dataset = write("memory://", LanceFileVersion::V2_0).await; - let err = dataset - .rewrite_columns(&["missing"], None) - .await - .unwrap_err(); - assert!(matches!(err, Error::InvalidInput { .. }), "{err}"); - - let legacy_dir = TempStrDir::default(); - let mut legacy = write(&legacy_dir, LanceFileVersion::Legacy).await; - let err = legacy.rewrite_columns(&["c"], None).await.unwrap_err(); - assert!(matches!(err, Error::NotSupported { .. }), "{err}"); - } - - #[tokio::test] - async fn rewrite_collapses_backfilled_files() { - use crate::dataset::NewColumnTransform; - let mut dataset = write("memory://", LanceFileVersion::V2_0).await; - // Two backfills, each appending one data file per fragment: the feature's - // target — a fragment split across many per-column files. - dataset - .add_columns( - NewColumnTransform::SqlExpressions(vec![("d".into(), "a + 1".into())]), - None, - None, - ) - .await - .unwrap(); - dataset - .add_columns( - NewColumnTransform::SqlExpressions(vec![("e".into(), "a + 2".into())]), - None, - None, - ) - .await - .unwrap(); - let before = dataset.scan().try_into_batch().await.unwrap(); - assert!( - dataset - .column_layout_stats() - .iter() - .all(|s| s.live_file_count == 3), - "each fragment carries three data files after two backfills: {:?}", - dataset.column_layout_stats() - ); - - // Repack every column into a single file per fragment. - dataset - .rewrite_columns(&["a", "b", "c", "d", "e"], None) - .await - .unwrap(); - - assert!( - dataset - .column_layout_stats() - .iter() - .all(|s| s.live_file_count == 1), - "rewrite collapsed each fragment to one data file: {:?}", - dataset.column_layout_stats() - ); - assert_eq!(dataset.scan().try_into_batch().await.unwrap(), before); - dataset.validate().await.unwrap(); - } - - #[tokio::test] - async fn rewrite_after_fragment_fully_deleted() { - let mut dataset = write("memory://", LanceFileVersion::V2_0).await; - // Deleting every row of a fragment removes it from the manifest, so the - // rewrite runs over the surviving fragment(s); it must still succeed and - // leave the live rows byte-for-byte unchanged. - dataset.delete("a < 4").await.unwrap(); - let before = dataset.scan().try_into_batch().await.unwrap(); - - dataset.rewrite_columns(&["c"], None).await.unwrap(); - - assert_eq!(dataset.scan().try_into_batch().await.unwrap(), before); - dataset.validate().await.unwrap(); - } - - /// A delete that lands on the fragment between the rewrite's read and its - /// commit, whether it removes one row or the whole fragment, must surface - /// as a retryable conflict. - #[rstest::rstest] - #[case::one_row("a = 1", 7)] - #[case::whole_fragment("a < 4", 4)] - #[tokio::test] - async fn rewrite_columns_conflicts_with_intervening_delete( - #[case] predicate: &str, - #[case] rows_left: usize, - ) { - use crate::dataset::index::DatasetIndexRemapperOptions; - use crate::dataset::optimize::{ - CompactionCommitter, CompactionExecutor, CompactionOptions, - }; - use std::borrow::Cow; - - let mut dataset = write("memory://", LanceFileVersion::V2_0).await; - - // Rewrite column "c" of fragment 0, reading at the current version. - let executor = RewriteColumnsExecutor::new(&["c"], None); - let fragment = dataset.get_fragment(0).unwrap(); - let stale = executor - .execute( - Cow::Borrowed(&dataset), - fragment, - &CompactionOptions::default(), - ) - .await - .unwrap() - .expect("fragment 0 has columns to repack"); - - // A delete lands on that same fragment before the rewrite is committed. - dataset.delete(predicate).await.unwrap(); - let after_delete = dataset.scan().try_into_batch().await.unwrap(); - assert_eq!(after_delete.num_rows(), rows_left); - - // The rewrite was read before the delete, so committing it must conflict - // rather than publish the stale layout and resurrect the deleted row. - let err = RewriteColumnsCommitter - .commit( - &mut dataset, - vec![Some(stale)], - Arc::new(DatasetIndexRemapperOptions::default()), - &CompactionOptions::default(), - ) - .await - .unwrap_err(); - assert!( - matches!(err, Error::RetryableCommitConflict { .. }), - "a rewrite read before the delete must not commit over it: {err}" - ); - - // The delete stands. - assert_eq!(dataset.scan().try_into_batch().await.unwrap(), after_delete); - assert_eq!( - dataset - .scan() - .filter(predicate) - .unwrap() - .try_into_batch() - .await - .unwrap() - .num_rows(), - 0 - ); - dataset.validate().await.unwrap(); - } - - /// `drop_columns` leaves a dropped column's id in the file that held it. - /// Rewriting the rest of that file's columns leaves it with nothing the - /// schema reads, so the rewrite must drop it, as `drop_columns` would. - #[tokio::test] - async fn rewrite_drops_file_left_with_only_dropped_columns() { - use crate::dataset::NewColumnTransform; - let mut dataset = write("memory://", LanceFileVersion::V2_0).await; - dataset - .add_columns( - NewColumnTransform::SqlExpressions(vec![("d".into(), "a + 1".into())]), - None, - None, - ) - .await - .unwrap(); - dataset.drop_columns(&["c"]).await.unwrap(); - let before = dataset.scan().try_into_batch().await.unwrap(); - - dataset - .rewrite_columns(&["a", "b", "d"], None) - .await - .unwrap(); - - for fragment in dataset.get_fragments() { - assert_eq!( - fragment.metadata().files.len(), - 1, - "only the new file is left: {:?}", - fragment.metadata().files - ); - } - assert_eq!(dataset.scan().try_into_batch().await.unwrap(), before); - dataset.validate().await.unwrap(); - } - - /// A file whose only other column was dropped already holds exactly the - /// requested columns, so there is nothing to rewrite. - #[tokio::test] - async fn rewrite_skips_file_whose_other_columns_were_dropped() { - let mut dataset = write("memory://", LanceFileVersion::V2_0).await; - dataset.drop_columns(&["c"]).await.unwrap(); - let version = dataset.manifest.version; - - dataset.rewrite_columns(&["a", "b"], None).await.unwrap(); - - assert_eq!(dataset.manifest.version, version, "nothing was committed"); - dataset.validate().await.unwrap(); - } - - /// A `drop_columns` that commits between a rewrite's read and its commit - /// would leave a file holding only the dropped column if the rewrite went - /// through, whether that is the rewritten column's new file or a file the - /// drop removed, so the commit must conflict. When the rewritten column - /// survived the drop, rerunning the rewrite on the new version succeeds. - #[rstest::rstest] - #[case::rewritten_column("c")] - #[case::other_column("d")] - #[tokio::test] - async fn rewrite_columns_conflicts_with_concurrent_drop(#[case] dropped: &str) { - use crate::dataset::NewColumnTransform; - use crate::dataset::index::DatasetIndexRemapperOptions; - use crate::dataset::optimize::{ - CompactionCommitter, CompactionExecutor, CompactionOptions, - }; - use std::borrow::Cow; - - let mut dataset = write("memory://", LanceFileVersion::V2_0).await; - // `d` gets a file of its own in every fragment. - dataset - .add_columns( - NewColumnTransform::SqlExpressions(vec![("d".into(), "a + 1".into())]), - None, - None, - ) - .await - .unwrap(); - - let executor = RewriteColumnsExecutor::new(&["c"], None); - let mut stale = Vec::new(); - for fragment in dataset.get_fragments() { - stale.push( - executor - .execute( - Cow::Borrowed(&dataset), - fragment, - &CompactionOptions::default(), - ) - .await - .unwrap(), - ); - } - - dataset.drop_columns(&[dropped]).await.unwrap(); - let version = dataset.manifest.version; - let expected = dataset.scan().try_into_batch().await.unwrap(); - - let err = RewriteColumnsCommitter - .commit( - &mut dataset, - stale, - Arc::new(DatasetIndexRemapperOptions::default()), - &CompactionOptions::default(), - ) - .await - .unwrap_err(); - assert!( - matches!(err, Error::RetryableCommitConflict { .. }), - "{err}" - ); - dataset.checkout_latest().await.unwrap(); - assert_eq!(dataset.manifest.version, version, "nothing was committed"); - dataset.validate().await.unwrap(); - - if dropped != "c" { - dataset.rewrite_columns(&["c"], None).await.unwrap(); - assert_eq!(dataset.scan().try_into_batch().await.unwrap(), expected); - dataset.validate().await.unwrap(); - } - } - - /// The concurrent-drop check must not trip on a file holding only - /// tombstones, which an in-place merge_insert leaves in its post-image for - /// the commit to drop: a rename racing such a commit has to go through. - #[tokio::test] - async fn rewrite_columns_commit_ignores_tombstoned_file_on_rename() { - use crate::dataset::index::DatasetIndexRemapperOptions; - use crate::dataset::optimize::{ - CompactionCommitter, CompactionExecutor, CompactionOptions, - }; - use crate::dataset::{ColumnAlteration, NewColumnTransform}; - use std::borrow::Cow; - - let mut dataset = write("memory://", LanceFileVersion::V2_0).await; - // `d` gets a file of its own in every fragment. - dataset - .add_columns( - NewColumnTransform::SqlExpressions(vec![("d".into(), "a + 1".into())]), - None, - None, - ) - .await - .unwrap(); - - // `d` already sits alone, so repack it together with `c`. - let executor = RewriteColumnsExecutor::new(&["c", "d"], None); - let mut results = Vec::new(); - for fragment in dataset.get_fragments() { - // The file `d` lived in, fully tombstoned, as merge_insert leaves it. - let mut tombstoned = fragment.metadata().files[1].clone(); - tombstoned.fields = vec![TOMBSTONE_FIELD_ID].into(); - let mut result = executor - .execute( - Cow::Borrowed(&dataset), - fragment, - &CompactionOptions::default(), - ) - .await - .unwrap() - .expect("c and d have files to repack"); - result.fragment.files.insert(1, tombstoned); - results.push(Some(result)); - } - - dataset - .alter_columns(&[ColumnAlteration::new("b".into()).rename("b2".into())]) - .await - .unwrap(); - let expected = dataset.scan().try_into_batch().await.unwrap(); - - RewriteColumnsCommitter - .commit( - &mut dataset, - results, - Arc::new(DatasetIndexRemapperOptions::default()), - &CompactionOptions::default(), - ) - .await - .unwrap(); - - for fragment in dataset.get_fragments() { - let layout: Vec> = fragment - .metadata() - .files - .iter() - .map(|file| file.fields.to_vec()) - .collect(); - assert_eq!( - layout, - vec![vec![0, 1, TOMBSTONE_FIELD_ID], vec![2, 3]], - "the rewrite landed and the commit dropped the tombstoned file" - ); - } - assert_eq!(dataset.scan().try_into_batch().await.unwrap(), expected); - dataset.validate().await.unwrap(); - } -} diff --git a/rust/lance/src/dataset/tests/dataset_overlay_index_masking.rs b/rust/lance/src/dataset/tests/dataset_overlay_index_masking.rs index 5c5fb92f374..6c78aa368ea 100644 --- a/rust/lance/src/dataset/tests/dataset_overlay_index_masking.rs +++ b/rust/lance/src/dataset/tests/dataset_overlay_index_masking.rs @@ -255,11 +255,10 @@ async fn test_overlay_stale_drop_and_new_match(#[values(false, true)] stable_row assert_eq!(ids_matching(&dataset, "age = 20").await, vec![2]); } -/// Horizontal compaction ([`Dataset::rewrite_columns`]) commits with an empty -/// `fields_modified`, so an existing data overlay must be carried forward and -/// keep shadowing the base — even when the rewrite repacks an unrelated column. +/// A column repack moves base values without changing them, so an overlay on +/// the fragment keeps shadowing the base it was written over. #[tokio::test] -async fn rewrite_columns_preserves_data_overlay() { +async fn repack_columns_preserves_data_overlay() { let dataset = create_base_dataset().await; // Fragment 0, offset 1 is id=1, age=10; the overlay changes its age to 999. let mut dataset = commit_overlay( @@ -273,8 +272,16 @@ async fn rewrite_columns_preserves_data_overlay() { .await; assert_eq!(ids_matching(&dataset, "age = 999").await, vec![1]); - // Repack an unrelated column; the overlay on `age` must survive. - dataset.rewrite_columns(&["id"], None).await.unwrap(); + // Split `id` into a file of its own; the overlay on `age` must survive. + let options = crate::dataset::optimize::CompactionOptions { + column_groups: vec![vec!["id".into()]], + scope: crate::dataset::optimize::CompactionScope::RepackColumns, + ..Default::default() + }; + crate::dataset::optimize::compact_files(&mut dataset, options, None) + .await + .unwrap(); + assert_eq!(dataset.get_fragment(0).unwrap().metadata().files.len(), 2); assert_eq!( dataset.get_fragment(0).unwrap().metadata().overlays.len(), diff --git a/rust/lance/src/index.rs b/rust/lance/src/index.rs index 70209f12ccf..9366343ce95 100644 --- a/rust/lance/src/index.rs +++ b/rust/lance/src/index.rs @@ -13749,6 +13749,7 @@ mod tests { .filter(|fragment| covered.contains(fragment.id as u32)) .cloned() .collect(), + kind: Default::default(), }], read_version: dataset.version().version, options: CompactionOptions::default(), @@ -13863,9 +13864,9 @@ mod tests { // fragment 2 when the remap is missing, which would leave this nothing // to ask about and hide whether the guard is sensitive on its own. let plan = CompactionPlan { - tasks: vec![TaskData { - fragments: dataset.fragments().as_ref().clone(), - }], + tasks: vec![TaskData::rewrite_fragments( + dataset.fragments().as_ref().clone(), + )], read_version: dataset.version().version, options: CompactionOptions::default(), }; diff --git a/rust/lance/src/index/frag_reuse.rs b/rust/lance/src/index/frag_reuse.rs index 25f260229f1..191a37d5b15 100644 --- a/rust/lance/src/index/frag_reuse.rs +++ b/rust/lance/src/index/frag_reuse.rs @@ -3410,6 +3410,7 @@ mod tests { read_version: dataset.manifest.version, original_fragments: vec![overlaid], row_addrs: Some(serialized), + repacked_files: None, }; let fragments_before: Vec = dataset.fragments().iter().map(|f| f.id).collect(); let error = crate::dataset::optimize::commit_compaction( diff --git a/rust/lance/src/io/commit/conflict_resolver.rs b/rust/lance/src/io/commit/conflict_resolver.rs index b5808c27357..3dc2c3b4173 100644 --- a/rust/lance/src/io/commit/conflict_resolver.rs +++ b/rust/lance/src/io/commit/conflict_resolver.rs @@ -18,7 +18,7 @@ use lance_index::frag_reuse::FRAG_REUSE_INDEX_NAME; use lance_index::mem_wal::{CompactedSsTable, MEM_WAL_INDEX_NAME}; use lance_select::{RowAddrTreeMap, RowSetOps}; use lance_table::format::IndexMetadata; -use lance_table::format::overlay::{OverlayCoverage, TOMBSTONE_FIELD_ID}; +use lance_table::format::overlay::OverlayCoverage; use lance_table::format::pb::fragment_reuse_index_details::{InlineContent, Transition}; use lance_table::system_index::frag_reuse::lineage::TaggedLineage; use lance_table::system_index::frag_reuse::metadata::is_tagged; @@ -681,34 +681,6 @@ impl<'a> TransactionRebase<'a> { ); } } - // An in-place column rewrite publishes each fragment's file - // list as it built it at the read version, so a file the - // projection removed (it held only dropped fields) would come - // back, and so would a new file holding only a dropped field. - // `fields_modified` can't catch this: a rewrite that moves - // values without changing them (horizontal compaction) leaves - // it empty. Apply the projection's own keep rule to our files. - // A file holding only tombstones is skipped: an in-place - // merge_insert leaves one behind and the commit drops it. - if matches!(self_update_mode, Some(UpdateMode::RewriteColumns)) { - let remaining: HashSet = - schema.fields_pre_order().map(|field| field.id).collect(); - let drops_a_file = self_updated_fragments.iter().any(|fragment| { - let spilled = fragment.spilled_row_lineage_field_ids(); - fragment.files.iter().any(|file| { - file.fields.iter().any(|id| *id != TOMBSTONE_FIELD_ID) - && !file - .fields - .iter() - .any(|id| remaining.contains(id) || spilled.contains(id)) - }) - }); - if drops_a_file { - return Err( - self.retryable_conflict_err(other_transaction, other_version) - ); - } - } Ok(()) } Operation::DataOverlay { groups } => { @@ -1088,7 +1060,7 @@ impl<'a> TransactionRebase<'a> { fragment.files.iter().any(|file| { file.fields.iter().any(|field| { *field - == TOMBSTONE_FIELD_ID + == lance_table::format::overlay::TOMBSTONE_FIELD_ID }) }) }) @@ -6275,7 +6247,7 @@ mod tests { ( "Moved DataReplacement vs DataReplacement of an overlapping field", Operation::DataReplacement { - replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01.clone())], + replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01)], data_change: false, }, Operation::DataReplacement { From f80007824e368e990ef19db8782ce0ce5347bde5 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Fri, 2 Oct 2026 23:49:14 +0800 Subject: [PATCH 11/24] refactor(compaction): keep scope a run option, warn once about stale column groups --- python/python/lance/dataset.py | 1 - rust/lance/src/dataset/optimize.rs | 31 +++++++++++-------- .../src/dataset/optimize/tests/repack.rs | 15 +++------ 3 files changed, 22 insertions(+), 25 deletions(-) diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index 8a552edc031..1a7e6ff85f2 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -7654,7 +7654,6 @@ def compact_files( ``lance.compaction.max_source_bytes``, ``lance.compaction.max_data_files_per_fragment``, ``lance.compaction.column_groups``, - ``lance.compaction.scope``, ``lance.compaction.data_storage_version``. Parameters diff --git a/rust/lance/src/dataset/optimize.rs b/rust/lance/src/dataset/optimize.rs index 5f14cb8d6b8..cfcd5f24750 100644 --- a/rust/lance/src/dataset/optimize.rs +++ b/rust/lance/src/dataset/optimize.rs @@ -433,7 +433,6 @@ impl CompactionOptions { /// - `lance.compaction.max_source_bytes` /// - `lance.compaction.max_overlays_per_fragment` /// - `lance.compaction.max_data_files_per_fragment` - /// - `lance.compaction.scope`: `all`, `rewrite_fragments` or `repack_columns` /// - `lance.compaction.data_storage_version` /// - `lance.compaction.column_groups`: groups separated by `;`, columns in /// a group by `,`, surrounding whitespace trimmed (`"b, c; d"` is @@ -586,9 +585,6 @@ impl CompactionOptions { })?), }; } - "scope" => { - self.scope = CompactionScope::try_from(value.as_str())?; - } "max_data_files_per_fragment" => { self.max_data_files_per_fragment = Some(value.parse().map_err(|_| { Error::invalid_input(format!( @@ -966,6 +962,20 @@ impl CompactionPlanner for DefaultCompactionPlanner { dataset.manifest.data_storage_format.lance_file_format(), write_version, )?; + // Warned about once here; the tasks skip the stale names silently. + for column in self.options.column_groups.iter().flatten() { + if !dataset + .schema() + .fields + .iter() + .any(|field| &field.name == column) + { + warn!( + "Ignoring column_groups entry \"{column}\": it is not a top-level column \ + of the dataset" + ); + } + } if !self.options.column_groups.is_empty() { // Surface a bad column-group config here rather than mid-write. column_group_schemas(dataset.schema(), &self.options.column_groups)?; @@ -2751,9 +2761,10 @@ async fn reserve_fragment_ids( /// are configured: the columns no group claims first (when any), then one /// schema per group. Columns keep the dataset's order within each file. /// -/// A name that is not a top-level column is skipped with a warning: the groups -/// usually come from table config, which `drop_columns` and a rename leave as -/// they are, and a stale name must not stop every later compaction. +/// A name that is not a top-level column is skipped: the groups usually come +/// from table config, which `drop_columns` and a rename leave as they are, and +/// a stale name must not stop every later compaction. The planner warns about +/// it. fn column_group_schemas( schema: &lance_core::datatypes::Schema, groups: &[Vec], @@ -2769,12 +2780,6 @@ fn column_group_schemas( "column_groups lists column \"{column}\" more than once" ))); } - if !names().any(|name| name == column) { - warn!( - "Ignoring column_groups entry \"{column}\": it is not a top-level column \ - of the dataset" - ); - } } let claimed = |name: &str| groups.iter().flatten().any(|column| column == name); let rest: Vec<&str> = names().filter(|name| !claimed(name)).collect(); diff --git a/rust/lance/src/dataset/optimize/tests/repack.rs b/rust/lance/src/dataset/optimize/tests/repack.rs index 595e2b9ef4c..c5c8622de19 100644 --- a/rust/lance/src/dataset/optimize/tests/repack.rs +++ b/rust/lance/src/dataset/optimize/tests/repack.rs @@ -543,19 +543,12 @@ fn max_data_files_per_fragment_must_be_positive() { #[test] fn repack_options_parse_from_config() { - let config = HashMap::from([ - ( - "lance.compaction.max_data_files_per_fragment".to_string(), - "4".to_string(), - ), - ( - "lance.compaction.scope".to_string(), - "repack_columns".to_string(), - ), - ]); + let config = HashMap::from([( + "lance.compaction.max_data_files_per_fragment".to_string(), + "4".to_string(), + )]); let options = CompactionOptions::from_dataset_config(&config).unwrap(); assert_eq!(options.max_data_files_per_fragment, Some(4)); - assert_eq!(options.scope, CompactionScope::RepackColumns); } /// The file holding a fragment's spilled row lineage keeps its columns: the From e0b560091ef785d40ecabd8679a193a5ff5793a5 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Sat, 3 Oct 2026 00:26:44 +0800 Subject: [PATCH 12/24] fix(compaction): address review of the repack planner --- .../src/transaction/index_maintenance.rs | 3 +- .../src/transaction/manifest_build.rs | 4 +- rust/lance/src/dataset/optimize.rs | 123 +++++++++--- rust/lance/src/dataset/optimize/repack.rs | 112 ++++++----- .../src/dataset/optimize/tests/repack.rs | 180 +++++++++++++++++- 5 files changed, 344 insertions(+), 78 deletions(-) diff --git a/rust/lance-table/src/transaction/index_maintenance.rs b/rust/lance-table/src/transaction/index_maintenance.rs index d1cd681338a..63471ec29db 100644 --- a/rust/lance-table/src/transaction/index_maintenance.rs +++ b/rust/lance-table/src/transaction/index_maintenance.rs @@ -100,8 +100,7 @@ impl Transaction { /// present in `previous_fragments`, the fields whose backing data file /// changed (`merge_rewritten_fields`); a data replacement that changes /// data rewrites the fields its new files carry, read through `schema`. - /// Any other - /// operation rewrites nothing. `previous_fragments` is the caller's + /// Any other operation rewrites nothing. `previous_fragments` is the caller's /// "before" list: the current manifest's for a commit, the read /// version's for a rebase. pub fn rewritten_physical_columns( diff --git a/rust/lance-table/src/transaction/manifest_build.rs b/rust/lance-table/src/transaction/manifest_build.rs index 976703f55e4..adb7351e622 100644 --- a/rust/lance-table/src/transaction/manifest_build.rs +++ b/rust/lance-table/src/transaction/manifest_build.rs @@ -1453,8 +1453,8 @@ impl Transaction { } } - // The replaced fields' coverage of the modified fragments was - // withdrawn by `prepare_indices`. + // When the data changed, the replaced fields' coverage of the + // modified fragments was withdrawn by `prepare_indices`. } Operation::DataOverlay { groups } => { // Stamp each overlay with the version this commit is producing. diff --git a/rust/lance/src/dataset/optimize.rs b/rust/lance/src/dataset/optimize.rs index cfcd5f24750..3141c6326e1 100644 --- a/rust/lance/src/dataset/optimize.rs +++ b/rust/lance/src/dataset/optimize.rs @@ -323,7 +323,12 @@ pub struct CompactionOptions { /// for a rewrite, so this is the trigger that collapses those files. The /// repack rewrites only the columns that move and keeps rows, fragment ids /// and indices as they are (see [`CompactionTaskKind::RepackColumns`]). - /// The count is [`FragmentColumnLayoutStats::live_file_count`]. + /// The count is [`FragmentColumnLayoutStats::live_file_count`]. The + /// columns of one file stay where they are unless the limit is 1, and + /// blob columns never move, so a fragment can stay above the limit; so + /// can one with more `column_groups` than the limit allows. The repack + /// writes each new file whole, with no `max_bytes_per_file` split. Not + /// planned under `ForceBinaryCopy`, since a repack reencodes. /// /// Defaults to `None` (no file-count trigger). Must be at least 1. /// @@ -351,19 +356,21 @@ pub struct CompactionOptions { /// Top-level columns to keep in their own data files. /// /// Each inner list becomes one data file per fragment holding exactly - /// those columns; every column not listed goes to one shared file. Empty - /// (the default) puts every column in a single file per fragment. + /// those columns. Empty (the default) means no groups. /// - /// Fragments a compaction rewrites are written in this layout. A fragment - /// it leaves alone whose files do not match the layout has its columns - /// repacked to match (see [`CompactionTaskKind::RepackColumns`]), unless - /// `scope` is [`CompactionScope::RewriteFragments`]. A name that is not a - /// top-level column of the dataset (say, one dropped or renamed since the - /// groups were configured) is ignored with a warning. Binary copy is - /// disabled when groups are set (`ForceBinaryCopy` is rejected), and - /// `max_bytes_per_file` is ignored so every group's files split at the - /// same rows. Set from dataset config with `lance.compaction.column_groups` - /// (see [`Self::from_dataset_config`]). + /// A fragment a compaction rewrites gets one file per group plus one file + /// for every column no group names. A fragment it leaves alone has a + /// group repacked into a file of its own when no single file holds + /// exactly that group's columns (see [`CompactionTaskKind::RepackColumns`]), + /// unless `scope` is [`CompactionScope::RewriteFragments`]; the columns no + /// group names are merged only when the fragment is over + /// `max_data_files_per_fragment`. A group holding a blob column is left + /// as it is. A name that is not a top-level column of the dataset (say, + /// one dropped or renamed since the groups were configured) is ignored + /// with a warning. Binary copy is disabled when groups are set + /// (`ForceBinaryCopy` is rejected), and `max_bytes_per_file` is ignored so + /// every group's files split at the same rows. Set from dataset config + /// with `lance.compaction.column_groups` (see [`Self::from_dataset_config`]). #[serde(default)] pub column_groups: Vec>, /// Transaction properties to store with this commit. @@ -1207,15 +1214,27 @@ impl DefaultCompactionPlanner { let groups = if self.options.column_groups.is_empty() { None } else { - Some(column_group_ids( + Some(claimed_column_ids( dataset.schema(), &self.options.column_groups, - )?) + )) }; let max_files = self.options.max_data_files_per_fragment; if groups.is_none() && max_files.is_none() { return Ok(Vec::new()); } + // A repack reencodes the columns it moves. `validate()` already + // rejects `column_groups` with ForceBinaryCopy. + if matches!( + self.options.compaction_mode(), + CompactionMode::ForceBinaryCopy + ) { + log::info!( + "not planning column repacks: compaction_mode=ForceBinaryCopy does not \ + reencode, and a repack does" + ); + return Ok(Vec::new()); + } let stats = dataset.column_layout_stats(); let tasks = dataset .manifest @@ -1457,7 +1476,8 @@ fn limit_tasks_to_source_budget( total_fragments += task.fragments.len(); total_rows = total_rows.saturating_add(live_rows); if options.max_source_bytes.is_some() { - total_bytes = total_bytes.saturating_add(task_source_bytes(&task, &schema_field_ids)?); + total_bytes = + total_bytes.saturating_add(task_source_bytes(&task, schema, &schema_field_ids)?); } let over_budget = options @@ -1493,19 +1513,42 @@ fn limit_tasks_to_source_budget( /// /// Files whose fields are all absent from `schema_field_ids` only back /// dropped columns; compaction does not read them, so they are neither -/// counted nor required to have a recorded size. +/// counted nor required to have a recorded size. A column repack reads only +/// the files holding the columns it moves, and no overlay, so only those +/// files count. /// Only sizes recorded in the manifest are used: a missing size is an error /// rather than a metadata request against object storage, which would turn /// planning into one round trip per file. Deletion files are not counted. -fn task_source_bytes(task: &TaskData, schema_field_ids: &HashSet) -> Result { +fn task_source_bytes( + task: &TaskData, + schema: &lance_core::datatypes::Schema, + schema_field_ids: &HashSet, +) -> Result { + // The field ids a task reads: every schema field, or a repack's columns. + let read_ids: HashSet = match &task.kind { + CompactionTaskKind::RewriteFragments => schema_field_ids.clone(), + CompactionTaskKind::RepackColumns { files } => { + let columns: Vec = files.iter().flatten().copied().collect(); + schema + .project_by_ids(&columns, true) + .field_ids() + .into_iter() + .collect() + } + }; + let reads_overlays = task.kind == CompactionTaskKind::RewriteFragments; let mut total_bytes = 0_u64; for fragment in &task.fragments { - let overlay_files = fragment.overlays.iter().map(|overlay| &overlay.data_file); + let overlay_files = fragment + .overlays + .iter() + .filter(|_| reads_overlays) + .map(|overlay| &overlay.data_file); for data_file in fragment.files.iter().chain(overlay_files) { if !data_file .fields .iter() - .any(|field_id| schema_field_ids.contains(field_id)) + .any(|field_id| read_ids.contains(field_id)) { continue; } @@ -2794,15 +2837,25 @@ fn column_group_schemas( .collect() } -/// [`column_group_schemas`] as the field ids of each file's top-level columns. -fn column_group_ids( +/// The field ids of each group's top-level columns, without the columns no +/// group names. A name that is not a top-level column is skipped, and so is a +/// group left empty. +fn claimed_column_ids( schema: &lance_core::datatypes::Schema, groups: &[Vec], -) -> Result>> { - Ok(column_group_schemas(schema, groups)? +) -> Vec> { + groups .iter() - .map(|group| group.fields.iter().map(|field| field.id).collect()) - .collect()) + .map(|group| { + schema + .fields + .iter() + .filter(|field| group.contains(&field.name)) + .map(|field| field.id) + .collect::>() + }) + .filter(|group| !group.is_empty()) + .collect() } /// Compaction's grouped write: fan one read of the merged rows out to one @@ -3721,7 +3774,7 @@ fn append_row_lineage_columns( /// Fragment rewrites commit as one `Operation::Rewrite`; column repacks then /// commit as one `Operation::DataReplacement` that moves values without /// changing them, so a run with both kinds of task makes two versions. A -/// planner never puts one fragment in both kinds of task. If the second +/// fragment in both kinds of task is rejected before either commit. If the second /// commit fails, the first stays committed and the error is returned; the /// repacks can be planned again. pub async fn commit_compaction( @@ -3733,6 +3786,22 @@ pub async fn commit_compaction( let (repacks, rewrites): (Vec<_>, Vec<_>) = completed_tasks .into_iter() .partition(|task| task.repacked_files.is_some()); + // Otherwise the rewrite commits and the repack then conflicts with it. + let rewritten: HashSet = rewrites + .iter() + .flat_map(|task| task.original_fragments.iter().map(|fragment| fragment.id)) + .collect(); + if let Some(fragment_id) = repacks + .iter() + .filter_map(|task| task.repacked_files.as_ref()) + .map(|repacked| repacked.fragment_id) + .find(|id| rewritten.contains(id)) + { + return Err(Error::invalid_input(format!( + "fragment {fragment_id} is both rewritten and repacked by this compaction; \ + a plan may put a fragment in one kind of task only" + ))); + } let mut metrics = commit_rewrites(dataset, rewrites, remap_options, options).await?; metrics += repack::commit_repacked_files(dataset, repacks, options).await?; Ok(metrics) diff --git a/rust/lance/src/dataset/optimize/repack.rs b/rust/lance/src/dataset/optimize/repack.rs index fc7a667cc0d..295c5ff1f6c 100644 --- a/rust/lance/src/dataset/optimize/repack.rs +++ b/rust/lance/src/dataset/optimize/repack.rs @@ -1,8 +1,9 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors -//! Column repacking, the compaction task that rewrites one fragment's columns -//! into fewer data files without moving rows. +//! Column repacking, the compaction task that rewrites some of one fragment's +//! columns into new data files without moving rows: fewer files, or the files +//! `column_groups` asks for. //! //! Every `add_columns` backfill gives each fragment one more data file, and a //! large fragment with few deletions never qualifies for a rewrite, so its @@ -27,11 +28,15 @@ use crate::dataset::transaction::{DataReplacementGroup, Operation, TransactionBu use crate::dataset::{Dataset, cleanup_data_fragments, versions}; use crate::{Error, Result}; -/// The field ids of `field` and everything beneath it. -fn subtree_ids(field: &LanceField, ids: &mut HashSet) { - ids.insert(field.id); +/// The ids of the leaf fields at and beneath `field`. A file holds a column's +/// data only through its leaves: a V2.0 file can keep a struct's header after +/// the struct's last child in it was dropped. +fn leaf_ids(field: &LanceField, ids: &mut HashSet) { + if field.children.is_empty() { + ids.insert(field.id); + } for child in &field.children { - subtree_ids(child, ids); + leaf_ids(child, ids); } } @@ -53,24 +58,25 @@ fn file_coverage(file: &DataFile, schema: &Schema) -> HashSet { /// The new files a repack of `fragment` should write, each the top-level /// field ids it holds, or `None` when no repack is due. /// -/// `groups` partitions the top-level columns into the files wanted, as -/// `column_group_schemas` orders them. Without groups, the columns of one -/// live file stay together (the file holding the fragment's spilled row -/// lineage, else the one holding the most bytes) and every other column goes -/// to one new file; when `max_files` is 1, every column goes to one file. A -/// group is written unless one live file holds exactly its columns. A column -/// the fragment has no data for (added as all nulls) is left out of its -/// group, and a group holding a blob column, or a column whose fields are -/// only partly in the fragment's files, is left as it is. +/// `groups` are the column groups, each the top-level field ids to keep in a +/// file of their own. A group is moved unless one live file holds exactly its +/// columns; moving it out leaves the other columns where they were. The +/// columns no group names are merged into one new file only when the +/// fragment is over `max_files`. Without groups and over `max_files`, the +/// columns of one live file stay where they are (the file holding the +/// fragment's spilled row lineage, else one holding a column that cannot +/// move, else the largest) and every other column goes to one new file; with +/// `max_files` 1 and every column movable, all of them go to one file. /// -/// The new files carry no row lineage, so a repack that would move every -/// column out of the file holding the fragment's spilled lineage is not -/// planned: that file would stay for the lineage alone, and a later rewrite -/// of the fragment is what reclaims it. +/// A column the fragment has no data for (added as all nulls) is left out of +/// its group, and a group holding a blob column, or a column whose fields are +/// only partly in the fragment's files, is left as it is. The new files carry +/// no row lineage, so a repack that would move every column out of the file +/// holding the fragment's spilled lineage is not planned: that file would +/// stay for the lineage alone until a rewrite of the fragment reclaims it. /// -/// A group out of place is a reason to repack; so is `live_files` (the -/// fragment's [`FragmentColumnLayoutStats::live_file_count`]) above -/// `max_files`. +/// `live_files` is the fragment's +/// [`FragmentColumnLayoutStats::live_file_count`]. /// /// [`FragmentColumnLayoutStats::live_file_count`]: crate::dataset::compaction_stats::FragmentColumnLayoutStats::live_file_count pub(super) fn plan_fragment_repack( @@ -101,63 +107,77 @@ pub(super) fn plan_fragment_repack( .map(|file| (file, file_coverage(file, schema))) .collect(); let covered: HashSet = live.iter().flat_map(|(_, ids)| ids).copied().collect(); - let column_ids: Vec<(i32, HashSet)> = schema + let column_leaves: Vec<(i32, HashSet)> = schema .fields .iter() .map(|field| { let mut ids = HashSet::new(); - subtree_ids(field, &mut ids); + leaf_ids(field, &mut ids); (field.id, ids) }) .collect(); - let present: BTreeSet = column_ids + let present: BTreeSet = column_leaves .iter() - .filter(|(_, ids)| !ids.is_disjoint(&covered)) + .filter(|(_, leaves)| !leaves.is_disjoint(&covered)) .map(|(column, _)| *column) .collect(); let movable: HashSet = schema .fields .iter() - .zip(&column_ids) - .filter(|(field, (_, ids))| !holds_blob(field) && ids.is_subset(&covered)) + .zip(&column_leaves) + .filter(|(field, (column, leaves))| { + present.contains(column) && !holds_blob(field) && leaves.is_subset(&covered) + }) .map(|(field, _)| field.id) .collect(); - // The top-level columns each live file holds any field of. + // The top-level columns each live file holds data for. let file_columns: Vec> = live .iter() .map(|(_, coverage)| { - column_ids + column_leaves .iter() - .filter(|(_, ids)| !ids.is_disjoint(coverage)) + .filter(|(_, leaves)| !leaves.is_disjoint(coverage)) .map(|(column, _)| *column) .collect() }) .collect(); + let holds_lineage = |file: &DataFile| file.fields.iter().any(|id| spilled.contains(id)); let wanted: Vec> = match groups { - Some(groups) => groups - .iter() - .map(|group| group.iter().copied().collect()) - .collect(), - None if max_files == Some(1) => vec![present.clone()], + Some(groups) => { + let mut wanted: Vec> = groups + .iter() + .map(|group| group.iter().copied().collect()) + .collect(); + let claimed: BTreeSet = wanted.iter().flatten().copied().collect(); + let rest: BTreeSet = present.difference(&claimed).copied().collect(); + let rest_files = file_columns + .iter() + .filter(|columns| !columns.is_disjoint(&rest)) + .count(); + if over_file_limit && rest_files > 1 { + wanted.push(rest); + } + wanted + } + None if max_files == Some(1) && present.iter().all(|c| movable.contains(c)) => { + vec![present.clone()] + } None => { let kept = live .iter() .zip(&file_columns) .max_by_key(|((file, _), columns)| { ( - file.fields.iter().any(|id| spilled.contains(id)), + holds_lineage(file), + columns.iter().any(|column| !movable.contains(column)), file.file_size_bytes.get().map(|size| size.get()), columns.len(), ) }) .map(|(_, columns)| columns.clone()) .unwrap_or_default(); - let rest = present - .iter() - .copied() - .filter(|column| !kept.contains(column) && movable.contains(column)) - .collect(); + let rest = present.difference(&kept).copied().collect(); vec![kept, rest] } }; @@ -179,7 +199,8 @@ pub(super) fn plan_fragment_repack( .collect(); let moved: HashSet = out_of_place.iter().flatten().copied().collect(); let strands_lineage = live.iter().zip(&file_columns).any(|((file, _), columns)| { - file.fields.iter().any(|id| spilled.contains(id)) + holds_lineage(file) + && !columns.is_empty() && columns.iter().all(|column| moved.contains(column)) }); (!out_of_place.is_empty() && !strands_lineage).then_some(out_of_place) @@ -236,7 +257,7 @@ pub(super) async fn execute_repack( } // Read the base values only: the overlays stay on the fragment and keep - // shadowing them, so the new files hold exactly what the old ones did. + // shadowing them, so every live row reads as it did. let mut base = fragment.clone(); base.overlays.clear(); let source = FileFragment::new(Arc::new(dataset.clone()), base); @@ -326,7 +347,8 @@ async fn write_one_file( /// Commit the files of repack results as one `DataReplacement` that moves /// values without changing them. Each result applies to its fragment as it /// stands at commit, so a concurrent change to other columns of the fragment -/// is kept; the conflict check retries one that touches the moved columns. +/// is kept; one that touches the moved columns fails the commit with a +/// retryable conflict. pub(super) async fn commit_repacked_files( dataset: &mut Dataset, results: Vec, diff --git a/rust/lance/src/dataset/optimize/tests/repack.rs b/rust/lance/src/dataset/optimize/tests/repack.rs index c5c8622de19..93571d178de 100644 --- a/rust/lance/src/dataset/optimize/tests/repack.rs +++ b/rust/lance/src/dataset/optimize/tests/repack.rs @@ -170,8 +170,8 @@ async fn repack_follows_column_groups() { }) .collect::>(); files.sort(); - // The columns no group claims share one file. - assert_eq!(files, vec![vec![0, 1, 3], vec![2, 4]], "{files:?}"); + // Only the group moves; the other columns stay where they were. + assert_eq!(files, vec![vec![0, 1], vec![2, 4], vec![3]], "{files:?}"); } assert_eq!(dataset.scan().try_into_batch().await.unwrap(), before); dataset.validate().await.unwrap(); @@ -604,3 +604,179 @@ fn repack_keeps_the_spilled_lineage_file() { None ); } + +/// Over the file limit with groups set, the columns no group names are merged +/// into one file as well. +#[tokio::test] +async fn repack_merges_unclaimed_columns_over_the_limit() { + let mut dataset = write_backfilled().await; + let before = dataset.scan().try_into_batch().await.unwrap(); + + compact_files(&mut dataset, repack_options(Some(2), vec![vec!["e"]]), None) + .await + .unwrap(); + + for files in layout(&dataset) { + let mut files = files + .into_iter() + .map(|fields| { + fields + .into_iter() + .filter(|id| *id != TOMBSTONE_FIELD_ID) + .collect::>() + }) + .collect::>(); + files.sort(); + assert_eq!(files, vec![vec![0, 1, 2, 3], vec![4]], "{files:?}"); + } + assert_eq!(dataset.scan().try_into_batch().await.unwrap(), before); + dataset.validate().await.unwrap(); +} + +fn lance_schema(arrow: Schema) -> lance_core::datatypes::Schema { + lance_core::datatypes::Schema::try_from(&arrow).unwrap() +} + +fn v2_0_file(path: &str, fields: Vec) -> lance_table::format::DataFile { + let indices = (0..fields.len() as i32).collect(); + lance_table::format::DataFile::new( + path, + fields, + indices, + lance_file::version::ConcreteFileVersion::V2_0, + None, + None, + ) +} + +/// A V2.0 file keeps a struct's header after the struct's last child in it +/// is dropped. That header holds no data, so the file does not hold the +/// struct: the struct sits alone in the file with its live child, and no +/// repack is planned (one would be planned again on every run). +#[test] +fn repack_ignores_struct_header_without_children() { + use arrow_schema::Fields; + let arrow = Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new( + "s", + DataType::Struct(Fields::from(vec![ + Field::new("x", DataType::Int32, true), + Field::new("y", DataType::Int32, true), + ])), + true, + ), + ]); + // a=0, s=1, s.x=2 (dropped), s.y=3. + let schema = lance_schema(arrow).project_by_ids(&[0, 1, 3], false); + let mut fragment = Fragment::new(0); + fragment.files = vec![ + v2_0_file("a.lance", vec![0, 1, 2]), + v2_0_file("s.lance", vec![1, 3]), + ]; + + assert_eq!( + crate::dataset::optimize::repack::plan_fragment_repack( + &schema, + &fragment, + 2, + Some(&[vec![1]]), + None + ), + None + ); + assert_eq!( + crate::dataset::optimize::repack::plan_fragment_repack( + &schema, + &fragment, + 2, + None, + Some(1) + ), + Some(vec![vec![0, 1]]) + ); +} + +/// A blob column never moves, so with a limit of 1 the other columns are +/// merged next to it instead of nothing being planned. +#[test] +fn repack_merges_around_a_blob_column() { + let mut blob = Field::new("img", DataType::LargeBinary, true); + blob.set_metadata(std::collections::HashMap::from([( + lance_arrow::BLOB_META_KEY.to_string(), + "true".to_string(), + )])); + let schema = lance_schema(Schema::new(vec![ + Field::new("id", DataType::Int32, true), + blob, + Field::new("d", DataType::Int32, true), + Field::new("e", DataType::Int32, true), + ])); + let mut fragment = Fragment::new(0); + fragment.files = vec![ + v2_0_file("base.lance", vec![0, 1]), + v2_0_file("d.lance", vec![2]), + v2_0_file("e.lance", vec![3]), + ]; + assert_eq!( + crate::dataset::optimize::repack::plan_fragment_repack( + &schema, + &fragment, + 3, + None, + Some(1) + ), + Some(vec![vec![2, 3]]) + ); +} + +#[tokio::test] +async fn repack_is_not_planned_under_force_binary_copy() { + let dataset = write_backfilled().await; + let options = CompactionOptions { + compaction_mode: Some(CompactionMode::ForceBinaryCopy), + ..repack_options(Some(1), vec![]) + }; + let plan = plan_compaction(&dataset, &options).await.unwrap(); + assert!(plan.tasks.is_empty(), "{plan:?}"); +} + +/// `max_source_bytes` counts only the files a repack reads. +#[tokio::test] +async fn repack_counts_only_the_files_it_reads_against_the_byte_budget() { + let dataset = write_backfilled().await; + let fragment = &dataset.manifest.fragments[0]; + let size = |index: usize| fragment.files[index].file_size_bytes.get().unwrap().get(); + // Each repack leaves the base file alone and merges d and e, so a budget + // for two of those fits both, though not one whole fragment. + let options = CompactionOptions { + max_source_bytes: Some(2 * (size(1) + size(2))), + ..repack_options(Some(2), vec![]) + }; + assert!(size(0) + size(1) + size(2) > 2 * (size(1) + size(2))); + let plan = plan_compaction(&dataset, &options).await.unwrap(); + assert_eq!(plan.tasks.len(), 2, "{plan:?}"); +} + +/// A fragment both rewritten and repacked is refused before either commit. +#[tokio::test] +async fn commit_rejects_a_fragment_in_both_kinds_of_task() { + let mut dataset = write_backfilled().await; + let version = dataset.manifest.version; + let repacks = stale_repack_results(&dataset, &repack_options(Some(1), vec![])).await; + let rewrite = CompactionTask { + task: TaskData::rewrite_fragments(vec![dataset.manifest.fragments[0].clone()]), + read_version: version, + options: CompactionOptions::default(), + } + .execute(&dataset) + .await + .unwrap(); + + let mut results = repacks; + results.push(rewrite); + let err = commit_results(&mut dataset, results).await.unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. }), "{err}"); + dataset.checkout_latest().await.unwrap(); + assert_eq!(dataset.manifest.version, version); +} From b047330f822b7c248aa5797b66b40a3cab8ce586 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Sat, 3 Oct 2026 00:59:28 +0800 Subject: [PATCH 13/24] fix(compaction): keep non-movable columns out of the merged repack file --- rust/lance/src/dataset/optimize.rs | 27 +++++++----- rust/lance/src/dataset/optimize/repack.rs | 26 +++++++---- .../src/dataset/optimize/tests/repack.rs | 44 ++++++++++++++++++- 3 files changed, 75 insertions(+), 22 deletions(-) diff --git a/rust/lance/src/dataset/optimize.rs b/rust/lance/src/dataset/optimize.rs index 3141c6326e1..5f173d14833 100644 --- a/rust/lance/src/dataset/optimize.rs +++ b/rust/lance/src/dataset/optimize.rs @@ -324,9 +324,10 @@ pub struct CompactionOptions { /// repack rewrites only the columns that move and keeps rows, fragment ids /// and indices as they are (see [`CompactionTaskKind::RepackColumns`]). /// The count is [`FragmentColumnLayoutStats::live_file_count`]. The - /// columns of one file stay where they are unless the limit is 1, and - /// blob columns never move, so a fragment can stay above the limit; so - /// can one with more `column_groups` than the limit allows. The repack + /// columns of one file stay where they are (unless the limit is 1 and + /// every column can move), and blob columns never move, so a fragment can + /// stay above the limit; so can one with more `column_groups` than the + /// limit allows. The repack /// writes each new file whole, with no `max_bytes_per_file` split. Not /// planned under `ForceBinaryCopy`, since a repack reencodes. /// @@ -358,8 +359,8 @@ pub struct CompactionOptions { /// Each inner list becomes one data file per fragment holding exactly /// those columns. Empty (the default) means no groups. /// - /// A fragment a compaction rewrites gets one file per group plus one file - /// for every column no group names. A fragment it leaves alone has a + /// A fragment a compaction rewrites gets one file per group plus one + /// shared file for the columns no group names. A fragment it leaves alone has a /// group repacked into a file of its own when no single file holds /// exactly that group's columns (see [`CompactionTaskKind::RepackColumns`]), /// unless `scope` is [`CompactionScope::RewriteFragments`]; the columns no @@ -1525,15 +1526,17 @@ fn task_source_bytes( schema_field_ids: &HashSet, ) -> Result { // The field ids a task reads: every schema field, or a repack's columns. - let read_ids: HashSet = match &task.kind { - CompactionTaskKind::RewriteFragments => schema_field_ids.clone(), + let read_ids: Cow<'_, HashSet> = match &task.kind { + CompactionTaskKind::RewriteFragments => Cow::Borrowed(schema_field_ids), CompactionTaskKind::RepackColumns { files } => { let columns: Vec = files.iter().flatten().copied().collect(); - schema - .project_by_ids(&columns, true) - .field_ids() - .into_iter() - .collect() + Cow::Owned( + schema + .project_by_ids(&columns, true) + .field_ids() + .into_iter() + .collect(), + ) } }; let reads_overlays = task.kind == CompactionTaskKind::RewriteFragments; diff --git a/rust/lance/src/dataset/optimize/repack.rs b/rust/lance/src/dataset/optimize/repack.rs index 295c5ff1f6c..dfe770818e2 100644 --- a/rust/lance/src/dataset/optimize/repack.rs +++ b/rust/lance/src/dataset/optimize/repack.rs @@ -28,9 +28,9 @@ use crate::dataset::transaction::{DataReplacementGroup, Operation, TransactionBu use crate::dataset::{Dataset, cleanup_data_fragments, versions}; use crate::{Error, Result}; -/// The ids of the leaf fields at and beneath `field`. A file holds a column's -/// data only through its leaves: a V2.0 file can keep a struct's header after -/// the struct's last child in it was dropped. +/// The ids of the leaf fields at and beneath `field`. The planner places a +/// column by where its leaves are: a V2.0 file can keep a struct's header +/// after the struct's last child in it was dropped. fn leaf_ids(field: &LanceField, ids: &mut HashSet) { if field.children.is_empty() { ids.insert(field.id); @@ -65,8 +65,9 @@ fn file_coverage(file: &DataFile, schema: &Schema) -> HashSet { /// fragment is over `max_files`. Without groups and over `max_files`, the /// columns of one live file stay where they are (the file holding the /// fragment's spilled row lineage, else one holding a column that cannot -/// move, else the largest) and every other column goes to one new file; with -/// `max_files` 1 and every column movable, all of them go to one file. +/// move, else the largest) and every other column that can move goes to one +/// new file; with `max_files` 1, every column movable and no spilled lineage, +/// all of them go to one file. /// /// A column the fragment has no data for (added as all nulls) is left out of /// its group, and a group holding a blob column, or a column whose fields are @@ -76,7 +77,9 @@ fn file_coverage(file: &DataFile, schema: &Schema) -> HashSet { /// stay for the lineage alone until a rewrite of the fragment reclaims it. /// /// `live_files` is the fragment's -/// [`FragmentColumnLayoutStats::live_file_count`]. +/// [`FragmentColumnLayoutStats::live_file_count`]. A V2.0 file left holding +/// only a struct header counts there but holds no column here, so it stays +/// until a rewrite of the fragment. /// /// [`FragmentColumnLayoutStats::live_file_count`]: crate::dataset::compaction_stats::FragmentColumnLayoutStats::live_file_count pub(super) fn plan_fragment_repack( @@ -160,7 +163,10 @@ pub(super) fn plan_fragment_repack( } wanted } - None if max_files == Some(1) && present.iter().all(|c| movable.contains(c)) => { + None if max_files == Some(1) + && present.iter().all(|c| movable.contains(c)) + && !live.iter().any(|(file, _)| holds_lineage(file)) => + { vec![present.clone()] } None => { @@ -177,7 +183,11 @@ pub(super) fn plan_fragment_repack( }) .map(|(_, columns)| columns.clone()) .unwrap_or_default(); - let rest = present.difference(&kept).copied().collect(); + let rest = present + .difference(&kept) + .copied() + .filter(|column| movable.contains(column)) + .collect(); vec![kept, rest] } }; diff --git a/rust/lance/src/dataset/optimize/tests/repack.rs b/rust/lance/src/dataset/optimize/tests/repack.rs index 93571d178de..a7a7ebe5a73 100644 --- a/rust/lance/src/dataset/optimize/tests/repack.rs +++ b/rust/lance/src/dataset/optimize/tests/repack.rs @@ -592,7 +592,8 @@ fn repack_keeps_the_spilled_lineage_file() { ), Some(vec![vec![1, 2]]) ); - // Moving every column would leave the lineage file behind, so no repack. + // With a limit of 1 as well: moving every column would leave the lineage + // file behind, so its column stays. assert_eq!( crate::dataset::optimize::repack::plan_fragment_repack( &schema, @@ -601,7 +602,7 @@ fn repack_keeps_the_spilled_lineage_file() { None, Some(1) ), - None + Some(vec![vec![1, 2]]) ); } @@ -780,3 +781,42 @@ async fn commit_rejects_a_fragment_in_both_kinds_of_task() { dataset.checkout_latest().await.unwrap(); assert_eq!(dataset.manifest.version, version); } + +/// A column that cannot move outside the kept file stays where it is; the +/// columns that can move are still merged. +#[test] +fn repack_merges_around_a_blob_outside_the_kept_file() { + use lance_table::format::{ROW_ID_FIELD_ID, RowIdMeta}; + let mut blob = Field::new("img", DataType::LargeBinary, true); + blob.set_metadata(std::collections::HashMap::from([( + lance_arrow::BLOB_META_KEY.to_string(), + "true".to_string(), + )])); + let schema = lance_schema(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + blob, + Field::new("d", DataType::Int32, true), + Field::new("e", DataType::Int32, true), + ])); + let mut fragment = Fragment::new(0); + fragment.files = vec![ + v2_0_file("lineage.lance", vec![0, ROW_ID_FIELD_ID]), + v2_0_file("img.lance", vec![1]), + v2_0_file("d.lance", vec![2]), + v2_0_file("e.lance", vec![3]), + ]; + fragment.row_id_meta = Some(RowIdMeta::Column); + for max_files in [1, 2] { + assert_eq!( + crate::dataset::optimize::repack::plan_fragment_repack( + &schema, + &fragment, + 4, + None, + Some(max_files) + ), + Some(vec![vec![2, 3]]), + "max_files={max_files}" + ); + } +} From e09c16c26764de0272a56ff52833e0bffb03ffbf Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Sat, 3 Oct 2026 01:18:50 +0800 Subject: [PATCH 14/24] fix(compaction): only merge files a repack empties --- rust/lance/src/dataset/optimize.rs | 13 ++-- rust/lance/src/dataset/optimize/repack.rs | 64 +++++++++++----- .../src/dataset/optimize/tests/repack.rs | 73 ++++++++++++++++++- 3 files changed, 122 insertions(+), 28 deletions(-) diff --git a/rust/lance/src/dataset/optimize.rs b/rust/lance/src/dataset/optimize.rs index 5f173d14833..20639c9773d 100644 --- a/rust/lance/src/dataset/optimize.rs +++ b/rust/lance/src/dataset/optimize.rs @@ -324,12 +324,13 @@ pub struct CompactionOptions { /// repack rewrites only the columns that move and keeps rows, fragment ids /// and indices as they are (see [`CompactionTaskKind::RepackColumns`]). /// The count is [`FragmentColumnLayoutStats::live_file_count`]. The - /// columns of one file stay where they are (unless the limit is 1 and - /// every column can move), and blob columns never move, so a fragment can - /// stay above the limit; so can one with more `column_groups` than the - /// limit allows. The repack - /// writes each new file whole, with no `max_bytes_per_file` split. Not - /// planned under `ForceBinaryCopy`, since a repack reencodes. + /// columns of one file stay where they are (unless the limit is 1, every + /// column can move and no file holds spilled row lineage), and a file + /// holding a blob column, a column only partly in the fragment's files, + /// or spilled row lineage is never emptied, so a fragment can stay above + /// the limit; so can one with more `column_groups` than the limit allows. + /// The repack writes each new file whole, with no `max_bytes_per_file` + /// split. Not planned under `ForceBinaryCopy`, since a repack reencodes. /// /// Defaults to `None` (no file-count trigger). Must be at least 1. /// diff --git a/rust/lance/src/dataset/optimize/repack.rs b/rust/lance/src/dataset/optimize/repack.rs index dfe770818e2..cd33b2fc99c 100644 --- a/rust/lance/src/dataset/optimize/repack.rs +++ b/rust/lance/src/dataset/optimize/repack.rs @@ -63,11 +63,12 @@ fn file_coverage(file: &DataFile, schema: &Schema) -> HashSet { /// columns; moving it out leaves the other columns where they were. The /// columns no group names are merged into one new file only when the /// fragment is over `max_files`. Without groups and over `max_files`, the -/// columns of one live file stay where they are (the file holding the -/// fragment's spilled row lineage, else one holding a column that cannot -/// move, else the largest) and every other column that can move goes to one -/// new file; with `max_files` 1, every column movable and no spilled lineage, -/// all of them go to one file. +/// files whose columns can all move, other than the one kept (the largest +/// when every file can be emptied), are merged +/// into one new file; a file holding a column that cannot move, or spilled +/// lineage, is left whole and is the file kept, so a repack never adds a +/// file. With `max_files` 1, every column movable and no file holding spilled +/// lineage, all columns go to one file. /// /// A column the fragment has no data for (added as all nulls) is left out of /// its group, and a group holding a blob column, or a column whose fields are @@ -170,25 +171,48 @@ pub(super) fn plan_fragment_repack( vec![present.clone()] } None => { - let kept = live + // The files a merge can empty: every column moves, no lineage. + let mergeable: Vec<(usize, &BTreeSet)> = live .iter() .zip(&file_columns) - .max_by_key(|((file, _), columns)| { - ( - holds_lineage(file), - columns.iter().any(|column| !movable.contains(column)), - file.file_size_bytes.get().map(|size| size.get()), - columns.len(), - ) + .enumerate() + .filter(|(_, ((file, _), columns))| { + !columns.is_empty() + && !holds_lineage(file) + && columns.iter().all(|column| movable.contains(column)) }) - .map(|(_, columns)| columns.clone()) - .unwrap_or_default(); - let rest = present - .difference(&kept) - .copied() - .filter(|column| movable.contains(column)) + .map(|(index, (_, columns))| (index, columns)) .collect(); - vec![kept, rest] + // A file that cannot be emptied stays anyway, so every mergeable + // file merges; otherwise the largest one stays. + let kept = if mergeable.len() < live.len() { + None + } else { + mergeable + .iter() + .max_by_key(|(index, columns)| { + let file = live[*index].0; + ( + file.file_size_bytes.get().map(|size| size.get()), + columns.len(), + ) + }) + .map(|(index, _)| *index) + }; + let rest: BTreeSet = mergeable + .iter() + .filter(|(index, _)| Some(*index) != kept) + .flat_map(|(_, columns)| columns.iter().copied()) + .collect(); + let rest_files = mergeable + .iter() + .filter(|(index, _)| Some(*index) != kept) + .count(); + if rest_files > 1 { + vec![rest] + } else { + Vec::new() + } } }; diff --git a/rust/lance/src/dataset/optimize/tests/repack.rs b/rust/lance/src/dataset/optimize/tests/repack.rs index a7a7ebe5a73..25f34670b39 100644 --- a/rust/lance/src/dataset/optimize/tests/repack.rs +++ b/rust/lance/src/dataset/optimize/tests/repack.rs @@ -782,8 +782,9 @@ async fn commit_rejects_a_fragment_in_both_kinds_of_task() { assert_eq!(dataset.manifest.version, version); } -/// A column that cannot move outside the kept file stays where it is; the -/// columns that can move are still merged. +/// A blob column outside the lineage file stays where it is; the columns +/// that can move are still merged. The lineage is what makes `lineage.lance` +/// the file kept. #[test] fn repack_merges_around_a_blob_outside_the_kept_file() { use lance_table::format::{ROW_ID_FIELD_ID, RowIdMeta}; @@ -820,3 +821,71 @@ fn repack_merges_around_a_blob_outside_the_kept_file() { ); } } + +/// A file holding a column that cannot move is never partly emptied: moving +/// its other columns out would add a file instead of removing one. +#[test] +fn repack_never_adds_a_file() { + use lance_table::format::{ROW_ID_FIELD_ID, RowIdMeta}; + let blob = |name: &str| { + let mut field = Field::new(name, DataType::LargeBinary, true); + field.set_metadata(std::collections::HashMap::from([( + lance_arrow::BLOB_META_KEY.to_string(), + "true".to_string(), + )])); + field + }; + let plan = |schema: &lance_core::datatypes::Schema, fragment: &Fragment, max| { + crate::dataset::optimize::repack::plan_fragment_repack( + schema, + fragment, + fragment.files.len(), + None, + Some(max), + ) + }; + + // Lineage in one file, a blob next to a movable column in the other. + let schema = lance_schema(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + blob("img"), + Field::new("d", DataType::Int32, true), + ])); + let mut fragment = Fragment::new(0); + fragment.files = vec![ + v2_0_file("lineage.lance", vec![0, ROW_ID_FIELD_ID]), + v2_0_file("mixed.lance", vec![1, 2]), + ]; + fragment.row_id_meta = Some(RowIdMeta::Column); + assert_eq!(plan(&schema, &fragment, 1), None); + + // Two files, each a blob next to a movable column. + let schema = lance_schema(Schema::new(vec![ + Field::new("id", DataType::Int32, true), + blob("img1"), + blob("img2"), + Field::new("cap", DataType::Utf8, true), + ])); + let mut fragment = Fragment::new(0); + fragment.files = vec![ + v2_0_file("base.lance", vec![0, 1]), + v2_0_file("extra.lance", vec![2, 3]), + ]; + assert_eq!(plan(&schema, &fragment, 1), None); + + // A blob next to a movable column, and two movable files: only those two + // merge. + let schema = lance_schema(Schema::new(vec![ + blob("img"), + Field::new("d", DataType::Int32, true), + Field::new("e", DataType::Int32, true), + Field::new("f", DataType::Int32, true), + ])); + let mut fragment = Fragment::new(0); + fragment.files = vec![ + v2_0_file("mixed.lance", vec![0, 1]), + v2_0_file("e.lance", vec![2]), + v2_0_file("f.lance", vec![3]), + ]; + assert_eq!(plan(&schema, &fragment, 2), Some(vec![vec![2, 3]])); +} From 8c3a2f194f96163df20eb57bf85539f92a270464 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Sat, 3 Oct 2026 01:40:26 +0800 Subject: [PATCH 15/24] fix(compaction): count struct headers when deciding which files a repack empties --- rust/lance/src/dataset/optimize.rs | 12 +- rust/lance/src/dataset/optimize/repack.rs | 133 +++++++++++------- .../src/dataset/optimize/tests/repack.rs | 125 +++++++++++++++- 3 files changed, 212 insertions(+), 58 deletions(-) diff --git a/rust/lance/src/dataset/optimize.rs b/rust/lance/src/dataset/optimize.rs index 20639c9773d..5348e83b707 100644 --- a/rust/lance/src/dataset/optimize.rs +++ b/rust/lance/src/dataset/optimize.rs @@ -323,12 +323,12 @@ pub struct CompactionOptions { /// for a rewrite, so this is the trigger that collapses those files. The /// repack rewrites only the columns that move and keeps rows, fragment ids /// and indices as they are (see [`CompactionTaskKind::RepackColumns`]). - /// The count is [`FragmentColumnLayoutStats::live_file_count`]. The - /// columns of one file stay where they are (unless the limit is 1, every - /// column can move and no file holds spilled row lineage), and a file - /// holding a blob column, a column only partly in the fragment's files, - /// or spilled row lineage is never emptied, so a fragment can stay above - /// the limit; so can one with more `column_groups` than the limit allows. + /// The count is [`FragmentColumnLayoutStats::live_file_count`]. Only + /// files the repack empties are merged: a file holding a blob column, a + /// column only partly in the fragment's files, or spilled row lineage + /// stays as it is, and when every file could go, the largest stays unless + /// the limit is 1. So a fragment can stay above the limit; so can one + /// with more `column_groups` than the limit allows. /// The repack writes each new file whole, with no `max_bytes_per_file` /// split. Not planned under `ForceBinaryCopy`, since a repack reencodes. /// diff --git a/rust/lance/src/dataset/optimize/repack.rs b/rust/lance/src/dataset/optimize/repack.rs index cd33b2fc99c..dfbde2b9bdb 100644 --- a/rust/lance/src/dataset/optimize/repack.rs +++ b/rust/lance/src/dataset/optimize/repack.rs @@ -62,13 +62,15 @@ fn file_coverage(file: &DataFile, schema: &Schema) -> HashSet { /// file of their own. A group is moved unless one live file holds exactly its /// columns; moving it out leaves the other columns where they were. The /// columns no group names are merged into one new file only when the -/// fragment is over `max_files`. Without groups and over `max_files`, the -/// files whose columns can all move, other than the one kept (the largest -/// when every file can be emptied), are merged -/// into one new file; a file holding a column that cannot move, or spilled -/// lineage, is left whole and is the file kept, so a repack never adds a -/// file. With `max_files` 1, every column movable and no file holding spilled -/// lineage, all columns go to one file. +/// fragment is over `max_files`. +/// +/// Without groups and over `max_files`, only files the move empties are +/// merged: files whose columns can all move, which hold no spilled lineage, +/// and whose other fields (a struct header) belong to columns that move too. +/// Any other file stays as it is. When every live file can be emptied, one +/// stays as well, the largest by recorded size that still leaves two to +/// merge, unless `max_files` is 1. Nothing is planned unless at least two +/// files would be emptied, so the file count drops. /// /// A column the fragment has no data for (added as all nulls) is left out of /// its group, and a group holding a blob column, or a column whose fields are @@ -120,6 +122,21 @@ pub(super) fn plan_fragment_repack( (field.id, ids) }) .collect(); + // Every schema field id under each top-level column. + let column_fields: std::collections::HashMap> = schema + .fields + .iter() + .map(|field| { + let ids = Schema { + fields: vec![field.clone()], + metadata: Default::default(), + } + .field_ids() + .into_iter() + .collect(); + (field.id, ids) + }) + .collect(); let present: BTreeSet = column_leaves .iter() .filter(|(_, leaves)| !leaves.is_disjoint(&covered)) @@ -146,6 +163,42 @@ pub(super) fn plan_fragment_repack( }) .collect(); let holds_lineage = |file: &DataFile| file.fields.iter().any(|id| spilled.contains(id)); + // The live files whose columns can all move and which hold no lineage. + let candidates: Vec = (0..live.len()) + .filter(|index| { + let columns = &file_columns[*index]; + !columns.is_empty() + && !holds_lineage(live[*index].0) + && columns.iter().all(|column| movable.contains(column)) + }) + .collect(); + // Of `files`, the ones a move of all their columns together leaves with no + // schema field, so the commit drops them. A file can also hold a field of + // a column it holds no data for (a struct header), which only moving that + // column from another file takes out. + let emptied = |mut files: Vec| loop { + let moved: HashSet = files + .iter() + .flat_map(|index| &file_columns[*index]) + .flat_map(|column| &column_fields[column]) + .copied() + .collect(); + let next: Vec = files + .iter() + .copied() + .filter(|index| { + live[*index] + .0 + .fields + .iter() + .all(|id| !schema_ids.contains(id) || moved.contains(id)) + }) + .collect(); + if next.len() == files.len() { + break files; + } + files = next; + }; let wanted: Vec> = match groups { Some(groups) => { @@ -164,52 +217,34 @@ pub(super) fn plan_fragment_repack( } wanted } - None if max_files == Some(1) - && present.iter().all(|c| movable.contains(c)) - && !live.iter().any(|(file, _)| holds_lineage(file)) => - { - vec![present.clone()] - } None => { - // The files a merge can empty: every column moves, no lineage. - let mergeable: Vec<(usize, &BTreeSet)> = live - .iter() - .zip(&file_columns) - .enumerate() - .filter(|(_, ((file, _), columns))| { - !columns.is_empty() - && !holds_lineage(file) - && columns.iter().all(|column| movable.contains(column)) - }) - .map(|(index, (_, columns))| (index, columns)) - .collect(); - // A file that cannot be emptied stays anyway, so every mergeable - // file merges; otherwise the largest one stays. - let kept = if mergeable.len() < live.len() { - None + let all = emptied(candidates); + // A file a move cannot empty stays anyway. Otherwise one file + // stays unless the limit is 1: the largest whose staying still + // leaves two files to merge. + let merged = if all.len() < live.len() || max_files == Some(1) { + all } else { - mergeable + let mut by_size = all.clone(); + by_size.sort_by_key(|index| { + std::cmp::Reverse(( + live[*index].0.file_size_bytes.get().map(|size| size.get()), + file_columns[*index].len(), + )) + }); + by_size .iter() - .max_by_key(|(index, columns)| { - let file = live[*index].0; - ( - file.file_size_bytes.get().map(|size| size.get()), - columns.len(), - ) - }) - .map(|(index, _)| *index) + .map(|kept| emptied(all.iter().copied().filter(|i| i != kept).collect())) + .find(|merged| merged.len() > 1) + .unwrap_or_default() }; - let rest: BTreeSet = mergeable - .iter() - .filter(|(index, _)| Some(*index) != kept) - .flat_map(|(_, columns)| columns.iter().copied()) - .collect(); - let rest_files = mergeable - .iter() - .filter(|(index, _)| Some(*index) != kept) - .count(); - if rest_files > 1 { - vec![rest] + if merged.len() > 1 { + vec![ + merged + .iter() + .flat_map(|index| file_columns[*index].iter().copied()) + .collect(), + ] } else { Vec::new() } diff --git a/rust/lance/src/dataset/optimize/tests/repack.rs b/rust/lance/src/dataset/optimize/tests/repack.rs index 25f34670b39..df34a600928 100644 --- a/rust/lance/src/dataset/optimize/tests/repack.rs +++ b/rust/lance/src/dataset/optimize/tests/repack.rs @@ -782,9 +782,8 @@ async fn commit_rejects_a_fragment_in_both_kinds_of_task() { assert_eq!(dataset.manifest.version, version); } -/// A blob column outside the lineage file stays where it is; the columns -/// that can move are still merged. The lineage is what makes `lineage.lance` -/// the file kept. +/// The lineage file and a blob file both stay as they are; the columns that +/// can move are still merged. #[test] fn repack_merges_around_a_blob_outside_the_kept_file() { use lance_table::format::{ROW_ID_FIELD_ID, RowIdMeta}; @@ -889,3 +888,123 @@ fn repack_never_adds_a_file() { ]; assert_eq!(plan(&schema, &fragment, 2), Some(vec![vec![2, 3]])); } + +/// A file holding a struct's header next to another column is emptied only +/// when the struct moves too, so a merge is planned around the file whose +/// staying leaves the others emptied. +#[test] +fn repack_moves_a_struct_with_its_header() { + use arrow_schema::Fields; + let plan = |schema: &lance_core::datatypes::Schema, fragment: &Fragment| { + crate::dataset::optimize::repack::plan_fragment_repack( + schema, + fragment, + fragment.files.len(), + None, + Some(2), + ) + }; + let child = |name: &str| Field::new(name, DataType::Int32, true); + + // a=0, s=1 {x=2 (dropped), y=3}, b=4, c=5. + let arrow = Schema::new(vec![ + child("a"), + Field::new( + "s", + DataType::Struct(Fields::from(vec![child("x"), child("y")])), + true, + ), + child("b"), + child("c"), + ]); + let schema = lance_schema(arrow).project_by_ids(&[0, 1, 3, 4, 5], false); + let mut fragment = Fragment::new(0); + fragment.files = vec![ + v2_0_file("a.lance", vec![0, 1, 2]), + v2_0_file("s.lance", vec![1, 3, 5]), + v2_0_file("b.lance", vec![4]), + ]; + // Keeping s.lance would leave a.lance alive on the header; keeping + // a.lance lets s.lance and b.lance both go. + assert_eq!(plan(&schema, &fragment), Some(vec![vec![1, 4, 5]])); + + // a=0, s=1 {x=2, y=3, z=4} (x, z dropped), b=5, c=6: both a.lance and + // b.lance hold the header. + let arrow = Schema::new(vec![ + child("a"), + Field::new( + "s", + DataType::Struct(Fields::from(vec![child("x"), child("y"), child("z")])), + true, + ), + child("b"), + child("c"), + ]); + let schema = lance_schema(arrow).project_by_ids(&[0, 1, 3, 5, 6], false); + let mut fragment = Fragment::new(0); + fragment.files = vec![ + v2_0_file("a.lance", vec![0, 1, 2]), + v2_0_file("b.lance", vec![5, 1, 4]), + v2_0_file("s.lance", vec![1, 3, 6]), + ]; + assert_eq!(plan(&schema, &fragment), Some(vec![vec![1, 5, 6]])); +} + +/// A struct-header file merged next to the struct reaches one file, and the +/// next run plans nothing. +#[tokio::test] +async fn repack_with_a_dropped_struct_child_converges() { + use arrow_array::{ArrayRef, StructArray}; + use arrow_schema::Fields; + let fields = Fields::from(vec![ + Field::new("x", DataType::Int32, true), + Field::new("y", DataType::Int32, true), + ]); + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("s", DataType::Struct(fields.clone()), true), + ])); + let values = + |offset: i32| Arc::new(Int32Array::from_iter_values(offset..offset + 4)) as ArrayRef; + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + values(0), + Arc::new(StructArray::new(fields, vec![values(10), values(20)], None)), + ], + ) + .unwrap(); + let mut dataset = Dataset::write( + RecordBatchIterator::new([Ok(batch)], schema), + "memory://", + Some(WriteParams { + data_storage_version: Some(LanceFileVersion::V2_0), + ..Default::default() + }), + ) + .await + .unwrap(); + dataset + .add_columns( + NewColumnTransform::SqlExpressions(vec![("b".into(), "a + 1".into())]), + None, + None, + ) + .await + .unwrap(); + dataset.drop_columns(&["s.x"]).await.unwrap(); + let before = dataset.scan().try_into_batch().await.unwrap(); + + compact_files(&mut dataset, repack_options(Some(1), vec![]), None) + .await + .unwrap(); + assert_eq!(live_file_counts(&dataset), vec![1]); + assert_eq!(dataset.scan().try_into_batch().await.unwrap(), before); + dataset.validate().await.unwrap(); + + let version = dataset.manifest.version; + compact_files(&mut dataset, repack_options(Some(1), vec![]), None) + .await + .unwrap(); + assert_eq!(dataset.manifest.version, version); +} From 513a52ba0b2e9be8e8fb173568e6b51950e4f800 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Sat, 3 Oct 2026 02:04:34 +0800 Subject: [PATCH 16/24] fix(compaction): apply the empties rule to group merges and leave kept and lineage files alone --- rust/lance/src/dataset/optimize.rs | 10 +- rust/lance/src/dataset/optimize/repack.rs | 99 ++++++++++++------- .../src/dataset/optimize/tests/repack.rs | 97 +++++++++++++++++- 3 files changed, 164 insertions(+), 42 deletions(-) diff --git a/rust/lance/src/dataset/optimize.rs b/rust/lance/src/dataset/optimize.rs index 5348e83b707..04b62dfb4a6 100644 --- a/rust/lance/src/dataset/optimize.rs +++ b/rust/lance/src/dataset/optimize.rs @@ -324,11 +324,11 @@ pub struct CompactionOptions { /// repack rewrites only the columns that move and keeps rows, fragment ids /// and indices as they are (see [`CompactionTaskKind::RepackColumns`]). /// The count is [`FragmentColumnLayoutStats::live_file_count`]. Only - /// files the repack empties are merged: a file holding a blob column, a - /// column only partly in the fragment's files, or spilled row lineage - /// stays as it is, and when every file could go, the largest stays unless - /// the limit is 1. So a fragment can stay above the limit; so can one - /// with more `column_groups` than the limit allows. + /// files the repack empties are merged, at least two at a time: a file + /// holding a blob column, a column only partly in the fragment's files, + /// or spilled row lineage stays, and when every file could go, one is + /// left untouched unless the limit is 1. So a fragment can stay above the + /// limit; so can one with more `column_groups` than the limit allows. /// The repack writes each new file whole, with no `max_bytes_per_file` /// split. Not planned under `ForceBinaryCopy`, since a repack reencodes. /// diff --git a/rust/lance/src/dataset/optimize/repack.rs b/rust/lance/src/dataset/optimize/repack.rs index dfbde2b9bdb..c332d593b6a 100644 --- a/rust/lance/src/dataset/optimize/repack.rs +++ b/rust/lance/src/dataset/optimize/repack.rs @@ -12,7 +12,7 @@ //! an `Operation::DataReplacement` with `data_change: false`: the fragment id, //! row addresses, overlays and index coverage are left as they are. -use std::collections::{BTreeSet, HashSet}; +use std::collections::{BTreeSet, HashMap, HashSet}; use std::sync::Arc; use lance_core::datatypes::{Field as LanceField, Schema}; @@ -28,6 +28,14 @@ use crate::dataset::transaction::{DataReplacementGroup, Operation, TransactionBu use crate::dataset::{Dataset, cleanup_data_fragments, versions}; use crate::{Error, Result}; +/// The ids of `field` and every field beneath it. +fn subtree_ids(field: &LanceField, ids: &mut HashSet) { + ids.insert(field.id); + for child in &field.children { + subtree_ids(child, ids); + } +} + /// The ids of the leaf fields at and beneath `field`. The planner places a /// column by where its leaves are: a V2.0 file can keep a struct's header /// after the struct's last child in it was dropped. @@ -60,17 +68,18 @@ fn file_coverage(file: &DataFile, schema: &Schema) -> HashSet { /// /// `groups` are the column groups, each the top-level field ids to keep in a /// file of their own. A group is moved unless one live file holds exactly its -/// columns; moving it out leaves the other columns where they were. The -/// columns no group names are merged into one new file only when the -/// fragment is over `max_files`. +/// columns; moving it out leaves the other columns where they were. Over +/// `max_files`, the files holding only columns no group names are merged as +/// well, under the rule below. /// -/// Without groups and over `max_files`, only files the move empties are -/// merged: files whose columns can all move, which hold no spilled lineage, +/// A merge only takes files it empties: files whose columns can all move, +/// which hold no spilled lineage and share no column with a file holding it, /// and whose other fields (a struct header) belong to columns that move too. -/// Any other file stays as it is. When every live file can be emptied, one -/// stays as well, the largest by recorded size that still leaves two to -/// merge, unless `max_files` is 1. Nothing is planned unless at least two -/// files would be emptied, so the file count drops. +/// The files left keep their other columns. Without groups, when every live +/// file can be emptied, one is left untouched as well, the largest by +/// recorded size that still leaves two to merge, unless `max_files` is 1. +/// Nothing is merged unless at least two files would be emptied, so a merge +/// lowers the file count. /// /// A column the fragment has no data for (added as all nulls) is left out of /// its group, and a group holding a blob column, or a column whose fields are @@ -81,8 +90,9 @@ fn file_coverage(file: &DataFile, schema: &Schema) -> HashSet { /// /// `live_files` is the fragment's /// [`FragmentColumnLayoutStats::live_file_count`]. A V2.0 file left holding -/// only a struct header counts there but holds no column here, so it stays -/// until a rewrite of the fragment. +/// only a struct header counts there but holds no column here: it goes when +/// a repack moves the struct, and otherwise stays until a rewrite of the +/// fragment. /// /// [`FragmentColumnLayoutStats::live_file_count`]: crate::dataset::compaction_stats::FragmentColumnLayoutStats::live_file_count pub(super) fn plan_fragment_repack( @@ -123,17 +133,12 @@ pub(super) fn plan_fragment_repack( }) .collect(); // Every schema field id under each top-level column. - let column_fields: std::collections::HashMap> = schema + let column_fields: HashMap> = schema .fields .iter() .map(|field| { - let ids = Schema { - fields: vec![field.clone()], - metadata: Default::default(), - } - .field_ids() - .into_iter() - .collect(); + let mut ids = HashSet::new(); + subtree_ids(field, &mut ids); (field.id, ids) }) .collect(); @@ -163,15 +168,29 @@ pub(super) fn plan_fragment_repack( }) .collect(); let holds_lineage = |file: &DataFile| file.fields.iter().any(|id| spilled.contains(id)); - // The live files whose columns can all move and which hold no lineage. + let lineage_columns: BTreeSet = live + .iter() + .zip(&file_columns) + .filter(|((file, _), _)| holds_lineage(file)) + .flat_map(|(_, columns)| columns.iter().copied()) + .collect(); + // The live files a merge may take: every column can move, and moving them + // takes nothing out of a file holding spilled lineage. let candidates: Vec = (0..live.len()) .filter(|index| { let columns = &file_columns[*index]; !columns.is_empty() && !holds_lineage(live[*index].0) + && columns.is_disjoint(&lineage_columns) && columns.iter().all(|column| movable.contains(column)) }) .collect(); + let columns_of = |files: &[usize]| -> BTreeSet { + files + .iter() + .flat_map(|index| file_columns[*index].iter().copied()) + .collect() + }; // Of `files`, the ones a move of all their columns together leaves with no // schema field, so the commit drops them. A file can also hold a field of // a column it holds no data for (a struct header), which only moving that @@ -207,13 +226,17 @@ pub(super) fn plan_fragment_repack( .map(|group| group.iter().copied().collect()) .collect(); let claimed: BTreeSet = wanted.iter().flatten().copied().collect(); - let rest: BTreeSet = present.difference(&claimed).copied().collect(); - let rest_files = file_columns - .iter() - .filter(|columns| !columns.is_disjoint(&rest)) - .count(); - if over_file_limit && rest_files > 1 { - wanted.push(rest); + if over_file_limit { + let merged = emptied( + candidates + .iter() + .copied() + .filter(|index| file_columns[*index].is_disjoint(&claimed)) + .collect(), + ); + if merged.len() > 1 { + wanted.push(columns_of(&merged)); + } } wanted } @@ -232,19 +255,25 @@ pub(super) fn plan_fragment_repack( file_columns[*index].len(), )) }); + // The file kept shares no column with the merge, so it is + // left untouched. by_size .iter() - .map(|kept| emptied(all.iter().copied().filter(|i| i != kept).collect())) + .map(|kept| { + emptied( + all.iter() + .copied() + .filter(|index| { + file_columns[*index].is_disjoint(&file_columns[*kept]) + }) + .collect(), + ) + }) .find(|merged| merged.len() > 1) .unwrap_or_default() }; if merged.len() > 1 { - vec![ - merged - .iter() - .flat_map(|index| file_columns[*index].iter().copied()) - .collect(), - ] + vec![columns_of(&merged)] } else { Vec::new() } diff --git a/rust/lance/src/dataset/optimize/tests/repack.rs b/rust/lance/src/dataset/optimize/tests/repack.rs index df34a600928..4fb97a24087 100644 --- a/rust/lance/src/dataset/optimize/tests/repack.rs +++ b/rust/lance/src/dataset/optimize/tests/repack.rs @@ -950,8 +950,8 @@ fn repack_moves_a_struct_with_its_header() { assert_eq!(plan(&schema, &fragment), Some(vec![vec![1, 5, 6]])); } -/// A struct-header file merged next to the struct reaches one file, and the -/// next run plans nothing. +/// A fragment whose first file still holds a dropped struct child reaches one +/// file, and the next run plans nothing. #[tokio::test] async fn repack_with_a_dropped_struct_child_converges() { use arrow_array::{ArrayRef, StructArray}; @@ -1008,3 +1008,96 @@ async fn repack_with_a_dropped_struct_child_converges() { .unwrap(); assert_eq!(dataset.manifest.version, version); } + +/// Struct-header layouts for the merges the other planner tests don't reach: +/// the unclaimed columns under `column_groups`, a struct split with the +/// lineage file, and the file kept when every file could go. +#[test] +fn repack_merges_only_files_it_empties() { + use arrow_schema::Fields; + use lance_table::format::{ROW_ID_FIELD_ID, RowIdMeta}; + let child = |name: &str| Field::new(name, DataType::Int32, true); + let plan = |schema: &lance_core::datatypes::Schema, + fragment: &Fragment, + groups: Option<&[Vec]>, + max| { + crate::dataset::optimize::repack::plan_fragment_repack( + schema, + fragment, + fragment.files.len(), + groups, + Some(max), + ) + }; + + // a=0, s=1 {x=2 (dropped), y=3, z=4 (dropped)}, b=5, groups [[s]]. + // a.lance and b.lance both keep s's header, so merging a and b would + // drop neither. + let arrow = Schema::new(vec![ + child("a"), + Field::new( + "s", + DataType::Struct(Fields::from(vec![child("x"), child("y"), child("z")])), + true, + ), + child("b"), + ]); + let schema = lance_schema(arrow).project_by_ids(&[0, 1, 3, 5], false); + let mut fragment = Fragment::new(0); + fragment.files = vec![ + v2_0_file("a.lance", vec![0, 1, 2]), + v2_0_file("b.lance", vec![5, 1, 4]), + v2_0_file("s.lance", vec![1, 3]), + ]; + assert_eq!(plan(&schema, &fragment, Some(&[vec![1]]), 2), None); + + // s=0 {y=1, z=2}, b=3, c=4. s is split between the lineage file and + // c1.lance; moving s would empty the lineage file's columns, so only b + // and c merge. + let arrow = Schema::new(vec![ + Field::new( + "s", + DataType::Struct(Fields::from(vec![child("y"), child("z")])), + true, + ), + child("b"), + child("c"), + ]); + let schema = lance_schema(arrow); + let mut fragment = Fragment::new(0); + fragment.files = vec![ + v2_0_file("lineage.lance", vec![0, 1, ROW_ID_FIELD_ID]), + v2_0_file("c1.lance", vec![0, 2]), + v2_0_file("b.lance", vec![3]), + v2_0_file("c.lance", vec![4]), + ]; + fragment.row_id_meta = Some(RowIdMeta::Column); + assert_eq!(plan(&schema, &fragment, None, 2), Some(vec![vec![3, 4]])); + + // a=0, s=1 {x=2, y=3}, b=4, c=5. k.lance is the largest and shares s with + // m1.lance; keeping k leaves s and m1 alone and merges b and c. + let arrow = Schema::new(vec![ + child("a"), + Field::new( + "s", + DataType::Struct(Fields::from(vec![child("x"), child("y")])), + true, + ), + child("b"), + child("c"), + ]); + let schema = lance_schema(arrow); + let sized = |path: &str, fields: Vec, size: u64| { + let mut file = v2_0_file(path, fields); + file.file_size_bytes = lance_io::utils::CachedFileSize::new(size); + file + }; + let mut fragment = Fragment::new(0); + fragment.files = vec![ + sized("k.lance", vec![1, 2], 100), + sized("m1.lance", vec![0, 1, 3], 10), + sized("m2.lance", vec![4], 10), + sized("m3.lance", vec![5], 10), + ]; + assert_eq!(plan(&schema, &fragment, None, 2), Some(vec![vec![4, 5]])); +} From 64c13f8034f7d3ff79751f7bf7b3abc41b0acf0e Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Sat, 3 Oct 2026 02:35:45 +0800 Subject: [PATCH 17/24] fix(compaction): fall back to merging every other file when the kept file shares a struct --- rust/lance/src/dataset/optimize.rs | 11 +-- rust/lance/src/dataset/optimize/repack.rs | 39 ++++---- .../src/dataset/optimize/tests/repack.rs | 89 +++++++++++++++++++ 3 files changed, 119 insertions(+), 20 deletions(-) diff --git a/rust/lance/src/dataset/optimize.rs b/rust/lance/src/dataset/optimize.rs index 04b62dfb4a6..49d425b00d3 100644 --- a/rust/lance/src/dataset/optimize.rs +++ b/rust/lance/src/dataset/optimize.rs @@ -324,11 +324,12 @@ pub struct CompactionOptions { /// repack rewrites only the columns that move and keeps rows, fragment ids /// and indices as they are (see [`CompactionTaskKind::RepackColumns`]). /// The count is [`FragmentColumnLayoutStats::live_file_count`]. Only - /// files the repack empties are merged, at least two at a time: a file - /// holding a blob column, a column only partly in the fragment's files, - /// or spilled row lineage stays, and when every file could go, one is - /// left untouched unless the limit is 1. So a fragment can stay above the - /// limit; so can one with more `column_groups` than the limit allows. + /// files the repack empties are merged, at least two at a time. A file + /// stays when it holds a blob column, a column only partly in the + /// fragment's files, or spilled row lineage, or shares a column with such + /// a file; without `column_groups`, when every file could go, the largest + /// stays unless the limit is 1. So a fragment can stay above the limit; + /// so can one with more `column_groups` than the limit allows. /// The repack writes each new file whole, with no `max_bytes_per_file` /// split. Not planned under `ForceBinaryCopy`, since a repack reencodes. /// diff --git a/rust/lance/src/dataset/optimize/repack.rs b/rust/lance/src/dataset/optimize/repack.rs index c332d593b6a..4baf7441f8d 100644 --- a/rust/lance/src/dataset/optimize/repack.rs +++ b/rust/lance/src/dataset/optimize/repack.rs @@ -75,11 +75,13 @@ fn file_coverage(file: &DataFile, schema: &Schema) -> HashSet { /// A merge only takes files it empties: files whose columns can all move, /// which hold no spilled lineage and share no column with a file holding it, /// and whose other fields (a struct header) belong to columns that move too. -/// The files left keep their other columns. Without groups, when every live -/// file can be emptied, one is left untouched as well, the largest by -/// recorded size that still leaves two to merge, unless `max_files` is 1. -/// Nothing is merged unless at least two files would be emptied, so a merge -/// lowers the file count. +/// The files left keep their columns, though a struct moved out of one +/// leaves a tombstone where its header was. Without groups, when every file +/// holding a column can be emptied, one of them stays unless `max_files` is +/// 1: the largest by recorded size, with the files sharing no column with it +/// merged, or if that leaves fewer than two, every other file. Nothing is +/// merged unless at least two files would be emptied, so a merge lowers the +/// file count. /// /// A column the fragment has no data for (added as all nulls) is left out of /// its group, and a group holding a blob column, or a column whose fields are @@ -242,10 +244,11 @@ pub(super) fn plan_fragment_repack( } None => { let all = emptied(candidates); + let holding = file_columns.iter().filter(|c| !c.is_empty()).count(); // A file a move cannot empty stays anyway. Otherwise one file // stays unless the limit is 1: the largest whose staying still // leaves two files to merge. - let merged = if all.len() < live.len() || max_files == Some(1) { + let merged = if all.len() < holding || max_files == Some(1) { all } else { let mut by_size = all.clone(); @@ -255,21 +258,27 @@ pub(super) fn plan_fragment_repack( file_columns[*index].len(), )) }); - // The file kept shares no column with the merge, so it is - // left untouched. + // Prefer merging only files that share no column with the + // file kept, which leaves it untouched; else merge every + // other file, as long as the kept file keeps a column. by_size .iter() - .map(|kept| { - emptied( + .find_map(|kept| { + let kept_columns = &file_columns[*kept]; + let untouched = emptied( all.iter() .copied() - .filter(|index| { - file_columns[*index].is_disjoint(&file_columns[*kept]) - }) + .filter(|index| file_columns[*index].is_disjoint(kept_columns)) .collect(), - ) + ); + if untouched.len() > 1 { + return Some(untouched); + } + let others = + emptied(all.iter().copied().filter(|index| index != kept).collect()); + (others.len() > 1 && !kept_columns.is_subset(&columns_of(&others))) + .then_some(others) }) - .find(|merged| merged.len() > 1) .unwrap_or_default() }; if merged.len() > 1 { diff --git a/rust/lance/src/dataset/optimize/tests/repack.rs b/rust/lance/src/dataset/optimize/tests/repack.rs index 4fb97a24087..9f52b502a33 100644 --- a/rust/lance/src/dataset/optimize/tests/repack.rs +++ b/rust/lance/src/dataset/optimize/tests/repack.rs @@ -1101,3 +1101,92 @@ fn repack_merges_only_files_it_empties() { ]; assert_eq!(plan(&schema, &fragment, None, 2), Some(vec![vec![4, 5]])); } + +/// When every file holds the struct, the largest file is kept with its other +/// column and the struct's other files merge; and when only some files share +/// a column with the largest, the merge falls back to every other file. +#[test] +fn repack_keeps_the_largest_file_that_shares_a_struct() { + use arrow_schema::Fields; + let child = |name: &str| Field::new(name, DataType::Int32, true); + let sized = |path: &str, fields: Vec, size: u64| { + let mut file = v2_0_file(path, fields); + file.file_size_bytes = lance_io::utils::CachedFileSize::new(size); + file + }; + let plan = |schema: &lance_core::datatypes::Schema, fragment: &Fragment| { + crate::dataset::optimize::repack::plan_fragment_repack( + schema, + fragment, + fragment.files.len(), + None, + Some(2), + ) + }; + // a=0, s=1 {x=2, y=3, z=4}, b=5. + let schema = lance_schema(Schema::new(vec![ + child("a"), + Field::new( + "s", + DataType::Struct(Fields::from(vec![child("x"), child("y"), child("z")])), + true, + ), + child("b"), + ])); + + let mut fragment = Fragment::new(0); + fragment.files = vec![ + sized("k.lance", vec![0, 1, 2], 100), + sized("m1.lance", vec![1, 3], 10), + sized("m2.lance", vec![1, 4], 10), + ]; + assert_eq!(plan(&schema, &fragment), Some(vec![vec![1]])); + + fragment.files = vec![ + sized("k.lance", vec![0, 1, 2], 100), + sized("m1.lance", vec![1, 3, 4], 10), + sized("m2.lance", vec![5], 10), + ]; + assert_eq!(plan(&schema, &fragment), Some(vec![vec![1, 5]])); +} + +/// A V2.0 file holding only a struct header holds no column, so it does not +/// stop the largest file from being kept; it goes when the struct moves. +#[test] +fn repack_keeps_the_largest_file_next_to_a_header_only_file() { + use arrow_schema::Fields; + let child = |name: &str| Field::new(name, DataType::Int32, true); + let sized = |path: &str, fields: Vec, size: u64| { + let mut file = v2_0_file(path, fields); + file.file_size_bytes = lance_io::utils::CachedFileSize::new(size); + file + }; + // a=0, s=1 {x=2 (dropped), y=3}, b=4. + let schema = lance_schema(Schema::new(vec![ + child("a"), + Field::new( + "s", + DataType::Struct(Fields::from(vec![child("x"), child("y")])), + true, + ), + child("b"), + ])) + .project_by_ids(&[0, 1, 3, 4], false); + let mut fragment = Fragment::new(0); + fragment.files = vec![ + sized("h.lance", vec![1, 2], 10), + sized("base.lance", vec![0], 100), + sized("sy.lance", vec![1, 3], 10), + sized("bf.lance", vec![4], 10), + ]; + assert_eq!( + crate::dataset::optimize::repack::plan_fragment_repack( + &schema, + &fragment, + 4, + None, + Some(2) + ), + Some(vec![vec![1, 4]]) + ); +} From 715801d85b517b70decf16f9666062c198950487 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Sat, 3 Oct 2026 02:55:45 +0800 Subject: [PATCH 18/24] fix(compaction): prefer a repack that leaves the largest file in place --- rust/lance/src/dataset/optimize.rs | 9 +-- rust/lance/src/dataset/optimize/repack.rs | 62 ++++++++++--------- .../src/dataset/optimize/tests/repack.rs | 45 ++++++++++++++ 3 files changed, 82 insertions(+), 34 deletions(-) diff --git a/rust/lance/src/dataset/optimize.rs b/rust/lance/src/dataset/optimize.rs index 49d425b00d3..7cff8745ce3 100644 --- a/rust/lance/src/dataset/optimize.rs +++ b/rust/lance/src/dataset/optimize.rs @@ -326,10 +326,11 @@ pub struct CompactionOptions { /// The count is [`FragmentColumnLayoutStats::live_file_count`]. Only /// files the repack empties are merged, at least two at a time. A file /// stays when it holds a blob column, a column only partly in the - /// fragment's files, or spilled row lineage, or shares a column with such - /// a file; without `column_groups`, when every file could go, the largest - /// stays unless the limit is 1. So a fragment can stay above the limit; - /// so can one with more `column_groups` than the limit allows. + /// fragment's files, or spilled row lineage, or shares a column with a + /// file holding spilled row lineage; without `column_groups`, when every + /// file holding a column could go, one stays (usually the largest) unless + /// the limit is 1. So a fragment can stay above the limit; so can one with + /// more `column_groups` than the limit allows. /// The repack writes each new file whole, with no `max_bytes_per_file` /// split. Not planned under `ForceBinaryCopy`, since a repack reencodes. /// diff --git a/rust/lance/src/dataset/optimize/repack.rs b/rust/lance/src/dataset/optimize/repack.rs index 4baf7441f8d..6f1c148eb24 100644 --- a/rust/lance/src/dataset/optimize/repack.rs +++ b/rust/lance/src/dataset/optimize/repack.rs @@ -75,13 +75,14 @@ fn file_coverage(file: &DataFile, schema: &Schema) -> HashSet { /// A merge only takes files it empties: files whose columns can all move, /// which hold no spilled lineage and share no column with a file holding it, /// and whose other fields (a struct header) belong to columns that move too. -/// The files left keep their columns, though a struct moved out of one -/// leaves a tombstone where its header was. Without groups, when every file -/// holding a column can be emptied, one of them stays unless `max_files` is -/// 1: the largest by recorded size, with the files sharing no column with it -/// merged, or if that leaves fewer than two, every other file. Nothing is -/// merged unless at least two files would be emptied, so a merge lowers the -/// file count. +/// The files left keep the columns that do not move. Without groups, when +/// every file holding a column can be emptied, one of them stays unless +/// `max_files` is 1. Trying the files in order of recorded size, the one kept +/// is the first whose merge (of the files sharing no column with it, else of +/// every other file, as long as it keeps a column) leaves the largest file +/// in place, or failing that the first with any merge. Nothing is merged +/// unless at least two files would be emptied, so a merge lowers the file +/// count. /// /// A column the fragment has no data for (added as all nulls) is left out of /// its group, and a group holding a blob column, or a column whose fields are @@ -246,8 +247,7 @@ pub(super) fn plan_fragment_repack( let all = emptied(candidates); let holding = file_columns.iter().filter(|c| !c.is_empty()).count(); // A file a move cannot empty stays anyway. Otherwise one file - // stays unless the limit is 1: the largest whose staying still - // leaves two files to merge. + // stays unless the limit is 1. let merged = if all.len() < holding || max_files == Some(1) { all } else { @@ -258,27 +258,29 @@ pub(super) fn plan_fragment_repack( file_columns[*index].len(), )) }); - // Prefer merging only files that share no column with the - // file kept, which leaves it untouched; else merge every - // other file, as long as the kept file keeps a column. - by_size - .iter() - .find_map(|kept| { - let kept_columns = &file_columns[*kept]; - let untouched = emptied( - all.iter() - .copied() - .filter(|index| file_columns[*index].is_disjoint(kept_columns)) - .collect(), - ); - if untouched.len() > 1 { - return Some(untouched); - } - let others = - emptied(all.iter().copied().filter(|index| index != kept).collect()); - (others.len() > 1 && !kept_columns.is_subset(&columns_of(&others))) - .then_some(others) - }) + // For each file kept, largest first: merge the files sharing + // no column with it, which leaves it untouched, else every + // other file, as long as the kept file keeps a column. A plan + // that leaves the largest file in place wins. + let plans = by_size.iter().flat_map(|kept| { + let untouched = emptied( + all.iter() + .copied() + .filter(|index| file_columns[*index].is_disjoint(&file_columns[*kept])) + .collect(), + ); + let others = emptied(all.iter().copied().filter(|i| i != kept).collect()); + let keeps_a_column = !file_columns[*kept].is_subset(&columns_of(&others)); + [Some(untouched), keeps_a_column.then_some(others)] + .into_iter() + .flatten() + .filter(|merged| merged.len() > 1) + }); + let largest = by_size[0]; + plans + .clone() + .find(|merged| !merged.contains(&largest)) + .or_else(|| plans.clone().next()) .unwrap_or_default() }; if merged.len() > 1 { diff --git a/rust/lance/src/dataset/optimize/tests/repack.rs b/rust/lance/src/dataset/optimize/tests/repack.rs index 9f52b502a33..18e08cb8642 100644 --- a/rust/lance/src/dataset/optimize/tests/repack.rs +++ b/rust/lance/src/dataset/optimize/tests/repack.rs @@ -1190,3 +1190,48 @@ fn repack_keeps_the_largest_file_next_to_a_header_only_file() { Some(vec![vec![1, 4]]) ); } + +/// A merge that leaves the largest file untouched wins over one that merges +/// it away, even when a smaller file is the one kept. +#[test] +fn repack_prefers_leaving_the_largest_file_untouched() { + use arrow_schema::Fields; + let child = |name: &str| Field::new(name, DataType::Int32, true); + let sized = |path: &str, fields: Vec, size: u64| { + let mut file = v2_0_file(path, fields); + file.file_size_bytes = lance_io::utils::CachedFileSize::new(size); + file + }; + // s=0 {x=1, y=2, z=3}, t=4 {p=5, q=6}, c=7. + let schema = lance_schema(Schema::new(vec![ + Field::new( + "s", + DataType::Struct(Fields::from(vec![child("x"), child("y"), child("z")])), + true, + ), + Field::new( + "t", + DataType::Struct(Fields::from(vec![child("p"), child("q")])), + true, + ), + child("c"), + ])); + let mut fragment = Fragment::new(0); + fragment.files = vec![ + sized("k.lance", vec![0, 1, 4, 5], 100), + sized("m1.lance", vec![0, 2, 7], 10), + sized("m2.lance", vec![0, 3], 10), + sized("m3.lance", vec![4, 6], 10), + ]; + // Keeping m3 merges m1 and m2 (s and c) and leaves k holding t. + assert_eq!( + crate::dataset::optimize::repack::plan_fragment_repack( + &schema, + &fragment, + 4, + None, + Some(3) + ), + Some(vec![vec![0, 7]]) + ); +} From 3cfb9799ac9bf2140bf8dc5472a1a9410735e772 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Sat, 3 Oct 2026 03:17:16 +0800 Subject: [PATCH 19/24] fix(compaction): plan nothing when no file holds a column --- rust/lance/src/dataset/optimize.rs | 4 +- rust/lance/src/dataset/optimize/repack.rs | 14 +++---- .../src/dataset/optimize/tests/repack.rs | 40 ++++++++++++++++++- 3 files changed, 47 insertions(+), 11 deletions(-) diff --git a/rust/lance/src/dataset/optimize.rs b/rust/lance/src/dataset/optimize.rs index 7cff8745ce3..7bbf68dbe6c 100644 --- a/rust/lance/src/dataset/optimize.rs +++ b/rust/lance/src/dataset/optimize.rs @@ -328,8 +328,8 @@ pub struct CompactionOptions { /// stays when it holds a blob column, a column only partly in the /// fragment's files, or spilled row lineage, or shares a column with a /// file holding spilled row lineage; without `column_groups`, when every - /// file holding a column could go, one stays (usually the largest) unless - /// the limit is 1. So a fragment can stay above the limit; so can one with + /// file holding a column could go, at least one stays (the largest when + /// some merge leaves it) unless the limit is 1. So a fragment can stay above the limit; so can one with /// more `column_groups` than the limit allows. /// The repack writes each new file whole, with no `max_bytes_per_file` /// split. Not planned under `ForceBinaryCopy`, since a repack reencodes. diff --git a/rust/lance/src/dataset/optimize/repack.rs b/rust/lance/src/dataset/optimize/repack.rs index 6f1c148eb24..2cee6e0cff0 100644 --- a/rust/lance/src/dataset/optimize/repack.rs +++ b/rust/lance/src/dataset/optimize/repack.rs @@ -76,11 +76,11 @@ fn file_coverage(file: &DataFile, schema: &Schema) -> HashSet { /// which hold no spilled lineage and share no column with a file holding it, /// and whose other fields (a struct header) belong to columns that move too. /// The files left keep the columns that do not move. Without groups, when -/// every file holding a column can be emptied, one of them stays unless -/// `max_files` is 1. Trying the files in order of recorded size, the one kept -/// is the first whose merge (of the files sharing no column with it, else of -/// every other file, as long as it keeps a column) leaves the largest file -/// in place, or failing that the first with any merge. Nothing is merged +/// every file holding a column can be emptied, at least one of them stays +/// unless `max_files` is 1. Trying the files largest first by recorded size, +/// the one kept is the first whose merge (of the files sharing no column with +/// it, else of every other file, as long as it keeps a column) leaves the +/// largest file in place, or failing that the first with any merge. Nothing is merged /// unless at least two files would be emptied, so a merge lowers the file /// count. /// @@ -248,7 +248,7 @@ pub(super) fn plan_fragment_repack( let holding = file_columns.iter().filter(|c| !c.is_empty()).count(); // A file a move cannot empty stays anyway. Otherwise one file // stays unless the limit is 1. - let merged = if all.len() < holding || max_files == Some(1) { + let merged = if all.len() < 2 || all.len() < holding || max_files == Some(1) { all } else { let mut by_size = all.clone(); @@ -259,7 +259,7 @@ pub(super) fn plan_fragment_repack( )) }); // For each file kept, largest first: merge the files sharing - // no column with it, which leaves it untouched, else every + // no column with it, which leaves its data in place, else every // other file, as long as the kept file keeps a column. A plan // that leaves the largest file in place wins. let plans = by_size.iter().flat_map(|kept| { diff --git a/rust/lance/src/dataset/optimize/tests/repack.rs b/rust/lance/src/dataset/optimize/tests/repack.rs index 18e08cb8642..6c75e82a249 100644 --- a/rust/lance/src/dataset/optimize/tests/repack.rs +++ b/rust/lance/src/dataset/optimize/tests/repack.rs @@ -1191,10 +1191,10 @@ fn repack_keeps_the_largest_file_next_to_a_header_only_file() { ); } -/// A merge that leaves the largest file untouched wins over one that merges +/// A merge that leaves the largest file in place wins over one that merges /// it away, even when a smaller file is the one kept. #[test] -fn repack_prefers_leaving_the_largest_file_untouched() { +fn repack_prefers_leaving_the_largest_file_in_place() { use arrow_schema::Fields; let child = |name: &str| Field::new(name, DataType::Int32, true); let sized = |path: &str, fields: Vec, size: u64| { @@ -1235,3 +1235,39 @@ fn repack_prefers_leaving_the_largest_file_untouched() { Some(vec![vec![0, 7]]) ); } + +/// A fragment whose files all hold only a struct header has no column to +/// move, and nothing is planned. +#[test] +fn repack_plans_nothing_for_header_only_files() { + use arrow_schema::Fields; + let child = |name: &str| Field::new(name, DataType::Int32, true); + // s=0 {a=1, b=2, c=3, d=4}; only d is left. + let schema = lance_schema(Schema::new(vec![Field::new( + "s", + DataType::Struct(Fields::from(vec![ + child("a"), + child("b"), + child("c"), + child("d"), + ])), + true, + )])) + .project_by_ids(&[0, 4], false); + let mut fragment = Fragment::new(0); + fragment.files = vec![ + v2_0_file("a.lance", vec![0, 1]), + v2_0_file("b.lance", vec![0, 2]), + v2_0_file("c.lance", vec![0, 3]), + ]; + assert_eq!( + crate::dataset::optimize::repack::plan_fragment_repack( + &schema, + &fragment, + 3, + None, + Some(2) + ), + None + ); +} From bf6f0a149094ba250465852bfd276a05360ebf29 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Sat, 3 Oct 2026 04:51:12 +0800 Subject: [PATCH 20/24] fix(compaction): box the compaction futures; fix the distributed repack test --- python/python/tests/test_optimize.py | 3 ++- rust/lance/src/dataset/optimize.rs | 12 +++++++++--- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/python/python/tests/test_optimize.py b/python/python/tests/test_optimize.py index fd2fccc3806..bcf74d4dd9e 100644 --- a/python/python/tests/test_optimize.py +++ b/python/python/tests/test_optimize.py @@ -192,8 +192,9 @@ def test_distributed_repack(tmp_path: Path): dataset = _backfilled(tmp_path) expected = dataset.to_table() + # c and d sit in separate files, so the group moves them into one. plan = Compaction.plan( - dataset, options=dict(column_groups=[["d"]], scope="repack_columns") + dataset, options=dict(column_groups=[["c", "d"]], scope="repack_columns") ) assert [task.kind for task in plan.tasks] == ["repack_columns"] * 2 results = [ diff --git a/rust/lance/src/dataset/optimize.rs b/rust/lance/src/dataset/optimize.rs index 7bbf68dbe6c..b1cda3dbc13 100644 --- a/rust/lance/src/dataset/optimize.rs +++ b/rust/lance/src/dataset/optimize.rs @@ -1297,7 +1297,7 @@ pub async fn compact_files( ) -> Result { info!(target: TRACE_DATASET_EVENTS, event=DATASET_COMPACTING_EVENT, uri = &dataset.uri); let planner = DefaultCompactionPlanner::new(options)?; - compact_files_with_planner(dataset, remap_options, &planner).await + Box::pin(compact_files_with_planner(dataset, remap_options, &planner)).await } pub async fn compact_files_with_planner( @@ -1305,7 +1305,13 @@ pub async fn compact_files_with_planner( remap_options: Option>, // These will be deprecated later planner: &dyn CompactionPlanner, ) -> Result { - compact_files_with_executor(dataset, remap_options, planner, &DefaultCompactionExecutor).await + Box::pin(compact_files_with_executor( + dataset, + remap_options, + planner, + &DefaultCompactionExecutor, + )) + .await } /// Plan with `planner`, run every task with `executor`, and commit the @@ -3808,7 +3814,7 @@ pub async fn commit_compaction( a plan may put a fragment in one kind of task only" ))); } - let mut metrics = commit_rewrites(dataset, rewrites, remap_options, options).await?; + let mut metrics = Box::pin(commit_rewrites(dataset, rewrites, remap_options, options)).await?; metrics += repack::commit_repacked_files(dataset, repacks, options).await?; Ok(metrics) } From fa4241fa2e184b2c0477aa5830012cae6f71dcc5 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Sat, 3 Oct 2026 14:30:49 +0800 Subject: [PATCH 21/24] feat(compaction): expose column layout stats in Python and Java; document repacking --- docs/src/guide/read_and_write.md | 40 ++++++++++++- java/lance-jni/src/blocking_dataset.rs | 57 +++++++++++++++++++ .../org/lance/ColumnLayoutStatistics.java | 57 +++++++++++++++++++ java/src/main/java/org/lance/Dataset.java | 16 ++++++ .../test/java/org/lance/CompactionTest.java | 8 +++ python/python/lance/dataset.py | 20 +++++++ python/python/lance/lance/__init__.pyi | 1 + python/python/tests/test_optimize.py | 5 ++ python/src/dataset.rs | 15 +++++ 9 files changed, 217 insertions(+), 2 deletions(-) create mode 100644 java/src/main/java/org/lance/ColumnLayoutStatistics.java diff --git a/docs/src/guide/read_and_write.md b/docs/src/guide/read_and_write.md index 40062b6ce33..e4e8a4d05d5 100644 --- a/docs/src/guide/read_and_write.md +++ b/docs/src/guide/read_and_write.md @@ -497,8 +497,9 @@ ordering of the data will be preserved. !!! note - Compaction creates a new version of the table. It does not delete the old - version of the table and the files referenced by it. + Compaction creates a new version of the table, or two when it both + rewrites fragments and repacks columns (see below). It does not delete the + old versions of the table and the files referenced by them. ```python import lance @@ -517,6 +518,41 @@ of this, it's recommended to rewrite files before re-building indices. +#### Repack columns + +Each `add_columns` backfill gives every fragment one more data file. A large +fragment with few deletions is never rewritten, so those files pile up and slow +down scans and random reads. Compaction can repack a fragment's columns into +fewer data files instead. A repack moves no rows, so fragment ids, row +addresses, deletions, data overlays and index coverage stay as they are. + +Repacking is off by default. Set either option, as a `compact_files` argument or +as table config: + +- `max_data_files_per_fragment` (`lance.compaction.max_data_files_per_fragment`): + repack a fragment that holds its columns in more files than this. +- `column_groups` (`lance.compaction.column_groups`, groups separated by `;` and + columns by `,`): keep these top-level columns together in a file of their + own. Fragments that compaction rewrites are written this way too. + +`scope` limits a run to one kind of work: `"rewrite_fragments"`, +`"repack_columns"`, or `"all"` (the default). + +```python +dataset.optimize.compact_files(max_data_files_per_fragment=2) + +# Only repack, and keep the embedding in a file of its own. +dataset.optimize.compact_files(column_groups=[["embedding"]], scope="repack_columns") + +# How many data files each fragment holds now. +dataset.stats.column_layout_stats() +``` + +A repack only merges files it can empty. Files holding a blob column, a column +only partly present in the fragment, or spilled row lineage stay as they are, +so a fragment can stay above the limit. Fragments with legacy (V1) data files +are not repacked. + ### Cleanup old versions Lance is an immutable format — every write creates a new version. The new version diff --git a/java/lance-jni/src/blocking_dataset.rs b/java/lance-jni/src/blocking_dataset.rs index 91e4af7b1f6..6be12e1d361 100644 --- a/java/lance-jni/src/blocking_dataset.rs +++ b/java/lance-jni/src/blocking_dataset.rs @@ -1952,6 +1952,63 @@ fn inner_get_fragment_statistics<'local>( )?) } +#[unsafe(no_mangle)] +pub extern "system" fn Java_org_lance_Dataset_nativeGetColumnLayoutStatistics<'a>( + mut env: JNIEnv<'a>, + jdataset: JObject, +) -> JObject<'a> { + ok_or_throw!(env, inner_get_column_layout_statistics(&mut env, jdataset)) +} + +/// `Dataset::column_layout_stats` as parallel Java primitive arrays. +fn inner_get_column_layout_statistics<'local>( + env: &mut JNIEnv<'local>, + jdataset: JObject, +) -> Result> { + let stats = { + let dataset = + unsafe { env.get_rust_field::<_, _, BlockingDataset>(jdataset, NATIVE_DATASET) }?; + dataset.inner.column_layout_stats() + }; + let len = i32::try_from(stats.len()).map_err(|_| { + Error::runtime_error(format!( + "Column layout statistics contain {} fragments, exceeding the Java array limit of {}", + stats.len(), + i32::MAX + )) + })?; + let to_int = |value: usize| { + i32::try_from(value) + .map_err(|_| Error::runtime_error(format!("{value} does not fit in a Java int"))) + }; + let fragment_ids: Vec = stats.iter().map(|s| s.fragment_id as i64).collect(); + let live_file_counts = stats + .iter() + .map(|s| to_int(s.live_file_count)) + .collect::>>()?; + let overlay_counts = stats + .iter() + .map(|s| to_int(s.overlay_count)) + .collect::>>()?; + + let jfragment_ids = env.new_long_array(len)?; + let jlive_file_counts = env.new_int_array(len)?; + let joverlay_counts = env.new_int_array(len)?; + env.set_long_array_region(&jfragment_ids, 0, &fragment_ids)?; + env.set_int_array_region(&jlive_file_counts, 0, &live_file_counts)?; + env.set_int_array_region(&joverlay_counts, 0, &overlay_counts)?; + + Ok(env.new_object( + "org/lance/ColumnLayoutStatistics", + "([J[I[I)V", + &[ + JValue::Object(&jfragment_ids), + JValue::Object(&jlive_file_counts), + JValue::Object(&joverlay_counts), + ], + )?) +} + #[unsafe(no_mangle)] pub extern "system" fn Java_org_lance_Dataset_getFragmentNative<'a>( mut env: JNIEnv<'a>, diff --git a/java/src/main/java/org/lance/ColumnLayoutStatistics.java b/java/src/main/java/org/lance/ColumnLayoutStatistics.java new file mode 100644 index 00000000000..fe9d30dd6fe --- /dev/null +++ b/java/src/main/java/org/lance/ColumnLayoutStatistics.java @@ -0,0 +1,57 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.lance; + +/** + * Per-fragment column-layout statistics of a dataset version as parallel primitive arrays: {@code + * fragmentIds[i]}, {@code liveFileCounts[i]} and {@code overlayCounts[i]} describe the same + * fragment. Returned by {@link Dataset#getColumnLayoutStatistics()}. Compaction repacks a + * fragment's columns into fewer files when its live file count is above {@code + * maxDataFilesPerFragment}. + */ +public final class ColumnLayoutStatistics { + private final long[] fragmentIds; + private final int[] liveFileCounts; + private final int[] overlayCounts; + + ColumnLayoutStatistics(long[] fragmentIds, int[] liveFileCounts, int[] overlayCounts) { + this.fragmentIds = fragmentIds; + this.liveFileCounts = liveFileCounts; + this.overlayCounts = overlayCounts; + } + + /** Fragment IDs in manifest order. */ + public long[] getFragmentIds() { + return fragmentIds; + } + + /** + * Number of data files holding at least one column of the schema, per fragment, aligned with + * {@link #getFragmentIds()}. A file left holding only tombstones or dropped columns, or kept only + * for the fragment's spilled row lineage, is not counted. + */ + public int[] getLiveFileCounts() { + return liveFileCounts; + } + + /** Number of data overlay files per fragment, aligned with {@link #getFragmentIds()}. */ + public int[] getOverlayCounts() { + return overlayCounts; + } + + /** Number of fragments described. */ + public int size() { + return fragmentIds.length; + } +} diff --git a/java/src/main/java/org/lance/Dataset.java b/java/src/main/java/org/lance/Dataset.java index 53fe78c9d65..a83010a733c 100644 --- a/java/src/main/java/org/lance/Dataset.java +++ b/java/src/main/java/org/lance/Dataset.java @@ -1575,6 +1575,22 @@ public FragmentStatistics getFragmentStatistics() { private native FragmentStatistics nativeGetFragmentStatistics(); + /** + * Get per-fragment column-layout statistics for this dataset version: how many data files hold a + * column of the schema, and how many overlay files each fragment carries. Only manifest metadata + * is read. + * + * @return column-layout statistics as parallel arrays, in manifest order + */ + public ColumnLayoutStatistics getColumnLayoutStatistics() { + try (LockManager.ReadLock readLock = lockManager.acquireReadLock()) { + Preconditions.checkArgument(nativeDatasetHandle != 0, "Dataset is closed"); + return nativeGetColumnLayoutStatistics(); + } + } + + private native ColumnLayoutStatistics nativeGetColumnLayoutStatistics(); + /** * Gets the arrow schema of the dataset. * diff --git a/java/src/test/java/org/lance/CompactionTest.java b/java/src/test/java/org/lance/CompactionTest.java index d37875a5c27..7f26048a780 100644 --- a/java/src/test/java/org/lance/CompactionTest.java +++ b/java/src/test/java/org/lance/CompactionTest.java @@ -42,6 +42,7 @@ import java.util.List; import java.util.Optional; +import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -285,6 +286,11 @@ public void testRepackColumns(@TempDir Path tempDir) throws Exception { new SqlExpressions.Builder().withExpression("double_id", "id * 2").build(), Optional.empty()); assertEquals(2, dataset.getFragments().get(0).metadata().getFiles().size()); + ColumnLayoutStatistics layout = dataset.getColumnLayoutStatistics(); + assertEquals(2, layout.size()); + assertArrayEquals(new long[] {0, 1}, layout.getFragmentIds()); + assertArrayEquals(new int[] {2, 2}, layout.getLiveFileCounts()); + assertArrayEquals(new int[] {0, 0}, layout.getOverlayCounts()); CompactionOptions options = CompactionOptions.builder() @@ -315,6 +321,8 @@ public void testRepackColumns(@TempDir Path tempDir) throws Exception { for (Fragment fragment : dataset.getFragments()) { assertEquals(1, fragment.metadata().getFiles().size()); } + assertArrayEquals( + new int[] {1, 1}, dataset.getColumnLayoutStatistics().getLiveFileCounts()); } } } diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index 1a7e6ff85f2..207458f18e4 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -8032,6 +8032,15 @@ class DatasetStats(TypedDict): num_small_files: int +class FragmentColumnLayoutStats(TypedDict): + fragment_id: int + #: Data files holding at least one column of the schema. This is the count + #: ``max_data_files_per_fragment`` in compaction is compared with. + live_file_count: int + #: Data overlay files attached to the fragment. + overlay_count: int + + class LanceStats: """ Statistics about a LanceDataset. @@ -8068,6 +8077,17 @@ def data_stats(self) -> DataStatistics: """ return self._ds.data_stats() + def column_layout_stats(self) -> List[FragmentColumnLayoutStats]: + """ + How many data files and overlay files each fragment holds, in manifest + fragment order. + + Compaction repacks a fragment's columns into fewer files when + ``live_file_count`` is above ``max_data_files_per_fragment``. Only + manifest metadata is read. + """ + return self._ds.column_layout_stats() + def write_dataset( data_obj: ReaderLike, diff --git a/python/python/lance/lance/__init__.pyi b/python/python/lance/lance/__init__.pyi index f338f16c9c7..1e9ea4316fa 100644 --- a/python/python/lance/lance/__init__.pyi +++ b/python/python/lance/lance/__init__.pyi @@ -503,6 +503,7 @@ class _Dataset: data_storage_version: Optional[str] = None, ) -> UpdateResult: ... def count_deleted_rows(self) -> int: ... + def column_layout_stats(self) -> List[Dict[str, int]]: ... def versions(self) -> List[Version]: ... def version_refs(self) -> List[VersionRef]: ... def version(self) -> int: ... diff --git a/python/python/tests/test_optimize.py b/python/python/tests/test_optimize.py index bcf74d4dd9e..ea554239d67 100644 --- a/python/python/tests/test_optimize.py +++ b/python/python/tests/test_optimize.py @@ -176,6 +176,10 @@ def test_compact_files_repacks_columns(tmp_path: Path): dataset = _backfilled(tmp_path) expected = dataset.to_table() assert _files_per_fragment(dataset) == [3, 3] + assert dataset.stats.column_layout_stats() == [ + {"fragment_id": 0, "live_file_count": 3, "overlay_count": 0}, + {"fragment_id": 1, "live_file_count": 3, "overlay_count": 0}, + ] metrics = dataset.optimize.compact_files( max_data_files_per_fragment=1, scope="repack_columns" @@ -184,6 +188,7 @@ def test_compact_files_repacks_columns(tmp_path: Path): assert metrics.files_added == 2 assert metrics.fragments_added == 0 assert _files_per_fragment(dataset) == [1, 1] + assert [s["live_file_count"] for s in dataset.stats.column_layout_stats()] == [1, 1] assert dataset.to_table() == expected assert [f.fragment_id for f in dataset.get_fragments()] == [0, 1] diff --git a/python/src/dataset.rs b/python/src/dataset.rs index eca9c2e756d..d781b2dee60 100644 --- a/python/src/dataset.rs +++ b/python/src/dataset.rs @@ -2103,6 +2103,21 @@ impl Dataset { .map_err(|err| PyIOError::new_err(err.to_string())) } + /// Per-fragment column-layout stats, in manifest fragment order. + fn column_layout_stats(&self, py: Python<'_>) -> PyResult>> { + self.ds + .column_layout_stats() + .into_iter() + .map(|stats| { + let dict = PyDict::new(py); + dict.set_item("fragment_id", stats.fragment_id)?; + dict.set_item("live_file_count", stats.live_file_count)?; + dict.set_item("overlay_count", stats.overlay_count)?; + Ok(dict.unbind()) + }) + .collect() + } + #[pyo3(signature=(new_bases, transaction_properties=None))] fn add_bases( &mut self, From dd5890416a571aa4c9300df2d00b2a70e5f3664f Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Sat, 3 Oct 2026 15:29:50 +0800 Subject: [PATCH 22/24] feat: report file sizes, fields per file and dead slot ratio in column_layout_stats --- docs/src/guide/read_and_write.md | 2 +- java/lance-jni/src/blocking_dataset.rs | 39 +++++-- .../org/lance/ColumnLayoutStatistics.java | 45 ++++++-- .../test/java/org/lance/CompactionTest.java | 9 +- python/python/lance/dataset.py | 13 ++- python/python/lance/lance/__init__.pyi | 2 +- python/python/tests/test_optimize.py | 16 ++- python/src/dataset.rs | 3 + rust/lance/src/dataset/compaction_stats.rs | 102 +++++++++++++++--- 9 files changed, 192 insertions(+), 39 deletions(-) diff --git a/docs/src/guide/read_and_write.md b/docs/src/guide/read_and_write.md index e4e8a4d05d5..ffeab642308 100644 --- a/docs/src/guide/read_and_write.md +++ b/docs/src/guide/read_and_write.md @@ -544,7 +544,7 @@ dataset.optimize.compact_files(max_data_files_per_fragment=2) # Only repack, and keep the embedding in a file of its own. dataset.optimize.compact_files(column_groups=[["embedding"]], scope="repack_columns") -# How many data files each fragment holds now. +# Per fragment: live data files, file sizes, fields per file, dead field slots. dataset.stats.column_layout_stats() ``` diff --git a/java/lance-jni/src/blocking_dataset.rs b/java/lance-jni/src/blocking_dataset.rs index 6be12e1d361..e464abf5bbf 100644 --- a/java/lance-jni/src/blocking_dataset.rs +++ b/java/lance-jni/src/blocking_dataset.rs @@ -1960,7 +1960,7 @@ pub extern "system" fn Java_org_lance_Dataset_nativeGetColumnLayoutStatistics<'a ok_or_throw!(env, inner_get_column_layout_statistics(&mut env, jdataset)) } -/// `Dataset::column_layout_stats` as parallel Java primitive arrays. +/// `Dataset::column_layout_stats` as parallel Java arrays. fn inner_get_column_layout_statistics<'local>( env: &mut JNIEnv<'local>, jdataset: JObject, @@ -1970,22 +1970,17 @@ fn inner_get_column_layout_statistics<'local>( unsafe { env.get_rust_field::<_, _, BlockingDataset>(jdataset, NATIVE_DATASET) }?; dataset.inner.column_layout_stats() }; - let len = i32::try_from(stats.len()).map_err(|_| { - Error::runtime_error(format!( - "Column layout statistics contain {} fragments, exceeding the Java array limit of {}", - stats.len(), - i32::MAX - )) - })?; let to_int = |value: usize| { i32::try_from(value) .map_err(|_| Error::runtime_error(format!("{value} does not fit in a Java int"))) }; + let len = to_int(stats.len())?; let fragment_ids: Vec = stats.iter().map(|s| s.fragment_id as i64).collect(); let live_file_counts = stats .iter() .map(|s| to_int(s.live_file_count)) .collect::>>()?; + let ratios: Vec = stats.iter().map(|s| s.tombstoned_field_ratio).collect(); let overlay_counts = stats .iter() .map(|s| to_int(s.overlay_count)) @@ -1993,17 +1988,43 @@ fn inner_get_column_layout_statistics<'local>( let jfragment_ids = env.new_long_array(len)?; let jlive_file_counts = env.new_int_array(len)?; + let jratios = env.new_double_array(len)?; let joverlay_counts = env.new_int_array(len)?; env.set_long_array_region(&jfragment_ids, 0, &fragment_ids)?; env.set_int_array_region(&jlive_file_counts, 0, &live_file_counts)?; + env.set_double_array_region(&jratios, 0, &ratios)?; env.set_int_array_region(&joverlay_counts, 0, &overlay_counts)?; + let jfile_sizes = env.new_object_array(len, "[J", JObject::null())?; + let jfields_per_file = env.new_object_array(len, "[I", JObject::null())?; + for (index, s) in stats.iter().enumerate() { + let sizes: Vec = s + .file_sizes + .iter() + .map(|size| size.map_or(-1, |size| size as i64)) + .collect(); + let fields = s + .fields_per_file + .iter() + .map(|count| to_int(*count)) + .collect::>>()?; + let jsizes = env.new_long_array(to_int(sizes.len())?)?; + env.set_long_array_region(&jsizes, 0, &sizes)?; + env.set_object_array_element(&jfile_sizes, index as i32, &jsizes)?; + let jfields = env.new_int_array(to_int(fields.len())?)?; + env.set_int_array_region(&jfields, 0, &fields)?; + env.set_object_array_element(&jfields_per_file, index as i32, &jfields)?; + } + Ok(env.new_object( "org/lance/ColumnLayoutStatistics", - "([J[I[I)V", + "([J[I[[J[[I[D[I)V", &[ JValue::Object(&jfragment_ids), JValue::Object(&jlive_file_counts), + JValue::Object(&jfile_sizes), + JValue::Object(&jfields_per_file), + JValue::Object(&jratios), JValue::Object(&joverlay_counts), ], )?) diff --git a/java/src/main/java/org/lance/ColumnLayoutStatistics.java b/java/src/main/java/org/lance/ColumnLayoutStatistics.java index fe9d30dd6fe..2c7031d7552 100644 --- a/java/src/main/java/org/lance/ColumnLayoutStatistics.java +++ b/java/src/main/java/org/lance/ColumnLayoutStatistics.java @@ -14,20 +14,31 @@ package org.lance; /** - * Per-fragment column-layout statistics of a dataset version as parallel primitive arrays: {@code - * fragmentIds[i]}, {@code liveFileCounts[i]} and {@code overlayCounts[i]} describe the same - * fragment. Returned by {@link Dataset#getColumnLayoutStatistics()}. Compaction repacks a - * fragment's columns into fewer files when its live file count is above {@code - * maxDataFilesPerFragment}. + * Per-fragment column-layout statistics of a dataset version as parallel arrays: index {@code i} of + * every getter describes the same fragment. Returned by {@link + * Dataset#getColumnLayoutStatistics()}. Compaction repacks a fragment's columns into fewer files + * when its live file count is above {@code maxDataFilesPerFragment}. */ public final class ColumnLayoutStatistics { private final long[] fragmentIds; private final int[] liveFileCounts; + private final long[][] fileSizes; + private final int[][] fieldsPerFile; + private final double[] tombstonedFieldRatios; private final int[] overlayCounts; - ColumnLayoutStatistics(long[] fragmentIds, int[] liveFileCounts, int[] overlayCounts) { + ColumnLayoutStatistics( + long[] fragmentIds, + int[] liveFileCounts, + long[][] fileSizes, + int[][] fieldsPerFile, + double[] tombstonedFieldRatios, + int[] overlayCounts) { this.fragmentIds = fragmentIds; this.liveFileCounts = liveFileCounts; + this.fileSizes = fileSizes; + this.fieldsPerFile = fieldsPerFile; + this.tombstonedFieldRatios = tombstonedFieldRatios; this.overlayCounts = overlayCounts; } @@ -45,6 +56,28 @@ public int[] getLiveFileCounts() { return liveFileCounts; } + /** + * Recorded size in bytes of each data file, per fragment, in the fragment's file order. {@code + * -1} when the manifest has no size for the file. + */ + public long[][] getFileSizes() { + return fileSizes; + } + + /** Number of schema fields each data file holds, per fragment, in the fragment's file order. */ + public int[][] getFieldsPerFile() { + return fieldsPerFile; + } + + /** + * Share of each fragment's field slots holding no live data: slots tombstoned by a column update + * or repack, or left by a dropped column. Spilled row lineage is not counted. A fragment rewrite + * reclaims these slots. + */ + public double[] getTombstonedFieldRatios() { + return tombstonedFieldRatios; + } + /** Number of data overlay files per fragment, aligned with {@link #getFragmentIds()}. */ public int[] getOverlayCounts() { return overlayCounts; diff --git a/java/src/test/java/org/lance/CompactionTest.java b/java/src/test/java/org/lance/CompactionTest.java index 7f26048a780..e0019af1890 100644 --- a/java/src/test/java/org/lance/CompactionTest.java +++ b/java/src/test/java/org/lance/CompactionTest.java @@ -290,6 +290,10 @@ public void testRepackColumns(@TempDir Path tempDir) throws Exception { assertEquals(2, layout.size()); assertArrayEquals(new long[] {0, 1}, layout.getFragmentIds()); assertArrayEquals(new int[] {2, 2}, layout.getLiveFileCounts()); + assertArrayEquals(new int[] {2, 1}, layout.getFieldsPerFile()[0]); + assertEquals(2, layout.getFileSizes()[0].length); + assertTrue(layout.getFileSizes()[0][0] > 0); + assertArrayEquals(new double[] {0.0, 0.0}, layout.getTombstonedFieldRatios()); assertArrayEquals(new int[] {0, 0}, layout.getOverlayCounts()); CompactionOptions options = @@ -321,8 +325,9 @@ public void testRepackColumns(@TempDir Path tempDir) throws Exception { for (Fragment fragment : dataset.getFragments()) { assertEquals(1, fragment.metadata().getFiles().size()); } - assertArrayEquals( - new int[] {1, 1}, dataset.getColumnLayoutStatistics().getLiveFileCounts()); + ColumnLayoutStatistics repacked = dataset.getColumnLayoutStatistics(); + assertArrayEquals(new int[] {1, 1}, repacked.getLiveFileCounts()); + assertArrayEquals(new int[] {3}, repacked.getFieldsPerFile()[0]); } } } diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index 207458f18e4..aa45b6fefb2 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -8037,6 +8037,15 @@ class FragmentColumnLayoutStats(TypedDict): #: Data files holding at least one column of the schema. This is the count #: ``max_data_files_per_fragment`` in compaction is compared with. live_file_count: int + #: Recorded size in bytes of each data file, in the fragment's file order; + #: ``None`` when the manifest has no size for the file. + file_sizes: List[Optional[int]] + #: Number of schema fields each data file holds, in the fragment's file order. + fields_per_file: List[int] + #: Share of the fragment's field slots holding no live data (tombstoned by a + #: column update or repack, or left by a dropped column). Spilled row + #: lineage is not counted. A fragment rewrite reclaims these slots. + tombstoned_field_ratio: float #: Data overlay files attached to the fragment. overlay_count: int @@ -8079,8 +8088,8 @@ def data_stats(self) -> DataStatistics: def column_layout_stats(self) -> List[FragmentColumnLayoutStats]: """ - How many data files and overlay files each fragment holds, in manifest - fragment order. + How each fragment's columns are laid out across data files, in + manifest fragment order. Compaction repacks a fragment's columns into fewer files when ``live_file_count`` is above ``max_data_files_per_fragment``. Only diff --git a/python/python/lance/lance/__init__.pyi b/python/python/lance/lance/__init__.pyi index 1e9ea4316fa..25ba1e4fc24 100644 --- a/python/python/lance/lance/__init__.pyi +++ b/python/python/lance/lance/__init__.pyi @@ -503,7 +503,7 @@ class _Dataset: data_storage_version: Optional[str] = None, ) -> UpdateResult: ... def count_deleted_rows(self) -> int: ... - def column_layout_stats(self) -> List[Dict[str, int]]: ... + def column_layout_stats(self) -> List[Dict[str, Any]]: ... def versions(self) -> List[Version]: ... def version_refs(self) -> List[VersionRef]: ... def version(self) -> int: ... diff --git a/python/python/tests/test_optimize.py b/python/python/tests/test_optimize.py index ea554239d67..f8ccd73e601 100644 --- a/python/python/tests/test_optimize.py +++ b/python/python/tests/test_optimize.py @@ -176,10 +176,14 @@ def test_compact_files_repacks_columns(tmp_path: Path): dataset = _backfilled(tmp_path) expected = dataset.to_table() assert _files_per_fragment(dataset) == [3, 3] - assert dataset.stats.column_layout_stats() == [ - {"fragment_id": 0, "live_file_count": 3, "overlay_count": 0}, - {"fragment_id": 1, "live_file_count": 3, "overlay_count": 0}, - ] + stats = dataset.stats.column_layout_stats() + assert [s["fragment_id"] for s in stats] == [0, 1] + for s in stats: + assert s["live_file_count"] == 3 + assert s["fields_per_file"] == [2, 1, 1] + assert all(size is not None for size in s["file_sizes"]) + assert s["tombstoned_field_ratio"] == 0.0 + assert s["overlay_count"] == 0 metrics = dataset.optimize.compact_files( max_data_files_per_fragment=1, scope="repack_columns" @@ -188,7 +192,9 @@ def test_compact_files_repacks_columns(tmp_path: Path): assert metrics.files_added == 2 assert metrics.fragments_added == 0 assert _files_per_fragment(dataset) == [1, 1] - assert [s["live_file_count"] for s in dataset.stats.column_layout_stats()] == [1, 1] + stats = dataset.stats.column_layout_stats() + assert [s["live_file_count"] for s in stats] == [1, 1] + assert [s["fields_per_file"] for s in stats] == [[4], [4]] assert dataset.to_table() == expected assert [f.fragment_id for f in dataset.get_fragments()] == [0, 1] diff --git a/python/src/dataset.rs b/python/src/dataset.rs index d781b2dee60..47663948771 100644 --- a/python/src/dataset.rs +++ b/python/src/dataset.rs @@ -2112,6 +2112,9 @@ impl Dataset { let dict = PyDict::new(py); dict.set_item("fragment_id", stats.fragment_id)?; dict.set_item("live_file_count", stats.live_file_count)?; + dict.set_item("file_sizes", stats.file_sizes)?; + dict.set_item("fields_per_file", stats.fields_per_file)?; + dict.set_item("tombstoned_field_ratio", stats.tombstoned_field_ratio)?; dict.set_item("overlay_count", stats.overlay_count)?; Ok(dict.unbind()) }) diff --git a/rust/lance/src/dataset/compaction_stats.rs b/rust/lance/src/dataset/compaction_stats.rs index 14b6ebd3f24..a690225aa41 100644 --- a/rust/lance/src/dataset/compaction_stats.rs +++ b/rust/lance/src/dataset/compaction_stats.rs @@ -11,19 +11,33 @@ use std::collections::HashSet; +use lance_table::format::overlay::TOMBSTONE_FIELD_ID; + use super::Dataset; /// Column-layout statistics for a single fragment. -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq)] pub struct FragmentColumnLayoutStats { /// The fragment these stats describe. pub fragment_id: u64, /// Number of data files holding at least one column of the dataset schema, /// the count `max_data_files_per_fragment` is compared with. A file left - /// holding only tombstones or dropped - /// columns, or kept only for the fragment's spilled row lineage, holds no - /// column a read touches and is not counted. + /// holding only tombstones or dropped columns, or kept only for the + /// fragment's spilled row lineage, holds no column a read touches and is + /// not counted. pub live_file_count: usize, + /// Recorded size in bytes of each data file, in the fragment's file order. + /// `None` when the manifest has no size for the file. + pub file_sizes: Vec>, + /// Number of fields of the dataset schema each data file holds, in the + /// fragment's file order. A file counted in `live_file_count` holds at + /// least one. + pub fields_per_file: Vec, + /// The share of the fragment's field slots that hold no live data: slots + /// tombstoned by a column update or repack, or left by a dropped column. + /// The reserved ids of spilled row lineage are not counted on either side. + /// A fragment rewrite reclaims these slots. + pub tombstoned_field_ratio: f64, /// Number of overlay files attached to the fragment. pub overlay_count: usize, } @@ -42,14 +56,42 @@ impl Dataset { self.manifest .fragments .iter() - .map(|fragment| FragmentColumnLayoutStats { - fragment_id: fragment.id, - live_file_count: fragment + .map(|fragment| { + let fields_per_file: Vec = fragment + .files + .iter() + .map(|file| { + file.fields + .iter() + .filter(|id| schema_ids.contains(id)) + .count() + }) + .collect(); + // Field slots of user columns, live or dead; the other negative + // ids are spilled row lineage and are left out. + let user_slots = fragment .files .iter() - .filter(|file| file.fields.iter().any(|id| schema_ids.contains(id))) - .count(), - overlay_count: fragment.overlays.len(), + .flat_map(|file| file.fields.iter()) + .filter(|id| **id >= 0 || **id == TOMBSTONE_FIELD_ID) + .count(); + let live_slots: usize = fields_per_file.iter().sum(); + FragmentColumnLayoutStats { + fragment_id: fragment.id, + live_file_count: fields_per_file.iter().filter(|count| **count > 0).count(), + file_sizes: fragment + .files + .iter() + .map(|file| file.file_size_bytes.get().map(|size| size.get())) + .collect(), + fields_per_file, + tombstoned_field_ratio: if user_slots == 0 { + 0.0 + } else { + (user_slots - live_slots) as f64 / user_slots as f64 + }, + overlay_count: fragment.overlays.len(), + } }) .collect() } @@ -112,6 +154,37 @@ mod tests { stats.iter().all(|s| s.live_file_count == 2), "each fragment should now have 2 data files: {stats:?}" ); + for s in &stats { + assert_eq!(s.fields_per_file, vec![1, 1]); + assert!(s.file_sizes.iter().all(Option::is_some), "{s:?}"); + assert_eq!(s.tombstoned_field_ratio, 0.0); + } + } + + /// A dropped column leaves its field id in the file that held it, which + /// counts as a dead slot. + #[tokio::test] + async fn column_layout_stats_counts_dropped_columns_as_dead_slots() { + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from_iter_values(0..4)), + Arc::new(Int32Array::from_iter_values(4..8)), + ], + ) + .unwrap(); + let reader = RecordBatchIterator::new([Ok(batch)], schema); + let mut dataset = Dataset::write(reader, "memory://", None).await.unwrap(); + dataset.drop_columns(&["b"]).await.unwrap(); + + let stats = dataset.column_layout_stats(); + assert_eq!(stats[0].live_file_count, 1); + assert_eq!(stats[0].fields_per_file, vec![1]); + assert_eq!(stats[0].tombstoned_field_ratio, 0.5); } /// Only files holding a schema column count. The extra files are added to @@ -121,9 +194,7 @@ mod tests { #[tokio::test] async fn column_layout_stats_skips_files_without_schema_columns() { use lance_file::version::ConcreteFileVersion; - use lance_table::format::{ - DataFile, ROW_ID_FIELD_ID, RowIdMeta, overlay::TOMBSTONE_FIELD_ID, - }; + use lance_table::format::{DataFile, ROW_ID_FIELD_ID, RowIdMeta}; let schema = Arc::new(ArrowSchema::new(vec![Field::new( "a", @@ -163,5 +234,10 @@ mod tests { let stats = dataset.column_layout_stats(); assert_eq!(stats[0].live_file_count, 1, "{stats:?}"); + assert_eq!(stats[0].fields_per_file, vec![1, 0, 0, 0]); + assert!(stats[0].file_sizes[0].is_some()); + assert_eq!(stats[0].file_sizes[1..], [None, None, None]); + // `a`, two tombstones and the dropped id; the row id slot is lineage. + assert_eq!(stats[0].tombstoned_field_ratio, 0.75); } } From 0f6c23ce4d6054f54d6a9c3520c0c58b52a82199 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Sat, 3 Oct 2026 15:43:19 +0800 Subject: [PATCH 23/24] fix(java): release per-fragment local refs in getColumnLayoutStatistics --- java/lance-jni/src/blocking_dataset.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/java/lance-jni/src/blocking_dataset.rs b/java/lance-jni/src/blocking_dataset.rs index e464abf5bbf..bee71622fb1 100644 --- a/java/lance-jni/src/blocking_dataset.rs +++ b/java/lance-jni/src/blocking_dataset.rs @@ -2011,9 +2011,11 @@ fn inner_get_column_layout_statistics<'local>( let jsizes = env.new_long_array(to_int(sizes.len())?)?; env.set_long_array_region(&jsizes, 0, &sizes)?; env.set_object_array_element(&jfile_sizes, index as i32, &jsizes)?; + env.delete_local_ref(jsizes)?; let jfields = env.new_int_array(to_int(fields.len())?)?; env.set_int_array_region(&jfields, 0, &fields)?; env.set_object_array_element(&jfields_per_file, index as i32, &jfields)?; + env.delete_local_ref(jfields)?; } Ok(env.new_object( From 104231c887613c80c7c94e0d213b1c649bfaea59 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Mon, 5 Oct 2026 00:07:37 +0800 Subject: [PATCH 24/24] chore: document DataReplacement group order in the proto --- protos/transaction.proto | 3 +++ 1 file changed, 3 insertions(+) diff --git a/protos/transaction.proto b/protos/transaction.proto index e3e626f62be..75c164f527a 100644 --- a/protos/transaction.proto +++ b/protos/transaction.proto @@ -364,6 +364,9 @@ message Transaction { // An operation that replaces the data in a region of the table with new data. message DataReplacement { + /* A fragment can be the target of several groups, one per new file. The + * groups apply in order, each to the fragment as the previous ones left it. + */ repeated DataReplacementGroup replacements = 1; /* Whether the new files change any value. Absent means true, which is what * a writer that predates this field meant. False means the values were only