diff --git a/docs/src/format/table/transaction.md b/docs/src/format/table/transaction.md index 85e4ca43ed4..ef3384c6254 100644 --- a/docs/src/format/table/transaction.md +++ b/docs/src/format/table/transaction.md @@ -505,13 +505,24 @@ The following operations are retryable conflicts with DataOverlay: overlay→base fold changes physical row addresses or consumes the overlays, so the overlay's offsets are no longer valid; the writer must re-read the new fragment, recompute, and retry. -- Merge (always) - -DataOverlay is compatible with another DataOverlay (any fields), Append, Delete, -and DataReplacement or a column rewrite (Update with `REWRITE_COLUMNS`) of the same -field, because all of these preserve physical row addresses: overlay offsets stay -valid, the overlay is newer and wins its covered cells, and the version gate -excludes those cells from any rebuilt index. +- Merge (always). +- A row-moving Update that touches an overlaid fragment — a delete-and-reinsert + update (any update that is not a `REWRITE_COLUMNS` column rewrite) relocates the + updated rows into new fragments, so the overlay's physical offsets no longer + address them; the writer must re-read and retry. + +DataOverlay is compatible with another DataOverlay (any fields), Append, Delete, a +`REWRITE_COLUMNS` column rewrite, and DataReplacement, because all of these +preserve physical row addresses: overlay offsets stay valid, the overlay is newer +and wins its covered cells, and the version gate excludes those cells from any +rebuilt index. + +When a DataReplacement or a `REWRITE_COLUMNS` update writes new base values for a +field, it supersedes any older overlay on that field: the writer tombstones the +overlay's entry for the rewritten field — replacing the field id with the obsolete +sentinel, as with obsolete base columns — so the fresh base values are not silently +shadowed. Overlay entries for other fields are preserved, and an overlay left with +no live fields is dropped. ### UpdateMemWalState diff --git a/java/lance-jni/src/fragment.rs b/java/lance-jni/src/fragment.rs index d6603925947..59ce5e553ff 100644 --- a/java/lance-jni/src/fragment.rs +++ b/java/lance-jni/src/fragment.rs @@ -828,6 +828,9 @@ impl FromJObjectWithEnv for JObject<'_> { row_id_meta, created_at_version_meta, last_updated_at_version_meta, + // Overlays are not exposed to Java yet, and the reverse conversion + // does not export them, so this round-trip is overlay-free. + overlays: vec![], }) } } diff --git a/python/src/fragment.rs b/python/src/fragment.rs index dbe5c426903..336c6a4cf51 100644 --- a/python/src/fragment.rs +++ b/python/src/fragment.rs @@ -825,6 +825,9 @@ impl FromPyObject<'_, '_> for PyLance { row_id_meta, last_updated_at_version_meta, created_at_version_meta, + // Overlays are not exposed to Python yet, and the reverse conversion + // does not export them, so this round-trip is overlay-free. + overlays: vec![], })) } } diff --git a/rust/lance-table/src/feature_flags.rs b/rust/lance-table/src/feature_flags.rs index 096f0da79e5..41b8e415f8e 100644 --- a/rust/lance-table/src/feature_flags.rs +++ b/rust/lance-table/src/feature_flags.rs @@ -20,8 +20,22 @@ pub const FLAG_TABLE_CONFIG: u64 = 8; pub const FLAG_BASE_PATHS: u64 = 16; /// Disable writing transaction file under _transaction/, this flag is set when we only want to write inline transaction in manifest pub const FLAG_DISABLE_TRANSACTION_FILE: u64 = 32; +/// Fragments contain data overlay files, which supply new values for a subset of +/// cells without rewriting base data files. A reader that does not understand +/// overlays must refuse the dataset, since ignoring an overlay would silently +/// return stale base values. +/// +/// Data overlay files are not yet a released feature: in release builds this flag +/// is treated as unknown (so a release reader/writer refuses an overlay dataset) +/// unless [`ENABLE_UNSTABLE_DATA_OVERLAY_FILES_ENV`] is set, which lets benchmarks opt in. +/// Debug builds always understand it so tests exercise the path. +pub const FLAG_UNSTABLE_DATA_OVERLAY_FILES: u64 = 64; /// The first bit that is unknown as a feature flag -pub const FLAG_UNKNOWN: u64 = 64; +pub const FLAG_UNKNOWN: u64 = 128; + +/// Environment variable that opts a release build into reading and writing data +/// overlay files before the feature is generally released. +pub const ENABLE_UNSTABLE_DATA_OVERLAY_FILES_ENV: &str = "LANCE_ENABLE_UNSTABLE_DATA_OVERLAY_FILES"; /// Set the reader and writer feature flags in the manifest based on the contents of the manifest. pub fn apply_feature_flags( @@ -71,18 +85,62 @@ pub fn apply_feature_flags( manifest.writer_feature_flags |= FLAG_BASE_PATHS; } + // Overlay files change cell values on read, so a reader that ignores them + // would return stale base values. Both readers and writers must understand + // them. + let has_overlays = manifest + .fragments + .iter() + .any(|frag| !frag.overlays.is_empty()); + if has_overlays { + manifest.reader_feature_flags |= FLAG_UNSTABLE_DATA_OVERLAY_FILES; + manifest.writer_feature_flags |= FLAG_UNSTABLE_DATA_OVERLAY_FILES; + } + if disable_transaction_file { manifest.writer_feature_flags |= FLAG_DISABLE_TRANSACTION_FILE; } Ok(()) } +/// Whether this build understands data overlay files: always in debug builds, +/// and in release builds only when [`ENABLE_UNSTABLE_DATA_OVERLAY_FILES_ENV`] is set. +fn data_overlay_files_enabled() -> bool { + cfg!(debug_assertions) || std::env::var_os(ENABLE_UNSTABLE_DATA_OVERLAY_FILES_ENV).is_some() +} + +/// Clear `flag` from `flags` when its gating feature is not enabled in this +/// build; leave it set otherwise. One call per unstable flag, so support for +/// several unstable features chains cleanly. +fn mark_supported(flags: &mut u64, flag: u64, feature_enabled: bool) { + if !feature_enabled { + *flags &= !flag; + } +} + +/// The feature-flag bits this build understands, given whether overlay support +/// is enabled. Split out from [`supported_flags`] so the policy is testable +/// without toggling the build profile or environment. +fn supported_flags_when(overlay_enabled: bool) -> u64 { + let mut supported = FLAG_UNKNOWN - 1; + mark_supported( + &mut supported, + FLAG_UNSTABLE_DATA_OVERLAY_FILES, + overlay_enabled, + ); + supported +} + +fn supported_flags() -> u64 { + supported_flags_when(data_overlay_files_enabled()) +} + pub fn can_read_dataset(reader_flags: u64) -> bool { - reader_flags < FLAG_UNKNOWN + reader_flags & !supported_flags() == 0 } pub fn can_write_dataset(writer_flags: u64) -> bool { - writer_flags < FLAG_UNKNOWN + writer_flags & !supported_flags() == 0 } pub fn has_deprecated_v2_feature_flag(writer_flags: u64) -> bool { @@ -103,6 +161,13 @@ mod tests { assert!(can_read_dataset(super::FLAG_TABLE_CONFIG)); assert!(can_read_dataset(super::FLAG_BASE_PATHS)); assert!(can_read_dataset(super::FLAG_DISABLE_TRANSACTION_FILE)); + // Overlay support is gated on the build profile / env opt-in, so the + // flag is readable exactly when overlays are enabled (see + // test_data_overlay_flag_release_gating for the full policy). + assert_eq!( + can_read_dataset(super::FLAG_UNSTABLE_DATA_OVERLAY_FILES), + data_overlay_files_enabled() + ); assert!(can_read_dataset( super::FLAG_DELETION_FILES | super::FLAG_STABLE_ROW_IDS @@ -111,6 +176,58 @@ mod tests { assert!(!can_read_dataset(super::FLAG_UNKNOWN)); } + #[test] + fn test_data_overlay_flag_release_gating() { + // Release default (overlays disabled): the overlay flag is treated as + // unknown so the dataset is refused, while other known flags still pass. + let supported = supported_flags_when(false); + assert_eq!(supported & FLAG_UNSTABLE_DATA_OVERLAY_FILES, 0); + assert_eq!(FLAG_DELETION_FILES & !supported, 0); + assert_ne!(FLAG_UNSTABLE_DATA_OVERLAY_FILES & !supported, 0); + // Enabled (debug or env opt-in): the overlay flag is understood. + let supported = supported_flags_when(true); + assert_eq!(FLAG_UNSTABLE_DATA_OVERLAY_FILES & !supported, 0); + } + + #[test] + fn test_apply_feature_flags_sets_overlay_flag() { + use crate::format::overlay::{DataOverlayFile, OverlayCoverage}; + use crate::format::{DataFile, DataStorageFormat, Fragment}; + use arrow_schema::{Field as ArrowField, Schema as ArrowSchema}; + use lance_core::datatypes::Schema; + use roaring::RoaringBitmap; + use std::collections::HashMap; + use std::sync::Arc; + + let arrow_schema = ArrowSchema::new(vec![ArrowField::new( + "id", + arrow_schema::DataType::Int64, + false, + )]); + let schema = Schema::try_from(&arrow_schema).unwrap(); + let mut fragment = Fragment::new(0); + fragment.overlays = vec![DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("o.lance", vec![0], None), + coverage: OverlayCoverage::dense(RoaringBitmap::from_iter([0u32])), + committed_version: 1, + }]; + let mut manifest = Manifest::new( + schema, + Arc::new(vec![fragment]), + DataStorageFormat::default(), + HashMap::new(), + ); + apply_feature_flags(&mut manifest, false, false).unwrap(); + assert_ne!( + manifest.reader_feature_flags & FLAG_UNSTABLE_DATA_OVERLAY_FILES, + 0 + ); + assert_ne!( + manifest.writer_feature_flags & FLAG_UNSTABLE_DATA_OVERLAY_FILES, + 0 + ); + } + #[test] fn test_write_check() { assert!(can_write_dataset(0)); @@ -120,6 +237,13 @@ mod tests { assert!(can_write_dataset(super::FLAG_TABLE_CONFIG)); assert!(can_write_dataset(super::FLAG_BASE_PATHS)); assert!(can_write_dataset(super::FLAG_DISABLE_TRANSACTION_FILE)); + // Overlay support is gated on the build profile / env opt-in, so the + // flag is writable exactly when overlays are enabled (see + // test_data_overlay_flag_release_gating for the full policy). + assert_eq!( + can_write_dataset(super::FLAG_UNSTABLE_DATA_OVERLAY_FILES), + data_overlay_files_enabled() + ); assert!(can_write_dataset( super::FLAG_DELETION_FILES | super::FLAG_STABLE_ROW_IDS diff --git a/rust/lance-table/src/format.rs b/rust/lance-table/src/format.rs index 842c76f1e58..5a5db7919d3 100644 --- a/rust/lance-table/src/format.rs +++ b/rust/lance-table/src/format.rs @@ -7,6 +7,7 @@ use uuid::Uuid; mod fragment; mod index; mod manifest; +pub mod overlay; mod transaction; pub use crate::rowids::version::{ diff --git a/rust/lance-table/src/format/fragment.rs b/rust/lance-table/src/format/fragment.rs index e9d9ce036ee..4929f54b7ff 100644 --- a/rust/lance-table/src/format/fragment.rs +++ b/rust/lance-table/src/format/fragment.rs @@ -13,6 +13,7 @@ use lance_io::utils::CachedFileSize; use object_store::path::Path; use serde::{Deserialize, Deserializer, Serialize, Serializer}; +use super::overlay::{DataOverlayFile, sort_overlays_newest_last}; use crate::format::pb; use crate::rowids::version::{ @@ -375,6 +376,15 @@ impl DataFileFieldInterner { .into_iter() .map(|f| self.intern_data_file(f)) .collect::>()?, + overlays: { + let mut overlays = p + .overlays + .into_iter() + .map(DataOverlayFile::try_from) + .collect::>>()?; + sort_overlays_newest_last(&mut overlays); + overlays + }, deletion_file: p.deletion_file.map(DeletionFile::try_from).transpose()?, row_id_meta: p.row_id_sequence.map(RowIdMeta::try_from).transpose()?, physical_rows, @@ -483,6 +493,12 @@ pub struct Fragment { /// Files within the fragment. pub files: Vec, + /// Overlay files supplying new values for a subset of cells without + /// rewriting the base data files. Order is significant: a later entry is + /// newer than an earlier one. See [`DataOverlayFile`] for resolution rules. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub overlays: Vec, + /// Optional file with deleted local row offsets. #[serde(skip_serializing_if = "Option::is_none")] pub deletion_file: Option, @@ -510,6 +526,7 @@ impl Fragment { Self { id, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: None, @@ -549,6 +566,7 @@ impl Fragment { Self { id, files: vec![DataFile::new_legacy(path, schema, None, None)], + overlays: vec![], deletion_file: None, physical_rows, row_id_meta: None, @@ -669,6 +687,15 @@ impl TryFrom for Fragment { .into_iter() .map(DataFile::try_from) .collect::>()?, + overlays: { + let mut overlays = p + .overlays + .into_iter() + .map(DataOverlayFile::try_from) + .collect::>>()?; + sort_overlays_newest_last(&mut overlays); + overlays + }, deletion_file: p.deletion_file.map(DeletionFile::try_from).transpose()?, row_id_meta: p.row_id_sequence.map(RowIdMeta::try_from).transpose()?, physical_rows, @@ -716,10 +743,7 @@ impl From<&Fragment> for pb::DataFragment { Self { id: f.id, files: f.files.iter().map(pb::DataFile::from).collect(), - // Overlay files are not produced by this version of the library; a - // dataset that uses them sets reader feature flag 64, which is - // rejected at the feature-flag layer (see lance-table feature_flags). - overlays: vec![], + overlays: f.overlays.iter().map(pb::DataOverlayFile::from).collect(), deletion_file, row_id_sequence, physical_rows: f.physical_rows.unwrap_or_default() as u64, @@ -732,12 +756,108 @@ impl From<&Fragment> for pb::DataFragment { #[cfg(test)] mod tests { use super::*; + use crate::format::overlay::OverlayCoverage; use arrow_schema::{ DataType, Field as ArrowField, Fields as ArrowFields, Schema as ArrowSchema, }; use object_store::path::Path; + use roaring::RoaringBitmap; use serde_json::{Value, json}; + #[test] + fn test_data_overlay_roundtrip() { + // A fragment carrying a dense overlay round-trips through protobuf and + // back, and the parsed coverage bitmap is recovered per field. + let mut bitmap = RoaringBitmap::new(); + bitmap.insert(1); + bitmap.insert(3); + + let overlay = DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("overlay-0.lance", vec![3], None), + coverage: OverlayCoverage::dense(bitmap.clone()), + committed_version: 7, + }; + let mut fragment = Fragment::new(0); + fragment.files = vec![DataFile::new_legacy_from_fields( + "base.lance", + vec![1, 3], + None, + )]; + fragment.overlays = vec![overlay]; + + let proto = pb::DataFragment::from(&fragment); + assert_eq!(proto.overlays.len(), 1); + let round_tripped = Fragment::try_from(proto).unwrap(); + assert_eq!(round_tripped, fragment); + + // Dense coverage applies to every field. + let recovered = round_tripped.overlays[0].coverage_for_field(0).unwrap(); + assert_eq!(*recovered, bitmap); + assert_eq!( + *round_tripped.overlays[0].coverage_for_field(5).unwrap(), + bitmap + ); + } + + #[test] + fn test_data_overlay_sparse_per_field_coverage() { + // A sparse overlay carries one bitmap per field, recovered by position. + let name_coverage = RoaringBitmap::from_iter([2u32, 3]); + let embedding_coverage = RoaringBitmap::from_iter([1u32]); + let overlay = DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("overlay-1.lance", vec![2, 4], None), + coverage: OverlayCoverage::sparse(vec![ + name_coverage.clone(), + embedding_coverage.clone(), + ]), + committed_version: 3, + }; + let mut fragment = Fragment::new(1); + fragment.overlays = vec![overlay]; + + let round_tripped = Fragment::try_from(pb::DataFragment::from(&fragment)).unwrap(); + assert_eq!( + *round_tripped.overlays[0].coverage_for_field(0).unwrap(), + name_coverage + ); + assert_eq!( + *round_tripped.overlays[0].coverage_for_field(1).unwrap(), + embedding_coverage + ); + } + + #[test] + fn test_overlays_sorted_newest_last_on_load() { + // Overlays load stable-sorted by committed_version (newest last), with + // list position preserved as the tiebreak for equal versions. + let mk = |version: u64, field: i32| DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("o.lance", vec![field], None), + coverage: OverlayCoverage::dense(RoaringBitmap::from_iter([0u32])), + committed_version: version, + }; + let mut fragment = Fragment::new(0); + // Written out of order: v5, v2, v2 (second), v3. + fragment.overlays = vec![mk(5, 1), mk(2, 2), mk(2, 3), mk(3, 4)]; + + let loaded = Fragment::try_from(pb::DataFragment::from(&fragment)).unwrap(); + let versions: Vec = loaded + .overlays + .iter() + .map(|o| o.committed_version) + .collect(); + assert_eq!(versions, vec![2, 2, 3, 5]); + // Stable: the two v2 overlays keep their original relative order (field 2 + // before field 3). + assert_eq!( + loaded.overlays[0].data_file.fields.as_ref(), + [2i32].as_slice() + ); + assert_eq!( + loaded.overlays[1].data_file.fields.as_ref(), + [3i32].as_slice() + ); + } + #[test] fn test_new_fragment() { let path = "foobar.lance"; diff --git a/rust/lance-table/src/format/manifest.rs b/rust/lance-table/src/format/manifest.rs index 9845061b7e4..cd0a403621f 100644 --- a/rust/lance-table/src/format/manifest.rs +++ b/rust/lance-table/src/format/manifest.rs @@ -1316,6 +1316,7 @@ mod tests { vec![0, 1, 2], None, )], + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: None, @@ -1328,6 +1329,7 @@ mod tests { DataFile::new_legacy_from_fields("path2", vec![0, 1, 43], None), DataFile::new_legacy_from_fields("path3", vec![2], None), ], + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: None, diff --git a/rust/lance-table/src/format/overlay.rs b/rust/lance-table/src/format/overlay.rs new file mode 100644 index 00000000000..5e729411502 --- /dev/null +++ b/rust/lance-table/src/format/overlay.rs @@ -0,0 +1,443 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Data overlay files. +//! +//! An overlay file supplies new values for a subset of `(physical offset, field)` +//! cells within a fragment, without rewriting the fragment's base data files. See +//! the Data Overlay Files specification for the full rules; the invariants this +//! module relies on are: +//! +//! - **Physical-offset coverage.** Coverage bitmaps index *physical* row offsets +//! (positions in the base data files, counting deleted rows), so they are stable +//! across deletions, like deletion vectors. +//! - **Rank-based values.** The overlay's `data_file` stores one value column per +//! field, with no row-offset key column. Within a value column, a covered +//! offset's value sits at its **rank** — the 0-based count of set bits below it +//! in that field's coverage bitmap. +//! - **Dense vs. sparse coverage.** A dense overlay shares one bitmap across every +//! field ([`OverlayCoverage::Shared`]); a sparse overlay carries one bitmap per +//! field ([`OverlayCoverage::PerField`]). +//! - **Parse once.** Bitmaps are parsed from their 32-bit Roaring encoding a single +//! time when the fragment loads and held behind an `Arc`, so cloning a fragment +//! is cheap. +//! - **Newest-last ordering.** A fragment's overlays are stored newest-last and +//! stable-sorted by `committed_version` on load (see [`sort_overlays_newest_last`]), +//! with list position breaking ties for equal versions. When two overlays cover +//! the same `(offset, field)`, the higher `committed_version` wins. +//! - **Field tombstones.** When new base values are written for a field (a +//! DataReplacement, or an in-place column rewrite), any overlay value for that +//! field is stale and must stop shadowing the fresh base. The field is marked +//! obsolete in the overlay's `data_file.fields` with [`TOMBSTONE_FIELD_ID`] +//! (the same sentinel used for obsolete base columns) rather than physically +//! removed, so the overlay's other fields — and its coverage positions — stay +//! intact (see [`tombstone_overlay_fields`]). + +use std::sync::Arc; + +use lance_core::Error; +use lance_core::deepsize::DeepSizeOf; +use lance_core::error::Result; +use roaring::RoaringBitmap; +use serde::{Deserialize, Serialize}; + +use object_store::path::Path; + +use super::DataFile; +use crate::format::pb; + +/// Field-id sentinel marking a tombstoned (obsolete) field within an overlay's +/// `data_file.fields`. Matches the tombstone convention for obsolete columns in +/// base data files; a tombstoned field's values are ignored on read. +pub const TOMBSTONE_FIELD_ID: i32 = -2; + +/// Which `(physical offset, field)` cells a [`DataOverlayFile`] provides values +/// for. +/// +/// Bitmaps are parsed from their 32-bit Roaring encoding once when the fragment +/// is loaded and held behind an `Arc` so cloning a fragment is cheap; use +/// [`DataOverlayFile::coverage_for_field`] to obtain the one that applies to a +/// given field. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(into = "OverlayCoverageBytes", try_from = "OverlayCoverageBytes")] +pub enum OverlayCoverage { + /// A single bitmap that applies to every field in the overlay's + /// `data_file.fields` (a dense / rectangular overlay): every covered offset + /// has a value for every field. + Shared(Arc), + /// One bitmap per field, in the same order as the overlay's + /// `data_file.fields` (a sparse overlay): different fields may cover + /// different offset sets. + PerField(Vec>), +} + +/// Serialized form of [`OverlayCoverage`] — each bitmap as its 32-bit Roaring +/// byte encoding. The in-memory form parses these once at load. +#[derive(Debug, Clone, Serialize, Deserialize)] +enum OverlayCoverageBytes { + Shared(Vec), + PerField(Vec>), +} + +// The bytes come from a persisted overlay (the protobuf manifest or a +// serialized fragment), so a decode failure is on-disk corruption, not caller +// input. `path` locates the overlay's data file when known (empty on the serde +// path, which deserializes coverage in isolation). +fn deserialize_roaring(bytes: &[u8], path: &Path) -> Result { + RoaringBitmap::deserialize_from(bytes).map_err(|e| { + Error::corrupt_file( + path.clone(), + format!("failed to deserialize overlay coverage bitmap: {e}"), + ) + }) +} + +fn serialize_roaring(bitmap: &RoaringBitmap) -> Vec { + let mut bytes = Vec::with_capacity(bitmap.serialized_size()); + // Writing to a Vec is infallible. + bitmap.serialize_into(&mut bytes).unwrap(); + bytes +} + +impl From for OverlayCoverageBytes { + fn from(coverage: OverlayCoverage) -> Self { + match coverage { + OverlayCoverage::Shared(bitmap) => Self::Shared(serialize_roaring(&bitmap)), + OverlayCoverage::PerField(bitmaps) => { + Self::PerField(bitmaps.iter().map(|b| serialize_roaring(b)).collect()) + } + } + } +} + +impl TryFrom for OverlayCoverage { + type Error = Error; + + fn try_from(bytes: OverlayCoverageBytes) -> Result { + // Serde deserializes the coverage in isolation, so the owning data + // file's path is not available here. + let path = Path::default(); + Ok(match bytes { + OverlayCoverageBytes::Shared(b) => { + Self::Shared(Arc::new(deserialize_roaring(&b, &path)?)) + } + OverlayCoverageBytes::PerField(bs) => Self::PerField( + bs.iter() + .map(|b| deserialize_roaring(b, &path).map(Arc::new)) + .collect::>()?, + ), + }) + } +} + +impl DeepSizeOf for OverlayCoverage { + fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { + // The same `Arc` is shared across every clone of a + // fragment, so mark each Arc's pointer and count its heap only the first + // time it is seen — otherwise walking many fragments double-counts the + // shared bitmaps. RoaringBitmap does not expose its allocation size; its + // serialized size is a cheap, close proxy for the heap it holds. + let bitmap_heap = |bitmap: &Arc, + context: &mut lance_core::deepsize::Context| { + if context.mark_seen(Arc::as_ptr(bitmap) as usize) { + std::mem::size_of::() + bitmap.serialized_size() + } else { + 0 + } + }; + match self { + Self::Shared(bitmap) => bitmap_heap(bitmap, context), + Self::PerField(bitmaps) => { + bitmaps.capacity() * std::mem::size_of::>() + + bitmaps + .iter() + .map(|b| bitmap_heap(b, context)) + .sum::() + } + } + } +} + +impl OverlayCoverage { + /// Build a dense coverage from a single bitmap shared across every field. + pub fn dense(bitmap: RoaringBitmap) -> Self { + Self::Shared(Arc::new(bitmap)) + } + + /// Build a sparse coverage from one bitmap per field. + pub fn sparse(bitmaps: Vec) -> Self { + Self::PerField(bitmaps.into_iter().map(Arc::new).collect()) + } +} + +/// An overlay file supplies new values for a subset of `(physical offset, field)` +/// cells within a fragment, without rewriting the fragment's base data files. See +/// the [module documentation](self) for the coverage, rank, and versioning rules. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, DeepSizeOf)] +pub struct DataOverlayFile { + /// The data file storing the overlay's new cell values. + pub data_file: DataFile, + /// Which cells this overlay provides values for. + pub coverage: OverlayCoverage, + /// The dataset version at which this overlay became effective (the version of + /// the commit that introduced it, stamped at commit time and re-stamped on + /// retry). Higher wins when two overlays cover the same `(offset, field)`. + pub committed_version: u64, +} + +impl DataOverlayFile { + /// The parsed coverage bitmap that applies to the field stored at + /// `field_pos` within `data_file.fields`. + /// + /// For a dense overlay the same shared bitmap is returned for every field; + /// for a sparse overlay the per-field bitmap at `field_pos` is returned. The + /// bitmap is already parsed, so this is a cheap `Arc` clone. + pub fn coverage_for_field(&self, field_pos: usize) -> Result> { + match &self.coverage { + OverlayCoverage::Shared(bitmap) => Ok(bitmap.clone()), + OverlayCoverage::PerField(bitmaps) => { + bitmaps.get(field_pos).cloned().ok_or_else(|| { + Error::invalid_input(format!( + "overlay per-field coverage has {} bitmaps but field position {} was requested", + bitmaps.len(), + field_pos + )) + }) + } + } + } +} + +/// Stable-sort a fragment's overlays newest-last by `committed_version`. The +/// stable sort preserves list position as the tiebreak for equal versions, so +/// resolution can rely on the ordering without re-checking. See the [module +/// documentation](self) for the ordering invariant. +pub fn sort_overlays_newest_last(overlays: &mut [DataOverlayFile]) { + overlays.sort_by_key(|overlay| overlay.committed_version); +} + +/// Verify a fragment's overlays are stored newest-last (non-decreasing +/// `committed_version`), the ordering invariant readers rely on for +/// resolution. Returns an error identifying the first out-of-order pair. +/// +/// [`sort_overlays_newest_last`] normalizes on load; this is the write-side +/// guard that rejects any commit path that assembled overlays out of order. See +/// the [module documentation](self) for the ordering invariant. +pub fn verify_overlays_newest_last(overlays: &[DataOverlayFile]) -> Result<()> { + for pair in overlays.windows(2) { + if pair[0].committed_version > pair[1].committed_version { + return Err(Error::invalid_input(format!( + "overlay files must be stored newest-last, but committed_version {} precedes {}", + pair[0].committed_version, pair[1].committed_version + ))); + } + } + Ok(()) +} + +/// Tombstone `fields` across a fragment's `overlays`, dropping any overlay left +/// with no live fields. +/// +/// Called when new base values are written for those fields (a DataReplacement, +/// or an in-place column rewrite): the stale overlay values must stop shadowing +/// the fresh base. Each matching field id is replaced with [`TOMBSTONE_FIELD_ID`] +/// in place, preserving the overlay's remaining fields and its coverage positions +/// (a per-field coverage bitmap stays aligned with `data_file.fields`). An overlay +/// whose fields are now all tombstoned is removed entirely. See the [module +/// documentation](self) for the tombstone invariant. +pub fn tombstone_overlay_fields(overlays: &mut Vec, fields: &[u32]) { + for overlay in overlays.iter_mut() { + let tombstoned: Vec = overlay + .data_file + .fields + .iter() + .map(|&field| { + if field >= 0 && fields.contains(&(field as u32)) { + TOMBSTONE_FIELD_ID + } else { + field + } + }) + .collect(); + overlay.data_file.fields = tombstoned.into(); + } + overlays.retain(|overlay| { + overlay + .data_file + .fields + .iter() + .any(|&field| field != TOMBSTONE_FIELD_ID) + }); +} + +impl From<&DataOverlayFile> for pb::DataOverlayFile { + fn from(overlay: &DataOverlayFile) -> Self { + let coverage = match &overlay.coverage { + OverlayCoverage::Shared(bitmap) => { + pb::data_overlay_file::Coverage::SharedOffsetBitmap(serialize_roaring(bitmap)) + } + OverlayCoverage::PerField(bitmaps) => { + pb::data_overlay_file::Coverage::FieldCoverage(pb::FieldCoverage { + offset_bitmaps: bitmaps.iter().map(|b| serialize_roaring(b)).collect(), + }) + } + }; + Self { + data_file: Some(pb::DataFile::from(&overlay.data_file)), + coverage: Some(coverage), + committed_version: overlay.committed_version, + } + } +} + +impl TryFrom for DataOverlayFile { + type Error = Error; + + fn try_from(proto: pb::DataOverlayFile) -> Result { + let data_file = proto + .data_file + .ok_or_else(|| Error::invalid_input("DataOverlayFile is missing its data_file"))?; + let path = Path::from(data_file.path.as_str()); + let coverage = match proto.coverage { + Some(pb::data_overlay_file::Coverage::SharedOffsetBitmap(bytes)) => { + OverlayCoverage::Shared(Arc::new(deserialize_roaring(&bytes, &path)?)) + } + Some(pb::data_overlay_file::Coverage::FieldCoverage(fc)) => OverlayCoverage::PerField( + fc.offset_bitmaps + .iter() + .map(|b| deserialize_roaring(b, &path).map(Arc::new)) + .collect::>()?, + ), + None => { + return Err(Error::invalid_input( + "DataOverlayFile is missing its coverage", + )); + } + }; + Ok(Self { + data_file: DataFile::try_from(data_file)?, + coverage, + committed_version: proto.committed_version, + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_data_overlay_missing_fields_error() { + // A DataOverlayFile proto missing its coverage or data_file is rejected. + let no_coverage = pb::DataOverlayFile { + data_file: Some(pb::DataFile::from(&DataFile::new_legacy_from_fields( + "overlay.lance", + vec![3], + None, + ))), + coverage: None, + committed_version: 1, + }; + let err = DataOverlayFile::try_from(no_coverage).unwrap_err(); + assert!(err.to_string().contains("missing its coverage"), "{err}"); + + let no_data_file = pb::DataOverlayFile { + data_file: None, + coverage: Some(pb::data_overlay_file::Coverage::SharedOffsetBitmap( + serialize_roaring(&RoaringBitmap::from_iter([0u32])), + )), + committed_version: 1, + }; + let err = DataOverlayFile::try_from(no_data_file).unwrap_err(); + assert!(err.to_string().contains("missing its data_file"), "{err}"); + } + + #[test] + fn test_overlay_coverage_serde_json_roundtrip() { + // The custom serde impl round-trips through JSON for dense/sparse, + // including empty bitmaps and a zero-bitmap sparse coverage. + for coverage in [ + OverlayCoverage::dense(RoaringBitmap::from_iter([1u32, 5, 100])), + OverlayCoverage::dense(RoaringBitmap::new()), + OverlayCoverage::sparse(vec![ + RoaringBitmap::from_iter([2u32, 3]), + RoaringBitmap::new(), + ]), + OverlayCoverage::sparse(vec![]), + ] { + let json = serde_json::to_string(&coverage).unwrap(); + let back: OverlayCoverage = serde_json::from_str(&json).unwrap(); + assert_eq!(back, coverage); + } + } + + #[test] + fn test_tombstone_overlay_fields() { + // An overlay covering fields [3, 5]: replacing field 5 tombstones just + // field 5's slot and keeps field 3. An overlay covering only field 5 is + // dropped entirely. An overlay touching no replaced field is untouched. + let mut overlays = vec![ + DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("a.lance", vec![3, 5], None), + coverage: OverlayCoverage::sparse(vec![ + RoaringBitmap::from_iter([0u32]), + RoaringBitmap::from_iter([1u32]), + ]), + committed_version: 1, + }, + DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("b.lance", vec![5], None), + coverage: OverlayCoverage::dense(RoaringBitmap::from_iter([0u32])), + committed_version: 1, + }, + DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("c.lance", vec![7], None), + coverage: OverlayCoverage::dense(RoaringBitmap::from_iter([0u32])), + committed_version: 1, + }, + ]; + + tombstone_overlay_fields(&mut overlays, &[5]); + + // The single-field overlay on field 5 is gone; the others remain. + assert_eq!(overlays.len(), 2); + // Field 3 preserved, field 5 tombstoned in place (coverage stays aligned). + assert_eq!( + overlays[0].data_file.fields.as_ref(), + &[3, TOMBSTONE_FIELD_ID] + ); + // The untouched overlay keeps its field. + assert_eq!(overlays[1].data_file.fields.as_ref(), &[7]); + } + + #[test] + fn test_verify_overlays_newest_last() { + let mk = |version: u64| DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("o.lance", vec![3], None), + coverage: OverlayCoverage::dense(RoaringBitmap::from_iter([0u32])), + committed_version: version, + }; + // Non-decreasing (including equal versions) is accepted. + assert!(verify_overlays_newest_last(&[]).is_ok()); + assert!(verify_overlays_newest_last(&[mk(1), mk(2), mk(2), mk(5)]).is_ok()); + // A newer version before an older one is rejected. + let err = verify_overlays_newest_last(&[mk(2), mk(1)]).unwrap_err(); + assert!(err.to_string().contains("newest-last"), "{err}"); + } + + #[test] + fn test_coverage_for_field_out_of_bounds() { + let overlay = DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("o.lance", vec![2, 4], None), + coverage: OverlayCoverage::sparse(vec![ + RoaringBitmap::from_iter([1u32]), + RoaringBitmap::from_iter([2u32]), + ]), + committed_version: 1, + }; + assert!(overlay.coverage_for_field(0).is_ok()); + assert!(overlay.coverage_for_field(1).is_ok()); + let err = overlay.coverage_for_field(5).unwrap_err(); + assert!(err.to_string().contains("field position"), "{err}"); + } +} diff --git a/rust/lance/src/dataset/files.rs b/rust/lance/src/dataset/files.rs index 848add7e4a8..2214822ed90 100644 --- a/rust/lance/src/dataset/files.rs +++ b/rust/lance/src/dataset/files.rs @@ -1036,6 +1036,7 @@ mod tests { // No base_id -> falls back to the dataset base_uri. mk_file("c.lance", None), ], + overlays: vec![], // Deletion files also carry a base_id when they originate from a // shallow clone, and must resolve against base_paths too. deletion_file: Some(DeletionFile { diff --git a/rust/lance/src/dataset/fragment.rs b/rust/lance/src/dataset/fragment.rs index a13de636a50..4df89ab71b8 100644 --- a/rust/lance/src/dataset/fragment.rs +++ b/rust/lance/src/dataset/fragment.rs @@ -911,6 +911,16 @@ impl FileFragment { projection: &Schema, read_config: FragReadConfig, ) -> Result { + // Overlay files supply newer cell values that must be merged on read. + // Until the scan/take merge path lands (the rest of OSS-1322 / OSS-1324), + // reading a fragment that has overlays would silently return stale base + // values, so we refuse rather than serve incorrect data. + if !self.metadata.overlays.is_empty() { + return Err(Error::not_supported( + "reading fragments with data overlay files is not yet supported \ + (overlay merge is in progress)", + )); + } let open_files = self.open_readers(projection, &read_config); let deletion_vec_load = self.get_deletion_vector(); diff --git a/rust/lance/src/dataset/optimize.rs b/rust/lance/src/dataset/optimize.rs index bb51056eff5..d352659bdcb 100644 --- a/rust/lance/src/dataset/optimize.rs +++ b/rust/lance/src/dataset/optimize.rs @@ -2221,6 +2221,7 @@ mod tests { let fragment = Fragment { id: 0, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: Some(0), diff --git a/rust/lance/src/dataset/schema_evolution.rs b/rust/lance/src/dataset/schema_evolution.rs index ce32362f324..40eac95e919 100644 --- a/rust/lance/src/dataset/schema_evolution.rs +++ b/rust/lance/src/dataset/schema_evolution.rs @@ -1957,6 +1957,7 @@ mod test { Ok(Some(Fragment { files: vec![], id: 0, + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: Some(50), diff --git a/rust/lance/src/dataset/transaction.rs b/rust/lance/src/dataset/transaction.rs index 14a5b22f5ab..372adee27db 100644 --- a/rust/lance/src/dataset/transaction.rs +++ b/rust/lance/src/dataset/transaction.rs @@ -32,7 +32,8 @@ use lance_table::rowids::read_row_ids; use lance_table::{ format::{ BasePath, DataFile, DataStorageFormat, Fragment, IndexFile, IndexMetadata, Manifest, - RowDatasetVersionMeta, RowDatasetVersionRun, RowDatasetVersionSequence, RowIdMeta, pb, + RowDatasetVersionMeta, RowDatasetVersionRun, RowDatasetVersionSequence, RowIdMeta, + overlay::DataOverlayFile, pb, }, io::{ commit::CommitHandler, @@ -258,6 +259,17 @@ pub struct Transaction { #[derive(Debug, Clone, DeepSizeOf, PartialEq)] pub struct DataReplacementGroup(pub u64, pub DataFile); +/// Overlay files to append to a single fragment, in order (the last entry is +/// newest). The overlays are appended to the fragment's existing `overlays` +/// list rather than replacing it, so overlays written by concurrent commits are +/// preserved. Each overlay's `committed_version` is stamped to the new dataset +/// version at commit time (re-stamped on retry). +#[derive(Debug, Clone, DeepSizeOf, PartialEq)] +pub struct DataOverlayGroup { + pub fragment_id: u64, + pub overlays: Vec, +} + /// An entry for a map update. If value is None, the key will be removed from the map. #[derive(Debug, Clone, DeepSizeOf, PartialEq)] pub struct UpdateMapEntry { @@ -367,6 +379,11 @@ pub enum Operation { DataReplacement { replacements: Vec, }, + /// Attach overlay files to fragments, supplying new values for a subset of + /// `(physical offset, field)` cells without rewriting the fragments' base + /// data files. See [`DataOverlayFile`] and the Data Overlay Files + /// specification for resolution, coverage, and versioning rules. + DataOverlay { groups: Vec }, /// Merge a new column in /// 'fragments' is the final fragments include all data files, the new fragments must align with old ones at rows. /// 'schema' is not forced to include existed columns, which means we could use Merge to drop column data @@ -499,6 +516,7 @@ impl std::fmt::Display for Operation { Self::Project { .. } => write!(f, "Project"), Self::UpdateConfig { .. } => write!(f, "UpdateConfig"), Self::DataReplacement { .. } => write!(f, "DataReplacement"), + Self::DataOverlay { .. } => write!(f, "DataOverlay"), Self::Clone { .. } => write!(f, "Clone"), Self::UpdateMemWalState { .. } => write!(f, "UpdateMemWalState"), Self::UpdateBases { .. } => write!(f, "UpdateBases"), @@ -1345,6 +1363,8 @@ impl PartialEq for Operation { (Self::Clone { .. }, Self::UpdateBases { .. }) => { std::mem::discriminant(self) == std::mem::discriminant(other) } + (Self::DataOverlay { groups: a }, Self::DataOverlay { groups: b }) => compare_vec(a, b), + (Self::DataOverlay { .. }, _) | (_, Self::DataOverlay { .. }) => false, } } } @@ -1521,6 +1541,7 @@ impl Operation { Self::Project { .. } => "Project", Self::UpdateConfig { .. } => "UpdateConfig", Self::DataReplacement { .. } => "DataReplacement", + Self::DataOverlay { .. } => "DataOverlay", Self::UpdateMemWalState { .. } => "UpdateMemWalState", Self::Clone { .. } => "Clone", Self::UpdateBases { .. } => "UpdateBases", @@ -1924,7 +1945,20 @@ impl Transaction { return None; } if let Some(updated) = updated_fragments.iter().find(|uf| uf.id == f.id) { - Some(updated.clone()) + let mut updated = updated.clone(); + // Carry forward the fragment's current overlays (which + // may include ones added by a concurrent commit). An + // in-place column rewrite then tombstones the overlaid + // fields it rewrote, since the fresh base values + // supersede them. + updated.overlays = f.overlays.clone(); + if matches!(update_mode, Some(RewriteColumns)) { + lance_table::format::overlay::tombstone_overlay_fields( + &mut updated.overlays, + fields_modified, + ); + } + Some(updated) } else { Some(f.clone()) } @@ -2271,6 +2305,15 @@ impl Transaction { "Expected to modify the fragment but no changes were made. This means the new data files does not align with any exiting datafiles. Please check if the schema of the new data files matches the schema of the old data files including the file major and minor versions", )); } + + // New base values for these fields supersede any overlay + // still shadowing them; tombstone the overlaid fields so the + // replacement is not silently masked. + lance_table::format::overlay::tombstone_overlay_fields( + &mut new_frag.overlays, + &replaced_fields, + ); + final_fragments.push(new_frag); } @@ -2302,6 +2345,53 @@ impl Transaction { &replaced_fields, ); } + Operation::DataOverlay { groups } => { + // Stamp each overlay with the version this commit is producing. + // build_manifest re-runs on every retry with an updated + // current_manifest, so this is naturally re-stamped on retry. + let new_version = current_manifest.map_or(1, |m| m.version + 1); + + let existing_fragments = maybe_existing_fragments?; + // Multiple groups may target the same fragment; merge them in + // order rather than letting a HashMap collapse drop all but the + // last group's overlays. + let mut overlays_by_fragment: HashMap> = HashMap::new(); + for group in groups { + overlays_by_fragment + .entry(group.fragment_id) + .or_default() + .extend(group.overlays.iter()); + } + + // Every group must target an existing fragment. Build a set of + // existing ids once so this is O(groups + fragments) rather than + // O(groups * fragments). + let existing_fragment_ids: HashSet = + existing_fragments.iter().map(|f| f.id).collect(); + for fragment_id in overlays_by_fragment.keys() { + if !existing_fragment_ids.contains(fragment_id) { + return Err(Error::invalid_input(format!( + "DataOverlay targets fragment {fragment_id}, which does not exist" + ))); + } + } + + for fragment in existing_fragments { + let mut fragment = fragment.clone(); + if let Some(new_overlays) = overlays_by_fragment.get(&fragment.id) { + // Appended (not replaced) so concurrently-written overlays + // survive; later entries are newer. + fragment + .overlays + .extend(new_overlays.iter().map(|&overlay| { + let mut overlay = overlay.clone(); + overlay.committed_version = new_version; + overlay + })); + } + final_fragments.push(fragment); + } + } Operation::UpdateMemWalState { merged_generations } => { update_mem_wal_index_merged_generations( &mut final_indices, @@ -2322,6 +2412,15 @@ impl Transaction { // Clean up data files that only contain tombstoned fields Self::remove_tombstoned_data_files(&mut final_fragments); + // Enforce the newest-last overlay ordering invariant at the write + // boundary. Load normalizes with a sort; this rejects any commit path + // that assembled a fragment's overlays out of order. + for fragment in &final_fragments { + if !fragment.overlays.is_empty() { + lance_table::format::overlay::verify_overlays_newest_last(&fragment.overlays)?; + } + } + let user_requested_version = match (&config.storage_format, config.use_legacy_format) { (Some(storage_format), _) => Some(storage_format.lance_file_version()?), (None, Some(true)) => Some(LanceFileVersion::Legacy), @@ -3063,6 +3162,34 @@ impl TryFrom for DataReplacementGroup { } } +impl From<&DataOverlayGroup> for pb::transaction::DataOverlayGroup { + fn from(group: &DataOverlayGroup) -> Self { + Self { + fragment_id: group.fragment_id, + overlays: group + .overlays + .iter() + .map(pb::DataOverlayFile::from) + .collect(), + } + } +} + +impl TryFrom for DataOverlayGroup { + type Error = Error; + + fn try_from(message: pb::transaction::DataOverlayGroup) -> Result { + Ok(Self { + fragment_id: message.fragment_id, + overlays: message + .overlays + .into_iter() + .map(DataOverlayFile::try_from) + .collect::>>()?, + }) + } +} + impl TryFrom for Transaction { type Error = Error; @@ -3345,16 +3472,14 @@ impl TryFrom for Transaction { })) => Operation::UpdateBases { new_bases: new_bases.into_iter().map(BasePath::from).collect(), }, - Some(pb::transaction::Operation::DataOverlay(_)) => { - // Overlay files are not supported by this version of the library. - // A dataset that uses them sets reader feature flag 64, which is - // already rejected at the feature-flag layer; reject here too so a - // transaction referencing the operation can never be applied. - return Err(Error::not_supported( - "data overlay files are not supported by this version of Lance \ - (reader feature flag 64)", - )); - } + Some(pb::transaction::Operation::DataOverlay(pb::transaction::DataOverlay { + groups, + })) => Operation::DataOverlay { + groups: groups + .into_iter() + .map(DataOverlayGroup::try_from) + .collect::>>()?, + }, None => { return Err(Error::internal( "Transaction message did not contain an operation".to_string(), @@ -3625,6 +3750,14 @@ impl From<&Transaction> for pb::Transaction { .collect(), }) } + Operation::DataOverlay { groups } => { + pb::transaction::Operation::DataOverlay(pb::transaction::DataOverlay { + groups: groups + .iter() + .map(pb::transaction::DataOverlayGroup::from) + .collect(), + }) + } Operation::UpdateMemWalState { merged_generations } => { pb::transaction::Operation::UpdateMemWalState(pb::transaction::UpdateMemWalState { merged_generations: merged_generations @@ -3903,6 +4036,7 @@ mod tests { use lance_core::{ROW_ADDR, ROW_CREATED_AT_VERSION, ROW_LAST_UPDATED_AT_VERSION}; use lance_file::version::LanceFileVersion; use lance_io::utils::CachedFileSize; + use lance_table::format::overlay::OverlayCoverage; use lance_table::format::{ RowDatasetVersionMeta, RowDatasetVersionRun, RowDatasetVersionSequence, RowIdMeta, }; @@ -4214,6 +4348,7 @@ mod tests { physical_rows: Some(100), row_id_meta: None, files: vec![], + overlays: vec![], deletion_file: None, last_updated_at_version_meta: None, created_at_version_meta: None, @@ -4246,6 +4381,7 @@ mod tests { physical_rows: Some(50), row_id_meta: Some(RowIdMeta::Inline(serialized)), files: vec![], + overlays: vec![], deletion_file: None, last_updated_at_version_meta: None, created_at_version_meta: None, @@ -4278,6 +4414,7 @@ mod tests { physical_rows: Some(50), // More physical rows than existing row IDs row_id_meta: Some(RowIdMeta::Inline(serialized)), files: vec![], + overlays: vec![], deletion_file: None, last_updated_at_version_meta: None, created_at_version_meta: None, @@ -4313,6 +4450,7 @@ mod tests { physical_rows: Some(50), // Less physical rows than existing row IDs row_id_meta: Some(RowIdMeta::Inline(serialized)), files: vec![], + overlays: vec![], deletion_file: None, last_updated_at_version_meta: None, created_at_version_meta: None, @@ -4341,6 +4479,7 @@ mod tests { physical_rows: Some(30), // No existing row IDs row_id_meta: None, files: vec![], + overlays: vec![], deletion_file: None, last_updated_at_version_meta: None, created_at_version_meta: None, @@ -4350,6 +4489,7 @@ mod tests { physical_rows: Some(25), // Partial existing row IDs row_id_meta: Some(RowIdMeta::Inline(serialized)), files: vec![], + overlays: vec![], deletion_file: None, last_updated_at_version_meta: None, created_at_version_meta: None, @@ -4394,6 +4534,7 @@ mod tests { physical_rows: None, row_id_meta: None, files: vec![], + overlays: vec![], deletion_file: None, last_updated_at_version_meta: None, created_at_version_meta: None, @@ -4901,6 +5042,7 @@ mod tests { let fragment = Fragment { id: 1, files: vec![data_file], + overlays: vec![], deletion_file: None, row_id_meta, physical_rows: Some(5), @@ -5170,6 +5312,7 @@ mod tests { None, )], physical_rows: Some(10), + overlays: vec![], deletion_file: None, row_id_meta: None, last_updated_at_version_meta: None, @@ -5261,6 +5404,7 @@ mod tests { let prev_fragment = Fragment { id: 0, files: vec![mk_file("before.lance")], + overlays: vec![], deletion_file: None, row_id_meta, physical_rows: Some(5), @@ -5333,6 +5477,7 @@ mod tests { let prev_fragment = Fragment { id: 0, files: vec![data_file.clone()], + overlays: vec![], deletion_file: None, row_id_meta: row_id_meta.clone(), physical_rows: Some(5), @@ -5352,6 +5497,7 @@ mod tests { let merged_fragment = Fragment { id: 0, files: vec![data_file], + overlays: vec![], deletion_file: None, row_id_meta, physical_rows: Some(5), @@ -5401,6 +5547,7 @@ mod tests { let prev_fragment = Fragment { id: 0, files: vec![mk_file("before.lance")], + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: Some(5), @@ -5465,6 +5612,7 @@ mod tests { let existing_fragment = Fragment { id: 0, files: vec![mk_file("existing.lance")], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&row_ids_0))), physical_rows: Some(3), @@ -5487,6 +5635,7 @@ mod tests { let new_fragment = Fragment { id: 1, files: vec![mk_file("new.lance")], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&row_ids_1))), physical_rows: Some(4), @@ -5549,6 +5698,7 @@ mod tests { let existing_fragment = Fragment { id: 1, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&existing_seq))), physical_rows: Some(3), @@ -5562,6 +5712,7 @@ mod tests { let new_fragment = Fragment { id: 10, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&new_seq))), physical_rows: Some(2), @@ -5604,6 +5755,7 @@ mod tests { Fragment { id: 1, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&frag_a_seq))), physical_rows: Some(2), @@ -5615,6 +5767,7 @@ mod tests { Fragment { id: 2, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&frag_b_seq))), physical_rows: Some(3), @@ -5630,6 +5783,7 @@ mod tests { let new_fragment = Fragment { id: 10, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&new_seq))), physical_rows: Some(2), @@ -5671,6 +5825,7 @@ mod tests { let existing_fragment = Fragment { id: 1, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&existing_seq))), physical_rows: Some(2), @@ -5685,6 +5840,7 @@ mod tests { let new_fragment = Fragment { id: 10, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&new_seq))), physical_rows: Some(2), @@ -5729,6 +5885,7 @@ mod tests { let existing_fragment = Fragment { id: 1, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&existing_seq))), physical_rows: Some(2), @@ -5742,6 +5899,7 @@ mod tests { let new_fragment = Fragment { id: 20, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&new_seq))), physical_rows: Some(4), @@ -5775,6 +5933,7 @@ mod tests { let existing_fragment = Fragment { id: 1, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&existing_seq))), physical_rows: Some(2), @@ -5786,6 +5945,7 @@ mod tests { let new_fragment = Fragment { id: 10, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&new_seq))), physical_rows: Some(1), @@ -5814,6 +5974,7 @@ mod tests { let existing_fragment = Fragment { id: 1, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&existing_seq))), physical_rows: Some(2), @@ -5824,6 +5985,7 @@ mod tests { let new_fragment = Fragment { id: 10, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: Some(3), @@ -5854,6 +6016,7 @@ mod tests { let existing_fragment = Fragment { id: 1, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&existing_seq))), physical_rows: Some(2), @@ -5867,6 +6030,7 @@ mod tests { let new_fragment = Fragment { id: 10, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&new_seq))), physical_rows: Some(1), @@ -5908,6 +6072,7 @@ mod tests { let in_range_frag = Fragment { id: 1, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&in_range_seq))), physical_rows: Some(2), @@ -5928,6 +6093,7 @@ mod tests { let out_of_range_frag = Fragment { id: 2, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&out_of_range_seq))), physical_rows: Some(2), @@ -5942,6 +6108,7 @@ mod tests { let new_frag = Fragment { id: 10, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&new_seq))), physical_rows: Some(2), @@ -5980,6 +6147,7 @@ mod tests { let existing = Fragment { id: 1, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&seq))), physical_rows: Some(3), @@ -5992,6 +6160,7 @@ mod tests { let new_frag = Fragment { id: 10, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&new_seq))), physical_rows: Some(2), @@ -6040,6 +6209,7 @@ mod tests { let src_frag = Fragment { id: 1, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&src_seq))), physical_rows: Some(100), @@ -6054,6 +6224,7 @@ mod tests { let new_frag = Fragment { id: 10, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&new_seq))), physical_rows: Some(100), @@ -6101,6 +6272,7 @@ mod tests { Fragment { id: 1, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&seq_a))), physical_rows: Some(3), @@ -6112,6 +6284,7 @@ mod tests { Fragment { id: 2, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&seq_b))), physical_rows: Some(3), @@ -6127,6 +6300,7 @@ mod tests { let new_frag = Fragment { id: 10, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&new_seq))), physical_rows: Some(2), @@ -6195,20 +6369,270 @@ mod tests { } #[test] - fn test_data_overlay_operation_rejected() { - // Overlay files are not supported by this version of the library. A - // transaction carrying the DataOverlay operation must be rejected rather - // than silently ignored, mirroring the feature-flag-64 rejection. + fn test_data_overlay_operation_roundtrips() { + // A DataOverlay operation survives the protobuf round-trip, preserving + // the target fragment, the overlay's coverage, and its committed_version. + let mut bitmap = roaring::RoaringBitmap::new(); + bitmap.insert(1); + bitmap.insert(4); + let overlay = DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("overlay-0.lance", vec![3], None), + coverage: OverlayCoverage::dense(bitmap.clone()), + committed_version: 6, + }; + let pb_overlay = pb::DataOverlayFile::from(&overlay); + let message = pb::Transaction { read_version: 1, uuid: Uuid::new_v4().to_string(), operation: Some(pb::transaction::Operation::DataOverlay( - pb::transaction::DataOverlay { groups: vec![] }, + pb::transaction::DataOverlay { + groups: vec![pb::transaction::DataOverlayGroup { + fragment_id: 7, + overlays: vec![pb_overlay], + }], + }, )), ..Default::default() }; - let result = Transaction::try_from(message); - assert!(matches!(result, Err(Error::NotSupported { .. }))); + let txn = Transaction::try_from(message).unwrap(); + match txn.operation { + Operation::DataOverlay { groups } => { + assert_eq!(groups.len(), 1); + assert_eq!(groups[0].fragment_id, 7); + assert_eq!(groups[0].overlays.len(), 1); + assert_eq!(groups[0].overlays[0].committed_version, 6); + assert_eq!( + *groups[0].overlays[0].coverage_for_field(0).unwrap(), + bitmap + ); + } + other => panic!("expected DataOverlay, got {other:?}"), + } + } + + fn overlay_with_field(field: i32, committed_version: u64) -> DataOverlayFile { + DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("o.lance", vec![field], None), + coverage: OverlayCoverage::dense(roaring::RoaringBitmap::from_iter([0u32])), + committed_version, + } + } + + #[test] + fn test_data_overlay_build_manifest_multi_fragment() { + // Overlays targeting two distinct fragments are each applied and stamped. + // A targeted fragment already carrying an overlay (committed at v3) gets + // the new overlay appended and stamped while its existing overlay is + // preserved, and a fragment the operation does not target is passed + // through with its existing overlays untouched. + let mut frag0 = Fragment::new(0); + frag0.overlays = vec![overlay_with_field(5, 3)]; // targeted, pre-existing at v3 + let frag1 = Fragment::new(1); + let mut frag2 = Fragment::new(2); + frag2.overlays = vec![overlay_with_field(9, 3)]; // untargeted, committed at v3 + let schema = ArrowSchema::new(vec![ArrowField::new("id", DataType::Int32, false)]); + let mut manifest = Manifest::new( + LanceSchema::try_from(&schema).unwrap(), + Arc::new(vec![frag0, frag1, frag2]), + lance_table::format::DataStorageFormat::new(LanceFileVersion::V2_0), + HashMap::new(), + ); + // The pre-existing overlays were committed at v3, so the current + // manifest must be at least that version; the new commit then stamps + // its overlay at v4, keeping the fragment's overlays newest-last. + manifest.version = 3; + + let txn = Transaction::new( + manifest.version, + Operation::DataOverlay { + groups: vec![ + DataOverlayGroup { + fragment_id: 0, + overlays: vec![overlay_with_field(1, 0)], + }, + DataOverlayGroup { + fragment_id: 1, + overlays: vec![overlay_with_field(2, 0)], + }, + ], + }, + None, + ); + + let (result, _) = txn + .build_manifest( + Some(&manifest), + vec![], + "txn", + &ManifestWriteConfig::default(), + ) + .unwrap(); + + let frag = |id: u64| { + result + .fragments + .iter() + .find(|f| f.id == id) + .unwrap_or_else(|| panic!("fragment {id} missing from result")) + }; + // The already-overlaid target keeps its v3 overlay and appends the new + // one, stamped to the new version. + assert_eq!(frag(0).overlays.len(), 2); + assert_eq!(frag(0).overlays[0].committed_version, 3); + assert_eq!(frag(0).overlays[1].committed_version, result.version); + // The fresh target gets its overlay, stamped to the new version. + assert_eq!(frag(1).overlays.len(), 1); + assert_eq!(frag(1).overlays[0].committed_version, result.version); + // The untargeted fragment is unchanged: same overlay, original version. + assert_eq!(frag(2).overlays.len(), 1); + assert_eq!(frag(2).overlays[0].committed_version, 3); + assert!(result.version > manifest.version); + } + + #[test] + fn test_data_replacement_tombstones_overlaid_fields() { + // A DataReplacement writing new base values for field 5 must stop any + // overlay from shadowing those cells: field 5 is tombstoned in place + // (preserving the overlay's field 3), and an overlay covering only field + // 5 is dropped entirely. + 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), + ]; + fragment.overlays = vec![ + DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("o35.lance", vec![3, 5], None), + coverage: OverlayCoverage::sparse(vec![ + roaring::RoaringBitmap::from_iter([0u32]), + roaring::RoaringBitmap::from_iter([0u32]), + ]), + committed_version: 3, + }, + DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("o5.lance", vec![5], None), + coverage: OverlayCoverage::dense(roaring::RoaringBitmap::from_iter([0u32])), + committed_version: 3, + }, + ]; + + let schema = ArrowSchema::new(vec![ArrowField::new("id", DataType::Int32, false)]); + let manifest = Manifest::new( + LanceSchema::try_from(&schema).unwrap(), + Arc::new(vec![fragment]), + lance_table::format::DataStorageFormat::new(LanceFileVersion::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), + )], + }, + None, + ); + + let (result, _) = txn + .build_manifest( + Some(&manifest), + vec![], + "txn", + &ManifestWriteConfig::default(), + ) + .unwrap(); + + let frag = &result.fragments[0]; + // The base data file for field 5 was swapped in. + assert!(frag.files.iter().any(|f| f.path == "f5-new.lance")); + // The [3, 5] overlay keeps field 3 and tombstones field 5; the [5]-only + // overlay is dropped. + assert_eq!(frag.overlays.len(), 1); + assert_eq!(frag.overlays[0].data_file.fields.as_ref(), &[3, -2]); + } + + #[test] + fn test_data_overlay_build_manifest_merges_duplicate_groups() { + // Two groups targeting the same fragment must both survive (a HashMap + // collapse would have dropped the first). + let manifest = sample_manifest(); + let txn = Transaction::new( + manifest.version, + Operation::DataOverlay { + groups: vec![ + DataOverlayGroup { + fragment_id: 0, + overlays: vec![overlay_with_field(1, 0)], + }, + DataOverlayGroup { + fragment_id: 0, + overlays: vec![overlay_with_field(2, 0)], + }, + ], + }, + None, + ); + + let (result, _) = txn + .build_manifest( + Some(&manifest), + vec![], + "txn", + &ManifestWriteConfig::default(), + ) + .unwrap(); + + let overlays = &result.fragments[0].overlays; + assert_eq!(overlays.len(), 2); + assert_eq!(overlays[0].data_file.fields.as_ref(), [1i32].as_slice()); + assert_eq!(overlays[1].data_file.fields.as_ref(), [2i32].as_slice()); + } + + #[test] + fn test_data_overlay_build_manifest_rejects_unknown_fragment() { + let manifest = sample_manifest(); + let txn = Transaction::new( + manifest.version, + Operation::DataOverlay { + groups: vec![DataOverlayGroup { + fragment_id: 99, + overlays: vec![overlay_with_field(1, 0)], + }], + }, + None, + ); + let err = txn + .build_manifest( + Some(&manifest), + vec![], + "txn", + &ManifestWriteConfig::default(), + ) + .unwrap_err(); + assert!(err.to_string().contains("does not exist"), "{err}"); + } + + #[test] + fn test_data_overlay_operation_eq() { + let overlay = |field: i32| Operation::DataOverlay { + groups: vec![DataOverlayGroup { + fragment_id: 0, + overlays: vec![overlay_with_field(field, 1)], + }], + }; + // Reflexive and value-based (the arm previously returned false for self). + assert_eq!(overlay(1), overlay(1)); + assert_ne!(overlay(1), overlay(2)); + // Not equal to a different operation kind (previously returned true vs Rewrite). + let rewrite = Operation::Rewrite { + groups: vec![], + rewritten_indices: vec![], + frag_reuse_index: None, + }; + assert_ne!(overlay(1), rewrite); } } diff --git a/rust/lance/src/dataset/write.rs b/rust/lance/src/dataset/write.rs index 80afa49ffdd..b35cfeb3eed 100644 --- a/rust/lance/src/dataset/write.rs +++ b/rust/lance/src/dataset/write.rs @@ -3980,6 +3980,7 @@ mod tests { let fragments = vec![Fragment { id: 0, files: vec![external_file, local_file], + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: Some(0), @@ -4060,6 +4061,7 @@ mod tests { let fragments = vec![Fragment { id: 0, files: vec![base1_file, base2_file, unknown_file], + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: Some(0), diff --git a/rust/lance/src/dataset/write/commit.rs b/rust/lance/src/dataset/write/commit.rs index f8d09d3c55e..e0e9b32e998 100644 --- a/rust/lance/src/dataset/write/commit.rs +++ b/rust/lance/src/dataset/write/commit.rs @@ -549,6 +549,7 @@ mod tests { file_size_bytes: CachedFileSize::new(100), base_id: None, }], + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: Some(10), diff --git a/rust/lance/src/io/commit.rs b/rust/lance/src/io/commit.rs index 5b2c0ac2944..35cd275b8ee 100644 --- a/rust/lance/src/io/commit.rs +++ b/rust/lance/src/io/commit.rs @@ -1690,6 +1690,7 @@ mod tests { DataFile::new_legacy_from_fields("path1", vec![0, 1, 2], None), DataFile::new_legacy_from_fields("unused", vec![9], None), ], + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: None, @@ -1702,6 +1703,7 @@ mod tests { DataFile::new_legacy_from_fields("path2", vec![0, 1, 2], None), DataFile::new_legacy_from_fields("path3", vec![2], None), ], + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: None, @@ -1739,6 +1741,7 @@ mod tests { vec![0, 1, 10], None, )], + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: None, @@ -1751,6 +1754,7 @@ mod tests { DataFile::new_legacy_from_fields("path2", vec![0, 1, 2], None), DataFile::new_legacy_from_fields("path3", vec![10], None), ], + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: None, @@ -1841,6 +1845,7 @@ mod tests { let fragment = Fragment { id: 0, files: vec![data_file], + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: Some(100), diff --git a/rust/lance/src/io/commit/conflict_resolver.rs b/rust/lance/src/io/commit/conflict_resolver.rs index d95821dd130..7be9c37230e 100644 --- a/rust/lance/src/io/commit/conflict_resolver.rs +++ b/rust/lance/src/io/commit/conflict_resolver.rs @@ -7,7 +7,7 @@ use crate::index::mem_wal::{load_mem_wal_index_details, new_mem_wal_index_meta}; use crate::io::deletion::read_dataset_deletion_file; use crate::{ Dataset, - dataset::transaction::{Operation, Transaction, UpdateMode}, + dataset::transaction::{DataOverlayGroup, Operation, Transaction, UpdateMode}, }; use futures::{StreamExt, TryStreamExt}; use lance_core::{Error, Result, utils::deletion::DeletionVector}; @@ -15,7 +15,9 @@ use lance_index::frag_reuse::FRAG_REUSE_INDEX_NAME; use lance_index::mem_wal::{MEM_WAL_INDEX_NAME, MergedGeneration}; use lance_select::{RowAddrTreeMap, RowSetOps}; use lance_table::format::IndexMetadata; +use lance_table::format::overlay::OverlayCoverage; use lance_table::{format::Fragment, io::deletion::write_deletion_file}; +use roaring::RoaringBitmap; use std::{ borrow::Cow, collections::{HashMap, HashSet}, @@ -137,6 +139,21 @@ impl<'a> TransactionRebase<'a> { conflicting_mem_wal_merged_gens: Vec::new(), }) } + Operation::DataOverlay { groups } => { + let modified_fragment_ids = + groups.iter().map(|g| g.fragment_id).collect::>(); + let initial_fragments = + initial_fragments_for_rebase(dataset, &transaction, &modified_fragment_ids) + .await; + Ok(Self { + transaction, + affected_rows, + initial_fragments, + modified_fragment_ids, + conflicting_frag_reuse_indices: Vec::new(), + conflicting_mem_wal_merged_gens: Vec::new(), + }) + } Operation::Merge { fragments, .. } => { let modified_fragment_ids = fragments.iter().map(|f| f.id).collect::>(); let initial_fragments = @@ -219,6 +236,9 @@ impl<'a> TransactionRebase<'a> { Operation::DataReplacement { .. } => { self.check_data_replacement_txn(other_transaction, other_version) } + Operation::DataOverlay { .. } => { + self.check_data_overlay_txn(other_transaction, other_version) + } Operation::Merge { .. } => self.check_merge_txn(other_transaction, other_version), Operation::Restore { .. } => self.check_restore_txn(other_transaction, other_version), Operation::ReserveFragments { .. } => { @@ -251,6 +271,10 @@ impl<'a> TransactionRebase<'a> { | Operation::Project { .. } | Operation::Append { .. } | Operation::UpdateConfig { .. } + // A concurrent overlay is inert against the rows we delete + // (deletions take precedence over overlays) and otherwise + // preserves physical offsets, so it never conflicts. + | Operation::DataOverlay { .. } | Operation::UpdateBases { .. } => Ok(()), Operation::Rewrite { groups, .. } => { if groups @@ -346,6 +370,8 @@ impl<'a> TransactionRebase<'a> { if let Operation::Update { inserted_rows_filter: self_inserted_rows_filter, merged_generations: self_merged_generations, + new_fragments: self_new_fragments, + update_mode: self_update_mode, .. } = &self.transaction.operation { @@ -399,6 +425,49 @@ impl<'a> TransactionRebase<'a> { | Operation::Clone { .. } | Operation::UpdateConfig { .. } | Operation::UpdateBases { .. } => Ok(()), + Operation::DataOverlay { groups } => { + // Our update recomputed rows from the pre-overlay base, so if + // it commits over an overlay it would silently undo the + // overlay's values for any cell it recomputed. A row-moving + // update (RewriteRows) relocates the rows it touches out to + // new fragments; only the rows it actually moved lose their + // overlay, so we conflict only when the moved rows intersect + // the overlay's coverage. An in-place column rewrite + // (RewriteColumns) preserves offsets and just tombstones the + // overlaid fields at build time, so it never conflicts. + let moves_rows = !self_new_fragments.is_empty() + && matches!(self_update_mode, Some(UpdateMode::RewriteRows) | None); + if !moves_rows { + return Ok(()); + } + // `affected_rows` holds the physical offsets (per fragment) + // this update moved. The overlay's coverage is in the same + // physical-offset space, so we can intersect the two in + // memory. Without affected rows we cannot be precise, so we + // fall back to a fragment-granular conflict. + for group in groups { + if !self.modified_fragment_ids.contains(&group.fragment_id) { + continue; + } + let Some(affected_rows) = self.affected_rows else { + return Err( + self.retryable_conflict_err(other_transaction, other_version) + ); + }; + let Some(moved) = + affected_rows.get_fragment_bitmap(group.fragment_id as u32) + else { + continue; + }; + let coverage = overlay_group_coverage(group); + if !(moved & &coverage).is_empty() { + return Err( + self.retryable_conflict_err(other_transaction, other_version) + ); + } + } + Ok(()) + } Operation::Append { .. } => { // If current transaction has primary key conflict detection, // we can't safely commit against an Append because we don't @@ -514,6 +583,10 @@ impl<'a> TransactionRebase<'a> { match &other_transaction.operation { Operation::Append { .. } | Operation::Clone { .. } + // An overlay committed after this index's version is newer than + // the index; the query path excludes its covered cells via the + // version gate, so the build does not conflict. + | Operation::DataOverlay { .. } | Operation::UpdateBases { .. } => Ok(()), Operation::CreateIndex { new_indices: created_indices, @@ -695,6 +768,20 @@ impl<'a> TransactionRebase<'a> { Ok(()) } } + Operation::DataOverlay { groups } => { + // Rewriting a fragment changes its physical row addresses, so + // an overlay addressed by physical offset on that fragment is + // invalidated and must be re-applied against the new base. + if groups + .iter() + .map(|g| g.fragment_id) + .any(|id| self.modified_fragment_ids.contains(&id)) + { + Err(self.retryable_conflict_err(other_transaction, other_version)) + } else { + Ok(()) + } + } Operation::Rewrite { groups, frag_reuse_index: committed_fri, @@ -874,6 +961,7 @@ impl<'a> TransactionRebase<'a> { | Operation::CreateIndex { .. } | Operation::Rewrite { .. } | Operation::DataReplacement { .. } + | Operation::DataOverlay { .. } | Operation::Merge { .. } | Operation::Restore { .. } | Operation::ReserveFragments { .. } @@ -907,7 +995,8 @@ impl<'a> TransactionRebase<'a> { | Operation::Merge { .. } | Operation::UpdateConfig { .. } | Operation::Clone { .. } - | Operation::DataReplacement { .. } => Ok(()), + | Operation::DataReplacement { .. } + | Operation::DataOverlay { .. } => Ok(()), } } @@ -923,6 +1012,9 @@ impl<'a> TransactionRebase<'a> { | Operation::UpdateConfig { .. } | Operation::ReserveFragments { .. } | Operation::Project { .. } + // Both a column replacement and an overlay preserve physical row + // addresses; the overlay is newer and wins its covered cells. + | Operation::DataOverlay { .. } | Operation::UpdateBases { .. } => Ok(()), Operation::Merge { .. } => { // Merge rewrites the whole fragment list; always conflict @@ -1055,6 +1147,119 @@ impl<'a> TransactionRebase<'a> { } } + /// Conflict checks for our DataOverlay transaction against a concurrent one. + /// + /// Overlays are intentionally permissive (see the Data Overlay Files spec): + /// they stack with other overlays and tolerate appends, index builds, data + /// replacement, deletes, and in-place column rewrites (Update with + /// `RewriteColumns`), because overlay coverage is addressed by physical offset + /// and the version gate keeps indexes correct. A concurrent operation + /// conflicts when it takes precedence over the overlay for cells the overlay + /// covers, dropping the overlay's values: retryably when it rewrites the + /// physical layout of one of our fragments (Rewrite, Merge) or re-creates the + /// covered rows from the pre-overlay base (a row-moving Update — checked + /// row-by-row in `finish_data_overlay`), or removes an overlaid fragment + /// outright (a Delete / Update that drops the fragment); and incompatibly for + /// whole-dataset replacements (Overwrite / Restore) and MemWAL state updates + /// (UpdateMemWalState), which do not rebase against data operations. + fn check_data_overlay_txn( + &mut self, + other_transaction: &Transaction, + other_version: u64, + ) -> Result<()> { + match &other_transaction.operation { + Operation::Append { .. } + | Operation::CreateIndex { .. } + | Operation::ReserveFragments { .. } + | Operation::Project { .. } + | Operation::UpdateConfig { .. } + | Operation::UpdateBases { .. } + | Operation::Clone { .. } + | Operation::DataReplacement { .. } + | Operation::DataOverlay { .. } => Ok(()), + // A concurrent Delete only tombstones rows via a deletion vector, + // which preserves physical offsets; the overlay value for a deleted + // offset is simply inert. Conflict only if the whole overlaid + // fragment was removed, orphaning the overlay. + Operation::Delete { + deleted_fragment_ids, + .. + } => { + if deleted_fragment_ids + .iter() + .any(|id| self.modified_fragment_ids.contains(id)) + { + Err(self.retryable_conflict_err(other_transaction, other_version)) + } else { + Ok(()) + } + } + // A concurrent Update that removed an overlaid fragment orphans the + // overlay outright — conflict. A row-moving update (RewriteRows) + // deletes the rows it touches and re-creates them in new fragments; + // the update took precedence and the re-created rows were computed + // from the pre-overlay base, so the overlay's values for those cells + // are lost. That is a per-row problem, not an offset one: only the + // moved rows are affected. Comparing the moved rows against the + // overlay's coverage needs the update's deletion vectors, so we mark + // the fragment here and verify row-by-row in `finish_data_overlay`. + // An in-place column rewrite (RewriteColumns) preserves rows and just + // tombstones the overlaid fields at build time, so it never conflicts. + Operation::Update { + removed_fragment_ids, + updated_fragments, + new_fragments, + update_mode, + .. + } => { + let removed_ours = removed_fragment_ids + .iter() + .any(|id| self.modified_fragment_ids.contains(id)); + if removed_ours { + return Err(self.retryable_conflict_err(other_transaction, other_version)); + } + let moves_rows = !new_fragments.is_empty() + && matches!(update_mode, Some(UpdateMode::RewriteRows) | None); + if moves_rows { + for updated in updated_fragments { + if let Some((_, needs_row_check)) = + self.initial_fragments.get_mut(&updated.id) + { + *needs_row_check = true; + } + } + } + Ok(()) + } + Operation::Rewrite { groups, .. } => { + // A rewrite (compaction / fold) of a fragment we are overlaying + // changes its physical row addresses, so our offsets would be + // invalid. Conflict only if it touches one of our fragments. + let touches_our_fragment = groups + .iter() + .flat_map(|g| g.old_fragments.iter()) + .any(|f| self.modified_fragment_ids.contains(&f.id)); + if touches_our_fragment { + Err(self.retryable_conflict_err(other_transaction, other_version)) + } else { + Ok(()) + } + } + Operation::Merge { .. } => { + // Merge rewrites the whole fragment list; always conflict. + Err(self.retryable_conflict_err(other_transaction, other_version)) + } + // Overwrite/Restore replace the dataset; UpdateMemWalState does not + // rebase against data operations (mirroring check_update_mem_wal_state_txn, + // which likewise treats a concurrent DataOverlay as incompatible). + Operation::Overwrite { .. } + | Operation::Restore { .. } + | Operation::UpdateMemWalState { .. } => { + Err(self.incompatible_conflict_err(other_transaction, other_version)) + } + } + } + fn check_merge_txn( &mut self, other_transaction: &Transaction, @@ -1072,7 +1277,8 @@ impl<'a> TransactionRebase<'a> { | Operation::Delete { .. } | Operation::Rewrite { .. } | Operation::Merge { .. } - | Operation::DataReplacement { .. } => { + | Operation::DataReplacement { .. } + | Operation::DataOverlay { .. } => { Err(self.retryable_conflict_err(other_transaction, other_version)) } Operation::Overwrite { .. } @@ -1096,6 +1302,7 @@ impl<'a> TransactionRebase<'a> { | Operation::CreateIndex { .. } | Operation::Rewrite { .. } | Operation::DataReplacement { .. } + | Operation::DataOverlay { .. } | Operation::Merge { .. } | Operation::Restore { .. } | Operation::ReserveFragments { .. } @@ -1124,6 +1331,7 @@ impl<'a> TransactionRebase<'a> { | Operation::CreateIndex { .. } | Operation::Rewrite { .. } | Operation::DataReplacement { .. } + | Operation::DataOverlay { .. } | Operation::Merge { .. } | Operation::ReserveFragments { .. } | Operation::Update { .. } @@ -1148,6 +1356,7 @@ impl<'a> TransactionRebase<'a> { | Operation::UpdateConfig { .. } | Operation::CreateIndex { .. } | Operation::DataReplacement { .. } + | Operation::DataOverlay { .. } | Operation::Rewrite { .. } | Operation::Clone { .. } | Operation::ReserveFragments { .. } @@ -1212,6 +1421,7 @@ impl<'a> TransactionRebase<'a> { | Operation::CreateIndex { .. } | Operation::Rewrite { .. } | Operation::DataReplacement { .. } + | Operation::DataOverlay { .. } | Operation::Merge { .. } | Operation::Restore { .. } | Operation::ReserveFragments { .. } @@ -1284,6 +1494,7 @@ impl<'a> TransactionRebase<'a> { | Operation::Overwrite { .. } | Operation::Delete { .. } | Operation::DataReplacement { .. } + | Operation::DataOverlay { .. } | Operation::Merge { .. } | Operation::Restore { .. } | Operation::Clone { .. } @@ -1374,6 +1585,7 @@ impl<'a> TransactionRebase<'a> { } Operation::CreateIndex { .. } => self.finish_create_index(dataset).await, Operation::Rewrite { .. } => self.finish_rewrite(dataset).await, + Operation::DataOverlay { .. } => self.finish_data_overlay(dataset).await, Operation::Append { .. } | Operation::Overwrite { .. } | Operation::DataReplacement { .. } @@ -1550,6 +1762,90 @@ impl<'a> TransactionRebase<'a> { } } + /// Verify no concurrent row-moving Update dropped the values of any cell + /// this overlay covers. `check_data_overlay_txn` flags (via the + /// `initial_fragments` needs-check bool) each overlaid fragment on which a + /// concurrent RewriteRows update relocated rows; here we read the deletion + /// vectors and conflict only when the moved rows intersect the overlay's + /// coverage. + /// + /// The moved rows are computed as the current deletion vector minus the + /// read-time one. In the rare case where both a concurrent Delete and a + /// concurrent Update touched the same flagged fragment, the Delete's rows are + /// also counted and may trigger an unnecessary retry — never data loss. Pure + /// concurrent deletes leave the fragment unflagged and are not examined here. + async fn finish_data_overlay(self, dataset: &Dataset) -> Result { + let fragments_to_check: HashSet = self + .initial_fragments + .iter() + .filter_map(|(id, (_, needs_check))| needs_check.then_some(*id)) + .collect(); + if fragments_to_check.is_empty() { + return Ok(Transaction { + read_version: dataset.manifest.version, + ..self.transaction + }); + } + + // Coverage (physical offsets, unioned across fields) per flagged fragment. + let Operation::DataOverlay { groups } = &self.transaction.operation else { + return Err(wrong_operation_err(&self.transaction.operation)); + }; + let mut coverage_by_fragment: HashMap = HashMap::new(); + for group in groups { + if !fragments_to_check.contains(&group.fragment_id) { + continue; + } + *coverage_by_fragment.entry(group.fragment_id).or_default() |= + overlay_group_coverage(group); + } + + for (fragment_id, coverage) in coverage_by_fragment { + let Some(current_fragment) = dataset + .fragments() + .as_slice() + .iter() + .find(|f| f.id == fragment_id) + else { + // The fragment is gone entirely; the overlay is orphaned. + return Err(crate::Error::retryable_commit_conflict_source( + dataset.manifest.version, + format!( + "This {} transaction was preempted: overlaid fragment {} was removed by a concurrent transaction. Please retry.", + self.transaction.uuid, fragment_id + ) + .into(), + )); + }; + let current_deletions = + read_fragment_deletion_bitmap(dataset, current_fragment).await?; + let initial_deletions = match self.initial_fragments.get(&fragment_id) { + Some((initial_fragment, _)) => { + read_fragment_deletion_bitmap(dataset, initial_fragment).await? + } + None => RoaringBitmap::new(), + }; + let moved_rows = ¤t_deletions - &initial_deletions; + let conflicting = &moved_rows & &coverage; + if !conflicting.is_empty() { + let sample: Vec = conflicting.iter().take(5).collect(); + return Err(crate::Error::retryable_commit_conflict_source( + dataset.manifest.version, + format!( + "This {} transaction was preempted by a concurrent update that moved overlaid rows on fragment {} (offsets {:?}). Please retry.", + self.transaction.uuid, fragment_id, sample.as_slice() + ) + .into(), + )); + } + } + + Ok(Transaction { + read_version: dataset.manifest.version, + ..self.transaction + }) + } + async fn finish_create_index(mut self, dataset: &Dataset) -> Result { if let Operation::CreateIndex { new_indices, @@ -1774,6 +2070,40 @@ async fn initial_fragments_for_rebase( .collect::>() } +/// Read a fragment's deletion vector as a bitmap of physical offsets, or an +/// empty bitmap when the fragment has no deletion file. +async fn read_fragment_deletion_bitmap( + dataset: &Dataset, + fragment: &Fragment, +) -> Result { + match &fragment.deletion_file { + Some(deletion_file) => { + let dv = read_dataset_deletion_file(dataset, fragment.id, deletion_file).await?; + Ok(RoaringBitmap::from(dv.as_ref())) + } + None => Ok(RoaringBitmap::new()), + } +} + +/// The physical offsets a group's overlays cover, unioned across every overlay +/// and every field. This is the set of cells whose values the overlay supplies, +/// used to test whether a concurrent row-moving Update actually invalidates the +/// overlay. +fn overlay_group_coverage(group: &DataOverlayGroup) -> RoaringBitmap { + let mut union = RoaringBitmap::new(); + for overlay in &group.overlays { + match &overlay.coverage { + OverlayCoverage::Shared(bitmap) => union |= bitmap.as_ref(), + OverlayCoverage::PerField(bitmaps) => { + for bitmap in bitmaps { + union |= bitmap.as_ref(); + } + } + } + } + union +} + fn wrong_operation_err(op: &Operation) -> Error { Error::internal(format!("function called against a wrong operation: {}", op)) } @@ -2781,6 +3111,442 @@ mod tests { } } + #[test] + fn test_data_overlay_conflicts() { + use crate::dataset::transaction::{DataOverlayGroup, UpdateMode}; + use ConflictResult::*; + use lance_table::format::overlay::{DataOverlayFile, OverlayCoverage}; + use roaring::RoaringBitmap; + + // Our transaction overlays fragment 1. + let overlay_op = |fragment_id: u64| Operation::DataOverlay { + groups: vec![DataOverlayGroup { + fragment_id, + overlays: vec![DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("overlay.lance", vec![0], None), + coverage: OverlayCoverage::dense(RoaringBitmap::from_iter([0u32])), + committed_version: 0, + }], + }], + }; + let update_removing = |removed_fragment_ids: Vec| Operation::Update { + removed_fragment_ids, + updated_fragments: vec![], + new_fragments: vec![], + fields_modified: vec![], + merged_generations: Vec::new(), + fields_for_preserving_frag_bitmap: vec![], + update_mode: None, + inserted_rows_filter: None, + updated_fragment_offsets: None, + }; + let delete = |updated: Vec, deleted: Vec| Operation::Delete { + updated_fragments: updated, + deleted_fragment_ids: deleted, + predicate: "x > 2".to_string(), + }; + // A row-moving update (RewriteRows) relocates the updated rows into + // new_fragments; an in-place column rewrite (RewriteColumns) leaves rows + // where they are. + let update_moving = |updated: Vec, new: Vec| Operation::Update { + removed_fragment_ids: vec![], + updated_fragments: updated, + new_fragments: new, + fields_modified: vec![], + merged_generations: Vec::new(), + fields_for_preserving_frag_bitmap: vec![], + update_mode: Some(UpdateMode::RewriteRows), + inserted_rows_filter: None, + updated_fragment_offsets: None, + }; + let update_rewrite_columns = |updated: Vec| Operation::Update { + removed_fragment_ids: vec![], + updated_fragments: updated, + new_fragments: vec![], + fields_modified: vec![0], + merged_generations: Vec::new(), + fields_for_preserving_frag_bitmap: vec![], + update_mode: Some(UpdateMode::RewriteColumns), + inserted_rows_filter: None, + updated_fragment_offsets: None, + }; + let rewrite_of = |old: &Fragment| Operation::Rewrite { + groups: vec![RewriteGroup { + old_fragments: vec![old.clone()], + new_fragments: vec![], + }], + rewritten_indices: vec![], + frag_reuse_index: None, + }; + + let fragment0 = Fragment::new(0); + let fragment1 = Fragment::new(1); + + // Each case is checked against our overlay on fragment 1. + let cases: Vec<(Operation, ConflictResult)> = vec![ + // Permissive: preserves physical offsets / leaves fragment 1 in place. + ( + Operation::Append { + fragments: vec![fragment0.clone()], + }, + Compatible, + ), + ( + Operation::CreateIndex { + new_indices: vec![], + removed_indices: vec![], + }, + Compatible, + ), + ( + Operation::DataReplacement { + replacements: vec![DataReplacementGroup( + 1, + DataFile::new_legacy_from_fields("r.lance", vec![0], None), + )], + }, + Compatible, + ), + // Another overlay on the same fragment stacks rather than conflicts. + (overlay_op(1), Compatible), + // A Delete only tombstones rows (deletion vector) on fragment 1, and + // an in-place column rewrite preserves offsets, so both are compatible. + (delete(vec![fragment1.clone()], vec![]), Compatible), + (update_rewrite_columns(vec![fragment1.clone()]), Compatible), + (update_removing(vec![2]), Compatible), + // ...but removing our overlaid fragment 1 orphans the overlay -> conflict. + (delete(vec![], vec![1]), Retryable), + (update_removing(vec![1]), Retryable), + // A row-moving update re-creates the rows it touches from the + // pre-overlay base. Whether that actually drops any overlaid cell is + // a per-row question answered in `finish_data_overlay` (see + // test_data_overlay_finish_conflicts_with_row_moving_update), so the + // check itself defers rather than conflicting; a moving update on any + // fragment is compatible at this stage. + ( + update_moving(vec![fragment1.clone()], vec![fragment0.clone()]), + Compatible, + ), + ( + update_moving(vec![fragment0.clone()], vec![fragment0.clone()]), + Compatible, + ), + // Rewriting fragment 1 invalidates its physical offsets -> conflict; + // a rewrite of a different fragment does not. + (rewrite_of(&fragment1), Retryable), + (rewrite_of(&fragment0), Compatible), + // Merge rewrites the whole fragment list; Restore replaces the dataset. + ( + Operation::Merge { + fragments: vec![fragment1.clone()], + schema: lance_core::datatypes::Schema::default(), + }, + Retryable, + ), + (Operation::Restore { version: 1 }, NotCompatible), + // Overwrite/Restore replace the dataset, and UpdateMemWalState does + // not rebase against data operations — all hard conflicts. + ( + Operation::Overwrite { + fragments: vec![fragment0.clone()], + schema: lance_core::datatypes::Schema::default(), + config_upsert_values: None, + initial_bases: None, + }, + NotCompatible, + ), + ( + Operation::UpdateMemWalState { + merged_generations: vec![], + }, + NotCompatible, + ), + ]; + + for (other, expected) in cases { + let mut rebase = TransactionRebase { + transaction: Transaction::new(0, overlay_op(1), None), + initial_fragments: HashMap::new(), + modified_fragment_ids: modified_fragment_ids(&overlay_op(1)) + .collect::>(), + affected_rows: None, + conflicting_frag_reuse_indices: Vec::new(), + conflicting_mem_wal_merged_gens: Vec::new(), + }; + let other_txn = Transaction::new(0, other.clone(), None); + let result = rebase.check_txn(&other_txn, 1); + match expected { + Compatible => assert!( + result.is_ok(), + "overlay should be compatible with {other:?}, got {result:?}" + ), + Retryable => assert!( + matches!(result, Err(Error::RetryableCommitConflict { .. })), + "overlay should retryably conflict with {other:?}, got {result:?}" + ), + NotCompatible => assert!( + matches!(result, Err(Error::IncompatibleTransaction { .. })), + "overlay should be incompatible with {other:?}, got {result:?}" + ), + } + } + } + + #[test] + fn test_rewrite_conflicts_with_data_overlay() { + // Reverse direction of test_data_overlay_conflicts: our transaction is a + // Rewrite and a concurrent DataOverlay has already committed. A rewrite + // changes the physical row addresses of the fragments it touches, so an + // overlay on one of those fragments is invalidated (retryable); an + // overlay on any other fragment is unaffected. + use crate::dataset::transaction::DataOverlayGroup; + use lance_table::format::overlay::{DataOverlayFile, OverlayCoverage}; + use roaring::RoaringBitmap; + + let overlay_on = |fragment_id: u64| Operation::DataOverlay { + groups: vec![DataOverlayGroup { + fragment_id, + overlays: vec![DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("overlay.lance", vec![0], None), + coverage: OverlayCoverage::dense(RoaringBitmap::from_iter([0u32])), + committed_version: 0, + }], + }], + }; + // Our transaction rewrites fragment 1. + let rewrite_op = Operation::Rewrite { + groups: vec![RewriteGroup { + old_fragments: vec![Fragment::new(1)], + new_fragments: vec![], + }], + rewritten_indices: vec![], + frag_reuse_index: None, + }; + + for (other, expect_conflict) in [(overlay_on(1), true), (overlay_on(0), false)] { + let mut rebase = TransactionRebase { + transaction: Transaction::new(0, rewrite_op.clone(), None), + initial_fragments: HashMap::new(), + modified_fragment_ids: modified_fragment_ids(&rewrite_op).collect::>(), + affected_rows: None, + conflicting_frag_reuse_indices: Vec::new(), + conflicting_mem_wal_merged_gens: Vec::new(), + }; + let other_txn = Transaction::new(0, other.clone(), None); + let result = rebase.check_txn(&other_txn, 1); + if expect_conflict { + assert!( + matches!(result, Err(Error::RetryableCommitConflict { .. })), + "rewrite of fragment 1 should retryably conflict with {other:?}, got {result:?}" + ); + } else { + assert!( + result.is_ok(), + "rewrite of fragment 1 should not conflict with {other:?}, got {result:?}" + ); + } + } + } + + #[test] + fn test_update_conflicts_with_data_overlay() { + // Reverse direction of test_data_overlay_conflicts: our transaction is an + // Update and a concurrent DataOverlay has already committed. A row-moving + // update relocates the rows it touches, so an overlay on one of those + // fragments can no longer be applied (retryable); an overlay on any other + // fragment, or an in-place column rewrite, is compatible. + use crate::dataset::transaction::{DataOverlayGroup, UpdateMode}; + use lance_table::format::overlay::{DataOverlayFile, OverlayCoverage}; + use roaring::RoaringBitmap; + + let overlay_on = |fragment_id: u64| Operation::DataOverlay { + groups: vec![DataOverlayGroup { + fragment_id, + overlays: vec![DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("overlay.lance", vec![0], None), + coverage: OverlayCoverage::dense(RoaringBitmap::from_iter([0u32])), + committed_version: 0, + }], + }], + }; + // Our update always touches fragment 1. + let update = + |update_mode: Option, new_fragments: Vec| Operation::Update { + removed_fragment_ids: vec![], + updated_fragments: vec![Fragment::new(1)], + new_fragments, + fields_modified: vec![0], + merged_generations: Vec::new(), + fields_for_preserving_frag_bitmap: vec![], + update_mode, + inserted_rows_filter: None, + updated_fragment_offsets: None, + }; + + // The overlay covers physical offset 0 of its fragment. Row addresses + // pack the fragment id in the high 32 bits and the offset in the low 32. + let rows_on = |fragment_id: u64, offsets: &[u32]| { + let mut map = RowAddrTreeMap::new(); + map.insert_bitmap( + fragment_id as u32, + RoaringBitmap::from_iter(offsets.iter().copied()), + ); + map + }; + + // (update, committed overlay, moved rows the update carries, expect conflict) + let cases = [ + // Row-moving update whose moved rows include the overlaid cell -> the + // update would undo the overlay, so conflict. + ( + update(Some(UpdateMode::RewriteRows), vec![Fragment::new(2)]), + overlay_on(1), + Some(rows_on(1, &[0])), + true, + ), + // ...but if the moved rows miss the overlaid cell, the overlay survives. + ( + update(Some(UpdateMode::RewriteRows), vec![Fragment::new(2)]), + overlay_on(1), + Some(rows_on(1, &[5])), + false, + ), + // An overlay on a fragment the update did not touch is fine. + ( + update(Some(UpdateMode::RewriteRows), vec![Fragment::new(2)]), + overlay_on(0), + Some(rows_on(1, &[0])), + false, + ), + // An in-place column rewrite preserves rows -> compatible. + ( + update(Some(UpdateMode::RewriteColumns), vec![]), + overlay_on(1), + Some(rows_on(1, &[0])), + false, + ), + // Without affected rows we cannot be precise, so a row-moving update + // on the overlaid fragment falls back to a conservative conflict. + ( + update(Some(UpdateMode::RewriteRows), vec![Fragment::new(2)]), + overlay_on(1), + None, + true, + ), + ]; + + for (update_op, other, affected_rows, expect_conflict) in cases { + let mut rebase = TransactionRebase { + transaction: Transaction::new(0, update_op.clone(), None), + initial_fragments: HashMap::new(), + modified_fragment_ids: modified_fragment_ids(&update_op).collect::>(), + affected_rows: affected_rows.as_ref(), + conflicting_frag_reuse_indices: Vec::new(), + conflicting_mem_wal_merged_gens: Vec::new(), + }; + let other_txn = Transaction::new(0, other.clone(), None); + let result = rebase.check_txn(&other_txn, 1); + if expect_conflict { + assert!( + matches!(result, Err(Error::RetryableCommitConflict { .. })), + "update should retryably conflict with {other:?}, got {result:?}" + ); + } else { + assert!( + result.is_ok(), + "update should be compatible with {other:?}, got {result:?}" + ); + } + } + } + + #[tokio::test] + #[rstest::rstest] + #[case::coverage_overlaps_moved_row(vec![0u32], true)] + #[case::coverage_disjoint_from_moved_row(vec![3u32], false)] + async fn test_data_overlay_finish_conflicts_with_row_moving_update( + #[case] coverage_offsets: Vec, + #[case] expect_conflict: bool, + ) { + // 5 rows in one fragment. A concurrent RewriteRows update moves row 0 out + // to a new fragment (deleting it from fragment 0). Our overlay on fragment + // 0 conflicts only when its coverage includes the moved row; the decision + // is made in finish, which reads the deletion vectors. + use crate::dataset::transaction::{DataOverlayGroup, UpdateMode}; + use lance_table::format::overlay::{DataOverlayFile, OverlayCoverage}; + use roaring::RoaringBitmap; + + let dataset = test_dataset(5, 1).await; + let mut fragment = dataset.fragments().as_slice()[0].clone(); + + let moved_fragment = Fragment::new(0) + .with_file( + "moved.lance", + vec![0], + vec![0], + &LanceFileVersion::Stable, + NonZero::new(10), + ) + .with_physical_rows(1); + let update_op = Operation::Update { + updated_fragments: vec![apply_deletion(&[0], &mut fragment, &dataset).await], + removed_fragment_ids: vec![], + new_fragments: vec![moved_fragment], + fields_modified: vec![], + merged_generations: Vec::new(), + fields_for_preserving_frag_bitmap: vec![], + update_mode: Some(UpdateMode::RewriteRows), + inserted_rows_filter: None, + updated_fragment_offsets: None, + }; + let update_txn = Transaction::new_from_version(dataset.manifest.version, update_op); + + let overlay_op = Operation::DataOverlay { + groups: vec![DataOverlayGroup { + fragment_id: 0, + overlays: vec![DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("overlay.lance", vec![0], None), + coverage: OverlayCoverage::dense(RoaringBitmap::from_iter(coverage_offsets)), + committed_version: 0, + }], + }], + }; + let overlay_txn = Transaction::new_from_version(dataset.manifest.version, overlay_op); + + // Commit the update so the latest dataset reflects the moved (deleted) row. + let latest_dataset = CommitBuilder::new(Arc::new(dataset.clone())) + .execute(update_txn.clone()) + .await + .unwrap(); + + let mut rebase = TransactionRebase::try_new(&dataset, overlay_txn.clone(), None) + .await + .unwrap(); + // The check defers the row-level decision to finish, flagging fragment 0. + rebase.check_txn(&update_txn, 1).unwrap(); + assert_eq!( + rebase + .initial_fragments + .iter() + .map(|(id, (_, needs_check))| (*id, *needs_check)) + .collect::>(), + vec![(0, true)], + ); + + let res = rebase.finish(&latest_dataset).await; + if expect_conflict { + assert!( + matches!(res, Err(crate::Error::RetryableCommitConflict { .. })), + "overlay covering the moved row should conflict, got {res:?}" + ); + } else { + assert!( + res.is_ok(), + "overlay disjoint from the moved row should succeed, got {res:?}" + ); + } + } + #[test] fn test_create_index_conflicts_only_on_same_name() { let index0 = IndexMetadata { @@ -3255,6 +4021,7 @@ mod tests { Operation::DataReplacement { replacements } => { Box::new(replacements.iter().map(|r| r.0)) } + Operation::DataOverlay { groups } => Box::new(groups.iter().map(|g| g.fragment_id)), } } diff --git a/rust/lance/src/utils/test.rs b/rust/lance/src/utils/test.rs index 3338eee07a8..f804a7cc38a 100644 --- a/rust/lance/src/utils/test.rs +++ b/rust/lance/src/utils/test.rs @@ -243,6 +243,7 @@ impl TestDatasetGenerator { Fragment { id: 0, files, + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: Some(batch.num_rows()),