From a78aefa8532abab4eaf226293145fbc367e9d9f4 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Tue, 30 Jun 2026 10:18:27 -0700 Subject: [PATCH 01/16] feat: resolve data overlay files on the take and scan read path Wire overlay resolution into FragmentReader so `take` and scan reads merge data overlay cells over the base data, removing the temporary guard that refused reads of fragments carrying overlays. Reads are lazy and rank-pushed: each overlay file is opened once (no value bytes read up front), requested offsets are routed to the newest covering overlay at their coverage rank, and only the touched ranks are fetched via the existing `take` primitive. Base and overlay IO run concurrently, then the column is assembled with `interleave`. This makes overlay reads O(requested rows) rather than O(coverage). Co-Authored-By: Claude Opus 4.8 (1M context) --- rust/lance/src/dataset.rs | 1 + rust/lance/src/dataset/fragment.rs | 518 ++++++++++++++++++++++++++- rust/lance/src/dataset/overlay.rs | 544 +++++++++++++++++++++++++++++ 3 files changed, 1049 insertions(+), 14 deletions(-) create mode 100644 rust/lance/src/dataset/overlay.rs 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..39f610a8c6a 100644 --- a/rust/lance/src/dataset/fragment.rs +++ b/rust/lance/src/dataset/fragment.rs @@ -65,6 +65,7 @@ use super::updater::Updater; use super::{NewColumnTransform, WriteParams, schema_evolution}; use crate::dataset::Dataset; use crate::dataset::fragment::session::FragmentSession; +use crate::dataset::overlay::{FieldOverlayPlan, load_overlay_plan, merge_overlay_batch}; use crate::io::deletion::read_dataset_deletion_file; /// Result of [`FileFragment::update_columns_with_offsets`]: updated fragment metadata, modified field ids, @@ -911,16 +912,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(); @@ -932,11 +923,24 @@ impl FileFragment { futures::future::Either::Right(futures::future::ready(Ok(None))) }; - let (opened_files, deletion_vec, row_id_sequence) = - join!(open_files, deletion_vec_load, row_id_load); + // Open overlay value-column readers concurrently with the base files + // (no value bytes are read yet — that happens per batch on read). + let overlay_plan_load = if self.metadata.overlays.is_empty() { + futures::future::Either::Left(futures::future::ready(Ok(Vec::new()))) + } else { + futures::future::Either::Right(load_overlay_plan(self, projection, &read_config)) + }; + + let (opened_files, deletion_vec, row_id_sequence, overlay_plans) = join!( + open_files, + deletion_vec_load, + row_id_load, + overlay_plan_load + ); let opened_files = opened_files?; let deletion_vec = deletion_vec?; let row_id_sequence = row_id_sequence?; + let overlay_plans = overlay_plans?; if opened_files.is_empty() && !read_config.has_system_cols() { return Err(Error::not_found(format!( @@ -959,6 +963,10 @@ impl FileFragment { Arc::new(self.metadata.clone()), )?; + if !overlay_plans.is_empty() { + reader.overlay_plans = Arc::new(overlay_plans); + } + if read_config.with_row_id { reader.with_row_id(); } @@ -996,7 +1004,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 +2272,11 @@ pub struct FragmentReader { // total number of physical rows in the fragment (all rows, ignoring deletions) num_physical_rows: usize, + + /// Overlay value columns for the projected fields, loaded newest-first. + /// Empty when the fragment has no data overlay files. Merged into base + /// batches (by physical offset) before deletion filtering on every read. + overlay_plans: Arc>, } // Custom clone impl needed because it is not easy to clone Box @@ -2292,6 +2305,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_plans: self.overlay_plans.clone(), } } } @@ -2359,6 +2373,7 @@ impl FragmentReader { created_at_sequence: None, num_rows, num_physical_rows, + overlay_plans: Arc::new(Vec::new()), }) } @@ -2616,6 +2631,50 @@ impl FragmentReader { Ok(result.project_by_schema(&output_schema)?) } + /// Merge data overlay values into a stream of base batches. + /// + /// Runs on physical rows in read order, *before* deletion filtering, so each + /// row can be addressed by its physical offset (from `params`) and deletions + /// take precedence naturally (an overlay value for a deleted row is dropped + /// with the row downstream). A no-op when the fragment has no overlays. + /// + /// Each batch's physical offsets are known from `params` and the task's + /// `num_rows` without reading any data, so the overlay value reads (only the + /// coverage ranks the batch actually touches) are issued concurrently with + /// the base read rather than after it. + fn merge_overlays( + &self, + merged: ReadBatchTaskStream, + params: &ReadBatchParams, + total_num_rows: u32, + ) -> ReadBatchTaskStream { + if self.overlay_plans.is_empty() { + return merged; + } + let offsets: Arc> = + Arc::new(params.to_offsets_total(total_num_rows).values().to_vec()); + let plans = self.overlay_plans.clone(); + let mut row_start = 0usize; + merged + .map(move |task| { + let num_rows = task.num_rows; + let start = row_start; + row_start += num_rows as usize; + let offsets = offsets.clone(); + let plans = plans.clone(); + let inner = task.task; + ReadBatchTask { + num_rows, + task: async move { + let slice = &offsets[start..start + num_rows as usize]; + merge_overlay_batch(inner, slice, &plans).await + } + .boxed(), + } + }) + .boxed() + } + async fn new_read_impl<'a, F>( &'a self, params: ReadBatchParams, @@ -2683,6 +2742,8 @@ impl FragmentReader { lance_table::utils::stream::merge_streams(read_streams) }; + let merged = self.merge_overlays(merged, ¶ms, total_num_rows); + // Add the row id column (if needed) and delete rows (if a deletion // vector is present). let config = RowIdAndDeletesConfig { @@ -2853,6 +2914,9 @@ 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); + // Add the row id column (if needed) and delete rows (if a deletion // vector is present). let config = RowIdAndDeletesConfig { @@ -2865,7 +2929,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 +3178,432 @@ 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}; + use arrow_schema::{DataType, Field as ArrowField, 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); + } + + #[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 value-column ranks + /// 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 ranks. + /// This is the regression guard for the lazy, rank-pushdown overlay read. + #[tokio::test] + async fn test_take_reads_only_needed_overlay_ranks() { + 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]); + } + } + #[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..98bc3488107 --- /dev/null +++ b/rust/lance/src/dataset/overlay.rs @@ -0,0 +1,544 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Resolution of data overlay files on read. +//! +//! An overlay supplies new values for a subset of `(physical offset, field)` +//! cells. To resolve a field's values for a set of physical row offsets, the +//! overlays that cover that field are walked **newest to oldest**: the first +//! overlay that covers an offset wins, and its value is taken at the offset's +//! **rank** (the 0-based count of set bits below it) in the field's coverage +//! bitmap. An offset that no overlay covers falls through to the base value. +//! +//! The offsets are supplied explicitly (one per base row), so a single code path +//! serves both the scan (a contiguous physical range) and `take` (arbitrary +//! physical offsets) read paths. +//! +//! Deletions take precedence over overlays, but that is handled downstream: the +//! merge runs on physical rows *before* the deletion filter, so an overlay value +//! for a deleted offset is computed and then dropped with the row — making it +//! inert, exactly as the specification requires, with no special handling here. + +use std::collections::{BTreeSet, HashMap}; +use std::sync::Arc; + +use arrow_array::{Array, ArrayRef, RecordBatch}; +use arrow_select::interleave::interleave; +use futures::StreamExt; +use lance_core::datatypes::Schema; +use lance_core::{Error, Result}; +use roaring::RoaringBitmap; + +use lance_table::format::overlay::DataOverlayFile; +use lance_table::utils::stream::ReadBatchFut; + +use crate::dataset::fragment::{FileFragment, FragReadConfig, GenericFileReader}; + +/// Order a fragment's overlays from newest to oldest for read resolution. +/// +/// Precedence is by `committed_version` (higher is newer); ties are broken by +/// position in the fragment's `overlays` list, where a later entry is newer. +/// Returns indices into `overlays`. +pub fn overlay_indices_newest_first(overlays: &[DataOverlayFile]) -> Vec { + let mut indices: Vec = (0..overlays.len()).collect(); + indices.sort_by(|&a, &b| { + overlays[b] + .committed_version + .cmp(&overlays[a].committed_version) + .then(b.cmp(&a)) + }); + indices +} + +/// How a batch of physical row offsets routes onto a field's overlays. +/// +/// Produced by [`route_overlays`] from the coverage bitmaps alone — before any +/// value column is read — so the caller can fetch only the ranks it actually +/// needs (see [`OverlayRouting::needed_ranks`]) instead of the whole column, and +/// then assemble the merged column with [`assemble_overlay_column`]. +pub struct OverlayRouting { + /// `interleave` source/position pairs, one per output row. Source `0` is the + /// base column (position = the row's index); source `k + 1` is overlay `k`'s + /// fetched values (position = the row's index within `needed_ranks[k]`). + indices: Vec<(usize, usize)>, + /// `needed_ranks[k]` is the sorted, deduplicated set of coverage ranks that + /// overlay `k` must supply for this batch — the indices to fetch from its + /// value column. + needed_ranks: Vec>, + /// Whether any row routes to an overlay at all (false ⇒ pure fall-through). + any_overlay: bool, +} + +impl OverlayRouting { + /// The ranks each overlay (newest-first) must fetch from its value column. + pub fn needed_ranks(&self) -> &[Vec] { + &self.needed_ranks + } + + /// True when no row is covered by any overlay, so the base column is the + /// answer unchanged and no value-column reads are needed. + pub fn all_fall_through(&self) -> bool { + !self.any_overlay + } +} + +/// Decide, for each physical offset in `offsets`, which source supplies its +/// value: the newest overlay whose coverage contains it (taken at the offset's +/// 0-based rank in that coverage), or the base column if none covers it. +/// +/// Reads only the coverage bitmaps (newest-first), so it can run before the +/// value columns are fetched and tells the caller exactly which ranks to fetch. +pub fn route_overlays( + offsets: &[u32], + coverages_newest_first: &[&RoaringBitmap], +) -> OverlayRouting { + let mut rank_sets: Vec> = vec![BTreeSet::new(); coverages_newest_first.len()]; + let mut raw: Vec> = Vec::with_capacity(offsets.len()); + for &offset in offsets { + let mut routed = None; + for (k, coverage) in coverages_newest_first.iter().enumerate() { + if coverage.contains(offset) { + // 0-based rank: number of set bits strictly below `offset`. + let rank = coverage.rank(offset) as u32 - 1; + rank_sets[k].insert(rank); + routed = Some((k, rank)); + break; + } + } + raw.push(routed); + } + + let needed_ranks: Vec> = rank_sets + .iter() + .map(|ranks| ranks.iter().copied().collect()) + .collect(); + let rank_positions: Vec> = needed_ranks + .iter() + .map(|ranks| ranks.iter().enumerate().map(|(pos, &r)| (r, pos)).collect()) + .collect(); + + let mut any_overlay = false; + let indices = raw + .into_iter() + .enumerate() + .map(|(i, routed)| match routed { + None => (0, i), + Some((k, rank)) => { + any_overlay = true; + (k + 1, rank_positions[k][&rank]) + } + }) + .collect(); + + OverlayRouting { + indices, + needed_ranks, + any_overlay, + } +} + +/// Assemble the merged column from `base` and the per-overlay values fetched for +/// the ranks [`route_overlays`] asked for. +/// +/// `fetched_newest_first[k]` holds overlay `k`'s values for `routing`'s +/// `needed_ranks[k]`, in that order. The result has the same length and data +/// type as `base`; a covered offset whose overlay value is NULL resolves **to** +/// NULL (distinct from a fall-through, which keeps its base value). +pub fn assemble_overlay_column( + base: &ArrayRef, + routing: &OverlayRouting, + fetched_newest_first: &[ArrayRef], +) -> Result { + if routing.all_fall_through() { + return Ok(base.clone()); + } + if fetched_newest_first.len() != routing.needed_ranks.len() { + return Err(Error::invalid_input(format!( + "overlay assembly got {} value columns but routing expects {}", + fetched_newest_first.len(), + routing.needed_ranks.len() + ))); + } + for (k, values) in fetched_newest_first.iter().enumerate() { + if values.len() != routing.needed_ranks[k].len() { + return Err(Error::invalid_input(format!( + "overlay value column {} has {} values but {} ranks were requested", + k, + values.len(), + routing.needed_ranks[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 field: which physical offsets it +/// covers, and an opened (but unread) reader over the overlay file from which the +/// field's value column is fetched by coverage rank at merge time. +#[derive(Debug, Clone)] +struct LoadedFieldOverlay { + /// Physical offsets this overlay covers for the field. + coverage: Arc, + /// Reader over the overlay data file, projected to this field; shared across + /// the fields that the same file covers. + reader: Arc, + /// Single-field projection used when fetching the value column. + field_projection: Arc, +} + +/// The overlays that apply to a single projected field, ordered newest-first. +/// `field_name` is the top-level read-batch column name the plan applies to. +#[derive(Debug, Clone)] +pub struct FieldOverlayPlan { + field_name: String, + overlays_newest_first: Vec, +} + +/// Open the overlay value-column readers for `fragment`'s projected fields, +/// ordered newest-first, ready to be merged into base batches on read. +/// +/// Each contributing overlay *file* is opened once (its metadata loaded), +/// projected to the fields it covers that the read also projects. The value +/// columns themselves are **not** read here — the per-batch merge fetches only +/// the coverage ranks it needs (see [`merge_overlay_batch`]), so a `take` of a +/// few rows no longer reads a whole overlay column. +/// +/// For each projected (top-level) field, the fragment's overlays are walked +/// newest-first; an overlay contributes if its `data_file.fields` includes the +/// field. Overlays on nested (non-top-level) fields are not yet supported and are +/// simply not matched here. +pub async fn load_overlay_plan( + fragment: &FileFragment, + projection: &Schema, + read_config: &FragReadConfig, +) -> Result> { + let order = overlay_indices_newest_first(&fragment.metadata.overlays); + + // Open each contributing overlay file once, concurrently. The reader is + // shared (via `Arc`) by every field that file covers. + // + // TODO(overlay perf): these reads use the default reader priority. Once + // we benchmark take/scan over overlays, decide whether overlay value + // reads should inherit `read_config.reader_priority` (or get a dedicated + // priority) so they schedule alongside the base reads. + let opened: Vec>> = + futures::future::try_join_all(order.iter().map(|&overlay_idx| async move { + let overlay = &fragment.metadata.overlays[overlay_idx]; + let covered: Vec = projection + .fields + .iter() + .filter(|f| overlay.data_file.fields.contains(&f.id)) + .cloned() + .collect(); + if covered.is_empty() { + return Ok::<_, Error>(None); + } + let schema = Schema { + fields: covered, + metadata: Default::default(), + }; + Ok(fragment + .open_reader(&overlay.data_file, Some(&schema), read_config) + .await? + .map(Arc::from)) + })) + .await?; + + let mut plans = Vec::new(); + for field in &projection.fields { + let mut overlays_newest_first = Vec::new(); + for (slot, &overlay_idx) in order.iter().enumerate() { + let overlay = &fragment.metadata.overlays[overlay_idx]; + let Some(field_pos) = overlay + .data_file + .fields + .iter() + .position(|&id| id == field.id) + else { + continue; + }; + let Some(reader) = &opened[slot] else { + continue; + }; + overlays_newest_first.push(LoadedFieldOverlay { + coverage: overlay.coverage_for_field(field_pos)?, + reader: reader.clone(), + field_projection: Arc::new(Schema { + fields: vec![field.clone()], + metadata: Default::default(), + }), + }); + } + if !overlays_newest_first.is_empty() { + plans.push(FieldOverlayPlan { + field_name: field.name.clone(), + overlays_newest_first, + }); + } + } + Ok(plans) +} + +/// Resolve overlays for one base batch: route each projected field against the +/// batch's physical `offsets`, fetch only the coverage ranks the batch touches +/// (concurrently with the base read), and assemble the merged columns. Fields +/// with no plan, and the row-id/row-address system columns, pass through. +pub async fn merge_overlay_batch( + base: ReadBatchFut, + offsets: &[u32], + plans: &[FieldOverlayPlan], +) -> Result { + let 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, &coverages); + if routing.all_fall_through() { + return Ok::<_, Error>((plan.field_name.as_str(), None)); + } + let fetched = futures::future::try_join_all( + plan.overlays_newest_first + .iter() + .zip(routing.needed_ranks()) + .map(|(overlay, ranks)| { + fetch_overlay_ranks( + overlay.reader.as_ref(), + overlay.field_projection.clone(), + ranks, + ) + }), + ) + .await?; + Ok((plan.field_name.as_str(), Some((routing, fetched)))) + })); + + // The base read and every overlay value read proceed concurrently. + let (batch, resolved) = futures::future::try_join(base, field_work).await?; + + let schema = batch.schema(); + let mut columns = batch.columns().to_vec(); + for (field_name, work) in resolved { + let Some((routing, fetched)) = work else { + continue; + }; + let Some(idx) = schema.index_of(field_name).ok() else { + // The plan's field is not in this batch's projection; skip it. + continue; + }; + columns[idx] = assemble_overlay_column(&columns[idx], &routing, &fetched)?; + } + Ok(RecordBatch::try_new(schema, columns)?) +} + +/// Fetch one overlay's value column at the given coverage `ranks` (sorted and +/// unique). Returns a column of `ranks.len()` values aligned with `ranks`; an +/// empty `ranks` reads nothing and yields an empty column. +async fn fetch_overlay_ranks( + reader: &dyn GenericFileReader, + projection: Arc, + ranks: &[u32], +) -> Result { + if ranks.is_empty() { + return Ok(arrow_array::new_empty_array( + &projection.fields[0].data_type(), + )); + } + let mut tasks = reader + .take_all_tasks(ranks, ranks.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 ranks from each overlay's *full* + /// value column (exactly what the rank-pushdown `take` does on disk), then + /// assemble. `overlays_newest_first` holds each overlay's `(coverage, full + /// value column indexed by rank)`. + 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.needed_ranks()) + .map(|((_, full), ranks)| { + let indices = UInt32Array::from(ranks.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_rank_addressing() { + // Base ages [30, 25, 40, 22]; overlay sets offset 1 -> 26 (rank 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_rank_addressing_multiple_offsets() { + // Coverage {0, 2, 3} -> values at ranks 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 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 physical offsets. The base + // rows correspond to offsets 5, 1, 8 (in that order); the overlay covers + // offsets {1, 8} with values at ranks 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 -> rank 0 (11); offset 8 -> rank 1 (88). + assert_i32_eq(&resolved, [Some(50), Some(11), Some(88)]); + } + + #[test] + fn test_routing_dedups_repeated_ranks() { + // A `take` may request the same offset twice; both rows must route to the + // same rank, and that rank is fetched only once. + let coverage = bitmap([2, 5]); + let routing = route_overlays(&[5, 2, 5], &[&coverage]); + // Offset 5 is rank 1, offset 2 is rank 0: distinct ranks {0, 1}, sorted. + assert_eq!(routing.needed_ranks(), &[vec![0, 1]]); + let full = i32_array([Some(20), Some(50)]); // values at ranks 0, 1 + let fetched = vec![ + arrow_select::take::take( + full.as_ref(), + &UInt32Array::from(routing.needed_ranks()[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 ranks is a caller bug. + let fetched = vec![i32_array([Some(9)])]; + assert!(assemble_overlay_column(&base, &routing, &fetched).is_err()); + } + + #[test] + fn test_overlay_ordering_newest_first() { + use lance_table::format::DataFile; + use lance_table::format::overlay::OverlayCoverage; + let mk = |version: u64| DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("o.lance", vec![1], None), + coverage: OverlayCoverage::dense(RoaringBitmap::new()), + committed_version: version, + }; + // List order [v2, v5, v3]; newest-first should be v5(idx1), v3(idx2), v2(idx0). + let overlays = vec![mk(2), mk(5), mk(3)]; + assert_eq!(overlay_indices_newest_first(&overlays), vec![1, 2, 0]); + + // Equal versions: later list position is newer. + let overlays = vec![mk(4), mk(4)]; + assert_eq!(overlay_indices_newest_first(&overlays), vec![1, 0]); + } +} From 2e83678edd31fdb0f87833864b350191b1973088 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Tue, 30 Jun 2026 12:30:33 -0700 Subject: [PATCH 02/16] test: cover data overlay read-path edge cases Adds tests for combinations the initial read-path suite missed: row-id pass-through alongside an overlay, an older overlay routed zero ranks (empty-ranks fetch branch), a newest NULL value shadowing an older non-null, per-field newest-wins across two sparse overlays, a plan present but all requested offsets uncovered (all-fall-through), and a dataset-level take spanning multiple overlaid fragments. Co-Authored-By: Claude Opus 4.8 (1M context) --- rust/lance/src/dataset/fragment.rs | 227 ++++++++++++++++++++++++++++- 1 file changed, 226 insertions(+), 1 deletion(-) diff --git a/rust/lance/src/dataset/fragment.rs b/rust/lance/src/dataset/fragment.rs index 39f610a8c6a..dca36c4424f 100644 --- a/rust/lance/src/dataset/fragment.rs +++ b/rust/lance/src/dataset/fragment.rs @@ -3184,7 +3184,9 @@ mod tests { mod overlay_read { use std::sync::Arc; - use arrow_array::{Array, ArrayRef, Int32Array, RecordBatch, RecordBatchIterator}; + use arrow_array::{ + Array, ArrayRef, Int32Array, RecordBatch, RecordBatchIterator, UInt64Array, + }; use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema}; use lance_core::datatypes::Schema; use lance_file::version::LanceFileVersion; @@ -3602,6 +3604,229 @@ mod tests { assert_eq!(val.value(0), values[0]); assert_eq!(val.value(1), values[1]); } + + /// 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 is routed zero ranks and its value column must not be + /// read (the empty-ranks branch of `fetch_overlay_ranks`). The result still + /// resolves to the newest overlay. + #[rstest] + #[tokio::test] + async fn test_take_older_overlay_contributes_no_ranks( + #[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 + // ranks 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.all_fall_through()` 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]); + } } #[rstest] From 38e2b31077cb0ab7e671f4a1cea5b1b4c8a81509 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Tue, 30 Jun 2026 16:11:44 -0700 Subject: [PATCH 03/16] perf: route overlays bitmap-major on contiguous scans MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `route_overlays` was offset-major: for every physical offset it probed each coverage bitmap newest-first. On a full scan that is `O(N_rows * K_overlays)` roaring `contains` probes — ~16M for a 1M-row, 16-overlay scan — of which 99% only confirm a fall-through to base. It dominated scan CPU (~87% of the scan thread) and made scan cost grow linearly with overlay count, badly so for scattered (stride) coverage. A scan reads a contiguous physical range, so offset `o` is output row `o - base`. The fast path now intersects each coverage with the batch's offset range (a container-level op that drops a non-overlapping batch in `O(containers)`) and routes only the in-range bits directly to rows. Those bits carry consecutive coverage ranks starting at the count of bits below `base`, so ranks need a single `rank` lookup rather than a probe per cell. `take` still supplies arbitrary offsets and keeps the general path. Microbenchmark of `route_overlays` over a 1M-row scan (ms/scan): overlays stride old -> new contiguous old -> new 1 13.3 -> 0.67 2.8 -> 0.55 4 46.3 -> 0.86 3.5 -> 0.66 16 178.5 -> 1.48 9.4 -> 0.96 64 829.0 -> 4.07 33.3 -> 2.09 A fuzz test asserts the fast path produces byte-identical routing to the general path across random bases, lengths, overlay counts, and densities. Co-Authored-By: Claude Opus 4.8 (1M context) --- rust/lance/src/dataset/overlay.rs | 130 ++++++++++++++++++++++++++++++ 1 file changed, 130 insertions(+) diff --git a/rust/lance/src/dataset/overlay.rs b/rust/lance/src/dataset/overlay.rs index 98bc3488107..a911121cf40 100644 --- a/rust/lance/src/dataset/overlay.rs +++ b/rust/lance/src/dataset/overlay.rs @@ -88,10 +88,97 @@ impl OverlayRouting { /// /// Reads only the coverage bitmaps (newest-first), so it can run before the /// value columns are fetched and tells the caller exactly which ranks to fetch. +/// +/// A scan reads a contiguous physical range, so when `offsets` is contiguous +/// ascending we take a bitmap-major fast path that visits only each coverage's +/// in-range bits — `O(covered + K)` — instead of probing every offset against +/// every coverage. `take` supplies arbitrary offsets and uses the general path. pub fn route_overlays( offsets: &[u32], coverages_newest_first: &[&RoaringBitmap], ) -> OverlayRouting { + match contiguous_base(offsets) { + Some(base) => route_contiguous(base, offsets.len(), coverages_newest_first), + None => route_arbitrary(offsets, coverages_newest_first), + } +} + +/// The starting offset if `offsets` is a contiguous ascending run +/// `[base, base + 1, ...]`, else `None` (including when empty). +fn contiguous_base(offsets: &[u32]) -> Option { + let base = *offsets.first()?; + offsets + .iter() + .enumerate() + .all(|(i, &offset)| offset as u64 == base as u64 + i as u64) + .then_some(base) +} + +/// Fast path for a contiguous batch: offset `o` is output row `o - base`, so a +/// coverage's bits route to rows directly without per-offset probing. +/// +/// For each coverage we intersect with the batch's offset range, which is a +/// container-level operation that drops a non-overlapping batch (e.g. a scan +/// batch past a contiguous coverage's bits) in `O(containers)` without touching +/// individual cells. The in-range bits then carry **consecutive** coverage ranks +/// starting at the count of bits below `base` (one `rank` lookup) — no bits lie +/// between them by construction — so ranks need no running count. Coverages are +/// processed newest-first with a "first claim wins" guard for precedence. +fn route_contiguous( + base: u32, + len: usize, + coverages_newest_first: &[&RoaringBitmap], +) -> OverlayRouting { + let mut needed_ranks: Vec> = vec![Vec::new(); coverages_newest_first.len()]; + let mut routed: Vec> = vec![None; len]; + let range_end = (base as u64 + len as u64).min(u32::MAX as u64) as u32; + let mut batch_range = RoaringBitmap::new(); + batch_range.insert_range(base..range_end); + + for (k, coverage) in coverages_newest_first.iter().enumerate() { + let in_range = *coverage & &batch_range; + if in_range.is_empty() { + continue; + } + // 0-based rank of the first in-range cell: the coverage bits below `base`. + let base_rank = if base == 0 { + 0 + } else { + coverage.rank(base - 1) as u32 + }; + for (i, offset) in in_range.iter().enumerate() { + let row = (offset - base) as usize; + if routed[row].is_none() { + routed[row] = Some((k, needed_ranks[k].len())); + needed_ranks[k].push(base_rank + i as u32); + } + } + } + + let mut any_overlay = false; + let indices = routed + .into_iter() + .enumerate() + .map(|(i, routed)| match routed { + None => (0, i), + Some((k, pos)) => { + any_overlay = true; + (k + 1, pos) + } + }) + .collect(); + + OverlayRouting { + indices, + needed_ranks, + any_overlay, + } +} + +/// General path for arbitrary (e.g. `take`) offsets: probe each offset against +/// the coverages newest-first. `take` batches are small, so the `O(N * K)` +/// probing here is not a bottleneck. +fn route_arbitrary(offsets: &[u32], coverages_newest_first: &[&RoaringBitmap]) -> OverlayRouting { let mut rank_sets: Vec> = vec![BTreeSet::new(); coverages_newest_first.len()]; let mut raw: Vec> = Vec::with_capacity(offsets.len()); for &offset in offsets { @@ -524,6 +611,49 @@ mod tests { 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 bases, lengths, overlay counts, and coverage densities — including + // bits outside the batch range — and compare both fields. + let mut state = 0x9e3779b97f4a7c15u64; + let mut next = || { + state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + (state >> 33) as u32 + }; + for _ in 0..500 { + let base = 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 base.saturating_sub(3)..base + len as u32 + 3 { + if next() % 100 < density { + b.insert(off); + } + } + b + }) + .collect(); + let refs: Vec<&RoaringBitmap> = coverages.iter().collect(); + let contiguous_offsets: Vec = (base..base + len as u32).collect(); + + let fast = route_contiguous(base, len, &refs); + let general = route_arbitrary(&contiguous_offsets, &refs); + assert_eq!(fast.indices, general.indices, "indices differ"); + assert_eq!( + fast.needed_ranks, general.needed_ranks, + "needed_ranks differ" + ); + assert_eq!(fast.any_overlay, general.any_overlay, "any_overlay differs"); + } + } + #[test] fn test_overlay_ordering_newest_first() { use lance_table::format::DataFile; From ec4e23f6555ec2ec34f584d16716866188b2133a Mon Sep 17 00:00:00 2001 From: Will Jones Date: Wed, 8 Jul 2026 10:05:50 -0700 Subject: [PATCH 04/16] refactor(overlay): clearer read-path terminology Review feedback: the "rank" language in the overlay read path is hard to follow. Name the three coordinate spaces explicitly and use them everywhere: - offset_in_frag: a row's physical position in the fragment - offset_in_batch: a row's position within the batch being assembled - offset_in_overlay: a value's position in an overlay's dense value column (what a roaring bitmap calls "rank") Renames `OverlayRouting::needed_ranks` -> `offsets_in_overlay`, `fetch_overlay_ranks` -> `fetch_overlay_values`, and the contiguous-path `base` -> `frag_start` (it was colliding with "base column"). Rewrites the module and function docs in plainer language. No behavior change. Co-Authored-By: Claude Opus 4.8 (1M context) --- rust/lance/src/dataset/fragment.rs | 55 ++-- rust/lance/src/dataset/overlay.rs | 395 ++++++++++++++++------------- 2 files changed, 254 insertions(+), 196 deletions(-) diff --git a/rust/lance/src/dataset/fragment.rs b/rust/lance/src/dataset/fragment.rs index dca36c4424f..2a1e8942b3a 100644 --- a/rust/lance/src/dataset/fragment.rs +++ b/rust/lance/src/dataset/fragment.rs @@ -2634,14 +2634,15 @@ impl FragmentReader { /// Merge data overlay values into a stream of base batches. /// /// Runs on physical rows in read order, *before* deletion filtering, so each - /// row can be addressed by its physical offset (from `params`) and deletions - /// take precedence naturally (an overlay value for a deleted row is dropped - /// with the row downstream). A no-op when the fragment has no overlays. + /// 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. /// - /// Each batch's physical offsets are known from `params` and the task's - /// `num_rows` without reading any data, so the overlay value reads (only the - /// coverage ranks the batch actually touches) are issued concurrently with - /// the base read rather than after it. + /// Each batch's `offset_in_frag` values are known from `params` and the task's + /// `num_rows` without reading any data, so the overlay reads (only the values + /// the batch actually needs) are issued concurrently with the base read rather + /// than after it. fn merge_overlays( &self, merged: ReadBatchTaskStream, @@ -2651,23 +2652,29 @@ impl FragmentReader { if self.overlay_plans.is_empty() { return merged; } - let offsets: Arc> = + // 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 each batch below slice out its own offsets without reading + // any data. Only paid when the fragment actually has overlays. + let offsets_in_frag: Arc> = Arc::new(params.to_offsets_total(total_num_rows).values().to_vec()); let plans = self.overlay_plans.clone(); - let mut row_start = 0usize; + // 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; merged .map(move |task| { let num_rows = task.num_rows; - let start = row_start; - row_start += num_rows as usize; - let offsets = offsets.clone(); + 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 slice = &offsets[start..start + num_rows as usize]; - merge_overlay_batch(inner, slice, &plans).await + let batch_offsets = &offsets_in_frag[start..start + num_rows as usize]; + merge_overlay_batch(inner, batch_offsets, &plans).await } .boxed(), } @@ -3525,13 +3532,13 @@ mod tests { assert_eq!(col(&batch, "val").values(), &expected); } - /// A `take` of a few rows must read only the overlay value-column ranks - /// 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 ranks. - /// This is the regression guard for the lazy, rank-pushdown overlay read. + /// 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_ranks() { + async fn test_take_reads_only_needed_overlay_values() { let version = LanceFileVersion::V2_1; const N: usize = 100_000; @@ -3647,18 +3654,18 @@ mod tests { } /// When the newest overlay covers every requested offset, an older overlay - /// in the same plan is routed zero ranks and its value column must not be - /// read (the empty-ranks branch of `fetch_overlay_ranks`). The result still + /// 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_ranks( + 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 - // ranks to fetch even though it is part of the field's plan. + // values to fetch even though it is part of the field's plan. let dataset = commit_overlay( dataset, "older", diff --git a/rust/lance/src/dataset/overlay.rs b/rust/lance/src/dataset/overlay.rs index a911121cf40..d8071e1db50 100644 --- a/rust/lance/src/dataset/overlay.rs +++ b/rust/lance/src/dataset/overlay.rs @@ -3,21 +3,36 @@ //! Resolution of data overlay files on read. //! -//! An overlay supplies new values for a subset of `(physical offset, field)` -//! cells. To resolve a field's values for a set of physical row offsets, the -//! overlays that cover that field are walked **newest to oldest**: the first -//! overlay that covers an offset wins, and its value is taken at the offset's -//! **rank** (the 0-based count of set bits below it) in the field's coverage -//! bitmap. An offset that no overlay covers falls through to the base value. +//! 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. //! -//! The offsets are supplied explicitly (one per base row), so a single code path -//! serves both the scan (a contiguous physical range) and `take` (arbitrary -//! physical offsets) read paths. +//! Three coordinate spaces show up throughout this module; keeping them straight +//! is most of the work: //! -//! Deletions take precedence over overlays, but that is handled downstream: the -//! merge runs on physical rows *before* the deletion filter, so an overlay value -//! for a deleted offset is computed and then dropped with the row — making it -//! inert, exactly as the specification requires, with no special handling here. +//! - `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; @@ -50,107 +65,120 @@ pub fn overlay_indices_newest_first(overlays: &[DataOverlayFile]) -> Vec indices } -/// How a batch of physical row offsets routes onto a field's overlays. +/// 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. /// -/// Produced by [`route_overlays`] from the coverage bitmaps alone — before any -/// value column is read — so the caller can fetch only the ranks it actually -/// needs (see [`OverlayRouting::needed_ranks`]) instead of the whole column, and -/// then assemble the merged column with [`assemble_overlay_column`]. +/// 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 (see [`OverlayRouting::offsets_in_overlay`]) rather than whole +/// columns, then build the merged column with [`assemble_overlay_column`]. pub struct OverlayRouting { - /// `interleave` source/position pairs, one per output row. Source `0` is the - /// base column (position = the row's index); source `k + 1` is overlay `k`'s - /// fetched values (position = the row's index within `needed_ranks[k]`). + /// 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)>, - /// `needed_ranks[k]` is the sorted, deduplicated set of coverage ranks that - /// overlay `k` must supply for this batch — the indices to fetch from its - /// value column. - needed_ranks: Vec>, - /// Whether any row routes to an overlay at all (false ⇒ pure fall-through). + /// 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). any_overlay: bool, } impl OverlayRouting { - /// The ranks each overlay (newest-first) must fetch from its value column. - pub fn needed_ranks(&self) -> &[Vec] { - &self.needed_ranks + /// Per overlay (newest-first), the `offset_in_overlay` values to fetch from its + /// value column. + pub fn offsets_in_overlay(&self) -> &[Vec] { + &self.offsets_in_overlay } - /// True when no row is covered by any overlay, so the base column is the - /// answer unchanged and no value-column reads are needed. + /// True when no row is covered by any overlay, so the base column is already the + /// answer and no overlay values need to be read. pub fn all_fall_through(&self) -> bool { !self.any_overlay } } -/// Decide, for each physical offset in `offsets`, which source supplies its -/// value: the newest overlay whose coverage contains it (taken at the offset's -/// 0-based rank in that coverage), or the base column if none covers it. +/// 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`. /// -/// Reads only the coverage bitmaps (newest-first), so it can run before the -/// value columns are fetched and tells the caller exactly which ranks to fetch. +/// 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 reads a contiguous physical range, so when `offsets` is contiguous -/// ascending we take a bitmap-major fast path that visits only each coverage's -/// in-range bits — `O(covered + K)` — instead of probing every offset against -/// every coverage. `take` supplies arbitrary offsets and uses the general path. +/// 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. pub fn route_overlays( - offsets: &[u32], + offsets_in_frag: &[u32], coverages_newest_first: &[&RoaringBitmap], ) -> OverlayRouting { - match contiguous_base(offsets) { - Some(base) => route_contiguous(base, offsets.len(), coverages_newest_first), - None => route_arbitrary(offsets, coverages_newest_first), + 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), } } -/// The starting offset if `offsets` is a contiguous ascending run -/// `[base, base + 1, ...]`, else `None` (including when empty). -fn contiguous_base(offsets: &[u32]) -> Option { - let base = *offsets.first()?; - offsets +/// 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 == base as u64 + i as u64) - .then_some(base) + .all(|(i, &offset)| offset as u64 == start as u64 + i as u64) + .then_some(start) } -/// Fast path for a contiguous batch: offset `o` is output row `o - base`, so a -/// coverage's bits route to rows directly without per-offset probing. +/// 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. /// -/// For each coverage we intersect with the batch's offset range, which is a -/// container-level operation that drops a non-overlapping batch (e.g. a scan -/// batch past a contiguous coverage's bits) in `O(containers)` without touching -/// individual cells. The in-range bits then carry **consecutive** coverage ranks -/// starting at the count of bits below `base` (one `rank` lookup) — no bits lie -/// between them by construction — so ranks need no running count. Coverages are -/// processed newest-first with a "first claim wins" guard for precedence. +/// 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( - base: u32, + frag_start: u32, len: usize, coverages_newest_first: &[&RoaringBitmap], ) -> OverlayRouting { - let mut needed_ranks: Vec> = vec![Vec::new(); coverages_newest_first.len()]; + 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 = (base as u64 + len as u64).min(u32::MAX as u64) as u32; + 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(base..range_end); + batch_range.insert_range(frag_start..range_end); for (k, coverage) in coverages_newest_first.iter().enumerate() { - let in_range = *coverage & &batch_range; - if in_range.is_empty() { + let covered_in_batch = *coverage & &batch_range; + if covered_in_batch.is_empty() { continue; } - // 0-based rank of the first in-range cell: the coverage bits below `base`. - let base_rank = if base == 0 { + // 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(base - 1) as u32 + coverage.rank(frag_start - 1) as u32 }; - for (i, offset) in in_range.iter().enumerate() { - let row = (offset - base) as usize; - if routed[row].is_none() { - routed[row] = Some((k, needed_ranks[k].len())); - needed_ranks[k].push(base_rank + i 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); } } } @@ -159,78 +187,91 @@ fn route_contiguous( let indices = routed .into_iter() .enumerate() - .map(|(i, routed)| match routed { - None => (0, i), - Some((k, pos)) => { + .map(|(offset_in_batch, routed)| match routed { + None => (0, offset_in_batch), + Some((k, fetch_pos)) => { any_overlay = true; - (k + 1, pos) + (k + 1, fetch_pos) } }) .collect(); OverlayRouting { indices, - needed_ranks, + offsets_in_overlay, any_overlay, } } -/// General path for arbitrary (e.g. `take`) offsets: probe each offset against -/// the coverages newest-first. `take` batches are small, so the `O(N * K)` -/// probing here is not a bottleneck. -fn route_arbitrary(offsets: &[u32], coverages_newest_first: &[&RoaringBitmap]) -> OverlayRouting { - let mut rank_sets: Vec> = vec![BTreeSet::new(); coverages_newest_first.len()]; - let mut raw: Vec> = Vec::with_capacity(offsets.len()); - for &offset in offsets { +/// 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) { - // 0-based rank: number of set bits strictly below `offset`. - let rank = coverage.rank(offset) as u32 - 1; - rank_sets[k].insert(rank); - routed = Some((k, rank)); + 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; } } - raw.push(routed); + routed_per_row.push(routed); } - let needed_ranks: Vec> = rank_sets + let offsets_in_overlay: Vec> = offset_sets .iter() - .map(|ranks| ranks.iter().copied().collect()) + .map(|offsets| offsets.iter().copied().collect()) .collect(); - let rank_positions: Vec> = needed_ranks + // 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(|ranks| ranks.iter().enumerate().map(|(pos, &r)| (r, pos)).collect()) + .map(|offsets| { + offsets + .iter() + .enumerate() + .map(|(pos, &o)| (o, pos)) + .collect() + }) .collect(); let mut any_overlay = false; - let indices = raw + let indices = routed_per_row .into_iter() .enumerate() - .map(|(i, routed)| match routed { - None => (0, i), - Some((k, rank)) => { + .map(|(offset_in_batch, routed)| match routed { + None => (0, offset_in_batch), + Some((k, offset_in_overlay)) => { any_overlay = true; - (k + 1, rank_positions[k][&rank]) + (k + 1, fetch_positions[k][&offset_in_overlay]) } }) .collect(); OverlayRouting { indices, - needed_ranks, + offsets_in_overlay, any_overlay, } } -/// Assemble the merged column from `base` and the per-overlay values fetched for -/// the ranks [`route_overlays`] asked for. +/// 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 -/// `needed_ranks[k]`, in that order. The result has the same length and data -/// type as `base`; a covered offset whose overlay value is NULL resolves **to** -/// NULL (distinct from a fall-through, which keeps its base value). +/// `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). pub fn assemble_overlay_column( base: &ArrayRef, routing: &OverlayRouting, @@ -239,20 +280,20 @@ pub fn assemble_overlay_column( if routing.all_fall_through() { return Ok(base.clone()); } - if fetched_newest_first.len() != routing.needed_ranks.len() { + 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.needed_ranks.len() + routing.offsets_in_overlay.len() ))); } for (k, values) in fetched_newest_first.iter().enumerate() { - if values.len() != routing.needed_ranks[k].len() { + if values.len() != routing.offsets_in_overlay[k].len() { return Err(Error::invalid_input(format!( - "overlay value column {} has {} values but {} ranks were requested", + "overlay value column {} has {} values but {} were requested", k, values.len(), - routing.needed_ranks[k].len() + routing.offsets_in_overlay[k].len() ))); } } @@ -265,12 +306,12 @@ pub fn assemble_overlay_column( interleave(&sources, &routing.indices).map_err(Error::from) } -/// One overlay's contribution to one projected field: which physical offsets it -/// covers, and an opened (but unread) reader over the overlay file from which the -/// field's value column is fetched by coverage rank at merge time. +/// One overlay's contribution to one projected field: the cells it covers, and an +/// opened (but unread) reader over the overlay file from which the field's values +/// are fetched by `offset_in_overlay` at merge time. #[derive(Debug, Clone)] struct LoadedFieldOverlay { - /// Physical offsets this overlay covers for the field. + /// The `offset_in_frag` cells this overlay covers for the field. coverage: Arc, /// Reader over the overlay data file, projected to this field; shared across /// the fields that the same file covers. @@ -287,14 +328,14 @@ pub struct FieldOverlayPlan { overlays_newest_first: Vec, } -/// Open the overlay value-column readers for `fragment`'s projected fields, -/// ordered newest-first, ready to be merged into base batches on read. +/// Open the overlay value-column readers for `fragment`'s projected fields, ordered +/// newest-first, ready to be merged into base batches on read. /// -/// Each contributing overlay *file* is opened once (its metadata loaded), -/// projected to the fields it covers that the read also projects. The value -/// columns themselves are **not** read here — the per-batch merge fetches only -/// the coverage ranks it needs (see [`merge_overlay_batch`]), so a `take` of a -/// few rows no longer reads a whole overlay column. +/// Each contributing overlay *file* is opened once (its metadata loaded), projected +/// to the fields it covers that the read also projects. The value columns +/// themselves are **not** read here — the per-batch merge fetches only the overlay +/// values it needs (see [`merge_overlay_batch`]), so a `take` of a few rows no +/// longer reads a whole overlay column. /// /// For each projected (top-level) field, the fragment's overlays are walked /// newest-first; an overlay contributes if its `data_file.fields` includes the @@ -373,12 +414,12 @@ pub async fn load_overlay_plan( } /// Resolve overlays for one base batch: route each projected field against the -/// batch's physical `offsets`, fetch only the coverage ranks the batch touches -/// (concurrently with the base read), and assemble the merged columns. Fields -/// with no plan, and the row-id/row-address system columns, pass through. +/// batch's `offsets_in_frag`, fetch only the overlay values the batch needs +/// (concurrently with the base read), and assemble the merged columns. Fields with +/// no plan, and the row-id/row-address system columns, pass through. pub async fn merge_overlay_batch( base: ReadBatchFut, - offsets: &[u32], + offsets_in_frag: &[u32], plans: &[FieldOverlayPlan], ) -> Result { let field_work = futures::future::try_join_all(plans.iter().map(|plan| async move { @@ -387,19 +428,19 @@ pub async fn merge_overlay_batch( .iter() .map(|overlay| overlay.coverage.as_ref()) .collect(); - let routing = route_overlays(offsets, &coverages); + let routing = route_overlays(offsets_in_frag, &coverages); if routing.all_fall_through() { return Ok::<_, Error>((plan.field_name.as_str(), None)); } let fetched = futures::future::try_join_all( plan.overlays_newest_first .iter() - .zip(routing.needed_ranks()) - .map(|(overlay, ranks)| { - fetch_overlay_ranks( + .zip(routing.offsets_in_overlay()) + .map(|(overlay, offsets_in_overlay)| { + fetch_overlay_values( overlay.reader.as_ref(), overlay.field_projection.clone(), - ranks, + offsets_in_overlay, ) }), ) @@ -425,21 +466,27 @@ pub async fn merge_overlay_batch( Ok(RecordBatch::try_new(schema, columns)?) } -/// Fetch one overlay's value column at the given coverage `ranks` (sorted and -/// unique). Returns a column of `ranks.len()` values aligned with `ranks`; an -/// empty `ranks` reads nothing and yields an empty column. -async fn fetch_overlay_ranks( +/// Fetch one overlay's values at the given `offsets_in_overlay` (sorted, unique): +/// the corresponding entries of its value column. Returns a column of +/// `offsets_in_overlay.len()` values in the same order; empty input reads nothing +/// and returns an empty column. +async fn fetch_overlay_values( reader: &dyn GenericFileReader, projection: Arc, - ranks: &[u32], + offsets_in_overlay: &[u32], ) -> Result { - if ranks.is_empty() { + 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(ranks, ranks.len() as u32, projection, None) + .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 { @@ -470,10 +517,10 @@ mod tests { } /// Drive the production flow purely in memory: route against the coverage - /// bitmaps, then fetch just the requested ranks from each overlay's *full* - /// value column (exactly what the rank-pushdown `take` does on disk), then - /// assemble. `overlays_newest_first` holds each overlay's `(coverage, full - /// value column indexed by rank)`. + /// 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], @@ -483,9 +530,9 @@ mod tests { let routing = route_overlays(offsets, &coverages); let fetched: Vec = overlays_newest_first .iter() - .zip(routing.needed_ranks()) - .map(|((_, full), ranks)| { - let indices = UInt32Array::from(ranks.clone()); + .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(); @@ -505,8 +552,9 @@ mod tests { } #[test] - fn test_single_overlay_rank_addressing() { - // Base ages [30, 25, 40, 22]; overlay sets offset 1 -> 26 (rank 0). + 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]); @@ -514,8 +562,8 @@ mod tests { } #[test] - fn test_rank_addressing_multiple_offsets() { - // Coverage {0, 2, 3} -> values at ranks 0,1,2. + 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]), @@ -527,7 +575,8 @@ mod tests { #[test] fn test_newest_overlay_wins() { - // Two overlays both cover offset 1; the newest (first in the slice) 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)])); @@ -569,29 +618,31 @@ mod tests { #[test] fn test_non_contiguous_offsets() { - // `take` supplies arbitrary, non-contiguous physical offsets. The base - // rows correspond to offsets 5, 1, 8 (in that order); the overlay covers - // offsets {1, 8} with values at ranks 0, 1. + // `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 -> rank 0 (11); offset 8 -> rank 1 (88). + // 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_ranks() { + fn test_routing_dedups_repeated_offsets() { // A `take` may request the same offset twice; both rows must route to the - // same rank, and that rank is fetched only once. + // same overlay value, and that value is fetched only once. let coverage = bitmap([2, 5]); let routing = route_overlays(&[5, 2, 5], &[&coverage]); - // Offset 5 is rank 1, offset 2 is rank 0: distinct ranks {0, 1}, sorted. - assert_eq!(routing.needed_ranks(), &[vec![0, 1]]); - let full = i32_array([Some(20), Some(50)]); // values at ranks 0, 1 + // 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![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.needed_ranks()[0].clone()), + &UInt32Array::from(routing.offsets_in_overlay()[0].clone()), None, ) .unwrap(), @@ -606,17 +657,17 @@ mod tests { 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 ranks is a caller bug. + // 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 bases, lengths, overlay counts, and coverage densities — including - // bits outside the batch range — and compare both fields. + // 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 @@ -625,14 +676,14 @@ mod tests { (state >> 33) as u32 }; for _ in 0..500 { - let base = next() % 64; + 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 base.saturating_sub(3)..base + len as u32 + 3 { + for off in frag_start.saturating_sub(3)..frag_start + len as u32 + 3 { if next() % 100 < density { b.insert(off); } @@ -641,14 +692,14 @@ mod tests { }) .collect(); let refs: Vec<&RoaringBitmap> = coverages.iter().collect(); - let contiguous_offsets: Vec = (base..base + len as u32).collect(); + let contiguous_offsets: Vec = (frag_start..frag_start + len as u32).collect(); - let fast = route_contiguous(base, len, &refs); + 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.needed_ranks, general.needed_ranks, - "needed_ranks differ" + fast.offsets_in_overlay, general.offsets_in_overlay, + "offsets_in_overlay differ" ); assert_eq!(fast.any_overlay, general.any_overlay, "any_overlay differs"); } From eed20327a00e9f219801944301dcc26fa2343dcb Mon Sep 17 00:00:00 2001 From: Will Jones Date: Wed, 8 Jul 2026 10:39:36 -0700 Subject: [PATCH 05/16] perf(overlay): open overlay readers lazily, pruned to the read's rows MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review feedback: prune overlay readers by row selection, not just projection. Overlay readers were opened eagerly in `FileFragment::open`, before the read's rows were known, so a `take` that misses an overlay's cells still paid to open its file. Split the work: - `plan_overlays` (at open): build the coverage plan from metadata only — no files opened, no IO. - `resolve_overlays` (at read, in `merge_overlays`): the read's offsets are known here, so open only the overlay files whose coverage intersects the requested rows; overlays disjoint from the read are dropped and never opened. Projection pruning (skip overlays covering no projected field) is unchanged. Trade-off: a reused `FragmentReader` reopens overlay readers per read rather than once at open. Acceptable — overlays are resolved per read operation, and the plan/state is shared via `Arc` so cloning a reader stays cheap. Test: `test_take_prunes_overlays_outside_row_selection` deletes an overlay's data file and shows a take missing its coverage still succeeds (never opened), while a take hitting it fails. Co-Authored-By: Claude Opus 4.8 (1M context) --- rust/lance/src/dataset/fragment.rs | 156 +++++++++++++------ rust/lance/src/dataset/overlay.rs | 232 ++++++++++++++++++++++------- 2 files changed, 292 insertions(+), 96 deletions(-) diff --git a/rust/lance/src/dataset/fragment.rs b/rust/lance/src/dataset/fragment.rs index 2a1e8942b3a..6c4af347bb2 100644 --- a/rust/lance/src/dataset/fragment.rs +++ b/rust/lance/src/dataset/fragment.rs @@ -65,7 +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::{FieldOverlayPlan, load_overlay_plan, merge_overlay_batch}; +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, @@ -650,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, @@ -923,24 +925,11 @@ impl FileFragment { futures::future::Either::Right(futures::future::ready(Ok(None))) }; - // Open overlay value-column readers concurrently with the base files - // (no value bytes are read yet — that happens per batch on read). - let overlay_plan_load = if self.metadata.overlays.is_empty() { - futures::future::Either::Left(futures::future::ready(Ok(Vec::new()))) - } else { - futures::future::Either::Right(load_overlay_plan(self, projection, &read_config)) - }; - - let (opened_files, deletion_vec, row_id_sequence, overlay_plans) = join!( - open_files, - deletion_vec_load, - row_id_load, - overlay_plan_load - ); + let (opened_files, deletion_vec, row_id_sequence) = + join!(open_files, deletion_vec_load, row_id_load); let opened_files = opened_files?; let deletion_vec = deletion_vec?; let row_id_sequence = row_id_sequence?; - let overlay_plans = overlay_plans?; if opened_files.is_empty() && !read_config.has_system_cols() { return Err(Error::not_found(format!( @@ -963,8 +952,17 @@ impl FileFragment { Arc::new(self.metadata.clone()), )?; - if !overlay_plans.is_empty() { - reader.overlay_plans = Arc::new(overlay_plans); + // 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 { @@ -2273,10 +2271,22 @@ pub struct FragmentReader { // total number of physical rows in the fragment (all rows, ignoring deletions) num_physical_rows: usize, - /// Overlay value columns for the projected fields, loaded newest-first. - /// Empty when the fragment has no data overlay files. Merged into base - /// batches (by physical offset) before deletion filtering on every read. - overlay_plans: Arc>, + /// 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 @@ -2305,7 +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_plans: self.overlay_plans.clone(), + overlay: self.overlay.clone(), } } } @@ -2373,7 +2383,7 @@ impl FragmentReader { created_at_sequence: None, num_rows, num_physical_rows, - overlay_plans: Arc::new(Vec::new()), + overlay: None, }) } @@ -2639,30 +2649,44 @@ impl FragmentReader { /// value for a deleted row is dropped along with the row downstream. A no-op /// when the fragment has no overlays. /// - /// Each batch's `offset_in_frag` values are known from `params` and the task's - /// `num_rows` without reading any data, so the overlay reads (only the values - /// the batch actually needs) are issued concurrently with the base read rather - /// than after it. - fn merge_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, - ) -> ReadBatchTaskStream { - if self.overlay_plans.is_empty() { - return merged; - } + ) -> 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 each batch below slice out its own offsets without reading - // any data. Only paid when the fragment actually has overlays. + // 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. let offsets_in_frag: Arc> = Arc::new(params.to_offsets_total(total_num_rows).values().to_vec()); - let plans = self.overlay_plans.clone(); + + // 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; - merged + let stream = merged .map(move |task| { let num_rows = task.num_rows; let start = rows_seen; @@ -2679,7 +2703,8 @@ impl FragmentReader { .boxed(), } }) - .boxed() + .boxed(); + Ok(stream) } async fn new_read_impl<'a, F>( @@ -2749,7 +2774,7 @@ impl FragmentReader { lance_table::utils::stream::merge_streams(read_streams) }; - let merged = self.merge_overlays(merged, ¶ms, total_num_rows); + 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). @@ -2922,7 +2947,9 @@ impl FragmentReader { }; let params = ReadBatchParams::Ranges(ranges); - let merged_stream = self.merge_overlays(merged_stream, ¶ms, total_num_rows); + 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). @@ -3612,6 +3639,51 @@ mod tests { 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. + assert!( + frag.take(&[5], &val_only).await.is_err(), + "take hitting the overlay's coverage should require its missing file", + ); + } + /// 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. diff --git a/rust/lance/src/dataset/overlay.rs b/rust/lance/src/dataset/overlay.rs index d8071e1db50..0dbf02572e0 100644 --- a/rust/lance/src/dataset/overlay.rs +++ b/rust/lance/src/dataset/overlay.rs @@ -44,6 +44,7 @@ use lance_core::datatypes::Schema; use lance_core::{Error, Result}; use roaring::RoaringBitmap; +use lance_table::format::DataFile; use lance_table::format::overlay::DataOverlayFile; use lance_table::utils::stream::ReadBatchFut; @@ -306,82 +307,117 @@ pub fn assemble_overlay_column( interleave(&sources, &routing.indices).map_err(Error::from) } -/// One overlay's contribution to one projected field: the cells it covers, and an -/// opened (but unread) reader over the overlay file from which the field's values -/// are fetched by `offset_in_overlay` at merge time. +/// One overlay's contribution to one projected field, with its file reader opened: +/// the cells it covers, and the reader from which the field's values are fetched by +/// `offset_in_overlay` at merge time. #[derive(Debug, Clone)] struct LoadedFieldOverlay { /// The `offset_in_frag` cells this overlay covers for the field. coverage: Arc, - /// Reader over the overlay data file, projected to this field; shared across - /// the fields that the same file covers. + /// Reader over the overlay data file, projected to the covered fields; shared + /// across the fields that the same file covers. reader: Arc, /// Single-field projection used when fetching the value column. field_projection: Arc, } -/// The overlays that apply to a single projected field, ordered newest-first. -/// `field_name` is the top-level read-batch column name the plan applies to. +/// The overlays that apply to a single projected field, ordered newest-first, with +/// readers opened and pruned to a specific read. `field_name` is the top-level +/// read-batch column name the plan applies to. Produced by [`resolve_overlays`] and +/// consumed by [`merge_overlay_batch`]. #[derive(Debug, Clone)] pub struct FieldOverlayPlan { field_name: String, overlays_newest_first: Vec, } -/// Open the overlay value-column readers for `fragment`'s projected fields, ordered -/// newest-first, ready to be merged into base batches on read. -/// -/// Each contributing overlay *file* is opened once (its metadata loaded), projected -/// to the fields it covers that the read also projects. The value columns -/// themselves are **not** read here — the per-batch merge fetches only the overlay -/// values it needs (see [`merge_overlay_batch`]), so a `take` of a few rows no -/// longer reads a whole overlay column. +/// 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 fields to project when the file is opened, so a + /// single reader serves every field the file contributes to. + open_projection: Arc, +} + +/// One overlay's contribution to one projected field, before the file is opened. +#[derive(Debug, Clone)] +struct PlannedFieldOverlay { + /// Index into [`OverlayReadPlanner::files`] of the file that supplies the value. + file: usize, + coverage: Arc, +} + +/// The overlays that apply to a single projected field, ordered newest-first, +/// before any file is opened. +#[derive(Debug, Clone)] +struct PlannedField { + field_name: String, + /// Single-field projection used when fetching this field's value column. + field_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 +/// [`FieldOverlayPlan`]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, + fields: Vec, +} + +impl OverlayReadPlanner { + /// True when no projected field has any overlay, so there is nothing to resolve. + pub fn is_empty(&self) -> bool { + self.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. /// /// For each projected (top-level) field, the fragment's overlays are walked /// newest-first; an overlay contributes if its `data_file.fields` includes the /// field. Overlays on nested (non-top-level) fields are not yet supported and are -/// simply not matched here. -pub async fn load_overlay_plan( - fragment: &FileFragment, - projection: &Schema, - read_config: &FragReadConfig, -) -> Result> { +/// simply not matched here. Each contributing overlay *file* appears once in +/// `files`, shared by every field it covers. +pub fn plan_overlays(fragment: &FileFragment, projection: &Schema) -> Result { let order = overlay_indices_newest_first(&fragment.metadata.overlays); - // Open each contributing overlay file once, concurrently. The reader is - // shared (via `Arc`) by every field that file covers. - // - // TODO(overlay perf): these reads use the default reader priority. Once - // we benchmark take/scan over overlays, decide whether overlay value - // reads should inherit `read_config.reader_priority` (or get a dedicated - // priority) so they schedule alongside the base reads. - let opened: Vec>> = - futures::future::try_join_all(order.iter().map(|&overlay_idx| async move { - let overlay = &fragment.metadata.overlays[overlay_idx]; - let covered: Vec = projection - .fields - .iter() - .filter(|f| overlay.data_file.fields.contains(&f.id)) - .cloned() - .collect(); - if covered.is_empty() { - return Ok::<_, Error>(None); - } - let schema = Schema { + // One entry per contributing overlay file, newest-first. `pos_to_file[pos]` maps + // a position in `order` to its index in `files` (None if it covers no projected + // field, so it is never referenced and never opened). + let mut files = Vec::new(); + let mut pos_to_file = vec![None; order.len()]; + for (pos, &overlay_idx) in order.iter().enumerate() { + let overlay = &fragment.metadata.overlays[overlay_idx]; + let covered: Vec = projection + .fields + .iter() + .filter(|f| overlay.data_file.fields.contains(&f.id)) + .cloned() + .collect(); + if covered.is_empty() { + continue; + } + pos_to_file[pos] = Some(files.len()); + files.push(PlannedOverlayFile { + data_file: overlay.data_file.clone(), + open_projection: Arc::new(Schema { fields: covered, metadata: Default::default(), - }; - Ok(fragment - .open_reader(&overlay.data_file, Some(&schema), read_config) - .await? - .map(Arc::from)) - })) - .await?; + }), + }); + } - let mut plans = Vec::new(); + let mut fields = Vec::new(); for field in &projection.fields { let mut overlays_newest_first = Vec::new(); - for (slot, &overlay_idx) in order.iter().enumerate() { + for (pos, &overlay_idx) in order.iter().enumerate() { let overlay = &fragment.metadata.overlays[overlay_idx]; let Some(field_pos) = overlay .data_file @@ -391,21 +427,94 @@ pub async fn load_overlay_plan( else { continue; }; - let Some(reader) = &opened[slot] else { + let Some(file) = pos_to_file[pos] else { continue; }; - overlays_newest_first.push(LoadedFieldOverlay { + overlays_newest_first.push(PlannedFieldOverlay { + file, coverage: overlay.coverage_for_field(field_pos)?, - reader: reader.clone(), + }); + } + if !overlays_newest_first.is_empty() { + fields.push(PlannedField { + field_name: field.name.clone(), field_projection: Arc::new(Schema { fields: vec![field.clone()], metadata: Default::default(), }), + overlays_newest_first, + }); + } + } + Ok(OverlayReadPlanner { files, fields }) +} + +/// 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 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 field in &planner.fields { + for overlay in &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 field that file covers. + // + // TODO(overlay perf): these reads use the default reader priority. Once we + // benchmark take/scan over overlays, decide whether overlay value reads should + // inherit `read_config.reader_priority` (or get a dedicated priority) so they + // schedule alongside the base reads. + 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 field in &planner.fields { + let mut overlays_newest_first = Vec::new(); + for overlay in &field.overlays_newest_first { + let Some(reader) = &opened[overlay.file] else { + continue; // pruned: coverage disjoint from the read + }; + overlays_newest_first.push(LoadedFieldOverlay { + coverage: overlay.coverage.clone(), + reader: reader.clone(), + field_projection: field.field_projection.clone(), }); } if !overlays_newest_first.is_empty() { plans.push(FieldOverlayPlan { - field_name: field.name.clone(), + field_name: field.field_name.clone(), overlays_newest_first, }); } @@ -413,6 +522,21 @@ pub async fn load_overlay_plan( 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 field against the /// batch's `offsets_in_frag`, fetch only the overlay values the batch needs /// (concurrently with the base read), and assemble the merged columns. Fields with From 9442fdfd14670c3aa0ba6c7debf3ef8a1cda3297 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Wed, 8 Jul 2026 11:22:01 -0700 Subject: [PATCH 06/16] refactor(overlay): index overlays by field id when planning MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `plan_overlays` walked every projected field against every overlay, which is `O(num_fields x num_overlays)` — costly when a fragment has thousands of fields. Rebuild it as a single newest-first pass over the overlays that indexes each covered field into a `field_id -> overlays` map, making planning `O(total overlay fields + projected fields)`. Overlays are already sorted newest-last on load (`sort_overlays_newest_last`), so the pass just iterates the slice in reverse for newest-first precedence — dropping the redundant `overlay_indices_newest_first` sort helper (its ordering rule stays covered by `test_overlays_sorted_newest_last_on_load`). A `debug_assert` documents the load-time invariant the reverse walk relies on. Co-Authored-By: Claude Opus 4.8 (1M context) --- rust/lance/src/dataset/overlay.rs | 129 ++++++++++++------------------ 1 file changed, 52 insertions(+), 77 deletions(-) diff --git a/rust/lance/src/dataset/overlay.rs b/rust/lance/src/dataset/overlay.rs index 0dbf02572e0..bc6699b5db6 100644 --- a/rust/lance/src/dataset/overlay.rs +++ b/rust/lance/src/dataset/overlay.rs @@ -45,27 +45,10 @@ use lance_core::{Error, Result}; use roaring::RoaringBitmap; use lance_table::format::DataFile; -use lance_table::format::overlay::DataOverlayFile; use lance_table::utils::stream::ReadBatchFut; use crate::dataset::fragment::{FileFragment, FragReadConfig, GenericFileReader}; -/// Order a fragment's overlays from newest to oldest for read resolution. -/// -/// Precedence is by `committed_version` (higher is newer); ties are broken by -/// position in the fragment's `overlays` list, where a later entry is newer. -/// Returns indices into `overlays`. -pub fn overlay_indices_newest_first(overlays: &[DataOverlayFile]) -> Vec { - let mut indices: Vec = (0..overlays.len()).collect(); - indices.sort_by(|&a, &b| { - overlays[b] - .committed_version - .cmp(&overlays[a].committed_version) - .then(b.cmp(&a)) - }); - indices -} - /// 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. @@ -380,62 +363,72 @@ impl OverlayReadPlanner { /// 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. /// -/// For each projected (top-level) field, the fragment's overlays are walked -/// newest-first; an overlay contributes if its `data_file.fields` includes the -/// field. Overlays on nested (non-top-level) fields are not yet supported and are -/// simply not matched here. Each contributing overlay *file* appears once in -/// `files`, shared by every field it covers. +/// Overlays are stored oldest-first (sorted newest-last on load, see +/// `sort_overlays_newest_last`), so walking them in reverse gives newest-first +/// precedence. A single pass — visiting only each overlay's own fields and matching +/// them against the projection through a field-id map — builds the per-field overlay +/// lists, so planning is `O(total overlay fields + projected fields)` rather than +/// `O(projected fields × overlays)` (a fragment can have thousands of fields). +/// +/// An overlay contributes to a projected field when its `data_file.fields` includes +/// that field's id. Overlays on nested (non-top-level) fields are not yet supported +/// and simply match no projected field. Each contributing overlay *file* appears once +/// in `files`, shared by every field it covers. pub fn plan_overlays(fragment: &FileFragment, projection: &Schema) -> Result { - let order = overlay_indices_newest_first(&fragment.metadata.overlays); - - // One entry per contributing overlay file, newest-first. `pos_to_file[pos]` maps - // a position in `order` to its index in `files` (None if it covers no projected - // field, so it is never referenced and never opened). + 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)" + ); + + // Projected fields indexed by id, so each overlay field is matched in O(1). + let projected: HashMap = + projection.fields.iter().map(|f| (f.id, f)).collect(); + + // Walk overlays newest-first, visiting only each overlay's own fields. Every + // covered field is appended to its entry in `field_overlays`; because we walk + // newest-first, each entry ends up ordered newest-first for free. let mut files = Vec::new(); - let mut pos_to_file = vec![None; order.len()]; - for (pos, &overlay_idx) in order.iter().enumerate() { - let overlay = &fragment.metadata.overlays[overlay_idx]; - let covered: Vec = projection - .fields - .iter() - .filter(|f| overlay.data_file.fields.contains(&f.id)) - .cloned() - .collect(); - if covered.is_empty() { + let mut field_overlays: HashMap> = HashMap::new(); + for overlay in overlays.iter().rev() { + // (field_id, field_pos) for each of this overlay's fields that we project. + let mut contributions = Vec::new(); + let mut covered_fields = Vec::new(); + for (field_pos, &field_id) in overlay.data_file.fields.iter().enumerate() { + if let Some(field) = projected.get(&field_id) { + contributions.push((field_id, field_pos)); + covered_fields.push((*field).clone()); + } + } + if contributions.is_empty() { continue; } - pos_to_file[pos] = Some(files.len()); + let file = files.len(); files.push(PlannedOverlayFile { data_file: overlay.data_file.clone(), open_projection: Arc::new(Schema { - fields: covered, + fields: covered_fields, metadata: Default::default(), }), }); + for (field_id, field_pos) in contributions { + field_overlays + .entry(field_id) + .or_default() + .push(PlannedFieldOverlay { + file, + coverage: overlay.coverage_for_field(field_pos)?, + }); + } } + // Emit one PlannedField per projected field that has overlays, in projection + // order for deterministic results. let mut fields = Vec::new(); for field in &projection.fields { - let mut overlays_newest_first = Vec::new(); - for (pos, &overlay_idx) in order.iter().enumerate() { - let overlay = &fragment.metadata.overlays[overlay_idx]; - let Some(field_pos) = overlay - .data_file - .fields - .iter() - .position(|&id| id == field.id) - else { - continue; - }; - let Some(file) = pos_to_file[pos] else { - continue; - }; - overlays_newest_first.push(PlannedFieldOverlay { - file, - coverage: overlay.coverage_for_field(field_pos)?, - }); - } - if !overlays_newest_first.is_empty() { + if let Some(overlays_newest_first) = field_overlays.remove(&field.id) { fields.push(PlannedField { field_name: field.name.clone(), field_projection: Arc::new(Schema { @@ -828,22 +821,4 @@ mod tests { assert_eq!(fast.any_overlay, general.any_overlay, "any_overlay differs"); } } - - #[test] - fn test_overlay_ordering_newest_first() { - use lance_table::format::DataFile; - use lance_table::format::overlay::OverlayCoverage; - let mk = |version: u64| DataOverlayFile { - data_file: DataFile::new_legacy_from_fields("o.lance", vec![1], None), - coverage: OverlayCoverage::dense(RoaringBitmap::new()), - committed_version: version, - }; - // List order [v2, v5, v3]; newest-first should be v5(idx1), v3(idx2), v2(idx0). - let overlays = vec![mk(2), mk(5), mk(3)]; - assert_eq!(overlay_indices_newest_first(&overlays), vec![1, 2, 0]); - - // Equal versions: later list position is newer. - let overlays = vec![mk(4), mk(4)]; - assert_eq!(overlay_indices_newest_first(&overlays), vec![1, 0]); - } } From b2a54c1f88ab390eb12fc048593a712b0227ccfb Mon Sep 17 00:00:00 2001 From: Will Jones Date: Wed, 8 Jul 2026 11:45:47 -0700 Subject: [PATCH 07/16] refactor(overlay): keep routing helpers private to the module MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `OverlayRouting` (with its methods), `route_overlays`, and `assemble_overlay_column` were `pub` but are only used within the `overlay` module — its sibling helpers (`route_contiguous`, `route_arbitrary`, `fetch_overlay_values`, ...) are already private. Narrow them to match, so the module's real entry points (`plan_overlays`, `resolve_overlays`, `merge_overlay_batch`, `OverlayReadPlanner`, `FieldOverlayPlan`) are clear at a glance. Co-Authored-By: Claude Opus 4.8 (1M context) --- rust/lance/src/dataset/overlay.rs | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/rust/lance/src/dataset/overlay.rs b/rust/lance/src/dataset/overlay.rs index bc6699b5db6..88bf0bdc838 100644 --- a/rust/lance/src/dataset/overlay.rs +++ b/rust/lance/src/dataset/overlay.rs @@ -57,7 +57,7 @@ use crate::dataset::fragment::{FileFragment, FragReadConfig, GenericFileReader}; /// column is read — so the caller can fetch only the overlay values it will /// actually use (see [`OverlayRouting::offsets_in_overlay`]) rather than whole /// columns, then build the merged column with [`assemble_overlay_column`]. -pub struct OverlayRouting { +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 @@ -75,13 +75,13 @@ pub struct OverlayRouting { impl OverlayRouting { /// Per overlay (newest-first), the `offset_in_overlay` values to fetch from its /// value column. - pub fn offsets_in_overlay(&self) -> &[Vec] { + fn offsets_in_overlay(&self) -> &[Vec] { &self.offsets_in_overlay } /// True when no row is covered by any overlay, so the base column is already the /// answer and no overlay values need to be read. - pub fn all_fall_through(&self) -> bool { + fn all_fall_through(&self) -> bool { !self.any_overlay } } @@ -96,7 +96,7 @@ impl OverlayRouting { /// 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. -pub fn route_overlays( +fn route_overlays( offsets_in_frag: &[u32], coverages_newest_first: &[&RoaringBitmap], ) -> OverlayRouting { @@ -256,7 +256,7 @@ fn route_arbitrary( /// `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). -pub fn assemble_overlay_column( +fn assemble_overlay_column( base: &ArrayRef, routing: &OverlayRouting, fetched_newest_first: &[ArrayRef], From c8d11461e11ef61077738ba69f893682f52cf490 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Wed, 8 Jul 2026 11:45:50 -0700 Subject: [PATCH 08/16] test(overlay): cover multi-batch scan offset slicing `merge_overlays` slices each batch out of the read's `offsets_in_frag` with a running `rows_seen` accumulator, but every existing scan test used single-batch 6-row fragments, so the cross-batch (`start > 0`) path was never exercised. Add a scan over a 10-row fragment at batch_size 4 with overlays in the 1st, 2nd, and 3rd batches; the test asserts the read actually spans multiple batches before checking alignment. Co-Authored-By: Claude Opus 4.8 (1M context) --- rust/lance/src/dataset/fragment.rs | 81 ++++++++++++++++++++++++++++++ 1 file changed, 81 insertions(+) diff --git a/rust/lance/src/dataset/fragment.rs b/rust/lance/src/dataset/fragment.rs index 6c4af347bb2..ef7f1683843 100644 --- a/rust/lance/src/dataset/fragment.rs +++ b/rust/lance/src/dataset/fragment.rs @@ -3906,6 +3906,87 @@ mod tests { 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); + } } #[rstest] From a2ca3e73c9c264dd9cae17bcf290d2b0c9eafe08 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Wed, 8 Jul 2026 11:55:40 -0700 Subject: [PATCH 09/16] test(overlay): cover empty take, string column, and projection pruning MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add three read-path cases that were untested: an empty `take` (plan present, no offsets to route), an end-to-end overlay on a variable-width `Utf8` column (exercises the string value-pushdown path through the real reader, plus a NULL override), and projection pruning doing zero IO — proven by deleting the overlay's data file and confirming a read that projects only the unrelated column still succeeds while projecting the overlaid column fails. Co-Authored-By: Claude Opus 4.8 (1M context) --- rust/lance/src/dataset/fragment.rs | 136 +++++++++++++++++++++++++++++ 1 file changed, 136 insertions(+) diff --git a/rust/lance/src/dataset/fragment.rs b/rust/lance/src/dataset/fragment.rs index ef7f1683843..4e139a62679 100644 --- a/rust/lance/src/dataset/fragment.rs +++ b/rust/lance/src/dataset/fragment.rs @@ -3987,6 +3987,142 @@ mod tests { .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. + assert!( + frag.take(&[0], &val_only).await.is_err(), + "projecting the overlaid column should require its missing file", + ); + } } #[rstest] From 1014d97f82d5490a68821af837b47f64db8fdb33 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Wed, 8 Jul 2026 12:06:38 -0700 Subject: [PATCH 10/16] fix(overlay): resolve overlays on struct and list columns MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit An overlay is written against the field ids the file stores, which for a struct or list column under the V2_1 structural encoding are its *leaf* child ids — the parent's top-level id is never recorded. `plan_overlays` matched only top-level ids, so a struct/list overlay was silently skipped on read (V2_1) and returned stale base values, while V2_0 — which does record the parent id — resolved it. The mismatch was untested. Map every id in the projection subtree (top-level fields and their nested descendants) to its top-level projected field, so a leaf id resolves to the column it belongs to; the whole column is then fetched and replaced as a unit. Add end-to-end tests for struct and list overlays on both formats. Co-Authored-By: Claude Opus 4.8 (1M context) --- rust/lance/src/dataset/fragment.rs | 158 +++++++++++++++++++++++++++++ rust/lance/src/dataset/overlay.rs | 64 ++++++++---- 2 files changed, 202 insertions(+), 20 deletions(-) diff --git a/rust/lance/src/dataset/fragment.rs b/rust/lance/src/dataset/fragment.rs index 4e139a62679..f1b807005a3 100644 --- a/rust/lance/src/dataset/fragment.rs +++ b/rust/lance/src/dataset/fragment.rs @@ -4123,6 +4123,164 @@ mod tests { "projecting the overlaid column should require its missing file", ); } + + /// 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, + ) { + use arrow_array::StructArray; + use arrow_schema::Fields; + + 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]); + } } #[rstest] diff --git a/rust/lance/src/dataset/overlay.rs b/rust/lance/src/dataset/overlay.rs index 88bf0bdc838..cb6e456156e 100644 --- a/rust/lance/src/dataset/overlay.rs +++ b/rust/lance/src/dataset/overlay.rs @@ -34,13 +34,13 @@ //! deleted row is simply dropped along with the row. This matches the spec with no //! special casing. -use std::collections::{BTreeSet, HashMap}; +use std::collections::{BTreeSet, HashMap, HashSet}; use std::sync::Arc; use arrow_array::{Array, ArrayRef, RecordBatch}; use arrow_select::interleave::interleave; use futures::StreamExt; -use lance_core::datatypes::Schema; +use lance_core::datatypes::{Field, Schema}; use lance_core::{Error, Result}; use roaring::RoaringBitmap; @@ -370,10 +370,13 @@ impl OverlayReadPlanner { /// lists, so planning is `O(total overlay fields + projected fields)` rather than /// `O(projected fields × overlays)` (a fragment can have thousands of fields). /// -/// An overlay contributes to a projected field when its `data_file.fields` includes -/// that field's id. Overlays on nested (non-top-level) fields are not yet supported -/// and simply match no projected field. Each contributing overlay *file* appears once -/// in `files`, shared by every field it covers. +/// An overlay contributes to a projected field when any id in its `data_file.fields` +/// belongs to that top-level field's subtree. Overlays are written against the field +/// ids the file actually stores, which for a struct or list column is its *leaf* +/// child ids (the V2_1 structural encoding does not record the parent id), so a leaf +/// is mapped back to its top-level projected ancestor — the whole column is then +/// fetched and replaced as a unit. Each contributing overlay *file* appears once in +/// `files`, shared by every top-level field it covers. pub fn plan_overlays(fragment: &FileFragment, projection: &Schema) -> Result { let overlays = &fragment.metadata.overlays; debug_assert!( @@ -383,23 +386,31 @@ pub fn plan_overlays(fragment: &FileFragment, projection: &Schema) -> Result = - projection.fields.iter().map(|f| (f.id, f)).collect(); + // Every id in the projection subtree (top-level fields *and* their nested + // descendants) mapped to its top-level projected field, so an overlay written + // against leaf ids still resolves to the column it belongs to, in O(1). + let mut top_level_of: HashMap = HashMap::new(); + for field in &projection.fields { + index_field_subtree(field, field, &mut top_level_of); + } // Walk overlays newest-first, visiting only each overlay's own fields. Every - // covered field is appended to its entry in `field_overlays`; because we walk - // newest-first, each entry ends up ordered newest-first for free. + // covered top-level field is appended to its entry in `field_overlays`; because we + // walk newest-first, each entry ends up ordered newest-first for free. let mut files = Vec::new(); let mut field_overlays: HashMap> = HashMap::new(); for overlay in overlays.iter().rev() { - // (field_id, field_pos) for each of this overlay's fields that we project. - let mut contributions = Vec::new(); - let mut covered_fields = Vec::new(); + // The distinct top-level fields this overlay covers, each with the + // `data_file.fields` position whose coverage to read (a struct/list stores + // several leaf ids all mapping to the same top-level field — the first wins, + // as whole-value replacement shares one coverage across the leaves). + let mut seen = HashSet::new(); + let mut contributions: Vec<(&Field, usize)> = Vec::new(); for (field_pos, &field_id) in overlay.data_file.fields.iter().enumerate() { - if let Some(field) = projected.get(&field_id) { - contributions.push((field_id, field_pos)); - covered_fields.push((*field).clone()); + if let Some(&top) = top_level_of.get(&field_id) + && seen.insert(top.id) + { + contributions.push((top, field_pos)); } } if contributions.is_empty() { @@ -409,13 +420,16 @@ pub fn plan_overlays(fragment: &FileFragment, projection: &Schema) -> Result Result(top: &'a Field, node: &'a Field, out: &mut HashMap) { + out.insert(node.id, top); + for child in &node.children { + index_field_subtree(top, child, out); + } +} + /// Open the overlay readers a specific read needs and return the per-field plans to /// merge, pruned to that read. /// From 76f9b936c4b9dfe08cacc60d72340cb44aa6a7bf Mon Sep 17 00:00:00 2001 From: Will Jones Date: Mon, 13 Jul 2026 10:44:31 -0700 Subject: [PATCH 11/16] fix(overlay): resolve overlays per leaf atom, splicing into columns MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The struct/list fix mapped an overlay's stored fields to their top-level projected column and replaced the whole column, which panicked whenever an overlay's fields did not equal that column's full leaf set — e.g. an overlay on sub-field `s.a` while the read projects the whole struct `s` (it fetched `s` from a file holding only `a`). Reported by review. Resolve per *atom* instead: a per-row field an overlay replaces as a unit (a primitive leaf, or a whole list/map; structs are recursed through). An overlay contributes to an atom when its stored leaf ids fall in the atom's leaf set; the atom's value is fetched (projected to exactly the atom) and spliced back into its output column, rebuilding the struct spine along the path. Overlays on different sub-fields of one struct resolve independently, newest-wins is per atom, and an overlay on a non-projected field is skipped (its file never opened). Adds `enumerate_atoms`, `collect_leaf_ids`, `descend_by_ids`, `splice_by_ids`. Also per review: clarify the reader-priority comment (priority 0 is correct because overlay reads are issued late, only when a ready consumer polls the batch), note the offsets-materialization follow-up, and fix a doc typo ("into" -> "onto"). Tests: sub-field overlay under a parent-struct projection (the panic), non-projected sibling skipped, three-level nesting, multiple sub-field overlays with per-atom newest-wins, and descend/splice unit tests (including struct-null preservation). Co-Authored-By: Claude Opus 4.8 (1M context) --- rust/lance/src/dataset/fragment.rs | 287 ++++++++++++++++- rust/lance/src/dataset/overlay.rs | 492 +++++++++++++++++++++-------- 2 files changed, 642 insertions(+), 137 deletions(-) diff --git a/rust/lance/src/dataset/fragment.rs b/rust/lance/src/dataset/fragment.rs index f1b807005a3..479ff7ae5ad 100644 --- a/rust/lance/src/dataset/fragment.rs +++ b/rust/lance/src/dataset/fragment.rs @@ -2641,7 +2641,7 @@ impl FragmentReader { Ok(result.project_by_schema(&output_schema)?) } - /// Merge data overlay values into a stream of base batches. + /// 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`, @@ -2667,6 +2667,11 @@ impl FragmentReader { // 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()); @@ -4281,6 +4286,286 @@ mod tests { assert_eq!(row1.values(), &[1, 10]); assert_eq!(row2.values(), &[77, 88, 99]); } + + use arrow_array::StructArray; + use arrow_schema::Fields; + + /// 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]); + } } #[rstest] diff --git a/rust/lance/src/dataset/overlay.rs b/rust/lance/src/dataset/overlay.rs index cb6e456156e..c2260d3b464 100644 --- a/rust/lance/src/dataset/overlay.rs +++ b/rust/lance/src/dataset/overlay.rs @@ -34,10 +34,11 @@ //! deleted row is simply dropped along with the row. This matches the spec with no //! special casing. -use std::collections::{BTreeSet, HashMap, HashSet}; +use std::collections::{BTreeSet, HashMap}; use std::sync::Arc; -use arrow_array::{Array, ArrayRef, RecordBatch}; +use arrow_array::{Array, ArrayRef, RecordBatch, StructArray}; +use arrow_schema::DataType; use arrow_select::interleave::interleave; use futures::StreamExt; use lance_core::datatypes::{Field, Schema}; @@ -290,28 +291,33 @@ fn assemble_overlay_column( interleave(&sources, &routing.indices).map_err(Error::from) } -/// One overlay's contribution to one projected field, with its file reader opened: -/// the cells it covers, and the reader from which the field's values are fetched by -/// `offset_in_overlay` at merge time. +/// One overlay's contribution to one projected atom, with its file reader opened. #[derive(Debug, Clone)] -struct LoadedFieldOverlay { - /// The `offset_in_frag` cells this overlay covers for the field. +struct LoadedAtomOverlay { + /// The `offset_in_frag` cells this overlay covers for the atom. coverage: Arc, - /// Reader over the overlay data file, projected to the covered fields; shared - /// across the fields that the same file covers. + /// Reader over the overlay data file, projected to the covered atoms; shared + /// across the atoms that the same file covers. reader: Arc, - /// Single-field projection used when fetching the value column. - field_projection: Arc, } -/// The overlays that apply to a single projected field, ordered newest-first, with -/// readers opened and pruned to a specific read. `field_name` is the top-level -/// read-batch column name the plan applies to. Produced by [`resolve_overlays`] and -/// consumed by [`merge_overlay_batch`]. +/// The overlays that apply to a single projected atom — 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 atoms). 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 FieldOverlayPlan { - field_name: String, - overlays_newest_first: Vec, +pub struct AtomOverlayPlan { + /// The top-level output column the atom 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 atom (empty when the atom *is* + /// the top-level column). Drives descending to, and splicing back, the atom. + ancestor_ids: Vec, + /// Projection of exactly the atom (its ancestor path pruned to the atom subtree), + /// used to fetch the atom'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 @@ -319,43 +325,43 @@ pub struct FieldOverlayPlan { #[derive(Debug, Clone)] struct PlannedOverlayFile { data_file: DataFile, - /// The covered ∩ projected fields to project when the file is opened, so a - /// single reader serves every field the file contributes to. + /// The covered ∩ projected atoms to project when the file is opened, so a single + /// reader serves every atom the file contributes to. open_projection: Arc, } -/// One overlay's contribution to one projected field, before the file is opened. +/// One overlay's contribution to one projected atom, before the file is opened. #[derive(Debug, Clone)] -struct PlannedFieldOverlay { +struct PlannedAtomOverlay { /// Index into [`OverlayReadPlanner::files`] of the file that supplies the value. file: usize, coverage: Arc, } -/// The overlays that apply to a single projected field, ordered newest-first, -/// before any file is opened. +/// The overlays that apply to a single projected atom, ordered newest-first, before +/// any file is opened. #[derive(Debug, Clone)] -struct PlannedField { - field_name: String, - /// Single-field projection used when fetching this field's value column. - field_projection: Arc, - overlays_newest_first: Vec, +struct PlannedAtom { + 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 -/// [`FieldOverlayPlan`]s for one specific read, opening only the files whose cells +/// [`AtomOverlayPlan`]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, - fields: Vec, + atoms: Vec, } impl OverlayReadPlanner { - /// True when no projected field has any overlay, so there is nothing to resolve. + /// True when no projected atom has any overlay, so there is nothing to resolve. pub fn is_empty(&self) -> bool { - self.fields.is_empty() + self.atoms.is_empty() } } @@ -365,18 +371,17 @@ impl OverlayReadPlanner { /// /// Overlays are stored oldest-first (sorted newest-last on load, see /// `sort_overlays_newest_last`), so walking them in reverse gives newest-first -/// precedence. A single pass — visiting only each overlay's own fields and matching -/// them against the projection through a field-id map — builds the per-field overlay -/// lists, so planning is `O(total overlay fields + projected fields)` rather than -/// `O(projected fields × overlays)` (a fragment can have thousands of fields). +/// precedence. /// -/// An overlay contributes to a projected field when any id in its `data_file.fields` -/// belongs to that top-level field's subtree. Overlays are written against the field -/// ids the file actually stores, which for a struct or list column is its *leaf* -/// child ids (the V2_1 structural encoding does not record the parent id), so a leaf -/// is mapped back to its top-level projected ancestor — the whole column is then -/// fetched and replaced as a unit. Each contributing overlay *file* appears once in -/// `files`, shared by every top-level field it covers. +/// Resolution is per *atom* — 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 atom 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 +/// atom when any id in its `data_file.fields` falls in that atom's leaf set. At merge +/// time the atom'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 atom it covers. pub fn plan_overlays(fragment: &FileFragment, projection: &Schema) -> Result { let overlays = &fragment.metadata.overlays; debug_assert!( @@ -386,84 +391,184 @@ pub fn plan_overlays(fragment: &FileFragment, projection: &Schema) -> Result = HashMap::new(); - for field in &projection.fields { - index_field_subtree(field, field, &mut top_level_of); + // The projection's atoms, and a leaf-id -> atom-index map so an overlay's stored + // leaf ids resolve to the atom they belong to in O(1). + struct AtomInfo<'a> { + top_field: &'a Field, + ancestor_ids: Vec, + atom_id: i32, + } + let mut atom_infos: Vec = Vec::new(); + let mut leaf_to_atom: HashMap = HashMap::new(); + for top in &projection.fields { + for (atom, ancestor_ids) in enumerate_atoms(top) { + let idx = atom_infos.len(); + let mut value_leaf_ids = Vec::new(); + collect_leaf_ids(atom, &mut value_leaf_ids); + for leaf in value_leaf_ids { + leaf_to_atom.insert(leaf, idx); + } + atom_infos.push(AtomInfo { + top_field: top, + ancestor_ids, + atom_id: atom.id, + }); + } } - // Walk overlays newest-first, visiting only each overlay's own fields. Every - // covered top-level field is appended to its entry in `field_overlays`; because we - // walk newest-first, each entry ends up ordered newest-first for free. + // Walk overlays newest-first. For each overlay, find the atoms it covers and push + // (newest-first, for free) into their per-atom overlay lists. let mut files = Vec::new(); - let mut field_overlays: HashMap> = HashMap::new(); + let mut atom_overlays: Vec> = vec![Vec::new(); atom_infos.len()]; for overlay in overlays.iter().rev() { - // The distinct top-level fields this overlay covers, each with the - // `data_file.fields` position whose coverage to read (a struct/list stores - // several leaf ids all mapping to the same top-level field — the first wins, - // as whole-value replacement shares one coverage across the leaves). - let mut seen = HashSet::new(); - let mut contributions: Vec<(&Field, usize)> = Vec::new(); + // atom index -> the `data_file.fields` position whose coverage to read. An + // overlay writes one value per row per atom, so its leaves share a coverage; + // the first leaf of each atom to appear wins. + let mut covered: HashMap = HashMap::new(); for (field_pos, &field_id) in overlay.data_file.fields.iter().enumerate() { - if let Some(&top) = top_level_of.get(&field_id) - && seen.insert(top.id) - { - contributions.push((top, field_pos)); + if let Some(&atom_idx) = leaf_to_atom.get(&field_id) { + covered.entry(atom_idx).or_insert(field_pos); } } - if contributions.is_empty() { + if covered.is_empty() { continue; } let file = files.len(); + let covered_ids: Vec = covered.keys().map(|&i| atom_infos[i].atom_id).collect(); files.push(PlannedOverlayFile { data_file: overlay.data_file.clone(), - open_projection: Arc::new(Schema { - fields: contributions - .iter() - .map(|(top, _)| (*top).clone()) - .collect(), - metadata: Default::default(), - }), + open_projection: Arc::new(projection.project_by_ids(&covered_ids, true)), }); - for (top, field_pos) in contributions { - field_overlays - .entry(top.id) - .or_default() - .push(PlannedFieldOverlay { - file, - coverage: overlay.coverage_for_field(field_pos)?, - }); + for (atom_idx, field_pos) in covered { + atom_overlays[atom_idx].push(PlannedAtomOverlay { + file, + coverage: overlay.coverage_for_field(field_pos)?, + }); } } - // Emit one PlannedField per projected field that has overlays, in projection - // order for deterministic results. - let mut fields = Vec::new(); - for field in &projection.fields { - if let Some(overlays_newest_first) = field_overlays.remove(&field.id) { - fields.push(PlannedField { - field_name: field.name.clone(), - field_projection: Arc::new(Schema { - fields: vec![field.clone()], - metadata: Default::default(), - }), - overlays_newest_first, - }); + // Emit one PlannedAtom per projected atom that has overlays, in atom order. + let mut atoms = Vec::new(); + for (idx, info) in atom_infos.iter().enumerate() { + let overlays_newest_first = std::mem::take(&mut atom_overlays[idx]); + if overlays_newest_first.is_empty() { + continue; } + atoms.push(PlannedAtom { + top_field: Arc::new(info.top_field.clone()), + ancestor_ids: info.ancestor_ids.clone(), + fetch_projection: Arc::new(projection.project_by_ids(&[info.atom_id], true)), + overlays_newest_first, + }); } - Ok(OverlayReadPlanner { files, fields }) + Ok(OverlayReadPlanner { files, atoms }) } -/// Map `node` and every descendant field id to `top`, the top-level projected field -/// the subtree belongs to. Used so an overlay written against a struct/list's leaf -/// ids resolves to the top-level column. -fn index_field_subtree<'a>(top: &'a Field, node: &'a Field, out: &mut HashMap) { - out.insert(node.id, top); - for child in &node.children { - index_field_subtree(top, child, out); +/// The per-row atoms 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 atom (values are one-per-row). A top-level primitive or +/// list yields a single atom with an empty path. +fn enumerate_atoms(top: &Field) -> Vec<(&Field, Vec)> { + fn recurse<'a>(field: &'a Field, path: &mut Vec, out: &mut Vec<(&'a Field, Vec)>) { + if matches!(field.data_type(), DataType::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 atom (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_atom`, 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_atom` (whole-column replacement). +fn splice_by_ids( + array: &ArrayRef, + field: &Field, + ancestor_ids: &[i32], + new_atom: ArrayRef, +) -> Result { + let Some((&id, rest)) = ancestor_ids.split_first() else { + return Ok(new_atom); + }; + 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_atom, + )?; + 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 @@ -480,14 +585,14 @@ pub async fn resolve_overlays( offsets_in_frag: &[u32], fragment: &FileFragment, read_config: &FragReadConfig, -) -> Result> { +) -> Result> { let read_offsets = read_offsets_bitmap(offsets_in_frag); - // A file is opened only if some field it covers has cells among the requested - // rows. This is the row-selection pruning: overlays outside the read are skipped. + // A file is opened only if some atom 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 field in &planner.fields { - for overlay in &field.overlays_newest_first { + for atom in &planner.atoms { + for overlay in &atom.overlays_newest_first { if !overlay.coverage.is_disjoint(&read_offsets) { file_needed[overlay.file] = true; } @@ -495,12 +600,14 @@ pub async fn resolve_overlays( } // Open each needed file once, concurrently. The reader is shared (via `Arc`) by - // every field that file covers. + // every atom that file covers. // - // TODO(overlay perf): these reads use the default reader priority. Once we - // benchmark take/scan over overlays, decide whether overlay value reads should - // inherit `read_config.reader_priority` (or get a dedicated priority) so they - // schedule alongside the base reads. + // 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]; @@ -517,21 +624,22 @@ pub async fn resolve_overlays( .await?; let mut plans = Vec::new(); - for field in &planner.fields { + for atom in &planner.atoms { let mut overlays_newest_first = Vec::new(); - for overlay in &field.overlays_newest_first { + for overlay in &atom.overlays_newest_first { let Some(reader) = &opened[overlay.file] else { continue; // pruned: coverage disjoint from the read }; - overlays_newest_first.push(LoadedFieldOverlay { + overlays_newest_first.push(LoadedAtomOverlay { coverage: overlay.coverage.clone(), reader: reader.clone(), - field_projection: field.field_projection.clone(), }); } if !overlays_newest_first.is_empty() { - plans.push(FieldOverlayPlan { - field_name: field.field_name.clone(), + plans.push(AtomOverlayPlan { + top_field: atom.top_field.clone(), + ancestor_ids: atom.ancestor_ids.clone(), + fetch_projection: atom.fetch_projection.clone(), overlays_newest_first, }); } @@ -554,16 +662,16 @@ fn read_offsets_bitmap(offsets_in_frag: &[u32]) -> RoaringBitmap { bitmap } -/// Resolve overlays for one base batch: route each projected field against the -/// batch's `offsets_in_frag`, fetch only the overlay values the batch needs -/// (concurrently with the base read), and assemble the merged columns. Fields with -/// no plan, and the row-id/row-address system columns, pass through. +/// Resolve overlays for one base batch: route each projected atom against the batch's +/// `offsets_in_frag`, fetch only the overlay values the batch needs (concurrently with +/// the base read), assemble the merged atom, and splice it into its output column. +/// Atoms with no covered rows, and columns with no plan, pass through. pub async fn merge_overlay_batch( base: ReadBatchFut, offsets_in_frag: &[u32], - plans: &[FieldOverlayPlan], + plans: &[AtomOverlayPlan], ) -> Result { - let field_work = futures::future::try_join_all(plans.iter().map(|plan| async move { + let atom_work = futures::future::try_join_all(plans.iter().map(|plan| async move { let coverages: Vec<&RoaringBitmap> = plan .overlays_newest_first .iter() @@ -571,46 +679,59 @@ pub async fn merge_overlay_batch( .collect(); let routing = route_overlays(offsets_in_frag, &coverages); if routing.all_fall_through() { - return Ok::<_, Error>((plan.field_name.as_str(), None)); + return Ok::<_, Error>((plan, None)); } + // Fetch each overlay's values and descend to the atom array. The fetch is + // projected to the atom's ancestor path, so the fetched column is the pruned + // top-level column; `descend_by_ids` walks it down to the atom. + let atom_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)| { - fetch_overlay_values( + .map(|(overlay, offsets_in_overlay)| async move { + let column = fetch_overlay_values( overlay.reader.as_ref(), - overlay.field_projection.clone(), + plan.fetch_projection.clone(), offsets_in_overlay, ) + .await?; + descend_by_ids(&column, atom_field, &plan.ancestor_ids) }), ) .await?; - Ok((plan.field_name.as_str(), Some((routing, fetched)))) + Ok((plan, Some((routing, fetched)))) })); // The base read and every overlay value read proceed concurrently. - let (batch, resolved) = futures::future::try_join(base, field_work).await?; + let (batch, resolved) = futures::future::try_join(base, atom_work).await?; let schema = batch.schema(); let mut columns = batch.columns().to_vec(); - for (field_name, work) in resolved { + for (plan, work) in resolved { let Some((routing, fetched)) = work else { continue; }; - let Some(idx) = schema.index_of(field_name).ok() else { - // The plan's field is not in this batch's projection; skip it. + 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; }; - columns[idx] = assemble_overlay_column(&columns[idx], &routing, &fetched)?; + let base_atom = descend_by_ids(&columns[idx], &plan.top_field, &plan.ancestor_ids)?; + let merged_atom = assemble_overlay_column(&base_atom, &routing, &fetched)?; + columns[idx] = splice_by_ids( + &columns[idx], + &plan.top_field, + &plan.ancestor_ids, + merged_atom, + )?; } 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. Returns a column of -/// `offsets_in_overlay.len()` values in the same order; empty input reads nothing -/// and returns an empty column. +/// 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, @@ -845,4 +966,103 @@ mod tests { 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)); + } } From 929e0e9681957a85f6f9b0722b15a208d011a840 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Mon, 13 Jul 2026 12:53:20 -0700 Subject: [PATCH 12/16] refactor(overlay): rename AtomOverlayPlan to LoadedAtom, inline routing accessors Self-review cleanups. `AtomOverlayPlan` is the opened counterpart of `PlannedAtom`, so name it `LoadedAtom` to match the `PlannedAtomOverlay` -> `LoadedAtomOverlay` pairing and stop it reading like a planning-phase type. Drop the `OverlayRouting::offsets_in_overlay`/`all_fall_through` getters: every caller is in-module and some already read the fields directly, so the thin forwarders only created a mixed convention. Co-Authored-By: Claude Opus 4.8 (1M context) --- rust/lance/src/dataset/overlay.rs | 45 +++++++++++-------------------- 1 file changed, 16 insertions(+), 29 deletions(-) diff --git a/rust/lance/src/dataset/overlay.rs b/rust/lance/src/dataset/overlay.rs index c2260d3b464..912b1a07f48 100644 --- a/rust/lance/src/dataset/overlay.rs +++ b/rust/lance/src/dataset/overlay.rs @@ -56,8 +56,8 @@ use crate::dataset::fragment::{FileFragment, FragReadConfig, GenericFileReader}; /// /// 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 (see [`OverlayRouting::offsets_in_overlay`]) rather than whole -/// columns, then build the merged column with [`assemble_overlay_column`]. +/// 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`; @@ -69,24 +69,11 @@ struct OverlayRouting { /// 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). + /// through to the base column, so the base is already the answer and no overlay + /// values need to be read). any_overlay: bool, } -impl OverlayRouting { - /// Per overlay (newest-first), the `offset_in_overlay` values to fetch from its - /// value column. - fn offsets_in_overlay(&self) -> &[Vec] { - &self.offsets_in_overlay - } - - /// True when no row is covered by any overlay, so the base column is already the - /// answer and no overlay values need to be read. - fn all_fall_through(&self) -> bool { - !self.any_overlay - } -} - /// 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`. /// @@ -254,7 +241,7 @@ fn route_arbitrary( /// `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 +/// `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( @@ -262,7 +249,7 @@ fn assemble_overlay_column( routing: &OverlayRouting, fetched_newest_first: &[ArrayRef], ) -> Result { - if routing.all_fall_through() { + if !routing.any_overlay { return Ok(base.clone()); } if fetched_newest_first.len() != routing.offsets_in_overlay.len() { @@ -307,7 +294,7 @@ struct LoadedAtomOverlay { /// and pruned to a specific read. Produced by [`resolve_overlays`] and consumed by /// [`merge_overlay_batch`]. #[derive(Debug, Clone)] -pub struct AtomOverlayPlan { +pub struct LoadedAtom { /// The top-level output column the atom lives in (its name locates the batch /// column; its field tree drives the descend/splice into that column). top_field: Arc, @@ -350,7 +337,7 @@ struct PlannedAtom { /// 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 -/// [`AtomOverlayPlan`]s for one specific read, opening only the files whose cells +/// [`LoadedAtom`]s for one specific read, opening only the files whose cells /// the read's rows actually touch. #[derive(Debug, Clone)] pub struct OverlayReadPlanner { @@ -585,7 +572,7 @@ pub async fn resolve_overlays( offsets_in_frag: &[u32], fragment: &FileFragment, read_config: &FragReadConfig, -) -> Result> { +) -> Result> { let read_offsets = read_offsets_bitmap(offsets_in_frag); // A file is opened only if some atom it covers has cells among the requested rows. @@ -636,7 +623,7 @@ pub async fn resolve_overlays( }); } if !overlays_newest_first.is_empty() { - plans.push(AtomOverlayPlan { + plans.push(LoadedAtom { top_field: atom.top_field.clone(), ancestor_ids: atom.ancestor_ids.clone(), fetch_projection: atom.fetch_projection.clone(), @@ -669,7 +656,7 @@ fn read_offsets_bitmap(offsets_in_frag: &[u32]) -> RoaringBitmap { pub async fn merge_overlay_batch( base: ReadBatchFut, offsets_in_frag: &[u32], - plans: &[AtomOverlayPlan], + plans: &[LoadedAtom], ) -> Result { let atom_work = futures::future::try_join_all(plans.iter().map(|plan| async move { let coverages: Vec<&RoaringBitmap> = plan @@ -678,7 +665,7 @@ pub async fn merge_overlay_batch( .map(|overlay| overlay.coverage.as_ref()) .collect(); let routing = route_overlays(offsets_in_frag, &coverages); - if routing.all_fall_through() { + if !routing.any_overlay { return Ok::<_, Error>((plan, None)); } // Fetch each overlay's values and descend to the atom array. The fetch is @@ -688,7 +675,7 @@ pub async fn merge_overlay_batch( let fetched = futures::future::try_join_all( plan.overlays_newest_first .iter() - .zip(routing.offsets_in_overlay()) + .zip(&routing.offsets_in_overlay) .map(|(overlay, offsets_in_overlay)| async move { let column = fetch_overlay_values( overlay.reader.as_ref(), @@ -792,7 +779,7 @@ mod tests { let routing = route_overlays(offsets, &coverages); let fetched: Vec = overlays_newest_first .iter() - .zip(routing.offsets_in_overlay()) + .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() @@ -899,12 +886,12 @@ mod tests { 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![0, 1]]); + 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()), + &UInt32Array::from(routing.offsets_in_overlay[0].clone()), None, ) .unwrap(), From 5ff87ce191d36e7b71d146edae345b833ebb7c00 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Mon, 13 Jul 2026 12:53:30 -0700 Subject: [PATCH 13/16] test(overlay): cover NULL base cells and map columns Add two end-to-end cases surfaced by self-review: - A NULL *base* cell (distinct from a NULL overlay value): a covered NULL base is overridden to the overlay's value, an uncovered NULL base falls through and stays NULL. - A top-level Map column, exercising a multi-leaf atom (key + value leaf ids both mapping back to the one Map atom). Maps require the 2.2+ file format, so this case runs at V2_2 rather than the V2_0/V2_1 parametrization. Also hoist the test module's StructArray/Fields imports to the top of `overlay_read` and update a stale reference to the removed `all_fall_through`. Co-Authored-By: Claude Opus 4.8 (1M context) --- rust/lance/src/dataset/fragment.rs | 155 +++++++++++++++++++++++++++-- 1 file changed, 147 insertions(+), 8 deletions(-) diff --git a/rust/lance/src/dataset/fragment.rs b/rust/lance/src/dataset/fragment.rs index 479ff7ae5ad..1794e7ea34c 100644 --- a/rust/lance/src/dataset/fragment.rs +++ b/rust/lance/src/dataset/fragment.rs @@ -3224,9 +3224,9 @@ mod tests { use std::sync::Arc; use arrow_array::{ - Array, ArrayRef, Int32Array, RecordBatch, RecordBatchIterator, UInt64Array, + Array, ArrayRef, Int32Array, RecordBatch, RecordBatchIterator, StructArray, UInt64Array, }; - use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema}; + 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}; @@ -3482,6 +3482,69 @@ mod tests { 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( @@ -3848,7 +3911,7 @@ mod tests { /// A fragment with an overlay plan, but a take that touches only uncovered /// offsets, must fall entirely through to the base values (the - /// `routing.all_fall_through()` early-return with a plan present). + /// `!routing.any_overlay` early-return with a plan present). #[rstest] #[tokio::test] async fn test_take_plan_present_all_offsets_uncovered( @@ -4138,9 +4201,6 @@ mod tests { async fn test_struct_overlay_end_to_end( #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_1)] version: LanceFileVersion, ) { - use arrow_array::StructArray; - use arrow_schema::Fields; - let struct_fields = Fields::from(vec![ ArrowField::new("x", DataType::Int32, true), ArrowField::new("y", DataType::Int32, true), @@ -4287,8 +4347,87 @@ mod tests { assert_eq!(row2.values(), &[77, 88, 99]); } - use arrow_array::StructArray; - use arrow_schema::Fields; + /// A top-level Map column resolves as a single atom even though its value + /// spans two leaves (key and value): both leaf ids map back to the one Map + /// atom, 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) { From 2198b534acdc9bf6357d0269dca516586f850905 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Mon, 13 Jul 2026 14:46:17 -0700 Subject: [PATCH 14/16] test(overlay): assert the missing-file error and cover projecting an intermediate struct Tighten the two pruning tests to assert the failure is a not-found error naming the deleted overlay file, not a bare is_err() that any unrelated failure would satisfy. Extend test_overlay_deeply_nested_subfield to also project the intermediate struct by id while the overlay targets a deeper leaf, covering the case where a top-level-only mapping would miss the overlay. Co-Authored-By: Claude Opus 4.8 (1M context) --- rust/lance/src/dataset/fragment.rs | 34 ++++++++++++++++++++++++------ 1 file changed, 28 insertions(+), 6 deletions(-) diff --git a/rust/lance/src/dataset/fragment.rs b/rust/lance/src/dataset/fragment.rs index 1794e7ea34c..7817757f519 100644 --- a/rust/lance/src/dataset/fragment.rs +++ b/rust/lance/src/dataset/fragment.rs @@ -3745,10 +3745,14 @@ mod tests { 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. + // 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!( - frag.take(&[5], &val_only).await.is_err(), - "take hitting the overlay's coverage should require its missing file", + 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}", ); } @@ -4185,10 +4189,14 @@ mod tests { .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. + // 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!( - frag.take(&[0], &val_only).await.is_err(), - "projecting the overlaid column should require its missing file", + 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}", ); } @@ -4704,6 +4712,20 @@ mod tests { // 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 atom. (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]); } } From f6602329642c6f05c9ef5891d4142778b9e62ed2 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Mon, 13 Jul 2026 15:23:15 -0700 Subject: [PATCH 15/16] refactor(overlay): rename "atom" to "atomic field" "atom" was coined terminology. Rename it to "atomic field", matching the existing `atomic layout` concept in lance-core (field.rs), which already denotes a field treated as one indivisible per-row unit whose interior children aren't addressed individually. Types: LoadedAtom -> LoadedAtomicField, PlannedAtom -> PlannedAtomicField, and the *Overlay/enumerate_/*_id counterparts. No behavior change. Co-Authored-By: Claude Opus 4.8 (1M context) --- rust/lance/src/dataset/fragment.rs | 9 +- rust/lance/src/dataset/overlay.rs | 204 +++++++++++++++-------------- 2 files changed, 112 insertions(+), 101 deletions(-) diff --git a/rust/lance/src/dataset/fragment.rs b/rust/lance/src/dataset/fragment.rs index 7817757f519..c19e9cd6ae3 100644 --- a/rust/lance/src/dataset/fragment.rs +++ b/rust/lance/src/dataset/fragment.rs @@ -4355,9 +4355,10 @@ mod tests { assert_eq!(row2.values(), &[77, 88, 99]); } - /// A top-level Map column resolves as a single atom even though its value - /// spans two leaves (key and value): both leaf ids map back to the one Map - /// atom, and the whole map value at a covered offset is replaced. Maps require + /// 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] @@ -4716,7 +4717,7 @@ mod tests { // 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 atom. (This is the case wjones127/westonpace flagged where 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(); diff --git a/rust/lance/src/dataset/overlay.rs b/rust/lance/src/dataset/overlay.rs index 912b1a07f48..d973e7d2588 100644 --- a/rust/lance/src/dataset/overlay.rs +++ b/rust/lance/src/dataset/overlay.rs @@ -278,33 +278,34 @@ fn assemble_overlay_column( interleave(&sources, &routing.indices).map_err(Error::from) } -/// One overlay's contribution to one projected atom, with its file reader opened. +/// One overlay's contribution to one projected atomic field, with its file reader opened. #[derive(Debug, Clone)] -struct LoadedAtomOverlay { - /// The `offset_in_frag` cells this overlay covers for the atom. +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 atoms; shared - /// across the atoms that the same file covers. + /// 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 atom — a per-row field an overlay +/// 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 atoms). Ordered newest-first, with readers opened +/// 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 LoadedAtom { - /// The top-level output column the atom lives in (its name locates the batch +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 atom (empty when the atom *is* - /// the top-level column). Drives descending to, and splicing back, the atom. + /// 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 atom (its ancestor path pruned to the atom subtree), - /// used to fetch the atom's values from the overlay file. + /// 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, + overlays_newest_first: Vec, } /// One overlay file that may contribute to a read, before it is opened. Opened @@ -312,43 +313,43 @@ pub struct LoadedAtom { #[derive(Debug, Clone)] struct PlannedOverlayFile { data_file: DataFile, - /// The covered ∩ projected atoms to project when the file is opened, so a single - /// reader serves every atom the file contributes to. + /// 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 atom, before the file is opened. +/// One overlay's contribution to one projected atomic field, before the file is opened. #[derive(Debug, Clone)] -struct PlannedAtomOverlay { +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 atom, ordered newest-first, before +/// The overlays that apply to a single projected atomic field, ordered newest-first, before /// any file is opened. #[derive(Debug, Clone)] -struct PlannedAtom { +struct PlannedAtomicField { top_field: Arc, ancestor_ids: Vec, fetch_projection: Arc, - overlays_newest_first: Vec, + 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 -/// [`LoadedAtom`]s for one specific read, opening only the files whose cells +/// [`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, - atoms: Vec, + atomic_fields: Vec, } impl OverlayReadPlanner { - /// True when no projected atom has any overlay, so there is nothing to resolve. + /// True when no projected atomic field has any overlay, so there is nothing to resolve. pub fn is_empty(&self) -> bool { - self.atoms.is_empty() + self.atomic_fields.is_empty() } } @@ -360,15 +361,16 @@ impl OverlayReadPlanner { /// `sort_overlays_newest_last`), so walking them in reverse gives newest-first /// precedence. /// -/// Resolution is per *atom* — a per-row field that an overlay replaces as a unit: a +/// 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 atom and can be overlaid independently of its +/// 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 -/// atom when any id in its `data_file.fields` falls in that atom's leaf set. At merge -/// time the atom'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 atom it covers. +/// 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!( @@ -378,84 +380,92 @@ pub fn plan_overlays(fragment: &FileFragment, projection: &Schema) -> Result atom-index map so an overlay's stored - // leaf ids resolve to the atom they belong to in O(1). - struct AtomInfo<'a> { + // 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, - atom_id: i32, + atomic_field_id: i32, } - let mut atom_infos: Vec = Vec::new(); - let mut leaf_to_atom: HashMap = HashMap::new(); + let mut atomic_field_infos: Vec = Vec::new(); + let mut leaf_to_atomic_field: HashMap = HashMap::new(); for top in &projection.fields { - for (atom, ancestor_ids) in enumerate_atoms(top) { - let idx = atom_infos.len(); + 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(atom, &mut value_leaf_ids); + collect_leaf_ids(atomic_field, &mut value_leaf_ids); for leaf in value_leaf_ids { - leaf_to_atom.insert(leaf, idx); + leaf_to_atomic_field.insert(leaf, idx); } - atom_infos.push(AtomInfo { + atomic_field_infos.push(AtomicFieldInfo { top_field: top, ancestor_ids, - atom_id: atom.id, + atomic_field_id: atomic_field.id, }); } } - // Walk overlays newest-first. For each overlay, find the atoms it covers and push - // (newest-first, for free) into their per-atom overlay lists. + // 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 atom_overlays: Vec> = vec![Vec::new(); atom_infos.len()]; + let mut atomic_field_overlays: Vec> = + vec![Vec::new(); atomic_field_infos.len()]; for overlay in overlays.iter().rev() { - // atom index -> the `data_file.fields` position whose coverage to read. An - // overlay writes one value per row per atom, so its leaves share a coverage; - // the first leaf of each atom to appear wins. + // 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(&atom_idx) = leaf_to_atom.get(&field_id) { - covered.entry(atom_idx).or_insert(field_pos); + 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| atom_infos[i].atom_id).collect(); + 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 (atom_idx, field_pos) in covered { - atom_overlays[atom_idx].push(PlannedAtomOverlay { + 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 PlannedAtom per projected atom that has overlays, in atom order. - let mut atoms = Vec::new(); - for (idx, info) in atom_infos.iter().enumerate() { - let overlays_newest_first = std::mem::take(&mut atom_overlays[idx]); + // 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; } - atoms.push(PlannedAtom { + 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.atom_id], true)), + fetch_projection: Arc::new(projection.project_by_ids(&[info.atomic_field_id], true)), overlays_newest_first, }); } - Ok(OverlayReadPlanner { files, atoms }) + Ok(OverlayReadPlanner { + files, + atomic_fields, + }) } -/// The per-row atoms of a projected top-level field, each with the child-id path from +/// 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 atom (values are one-per-row). A top-level primitive or -/// list yields a single atom with an empty path. -fn enumerate_atoms(top: &Field) -> Vec<(&Field, Vec)> { +/// 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 matches!(field.data_type(), DataType::Struct(_)) { for child in &field.children { @@ -474,7 +484,7 @@ fn enumerate_atoms(top: &Field) -> Vec<(&Field, Vec)> { } /// Collect the leaf field ids in `field`'s subtree — the ids an overlay stores for -/// this atom (its own id if primitive; its item leaves if a list/map). +/// 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); @@ -514,17 +524,17 @@ fn descend_by_ids(array: &ArrayRef, field: &Field, ancestor_ids: &[i32]) -> Resu Ok(arr) } -/// Rebuild `array` with the array at `ancestor_ids` replaced by `new_atom`, cloning +/// 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_atom` (whole-column replacement). +/// 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_atom: ArrayRef, + new_atomic_field: ArrayRef, ) -> Result { let Some((&id, rest)) = ancestor_ids.split_first() else { - return Ok(new_atom); + return Ok(new_atomic_field); }; let child_pos = field .children @@ -551,7 +561,7 @@ fn splice_by_ids( &children[child_pos], &field.children[child_pos], rest, - new_atom, + new_atomic_field, )?; Ok(Arc::new(StructArray::try_new_with_length( fields, children, nulls, len, @@ -572,14 +582,14 @@ pub async fn resolve_overlays( offsets_in_frag: &[u32], fragment: &FileFragment, read_config: &FragReadConfig, -) -> Result> { +) -> Result> { let read_offsets = read_offsets_bitmap(offsets_in_frag); - // A file is opened only if some atom it covers has cells among the requested rows. + // 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 atom in &planner.atoms { - for overlay in &atom.overlays_newest_first { + 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; } @@ -587,7 +597,7 @@ pub async fn resolve_overlays( } // Open each needed file once, concurrently. The reader is shared (via `Arc`) by - // every atom that file covers. + // 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 @@ -611,22 +621,22 @@ pub async fn resolve_overlays( .await?; let mut plans = Vec::new(); - for atom in &planner.atoms { + for atomic_field in &planner.atomic_fields { let mut overlays_newest_first = Vec::new(); - for overlay in &atom.overlays_newest_first { + 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(LoadedAtomOverlay { + overlays_newest_first.push(LoadedAtomicFieldOverlay { coverage: overlay.coverage.clone(), reader: reader.clone(), }); } if !overlays_newest_first.is_empty() { - plans.push(LoadedAtom { - top_field: atom.top_field.clone(), - ancestor_ids: atom.ancestor_ids.clone(), - fetch_projection: atom.fetch_projection.clone(), + 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, }); } @@ -649,16 +659,16 @@ fn read_offsets_bitmap(offsets_in_frag: &[u32]) -> RoaringBitmap { bitmap } -/// Resolve overlays for one base batch: route each projected atom against the batch's +/// 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 atom, and splice it into its output column. -/// Atoms with no covered rows, and columns with no plan, pass through. +/// 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: &[LoadedAtom], + plans: &[LoadedAtomicField], ) -> Result { - let atom_work = futures::future::try_join_all(plans.iter().map(|plan| async move { + let atomic_field_work = futures::future::try_join_all(plans.iter().map(|plan| async move { let coverages: Vec<&RoaringBitmap> = plan .overlays_newest_first .iter() @@ -668,10 +678,10 @@ pub async fn merge_overlay_batch( if !routing.any_overlay { return Ok::<_, Error>((plan, None)); } - // Fetch each overlay's values and descend to the atom array. The fetch is - // projected to the atom's ancestor path, so the fetched column is the pruned - // top-level column; `descend_by_ids` walks it down to the atom. - let atom_field = &plan.fetch_projection.fields[0]; + // 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() @@ -683,7 +693,7 @@ pub async fn merge_overlay_batch( offsets_in_overlay, ) .await?; - descend_by_ids(&column, atom_field, &plan.ancestor_ids) + descend_by_ids(&column, atomic_field, &plan.ancestor_ids) }), ) .await?; @@ -691,7 +701,7 @@ pub async fn merge_overlay_batch( })); // The base read and every overlay value read proceed concurrently. - let (batch, resolved) = futures::future::try_join(base, atom_work).await?; + let (batch, resolved) = futures::future::try_join(base, atomic_field_work).await?; let schema = batch.schema(); let mut columns = batch.columns().to_vec(); @@ -703,13 +713,13 @@ pub async fn merge_overlay_batch( // The plan's column is not in this batch's projection; skip it. continue; }; - let base_atom = descend_by_ids(&columns[idx], &plan.top_field, &plan.ancestor_ids)?; - let merged_atom = assemble_overlay_column(&base_atom, &routing, &fetched)?; + 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_atom, + merged_atomic_field, )?; } Ok(RecordBatch::try_new(schema, columns)?) From a1bce9abaf60004334199a3530db151ff3b49386 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Mon, 13 Jul 2026 15:36:30 -0700 Subject: [PATCH 16/16] perf(overlay): avoid materializing arrow types when enumerating atomic fields enumerate_atomic_fields tested `matches!(field.data_type(), DataType::Struct(_))` at every node. Field::data_type() rebuilds the whole arrow DataType of a struct's subtree on each call, so enumerating a wide/deeply-nested projection was superlinear (~O(N * depth)) with heavy transient allocation. Use the O(1) `logical_type.is_struct()` check instead; same semantics, no allocation. Co-Authored-By: Claude Opus 4.8 (1M context) --- rust/lance/src/dataset/overlay.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/rust/lance/src/dataset/overlay.rs b/rust/lance/src/dataset/overlay.rs index d973e7d2588..d73329d8a41 100644 --- a/rust/lance/src/dataset/overlay.rs +++ b/rust/lance/src/dataset/overlay.rs @@ -38,7 +38,6 @@ use std::collections::{BTreeSet, HashMap}; use std::sync::Arc; use arrow_array::{Array, ArrayRef, RecordBatch, StructArray}; -use arrow_schema::DataType; use arrow_select::interleave::interleave; use futures::StreamExt; use lance_core::datatypes::{Field, Schema}; @@ -467,7 +466,7 @@ pub fn plan_overlays(fragment: &FileFragment, projection: &Schema) -> Result Vec<(&Field, Vec)> { fn recurse<'a>(field: &'a Field, path: &mut Vec, out: &mut Vec<(&'a Field, Vec)>) { - if matches!(field.data_type(), DataType::Struct(_)) { + if field.logical_type.is_struct() { for child in &field.children { path.push(child.id); recurse(child, path, out);