diff --git a/rust/lance/src/dataset.rs b/rust/lance/src/dataset.rs index 72b13e903b5..cf2695fc668 100644 --- a/rust/lance/src/dataset.rs +++ b/rust/lance/src/dataset.rs @@ -80,6 +80,7 @@ pub mod index; pub mod mem_wal; mod metadata; pub mod optimize; +pub(crate) mod overlay; pub mod progress; pub mod refs; pub(crate) mod rowids; diff --git a/rust/lance/src/dataset/fragment.rs b/rust/lance/src/dataset/fragment.rs index 4df89ab71b8..c19e9cd6ae3 100644 --- a/rust/lance/src/dataset/fragment.rs +++ b/rust/lance/src/dataset/fragment.rs @@ -65,6 +65,9 @@ use super::updater::Updater; use super::{NewColumnTransform, WriteParams, schema_evolution}; use crate::dataset::Dataset; use crate::dataset::fragment::session::FragmentSession; +use crate::dataset::overlay::{ + OverlayReadPlanner, merge_overlay_batch, plan_overlays, resolve_overlays, +}; use crate::io::deletion::read_dataset_deletion_file; /// Result of [`FileFragment::update_columns_with_offsets`]: updated fragment metadata, modified field ids, @@ -649,7 +652,7 @@ impl GenericFileReader for NullReader { } } -#[derive(Debug, Default)] +#[derive(Debug, Default, Clone)] pub struct FragReadConfig { // Add the row id column pub with_row_id: bool, @@ -911,16 +914,6 @@ 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(); @@ -959,6 +952,19 @@ impl FileFragment { Arc::new(self.metadata.clone()), )?; + // Plan overlay resolution from coverage metadata (no files opened here); the + // readers are opened lazily on read, pruned to the rows each read touches. + if !self.metadata.overlays.is_empty() { + let planner = plan_overlays(self, projection)?; + if !planner.is_empty() { + reader.overlay = Some(OverlayReadState { + planner: Arc::new(planner), + fragment: Arc::new(self.clone()), + read_config: Arc::new(read_config.clone()), + }); + } + } + if read_config.with_row_id { reader.with_row_id(); } @@ -996,7 +1002,7 @@ impl FileFragment { selected_columns.saturating_mul(4) < total_columns } - async fn open_reader( + pub(super) async fn open_reader( &self, data_file: &DataFile, projection: Option<&Schema>, @@ -2264,6 +2270,23 @@ pub struct FragmentReader { // total number of physical rows in the fragment (all rows, ignoring deletions) num_physical_rows: usize, + + /// Read-time state for resolving data overlay files: the coverage plan plus + /// what is needed to open overlay readers. `None` when the fragment has no + /// overlays. Overlays are merged into base batches (by `offset_in_frag`) before + /// deletion filtering, opening only the files each read's rows touch. + overlay: Option, +} + +/// What [`FragmentReader`] needs to resolve overlays at read time: the coverage +/// plan (from metadata, cheap to build), and the fragment + config needed to open +/// overlay readers once the read's rows — and therefore which files it touches — +/// are known. All `Arc` so cloning a reader stays cheap. +#[derive(Clone, Debug)] +struct OverlayReadState { + planner: Arc, + fragment: Arc, + read_config: Arc, } // Custom clone impl needed because it is not easy to clone Box @@ -2292,6 +2315,7 @@ impl Clone for FragmentReader { created_at_sequence: self.created_at_sequence.clone(), num_rows: self.num_rows, num_physical_rows: self.num_physical_rows, + overlay: self.overlay.clone(), } } } @@ -2359,6 +2383,7 @@ impl FragmentReader { created_at_sequence: None, num_rows, num_physical_rows, + overlay: None, }) } @@ -2616,6 +2641,77 @@ impl FragmentReader { Ok(result.project_by_schema(&output_schema)?) } + /// Merge data overlay values onto a stream of base batches. + /// + /// Runs on physical rows in read order, *before* deletion filtering, so each + /// row can be addressed by its position in the fragment (its `offset_in_frag`, + /// derived from `params`) and deletions take precedence naturally: an overlay + /// value for a deleted row is dropped along with the row downstream. A no-op + /// when the fragment has no overlays. + /// + /// The read's `offset_in_frag` values are known from `params` up front, so + /// overlays are resolved here to just the files this read's rows touch — an + /// overlay whose cells fall outside the read is not opened at all. Within each + /// batch, the overlay reads (only the values that batch needs) are then issued + /// concurrently with the base read rather than after it. + async fn merge_overlays( + &self, + merged: ReadBatchTaskStream, + params: &ReadBatchParams, + total_num_rows: u32, + ) -> Result { + let Some(overlay) = &self.overlay else { + return Ok(merged); + }; + // The offset_in_frag of every row this read will return, materialized once. + // Cost is one u32 per output row (a whole-fragment scan is 4 bytes/row), and + // it lets us both prune overlays to the read and slice each batch's offsets + // below without reading any data. Only paid when the fragment has overlays. + // + // TODO(overlay perf): this could be avoided by teaching `ReadBatchParams` to + // yield a coverage bitmap directly (for pruning) and to slice per batch (for + // the routing below), or by moving `ReadBatchParams` to a roaring bitmap + // wholesale — a larger refactor tracked separately. + let offsets_in_frag: Arc> = + Arc::new(params.to_offsets_total(total_num_rows).values().to_vec()); + + // Open only the overlay readers this read touches (pruned by row selection). + let plans = resolve_overlays( + &overlay.planner, + &offsets_in_frag, + &overlay.fragment, + &overlay.read_config, + ) + .await?; + if plans.is_empty() { + return Ok(merged); + } + let plans = Arc::new(plans); + + // Batches arrive in physical read order, so a running total of the rows seen + // so far gives each batch its starting offset_in_batch into `offsets_in_frag`. + let mut rows_seen = 0usize; + let stream = merged + .map(move |task| { + let num_rows = task.num_rows; + let start = rows_seen; + rows_seen += num_rows as usize; + let offsets_in_frag = offsets_in_frag.clone(); + let plans = plans.clone(); + let inner = task.task; + ReadBatchTask { + num_rows, + task: async move { + let batch_offsets = &offsets_in_frag[start..start + num_rows as usize]; + merge_overlay_batch(inner, batch_offsets, &plans).await + } + .boxed(), + } + }) + .boxed(); + Ok(stream) + } + async fn new_read_impl<'a, F>( &'a self, params: ReadBatchParams, @@ -2683,6 +2779,8 @@ impl FragmentReader { lance_table::utils::stream::merge_streams(read_streams) }; + let merged = self.merge_overlays(merged, ¶ms, total_num_rows).await?; + // Add the row id column (if needed) and delete rows (if a deletion // vector is present). let config = RowIdAndDeletesConfig { @@ -2853,6 +2951,11 @@ impl FragmentReader { lance_table::utils::stream::merge_streams(read_streams) }; + let params = ReadBatchParams::Ranges(ranges); + let merged_stream = self + .merge_overlays(merged_stream, ¶ms, total_num_rows) + .await?; + // Add the row id column (if needed) and delete rows (if a deletion // vector is present). let config = RowIdAndDeletesConfig { @@ -2865,7 +2968,7 @@ impl FragmentReader { with_row_created_at_version: self.with_row_created_at_version, last_updated_at_sequence: self.last_updated_at_sequence.clone(), created_at_sequence: self.created_at_sequence.clone(), - params: ReadBatchParams::Ranges(ranges), + params, total_num_rows, }; let output_schema = Arc::new(self.output_schema.clone()); @@ -3114,6 +3217,1519 @@ mod tests { Dataset::open(test_uri).await.unwrap() } + /// End-to-end tests for reading data overlay files (OSS-1324): overlays are + /// written, committed via the `DataOverlay` transaction, and then resolved on + /// the `take` and scan read paths. + mod overlay_read { + use std::sync::Arc; + + use arrow_array::{ + Array, ArrayRef, Int32Array, RecordBatch, RecordBatchIterator, StructArray, UInt64Array, + }; + use arrow_schema::{DataType, Field as ArrowField, Fields, Schema as ArrowSchema}; + use lance_core::datatypes::Schema; + use lance_file::version::LanceFileVersion; + use lance_file::writer::{FileWriter, FileWriterOptions}; + use lance_io::utils::CachedFileSize; + use lance_table::format::DataFile; + use lance_table::format::overlay::{DataOverlayFile, OverlayCoverage}; + use object_store::path::Path; + use roaring::RoaringBitmap; + use rstest::rstest; + + use crate::dataset::transaction::{DataOverlayGroup, Operation}; + use crate::dataset::{Dataset, WriteDestination, WriteParams}; + + fn bitmap(offsets: impl IntoIterator) -> RoaringBitmap { + RoaringBitmap::from_iter(offsets) + } + + fn i32_array(values: impl IntoIterator>) -> ArrayRef { + Arc::new(Int32Array::from_iter(values)) + } + + /// Two-fragment Int32 dataset: `id` (field 0) = 0..12 and `val` (field 1) + /// = id * 10, written 6 rows per file (fragments 0 and 1). + /// + /// Uses an in-memory store so the test can write overlay files with a + /// store-relative `data/.lance` path and commit against the returned + /// dataset directly. + async fn create_base_dataset(version: LanceFileVersion) -> Dataset { + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", DataType::Int32, true), + ArrowField::new("val", DataType::Int32, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from_iter_values(0..12)), + Arc::new(Int32Array::from_iter_values((0..12).map(|v| v * 10))), + ], + ) + .unwrap(); + let write_params = WriteParams { + max_rows_per_file: 6, + max_rows_per_group: 6, + data_storage_version: Some(version), + ..Default::default() + }; + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + Dataset::write(reader, "memory://", Some(write_params)) + .await + .unwrap() + } + + /// Write an overlay file covering `fields` (dataset field ids) of + /// `fragment_id` with the given coverage and per-field value columns, then + /// commit it as a `DataOverlay` transaction. `name` makes the file unique. + #[allow(clippy::too_many_arguments)] + async fn commit_overlay( + dataset: Dataset, + name: &str, + fragment_id: u64, + fields: &[i32], + coverage: OverlayCoverage, + columns: Vec, + version: LanceFileVersion, + ) -> Dataset { + let read_version = dataset.version().version; + let overlay_schema = dataset.schema().project_by_ids(fields, true); + + let filename = format!("{name}.lance"); + let path = Path::from(format!("data/{filename}")); + let obj_writer = dataset.object_store.create(&path).await.unwrap(); + let mut writer = FileWriter::try_new( + obj_writer, + overlay_schema, + FileWriterOptions { + format_version: Some(version), + ..Default::default() + }, + ) + .unwrap(); + let (major, minor) = writer.version().to_numbers(); + for (column_index, array) in columns.into_iter().enumerate() { + writer.write_column(column_index, array).await.unwrap(); + } + let summary = writer.finish().await.unwrap(); + + let mut data_file = DataFile::new_unstarted(filename, major, minor); + data_file.fields = writer + .field_id_to_column_indices() + .iter() + .map(|(field_id, _)| *field_id as i32) + .collect::>() + .into(); + data_file.column_indices = writer + .field_id_to_column_indices() + .iter() + .map(|(_, column_index)| *column_index as i32) + .collect::>() + .into(); + data_file.file_size_bytes = CachedFileSize::new(summary.size_bytes); + + let overlay = DataOverlayFile { + data_file, + coverage, + committed_version: 0, + }; + Dataset::commit( + WriteDestination::Dataset(Arc::new(dataset)), + Operation::DataOverlay { + groups: vec![DataOverlayGroup { + fragment_id, + overlays: vec![overlay], + }], + }, + Some(read_version), + None, + None, + Arc::new(Default::default()), + false, + ) + .await + .unwrap() + } + + fn full_schema(dataset: &Dataset) -> Schema { + dataset.schema().clone() + } + + fn col(batch: &RecordBatch, name: &str) -> Int32Array { + let idx = batch.schema().index_of(name).unwrap(); + batch + .column(idx) + .as_any() + .downcast_ref::() + .unwrap() + .clone() + } + + #[rstest] + #[tokio::test] + async fn test_take_covered_and_uncovered( + #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_1)] version: LanceFileVersion, + ) { + let dataset = create_base_dataset(version).await; + // Overlay fragment 0's `val` at physical offsets {1, 4}. + let dataset = commit_overlay( + dataset, + "ov", + 0, + &[1], + OverlayCoverage::dense(bitmap([1, 4])), + vec![i32_array([Some(111), Some(444)])], + version, + ) + .await; + + let frag = dataset.get_fragment(0).unwrap(); + let batch = frag + .take(&[0, 1, 2, 4], &full_schema(&dataset)) + .await + .unwrap(); + // Offsets 1 and 4 take overlay values; 0 and 2 fall through to base. + assert_eq!(col(&batch, "val").values(), &[0, 111, 20, 444]); + // The unrelated `id` column is untouched. + assert_eq!(col(&batch, "id").values(), &[0, 1, 2, 4]); + } + + #[rstest] + #[tokio::test] + async fn test_take_newest_overlay_wins( + #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_1)] version: LanceFileVersion, + ) { + let dataset = create_base_dataset(version).await; + let dataset = commit_overlay( + dataset, + "older", + 0, + &[1], + OverlayCoverage::dense(bitmap([1, 4])), + vec![i32_array([Some(111), Some(444)])], + version, + ) + .await; + // A newer overlay (later commit -> higher committed_version) re-covers + // offset 1. + let dataset = commit_overlay( + dataset, + "newer", + 0, + &[1], + OverlayCoverage::dense(bitmap([1])), + vec![i32_array([Some(999)])], + version, + ) + .await; + + let frag = dataset.get_fragment(0).unwrap(); + let batch = frag.take(&[1, 4], &full_schema(&dataset)).await.unwrap(); + // Offset 1 -> newest overlay (999); offset 4 -> only older covers it. + assert_eq!(col(&batch, "val").values(), &[999, 444]); + } + + #[rstest] + #[tokio::test] + async fn test_take_per_field_coverage( + #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_1)] version: LanceFileVersion, + ) { + let dataset = create_base_dataset(version).await; + // Sparse overlay: `id` covers {2}, `val` covers {2, 3} — different + // offset sets and therefore unequal-length value columns. + let dataset = commit_overlay( + dataset, + "sparse", + 0, + &[0, 1], + OverlayCoverage::sparse(vec![bitmap([2]), bitmap([2, 3])]), + vec![i32_array([Some(777)]), i32_array([Some(220), Some(330)])], + version, + ) + .await; + + let frag = dataset.get_fragment(0).unwrap(); + let batch = frag.take(&[2, 3], &full_schema(&dataset)).await.unwrap(); + // id: offset 2 covered (777), offset 3 falls through (3). + assert_eq!(col(&batch, "id").values(), &[777, 3]); + // val: both offsets covered (220, 330). + assert_eq!(col(&batch, "val").values(), &[220, 330]); + } + + #[rstest] + #[tokio::test] + async fn test_take_null_override( + #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_1)] version: LanceFileVersion, + ) { + let dataset = create_base_dataset(version).await; + let dataset = commit_overlay( + dataset, + "nullov", + 0, + &[1], + OverlayCoverage::dense(bitmap([0])), + vec![i32_array([None])], + version, + ) + .await; + + let frag = dataset.get_fragment(0).unwrap(); + let batch = frag.take(&[0, 1], &full_schema(&dataset)).await.unwrap(); + let val = col(&batch, "val"); + // Offset 0 is covered with a NULL value -> resolves to NULL; offset 1 + // falls through to the base value. + assert!(val.is_null(0)); + assert_eq!(val.value(1), 10); + } + + /// Overlays interact correctly with NULL *base* cells (distinct from a NULL + /// overlay value): a covered row whose base value is NULL is overridden to the + /// overlay's non-null value, while an uncovered NULL base cell falls through + /// and stays NULL. + #[rstest] + #[tokio::test] + async fn test_take_null_base_cell( + #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_1)] version: LanceFileVersion, + ) { + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", DataType::Int32, true), + ArrowField::new("val", DataType::Int32, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from_iter_values(0..6)), + // `val` is NULL at offsets 1 and 3. + Arc::new(Int32Array::from_iter([ + Some(0), + None, + Some(20), + None, + Some(40), + Some(50), + ])), + ], + ) + .unwrap(); + let write_params = WriteParams { + max_rows_per_file: 6, + max_rows_per_group: 6, + data_storage_version: Some(version), + ..Default::default() + }; + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + let dataset = Dataset::write(reader, "memory://", Some(write_params)) + .await + .unwrap(); + + // Cover offset 1 (NULL base) and offset 4 (non-null base); leave offset + // 3's NULL base uncovered. + let dataset = commit_overlay( + dataset, + "nullbase", + 0, + &[1], + OverlayCoverage::dense(bitmap([1, 4])), + vec![i32_array([Some(111), Some(444)])], + version, + ) + .await; + + let frag = dataset.get_fragment(0).unwrap(); + let batch = frag.take(&[1, 3, 4], &full_schema(&dataset)).await.unwrap(); + let val = col(&batch, "val"); + // Offset 1: NULL base overridden to 111. Offset 3: uncovered NULL base + // stays NULL. Offset 4: non-null base overridden to 444. + assert_eq!(val.value(0), 111); + assert!(val.is_null(1)); + assert_eq!(val.value(2), 444); + } + + #[rstest] + #[tokio::test] + async fn test_overlay_on_deleted_row_is_inert( + #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_1)] version: LanceFileVersion, + ) { + let mut dataset = create_base_dataset(version).await; + // Delete global row 1 (fragment 0, physical offset 1). + dataset.delete("id = 1").await.unwrap(); + // Overlay covers the deleted offset 1 and the live offset 4. + let dataset = commit_overlay( + dataset, + "delov", + 0, + &[1], + OverlayCoverage::dense(bitmap([1, 4])), + vec![i32_array([Some(111), Some(444)])], + version, + ) + .await; + + // Scan fragment 0: row 1 is gone, and offset 4's overlay value survives + // even though the deletion shifts logical positions — coverage is keyed + // by physical offset. + let frag = dataset.get_fragment(0).unwrap(); + let mut scanner = frag.scan(); + let batch = scanner + .project(&["id", "val"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(col(&batch, "id").values(), &[0, 2, 3, 4, 5]); + assert_eq!(col(&batch, "val").values(), &[0, 20, 30, 444, 50]); + } + + #[rstest] + #[tokio::test] + async fn test_scan_multi_fragment_overlays( + #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_1)] version: LanceFileVersion, + ) { + let dataset = create_base_dataset(version).await; + // Overlay fragment 0 at offset 0 and fragment 1 at offset 0 (global + // row 6). Each fragment's coverage is independent. + let dataset = commit_overlay( + dataset, + "frag0", + 0, + &[1], + OverlayCoverage::dense(bitmap([0])), + vec![i32_array([Some(1000)])], + version, + ) + .await; + let dataset = commit_overlay( + dataset, + "frag1", + 1, + &[1], + OverlayCoverage::dense(bitmap([0])), + vec![i32_array([Some(6000)])], + version, + ) + .await; + + let batch = dataset + .scan() + .project(&["id", "val"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(batch.num_rows(), 12); + let expected: Vec = (0..12) + .map(|i| match i { + 0 => 1000, + 6 => 6000, + other => other * 10, + }) + .collect(); + assert_eq!(col(&batch, "val").values(), &expected); + } + + /// A `take` of a few rows must read only the overlay values those rows + /// touch — not the whole column. Uses v2.1 (which slices pages on read) and + /// an incompressible, all-covering overlay, so reading the full column would + /// be far more bytes than reading a couple of values. This is the regression + /// guard for the lazy, value-pushdown overlay read. + #[tokio::test] + async fn test_take_reads_only_needed_overlay_values() { + let version = LanceFileVersion::V2_1; + const N: usize = 100_000; + + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", DataType::Int32, true), + ArrowField::new("val", DataType::Int32, true), + ])); + let base = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from_iter_values(0..N as i32)), + Arc::new(Int32Array::from_iter_values((0..N as i32).map(|v| v * 10))), + ], + ) + .unwrap(); + let write_params = WriteParams { + max_rows_per_file: N, + max_rows_per_group: N, + data_storage_version: Some(version), + ..Default::default() + }; + let reader = RecordBatchIterator::new(vec![Ok(base)], schema.clone()); + let dataset = Dataset::write(reader, "memory://", Some(write_params)) + .await + .unwrap(); + + // Overlay `val` over ALL N offsets with incompressible values, so the + // value column is ~N*4 bytes on disk. + let values: Vec = (0..N as u64) + .map(|i| { + let mut x = i; + x ^= x >> 33; + x = x.wrapping_mul(0xff51_afd7_ed55_8ccd); + x ^= x >> 33; + x as i32 + }) + .collect(); + let dataset = commit_overlay( + dataset, + "big", + 0, + &[1], + OverlayCoverage::dense(bitmap(0..N as u32)), + vec![Arc::new(Int32Array::from(values.clone())) as ArrayRef], + version, + ) + .await; + + let frag = dataset.get_fragment(0).unwrap(); + let val_only = dataset.schema().project_by_ids(&[1], true); + + // Measure only the reads that resolve the take. + dataset.object_store.io_stats_incremental(); + let batch = frag.take(&[0, 1], &val_only).await.unwrap(); + let io = dataset.object_store.io_stats_incremental(); + + // The overlay's `val` column alone is N*4 bytes; resolving two adjacent + // offsets must read only a small fraction of it. + let full_column_bytes = (N * std::mem::size_of::()) as u64; + assert!( + io.read_bytes > 0 && io.read_bytes < full_column_bytes / 4, + "take read {} bytes; expected far less than the {}-byte overlay \ + column (a take must not read the whole value column)", + io.read_bytes, + full_column_bytes, + ); + + // ...and it still resolves correctly. + let val = col(&batch, "val"); + assert_eq!(val.value(0), values[0]); + assert_eq!(val.value(1), values[1]); + } + + /// Row-selection pruning: an overlay whose coverage is disjoint from the + /// requested rows must not be opened at all. Proven by deleting the overlay's + /// data file — a `take` that misses its coverage still succeeds (the file is + /// never touched), while a `take` that hits it then fails because the file is + /// genuinely needed. + #[rstest] + #[tokio::test] + async fn test_take_prunes_overlays_outside_row_selection( + #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_1)] version: LanceFileVersion, + ) { + let dataset = create_base_dataset(version).await; + // Overlay on fragment 0 (offsets 0..6) covering only offset_in_frag 5. + let dataset = commit_overlay( + dataset, + "miss", + 0, + &[1], + OverlayCoverage::dense(bitmap([5])), + vec![i32_array([Some(5000)])], + version, + ) + .await; + + // Delete the overlay's data file: opening it now fails. + dataset + .object_store + .delete(&Path::from("data/miss.lance")) + .await + .unwrap(); + + let frag = dataset.get_fragment(0).unwrap(); + let val_only = dataset.schema().project_by_ids(&[1], true); + + // A take that misses the overlay's coverage must not open it, so it + // succeeds and returns base values (val = offset * 10). + let batch = frag.take(&[0, 1], &val_only).await.unwrap(); + assert_eq!(col(&batch, "val").values(), &[0, 10]); + + // A take that hits the coverage does need the file, so it now fails with + // a not-found error naming the missing overlay file. + let err = frag.take(&[5], &val_only).await.unwrap_err(); + let err = format!("{err:?}"); + assert!( + err.contains("miss.lance") && err.to_lowercase().contains("not found"), + "take hitting the overlay's coverage should fail with a not-found error \ + for its missing file, got: {err}", + ); + } + + /// The overlay merge runs before `wrap_with_row_id_and_delete`, so the + /// `_rowid` system column must coexist with overlay-resolved data columns: + /// the row ids are unaffected by the merge and the overlay value still wins. + #[rstest] + #[tokio::test] + async fn test_scan_with_row_id_alongside_overlay( + #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_1)] version: LanceFileVersion, + ) { + let dataset = create_base_dataset(version).await; + let dataset = commit_overlay( + dataset, + "rowidov", + 0, + &[1], + OverlayCoverage::dense(bitmap([0])), + vec![i32_array([Some(1000)])], + version, + ) + .await; + + let frag = dataset.get_fragment(0).unwrap(); + let batch = frag + .scan() + .with_row_id() + .project(&["id", "val"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + // Overlay value resolves... + assert_eq!(col(&batch, "val").values()[0], 1000); + assert_eq!(&col(&batch, "val").values()[1..], &[10, 20, 30, 40, 50]); + // ...and the row ids for fragment 0 are the untouched physical offsets. + let row_ids = batch + .column(batch.schema().index_of("_rowid").unwrap()) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(row_ids.values(), &[0, 1, 2, 3, 4, 5]); + } + + /// When the newest overlay covers every requested offset, an older overlay + /// in the same plan needs zero values and its value column must not be read + /// (the empty-input branch of `fetch_overlay_values`). The result still + /// resolves to the newest overlay. + #[rstest] + #[tokio::test] + async fn test_take_older_overlay_contributes_no_values( + #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_1)] version: LanceFileVersion, + ) { + let dataset = create_base_dataset(version).await; + // Older covers {1, 4}; newer re-covers {1}. A take of only offset 1 + // routes entirely to the newer overlay, leaving the older one with no + // values to fetch even though it is part of the field's plan. + let dataset = commit_overlay( + dataset, + "older", + 0, + &[1], + OverlayCoverage::dense(bitmap([1, 4])), + vec![i32_array([Some(111), Some(444)])], + version, + ) + .await; + let dataset = commit_overlay( + dataset, + "newer", + 0, + &[1], + OverlayCoverage::dense(bitmap([1])), + vec![i32_array([Some(999)])], + version, + ) + .await; + + let frag = dataset.get_fragment(0).unwrap(); + let batch = frag.take(&[1], &full_schema(&dataset)).await.unwrap(); + assert_eq!(col(&batch, "val").values(), &[999]); + } + + /// A newest overlay whose value is NULL must shadow an older overlay's + /// non-null value at the same offset — the merge resolves to NULL, it does + /// not fall back to the older overlay. + #[rstest] + #[tokio::test] + async fn test_take_newest_null_shadows_older( + #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_1)] version: LanceFileVersion, + ) { + let dataset = create_base_dataset(version).await; + let dataset = commit_overlay( + dataset, + "older", + 0, + &[1], + OverlayCoverage::dense(bitmap([1])), + vec![i32_array([Some(111)])], + version, + ) + .await; + let dataset = commit_overlay( + dataset, + "newer_null", + 0, + &[1], + OverlayCoverage::dense(bitmap([1])), + vec![i32_array([None])], + version, + ) + .await; + + let frag = dataset.get_fragment(0).unwrap(); + let batch = frag.take(&[1], &full_schema(&dataset)).await.unwrap(); + let val = col(&batch, "val"); + assert!(val.is_null(0), "newest NULL must win over older 111"); + } + + /// Newest-wins is resolved independently per field across multiple sparse + /// overlays: for the same offset, `id` can resolve to one overlay while + /// `val` resolves to the other, depending on which overlay newly covers + /// that field at that offset. + #[rstest] + #[tokio::test] + async fn test_take_multi_sparse_per_field_newest_wins( + #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_1)] version: LanceFileVersion, + ) { + let dataset = create_base_dataset(version).await; + // Older: id covers {3}, val covers {2}. + let dataset = commit_overlay( + dataset, + "older", + 0, + &[0, 1], + OverlayCoverage::sparse(vec![bitmap([3]), bitmap([2])]), + vec![i32_array([Some(7773)]), i32_array([Some(2772)])], + version, + ) + .await; + // Newer: id covers {2}, val covers {3} — the mirror image. + let dataset = commit_overlay( + dataset, + "newer", + 0, + &[0, 1], + OverlayCoverage::sparse(vec![bitmap([2]), bitmap([3])]), + vec![i32_array([Some(9992)]), i32_array([Some(9993)])], + version, + ) + .await; + + let frag = dataset.get_fragment(0).unwrap(); + let batch = frag.take(&[2, 3], &full_schema(&dataset)).await.unwrap(); + // id: offset 2 -> newer (9992), offset 3 -> older (7773). + assert_eq!(col(&batch, "id").values(), &[9992, 7773]); + // val: offset 2 -> older (2772), offset 3 -> newer (9993). + assert_eq!(col(&batch, "val").values(), &[2772, 9993]); + } + + /// A fragment with an overlay plan, but a take that touches only uncovered + /// offsets, must fall entirely through to the base values (the + /// `!routing.any_overlay` early-return with a plan present). + #[rstest] + #[tokio::test] + async fn test_take_plan_present_all_offsets_uncovered( + #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_1)] version: LanceFileVersion, + ) { + let dataset = create_base_dataset(version).await; + let dataset = commit_overlay( + dataset, + "ov", + 0, + &[1], + OverlayCoverage::dense(bitmap([1, 4])), + vec![i32_array([Some(111), Some(444)])], + version, + ) + .await; + + let frag = dataset.get_fragment(0).unwrap(); + // None of {0, 2, 5} are covered: the plan exists but contributes nothing. + let batch = frag.take(&[0, 2, 5], &full_schema(&dataset)).await.unwrap(); + assert_eq!(col(&batch, "val").values(), &[0, 20, 50]); + assert_eq!(col(&batch, "id").values(), &[0, 2, 5]); + } + + /// A dataset-level `take` spanning multiple fragments, each with its own + /// overlay, routes every global row index to the right fragment's overlay. + #[rstest] + #[tokio::test] + async fn test_dataset_take_multi_fragment_overlays( + #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_1)] version: LanceFileVersion, + ) { + let dataset = create_base_dataset(version).await; + let dataset = commit_overlay( + dataset, + "frag0", + 0, + &[1], + OverlayCoverage::dense(bitmap([0])), + vec![i32_array([Some(1000)])], + version, + ) + .await; + let dataset = commit_overlay( + dataset, + "frag1", + 1, + &[1], + OverlayCoverage::dense(bitmap([0])), + vec![i32_array([Some(6000)])], + version, + ) + .await; + + // Global rows 0 and 6 are the overlaid offset-0 rows of fragments 0 and + // 1; rows 1 and 7 fall through to base. + let batch = dataset + .take(&[0, 1, 6, 7], full_schema(&dataset)) + .await + .unwrap(); + assert_eq!(col(&batch, "id").values(), &[0, 1, 6, 7]); + assert_eq!(col(&batch, "val").values(), &[1000, 10, 6000, 70]); + } + + /// A scan whose read splits into multiple batches must slice + /// `offsets_in_frag` per batch correctly — the running `rows_seen` + /// accumulator in `merge_overlays` gives each batch its start. Every other + /// scan test uses single-batch fragments, so this is the only guard for the + /// cross-batch (`start > 0`) path. + #[rstest] + #[tokio::test] + async fn test_scan_multi_batch_overlay_slicing( + #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_1)] version: LanceFileVersion, + ) { + use futures::TryStreamExt; + + // One fragment of 10 rows so the read can be chunked below. + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", DataType::Int32, true), + ArrowField::new("val", DataType::Int32, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from_iter_values(0..10)), + Arc::new(Int32Array::from_iter_values((0..10).map(|v| v * 10))), + ], + ) + .unwrap(); + let write_params = WriteParams { + max_rows_per_file: 100, + max_rows_per_group: 100, + data_storage_version: Some(version), + ..Default::default() + }; + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + let dataset = Dataset::write(reader, "memory://", Some(write_params)) + .await + .unwrap(); + + // Overlay one offset in each batch that batch_size 4 produces (batches + // [0,4), [4,8), [8,10)): offsets 1, 5, 9 with distinct values. A wrong + // per-batch slice would misalign these. + let dataset = commit_overlay( + dataset, + "multibatch", + 0, + &[1], + OverlayCoverage::dense(bitmap([1, 5, 9])), + vec![i32_array([Some(111), Some(555), Some(999)])], + version, + ) + .await; + + let frag = dataset.get_fragment(0).unwrap(); + let mut scanner = frag.scan(); + scanner.batch_size(4).project(&["val"]).unwrap(); + let batches: Vec = scanner + .try_into_stream() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + // Guard the guard: the read must actually span multiple batches, else + // this would not exercise the cross-batch slice at all. + assert!( + batches.len() > 1, + "expected a multi-batch scan, got {} batch(es)", + batches.len() + ); + + let merged = + arrow_select::concat::concat_batches(&batches[0].schema(), &batches).unwrap(); + let expected: Vec = (0..10) + .map(|i| match i { + 1 => 111, + 5 => 555, + 9 => 999, + other => other * 10, + }) + .collect(); + assert_eq!(col(&merged, "val").values(), &expected); + } + + /// An empty selection must not trip over the overlay path: the plan exists + /// but there are no offsets to route, so the result is an empty batch. + #[rstest] + #[tokio::test] + async fn test_take_empty_selection( + #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_1)] version: LanceFileVersion, + ) { + let dataset = create_base_dataset(version).await; + let dataset = commit_overlay( + dataset, + "ov", + 0, + &[1], + OverlayCoverage::dense(bitmap([1, 4])), + vec![i32_array([Some(111), Some(444)])], + version, + ) + .await; + + let frag = dataset.get_fragment(0).unwrap(); + let batch = frag.take(&[], &full_schema(&dataset)).await.unwrap(); + assert_eq!(batch.num_rows(), 0); + } + + /// Overlays resolve variable-width columns end-to-end, not just fixed-width + /// ones: the value column is fetched through the real file reader (a + /// different value-pushdown path than the fixed-width case) and assembled. + #[rstest] + #[tokio::test] + async fn test_string_overlay_end_to_end( + #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_1)] version: LanceFileVersion, + ) { + use arrow_array::StringArray; + + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", DataType::Int32, true), + ArrowField::new("name", DataType::Utf8, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from_iter_values(0..6)), + Arc::new(StringArray::from(vec!["a", "b", "c", "d", "e", "f"])), + ], + ) + .unwrap(); + let write_params = WriteParams { + max_rows_per_file: 6, + max_rows_per_group: 6, + data_storage_version: Some(version), + ..Default::default() + }; + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + let dataset = Dataset::write(reader, "memory://", Some(write_params)) + .await + .unwrap(); + + // Overlay `name` at offsets {1, 4}, one of the values NULL. + let dataset = commit_overlay( + dataset, + "strov", + 0, + &[1], + OverlayCoverage::dense(bitmap([1, 4])), + vec![Arc::new(StringArray::from(vec![Some("B"), None])) as ArrayRef], + version, + ) + .await; + + let frag = dataset.get_fragment(0).unwrap(); + let batch = frag.take(&[0, 1, 4], &full_schema(&dataset)).await.unwrap(); + let name = batch + .column(batch.schema().index_of("name").unwrap()) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(name.value(0), "a"); // falls through to base + assert_eq!(name.value(1), "B"); // overlay value + assert!(name.is_null(2)); // overlay NULL wins + } + + /// Projection pruning must do NO IO to overlay files whose fields are not + /// projected. Proven the same way as row-selection pruning: delete the + /// overlay's data file, then read projecting only the *unrelated* `id` + /// column — it must succeed (the `val` overlay file is never opened), while + /// projecting the overlaid `val` column then fails because its file is gone. + #[rstest] + #[tokio::test] + async fn test_projection_prunes_overlay_files_no_io( + #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_1)] version: LanceFileVersion, + ) { + let dataset = create_base_dataset(version).await; + // Overlay covers `val` (field 1) only. + let dataset = commit_overlay( + dataset, + "valov", + 0, + &[1], + OverlayCoverage::dense(bitmap([0, 1])), + vec![i32_array([Some(1000), Some(1010)])], + version, + ) + .await; + + // Delete the overlay's data file: opening it now fails. + dataset + .object_store + .delete(&Path::from("data/valov.lance")) + .await + .unwrap(); + + let frag = dataset.get_fragment(0).unwrap(); + let id_only = dataset.schema().project_by_ids(&[0], true); + let val_only = dataset.schema().project_by_ids(&[1], true); + + // Projecting only `id` must not open the `val` overlay file, so it + // succeeds and returns untouched base values. + let batch = frag.take(&[0, 1], &id_only).await.unwrap(); + assert_eq!(col(&batch, "id").values(), &[0, 1]); + // A scan projecting only `id` must likewise never touch the file. + let batch = frag + .scan() + .project(&["id"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(col(&batch, "id").values(), &[0, 1, 2, 3, 4, 5]); + + // Projecting the overlaid `val` column does need the file, so it fails + // with a not-found error naming the missing overlay file. + let err = frag.take(&[0], &val_only).await.unwrap_err(); + let err = format!("{err:?}"); + assert!( + err.contains("valov.lance") && err.to_lowercase().contains("not found"), + "projecting the overlaid column should fail with a not-found error \ + for its missing file, got: {err}", + ); + } + + /// A top-level struct column resolves through overlays: the overlay stores + /// the struct's leaf columns (under V2_1 those are the only ids in + /// `data_file.fields`), and `plan_overlays` maps them back to the top-level + /// struct so the whole value is fetched and replaced as a unit. + #[rstest] + #[tokio::test] + async fn test_struct_overlay_end_to_end( + #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_1)] version: LanceFileVersion, + ) { + let struct_fields = Fields::from(vec![ + ArrowField::new("x", DataType::Int32, true), + ArrowField::new("y", DataType::Int32, true), + ]); + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", DataType::Int32, true), + ArrowField::new("info", DataType::Struct(struct_fields.clone()), true), + ])); + let info = Arc::new(StructArray::new( + struct_fields.clone(), + vec![ + Arc::new(Int32Array::from_iter_values(0..6)), + Arc::new(Int32Array::from_iter_values((0..6).map(|v| v * 100))), + ], + None, + )); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from_iter_values(0..6)), info], + ) + .unwrap(); + let write_params = WriteParams { + max_rows_per_file: 6, + max_rows_per_group: 6, + data_storage_version: Some(version), + ..Default::default() + }; + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + let dataset = Dataset::write(reader, "memory://", Some(write_params)) + .await + .unwrap(); + + // Overlay the whole `info` struct (top-level field id 1) at offset 2. + let overlay_info = Arc::new(StructArray::new( + struct_fields, + vec![ + Arc::new(Int32Array::from(vec![777])), + Arc::new(Int32Array::from(vec![888])), + ], + None, + )) as ArrayRef; + let dataset = commit_overlay( + dataset, + "structov", + 0, + &[1], + OverlayCoverage::dense(bitmap([2])), + vec![overlay_info], + version, + ) + .await; + + let frag = dataset.get_fragment(0).unwrap(); + let batch = frag.take(&[1, 2], &full_schema(&dataset)).await.unwrap(); + let info = batch + .column(batch.schema().index_of("info").unwrap()) + .as_any() + .downcast_ref::() + .unwrap(); + let x = info + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let y = info + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + // Offset 1 falls through to base {1, 100}; offset 2 takes the overlay. + assert_eq!(x.values(), &[1, 777]); + assert_eq!(y.values(), &[100, 888]); + } + + /// A top-level list column resolves through overlays the same way — the + /// overlay's leaf (item) id maps back to the top-level list, and the whole + /// list value at a covered offset is replaced. + #[rstest] + #[tokio::test] + async fn test_list_overlay_end_to_end( + #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_1)] version: LanceFileVersion, + ) { + use arrow_array::ListArray; + use arrow_array::types::Int32Type; + + let item = Arc::new(ArrowField::new("item", DataType::Int32, true)); + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", DataType::Int32, true), + ArrowField::new("tags", DataType::List(item.clone()), true), + ])); + let base_tags = ListArray::from_iter_primitive::( + (0..6i32).map(|i| Some(vec![Some(i), Some(i * 10)])), + ); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from_iter_values(0..6)), + Arc::new(base_tags), + ], + ) + .unwrap(); + let write_params = WriteParams { + max_rows_per_file: 6, + max_rows_per_group: 6, + data_storage_version: Some(version), + ..Default::default() + }; + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + let dataset = Dataset::write(reader, "memory://", Some(write_params)) + .await + .unwrap(); + + // Overlay `tags` (top-level field id 1) at offset 2 with a new list. + let overlay_tags = + ListArray::from_iter_primitive::(std::iter::once(Some(vec![ + Some(77), + Some(88), + Some(99), + ]))); + let dataset = commit_overlay( + dataset, + "listov", + 0, + &[1], + OverlayCoverage::dense(bitmap([2])), + vec![Arc::new(overlay_tags) as ArrayRef], + version, + ) + .await; + + let frag = dataset.get_fragment(0).unwrap(); + let batch = frag.take(&[1, 2], &full_schema(&dataset)).await.unwrap(); + let tags = batch + .column(batch.schema().index_of("tags").unwrap()) + .as_any() + .downcast_ref::() + .unwrap(); + let row1 = tags.value(0); + let row1 = row1.as_any().downcast_ref::().unwrap(); + let row2 = tags.value(1); + let row2 = row2.as_any().downcast_ref::().unwrap(); + // Offset 1 falls through to base [1, 10]; offset 2 takes the overlay. + assert_eq!(row1.values(), &[1, 10]); + assert_eq!(row2.values(), &[77, 88, 99]); + } + + /// A top-level Map column resolves as a single atomic field even though its + /// value spans two leaves (key and value): both leaf ids map back to the one + /// Map atomic field, and the whole map value at a covered offset is replaced. + /// Maps require + /// the 2.2+ file format, so this runs only at V2_2 (unlike the V2_0/V2_1 + /// parametrized tests). + #[tokio::test] + async fn test_map_overlay_end_to_end() { + use arrow_array::MapArray; + use arrow_array::builder::{Int32Builder, MapBuilder}; + + let version = LanceFileVersion::V2_2; + + // Base row i holds the single entry {i: i * 10}. + let mut builder = MapBuilder::new(None, Int32Builder::new(), Int32Builder::new()); + for i in 0..6i32 { + builder.keys().append_value(i); + builder.values().append_value(i * 10); + builder.append(true).unwrap(); + } + let base_attrs = builder.finish(); + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", DataType::Int32, true), + ArrowField::new("attrs", base_attrs.data_type().clone(), true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from_iter_values(0..6)), + Arc::new(base_attrs), + ], + ) + .unwrap(); + let write_params = WriteParams { + max_rows_per_file: 6, + max_rows_per_group: 6, + data_storage_version: Some(version), + ..Default::default() + }; + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + let dataset = Dataset::write(reader, "memory://", Some(write_params)) + .await + .unwrap(); + + // Overlay `attrs` (top-level field id 1) at offset 2 with a two-entry map. + let mut ov = MapBuilder::new(None, Int32Builder::new(), Int32Builder::new()); + ov.keys().append_value(7); + ov.values().append_value(77); + ov.keys().append_value(8); + ov.values().append_value(88); + ov.append(true).unwrap(); + let overlay_attrs = ov.finish(); + let dataset = commit_overlay( + dataset, + "mapov", + 0, + &[1], + OverlayCoverage::dense(bitmap([2])), + vec![Arc::new(overlay_attrs) as ArrayRef], + version, + ) + .await; + + let frag = dataset.get_fragment(0).unwrap(); + let batch = frag.take(&[1, 2], &full_schema(&dataset)).await.unwrap(); + let attrs = batch + .column(batch.schema().index_of("attrs").unwrap()) + .as_any() + .downcast_ref::() + .unwrap(); + + let entries = |i: usize| -> (Vec, Vec) { + let row = attrs.value(i); + let keys = row.column(0).as_any().downcast_ref::().unwrap(); + let vals = row.column(1).as_any().downcast_ref::().unwrap(); + (keys.values().to_vec(), vals.values().to_vec()) + }; + // Offset 1 falls through to the base entry {1: 10}; offset 2 takes the + // overlay map {7: 77, 8: 88}. + assert_eq!(entries(0), (vec![1], vec![10])); + assert_eq!(entries(1), (vec![7, 8], vec![77, 88])); + } + + /// Base `id` + a struct `s { a, b }` (6 rows). Field ids: s=1, a=2, b=3. + async fn create_struct_dataset(version: LanceFileVersion) -> (Dataset, Fields) { + let s_fields = Fields::from(vec![ + ArrowField::new("a", DataType::Int32, true), + ArrowField::new("b", DataType::Int32, true), + ]); + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", DataType::Int32, true), + ArrowField::new("s", DataType::Struct(s_fields.clone()), true), + ])); + let s = Arc::new(StructArray::new( + s_fields.clone(), + vec![ + Arc::new(Int32Array::from_iter_values(0..6)), + Arc::new(Int32Array::from_iter_values((0..6).map(|v| v * 100))), + ], + None, + )); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from_iter_values(0..6)), s], + ) + .unwrap(); + let write_params = WriteParams { + max_rows_per_file: 6, + max_rows_per_group: 6, + data_storage_version: Some(version), + ..Default::default() + }; + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + let dataset = Dataset::write(reader, "memory://", Some(write_params)) + .await + .unwrap(); + (dataset, s_fields) + } + + fn struct_col<'a>(batch: &'a RecordBatch, name: &str) -> &'a StructArray { + batch + .column(batch.schema().index_of(name).unwrap()) + .as_any() + .downcast_ref::() + .unwrap() + } + + fn i32_child(s: &StructArray, i: usize) -> Int32Array { + s.column(i) + .as_any() + .downcast_ref::() + .unwrap() + .clone() + } + + /// The reviewer's core case (r3553495147): an overlay stores only sub-field + /// `s.a`, but the read projects the whole struct `s`. The overlay must splice + /// into `a` and leave `b` untouched (previously this panicked because the merge + /// fetched the whole `s` from an overlay file holding only `a`). + #[rstest] + #[tokio::test] + async fn test_overlay_subfield_projecting_parent_struct( + #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_1)] version: LanceFileVersion, + ) { + let (dataset, _) = create_struct_dataset(version).await; + // Overlay ONLY `s.a` (field id 2) at offset 2. + let a_only = Fields::from(vec![ArrowField::new("a", DataType::Int32, true)]); + let overlay = Arc::new(StructArray::new( + a_only, + vec![Arc::new(Int32Array::from(vec![777]))], + None, + )) as ArrayRef; + let dataset = commit_overlay( + dataset, + "aov", + 0, + &[2], + OverlayCoverage::dense(bitmap([2])), + vec![overlay], + version, + ) + .await; + + let frag = dataset.get_fragment(0).unwrap(); + let batch = frag.take(&[1, 2], &full_schema(&dataset)).await.unwrap(); + let s = struct_col(&batch, "s"); + // a: offset 1 base (1), offset 2 overlaid (777). + assert_eq!(i32_child(s, 0).values(), &[1, 777]); + // b: untouched base (100, 200). + assert_eq!(i32_child(s, 1).values(), &[100, 200]); + } + + /// An overlay on a non-projected sibling leaf must be skipped and its file + /// never opened: overlay covers `s.b`, but the read projects only `s.a`. + #[rstest] + #[tokio::test] + async fn test_overlay_nonprojected_sibling_skipped( + #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_1)] version: LanceFileVersion, + ) { + let (dataset, _) = create_struct_dataset(version).await; + let b_only = Fields::from(vec![ArrowField::new("b", DataType::Int32, true)]); + let overlay = Arc::new(StructArray::new( + b_only, + vec![Arc::new(Int32Array::from(vec![888]))], + None, + )) as ArrayRef; + let dataset = commit_overlay( + dataset, + "bov", + 0, + &[3], + OverlayCoverage::dense(bitmap([2])), + vec![overlay], + version, + ) + .await; + // Delete the overlay file: if projecting only `s.a` opened it, this fails. + dataset + .object_store + .delete(&Path::from("data/bov.lance")) + .await + .unwrap(); + + let frag = dataset.get_fragment(0).unwrap(); + let a_only = dataset.schema().project_by_ids(&[2], true); + let batch = frag.take(&[1, 2], &a_only).await.unwrap(); + let s = struct_col(&batch, "s"); + // Only `a` is projected, unchanged base values. + assert_eq!(i32_child(s, 0).values(), &[1, 2]); + } + + /// Two overlays target different sub-fields of the same struct, and a third + /// re-overlays `s.a`. Each leaf resolves independently and newest wins on `a`. + #[rstest] + #[tokio::test] + async fn test_overlay_multiple_subfields_newest_wins( + #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_1)] version: LanceFileVersion, + ) { + let (dataset, _) = create_struct_dataset(version).await; + let a_field = Fields::from(vec![ArrowField::new("a", DataType::Int32, true)]); + let b_field = Fields::from(vec![ArrowField::new("b", DataType::Int32, true)]); + // Older: a := 700 at offset 2. + let dataset = commit_overlay( + dataset, + "a_old", + 0, + &[2], + OverlayCoverage::dense(bitmap([2])), + vec![Arc::new(StructArray::new( + a_field.clone(), + vec![Arc::new(Int32Array::from(vec![700]))], + None, + )) as ArrayRef], + version, + ) + .await; + // b := 800 at offset 2. + let dataset = commit_overlay( + dataset, + "b_ov", + 0, + &[3], + OverlayCoverage::dense(bitmap([2])), + vec![Arc::new(StructArray::new( + b_field, + vec![Arc::new(Int32Array::from(vec![800]))], + None, + )) as ArrayRef], + version, + ) + .await; + // Newest: a := 999 at offset 2 (shadows the older `a` overlay). + let dataset = commit_overlay( + dataset, + "a_new", + 0, + &[2], + OverlayCoverage::dense(bitmap([2])), + vec![Arc::new(StructArray::new( + a_field, + vec![Arc::new(Int32Array::from(vec![999]))], + None, + )) as ArrayRef], + version, + ) + .await; + + let frag = dataset.get_fragment(0).unwrap(); + let batch = frag.take(&[2], &full_schema(&dataset)).await.unwrap(); + let s = struct_col(&batch, "s"); + assert_eq!(i32_child(s, 0).values(), &[999]); // newest `a` wins + assert_eq!(i32_child(s, 1).values(), &[800]); // `b` from its own overlay + } + + /// Three levels of nesting: `outer { middle { a, b } }`. An overlay on the + /// deep leaf `outer.middle.a` splices correctly when the whole `outer` is read. + #[rstest] + #[tokio::test] + async fn test_overlay_deeply_nested_subfield( + #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_1)] version: LanceFileVersion, + ) { + let mid_fields = Fields::from(vec![ + ArrowField::new("a", DataType::Int32, true), + ArrowField::new("b", DataType::Int32, true), + ]); + let outer_fields = Fields::from(vec![ArrowField::new( + "middle", + DataType::Struct(mid_fields.clone()), + true, + )]); + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", DataType::Int32, true), + ArrowField::new("outer", DataType::Struct(outer_fields.clone()), true), + ])); + // Field ids: outer=1, middle=2, a=3, b=4. + let middle = Arc::new(StructArray::new( + mid_fields.clone(), + vec![ + Arc::new(Int32Array::from_iter_values(0..6)), + Arc::new(Int32Array::from_iter_values((0..6).map(|v| v * 100))), + ], + None, + )); + let outer = Arc::new(StructArray::new(outer_fields, vec![middle], None)); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from_iter_values(0..6)), outer], + ) + .unwrap(); + let write_params = WriteParams { + max_rows_per_file: 6, + max_rows_per_group: 6, + data_storage_version: Some(version), + ..Default::default() + }; + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + let dataset = Dataset::write(reader, "memory://", Some(write_params)) + .await + .unwrap(); + + // Overlay the deep leaf `outer.middle.a` (field id 3) at offset 2. + let a_leaf = Fields::from(vec![ArrowField::new("a", DataType::Int32, true)]); + let mid_a = Fields::from(vec![ArrowField::new( + "middle", + DataType::Struct(a_leaf.clone()), + true, + )]); + let overlay = Arc::new(StructArray::new( + mid_a, + vec![Arc::new(StructArray::new( + a_leaf, + vec![Arc::new(Int32Array::from(vec![777]))], + None, + ))], + None, + )) as ArrayRef; + let dataset = commit_overlay( + dataset, + "deepov", + 0, + &[3], + OverlayCoverage::dense(bitmap([2])), + vec![overlay], + version, + ) + .await; + + let frag = dataset.get_fragment(0).unwrap(); + let batch = frag.take(&[1, 2], &full_schema(&dataset)).await.unwrap(); + let outer = struct_col(&batch, "outer"); + let middle = outer + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + // a: offset 1 base (1), offset 2 overlaid (777); b untouched. + assert_eq!(i32_child(middle, 0).values(), &[1, 777]); + assert_eq!(i32_child(middle, 1).values(), &[100, 200]); + + // Projecting the *intermediate* struct `outer.middle` (field id 2) while + // the overlay targets a deeper field (id 3) must still apply: the + // overlay's leaf id falls inside the projected subtree, so it maps to a + // projected atomic field. (This is the case wjones127/westonpace flagged where a + // top-level-only mapping would miss the overlay.) + let middle_only = dataset.schema().project_by_ids(&[2], true); + let batch = frag.take(&[2], &middle_only).await.unwrap(); + let middle = struct_col(&batch, "outer") + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(i32_child(middle, 0).values(), &[777]); + } + } + #[rstest] #[tokio::test] async fn test_fragment_scan( diff --git a/rust/lance/src/dataset/overlay.rs b/rust/lance/src/dataset/overlay.rs new file mode 100644 index 00000000000..d73329d8a41 --- /dev/null +++ b/rust/lance/src/dataset/overlay.rs @@ -0,0 +1,1064 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Resolution of data overlay files on read. +//! +//! An overlay supplies replacement values for some `(row, field)` cells without +//! rewriting the base data. Resolving a read means, for each row we return, +//! deciding whether its value comes from the base column or from an overlay. +//! +//! Three coordinate spaces show up throughout this module; keeping them straight +//! is most of the work: +//! +//! - `offset_in_frag`: a row's physical position in the fragment (0-based over all +//! physical rows, ignoring deletions). This is how a cell is addressed on disk +//! and in an overlay's coverage bitmap. +//! - `offset_in_batch`: a row's position within the batch we are currently +//! assembling (0-based). The output column is indexed by this. +//! - `offset_in_overlay`: the position of a value in an overlay's value column. +//! An overlay stores its values densely — one per covered cell, in ascending +//! `offset_in_frag` order — so a covered cell's value is found by counting how +//! many covered cells come before it. (That count is what a roaring bitmap calls +//! the cell's "rank".) +//! +//! For a given field, the overlays covering it are consulted newest to oldest: the +//! first overlay that covers a row wins, and its value is read at that row's +//! `offset_in_overlay`. A row that no overlay covers keeps its base value. +//! +//! The rows to resolve are passed in as a list of `offset_in_frag` (one per output +//! row), so a single code path serves both scans (a contiguous range of offsets) +//! and `take` (arbitrary offsets). +//! +//! Deletions win over overlays, but nothing here handles that: the merge runs on +//! physical rows *before* deletions are applied, so an overlay value computed for a +//! deleted row is simply dropped along with the row. This matches the spec with no +//! special casing. + +use std::collections::{BTreeSet, HashMap}; +use std::sync::Arc; + +use arrow_array::{Array, ArrayRef, RecordBatch, StructArray}; +use arrow_select::interleave::interleave; +use futures::StreamExt; +use lance_core::datatypes::{Field, Schema}; +use lance_core::{Error, Result}; +use roaring::RoaringBitmap; + +use lance_table::format::DataFile; +use lance_table::utils::stream::ReadBatchFut; + +use crate::dataset::fragment::{FileFragment, FragReadConfig, GenericFileReader}; + +/// The plan for merging one field's overlays into one batch: which source (base or +/// a particular overlay) supplies each output row, and which overlay values must be +/// fetched to do it. +/// +/// Built by [`route_overlays`] from the coverage bitmaps alone — before any value +/// column is read — so the caller can fetch only the overlay values it will +/// actually use (its `offsets_in_overlay`) rather than whole columns, then build the +/// merged column with [`assemble_overlay_column`]. +struct OverlayRouting { + /// One `(source, position)` pair per output row, ready to hand to `interleave`. + /// Source `0` is the base column, with `position` = the row's `offset_in_batch`; + /// source `k + 1` is overlay `k`'s fetched values, with `position` = the row's + /// index into those fetched values. + indices: Vec<(usize, usize)>, + /// Per overlay (newest-first): the sorted, deduplicated `offset_in_overlay` + /// values this batch needs from that overlay — i.e. exactly which entries of its + /// value column to fetch. + offsets_in_overlay: Vec>, + /// Whether any row is covered by an overlay at all (false ⇒ every row falls + /// through to the base column, so the base is already the answer and no overlay + /// values need to be read). + any_overlay: bool, +} + +/// For each row in `offsets_in_frag`, decide whether its value comes from the base +/// column or from an overlay — and if from an overlay, at which `offset_in_overlay`. +/// +/// Only the coverage bitmaps are consulted (newest-first), so this runs before any +/// value column is read and reports exactly which overlay values the caller must +/// fetch. +/// +/// A scan asks for a contiguous, ascending range of offsets, which enables a faster +/// bitmap-driven path ([`route_contiguous`]); `take` asks for arbitrary offsets and +/// uses the general path ([`route_arbitrary`]). Both produce identical routing. +fn route_overlays( + offsets_in_frag: &[u32], + coverages_newest_first: &[&RoaringBitmap], +) -> OverlayRouting { + match contiguous_frag_start(offsets_in_frag) { + Some(frag_start) => { + route_contiguous(frag_start, offsets_in_frag.len(), coverages_newest_first) + } + None => route_arbitrary(offsets_in_frag, coverages_newest_first), + } +} + +/// If `offsets_in_frag` is a contiguous ascending run `[start, start + 1, ...]`, +/// return `start`; otherwise `None` (including when empty). +fn contiguous_frag_start(offsets_in_frag: &[u32]) -> Option { + let start = *offsets_in_frag.first()?; + offsets_in_frag + .iter() + .enumerate() + .all(|(i, &offset)| offset as u64 == start as u64 + i as u64) + .then_some(start) +} + +/// Fast path for a scan, where the batch is a contiguous run of offsets starting at +/// `frag_start`. Because the offsets are contiguous, a row's `offset_in_batch` is +/// just `offset_in_frag - frag_start`, so a coverage's set bits map straight to +/// output rows — no need to test each row against each coverage. +/// +/// For each coverage we intersect it with the batch's offset range. Roaring does +/// this a block at a time, so a coverage that does not overlap the batch (e.g. a +/// scan batch past the last cell this overlay touches) is skipped cheaply without +/// inspecting individual bits. +/// +/// Within the batch a coverage's cells appear in ascending order, so their +/// `offset_in_overlay` values are consecutive: the first in-batch cell sits at +/// `offset_in_overlay = ` (a single +/// `rank` lookup), and each following cell is one more. Coverages are applied +/// newest-first, and the first overlay to claim a row wins. +fn route_contiguous( + frag_start: u32, + len: usize, + coverages_newest_first: &[&RoaringBitmap], +) -> OverlayRouting { + let mut offsets_in_overlay: Vec> = vec![Vec::new(); coverages_newest_first.len()]; + // Indexed by offset_in_batch: which (overlay, fetch position) supplies the row. + let mut routed: Vec> = vec![None; len]; + let range_end = (frag_start as u64 + len as u64).min(u32::MAX as u64) as u32; + let mut batch_range = RoaringBitmap::new(); + batch_range.insert_range(frag_start..range_end); + + for (k, coverage) in coverages_newest_first.iter().enumerate() { + let covered_in_batch = *coverage & &batch_range; + if covered_in_batch.is_empty() { + continue; + } + // offset_in_overlay of this coverage's first in-batch cell = the number of + // its cells that lie before the batch. + let first_offset_in_overlay = if frag_start == 0 { + 0 + } else { + coverage.rank(frag_start - 1) as u32 + }; + for (nth_in_batch, offset_in_frag) in covered_in_batch.iter().enumerate() { + let offset_in_batch = (offset_in_frag - frag_start) as usize; + if routed[offset_in_batch].is_none() { + routed[offset_in_batch] = Some((k, offsets_in_overlay[k].len())); + offsets_in_overlay[k].push(first_offset_in_overlay + nth_in_batch as u32); + } + } + } + + let mut any_overlay = false; + let indices = routed + .into_iter() + .enumerate() + .map(|(offset_in_batch, routed)| match routed { + None => (0, offset_in_batch), + Some((k, fetch_pos)) => { + any_overlay = true; + (k + 1, fetch_pos) + } + }) + .collect(); + + OverlayRouting { + indices, + offsets_in_overlay, + any_overlay, + } +} + +/// General path for arbitrary offsets (e.g. `take`): test each row's +/// `offset_in_frag` against the coverages newest-first. `take` batches are small, +/// so this `O(rows * overlays)` probing is not a concern. +fn route_arbitrary( + offsets_in_frag: &[u32], + coverages_newest_first: &[&RoaringBitmap], +) -> OverlayRouting { + // Per overlay: the distinct offset_in_overlay values this batch needs, sorted. + let mut offset_sets: Vec> = vec![BTreeSet::new(); coverages_newest_first.len()]; + // Per output row: the (overlay, offset_in_overlay) that supplies it, if any. + let mut routed_per_row: Vec> = Vec::with_capacity(offsets_in_frag.len()); + for &offset_in_frag in offsets_in_frag { + let mut routed = None; + for (k, coverage) in coverages_newest_first.iter().enumerate() { + if coverage.contains(offset_in_frag) { + // offset_in_overlay = number of covered cells before this one. + let offset_in_overlay = coverage.rank(offset_in_frag) as u32 - 1; + offset_sets[k].insert(offset_in_overlay); + routed = Some((k, offset_in_overlay)); + break; + } + } + routed_per_row.push(routed); + } + + let offsets_in_overlay: Vec> = offset_sets + .iter() + .map(|offsets| offsets.iter().copied().collect()) + .collect(); + // For each overlay, map an offset_in_overlay to its position in the fetched + // (sorted, deduplicated) value list. + let fetch_positions: Vec> = offsets_in_overlay + .iter() + .map(|offsets| { + offsets + .iter() + .enumerate() + .map(|(pos, &o)| (o, pos)) + .collect() + }) + .collect(); + + let mut any_overlay = false; + let indices = routed_per_row + .into_iter() + .enumerate() + .map(|(offset_in_batch, routed)| match routed { + None => (0, offset_in_batch), + Some((k, offset_in_overlay)) => { + any_overlay = true; + (k + 1, fetch_positions[k][&offset_in_overlay]) + } + }) + .collect(); + + OverlayRouting { + indices, + offsets_in_overlay, + any_overlay, + } +} + +/// Build the merged column from `base` and the overlay values fetched for the +/// `offset_in_overlay` values [`route_overlays`] asked for. +/// +/// `fetched_newest_first[k]` holds overlay `k`'s values for `routing`'s +/// `offsets_in_overlay[k]`, in that order. The result has the same length and +/// type as `base`. A covered row whose overlay value is NULL resolves **to** NULL +/// (distinct from a fall-through, which keeps the base value). +fn assemble_overlay_column( + base: &ArrayRef, + routing: &OverlayRouting, + fetched_newest_first: &[ArrayRef], +) -> Result { + if !routing.any_overlay { + return Ok(base.clone()); + } + if fetched_newest_first.len() != routing.offsets_in_overlay.len() { + return Err(Error::invalid_input(format!( + "overlay assembly got {} value columns but routing expects {}", + fetched_newest_first.len(), + routing.offsets_in_overlay.len() + ))); + } + for (k, values) in fetched_newest_first.iter().enumerate() { + if values.len() != routing.offsets_in_overlay[k].len() { + return Err(Error::invalid_input(format!( + "overlay value column {} has {} values but {} were requested", + k, + values.len(), + routing.offsets_in_overlay[k].len() + ))); + } + } + + let mut sources: Vec<&dyn Array> = Vec::with_capacity(fetched_newest_first.len() + 1); + sources.push(base.as_ref()); + for values in fetched_newest_first { + sources.push(values.as_ref()); + } + interleave(&sources, &routing.indices).map_err(Error::from) +} + +/// One overlay's contribution to one projected atomic field, with its file reader opened. +#[derive(Debug, Clone)] +struct LoadedAtomicFieldOverlay { + /// The `offset_in_frag` cells this overlay covers for the atomic field. + coverage: Arc, + /// Reader over the overlay data file, projected to the covered atomic fields; shared + /// across the atomic fields that the same file covers. + reader: Arc, +} + +/// The overlays that apply to a single projected atomic field — a per-row field an overlay +/// can replace as a unit (a primitive leaf, or a whole list/map field; structs are +/// recursed through, not treated as atomic fields). Ordered newest-first, with readers opened +/// and pruned to a specific read. Produced by [`resolve_overlays`] and consumed by +/// [`merge_overlay_batch`]. +#[derive(Debug, Clone)] +pub struct LoadedAtomicField { + /// The top-level output column the atomic field lives in (its name locates the batch + /// column; its field tree drives the descend/splice into that column). + top_field: Arc, + /// Child field ids from `top_field` down to the atomic field (empty when the atomic + /// field *is* the top-level column). Drives descending to, and splicing back, the + /// atomic field. + ancestor_ids: Vec, + /// Projection of exactly the atomic field (its ancestor path pruned to the atomic + /// field subtree), used to fetch the atomic field's values from the overlay file. + fetch_projection: Arc, + overlays_newest_first: Vec, +} + +/// One overlay file that may contribute to a read, before it is opened. Opened +/// lazily by [`resolve_overlays`], and only if the read actually touches it. +#[derive(Debug, Clone)] +struct PlannedOverlayFile { + data_file: DataFile, + /// The covered ∩ projected atomic fields to project when the file is opened, so a single + /// reader serves every atomic field the file contributes to. + open_projection: Arc, +} + +/// One overlay's contribution to one projected atomic field, before the file is opened. +#[derive(Debug, Clone)] +struct PlannedAtomicFieldOverlay { + /// Index into [`OverlayReadPlanner::files`] of the file that supplies the value. + file: usize, + coverage: Arc, +} + +/// The overlays that apply to a single projected atomic field, ordered newest-first, before +/// any file is opened. +#[derive(Debug, Clone)] +struct PlannedAtomicField { + top_field: Arc, + ancestor_ids: Vec, + fetch_projection: Arc, + overlays_newest_first: Vec, +} + +/// A fragment's overlay-resolution plan for a projection, derived from coverage +/// metadata alone — no file opened, no IO. [`resolve_overlays`] turns it into opened +/// [`LoadedAtomicField`]s for one specific read, opening only the files whose cells +/// the read's rows actually touch. +#[derive(Debug, Clone)] +pub struct OverlayReadPlanner { + files: Vec, + atomic_fields: Vec, +} + +impl OverlayReadPlanner { + /// True when no projected atomic field has any overlay, so there is nothing to resolve. + pub fn is_empty(&self) -> bool { + self.atomic_fields.is_empty() + } +} + +/// Plan `fragment`'s overlay resolution for a projection from coverage metadata +/// alone. No files are opened here (see [`resolve_overlays`]) — this only reads the +/// already-parsed coverage bitmaps, so it is cheap enough to run on every open. +/// +/// Overlays are stored oldest-first (sorted newest-last on load, see +/// `sort_overlays_newest_last`), so walking them in reverse gives newest-first +/// precedence. +/// +/// Resolution is per *atomic field* — a per-row field that an overlay replaces as a unit: a +/// primitive leaf, or a whole list/map field. Structs are internal nodes, so each +/// leaf of a struct is its own atomic field and can be overlaid independently of its +/// siblings. An overlay is written against the leaf ids it stores (the V2_1 +/// structural encoding records only leaves), so an overlay contributes to a projected +/// atomic field when any id in its `data_file.fields` falls in that atomic field's leaf +/// set. At merge time the atomic field's value is fetched and spliced into its output +/// column, so an overlay on a sub-field never disturbs the column's other leaves. Each +/// contributing overlay *file* appears once in `files`, shared by every atomic field it +/// covers. +pub fn plan_overlays(fragment: &FileFragment, projection: &Schema) -> Result { + let overlays = &fragment.metadata.overlays; + debug_assert!( + overlays + .windows(2) + .all(|w| w[0].committed_version <= w[1].committed_version), + "overlays must be sorted newest-last (see sort_overlays_newest_last)" + ); + + // The projection's atomic fields, and a leaf-id -> atomic-field-index map so an + // overlay's stored leaf ids resolve to the atomic field they belong to in O(1). + struct AtomicFieldInfo<'a> { + top_field: &'a Field, + ancestor_ids: Vec, + atomic_field_id: i32, + } + let mut atomic_field_infos: Vec = Vec::new(); + let mut leaf_to_atomic_field: HashMap = HashMap::new(); + for top in &projection.fields { + for (atomic_field, ancestor_ids) in enumerate_atomic_fields(top) { + let idx = atomic_field_infos.len(); + let mut value_leaf_ids = Vec::new(); + collect_leaf_ids(atomic_field, &mut value_leaf_ids); + for leaf in value_leaf_ids { + leaf_to_atomic_field.insert(leaf, idx); + } + atomic_field_infos.push(AtomicFieldInfo { + top_field: top, + ancestor_ids, + atomic_field_id: atomic_field.id, + }); + } + } + + // Walk overlays newest-first. For each overlay, find the atomic fields it covers and push + // (newest-first, for free) into their per-atomic field overlay lists. + let mut files = Vec::new(); + let mut atomic_field_overlays: Vec> = + vec![Vec::new(); atomic_field_infos.len()]; + for overlay in overlays.iter().rev() { + // atomic field index -> the `data_file.fields` position whose coverage to read. An + // overlay writes one value per row per atomic field, so its leaves share a coverage; + // the first leaf of each atomic field to appear wins. + let mut covered: HashMap = HashMap::new(); + for (field_pos, &field_id) in overlay.data_file.fields.iter().enumerate() { + if let Some(&atomic_field_idx) = leaf_to_atomic_field.get(&field_id) { + covered.entry(atomic_field_idx).or_insert(field_pos); + } + } + if covered.is_empty() { + continue; + } + let file = files.len(); + let covered_ids: Vec = covered + .keys() + .map(|&i| atomic_field_infos[i].atomic_field_id) + .collect(); + files.push(PlannedOverlayFile { + data_file: overlay.data_file.clone(), + open_projection: Arc::new(projection.project_by_ids(&covered_ids, true)), + }); + for (atomic_field_idx, field_pos) in covered { + atomic_field_overlays[atomic_field_idx].push(PlannedAtomicFieldOverlay { + file, + coverage: overlay.coverage_for_field(field_pos)?, + }); + } + } + + // Emit one PlannedAtomicField per projected atomic field that has overlays, in + // atomic field order. + let mut atomic_fields = Vec::new(); + for (idx, info) in atomic_field_infos.iter().enumerate() { + let overlays_newest_first = std::mem::take(&mut atomic_field_overlays[idx]); + if overlays_newest_first.is_empty() { + continue; + } + atomic_fields.push(PlannedAtomicField { + top_field: Arc::new(info.top_field.clone()), + ancestor_ids: info.ancestor_ids.clone(), + fetch_projection: Arc::new(projection.project_by_ids(&[info.atomic_field_id], true)), + overlays_newest_first, + }); + } + Ok(OverlayReadPlanner { + files, + atomic_fields, + }) +} + +/// The per-row atomic fields of a projected top-level field, each with the child-id path from +/// the top-level field down to it. Structs are recursed through; a primitive leaf or a +/// whole list/map field is an atomic field (values are one-per-row). A top-level primitive or +/// list yields a single atomic field with an empty path. +fn enumerate_atomic_fields(top: &Field) -> Vec<(&Field, Vec)> { + fn recurse<'a>(field: &'a Field, path: &mut Vec, out: &mut Vec<(&'a Field, Vec)>) { + if field.logical_type.is_struct() { + for child in &field.children { + path.push(child.id); + recurse(child, path, out); + path.pop(); + } + } else { + out.push((field, path.clone())); + } + } + let mut out = Vec::new(); + let mut path = Vec::new(); + recurse(top, &mut path, &mut out); + out +} + +/// Collect the leaf field ids in `field`'s subtree — the ids an overlay stores for +/// this atomic field (its own id if primitive; its item leaves if a list/map). +fn collect_leaf_ids(field: &Field, out: &mut Vec) { + if field.children.is_empty() { + out.push(field.id); + } else { + for child in &field.children { + collect_leaf_ids(child, out); + } + } +} + +/// Follow a path of child field ids from `field` down through nested structs, taking +/// the corresponding child array at each step. Returns the array at the end of the +/// path (the whole `array` when `ancestor_ids` is empty). +fn descend_by_ids(array: &ArrayRef, field: &Field, ancestor_ids: &[i32]) -> Result { + let mut arr = array.clone(); + let mut fld = field; + for &id in ancestor_ids { + let child_pos = fld + .children + .iter() + .position(|c| c.id == id) + .ok_or_else(|| { + Error::invalid_input(format!( + "overlay descend: field id {id} not found under '{}'", + fld.name + )) + })?; + let structs = arr.as_any().downcast_ref::().ok_or_else(|| { + Error::invalid_input(format!( + "overlay descend: expected a struct at '{}'", + fld.name + )) + })?; + arr = structs.column(child_pos).clone(); + fld = &fld.children[child_pos]; + } + Ok(arr) +} + +/// Rebuild `array` with the array at `ancestor_ids` replaced by `new_atomic_field`, cloning +/// the struct spine along the path and preserving each struct's null buffer and other +/// children. With an empty path this is just `new_atomic_field` (whole-column replacement). +fn splice_by_ids( + array: &ArrayRef, + field: &Field, + ancestor_ids: &[i32], + new_atomic_field: ArrayRef, +) -> Result { + let Some((&id, rest)) = ancestor_ids.split_first() else { + return Ok(new_atomic_field); + }; + let child_pos = field + .children + .iter() + .position(|c| c.id == id) + .ok_or_else(|| { + Error::invalid_input(format!( + "overlay splice: field id {id} not found under '{}'", + field.name + )) + })?; + let structs = array + .as_any() + .downcast_ref::() + .ok_or_else(|| { + Error::invalid_input(format!( + "overlay splice: expected a struct at '{}'", + field.name + )) + })?; + let len = structs.len(); + let (fields, mut children, nulls) = structs.clone().into_parts(); + children[child_pos] = splice_by_ids( + &children[child_pos], + &field.children[child_pos], + rest, + new_atomic_field, + )?; + Ok(Arc::new(StructArray::try_new_with_length( + fields, children, nulls, len, + )?)) +} + +/// Open the overlay readers a specific read needs and return the per-field plans to +/// merge, pruned to that read. +/// +/// `offsets_in_frag` are the rows the read will return. An overlay whose coverage is +/// disjoint from those rows contributes nothing, so it is dropped and its file is +/// never opened — a `take` that misses an overlay's cells pays no IO for it. Each +/// surviving file is opened once, concurrently, projected to the covered fields; the +/// value bytes are still not read here (the per-batch [`merge_overlay_batch`] fetches +/// only the values it needs). +pub async fn resolve_overlays( + planner: &OverlayReadPlanner, + offsets_in_frag: &[u32], + fragment: &FileFragment, + read_config: &FragReadConfig, +) -> Result> { + let read_offsets = read_offsets_bitmap(offsets_in_frag); + + // A file is opened only if some atomic field it covers has cells among the requested rows. + // This is the row-selection pruning: overlays outside the read are skipped. + let mut file_needed = vec![false; planner.files.len()]; + for atomic_field in &planner.atomic_fields { + for overlay in &atomic_field.overlays_newest_first { + if !overlay.coverage.is_disjoint(&read_offsets) { + file_needed[overlay.file] = true; + } + } + } + + // Open each needed file once, concurrently. The reader is shared (via `Arc`) by + // every atomic field that file covers. + // + // These reads use priority 0 (highest): they are issued only when a ready + // consumer polls the batch task (see `merge_overlay_batch`), so we have already + // committed to reading this batch and the overlay reads cannot clog the + // backpressure queue ahead of work we are not ready for. (A future optimization + // could start the overlay fetches earlier to fill compute bubbles, which would + // want a priority tied to the base read.) + let opened: Vec>> = + futures::future::try_join_all(planner.files.iter().enumerate().map(|(i, file)| { + let needed = file_needed[i]; + async move { + if !needed { + return Ok::<_, Error>(None); + } + Ok(fragment + .open_reader(&file.data_file, Some(&file.open_projection), read_config) + .await? + .map(Arc::from)) + } + })) + .await?; + + let mut plans = Vec::new(); + for atomic_field in &planner.atomic_fields { + let mut overlays_newest_first = Vec::new(); + for overlay in &atomic_field.overlays_newest_first { + let Some(reader) = &opened[overlay.file] else { + continue; // pruned: coverage disjoint from the read + }; + overlays_newest_first.push(LoadedAtomicFieldOverlay { + coverage: overlay.coverage.clone(), + reader: reader.clone(), + }); + } + if !overlays_newest_first.is_empty() { + plans.push(LoadedAtomicField { + top_field: atomic_field.top_field.clone(), + ancestor_ids: atomic_field.ancestor_ids.clone(), + fetch_projection: atomic_field.fetch_projection.clone(), + overlays_newest_first, + }); + } + } + Ok(plans) +} + +/// The set of `offset_in_frag` a read will return, as a bitmap for cheap +/// intersection against overlay coverages. Contiguous scans build a single range; +/// arbitrary `take` offsets (small batches) are inserted individually. +fn read_offsets_bitmap(offsets_in_frag: &[u32]) -> RoaringBitmap { + let mut bitmap = RoaringBitmap::new(); + match contiguous_frag_start(offsets_in_frag) { + Some(start) => { + let end = (start as u64 + offsets_in_frag.len() as u64).min(u32::MAX as u64) as u32; + bitmap.insert_range(start..end); + } + None => bitmap.extend(offsets_in_frag.iter().copied()), + } + bitmap +} + +/// Resolve overlays for one base batch: route each projected atomic field against the batch's +/// `offsets_in_frag`, fetch only the overlay values the batch needs (concurrently with +/// the base read), assemble the merged atomic field, and splice it into its output column. +/// AtomicFields with no covered rows, and columns with no plan, pass through. +pub async fn merge_overlay_batch( + base: ReadBatchFut, + offsets_in_frag: &[u32], + plans: &[LoadedAtomicField], +) -> Result { + let atomic_field_work = futures::future::try_join_all(plans.iter().map(|plan| async move { + let coverages: Vec<&RoaringBitmap> = plan + .overlays_newest_first + .iter() + .map(|overlay| overlay.coverage.as_ref()) + .collect(); + let routing = route_overlays(offsets_in_frag, &coverages); + if !routing.any_overlay { + return Ok::<_, Error>((plan, None)); + } + // Fetch each overlay's values and descend to the atomic field array. The fetch is + // projected to the atomic field's ancestor path, so the fetched column is the pruned + // top-level column; `descend_by_ids` walks it down to the atomic field. + let atomic_field = &plan.fetch_projection.fields[0]; + let fetched = futures::future::try_join_all( + plan.overlays_newest_first + .iter() + .zip(&routing.offsets_in_overlay) + .map(|(overlay, offsets_in_overlay)| async move { + let column = fetch_overlay_values( + overlay.reader.as_ref(), + plan.fetch_projection.clone(), + offsets_in_overlay, + ) + .await?; + descend_by_ids(&column, atomic_field, &plan.ancestor_ids) + }), + ) + .await?; + Ok((plan, Some((routing, fetched)))) + })); + + // The base read and every overlay value read proceed concurrently. + let (batch, resolved) = futures::future::try_join(base, atomic_field_work).await?; + + let schema = batch.schema(); + let mut columns = batch.columns().to_vec(); + for (plan, work) in resolved { + let Some((routing, fetched)) = work else { + continue; + }; + let Some(idx) = schema.index_of(&plan.top_field.name).ok() else { + // The plan's column is not in this batch's projection; skip it. + continue; + }; + let base_atomic_field = descend_by_ids(&columns[idx], &plan.top_field, &plan.ancestor_ids)?; + let merged_atomic_field = assemble_overlay_column(&base_atomic_field, &routing, &fetched)?; + columns[idx] = splice_by_ids( + &columns[idx], + &plan.top_field, + &plan.ancestor_ids, + merged_atomic_field, + )?; + } + Ok(RecordBatch::try_new(schema, columns)?) +} + +/// Fetch one overlay's values at the given `offsets_in_overlay` (sorted, unique): +/// the corresponding entries of its value column, as the top-level column pruned to +/// `projection`. Returns `offsets_in_overlay.len()` rows in the same order; empty +/// input reads nothing and returns an empty column. +async fn fetch_overlay_values( + reader: &dyn GenericFileReader, + projection: Arc, + offsets_in_overlay: &[u32], +) -> Result { + if offsets_in_overlay.is_empty() { + return Ok(arrow_array::new_empty_array( + &projection.fields[0].data_type(), + )); + } + let mut tasks = reader + .take_all_tasks( + offsets_in_overlay, + offsets_in_overlay.len() as u32, + projection, + None, + ) + .await?; + let mut chunks: Vec = Vec::new(); + while let Some(task) = tasks.next().await { + let batch = task.task.await?; + chunks.push(batch.column(0).clone()); + } + let chunk_refs: Vec<&dyn arrow_array::Array> = chunks.iter().map(|a| a.as_ref()).collect(); + Ok(arrow_select::concat::concat(&chunk_refs)?) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow_array::{Int32Array, StringArray, UInt32Array}; + use std::sync::Arc; + + fn i32_array(values: impl IntoIterator>) -> ArrayRef { + Arc::new(Int32Array::from_iter(values)) + } + + fn bitmap(offsets: impl IntoIterator) -> RoaringBitmap { + RoaringBitmap::from_iter(offsets) + } + + /// Physical offsets for a contiguous range `[start, start + len)`. + fn offsets(start: u32, len: usize) -> Vec { + (start..start + len as u32).collect() + } + + /// Drive the production flow purely in memory: route against the coverage + /// bitmaps, then fetch just the requested `offset_in_overlay` entries from each + /// overlay's *full* value column (exactly what the value-pushdown `take` does on + /// disk), then assemble. `overlays_newest_first` holds each overlay's + /// `(coverage, full value column indexed by offset_in_overlay)`. + fn resolve( + base: &ArrayRef, + offsets: &[u32], + overlays_newest_first: &[(RoaringBitmap, ArrayRef)], + ) -> ArrayRef { + let coverages: Vec<&RoaringBitmap> = overlays_newest_first.iter().map(|(c, _)| c).collect(); + let routing = route_overlays(offsets, &coverages); + let fetched: Vec = overlays_newest_first + .iter() + .zip(&routing.offsets_in_overlay) + .map(|((_, full), offsets_in_overlay)| { + let indices = UInt32Array::from(offsets_in_overlay.clone()); + arrow_select::take::take(full.as_ref(), &indices, None).unwrap() + }) + .collect(); + assemble_overlay_column(base, &routing, &fetched).unwrap() + } + + fn assert_i32_eq(actual: &ArrayRef, expected: impl IntoIterator>) { + let actual = actual.as_any().downcast_ref::().unwrap(); + assert_eq!(actual, &Int32Array::from_iter(expected)); + } + + #[test] + fn test_no_overlays_returns_base() { + let base = i32_array([Some(1), Some(2), Some(3)]); + let resolved = resolve(&base, &offsets(0, 3), &[]); + assert_i32_eq(&resolved, [Some(1), Some(2), Some(3)]); + } + + #[test] + fn test_single_overlay_value_offset() { + // Base ages [30, 25, 40, 22]; overlay sets offset_in_frag 1 -> 26, whose + // value sits at offset_in_overlay 0. + let base = i32_array([Some(30), Some(25), Some(40), Some(22)]); + let overlay = (bitmap([1]), i32_array([Some(26)])); + let resolved = resolve(&base, &offsets(0, 4), &[overlay]); + assert_i32_eq(&resolved, [Some(30), Some(26), Some(40), Some(22)]); + } + + #[test] + fn test_value_offsets_multiple_cells() { + // Coverage {0, 2, 3} -> values at offset_in_overlay 0, 1, 2. + let base = i32_array([Some(10), Some(11), Some(12), Some(13)]); + let overlay = ( + bitmap([0, 2, 3]), + i32_array([Some(100), Some(120), Some(130)]), + ); + let resolved = resolve(&base, &offsets(0, 4), &[overlay]); + assert_i32_eq(&resolved, [Some(100), Some(11), Some(120), Some(130)]); + } + + #[test] + fn test_newest_overlay_wins() { + // Two overlays both cover offset_in_frag 1; the newest (first in the slice) + // wins. + let base = i32_array([Some(0), Some(1), Some(2)]); + let newest = (bitmap([1]), i32_array([Some(999)])); + let older = (bitmap([1, 2]), i32_array([Some(111), Some(222)])); + let resolved = resolve(&base, &offsets(0, 3), &[newest, older]); + // offset 1 -> newest (999); offset 2 -> only older covers it (222). + assert_i32_eq(&resolved, [Some(0), Some(999), Some(222)]); + } + + #[test] + fn test_null_override_vs_fall_through() { + // A covered offset with a NULL value overrides the cell to NULL; an + // absent offset falls through to the base. + let base = i32_array([Some(1), Some(2), Some(3)]); + let overlay = (bitmap([0]), i32_array([None])); + let resolved = resolve(&base, &offsets(0, 3), &[overlay]); + assert_i32_eq(&resolved, [None, Some(2), Some(3)]); + } + + #[test] + fn test_physical_start_offset() { + // The batch covers physical rows [10, 13); the overlay covers offset 11. + let base = i32_array([Some(0), Some(0), Some(0)]); + let overlay = (bitmap([11]), i32_array([Some(7)])); + let resolved = resolve(&base, &offsets(10, 3), &[overlay]); + assert_i32_eq(&resolved, [Some(0), Some(7), Some(0)]); + } + + #[test] + fn test_string_column_merge() { + let base: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "c"])); + let overlay = ( + bitmap([0, 2]), + Arc::new(StringArray::from(vec!["A", "C"])) as ArrayRef, + ); + let resolved = resolve(&base, &offsets(0, 3), &[overlay]); + let expected: ArrayRef = Arc::new(StringArray::from(vec!["A", "b", "C"])); + assert_eq!(&resolved, &expected); + } + + #[test] + fn test_non_contiguous_offsets() { + // `take` supplies arbitrary, non-contiguous offsets_in_frag. The base rows + // correspond to offsets 5, 1, 8 (in that order); the overlay covers offsets + // {1, 8}, whose values sit at offset_in_overlay 0, 1. + let base = i32_array([Some(50), Some(10), Some(80)]); + let overlay = (bitmap([1, 8]), i32_array([Some(11), Some(88)])); + let resolved = resolve(&base, &[5, 1, 8], &[overlay]); + // offset 5 uncovered -> base 50; offset 1 -> offset_in_overlay 0 (11); + // offset 8 -> offset_in_overlay 1 (88). + assert_i32_eq(&resolved, [Some(50), Some(11), Some(88)]); + } + + #[test] + fn test_routing_dedups_repeated_offsets() { + // A `take` may request the same offset twice; both rows must route to the + // same overlay value, and that value is fetched only once. + let coverage = bitmap([2, 5]); + let routing = route_overlays(&[5, 2, 5], &[&coverage]); + // offset_in_frag 5 is offset_in_overlay 1, offset_in_frag 2 is + // offset_in_overlay 0: distinct values {0, 1}, sorted. + assert_eq!(routing.offsets_in_overlay, vec![vec![0, 1]]); + let full = i32_array([Some(20), Some(50)]); // values at offset_in_overlay 0, 1 + let fetched = vec![ + arrow_select::take::take( + full.as_ref(), + &UInt32Array::from(routing.offsets_in_overlay[0].clone()), + None, + ) + .unwrap(), + ]; + let base = i32_array([Some(0), Some(0), Some(0)]); + let resolved = assemble_overlay_column(&base, &routing, &fetched).unwrap(); + assert_i32_eq(&resolved, [Some(50), Some(20), Some(50)]); + } + + #[test] + fn test_assemble_value_count_mismatch_errors() { + let coverage = bitmap([0, 1]); + let routing = route_overlays(&[0, 1], &[&coverage]); + let base = i32_array([Some(1), Some(2)]); + // One value supplied for two requested offsets is a caller bug. + let fetched = vec![i32_array([Some(9)])]; + assert!(assemble_overlay_column(&base, &routing, &fetched).is_err()); + } + + #[test] + fn test_contiguous_fast_path_matches_general() { + // The contiguous fast path must produce byte-for-byte identical routing to + // the general offset-major path for any contiguous batch. Fuzz a range of + // fragment starts, lengths, overlay counts, and coverage densities — + // including bits outside the batch range — and compare both paths. + let mut state = 0x9e3779b97f4a7c15u64; + let mut next = || { + state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + (state >> 33) as u32 + }; + for _ in 0..500 { + let frag_start = next() % 64; + let len = (next() % 48 + 1) as usize; + let num_overlays = (next() % 5) as usize; + let coverages: Vec = (0..num_overlays) + .map(|_| { + let density = next() % 101; + let mut b = RoaringBitmap::new(); + for off in frag_start.saturating_sub(3)..frag_start + len as u32 + 3 { + if next() % 100 < density { + b.insert(off); + } + } + b + }) + .collect(); + let refs: Vec<&RoaringBitmap> = coverages.iter().collect(); + let contiguous_offsets: Vec = (frag_start..frag_start + len as u32).collect(); + + let fast = route_contiguous(frag_start, len, &refs); + let general = route_arbitrary(&contiguous_offsets, &refs); + assert_eq!(fast.indices, general.indices, "indices differ"); + assert_eq!( + fast.offsets_in_overlay, general.offsets_in_overlay, + "offsets_in_overlay differ" + ); + assert_eq!(fast.any_overlay, general.any_overlay, "any_overlay differs"); + } + } + + /// `outer { middle { a, b } }` for exercising the descend/splice helpers. + fn nested_struct() -> (Schema, ArrayRef) { + use arrow_schema::{DataType, Field as ArrowField, Fields, Schema as ArrowSchema}; + let mid = Fields::from(vec![ + ArrowField::new("a", DataType::Int32, true), + ArrowField::new("b", DataType::Int32, true), + ]); + let outer_fields = + Fields::from(vec![ArrowField::new("middle", DataType::Struct(mid), true)]); + let arrow_schema = ArrowSchema::new(vec![ArrowField::new( + "outer", + DataType::Struct(outer_fields), + true, + )]); + let schema = Schema::try_from(&arrow_schema).unwrap(); + let middle = StructArray::from(vec![ + ( + Arc::new(ArrowField::new("a", DataType::Int32, true)), + i32_array([Some(1), Some(2), Some(3)]), + ), + ( + Arc::new(ArrowField::new("b", DataType::Int32, true)), + i32_array([Some(10), Some(20), Some(30)]), + ), + ]); + let outer: ArrayRef = Arc::new(StructArray::from(vec![( + Arc::new(ArrowField::new("middle", middle.data_type().clone(), true)), + Arc::new(middle) as ArrayRef, + )])); + (schema, outer) + } + + #[test] + fn test_descend_and_splice_roundtrip() { + let (schema, outer_arr) = nested_struct(); + let outer_field = &schema.fields[0]; + let middle_id = outer_field.children[0].id; + let a_id = outer_field.children[0].children[0].id; + let path = [middle_id, a_id]; + + // Descend to the deep leaf `outer.middle.a`. + let a = descend_by_ids(&outer_arr, outer_field, &path).unwrap(); + assert_i32_eq(&a, [Some(1), Some(2), Some(3)]); + + // Splice a replacement in; only `a` changes, `b` is preserved. + let spliced = splice_by_ids( + &outer_arr, + outer_field, + &path, + i32_array([Some(7), Some(8), Some(9)]), + ) + .unwrap(); + let middle = spliced + .as_any() + .downcast_ref::() + .unwrap() + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .clone(); + assert_i32_eq(&middle.column(0).clone(), [Some(7), Some(8), Some(9)]); + assert_i32_eq(&middle.column(1).clone(), [Some(10), Some(20), Some(30)]); + } + + #[test] + fn test_splice_preserves_struct_nulls() { + use arrow_buffer::NullBuffer; + let (schema, base) = nested_struct(); + let outer_field = &schema.fields[0]; + // Rebuild `outer` with a null at row 1 (a null struct value). + let base = base.as_any().downcast_ref::().unwrap(); + let (fields, children, _) = base.clone().into_parts(); + let outer_arr: ArrayRef = Arc::new( + StructArray::try_new( + fields, + children, + Some(NullBuffer::from(vec![true, false, true])), + ) + .unwrap(), + ); + let path = [ + outer_field.children[0].id, + outer_field.children[0].children[0].id, + ]; + let spliced = splice_by_ids( + &outer_arr, + outer_field, + &path, + i32_array([Some(7), Some(8), Some(9)]), + ) + .unwrap(); + let spliced = spliced.as_any().downcast_ref::().unwrap(); + // The outer struct's null buffer survives the splice. + assert!(!spliced.is_null(0)); + assert!(spliced.is_null(1)); + assert!(!spliced.is_null(2)); + } +}