From fbd230ad9f3e73a54e5d84021606d271fbdc4d82 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Fri, 19 Jun 2026 15:55:38 -0700 Subject: [PATCH 01/13] docs: specify data overlay files for the table format Add a specification for data overlay files: small files attached to a fragment that supply new values for a subset of (row offset, field) cells without rewriting the base data files, for cheap cell-level updates. - protos/table.proto: rework DataOverlayFile with a dense/sparse coverage oneof (shared_offset_bitmap vs new FieldCoverage), rename read_version to committed_version (effective, commit-stamped), and document rank-based addressing with no offset column. Document reader feature flag 64. - docs: add data_overlay_file.md (full spec, worked example, guidance stub) and link it from the table format overview. Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/src/format/table/.pages | 1 + docs/src/format/table/data_overlay_file.md | 390 +++++++++++++++++++++ docs/src/format/table/index.md | 29 ++ protos/table.proto | 74 ++++ 4 files changed, 494 insertions(+) create mode 100644 docs/src/format/table/data_overlay_file.md diff --git a/docs/src/format/table/.pages b/docs/src/format/table/.pages index 16c20058608..5b0cb0e95e6 100644 --- a/docs/src/format/table/.pages +++ b/docs/src/format/table/.pages @@ -6,4 +6,5 @@ nav: - Layout: layout.md - Branch & Tag: branch_tag.md - Row ID & Lineage: row_id_lineage.md + - Data Overlay Files: data_overlay_file.md - MemTable & WAL: mem_wal.md diff --git a/docs/src/format/table/data_overlay_file.md b/docs/src/format/table/data_overlay_file.md new file mode 100644 index 00000000000..5540f860018 --- /dev/null +++ b/docs/src/format/table/data_overlay_file.md @@ -0,0 +1,390 @@ +# Data Overlay Files + +!!! note "Overlay files require feature flag 64 (data overlay files)" + + A reader or writer that does not understand overlay files must refuse a + dataset that uses them. Silently ignoring an overlay would return stale base + values, which is a correctness bug rather than a degraded experience. + +Overlay files supply new values for a subset of `(row offset, field)` cells +within a fragment **without rewriting the fragment's base data files**. They make +updates cheap when only a small fraction of rows and/or columns change: instead +of rewriting whole columns or moving rows to a new fragment, a writer appends a +small file carrying just the changed cells. + +This is Lance's third mechanism for changing data in place, alongside +[deletion files](index.md#deletion-files) (which remove rows) and +[data evolution](index.md#data-evolution) (which adds or rewrites whole columns). +An overlay changes individual cells. + +## Concepts + +### Coverage and resolution + +Each overlay declares which cells it provides through a **coverage** bitmap (or, +for sparse overlays, one bitmap per field). The bitmaps index **physical row +offsets** — positions in the base data files, counting deleted rows — so they are +stable across deletions, exactly like deletion vectors. + +To resolve a cell `(offset, field)` on read, walk the fragment's overlays from +**newest to oldest**. The first overlay that covers `(offset, field)` wins; its +value is used. If no overlay covers the cell, the value falls through to the base +data file (or is `NULL` if no base data file holds that field). + +Precedence among overlays is determined by: + +1. `committed_version` — higher wins (see [Versioning](#versioning-and-ordering)). +2. Position in `DataFragment.overlays` as a tiebreaker — a later entry is newer. + +A covered offset whose value is `NULL` overrides the cell **to** `NULL`. This is +distinct from an offset that is simply absent from the bitmap, which falls +through to the base. Coverage, not value-nullness, decides whether an overlay +applies. + +### Interaction with deletions + +Deletions take precedence over overlays. If a row offset is marked deleted in the +fragment's deletion file, any overlay value for that offset is dead and is +ignored, regardless of commit order. This keeps the invariant simple: a deletion +is the final word on a row, so a concurrent overlay against a row that was +deleted needs no special conflict handling — its values are merely inert. + +### Physical layout + +An overlay's data file stores **one value column per field**, in the order of +`data_file.fields`. It does **not** store a row-offset key column. The position of +a covered offset's value within its column is the **rank** of that offset in the +field's coverage bitmap — the number of set bits below it. For a Roaring bitmap +this is an O(1) operation, so random access to any cell is a rank computation +followed by a single value fetch, with no offset column to read and no binary +search. + +Because different fields may cover different offset sets, the value columns of a +single sparse overlay may have **different lengths**. The Lance file format +permits columns of differing item counts within one file, so a sparse overlay is +representable as a single file. (See [Writer support](#writer-support) for the +current implementation status.) + +### Dense vs. sparse overlays + +A single overlay is one of two shapes: + +- **Dense (rectangular).** One `shared_offset_bitmap` applies to every field. Every + covered offset has a value for every field. This is the common case for a plain + `UPDATE`, where one `SET` list is applied to one set of rows. +- **Sparse.** A `FieldCoverage` carries one bitmap per field, used when different + fields cover different offset sets — for example a `MERGE` with multiple + `WHEN MATCHED` branches, where different rows update different columns. A dense + overlay would have to widen to the bounding rectangle and fill the untouched + cells with their current values (post-images), which for wide columns such as + embeddings means re-storing data that did not change. A sparse overlay stores + exactly the changed cells. + +A writer may always express a non-rectangular update as **multiple dense overlays +in one transaction** (one per coverage group) instead of a single sparse overlay. + +## Protobuf + +
+DataOverlayFile protobuf message + +```protobuf +%%% proto.message.DataOverlayFile %%% +``` + +
+ +
+FieldCoverage protobuf message + +```protobuf +%%% proto.message.FieldCoverage %%% +``` + +
+ +## Versioning and ordering + +Overlays reuse the dataset version as their ordering clock rather than +introducing a separate generation counter. + +`committed_version` is the dataset version at which an overlay **became +effective** — the version of the commit that introduced it, **not** the version +it was read from. It is stamped at commit time and re-stamped if the commit is +retried, in the same way as the created-at / last-updated-at version sequences. + +This single value drives every ordering decision: + +- **Overlay vs. overlay** (read precedence): higher `committed_version` wins. +- **Overlay vs. index** (query correctness): an index records the + `dataset_version` it was built from. An index whose `dataset_version >= + committed_version` already incorporates the overlay. An overlay whose + `committed_version > index.dataset_version` is newer than the index and its + cells must be excluded from index results and re-evaluated. +- **Scheduler signal**: the gap between an overlay's `committed_version` and an + index's `dataset_version`, or between an overlay and the base, is a staleness + measure the compaction scheduler can use. + +!!! note "Why effective version, not read version" + + Suppose an overlay reads version 5 and commits at version 6, while an index + is built reading version 5 (before the overlay) and commits at version 7 with + `dataset_version = 5`. If the overlay stored its *read* version (5), the test + `5 > 5` is false, the row would not be excluded, and the index — which never + saw the overlay — would return a stale result. Storing the *effective* + version (6) makes `6 > 5` true, the cell is excluded and re-evaluated, and the + result is correct. + +## Index integration + +Building an index over a fragment that has overlays does **not** require dropping +the fragment from the index's coverage. The fragment stays indexed, and the query +path reconciles overlays at query time using an **exclusion set**. + +The exclusion set for an index on field `F` is the union of the coverage bitmaps, +restricted to field `F`, of every overlay whose `committed_version > +index.dataset_version`. The exclusion is **field-aware**: an overlay that touches +only unrelated columns does not exclude anything from the index on `F`. + +The query then proceeds as: + +1. Run the index search as usual, producing candidate rows. +2. Remove any candidate in the exclusion set. (Its indexed value may be stale.) +3. **Re-evaluate** the excluded rows against their current values — the same flat + path already used for the unindexed tail of fragments. For a scalar predicate + this re-applies the filter; for a vector query it re-scores the row's current + vector. Rows that still match are added back to the result. + +Step 3 is what makes exclusion correct rather than merely safe: removing a row +from index candidates without re-evaluating it would silently drop a row that +should match under its new value. + +### Correctness invariant + +> For every indexed field `F` and every row offset `o` in a fragment the index +> covers, the index's entry for `(o, F)` is trusted unless `o` is excluded. +> `o` is excluded iff some overlay with `committed_version > index.dataset_version` +> covers `(o, F)`. + +The write and compaction paths together preserve this: + +- **Writes** change a cell only by adding an overlay, and that overlay's + `committed_version` exceeds the version of any pre-existing index — so the + change is always covered by an exclusion. +- **Compaction** may remove an overlay only if the index no longer relies on it + (see below). + +## Compaction + +Overlays accumulate read cost — every overlay is a bitmap to test and a possible +file to open. Compaction bounds that cost in two modes: + +- **Overlay → overlay.** Merge several overlays into fewer, computing the + post-image per `(offset, field)` by walking the merged overlays newest-first. + The merged overlay takes the **maximum** `committed_version` of its inputs, so + the exclusion semantics are preserved. Indexes are unaffected. This is cheap and + does not touch the base. +- **Overlay → base.** Fold overlays into a fresh base data file, computing the + post-image for every covered cell, then clear the overlays. The base is + complete, so every post-image is well defined. Overlay offsets are physical, so + they cannot survive a rewrite that reorders rows; folding therefore materializes + values rather than carrying overlays forward. + +!!! warning "Folding an indexed field must update its index" + + An overlay→base fold removes the overlay, which removes the exclusion signal + that kept an index correct. Folding an overlay that covers an indexed field + `F` is therefore equivalent to a column rewrite of `F` and must, in the same + commit, either rebuild the index to a `dataset_version` at least the folded + overlay's `committed_version`, or remove the fragment from the index's + coverage so the rows fall to the flat path. Otherwise the index would serve + stale values with no overlay to exclude them. This is the same rule that + already governs rewriting a column that an index is built on. + +When a fragment with overlays is compacted by a row-rewriting operation +(`RewriteRows`, which produces new fragments with new row addresses), the +overlays are folded into the new base as part of the rewrite, and existing +[fragment-reuse remapping](row_id_lineage.md) handles the row-address changes as +it does today. + +## Row lineage + +An overlay write updates the `last_updated_at_version` of every covered row, so +change-data-feed and time-travel queries observe the update. Because overlays are +addressed by physical offset, they do **not** require stable row IDs to be +enabled; lineage updates apply only when those features are on. + +## Worked example + +A table `users` with stable row IDs enabled and these fields: + +| field id | name | type | +|----------|-----------|-------------------------| +| 1 | id | `int32` (primary key) | +| 2 | name | `utf8` | +| 3 | age | `int32` | +| 4 | embedding | `fixed_size_list`| + +Created at version 1 as a single fragment `0` with one base data file +`data/file0.lance` holding all four columns. `physical_rows = 4`: + +| offset | id | name | age | embedding | +|--------|----|-------|-----|------------------| +| 0 | 1 | Alice | 30 | … | +| 1 | 2 | Bob | 25 | … | +| 2 | 3 | Carol | 40 | … | +| 3 | 4 | Dave | 22 | … | + +A BTree scalar index on `age` is built at version 1, covering fragment `0` +(`dataset_version = 1`). + +### Step 1 — write an overlay + +```sql +UPDATE users SET age = 26 WHERE id = 2; -- Bob, offset 1 +``` + +This touches one field (`age`) for one row, so the writer emits a dense overlay +and commits it as version 2. Fragment `0` gains: + +```text +DataOverlayFile { + data_file: { path: "data/overlay-.lance", fields: [3], column_indices: [0] } + coverage: shared_offset_bitmap = {1} + committed_version: 2 +} +``` + +The overlay file stores a single `age` column with one value, `[26]`, at +rank `{1}.rank(1) = 0`. `last_updated_at_version[1]` is set to 2. + +### Step 2 — read + +`SELECT id, age FROM users` reads base ages `[30, 25, 40, 22]`. For `age` +(field 3), the overlay covers offset 1, so `age[1]` is replaced with the overlay +value at position `{1}.rank(1) = 0` → `26`. Result ages: `[30, 26, 40, 22]`. + +### Step 3 — index query + +```sql +SELECT * FROM users WHERE age = 26; +``` + +The `age` index was built at `dataset_version = 1`; the overlay's +`committed_version` is 2. Since `2 > 1`, the overlay's coverage for `age`, `{1}`, +is the exclusion set for this query. + +- The index (built at v1) holds Bob's *old* `age = 25`, so a lookup for `26` + returns nothing from the index. +- Offset 1 is in the exclusion set, so it is re-evaluated on the flat path. Its + current `age` (26, via the overlay) matches `age = 26`, so Bob is returned. + +The mirror case `WHERE age = 25` shows exclusion preventing a stale hit: the index +returns offset 1 (stale `25`), but offset 1 is excluded, re-evaluated to `26`, and +correctly dropped. + +### Step 4 — a second, non-rectangular write + +```sql +MERGE INTO users USING staged ON users.id = staged.id +WHEN MATCHED AND staged.kind = 'rename' THEN UPDATE SET name = staged.name -- Carol(2), Dave(3) +WHEN MATCHED AND staged.kind = 'revec' THEN UPDATE SET embedding = staged.embedding -- Bob(1) +``` + +`name` is updated for offsets `{2, 3}` and `embedding` for offset `{1}` — different +fields over different rows. This is a sparse overlay, committed as version 3: + +```text +DataOverlayFile { + data_file: { path: "data/overlay-.lance", fields: [2, 4], column_indices: [0, 1] } + coverage: field_coverage { offset_bitmaps: [ {2,3}, {1} ] } + // name (field 2) ^ ^ embedding (field 4) + committed_version: 3 +} +``` + +The file's `name` column has **two** values (`["Caroline", "David"]`, at +ranks 0 and 1 of `{2,3}`) and its `embedding` column has **one** value (at rank 0 +of `{1}`) — columns of different lengths in one file. + +### Step 5 — read after the second write + +`SELECT name, age, embedding FROM users` resolves each field independently, +newest overlay first: + +- `name`: the v3 overlay covers `{2,3}` → `["Alice", "Bob", "Caroline", "David"]`. +- `age`: the v3 overlay does not cover `age`; the v2 overlay still applies at + offset 1 → `[30, 26, 40, 22]`. +- `embedding`: the v3 overlay covers `{1}` → Bob's vector is the new one, others + from base. + +Overlays from different versions coexist and apply per field. + +### Step 6 — compaction (overlay → base) + +The scheduler folds both overlays into fragment `0` at version 4, computing +post-images for `age`, `name`, and `embedding`, and writing a new base data file +`data/file1.lance` with those columns. In the old file, fields 2, 3, and 4 are +tombstoned (`-2`); field 1 (`id`) remains. The fragment's `overlays` list is +cleared. Row addresses are preserved (a column rewrite, not a row rewrite), so +stable row IDs and the deletion vector are untouched. + +Because the fold removed the overlay that was excluding offset 1 from the `age` +index, the same commit must reconcile that index: either rebuild it at +`dataset_version >= 2`, or drop fragment `0` from its coverage so `age` queries +fall to the flat path. After a rebuild at version 4, no overlay remains and the +`age` index directly returns `26` for Bob with no exclusion needed. + +## Guidance + +!!! note "This section is a stub." + + The following are implementation considerations, not part of the on-disk + specification. + +### When to overlay vs. rewrite a column vs. move rows + +*(To be expanded.)* The choice between appending an overlay, rewriting a full +column (data evolution), and moving updated rows to a new fragment depends on the +fraction of rows changed, the fraction of columns changed, column width, the +presence of indexes on the changed columns, and the accumulated overlay read +cost. Roughly: few rows changed favors overlays; most rows in a few columns +favors a column rewrite; most columns changed favors moving rows to a new +fragment. + +### Writer support + +*(To be expanded.)* Dense (rectangular) overlays write with the existing +equal-length file writer today. Sparse overlays stored as a **single** file +require the writer to emit columns of independent lengths, which the current v2 +writer does not yet do (it advances all columns from one global row counter). +Until that support lands, a writer can express a sparse update as multiple dense +overlays in one transaction. + +### Scheduling compaction + +*(To be expanded.)* The overlay→overlay and overlay→base modes have very +different costs; a cost/benefit scheduler decides when each is worthwhile, using +the version gap as a staleness signal. + +### Open questions + +*(To be resolved.)* + +- **Per-fragment vs. per-table overlays.** Overlays are attached per fragment. + Should there be a table-level overlay concept, and how would it interact with + fragment-level row addressing? +- **Relationship to LSM.** Overlays plus compaction resemble an LSM tree (newest + layer wins, periodic merge). How far should that analogy be taken, and what do + we deliberately do differently given Lance's random-access requirements? +- **Coverage bitmap spill.** Coverage bitmaps live inline in the manifest. Very + large coverage (an overlay touching many rows) may warrant external spill, as + the row-ID and last-updated-at sequences already do above a size threshold. + +## Related specifications + +- [Table format overview](index.md) +- [Transactions](transaction.md) +- [Row ID & Lineage](row_id_lineage.md) +- [Index Formats](../index/index.md) +- [Format Versioning](versioning.md) diff --git a/docs/src/format/table/index.md b/docs/src/format/table/index.md index 94ea4b90dc9..ce4d0b26613 100644 --- a/docs/src/format/table/index.md +++ b/docs/src/format/table/index.md @@ -168,6 +168,35 @@ However, this invalidates row addresses and requires rebuilding indices, which c +## Data Overlay Files + +!!! note "Overlay files require feature flag 64 (data overlay files)" + +Overlay files supply new values for a subset of `(row offset, field)` cells within +a fragment without rewriting the base data files. They make updates cheap when only +a small percentage of rows and/or columns change: a writer appends a small file +carrying just the changed cells instead of rewriting whole columns or moving rows +to a new fragment. + +On read, each cell is resolved by consulting the fragment's overlays from newest to +oldest; the first overlay covering that `(offset, field)` wins, otherwise the value +falls through to the base data file. Indices keep covering the fragment and reconcile +overlays at query time through a field-aware exclusion set. + +For the full specification — coverage and resolution rules, dense vs. sparse layout, +versioning, index integration, compaction, and a worked example — see the +[Data Overlay Files Specification](data_overlay_file.md). + +
+DataOverlayFile protobuf message + +```protobuf +%%% proto.message.DataOverlayFile %%% +``` + +
+ + ## Related Specifications ### Storage Layout diff --git a/protos/table.proto b/protos/table.proto index d298809d5d8..cc8b477a6a6 100644 --- a/protos/table.proto +++ b/protos/table.proto @@ -113,6 +113,11 @@ message Manifest { // * 2: row ids are stable and stored as part of the fragment metadata. // * 4: use v2 format (deprecated) // * 8: table config is present + // * 16: data files use multiple base paths (shallow clone / multi-base) + // * 32: the transaction file under _transactions is not written (inline only) + // * 64: data overlay files are present (see DataOverlayFile). Readers that do + // not understand overlays must refuse the dataset, since ignoring an overlay + // would silently return stale base values. uint64 reader_feature_flags = 9; // Feature flags for writers. @@ -311,6 +316,15 @@ message DataFragment { repeated DataFile files = 2; + // Optional overlay files for this fragment, which supply new values for a + // subset of cells without rewriting the base data files. This MUST be empty + // if the data overlay files feature flag (64) is not set in the manifest. + // + // Order is significant: a later entry is newer than an earlier one. When two + // overlays cover the same (offset, field) and share a `committed_version`, the + // later entry wins. See DataOverlayFile for the full resolution rules. + repeated DataOverlayFile overlays = 11; + // File that indicates which rows, if any, should be considered deleted. DeletionFile deletion_file = 3; @@ -433,6 +447,66 @@ message DataFile { optional uint32 base_id = 7; } // DataFile +// An overlay file supplies new values for a subset of (row offset, field) cells +// within a fragment, without rewriting the fragment's base data files. It is +// used for efficient updates when only a small fraction of rows and/or columns +// change. +// +// On read, a cell is resolved by consulting the fragment's overlays from newest +// to oldest: the first overlay that covers that (offset, field) wins; if none +// cover it, the value falls through to the base data file. Because deletions +// take precedence over overlays, an overlay value for an offset that is also +// marked deleted is dead and is ignored. +// +// The overlay's data file does NOT store a row-offset key column. Within a value +// column, the position of a covered offset's value is the rank (0-based count of +// set bits below it) of that offset within the field's coverage bitmap. Because +// fields may cover different offset sets, the value columns of a single overlay +// data file may have different lengths (which the Lance file format permits). +message DataOverlayFile { + // The data file storing the overlay's new cell values, one value column per + // field in `data_file.fields`. No row-offset key column is stored. + DataFile data_file = 1; + + // Which (offset, field) cells this overlay provides values for. + oneof coverage { + // A single 32-bit Roaring bitmap of physical row offsets that applies to + // every field in `data_file.fields` (a "dense" / rectangular overlay). + // Every covered offset has a value for every field. This is the common case + // for a plain UPDATE, where one SET list is applied to one set of rows. + bytes shared_offset_bitmap = 2; + // Per-field coverage for a "sparse" overlay, used when different fields cover + // different offset sets (e.g. a MERGE with multiple WHEN MATCHED branches). + FieldCoverage field_coverage = 4; + } + + // The dataset version at which this overlay became effective: the version of + // the commit that introduced it, NOT the version it was read from. It is + // stamped at commit time and re-stamped if the commit is retried, in the same + // way as the created-at / last-updated-at version sequences. + // + // This drives two orderings: + // * Versus index builds: an index whose `dataset_version` >= this value + // already incorporates this overlay. Otherwise the overlay's covered cells + // are excluded from index results for the affected fields and re-evaluated + // against their current values (see the Data Overlay Files specification). + // * Versus other overlays: when two overlays cover the same (offset, field), + // the one with the higher `committed_version` wins. Overlays that share a + // `committed_version` are ordered by their position in + // `DataFragment.overlays`, where a later entry is newer and wins. + uint64 committed_version = 3; +} + +// Per-field coverage for a sparse overlay. +message FieldCoverage { + // One entry per field in the overlay's `data_file.fields`, in the same order. + // Each is a 32-bit Roaring bitmap of the physical row offsets covered for that + // field. An offset present in a field's bitmap but mapped to a NULL value + // means the cell is overridden to NULL (distinct from an offset that is absent, + // which falls through to the base data file). + repeated bytes offset_bitmaps = 1; +} + // Deletion File // // The path of the deletion file is constructed as: From 6aaca5eb55b557f6cbea157dfcbb2cb20e771d52 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Mon, 22 Jun 2026 12:47:14 -0700 Subject: [PATCH 02/13] feat: add DataOverlay transaction operation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add the `DataOverlay` operation (and `DataOverlayGroup`) to attach overlay files to fragments without rewriting their base data. Mirrors the `DataReplacement` batch shape, appends to each fragment's `overlays` list, and documents permissive conflict semantics: concurrent overlays, appends, deletes, and column rewrites are compatible; row-rewrites, compaction, and overlay->base folds conflict. committed_version is left 0 by the writer and stamped at commit time. Proto only — Rust/Python bindings deferred. Co-Authored-By: Claude Opus 4.8 (1M context) --- protos/transaction.proto | 39 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/protos/transaction.proto b/protos/transaction.proto index e72e95025a4..bfc0eee354b 100644 --- a/protos/transaction.proto +++ b/protos/transaction.proto @@ -315,6 +315,44 @@ message Transaction { repeated DataReplacementGroup replacements = 1; } + // Overlay files to append to a single fragment, in order (the last entry is + // newest). The overlays are appended to the fragment's existing `overlays` + // list; they do not replace it, so overlays written by concurrent commits are + // preserved. + message DataOverlayGroup { + uint64 fragment_id = 1; + // Each DataOverlayFile.committed_version is left 0 by the writer and stamped + // to the new dataset version at commit time (re-stamped on retry), in the + // same way as the created-at / last-updated-at version sequences. The fields + // touched are read from each overlay's `data_file.fields`. + repeated DataOverlayFile overlays = 2; + } + + // Attach overlay files to fragments, supplying new values for a subset of + // (row offset, field) cells without rewriting the fragments' base data files. + // See the DataOverlayFile message in table.proto and the Data Overlay Files + // specification for resolution, coverage, and versioning rules. + // + // Conflict semantics (intentionally permissive, like DataReplacement). Against + // a concurrent operation that touches one of the same fragments: + // * Another DataOverlay (any fields): COMPATIBLE. Overlays stack; when two + // overlays cover the same (offset, field) the one with the higher + // `committed_version` wins, so independent backfills never conflict. + // * Append / new fragments: COMPATIBLE. + // * Delete: COMPATIBLE. A deletion takes precedence over an overlay, so an + // overlay value for a deleted offset is inert (no special handling needed). + // * DataReplacement or column-rewrite (Update with REWRITE_COLUMNS) of the + // same field: COMPATIBLE. Both preserve physical row addresses, so overlay + // offsets stay valid; the overlay is newer and wins its covered cells, and + // the version gate excludes those cells from any rebuilt index. + // * Row-rewrite, compaction, or an overlay->base fold of the fragment: + // CONFLICT. These change physical row addresses or consume the overlays, so + // the overlay's offsets are no longer valid. The writer must re-read the new + // fragment, recompute, and retry. + message DataOverlay { + repeated DataOverlayGroup groups = 1; + } + // Update the merged generations in MemWAL index. // This operation is used during merge-insert to atomically record which // generations have been merged to the base table. @@ -346,6 +384,7 @@ message Transaction { UpdateMemWalState update_mem_wal_state = 112; Clone clone = 113; UpdateBases update_bases = 114; + DataOverlay data_overlay = 115; } // Fields 200/202 (`blob_append` / `blob_overwrite`) previously represented blob dataset ops. From a88fbe2a21908d191ce181f7e75edc4038ea66b0 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Mon, 22 Jun 2026 14:29:29 -0700 Subject: [PATCH 03/13] chore: compile against DataOverlay proto, reject unsupported overlays The table/transaction proto changes generate new fields and an Operation variant. This wires the minimum needed to compile without implementing overlay support: - Emit empty `overlays` when converting fragments to proto. - Reject the `DataOverlay` transaction operation with NotSupported on read. Datasets that use overlays set reader feature flag 64, which already falls in the unknown-flag range rejected by `can_read_dataset`, so the library refuses them at the feature-flag layer. Co-Authored-By: Claude Opus 4.8 (1M context) --- rust/lance-table/benches/manifest_intern.rs | 2 ++ rust/lance-table/src/format/fragment.rs | 4 +++ rust/lance/src/dataset/transaction.rs | 28 +++++++++++++++++++++ 3 files changed, 34 insertions(+) diff --git a/rust/lance-table/benches/manifest_intern.rs b/rust/lance-table/benches/manifest_intern.rs index 78b7e352207..81bd57c1a22 100644 --- a/rust/lance-table/benches/manifest_intern.rs +++ b/rust/lance-table/benches/manifest_intern.rs @@ -59,6 +59,7 @@ fn make_uniform_pb_fragments(n: u64, num_fields: usize) -> Vec file_size_bytes: 0, base_id: None, }], + overlays: vec![], deletion_file: None, row_id_sequence: None, physical_rows: 1000, @@ -135,6 +136,7 @@ fn make_diverse_pb_fragments( file_size_bytes: 0, base_id: None, }], + overlays: vec![], deletion_file: None, row_id_sequence: None, physical_rows: 1000, diff --git a/rust/lance-table/src/format/fragment.rs b/rust/lance-table/src/format/fragment.rs index 431e466dbd4..e9d9ce036ee 100644 --- a/rust/lance-table/src/format/fragment.rs +++ b/rust/lance-table/src/format/fragment.rs @@ -716,6 +716,10 @@ impl From<&Fragment> for pb::DataFragment { Self { id: f.id, files: f.files.iter().map(pb::DataFile::from).collect(), + // Overlay files are not produced by this version of the library; a + // dataset that uses them sets reader feature flag 64, which is + // rejected at the feature-flag layer (see lance-table feature_flags). + overlays: vec![], deletion_file, row_id_sequence, physical_rows: f.physical_rows.unwrap_or_default() as u64, diff --git a/rust/lance/src/dataset/transaction.rs b/rust/lance/src/dataset/transaction.rs index 4555cd7ee6c..dcf0682812a 100644 --- a/rust/lance/src/dataset/transaction.rs +++ b/rust/lance/src/dataset/transaction.rs @@ -3289,6 +3289,16 @@ impl TryFrom for Transaction { })) => Operation::UpdateBases { new_bases: new_bases.into_iter().map(BasePath::from).collect(), }, + Some(pb::transaction::Operation::DataOverlay(_)) => { + // Overlay files are not supported by this version of the library. + // A dataset that uses them sets reader feature flag 64, which is + // already rejected at the feature-flag layer; reject here too so a + // transaction referencing the operation can never be applied. + return Err(Error::not_supported( + "data overlay files are not supported by this version of Lance \ + (reader feature flag 64)", + )); + } None => { return Err(Error::internal( "Transaction message did not contain an operation".to_string(), @@ -6127,4 +6137,22 @@ mod tests { assert!(!left.modifies_same_metadata(&different_key)); assert!(left.modifies_same_metadata(&replace)); } + + #[test] + fn test_data_overlay_operation_rejected() { + // Overlay files are not supported by this version of the library. A + // transaction carrying the DataOverlay operation must be rejected rather + // than silently ignored, mirroring the feature-flag-64 rejection. + let message = pb::Transaction { + read_version: 1, + uuid: Uuid::new_v4().to_string(), + operation: Some(pb::transaction::Operation::DataOverlay( + pb::transaction::DataOverlay { groups: vec![] }, + )), + ..Default::default() + }; + + let result = Transaction::try_from(message); + assert!(matches!(result, Err(Error::NotSupported { .. }))); + } } From 7243137593e92128ffbc0d7a8ab04fbfed121366 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Mon, 22 Jun 2026 19:37:13 -0700 Subject: [PATCH 04/13] feat(file): v2 writer/reader support columns of unequal length The v2 file writer advanced every column from a single global row counter, so a single file could only hold columns of equal length. Sparse data overlay files need columns whose item counts differ within one file (each field covers a different set of rows). Add `FileWriter::write_columns`, which writes a set of `(field, array)` pairs and advances each field's row counter independently, leaving other fields untouched. A field never written ends up as a zero-length column. `write_batch` is unchanged: it still advances all fields together, so ordinary rectangular files round-trip exactly as before. Per-column lengths were already derivable from page metadata; expose them via `FileReader::column_num_rows`. The reader already schedules each column from its own pages, so reading a column at its own length and random access within it work without further changes. Part of the Data Overlay Files feature (OSS-1323). Co-Authored-By: Claude Opus 4.8 (1M context) --- rust/lance-file/src/reader.rs | 15 ++ rust/lance-file/src/writer.rs | 343 ++++++++++++++++++++++++++++++++-- 2 files changed, 344 insertions(+), 14 deletions(-) diff --git a/rust/lance-file/src/reader.rs b/rust/lance-file/src/reader.rs index c454f73819e..45d50541879 100644 --- a/rust/lance-file/src/reader.rs +++ b/rust/lance-file/src/reader.rs @@ -491,6 +491,21 @@ impl FileReader { self.num_rows } + /// The number of rows stored in a single physical column. + /// + /// For ordinary (rectangular) files every column has the same length, equal + /// to [`num_rows`](Self::num_rows). Files written with + /// [`FileWriter::write_columns`](crate::writer::FileWriter::write_columns) + /// may have columns of differing lengths; this returns the length of one + /// such column, derived by summing its pages' row counts. Returns `None` if + /// `column_index` is out of bounds. + pub fn column_num_rows(&self, column_index: usize) -> Option { + self.metadata + .column_metadatas + .get(column_index) + .map(|col| col.pages.iter().map(|page| page.length).sum()) + } + pub fn metadata(&self) -> &Arc { &self.metadata } diff --git a/rust/lance-file/src/writer.rs b/rust/lance-file/src/writer.rs index 12bd50df6fe..63ff7a95314 100644 --- a/rust/lance-file/src/writer.rs +++ b/rust/lance-file/src/writer.rs @@ -6,7 +6,7 @@ use std::collections::HashMap; use std::sync::Arc; use std::sync::atomic::AtomicBool; -use arrow_array::RecordBatch; +use arrow_array::{ArrayRef, RecordBatch}; use arrow_data::ArrayData; use bytes::{Buf, BufMut, Bytes, BytesMut}; @@ -221,6 +221,11 @@ pub struct FileWriter { field_id_to_column_indices: Vec<(u32, u32)>, num_columns: u32, rows_written: u64, + // The number of rows written for each top-level field (i.e. each entry in + // `column_writers`). With `write_batch` every field advances together and + // these are all equal, but `write_columns` advances fields independently, so + // a single file may end up with columns of differing item counts. + field_rows_written: Vec, global_buffers: Vec<(u64, u64)>, schema_metadata: HashMap, options: FileWriterOptions, @@ -277,6 +282,7 @@ impl FileWriter { column_metadata: Vec::new(), num_columns: 0, rows_written: 0, + field_rows_written: Vec::new(), field_id_to_column_indices: Vec::new(), global_buffers: Vec::new(), schema_metadata: HashMap::new(), @@ -467,6 +473,7 @@ impl FileWriter { BatchEncoder::try_new(&schema, encoding_strategy.as_ref(), &encoding_options)?; self.num_columns = encoder.num_columns(); + self.field_rows_written = vec![0; encoder.field_encoders.len()]; self.column_writers = encoder.field_encoders; self.column_metadata = vec![initial_column_metadata(); self.num_columns as usize]; self.field_id_to_column_indices = encoder.field_id_to_column_index; @@ -490,13 +497,14 @@ impl FileWriter { batch: &RecordBatch, external_buffers: &mut OutOfLineBuffers, ) -> Result>> { - self.schema + let items = self + .schema .as_ref() .unwrap() .fields .iter() - .zip(self.column_writers.iter_mut()) - .map(|(field, column_writer)| { + .enumerate() + .map(|(field_idx, field)| { let array = batch .column_by_name(&field.name) @@ -507,19 +515,53 @@ impl FileWriter { ) .into(), ))?; + Ok((field_idx, array.clone())) + }) + .collect::>>()?; + self.encode_columns(&items, external_buffers) + } + + /// Encode a set of `(field index, array)` pairs, each advancing only its own + /// column. The returned tasks must be written before the per-field row + /// counters are advanced (see `advance_columns`). + fn encode_columns( + &mut self, + items: &[(usize, ArrayRef)], + external_buffers: &mut OutOfLineBuffers, + ) -> Result>> { + // Snapshot the starting row number of each field before borrowing the + // column writers mutably below. + let row_numbers = items + .iter() + .map(|(field_idx, _)| self.field_rows_written[*field_idx]) + .collect::>(); + items + .iter() + .zip(row_numbers) + .map(|((field_idx, array), row_number)| { let repdef = RepDefBuilder::default(); let num_rows = array.len() as u64; - column_writer.maybe_encode( + self.column_writers[*field_idx].maybe_encode( array.clone(), external_buffers, repdef, - self.rows_written, + row_number, num_rows, ) }) .collect::>>() } + /// Advance the per-field row counters after a set of columns has been + /// written, keeping `rows_written` (the file's logical length) in sync as the + /// maximum column length. + fn advance_columns(&mut self, items: &[(usize, ArrayRef)]) { + for (field_idx, array) in items { + self.field_rows_written[*field_idx] += array.len() as u64; + } + self.rows_written = self.field_rows_written.iter().copied().max().unwrap_or(0); + } + /// Schedule a batch of data to be written to the file /// /// Note: the future returned by this method may complete before the data has been fully @@ -557,18 +599,100 @@ impl FileWriter { .flatten() .collect::>(); - self.rows_written = match self.rows_written.checked_add(batch.num_rows() as u64) { - Some(rows_written) => rows_written, - None => { - return Err(Error::invalid_input_source(format!("cannot write batch with {} rows because {} rows have already been written and Lance files cannot contain more than 2^64 rows", num_rows, self.rows_written).into())); - } - }; + // `write_batch` advances every field by the same amount, keeping all + // columns equal length. Guard against overflowing the row counter. + if self.rows_written.checked_add(num_rows).is_none() { + return Err(Error::invalid_input_source(format!("cannot write batch with {} rows because {} rows have already been written and Lance files cannot contain more than 2^64 rows", num_rows, self.rows_written).into())); + } + for field_rows in self.field_rows_written.iter_mut() { + *field_rows += num_rows; + } + self.rows_written = self.field_rows_written.iter().copied().max().unwrap_or(0); self.write_pages(encoding_tasks).await?; Ok(()) } + /// Write a set of columns whose lengths may differ from one another. + /// + /// Unlike [`write_batch`](Self::write_batch), which advances every column + /// from a single shared row counter, this method advances each column + /// independently. The result is a single file whose columns may have + /// different item counts — the physical layout used by sparse data overlay + /// files, where each field covers a different set of rows. + /// + /// `columns` is a list of `(field index, array)` pairs, where the field + /// index refers to a top-level field in the writer's schema (the same order + /// as the schema's fields). A field may be written across multiple calls; + /// its values are appended. A field that is never written ends up as a + /// zero-length column. The writer must have been created with an explicit + /// schema (via [`try_new`](Self::try_new)); a lazy schema cannot be inferred + /// here because individual calls need not cover every field. + /// + /// ``` + /// # use arrow_array::{ArrayRef, Int32Array}; + /// # use std::sync::Arc; + /// # use lance_file::writer::FileWriter; + /// # async fn example(writer: &mut FileWriter) -> lance_core::Result<()> { + /// // Field 0 gets three values, field 1 gets one — a non-rectangular file. + /// let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); + /// let b: ArrayRef = Arc::new(Int32Array::from(vec![10])); + /// writer.write_columns(vec![(0, a), (1, b)]).await?; + /// # Ok(()) + /// # } + /// ``` + pub async fn write_columns(&mut self, columns: Vec<(usize, ArrayRef)>) -> Result<()> { + let schema = self.schema.as_ref().ok_or_else(|| { + Error::invalid_input_source( + "write_columns requires the writer to be created with an explicit schema".into(), + ) + })?; + // Validate field indices, lengths, and nullability up front. + for (field_idx, array) in &columns { + let field = schema.fields.get(*field_idx).ok_or_else(|| { + Error::invalid_input_source( + format!( + "write_columns: field index {} is out of bounds (schema has {} fields)", + field_idx, + schema.fields.len() + ) + .into(), + ) + })?; + if array.len() as u64 > u32::MAX as u64 { + return Err(Error::invalid_input_source( + "cannot write Lance files with more than 2^32 rows".into(), + )); + } + Self::verify_field_nullability(&array.to_data(), field)?; + } + // Skip empty arrays: a never-advanced field simply remains a zero-length + // column, which the encoders handle at `finish` time. + let columns = columns + .into_iter() + .filter(|(_, array)| !array.is_empty()) + .collect::>(); + if columns.is_empty() { + return Ok(()); + } + + let mut external_buffers = + OutOfLineBuffers::new(self.tell().await?, PAGE_BUFFER_ALIGNMENT as u64); + let encoding_tasks = self.encode_columns(&columns, &mut external_buffers)?; + for external_buffer in external_buffers.take_buffers() { + Self::do_write_buffer(&mut self.writer, &external_buffer).await?; + } + let encoding_tasks = encoding_tasks + .into_iter() + .flatten() + .collect::>(); + + self.advance_columns(&columns); + self.write_pages(encoding_tasks).await?; + Ok(()) + } + async fn write_column_metadata( &mut self, metadata: pbfile::ColumnMetadata, @@ -974,11 +1098,11 @@ mod tests { use std::collections::HashMap; use std::sync::Arc; - use crate::reader::{FileReader, FileReaderOptions, describe_encoding}; + use crate::reader::{FileReader, FileReaderOptions, ReaderProjection, describe_encoding}; use crate::testing::FsFixture; use crate::writer::{ENV_LANCE_FILE_WRITER_MAX_PAGE_BYTES, FileWriter, FileWriterOptions}; use arrow_array::builder::{Float32Builder, Int32Builder}; - use arrow_array::{Int32Array, RecordBatch, UInt64Array}; + use arrow_array::{ArrayRef, Int32Array, RecordBatch, UInt64Array}; use arrow_array::{RecordBatchReader, StringArray, types::Float64Type}; use arrow_schema::{DataType, Field, Field as ArrowField, Schema, Schema as ArrowSchema}; use lance_core::cache::LanceCache; @@ -990,6 +1114,7 @@ mod tests { use lance_encoding::version::LanceFileVersion; use lance_io::object_store::ObjectStore; use lance_io::utils::CachedFileSize; + use rstest::rstest; #[tokio::test] async fn test_basic_write() { @@ -1040,6 +1165,196 @@ mod tests { file_writer.finish().await.unwrap(); } + /// Read a single column back at an explicit range/index set, returning its + /// `Int32` values. Reading one column at a time is how unequal-length files + /// are consumed: a global full-scan would conflate columns of different + /// lengths into one (impossible) rectangular batch. + async fn read_int32_column( + reader: &FileReader, + schema: &LanceSchema, + version: LanceFileVersion, + name: &str, + params: lance_io::ReadBatchParams, + ) -> Vec> { + use futures::TryStreamExt; + use lance_encoding::decoder::FilterExpression; + + let projection = ReaderProjection::from_column_names(version, schema, &[name]).unwrap(); + let batches: Vec = reader + .read_stream_projected(params, 1024, 16, projection, FilterExpression::no_filter()) + .await + .unwrap() + .try_collect() + .await + .unwrap(); + batches + .iter() + .flat_map(|b| { + b.column(0) + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .collect::>() + }) + .collect() + } + + /// A single file may hold columns of differing item counts (no shared global + /// row counter). This is the physical layout used by sparse data overlay + /// files, where each field covers a different set of rows. + #[rstest] + #[tokio::test] + async fn test_write_columns_unequal_lengths( + #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_1)] version: LanceFileVersion, + ) { + use lance_io::ReadBatchParams; + + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("a", DataType::Int32, true), + ArrowField::new("b", DataType::Int32, true), + ArrowField::new("c", DataType::Int32, true), + ])); + let lance_schema = LanceSchema::try_from(arrow_schema.as_ref()).unwrap(); + + let fs = FsFixture::default(); + let options = FileWriterOptions { + format_version: Some(version), + ..Default::default() + }; + let mut writer = FileWriter::try_new( + fs.object_store.create(&fs.tmp_path).await.unwrap(), + lance_schema.clone(), + options, + ) + .unwrap(); + + // Field "a" gets 5 values across two calls (appending), field "b" gets a + // single value, and field "c" is never written (a zero-length column). + let a1: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); + let b: ArrayRef = Arc::new(Int32Array::from(vec![10])); + writer.write_columns(vec![(0, a1), (1, b)]).await.unwrap(); + let a2: ArrayRef = Arc::new(Int32Array::from(vec![4, 5])); + // A zero-length array for an otherwise-unwritten field is a no-op. + let c_empty: ArrayRef = Arc::new(Int32Array::from(Vec::::new())); + writer + .write_columns(vec![(0, a2), (2, c_empty)]) + .await + .unwrap(); + + let summary = writer.finish().await.unwrap(); + // The file's logical length is the longest column. + assert_eq!(summary.num_rows, 5); + + let file_scheduler = fs + .scheduler + .open_file(&fs.tmp_path, &CachedFileSize::unknown()) + .await + .unwrap(); + let reader = FileReader::try_open( + file_scheduler, + None, + Arc::::default(), + &LanceCache::no_cache(), + FileReaderOptions::default(), + ) + .await + .unwrap(); + + // Per-column row counts are recorded in / derivable from file metadata. + assert_eq!(reader.num_rows(), 5); + assert_eq!(reader.column_num_rows(0), Some(5)); + assert_eq!(reader.column_num_rows(1), Some(1)); + assert_eq!(reader.column_num_rows(2), Some(0)); + assert_eq!(reader.column_num_rows(3), None); + + // Each column reads back independently at its own length. + assert_eq!( + read_int32_column( + &reader, + &lance_schema, + version, + "a", + ReadBatchParams::Range(0..5) + ) + .await, + vec![Some(1), Some(2), Some(3), Some(4), Some(5)], + ); + assert_eq!( + read_int32_column( + &reader, + &lance_schema, + version, + "b", + ReadBatchParams::Range(0..1) + ) + .await, + vec![Some(10)], + ); + + // Random access by position within the longer column returns the right + // value even though other columns are shorter. (The take path requires + // strictly increasing indices.) + assert_eq!( + read_int32_column( + &reader, + &lance_schema, + version, + "a", + ReadBatchParams::Indices(arrow_array::UInt32Array::from(vec![0, 2, 4])), + ) + .await, + vec![Some(1), Some(3), Some(5)], + ); + } + + /// Files written the ordinary (rectangular) way keep equal column lengths, + /// so the unequal-length support is backwards compatible. + #[tokio::test] + async fn test_write_batch_keeps_equal_lengths() { + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("a", DataType::Int32, true), + ArrowField::new("b", DataType::Int32, true), + ])); + let lance_schema = LanceSchema::try_from(arrow_schema.as_ref()).unwrap(); + + let fs = FsFixture::default(); + let mut writer = FileWriter::try_new( + fs.object_store.create(&fs.tmp_path).await.unwrap(), + lance_schema, + FileWriterOptions::default(), + ) + .unwrap(); + let batch = RecordBatch::try_new( + arrow_schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(Int32Array::from(vec![4, 5, 6])), + ], + ) + .unwrap(); + writer.write_batch(&batch).await.unwrap(); + let summary = writer.finish().await.unwrap(); + assert_eq!(summary.num_rows, 3); + + let file_scheduler = fs + .scheduler + .open_file(&fs.tmp_path, &CachedFileSize::unknown()) + .await + .unwrap(); + let reader = FileReader::try_open( + file_scheduler, + None, + Arc::::default(), + &LanceCache::no_cache(), + FileReaderOptions::default(), + ) + .await + .unwrap(); + assert_eq!(reader.column_num_rows(0), Some(3)); + assert_eq!(reader.column_num_rows(1), Some(3)); + } + #[tokio::test] async fn test_max_page_bytes_enforced() { let arrow_field = Field::new("data", DataType::UInt64, false); From 1dba164882ddfc54110f0330c028424f5d28b501 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Tue, 23 Jun 2026 16:21:05 -0700 Subject: [PATCH 05/13] refactor(file): write_column singular API + reader length validation Address review on the unequal-length-columns work: - Rename `FileWriter::write_columns(Vec<(usize, ArrayRef)>)` to the simpler `write_column(column_index, array)`; callers loop instead of batching. - Drop doc comments from internal-only helpers. - Reader: validate that a projection's top-level columns share a length and reject mismatches up front with a descriptive error, instead of panicking in the decoder. `RangeFull`/`RangeFrom` now resolve to the projected common length rather than the file's longest column, so a single short column reads back at its own length. Rectangular files are unchanged. Adds `test_read_unequal_length_projection` (V2_0 + V2_1). --- rust/lance-file/src/reader.rs | 109 +++++++++++++--- rust/lance-file/src/writer.rs | 237 ++++++++++++++++++++++++++-------- 2 files changed, 272 insertions(+), 74 deletions(-) diff --git a/rust/lance-file/src/reader.rs b/rust/lance-file/src/reader.rs index 45d50541879..98135b60f7e 100644 --- a/rust/lance-file/src/reader.rs +++ b/rust/lance-file/src/reader.rs @@ -495,7 +495,7 @@ impl FileReader { /// /// For ordinary (rectangular) files every column has the same length, equal /// to [`num_rows`](Self::num_rows). Files written with - /// [`FileWriter::write_columns`](crate::writer::FileWriter::write_columns) + /// [`FileWriter::write_column`](crate::writer::FileWriter::write_column) /// may have columns of differing lengths; this returns the length of one /// such column, derived by summing its pages' row counts. Returns `None` if /// `column_index` is out of bounds. @@ -506,6 +506,72 @@ impl FileReader { .map(|col| col.pages.iter().map(|page| page.length).sum()) } + // Number of physical columns a field (and its descendants) contributes to a + // projection, mirroring `ReaderProjection::from_field_ids_helper`. Used to + // partition a projection's flat `column_indices` by top-level field. + fn count_projected_columns(field: &Field, is_structural: bool) -> usize { + let contributes = !is_structural + || field.children.is_empty() + || field.is_blob() + || field.is_packed_struct(); + let recurse = !is_structural || (!field.is_blob() && !field.is_packed_struct()); + let mut count = contributes as usize; + if recurse { + for child in &field.children { + count += Self::count_projected_columns(child, is_structural); + } + } + count + } + + // The top-level row count of each projected top-level field, in projection + // order. A field's length is the page-row sum of its root (first) physical + // column, which equals the field's top-level row count for primitive, + // struct, and list fields in both v2.0 and v2.1. Ordinary files have equal + // lengths across columns; files written with `write_column` may not. + fn projected_field_lengths(&self, projection: &ReaderProjection) -> Vec { + let is_structural = self.metadata.version() >= LanceFileVersion::V2_1; + let mut lengths = Vec::with_capacity(projection.schema.fields.len()); + let mut cursor = 0usize; + for field in &projection.schema.fields { + let n = Self::count_projected_columns(field, is_structural); + if n == 0 { + continue; + } + let Some(&root_column) = projection.column_indices.get(cursor) else { + break; + }; + lengths.push(self.column_num_rows(root_column as usize).unwrap_or(0)); + cursor += n; + } + lengths + } + + // The reader combines a projection's columns into rectangular batches, so + // they must all have the same length. Returns that common length, or a + // descriptive error (naming the per-column lengths) when they differ. + // Ordinary files always pass; only files written with `write_column` whose + // columns ended up unequal can fail, and those must be read separately. + fn verify_uniform_lengths(projection: &ReaderProjection, field_lengths: &[u64]) -> Result { + let first = field_lengths.first().copied().unwrap_or(0); + if field_lengths.iter().all(|&len| len == first) { + return Ok(first); + } + let columns = projection + .schema + .fields + .iter() + .map(|f| f.name.as_str()) + .zip(field_lengths.iter().copied()) + .map(|(name, len)| format!("{name}={len}")) + .collect::>() + .join(", "); + Err(Error::invalid_input(format!( + "cannot read columns of differing lengths together ({columns}); \ + read each column (or equal-length group) separately" + ))) + } + pub fn metadata(&self) -> &Arc { &self.metadata } @@ -1236,11 +1302,18 @@ impl FileReader { ) -> Result + Send>>> { let projection = projection.unwrap_or_else(|| self.base_projection.clone()); Self::validate_projection(&projection, &self.metadata)?; + // All projected columns must share a length: the reader combines them + // into rectangular batches. Ordinary files satisfy this (every column + // has `self.num_rows` rows); files written with `FileWriter::write_column` + // may not, and such columns must be read separately. `read_len` is that + // common length, which `RangeFull`/`RangeFrom` resolve against (rather + // than `self.num_rows`, the file's longest column). + let field_lengths = self.projected_field_lengths(&projection); + let read_len = Self::verify_uniform_lengths(&projection, &field_lengths)?; let verify_bound = |params: &ReadBatchParams, bound: u64, inclusive: bool| { - if bound > self.num_rows || bound == self.num_rows && inclusive { + if bound > read_len || (bound == read_len && inclusive) { Err(Error::invalid_input(format!( - "cannot read {:?} from file with {} rows", - params, self.num_rows + "cannot read {params:?} from columns with {read_len} rows" ))) } else { Ok(()) @@ -1282,13 +1355,8 @@ impl FileReader { } ReadBatchParams::RangeFrom(range) => { verify_bound(¶ms, range.start as u64, true)?; - self.read_range( - range.start as u64..self.num_rows, - batch_size, - projection, - filter, - ) - .await + self.read_range(range.start as u64..read_len, batch_size, projection, filter) + .await } ReadBatchParams::RangeTo(range) => { verify_bound(¶ms, range.end as u64, false)?; @@ -1296,7 +1364,7 @@ impl FileReader { .await } ReadBatchParams::RangeFull => { - self.read_range(0..self.num_rows, batch_size, projection, filter) + self.read_range(0..read_len, batch_size, projection, filter) .await } } @@ -1491,11 +1559,18 @@ impl FileReader { ) -> Result> { let projection = projection.unwrap_or_else(|| self.base_projection.clone()); Self::validate_projection(&projection, &self.metadata)?; + // All projected columns must share a length: the reader combines them + // into rectangular batches. Ordinary files satisfy this (every column + // has `self.num_rows` rows); files written with `FileWriter::write_column` + // may not, and such columns must be read separately. `read_len` is that + // common length, which `RangeFull`/`RangeFrom` resolve against (rather + // than `self.num_rows`, the file's longest column). + let field_lengths = self.projected_field_lengths(&projection); + let read_len = Self::verify_uniform_lengths(&projection, &field_lengths)?; let verify_bound = |params: &ReadBatchParams, bound: u64, inclusive: bool| { - if bound > self.num_rows || bound == self.num_rows && inclusive { + if bound > read_len || (bound == read_len && inclusive) { Err(Error::invalid_input(format!( - "cannot read {:?} from file with {} rows", - params, self.num_rows + "cannot read {params:?} from columns with {read_len} rows" ))) } else { Ok(()) @@ -1536,7 +1611,7 @@ impl FileReader { ReadBatchParams::RangeFrom(range) => { verify_bound(¶ms, range.start as u64, true)?; self.read_range_blocking( - range.start as u64..self.num_rows, + range.start as u64..read_len, batch_size, projection, filter, @@ -1547,7 +1622,7 @@ impl FileReader { self.read_range_blocking(0..range.end as u64, batch_size, projection, filter) } ReadBatchParams::RangeFull => { - self.read_range_blocking(0..self.num_rows, batch_size, projection, filter) + self.read_range_blocking(0..read_len, batch_size, projection, filter) } } } diff --git a/rust/lance-file/src/writer.rs b/rust/lance-file/src/writer.rs index 63ff7a95314..96d282e978b 100644 --- a/rust/lance-file/src/writer.rs +++ b/rust/lance-file/src/writer.rs @@ -223,7 +223,7 @@ pub struct FileWriter { rows_written: u64, // The number of rows written for each top-level field (i.e. each entry in // `column_writers`). With `write_batch` every field advances together and - // these are all equal, but `write_columns` advances fields independently, so + // these are all equal, but `write_column` advances one field at a time, so // a single file may end up with columns of differing item counts. field_rows_written: Vec, global_buffers: Vec<(u64, u64)>, @@ -521,9 +521,9 @@ impl FileWriter { self.encode_columns(&items, external_buffers) } - /// Encode a set of `(field index, array)` pairs, each advancing only its own - /// column. The returned tasks must be written before the per-field row - /// counters are advanced (see `advance_columns`). + // Encode a set of `(field index, array)` pairs, each advancing only its own + // column. The returned tasks must be written before the per-field row + // counters are advanced (see `advance_columns`). fn encode_columns( &mut self, items: &[(usize, ArrayRef)], @@ -552,9 +552,9 @@ impl FileWriter { .collect::>>() } - /// Advance the per-field row counters after a set of columns has been - /// written, keeping `rows_written` (the file's logical length) in sync as the - /// maximum column length. + // Advance the per-field row counters after a set of columns has been + // written, keeping `rows_written` (the file's logical length) in sync as the + // maximum column length. fn advance_columns(&mut self, items: &[(usize, ArrayRef)]) { for (field_idx, array) in items { self.field_rows_written[*field_idx] += array.len() as u64; @@ -614,21 +614,20 @@ impl FileWriter { Ok(()) } - /// Write a set of columns whose lengths may differ from one another. + /// Write a single column, advancing only that column's row counter. /// /// Unlike [`write_batch`](Self::write_batch), which advances every column - /// from a single shared row counter, this method advances each column - /// independently. The result is a single file whose columns may have - /// different item counts — the physical layout used by sparse data overlay - /// files, where each field covers a different set of rows. + /// from a single shared row counter, this method advances one column + /// independently. Used across calls it produces a single file whose columns + /// may have different item counts — the physical layout used by sparse data + /// overlay files, where each field covers a different set of rows. /// - /// `columns` is a list of `(field index, array)` pairs, where the field - /// index refers to a top-level field in the writer's schema (the same order - /// as the schema's fields). A field may be written across multiple calls; - /// its values are appended. A field that is never written ends up as a - /// zero-length column. The writer must have been created with an explicit - /// schema (via [`try_new`](Self::try_new)); a lazy schema cannot be inferred - /// here because individual calls need not cover every field. + /// `column_index` refers to a top-level field in the writer's schema (the + /// same order as the schema's fields). A column may be written across + /// multiple calls; its values are appended. A field that is never written + /// ends up as a zero-length column. The writer must have been created with + /// an explicit schema (via [`try_new`](Self::try_new)); a lazy schema cannot + /// be inferred here because individual calls need not cover every field. /// /// ``` /// # use arrow_array::{ArrayRef, Int32Array}; @@ -636,47 +635,41 @@ impl FileWriter { /// # use lance_file::writer::FileWriter; /// # async fn example(writer: &mut FileWriter) -> lance_core::Result<()> { /// // Field 0 gets three values, field 1 gets one — a non-rectangular file. - /// let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); - /// let b: ArrayRef = Arc::new(Int32Array::from(vec![10])); - /// writer.write_columns(vec![(0, a), (1, b)]).await?; + /// writer.write_column(0, Arc::new(Int32Array::from(vec![1, 2, 3]))).await?; + /// writer.write_column(1, Arc::new(Int32Array::from(vec![10]))).await?; /// # Ok(()) /// # } /// ``` - pub async fn write_columns(&mut self, columns: Vec<(usize, ArrayRef)>) -> Result<()> { + pub async fn write_column(&mut self, column_index: usize, array: ArrayRef) -> Result<()> { let schema = self.schema.as_ref().ok_or_else(|| { Error::invalid_input_source( - "write_columns requires the writer to be created with an explicit schema".into(), + "write_column requires the writer to be created with an explicit schema".into(), ) })?; - // Validate field indices, lengths, and nullability up front. - for (field_idx, array) in &columns { - let field = schema.fields.get(*field_idx).ok_or_else(|| { - Error::invalid_input_source( - format!( - "write_columns: field index {} is out of bounds (schema has {} fields)", - field_idx, - schema.fields.len() - ) - .into(), + let field = schema.fields.get(column_index).ok_or_else(|| { + Error::invalid_input_source( + format!( + "write_column: field index {} is out of bounds (schema has {} fields)", + column_index, + schema.fields.len() ) - })?; - if array.len() as u64 > u32::MAX as u64 { - return Err(Error::invalid_input_source( - "cannot write Lance files with more than 2^32 rows".into(), - )); - } - Self::verify_field_nullability(&array.to_data(), field)?; + .into(), + ) + })?; + if array.len() as u64 > u32::MAX as u64 { + return Err(Error::invalid_input_source( + "cannot write Lance files with more than 2^32 rows".into(), + )); } - // Skip empty arrays: a never-advanced field simply remains a zero-length - // column, which the encoders handle at `finish` time. - let columns = columns - .into_iter() - .filter(|(_, array)| !array.is_empty()) - .collect::>(); - if columns.is_empty() { + Self::verify_field_nullability(&array.to_data(), field)?; + + // A never-advanced field simply remains a zero-length column, which the + // encoders handle at `finish` time. + if array.is_empty() { return Ok(()); } + let columns = [(column_index, array)]; let mut external_buffers = OutOfLineBuffers::new(self.tell().await?, PAGE_BUFFER_ALIGNMENT as u64); let encoding_tasks = self.encode_columns(&columns, &mut external_buffers)?; @@ -1165,10 +1158,10 @@ mod tests { file_writer.finish().await.unwrap(); } - /// Read a single column back at an explicit range/index set, returning its - /// `Int32` values. Reading one column at a time is how unequal-length files - /// are consumed: a global full-scan would conflate columns of different - /// lengths into one (impossible) rectangular batch. + // Read a single column back at an explicit range/index set, returning its + // `Int32` values. Reading one column (or an equal-length group) at a time is + // how unequal-length files are consumed: a full scan across columns of + // differing lengths cannot form a single rectangular batch. async fn read_int32_column( reader: &FileReader, schema: &LanceSchema, @@ -1233,14 +1226,13 @@ mod tests { // single value, and field "c" is never written (a zero-length column). let a1: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); let b: ArrayRef = Arc::new(Int32Array::from(vec![10])); - writer.write_columns(vec![(0, a1), (1, b)]).await.unwrap(); + writer.write_column(0, a1).await.unwrap(); + writer.write_column(1, b).await.unwrap(); let a2: ArrayRef = Arc::new(Int32Array::from(vec![4, 5])); + writer.write_column(0, a2).await.unwrap(); // A zero-length array for an otherwise-unwritten field is a no-op. let c_empty: ArrayRef = Arc::new(Int32Array::from(Vec::::new())); - writer - .write_columns(vec![(0, a2), (2, c_empty)]) - .await - .unwrap(); + writer.write_column(2, c_empty).await.unwrap(); let summary = writer.finish().await.unwrap(); // The file's logical length is the longest column. @@ -1308,6 +1300,137 @@ mod tests { ); } + /// Reading an unequal-length file: + /// - a projection whose columns are equal length full-scans normally; + /// - a full scan across columns of differing length is rejected up front, + /// before any batch is produced (even though a prefix would be rectangular); + /// - a bounded read is valid as long as every projected column covers it; + /// - a single-column `RangeFull` resolves to that column's own length, not + /// the file's (maximum) length. + #[rstest] + #[tokio::test] + async fn test_read_unequal_length_projection( + #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_1)] version: LanceFileVersion, + ) { + use futures::TryStreamExt; + use lance_encoding::decoder::FilterExpression; + use lance_io::ReadBatchParams; + + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("a", DataType::Int32, true), + ArrowField::new("b", DataType::Int32, true), + ArrowField::new("c", DataType::Int32, true), + ])); + let lance_schema = LanceSchema::try_from(arrow_schema.as_ref()).unwrap(); + let fs = FsFixture::default(); + let options = FileWriterOptions { + format_version: Some(version), + ..Default::default() + }; + let mut writer = FileWriter::try_new( + fs.object_store.create(&fs.tmp_path).await.unwrap(), + lance_schema.clone(), + options, + ) + .unwrap(); + // "a" and "b" are equal length (5); "c" is shorter (1). + writer + .write_column(0, Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5]))) + .await + .unwrap(); + writer + .write_column(1, Arc::new(Int32Array::from(vec![6, 7, 8, 9, 10]))) + .await + .unwrap(); + writer + .write_column(2, Arc::new(Int32Array::from(vec![100]))) + .await + .unwrap(); + writer.finish().await.unwrap(); + + let file_scheduler = fs + .scheduler + .open_file(&fs.tmp_path, &CachedFileSize::unknown()) + .await + .unwrap(); + let reader = FileReader::try_open( + file_scheduler, + None, + Arc::::default(), + &LanceCache::no_cache(), + FileReaderOptions::default(), + ) + .await + .unwrap(); + + let read = |names: &'static [&'static str], params: ReadBatchParams| { + let projection = + ReaderProjection::from_column_names(version, &lance_schema, names).unwrap(); + async { + match reader + .read_stream_projected( + params, + 1024, + 16, + projection, + FilterExpression::no_filter(), + ) + .await + { + Ok(stream) => stream.try_collect::>().await, + Err(e) => Err(e), + } + } + }; + let col_values = |batches: &[RecordBatch], idx: usize| -> Vec> { + batches + .iter() + .flat_map(|b| { + b.column(idx) + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .collect::>() + }) + .collect() + }; + + // Equal-length projection [a, b] full-scans into rectangular batches. + let batches = read(&["a", "b"], ReadBatchParams::RangeFull).await.unwrap(); + assert_eq!( + col_values(&batches, 0), + vec![Some(1), Some(2), Some(3), Some(4), Some(5)] + ); + assert_eq!( + col_values(&batches, 1), + vec![Some(6), Some(7), Some(8), Some(9), Some(10)] + ); + + // A mismatched-length projection [a, c] (5 vs 1) is rejected before any + // batch is yielded, regardless of the read params — its columns cannot + // be combined into rectangular batches. + assert!( + read(&["a", "c"], ReadBatchParams::RangeFull).await.is_err(), + "full scan across unequal-length columns must error" + ); + assert!( + read(&["a", "c"], ReadBatchParams::Range(0..1)) + .await + .is_err(), + "even a common-prefix read of unequal-length columns must error" + ); + + // A single-column RangeFull resolves to that column's own length. + let batches = read(&["c"], ReadBatchParams::RangeFull).await.unwrap(); + assert_eq!(col_values(&batches, 0), vec![Some(100)]); + let batches = read(&["a"], ReadBatchParams::RangeFull).await.unwrap(); + assert_eq!( + col_values(&batches, 0), + vec![Some(1), Some(2), Some(3), Some(4), Some(5)] + ); + } + /// Files written the ordinary (rectangular) way keep equal column lengths, /// so the unequal-length support is backwards compatible. #[tokio::test] From 34a19affaaae83aa67916d6252853f3c8336b958 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Mon, 22 Jun 2026 20:06:05 -0700 Subject: [PATCH 06/13] feat: data overlay file model, feature flag, and DataOverlay transaction MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds the in-memory + commit machinery for data overlay files (per the spec in #7381), the foundation the scanner/take/index/compaction work builds on. - `DataOverlayFile` / `OverlayCoverage` (dense `shared_offset_bitmap` and sparse per-field) with protobuf round-trip, attached to `Fragment.overlays`. - Reader feature flag 64 (`FLAG_DATA_OVERLAY_FILES`): set whenever any fragment carries overlays, so a reader that does not understand them refuses the dataset instead of returning stale base values. - `Operation::DataOverlay` transaction op: appends overlays to a fragment's list (preserving concurrently-written overlays) and stamps each overlay's `committed_version` to the new dataset version at commit time (re-stamped on retry). Conflict rules mirror DataReplacement — permissive against appends, deletes, column rewrites, index builds, and other overlays; conflicts only with row-rewriting compaction of the same fragment. Scan-side merge, take, and end-to-end write+read tests follow in the same PR branch. Part of the Data Overlay Files feature (OSS-1322). Co-Authored-By: Claude Opus 4.8 (1M context) --- rust/lance-table/src/feature_flags.rs | 22 +- rust/lance-table/src/format/fragment.rs | 222 +++++++++++++++++- rust/lance-table/src/format/manifest.rs | 2 + rust/lance/src/dataset/files.rs | 1 + rust/lance/src/dataset/optimize.rs | 1 + rust/lance/src/dataset/schema_evolution.rs | 1 + rust/lance/src/dataset/transaction.rs | 202 ++++++++++++++-- rust/lance/src/dataset/write.rs | 1 + rust/lance/src/dataset/write/commit.rs | 1 + rust/lance/src/io/commit.rs | 5 + rust/lance/src/io/commit/conflict_resolver.rs | 111 ++++++++- rust/lance/src/utils/test.rs | 1 + 12 files changed, 544 insertions(+), 26 deletions(-) diff --git a/rust/lance-table/src/feature_flags.rs b/rust/lance-table/src/feature_flags.rs index 096f0da79e5..1e0be5a3d06 100644 --- a/rust/lance-table/src/feature_flags.rs +++ b/rust/lance-table/src/feature_flags.rs @@ -20,8 +20,13 @@ pub const FLAG_TABLE_CONFIG: u64 = 8; pub const FLAG_BASE_PATHS: u64 = 16; /// Disable writing transaction file under _transaction/, this flag is set when we only want to write inline transaction in manifest pub const FLAG_DISABLE_TRANSACTION_FILE: u64 = 32; +/// Fragments contain data overlay files, which supply new values for a subset of +/// cells without rewriting base data files. A reader that does not understand +/// overlays must refuse the dataset, since ignoring an overlay would silently +/// return stale base values. +pub const FLAG_DATA_OVERLAY_FILES: u64 = 64; /// The first bit that is unknown as a feature flag -pub const FLAG_UNKNOWN: u64 = 64; +pub const FLAG_UNKNOWN: u64 = 128; /// Set the reader and writer feature flags in the manifest based on the contents of the manifest. pub fn apply_feature_flags( @@ -71,6 +76,18 @@ pub fn apply_feature_flags( manifest.writer_feature_flags |= FLAG_BASE_PATHS; } + // Overlay files change cell values on read, so a reader that ignores them + // would return stale base values. Both readers and writers must understand + // them. + let has_overlays = manifest + .fragments + .iter() + .any(|frag| !frag.overlays.is_empty()); + if has_overlays { + manifest.reader_feature_flags |= FLAG_DATA_OVERLAY_FILES; + manifest.writer_feature_flags |= FLAG_DATA_OVERLAY_FILES; + } + if disable_transaction_file { manifest.writer_feature_flags |= FLAG_DISABLE_TRANSACTION_FILE; } @@ -103,6 +120,7 @@ mod tests { assert!(can_read_dataset(super::FLAG_TABLE_CONFIG)); assert!(can_read_dataset(super::FLAG_BASE_PATHS)); assert!(can_read_dataset(super::FLAG_DISABLE_TRANSACTION_FILE)); + assert!(can_read_dataset(super::FLAG_DATA_OVERLAY_FILES)); assert!(can_read_dataset( super::FLAG_DELETION_FILES | super::FLAG_STABLE_ROW_IDS @@ -120,12 +138,14 @@ mod tests { assert!(can_write_dataset(super::FLAG_TABLE_CONFIG)); assert!(can_write_dataset(super::FLAG_BASE_PATHS)); assert!(can_write_dataset(super::FLAG_DISABLE_TRANSACTION_FILE)); + assert!(can_write_dataset(super::FLAG_DATA_OVERLAY_FILES)); assert!(can_write_dataset( super::FLAG_DELETION_FILES | super::FLAG_STABLE_ROW_IDS | super::FLAG_USE_V2_FORMAT_DEPRECATED | super::FLAG_TABLE_CONFIG | super::FLAG_BASE_PATHS + | super::FLAG_DATA_OVERLAY_FILES )); assert!(!can_write_dataset(super::FLAG_UNKNOWN)); } diff --git a/rust/lance-table/src/format/fragment.rs b/rust/lance-table/src/format/fragment.rs index e9d9ce036ee..6b0721cf75d 100644 --- a/rust/lance-table/src/format/fragment.rs +++ b/rust/lance-table/src/format/fragment.rs @@ -11,6 +11,7 @@ use lance_file::format::{MAJOR_VERSION, MINOR_VERSION}; use lance_file::version::LanceFileVersion; use lance_io::utils::CachedFileSize; use object_store::path::Path; +use roaring::RoaringBitmap; use serde::{Deserialize, Deserializer, Serialize, Serializer}; use crate::format::pb; @@ -233,6 +234,145 @@ impl TryFrom for DataFile { } } +/// Which `(physical offset, field)` cells a [`DataOverlayFile`] provides values +/// for. +/// +/// The coverage bitmaps index **physical** row offsets (positions in the base +/// data files, counting deleted rows), so they are stable across deletions, like +/// deletion vectors. Bitmaps are stored as serialized 32-bit Roaring bitmaps; use +/// [`DataOverlayFile::coverage_for_field`] to obtain the parsed bitmap that +/// applies to a given field. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, DeepSizeOf)] +pub enum OverlayCoverage { + /// A single bitmap that applies to every field in the overlay's + /// `data_file.fields` (a dense / rectangular overlay): every covered offset + /// has a value for every field. + Shared(Vec), + /// One bitmap per field, in the same order as the overlay's + /// `data_file.fields` (a sparse overlay): different fields may cover + /// different offset sets. + PerField(Vec>), +} + +fn deserialize_roaring(bytes: &[u8]) -> Result { + RoaringBitmap::deserialize_from(bytes).map_err(|e| { + Error::invalid_input(format!( + "failed to deserialize overlay coverage bitmap: {e}" + )) + }) +} + +fn serialize_roaring(bitmap: &RoaringBitmap) -> Vec { + let mut bytes = Vec::with_capacity(bitmap.serialized_size()); + // Writing to a Vec is infallible. + bitmap.serialize_into(&mut bytes).unwrap(); + bytes +} + +impl OverlayCoverage { + /// Build a dense coverage from a single bitmap. + pub fn dense(bitmap: &RoaringBitmap) -> Self { + Self::Shared(serialize_roaring(bitmap)) + } + + /// Build a sparse coverage from one bitmap per field. + pub fn sparse(bitmaps: &[RoaringBitmap]) -> Self { + Self::PerField(bitmaps.iter().map(serialize_roaring).collect()) + } +} + +/// An overlay file supplies new values for a subset of `(physical offset, field)` +/// cells within a fragment, without rewriting the fragment's base data files. See +/// the Data Overlay Files specification for the full resolution, coverage, and +/// versioning rules. +/// +/// The overlay's `data_file` stores one value column per field in +/// `data_file.fields`, with **no** row-offset key column. Within a value column, +/// the position of a covered offset's value is the **rank** (0-based count of set +/// bits below it) of that offset in the field's coverage bitmap. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, DeepSizeOf)] +pub struct DataOverlayFile { + /// The data file storing the overlay's new cell values. + pub data_file: DataFile, + /// Which cells this overlay provides values for. + pub coverage: OverlayCoverage, + /// The dataset version at which this overlay became effective (the version of + /// the commit that introduced it, stamped at commit time and re-stamped on + /// retry). Higher wins when two overlays cover the same `(offset, field)`. + pub committed_version: u64, +} + +impl DataOverlayFile { + /// The parsed coverage bitmap that applies to the field stored at + /// `field_pos` within `data_file.fields`. + /// + /// For a dense overlay the same shared bitmap is returned for every field; + /// for a sparse overlay the per-field bitmap at `field_pos` is returned. + pub fn coverage_for_field(&self, field_pos: usize) -> Result { + match &self.coverage { + OverlayCoverage::Shared(bytes) => deserialize_roaring(bytes), + OverlayCoverage::PerField(bitmaps) => { + let bytes = bitmaps.get(field_pos).ok_or_else(|| { + Error::invalid_input(format!( + "overlay field_coverage has {} bitmaps but field position {} was requested", + bitmaps.len(), + field_pos + )) + })?; + deserialize_roaring(bytes) + } + } + } +} + +impl From<&DataOverlayFile> for pb::DataOverlayFile { + fn from(overlay: &DataOverlayFile) -> Self { + let coverage = match &overlay.coverage { + OverlayCoverage::Shared(bytes) => { + pb::data_overlay_file::Coverage::SharedOffsetBitmap(bytes.clone()) + } + OverlayCoverage::PerField(bitmaps) => { + pb::data_overlay_file::Coverage::FieldCoverage(pb::FieldCoverage { + offset_bitmaps: bitmaps.clone(), + }) + } + }; + Self { + data_file: Some(pb::DataFile::from(&overlay.data_file)), + coverage: Some(coverage), + committed_version: overlay.committed_version, + } + } +} + +impl TryFrom for DataOverlayFile { + type Error = Error; + + fn try_from(proto: pb::DataOverlayFile) -> Result { + let data_file = proto + .data_file + .ok_or_else(|| Error::invalid_input("DataOverlayFile is missing its data_file"))?; + let coverage = match proto.coverage { + Some(pb::data_overlay_file::Coverage::SharedOffsetBitmap(bytes)) => { + OverlayCoverage::Shared(bytes) + } + Some(pb::data_overlay_file::Coverage::FieldCoverage(fc)) => { + OverlayCoverage::PerField(fc.offset_bitmaps) + } + None => { + return Err(Error::invalid_input( + "DataOverlayFile is missing its coverage", + )); + } + }; + Ok(Self { + data_file: DataFile::try_from(data_file)?, + coverage, + committed_version: proto.committed_version, + }) + } +} + /// Interns repeated data so that fragments with identical content share a /// single heap allocation via `Arc`. /// @@ -375,6 +515,11 @@ impl DataFileFieldInterner { .into_iter() .map(|f| self.intern_data_file(f)) .collect::>()?, + overlays: p + .overlays + .into_iter() + .map(DataOverlayFile::try_from) + .collect::>()?, deletion_file: p.deletion_file.map(DeletionFile::try_from).transpose()?, row_id_meta: p.row_id_sequence.map(RowIdMeta::try_from).transpose()?, physical_rows, @@ -483,6 +628,12 @@ pub struct Fragment { /// Files within the fragment. pub files: Vec, + /// Overlay files supplying new values for a subset of cells without + /// rewriting the base data files. Order is significant: a later entry is + /// newer than an earlier one. See [`DataOverlayFile`] for resolution rules. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub overlays: Vec, + /// Optional file with deleted local row offsets. #[serde(skip_serializing_if = "Option::is_none")] pub deletion_file: Option, @@ -510,6 +661,7 @@ impl Fragment { Self { id, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: None, @@ -549,6 +701,7 @@ impl Fragment { Self { id, files: vec![DataFile::new_legacy(path, schema, None, None)], + overlays: vec![], deletion_file: None, physical_rows, row_id_meta: None, @@ -669,6 +822,11 @@ impl TryFrom for Fragment { .into_iter() .map(DataFile::try_from) .collect::>()?, + overlays: p + .overlays + .into_iter() + .map(DataOverlayFile::try_from) + .collect::>()?, deletion_file: p.deletion_file.map(DeletionFile::try_from).transpose()?, row_id_meta: p.row_id_sequence.map(RowIdMeta::try_from).transpose()?, physical_rows, @@ -716,10 +874,7 @@ impl From<&Fragment> for pb::DataFragment { Self { id: f.id, files: f.files.iter().map(pb::DataFile::from).collect(), - // Overlay files are not produced by this version of the library; a - // dataset that uses them sets reader feature flag 64, which is - // rejected at the feature-flag layer (see lance-table feature_flags). - overlays: vec![], + overlays: f.overlays.iter().map(pb::DataOverlayFile::from).collect(), deletion_file, row_id_sequence, physical_rows: f.physical_rows.unwrap_or_default() as u64, @@ -738,6 +893,65 @@ mod tests { use object_store::path::Path; use serde_json::{Value, json}; + #[test] + fn test_data_overlay_roundtrip() { + // A fragment carrying a dense overlay round-trips through protobuf and + // back, and the parsed coverage bitmap is recovered per field. + let mut bitmap = RoaringBitmap::new(); + bitmap.insert(1); + bitmap.insert(3); + + let overlay = DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("overlay-0.lance", vec![3], None), + coverage: OverlayCoverage::dense(&bitmap), + committed_version: 7, + }; + let mut fragment = Fragment::new(0); + fragment.files = vec![DataFile::new_legacy_from_fields( + "base.lance", + vec![1, 3], + None, + )]; + fragment.overlays = vec![overlay]; + + let proto = pb::DataFragment::from(&fragment); + assert_eq!(proto.overlays.len(), 1); + let round_tripped = Fragment::try_from(proto).unwrap(); + assert_eq!(round_tripped, fragment); + + // Dense coverage applies to every field. + let recovered = round_tripped.overlays[0].coverage_for_field(0).unwrap(); + assert_eq!(recovered, bitmap); + assert_eq!( + round_tripped.overlays[0].coverage_for_field(5).unwrap(), + bitmap + ); + } + + #[test] + fn test_data_overlay_sparse_per_field_coverage() { + // A sparse overlay carries one bitmap per field, recovered by position. + let name_coverage = RoaringBitmap::from_iter([2u32, 3]); + let embedding_coverage = RoaringBitmap::from_iter([1u32]); + let overlay = DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("overlay-1.lance", vec![2, 4], None), + coverage: OverlayCoverage::sparse(&[name_coverage.clone(), embedding_coverage.clone()]), + committed_version: 3, + }; + let mut fragment = Fragment::new(1); + fragment.overlays = vec![overlay]; + + let round_tripped = Fragment::try_from(pb::DataFragment::from(&fragment)).unwrap(); + assert_eq!( + round_tripped.overlays[0].coverage_for_field(0).unwrap(), + name_coverage + ); + assert_eq!( + round_tripped.overlays[0].coverage_for_field(1).unwrap(), + embedding_coverage + ); + } + #[test] fn test_new_fragment() { let path = "foobar.lance"; diff --git a/rust/lance-table/src/format/manifest.rs b/rust/lance-table/src/format/manifest.rs index 9845061b7e4..cd0a403621f 100644 --- a/rust/lance-table/src/format/manifest.rs +++ b/rust/lance-table/src/format/manifest.rs @@ -1316,6 +1316,7 @@ mod tests { vec![0, 1, 2], None, )], + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: None, @@ -1328,6 +1329,7 @@ mod tests { DataFile::new_legacy_from_fields("path2", vec![0, 1, 43], None), DataFile::new_legacy_from_fields("path3", vec![2], None), ], + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: None, diff --git a/rust/lance/src/dataset/files.rs b/rust/lance/src/dataset/files.rs index 848add7e4a8..2214822ed90 100644 --- a/rust/lance/src/dataset/files.rs +++ b/rust/lance/src/dataset/files.rs @@ -1036,6 +1036,7 @@ mod tests { // No base_id -> falls back to the dataset base_uri. mk_file("c.lance", None), ], + overlays: vec![], // Deletion files also carry a base_id when they originate from a // shallow clone, and must resolve against base_paths too. deletion_file: Some(DeletionFile { diff --git a/rust/lance/src/dataset/optimize.rs b/rust/lance/src/dataset/optimize.rs index 87dda8e7e57..25168b54780 100644 --- a/rust/lance/src/dataset/optimize.rs +++ b/rust/lance/src/dataset/optimize.rs @@ -2125,6 +2125,7 @@ mod tests { let fragment = Fragment { id: 0, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: Some(0), diff --git a/rust/lance/src/dataset/schema_evolution.rs b/rust/lance/src/dataset/schema_evolution.rs index 5ef35a33ab7..8c1595c4daa 100644 --- a/rust/lance/src/dataset/schema_evolution.rs +++ b/rust/lance/src/dataset/schema_evolution.rs @@ -1941,6 +1941,7 @@ mod test { Ok(Some(Fragment { files: vec![], id: 0, + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: Some(50), diff --git a/rust/lance/src/dataset/transaction.rs b/rust/lance/src/dataset/transaction.rs index dcf0682812a..32d1833f9bb 100644 --- a/rust/lance/src/dataset/transaction.rs +++ b/rust/lance/src/dataset/transaction.rs @@ -31,8 +31,9 @@ use lance_table::feature_flags::{FLAG_STABLE_ROW_IDS, apply_feature_flags}; use lance_table::rowids::read_row_ids; use lance_table::{ format::{ - BasePath, DataFile, DataStorageFormat, Fragment, IndexFile, IndexMetadata, Manifest, - RowDatasetVersionMeta, RowDatasetVersionRun, RowDatasetVersionSequence, RowIdMeta, pb, + BasePath, DataFile, DataOverlayFile, DataStorageFormat, Fragment, IndexFile, IndexMetadata, + Manifest, RowDatasetVersionMeta, RowDatasetVersionRun, RowDatasetVersionSequence, + RowIdMeta, pb, }, io::{ commit::CommitHandler, @@ -258,6 +259,17 @@ pub struct Transaction { #[derive(Debug, Clone, DeepSizeOf, PartialEq)] pub struct DataReplacementGroup(pub u64, pub DataFile); +/// Overlay files to append to a single fragment, in order (the last entry is +/// newest). The overlays are appended to the fragment's existing `overlays` +/// list rather than replacing it, so overlays written by concurrent commits are +/// preserved. Each overlay's `committed_version` is stamped to the new dataset +/// version at commit time (re-stamped on retry). +#[derive(Debug, Clone, DeepSizeOf, PartialEq)] +pub struct DataOverlayGroup { + pub fragment_id: u64, + pub overlays: Vec, +} + /// An entry for a map update. If value is None, the key will be removed from the map. #[derive(Debug, Clone, DeepSizeOf, PartialEq)] pub struct UpdateMapEntry { @@ -367,6 +379,11 @@ pub enum Operation { DataReplacement { replacements: Vec, }, + /// Attach overlay files to fragments, supplying new values for a subset of + /// `(row offset, field)` cells without rewriting the fragments' base data + /// files. See [`DataOverlayFile`] and the Data Overlay Files specification + /// for resolution, coverage, and versioning rules. + DataOverlay { groups: Vec }, /// Merge a new column in /// 'fragments' is the final fragments include all data files, the new fragments must align with old ones at rows. /// 'schema' is not forced to include existed columns, which means we could use Merge to drop column data @@ -499,6 +516,7 @@ impl std::fmt::Display for Operation { Self::Project { .. } => write!(f, "Project"), Self::UpdateConfig { .. } => write!(f, "UpdateConfig"), Self::DataReplacement { .. } => write!(f, "DataReplacement"), + Self::DataOverlay { .. } => write!(f, "DataOverlay"), Self::Clone { .. } => write!(f, "Clone"), Self::UpdateMemWalState { .. } => write!(f, "UpdateMemWalState"), Self::UpdateBases { .. } => write!(f, "UpdateBases"), @@ -1345,6 +1363,16 @@ impl PartialEq for Operation { (Self::Clone { .. }, Self::UpdateBases { .. }) => { std::mem::discriminant(self) == std::mem::discriminant(other) } + // Data overlays are intentionally permissive, like DataReplacement. + // Two overlays stack (the higher committed_version wins each covered + // cell), and overlays are compatible with appends, deletes, column + // rewrites, and concurrent overlays. Only an operation that rewrites + // rows or otherwise invalidates physical offsets (Rewrite, which + // covers compaction and overlay->base folds) conflicts, since the + // overlay's physical offsets would no longer be valid. + (Self::DataOverlay { .. }, Self::Rewrite { .. }) + | (Self::Rewrite { .. }, Self::DataOverlay { .. }) => true, + (Self::DataOverlay { .. }, _) | (_, Self::DataOverlay { .. }) => false, } } } @@ -1521,6 +1549,7 @@ impl Operation { Self::Project { .. } => "Project", Self::UpdateConfig { .. } => "UpdateConfig", Self::DataReplacement { .. } => "DataReplacement", + Self::DataOverlay { .. } => "DataOverlay", Self::UpdateMemWalState { .. } => "UpdateMemWalState", Self::Clone { .. } => "Clone", Self::UpdateBases { .. } => "UpdateBases", @@ -2300,6 +2329,42 @@ impl Transaction { &replaced_fields, ); } + Operation::DataOverlay { groups } => { + // Stamp each overlay with the version this commit is producing. + // build_manifest re-runs on every retry with an updated + // current_manifest, so this is naturally re-stamped on retry. + let new_version = current_manifest.map_or(1, |m| m.version + 1); + + let existing_fragments = maybe_existing_fragments?; + let overlays_by_fragment: HashMap> = groups + .iter() + .map(|g| (g.fragment_id, &g.overlays)) + .collect(); + + // Every group must target an existing fragment. + for fragment_id in overlays_by_fragment.keys() { + if !existing_fragments.iter().any(|f| f.id == *fragment_id) { + return Err(Error::invalid_input(format!( + "DataOverlay targets fragment {fragment_id}, which does not exist" + ))); + } + } + + for fragment in existing_fragments { + let mut fragment = fragment.clone(); + if let Some(new_overlays) = overlays_by_fragment.get(&fragment.id) { + // Appended (not replaced) so concurrently-written overlays + // survive; later entries are newer. + fragment.overlays.extend(new_overlays.iter().cloned().map( + |mut overlay| { + overlay.committed_version = new_version; + overlay + }, + )); + } + final_fragments.push(fragment); + } + } Operation::UpdateMemWalState { merged_generations } => { update_mem_wal_index_merged_generations( &mut final_indices, @@ -3007,6 +3072,34 @@ impl TryFrom for DataReplacementGroup { } } +impl From<&DataOverlayGroup> for pb::transaction::DataOverlayGroup { + fn from(group: &DataOverlayGroup) -> Self { + Self { + fragment_id: group.fragment_id, + overlays: group + .overlays + .iter() + .map(pb::DataOverlayFile::from) + .collect(), + } + } +} + +impl TryFrom for DataOverlayGroup { + type Error = Error; + + fn try_from(message: pb::transaction::DataOverlayGroup) -> Result { + Ok(Self { + fragment_id: message.fragment_id, + overlays: message + .overlays + .into_iter() + .map(DataOverlayFile::try_from) + .collect::>>()?, + }) + } +} + impl TryFrom for Transaction { type Error = Error; @@ -3289,16 +3382,14 @@ impl TryFrom for Transaction { })) => Operation::UpdateBases { new_bases: new_bases.into_iter().map(BasePath::from).collect(), }, - Some(pb::transaction::Operation::DataOverlay(_)) => { - // Overlay files are not supported by this version of the library. - // A dataset that uses them sets reader feature flag 64, which is - // already rejected at the feature-flag layer; reject here too so a - // transaction referencing the operation can never be applied. - return Err(Error::not_supported( - "data overlay files are not supported by this version of Lance \ - (reader feature flag 64)", - )); - } + Some(pb::transaction::Operation::DataOverlay(pb::transaction::DataOverlay { + groups, + })) => Operation::DataOverlay { + groups: groups + .into_iter() + .map(DataOverlayGroup::try_from) + .collect::>>()?, + }, None => { return Err(Error::internal( "Transaction message did not contain an operation".to_string(), @@ -3569,6 +3660,14 @@ impl From<&Transaction> for pb::Transaction { .collect(), }) } + Operation::DataOverlay { groups } => { + pb::transaction::Operation::DataOverlay(pb::transaction::DataOverlay { + groups: groups + .iter() + .map(pb::transaction::DataOverlayGroup::from) + .collect(), + }) + } Operation::UpdateMemWalState { merged_generations } => { pb::transaction::Operation::UpdateMemWalState(pb::transaction::UpdateMemWalState { merged_generations: merged_generations @@ -4158,6 +4257,7 @@ mod tests { physical_rows: Some(100), row_id_meta: None, files: vec![], + overlays: vec![], deletion_file: None, last_updated_at_version_meta: None, created_at_version_meta: None, @@ -4190,6 +4290,7 @@ mod tests { physical_rows: Some(50), row_id_meta: Some(RowIdMeta::Inline(serialized)), files: vec![], + overlays: vec![], deletion_file: None, last_updated_at_version_meta: None, created_at_version_meta: None, @@ -4222,6 +4323,7 @@ mod tests { physical_rows: Some(50), // More physical rows than existing row IDs row_id_meta: Some(RowIdMeta::Inline(serialized)), files: vec![], + overlays: vec![], deletion_file: None, last_updated_at_version_meta: None, created_at_version_meta: None, @@ -4257,6 +4359,7 @@ mod tests { physical_rows: Some(50), // Less physical rows than existing row IDs row_id_meta: Some(RowIdMeta::Inline(serialized)), files: vec![], + overlays: vec![], deletion_file: None, last_updated_at_version_meta: None, created_at_version_meta: None, @@ -4285,6 +4388,7 @@ mod tests { physical_rows: Some(30), // No existing row IDs row_id_meta: None, files: vec![], + overlays: vec![], deletion_file: None, last_updated_at_version_meta: None, created_at_version_meta: None, @@ -4294,6 +4398,7 @@ mod tests { physical_rows: Some(25), // Partial existing row IDs row_id_meta: Some(RowIdMeta::Inline(serialized)), files: vec![], + overlays: vec![], deletion_file: None, last_updated_at_version_meta: None, created_at_version_meta: None, @@ -4338,6 +4443,7 @@ mod tests { physical_rows: None, row_id_meta: None, files: vec![], + overlays: vec![], deletion_file: None, last_updated_at_version_meta: None, created_at_version_meta: None, @@ -4845,6 +4951,7 @@ mod tests { let fragment = Fragment { id: 1, files: vec![data_file], + overlays: vec![], deletion_file: None, row_id_meta, physical_rows: Some(5), @@ -5114,6 +5221,7 @@ mod tests { None, )], physical_rows: Some(10), + overlays: vec![], deletion_file: None, row_id_meta: None, last_updated_at_version_meta: None, @@ -5205,6 +5313,7 @@ mod tests { let prev_fragment = Fragment { id: 0, files: vec![mk_file("before.lance")], + overlays: vec![], deletion_file: None, row_id_meta, physical_rows: Some(5), @@ -5277,6 +5386,7 @@ mod tests { let prev_fragment = Fragment { id: 0, files: vec![data_file.clone()], + overlays: vec![], deletion_file: None, row_id_meta: row_id_meta.clone(), physical_rows: Some(5), @@ -5296,6 +5406,7 @@ mod tests { let merged_fragment = Fragment { id: 0, files: vec![data_file], + overlays: vec![], deletion_file: None, row_id_meta, physical_rows: Some(5), @@ -5345,6 +5456,7 @@ mod tests { let prev_fragment = Fragment { id: 0, files: vec![mk_file("before.lance")], + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: Some(5), @@ -5409,6 +5521,7 @@ mod tests { let existing_fragment = Fragment { id: 0, files: vec![mk_file("existing.lance")], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&row_ids_0))), physical_rows: Some(3), @@ -5431,6 +5544,7 @@ mod tests { let new_fragment = Fragment { id: 1, files: vec![mk_file("new.lance")], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&row_ids_1))), physical_rows: Some(4), @@ -5493,6 +5607,7 @@ mod tests { let existing_fragment = Fragment { id: 1, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&existing_seq))), physical_rows: Some(3), @@ -5506,6 +5621,7 @@ mod tests { let new_fragment = Fragment { id: 10, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&new_seq))), physical_rows: Some(2), @@ -5548,6 +5664,7 @@ mod tests { Fragment { id: 1, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&frag_a_seq))), physical_rows: Some(2), @@ -5559,6 +5676,7 @@ mod tests { Fragment { id: 2, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&frag_b_seq))), physical_rows: Some(3), @@ -5574,6 +5692,7 @@ mod tests { let new_fragment = Fragment { id: 10, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&new_seq))), physical_rows: Some(2), @@ -5615,6 +5734,7 @@ mod tests { let existing_fragment = Fragment { id: 1, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&existing_seq))), physical_rows: Some(2), @@ -5629,6 +5749,7 @@ mod tests { let new_fragment = Fragment { id: 10, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&new_seq))), physical_rows: Some(2), @@ -5673,6 +5794,7 @@ mod tests { let existing_fragment = Fragment { id: 1, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&existing_seq))), physical_rows: Some(2), @@ -5686,6 +5808,7 @@ mod tests { let new_fragment = Fragment { id: 20, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&new_seq))), physical_rows: Some(4), @@ -5719,6 +5842,7 @@ mod tests { let existing_fragment = Fragment { id: 1, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&existing_seq))), physical_rows: Some(2), @@ -5730,6 +5854,7 @@ mod tests { let new_fragment = Fragment { id: 10, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&new_seq))), physical_rows: Some(1), @@ -5758,6 +5883,7 @@ mod tests { let existing_fragment = Fragment { id: 1, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&existing_seq))), physical_rows: Some(2), @@ -5768,6 +5894,7 @@ mod tests { let new_fragment = Fragment { id: 10, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: Some(3), @@ -5798,6 +5925,7 @@ mod tests { let existing_fragment = Fragment { id: 1, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&existing_seq))), physical_rows: Some(2), @@ -5811,6 +5939,7 @@ mod tests { let new_fragment = Fragment { id: 10, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&new_seq))), physical_rows: Some(1), @@ -5852,6 +5981,7 @@ mod tests { let in_range_frag = Fragment { id: 1, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&in_range_seq))), physical_rows: Some(2), @@ -5872,6 +6002,7 @@ mod tests { let out_of_range_frag = Fragment { id: 2, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&out_of_range_seq))), physical_rows: Some(2), @@ -5886,6 +6017,7 @@ mod tests { let new_frag = Fragment { id: 10, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&new_seq))), physical_rows: Some(2), @@ -5924,6 +6056,7 @@ mod tests { let existing = Fragment { id: 1, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&seq))), physical_rows: Some(3), @@ -5936,6 +6069,7 @@ mod tests { let new_frag = Fragment { id: 10, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&new_seq))), physical_rows: Some(2), @@ -5984,6 +6118,7 @@ mod tests { let src_frag = Fragment { id: 1, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&src_seq))), physical_rows: Some(100), @@ -5998,6 +6133,7 @@ mod tests { let new_frag = Fragment { id: 10, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&new_seq))), physical_rows: Some(100), @@ -6045,6 +6181,7 @@ mod tests { Fragment { id: 1, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&seq_a))), physical_rows: Some(3), @@ -6056,6 +6193,7 @@ mod tests { Fragment { id: 2, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&seq_b))), physical_rows: Some(3), @@ -6071,6 +6209,7 @@ mod tests { let new_frag = Fragment { id: 10, files: vec![], + overlays: vec![], deletion_file: None, row_id_meta: Some(RowIdMeta::Inline(write_row_ids(&new_seq))), physical_rows: Some(2), @@ -6139,20 +6278,45 @@ mod tests { } #[test] - fn test_data_overlay_operation_rejected() { - // Overlay files are not supported by this version of the library. A - // transaction carrying the DataOverlay operation must be rejected rather - // than silently ignored, mirroring the feature-flag-64 rejection. + fn test_data_overlay_operation_roundtrips() { + // A DataOverlay operation survives the protobuf round-trip, preserving + // the target fragment, the overlay's coverage, and its committed_version. + use lance_table::format::{DataOverlayFile, OverlayCoverage}; + + let mut bitmap = roaring::RoaringBitmap::new(); + bitmap.insert(1); + bitmap.insert(4); + let overlay = DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("overlay-0.lance", vec![3], None), + coverage: OverlayCoverage::dense(&bitmap), + committed_version: 6, + }; + let pb_overlay = pb::DataOverlayFile::from(&overlay); + let message = pb::Transaction { read_version: 1, uuid: Uuid::new_v4().to_string(), operation: Some(pb::transaction::Operation::DataOverlay( - pb::transaction::DataOverlay { groups: vec![] }, + pb::transaction::DataOverlay { + groups: vec![pb::transaction::DataOverlayGroup { + fragment_id: 7, + overlays: vec![pb_overlay], + }], + }, )), ..Default::default() }; - let result = Transaction::try_from(message); - assert!(matches!(result, Err(Error::NotSupported { .. }))); + let txn = Transaction::try_from(message).unwrap(); + match txn.operation { + Operation::DataOverlay { groups } => { + assert_eq!(groups.len(), 1); + assert_eq!(groups[0].fragment_id, 7); + assert_eq!(groups[0].overlays.len(), 1); + assert_eq!(groups[0].overlays[0].committed_version, 6); + assert_eq!(groups[0].overlays[0].coverage_for_field(0).unwrap(), bitmap); + } + other => panic!("expected DataOverlay, got {other:?}"), + } } } diff --git a/rust/lance/src/dataset/write.rs b/rust/lance/src/dataset/write.rs index ff0a119158c..47c4add18bd 100644 --- a/rust/lance/src/dataset/write.rs +++ b/rust/lance/src/dataset/write.rs @@ -3562,6 +3562,7 @@ mod tests { let fragments = vec![Fragment { id: 0, files: vec![external_file, local_file], + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: Some(0), diff --git a/rust/lance/src/dataset/write/commit.rs b/rust/lance/src/dataset/write/commit.rs index baad71b3e39..7aed72a2d71 100644 --- a/rust/lance/src/dataset/write/commit.rs +++ b/rust/lance/src/dataset/write/commit.rs @@ -551,6 +551,7 @@ mod tests { file_size_bytes: CachedFileSize::new(100), base_id: None, }], + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: Some(10), diff --git a/rust/lance/src/io/commit.rs b/rust/lance/src/io/commit.rs index ce0d29d550b..8980dd4f1ed 100644 --- a/rust/lance/src/io/commit.rs +++ b/rust/lance/src/io/commit.rs @@ -1687,6 +1687,7 @@ mod tests { DataFile::new_legacy_from_fields("path1", vec![0, 1, 2], None), DataFile::new_legacy_from_fields("unused", vec![9], None), ], + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: None, @@ -1699,6 +1700,7 @@ mod tests { DataFile::new_legacy_from_fields("path2", vec![0, 1, 2], None), DataFile::new_legacy_from_fields("path3", vec![2], None), ], + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: None, @@ -1736,6 +1738,7 @@ mod tests { vec![0, 1, 10], None, )], + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: None, @@ -1748,6 +1751,7 @@ mod tests { DataFile::new_legacy_from_fields("path2", vec![0, 1, 2], None), DataFile::new_legacy_from_fields("path3", vec![10], None), ], + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: None, @@ -1838,6 +1842,7 @@ mod tests { let fragment = Fragment { id: 0, files: vec![data_file], + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: Some(100), diff --git a/rust/lance/src/io/commit/conflict_resolver.rs b/rust/lance/src/io/commit/conflict_resolver.rs index dc898534c89..1af51506f8e 100644 --- a/rust/lance/src/io/commit/conflict_resolver.rs +++ b/rust/lance/src/io/commit/conflict_resolver.rs @@ -137,6 +137,21 @@ impl<'a> TransactionRebase<'a> { conflicting_mem_wal_merged_gens: Vec::new(), }) } + Operation::DataOverlay { groups } => { + let modified_fragment_ids = + groups.iter().map(|g| g.fragment_id).collect::>(); + let initial_fragments = + initial_fragments_for_rebase(dataset, &transaction, &modified_fragment_ids) + .await; + Ok(Self { + transaction, + affected_rows, + initial_fragments, + modified_fragment_ids, + conflicting_frag_reuse_indices: Vec::new(), + conflicting_mem_wal_merged_gens: Vec::new(), + }) + } Operation::Merge { fragments, .. } => { let modified_fragment_ids = fragments.iter().map(|f| f.id).collect::>(); let initial_fragments = @@ -203,6 +218,9 @@ impl<'a> TransactionRebase<'a> { Operation::DataReplacement { .. } => { self.check_data_replacement_txn(other_transaction, other_version) } + Operation::DataOverlay { .. } => { + self.check_data_overlay_txn(other_transaction, other_version) + } Operation::Merge { .. } => self.check_merge_txn(other_transaction, other_version), Operation::Restore { .. } => self.check_restore_txn(other_transaction, other_version), Operation::ReserveFragments { .. } => { @@ -235,6 +253,10 @@ impl<'a> TransactionRebase<'a> { | Operation::Project { .. } | Operation::Append { .. } | Operation::UpdateConfig { .. } + // A concurrent overlay is inert against the rows we delete + // (deletions take precedence over overlays) and otherwise + // preserves physical offsets, so it never conflicts. + | Operation::DataOverlay { .. } | Operation::UpdateBases { .. } => Ok(()), Operation::Rewrite { groups, .. } => { if groups @@ -382,6 +404,9 @@ impl<'a> TransactionRebase<'a> { | Operation::Project { .. } | Operation::Clone { .. } | Operation::UpdateConfig { .. } + // A concurrent overlay preserves physical offsets and is newer + // than this update, so it wins its covered cells without conflict. + | Operation::DataOverlay { .. } | Operation::UpdateBases { .. } => Ok(()), Operation::Append { .. } => { // If current transaction has primary key conflict detection, @@ -498,6 +523,10 @@ impl<'a> TransactionRebase<'a> { match &other_transaction.operation { Operation::Append { .. } | Operation::Clone { .. } + // An overlay committed after this index's version is newer than + // the index; the query path excludes its covered cells via the + // version gate, so the build does not conflict. + | Operation::DataOverlay { .. } | Operation::UpdateBases { .. } => Ok(()), Operation::CreateIndex { new_indices: created_indices, @@ -679,6 +708,20 @@ impl<'a> TransactionRebase<'a> { Ok(()) } } + Operation::DataOverlay { groups } => { + // Rewriting a fragment changes its physical row addresses, so + // an overlay addressed by physical offset on that fragment is + // invalidated and must be re-applied against the new base. + if groups + .iter() + .map(|g| g.fragment_id) + .any(|id| self.modified_fragment_ids.contains(&id)) + { + Err(self.retryable_conflict_err(other_transaction, other_version)) + } else { + Ok(()) + } + } Operation::Rewrite { groups, frag_reuse_index: committed_fri, @@ -858,6 +901,7 @@ impl<'a> TransactionRebase<'a> { | Operation::CreateIndex { .. } | Operation::Rewrite { .. } | Operation::DataReplacement { .. } + | Operation::DataOverlay { .. } | Operation::Merge { .. } | Operation::Restore { .. } | Operation::ReserveFragments { .. } @@ -891,7 +935,8 @@ impl<'a> TransactionRebase<'a> { | Operation::Merge { .. } | Operation::UpdateConfig { .. } | Operation::Clone { .. } - | Operation::DataReplacement { .. } => Ok(()), + | Operation::DataReplacement { .. } + | Operation::DataOverlay { .. } => Ok(()), } } @@ -907,6 +952,9 @@ impl<'a> TransactionRebase<'a> { | Operation::UpdateConfig { .. } | Operation::ReserveFragments { .. } | Operation::Project { .. } + // Both a column replacement and an overlay preserve physical row + // addresses; the overlay is newer and wins its covered cells. + | Operation::DataOverlay { .. } | Operation::UpdateBases { .. } => Ok(()), Operation::Merge { .. } => { // Merge rewrites the whole fragment list; always conflict @@ -1009,6 +1057,57 @@ impl<'a> TransactionRebase<'a> { } } + /// Conflict checks for our DataOverlay transaction against a concurrent one. + /// + /// Overlays are intentionally permissive (see the Data Overlay Files spec): + /// they stack with other overlays and tolerate appends, deletes, column + /// rewrites, and index builds, because overlay coverage is addressed by + /// physical offset and the version gate keeps indexes correct. The only + /// concurrent operations that invalidate an overlay are those that rewrite + /// rows or consume the overlays on one of our fragments (Rewrite / Merge), + /// and the whole-dataset replacements (Overwrite / Restore). + fn check_data_overlay_txn( + &mut self, + other_transaction: &Transaction, + other_version: u64, + ) -> Result<()> { + match &other_transaction.operation { + Operation::Append { .. } + | Operation::Delete { .. } + | Operation::Update { .. } + | Operation::CreateIndex { .. } + | Operation::ReserveFragments { .. } + | Operation::Project { .. } + | Operation::UpdateConfig { .. } + | Operation::UpdateBases { .. } + | Operation::Clone { .. } + | Operation::UpdateMemWalState { .. } + | Operation::DataReplacement { .. } + | Operation::DataOverlay { .. } => Ok(()), + Operation::Rewrite { groups, .. } => { + // A rewrite (compaction / fold) of a fragment we are overlaying + // changes its physical row addresses, so our offsets would be + // invalid. Conflict only if it touches one of our fragments. + let touches_our_fragment = groups + .iter() + .flat_map(|g| g.old_fragments.iter()) + .any(|f| self.modified_fragment_ids.contains(&f.id)); + if touches_our_fragment { + Err(self.retryable_conflict_err(other_transaction, other_version)) + } else { + Ok(()) + } + } + Operation::Merge { .. } => { + // Merge rewrites the whole fragment list; always conflict. + Err(self.retryable_conflict_err(other_transaction, other_version)) + } + Operation::Overwrite { .. } | Operation::Restore { .. } => { + Err(self.incompatible_conflict_err(other_transaction, other_version)) + } + } + } + fn check_merge_txn( &mut self, other_transaction: &Transaction, @@ -1026,7 +1125,8 @@ impl<'a> TransactionRebase<'a> { | Operation::Delete { .. } | Operation::Rewrite { .. } | Operation::Merge { .. } - | Operation::DataReplacement { .. } => { + | Operation::DataReplacement { .. } + | Operation::DataOverlay { .. } => { Err(self.retryable_conflict_err(other_transaction, other_version)) } Operation::Overwrite { .. } @@ -1050,6 +1150,7 @@ impl<'a> TransactionRebase<'a> { | Operation::CreateIndex { .. } | Operation::Rewrite { .. } | Operation::DataReplacement { .. } + | Operation::DataOverlay { .. } | Operation::Merge { .. } | Operation::Restore { .. } | Operation::ReserveFragments { .. } @@ -1078,6 +1179,7 @@ impl<'a> TransactionRebase<'a> { | Operation::CreateIndex { .. } | Operation::Rewrite { .. } | Operation::DataReplacement { .. } + | Operation::DataOverlay { .. } | Operation::Merge { .. } | Operation::ReserveFragments { .. } | Operation::Update { .. } @@ -1102,6 +1204,7 @@ impl<'a> TransactionRebase<'a> { | Operation::UpdateConfig { .. } | Operation::CreateIndex { .. } | Operation::DataReplacement { .. } + | Operation::DataOverlay { .. } | Operation::Rewrite { .. } | Operation::Clone { .. } | Operation::ReserveFragments { .. } @@ -1166,6 +1269,7 @@ impl<'a> TransactionRebase<'a> { | Operation::CreateIndex { .. } | Operation::Rewrite { .. } | Operation::DataReplacement { .. } + | Operation::DataOverlay { .. } | Operation::Merge { .. } | Operation::Restore { .. } | Operation::ReserveFragments { .. } @@ -1238,6 +1342,7 @@ impl<'a> TransactionRebase<'a> { | Operation::Overwrite { .. } | Operation::Delete { .. } | Operation::DataReplacement { .. } + | Operation::DataOverlay { .. } | Operation::Merge { .. } | Operation::Restore { .. } | Operation::Clone { .. } @@ -1331,6 +1436,7 @@ impl<'a> TransactionRebase<'a> { Operation::Append { .. } | Operation::Overwrite { .. } | Operation::DataReplacement { .. } + | Operation::DataOverlay { .. } | Operation::Merge { .. } | Operation::Restore { .. } | Operation::ReserveFragments { .. } @@ -3208,6 +3314,7 @@ mod tests { Operation::DataReplacement { replacements } => { Box::new(replacements.iter().map(|r| r.0)) } + Operation::DataOverlay { groups } => Box::new(groups.iter().map(|g| g.fragment_id)), } } diff --git a/rust/lance/src/utils/test.rs b/rust/lance/src/utils/test.rs index 3338eee07a8..f804a7cc38a 100644 --- a/rust/lance/src/utils/test.rs +++ b/rust/lance/src/utils/test.rs @@ -243,6 +243,7 @@ impl TestDatasetGenerator { Fragment { id: 0, files, + overlays: vec![], deletion_file: None, row_id_meta: None, physical_rows: Some(batch.num_rows()), From 107587ef686e21a689ee5c58718cde4d7cfe08fa Mon Sep 17 00:00:00 2001 From: Will Jones Date: Mon, 22 Jun 2026 20:09:39 -0700 Subject: [PATCH 07/13] feat: refuse reads of fragments with overlays until merge lands Now that overlays can be committed, a scan or take over a fragment that has overlays would silently return stale base values, since the read-path merge is not implemented yet. Refuse such reads at `FileFragment::open` with a clear error instead of serving incorrect data. Lifted once the scan/take merge lands (rest of OSS-1322 / OSS-1324). Co-Authored-By: Claude Opus 4.8 (1M context) --- rust/lance/src/dataset/fragment.rs | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/rust/lance/src/dataset/fragment.rs b/rust/lance/src/dataset/fragment.rs index eb165e5f612..58a3ce0a9b6 100644 --- a/rust/lance/src/dataset/fragment.rs +++ b/rust/lance/src/dataset/fragment.rs @@ -900,6 +900,16 @@ impl FileFragment { projection: &Schema, read_config: FragReadConfig, ) -> Result { + // Overlay files supply newer cell values that must be merged on read. + // Until the scan/take merge path lands (the rest of OSS-1322 / OSS-1324), + // reading a fragment that has overlays would silently return stale base + // values, so we refuse rather than serve incorrect data. + if !self.metadata.overlays.is_empty() { + return Err(Error::not_supported( + "reading fragments with data overlay files is not yet supported \ + (overlay merge is in progress)", + )); + } let open_files = self.open_readers(projection, &read_config); let deletion_vec_load = self.get_deletion_vector(); From a1157a907456f47f20b0ad1e4d660af5d8a1aa05 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Tue, 23 Jun 2026 17:17:38 -0700 Subject: [PATCH 08/13] refactor(table): parse overlay coverage on load and gate the overlay flag OverlayCoverage now holds parsed Arcs, deserialized once when a fragment is loaded and shared cheaply on clone, instead of re-deserializing the serialized bytes on every coverage_for_field call. Treat the data overlay feature flag (64) as unknown in release builds unless LANCE_ENABLE_DATA_OVERLAY_FILES is set, so the unreleased feature is neither emitted by release writers nor accepted by release readers; debug builds understand it so tests exercise the path. Stable-sort a fragment's overlays by committed_version (newest last) on load so resolution can assume the ordering. Adds tests for the missing-coverage/data_file errors, the release gating, and the load-time sort. Co-Authored-By: Claude Opus 4.8 (1M context) --- rust/lance-table/src/feature_flags.rs | 47 +++++- rust/lance-table/src/format/fragment.rs | 210 +++++++++++++++++++----- rust/lance/src/dataset/transaction.rs | 7 +- 3 files changed, 222 insertions(+), 42 deletions(-) diff --git a/rust/lance-table/src/feature_flags.rs b/rust/lance-table/src/feature_flags.rs index 1e0be5a3d06..40cc38aeeda 100644 --- a/rust/lance-table/src/feature_flags.rs +++ b/rust/lance-table/src/feature_flags.rs @@ -24,10 +24,19 @@ pub const FLAG_DISABLE_TRANSACTION_FILE: u64 = 32; /// cells without rewriting base data files. A reader that does not understand /// overlays must refuse the dataset, since ignoring an overlay would silently /// return stale base values. +/// +/// Data overlay files are not yet a released feature: in release builds this flag +/// is treated as unknown (so a release reader/writer refuses an overlay dataset) +/// unless [`ENABLE_DATA_OVERLAY_FILES_ENV`] is set, which lets benchmarks opt in. +/// Debug builds always understand it so tests exercise the path. pub const FLAG_DATA_OVERLAY_FILES: u64 = 64; /// The first bit that is unknown as a feature flag pub const FLAG_UNKNOWN: u64 = 128; +/// Environment variable that opts a release build into reading and writing data +/// overlay files before the feature is generally released. +pub const ENABLE_DATA_OVERLAY_FILES_ENV: &str = "LANCE_ENABLE_DATA_OVERLAY_FILES"; + /// Set the reader and writer feature flags in the manifest based on the contents of the manifest. pub fn apply_feature_flags( manifest: &mut Manifest, @@ -94,12 +103,33 @@ pub fn apply_feature_flags( Ok(()) } +/// Whether this build understands data overlay files: always in debug builds, +/// and in release builds only when [`ENABLE_DATA_OVERLAY_FILES_ENV`] is set. +fn data_overlay_files_enabled() -> bool { + cfg!(debug_assertions) || std::env::var_os(ENABLE_DATA_OVERLAY_FILES_ENV).is_some() +} + +/// The feature-flag bits this build understands, given whether overlay support +/// is enabled. Split out from [`supported_flags`] so the policy is testable +/// without toggling the build profile or environment. +fn supported_flags_when(overlay_enabled: bool) -> u64 { + let mut supported = FLAG_UNKNOWN - 1; + if !overlay_enabled { + supported &= !FLAG_DATA_OVERLAY_FILES; + } + supported +} + +fn supported_flags() -> u64 { + supported_flags_when(data_overlay_files_enabled()) +} + pub fn can_read_dataset(reader_flags: u64) -> bool { - reader_flags < FLAG_UNKNOWN + reader_flags & !supported_flags() == 0 } pub fn can_write_dataset(writer_flags: u64) -> bool { - writer_flags < FLAG_UNKNOWN + writer_flags & !supported_flags() == 0 } pub fn has_deprecated_v2_feature_flag(writer_flags: u64) -> bool { @@ -129,6 +159,19 @@ mod tests { assert!(!can_read_dataset(super::FLAG_UNKNOWN)); } + #[test] + fn test_data_overlay_flag_release_gating() { + // Release default (overlays disabled): the overlay flag is treated as + // unknown so the dataset is refused, while other known flags still pass. + let supported = supported_flags_when(false); + assert_eq!(supported & FLAG_DATA_OVERLAY_FILES, 0); + assert_eq!(FLAG_DELETION_FILES & !supported, 0); + assert_ne!(FLAG_DATA_OVERLAY_FILES & !supported, 0); + // Enabled (debug or env opt-in): the overlay flag is understood. + let supported = supported_flags_when(true); + assert_eq!(FLAG_DATA_OVERLAY_FILES & !supported, 0); + } + #[test] fn test_write_check() { assert!(can_write_dataset(0)); diff --git a/rust/lance-table/src/format/fragment.rs b/rust/lance-table/src/format/fragment.rs index 6b0721cf75d..c23f312d3ce 100644 --- a/rust/lance-table/src/format/fragment.rs +++ b/rust/lance-table/src/format/fragment.rs @@ -239,18 +239,28 @@ impl TryFrom for DataFile { /// /// The coverage bitmaps index **physical** row offsets (positions in the base /// data files, counting deleted rows), so they are stable across deletions, like -/// deletion vectors. Bitmaps are stored as serialized 32-bit Roaring bitmaps; use -/// [`DataOverlayFile::coverage_for_field`] to obtain the parsed bitmap that +/// deletion vectors. Bitmaps are parsed from their 32-bit Roaring encoding once +/// when the fragment is loaded and held behind an `Arc` so cloning a fragment is +/// cheap; use [`DataOverlayFile::coverage_for_field`] to obtain the one that /// applies to a given field. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, DeepSizeOf)] +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(into = "OverlayCoverageBytes", try_from = "OverlayCoverageBytes")] pub enum OverlayCoverage { /// A single bitmap that applies to every field in the overlay's /// `data_file.fields` (a dense / rectangular overlay): every covered offset /// has a value for every field. - Shared(Vec), + Shared(Arc), /// One bitmap per field, in the same order as the overlay's /// `data_file.fields` (a sparse overlay): different fields may cover /// different offset sets. + PerField(Vec>), +} + +/// Serialized form of [`OverlayCoverage`] — each bitmap as its 32-bit Roaring +/// byte encoding. The in-memory form parses these once at load. +#[derive(Debug, Clone, Serialize, Deserialize)] +enum OverlayCoverageBytes { + Shared(Vec), PerField(Vec>), } @@ -269,15 +279,58 @@ fn serialize_roaring(bitmap: &RoaringBitmap) -> Vec { bytes } +impl From for OverlayCoverageBytes { + fn from(coverage: OverlayCoverage) -> Self { + match coverage { + OverlayCoverage::Shared(bitmap) => Self::Shared(serialize_roaring(&bitmap)), + OverlayCoverage::PerField(bitmaps) => { + Self::PerField(bitmaps.iter().map(|b| serialize_roaring(b)).collect()) + } + } + } +} + +impl TryFrom for OverlayCoverage { + type Error = Error; + + fn try_from(bytes: OverlayCoverageBytes) -> Result { + Ok(match bytes { + OverlayCoverageBytes::Shared(b) => Self::Shared(Arc::new(deserialize_roaring(&b)?)), + OverlayCoverageBytes::PerField(bs) => Self::PerField( + bs.iter() + .map(|b| deserialize_roaring(b).map(Arc::new)) + .collect::>()?, + ), + }) + } +} + +impl DeepSizeOf for OverlayCoverage { + fn deep_size_of_children(&self, _context: &mut lance_core::deepsize::Context) -> usize { + // RoaringBitmap does not expose its allocation size; its serialized size + // is a cheap, close proxy for the heap it holds. + let bitmap_heap = |bitmap: &RoaringBitmap| { + std::mem::size_of::() + bitmap.serialized_size() + }; + match self { + Self::Shared(bitmap) => bitmap_heap(bitmap), + Self::PerField(bitmaps) => { + bitmaps.capacity() * std::mem::size_of::>() + + bitmaps.iter().map(|b| bitmap_heap(b)).sum::() + } + } + } +} + impl OverlayCoverage { - /// Build a dense coverage from a single bitmap. - pub fn dense(bitmap: &RoaringBitmap) -> Self { - Self::Shared(serialize_roaring(bitmap)) + /// Build a dense coverage from a single bitmap shared across every field. + pub fn dense(bitmap: RoaringBitmap) -> Self { + Self::Shared(Arc::new(bitmap)) } /// Build a sparse coverage from one bitmap per field. - pub fn sparse(bitmaps: &[RoaringBitmap]) -> Self { - Self::PerField(bitmaps.iter().map(serialize_roaring).collect()) + pub fn sparse(bitmaps: Vec) -> Self { + Self::PerField(bitmaps.into_iter().map(Arc::new).collect()) } } @@ -307,33 +360,42 @@ impl DataOverlayFile { /// `field_pos` within `data_file.fields`. /// /// For a dense overlay the same shared bitmap is returned for every field; - /// for a sparse overlay the per-field bitmap at `field_pos` is returned. - pub fn coverage_for_field(&self, field_pos: usize) -> Result { + /// for a sparse overlay the per-field bitmap at `field_pos` is returned. The + /// bitmap is already parsed, so this is a cheap `Arc` clone. + pub fn coverage_for_field(&self, field_pos: usize) -> Result> { match &self.coverage { - OverlayCoverage::Shared(bytes) => deserialize_roaring(bytes), + OverlayCoverage::Shared(bitmap) => Ok(bitmap.clone()), OverlayCoverage::PerField(bitmaps) => { - let bytes = bitmaps.get(field_pos).ok_or_else(|| { + bitmaps.get(field_pos).cloned().ok_or_else(|| { Error::invalid_input(format!( "overlay field_coverage has {} bitmaps but field position {} was requested", bitmaps.len(), field_pos )) - })?; - deserialize_roaring(bytes) + }) } } } } +/// Overlays are stored newest-last: a later list position is newer, and ties in +/// `committed_version` are broken by position. Loading stable-sorts by +/// `committed_version` so resolution can rely on the ordering without +/// re-checking; the stable sort preserves the position tiebreak for equal +/// versions. +fn sort_overlays_newest_last(overlays: &mut [DataOverlayFile]) { + overlays.sort_by_key(|overlay| overlay.committed_version); +} + impl From<&DataOverlayFile> for pb::DataOverlayFile { fn from(overlay: &DataOverlayFile) -> Self { let coverage = match &overlay.coverage { - OverlayCoverage::Shared(bytes) => { - pb::data_overlay_file::Coverage::SharedOffsetBitmap(bytes.clone()) + OverlayCoverage::Shared(bitmap) => { + pb::data_overlay_file::Coverage::SharedOffsetBitmap(serialize_roaring(bitmap)) } OverlayCoverage::PerField(bitmaps) => { pb::data_overlay_file::Coverage::FieldCoverage(pb::FieldCoverage { - offset_bitmaps: bitmaps.clone(), + offset_bitmaps: bitmaps.iter().map(|b| serialize_roaring(b)).collect(), }) } }; @@ -354,11 +416,14 @@ impl TryFrom for DataOverlayFile { .ok_or_else(|| Error::invalid_input("DataOverlayFile is missing its data_file"))?; let coverage = match proto.coverage { Some(pb::data_overlay_file::Coverage::SharedOffsetBitmap(bytes)) => { - OverlayCoverage::Shared(bytes) - } - Some(pb::data_overlay_file::Coverage::FieldCoverage(fc)) => { - OverlayCoverage::PerField(fc.offset_bitmaps) + OverlayCoverage::Shared(Arc::new(deserialize_roaring(&bytes)?)) } + Some(pb::data_overlay_file::Coverage::FieldCoverage(fc)) => OverlayCoverage::PerField( + fc.offset_bitmaps + .iter() + .map(|b| deserialize_roaring(b).map(Arc::new)) + .collect::>()?, + ), None => { return Err(Error::invalid_input( "DataOverlayFile is missing its coverage", @@ -515,11 +580,15 @@ impl DataFileFieldInterner { .into_iter() .map(|f| self.intern_data_file(f)) .collect::>()?, - overlays: p - .overlays - .into_iter() - .map(DataOverlayFile::try_from) - .collect::>()?, + overlays: { + let mut overlays = p + .overlays + .into_iter() + .map(DataOverlayFile::try_from) + .collect::>>()?; + sort_overlays_newest_last(&mut overlays); + overlays + }, deletion_file: p.deletion_file.map(DeletionFile::try_from).transpose()?, row_id_meta: p.row_id_sequence.map(RowIdMeta::try_from).transpose()?, physical_rows, @@ -822,11 +891,15 @@ impl TryFrom for Fragment { .into_iter() .map(DataFile::try_from) .collect::>()?, - overlays: p - .overlays - .into_iter() - .map(DataOverlayFile::try_from) - .collect::>()?, + overlays: { + let mut overlays = p + .overlays + .into_iter() + .map(DataOverlayFile::try_from) + .collect::>>()?; + sort_overlays_newest_last(&mut overlays); + overlays + }, deletion_file: p.deletion_file.map(DeletionFile::try_from).transpose()?, row_id_meta: p.row_id_sequence.map(RowIdMeta::try_from).transpose()?, physical_rows, @@ -903,7 +976,7 @@ mod tests { let overlay = DataOverlayFile { data_file: DataFile::new_legacy_from_fields("overlay-0.lance", vec![3], None), - coverage: OverlayCoverage::dense(&bitmap), + coverage: OverlayCoverage::dense(bitmap.clone()), committed_version: 7, }; let mut fragment = Fragment::new(0); @@ -921,9 +994,9 @@ mod tests { // Dense coverage applies to every field. let recovered = round_tripped.overlays[0].coverage_for_field(0).unwrap(); - assert_eq!(recovered, bitmap); + assert_eq!(*recovered, bitmap); assert_eq!( - round_tripped.overlays[0].coverage_for_field(5).unwrap(), + *round_tripped.overlays[0].coverage_for_field(5).unwrap(), bitmap ); } @@ -935,7 +1008,10 @@ mod tests { let embedding_coverage = RoaringBitmap::from_iter([1u32]); let overlay = DataOverlayFile { data_file: DataFile::new_legacy_from_fields("overlay-1.lance", vec![2, 4], None), - coverage: OverlayCoverage::sparse(&[name_coverage.clone(), embedding_coverage.clone()]), + coverage: OverlayCoverage::sparse(vec![ + name_coverage.clone(), + embedding_coverage.clone(), + ]), committed_version: 3, }; let mut fragment = Fragment::new(1); @@ -943,15 +1019,73 @@ mod tests { let round_tripped = Fragment::try_from(pb::DataFragment::from(&fragment)).unwrap(); assert_eq!( - round_tripped.overlays[0].coverage_for_field(0).unwrap(), + *round_tripped.overlays[0].coverage_for_field(0).unwrap(), name_coverage ); assert_eq!( - round_tripped.overlays[0].coverage_for_field(1).unwrap(), + *round_tripped.overlays[0].coverage_for_field(1).unwrap(), embedding_coverage ); } + #[test] + fn test_data_overlay_missing_fields_error() { + // A DataOverlayFile proto missing its coverage or data_file is rejected. + let no_coverage = pb::DataOverlayFile { + data_file: Some(pb::DataFile::from(&DataFile::new_legacy_from_fields( + "overlay.lance", + vec![3], + None, + ))), + coverage: None, + committed_version: 1, + }; + let err = DataOverlayFile::try_from(no_coverage).unwrap_err(); + assert!(err.to_string().contains("missing its coverage"), "{err}"); + + let no_data_file = pb::DataOverlayFile { + data_file: None, + coverage: Some(pb::data_overlay_file::Coverage::SharedOffsetBitmap( + serialize_roaring(&RoaringBitmap::from_iter([0u32])), + )), + committed_version: 1, + }; + let err = DataOverlayFile::try_from(no_data_file).unwrap_err(); + assert!(err.to_string().contains("missing its data_file"), "{err}"); + } + + #[test] + fn test_overlays_sorted_newest_last_on_load() { + // Overlays load stable-sorted by committed_version (newest last), with + // list position preserved as the tiebreak for equal versions. + let mk = |version: u64, field: i32| DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("o.lance", vec![field], None), + coverage: OverlayCoverage::dense(RoaringBitmap::from_iter([0u32])), + committed_version: version, + }; + let mut fragment = Fragment::new(0); + // Written out of order: v5, v2, v2 (second), v3. + fragment.overlays = vec![mk(5, 1), mk(2, 2), mk(2, 3), mk(3, 4)]; + + let loaded = Fragment::try_from(pb::DataFragment::from(&fragment)).unwrap(); + let versions: Vec = loaded + .overlays + .iter() + .map(|o| o.committed_version) + .collect(); + assert_eq!(versions, vec![2, 2, 3, 5]); + // Stable: the two v2 overlays keep their original relative order (field 2 + // before field 3). + assert_eq!( + loaded.overlays[0].data_file.fields.as_ref(), + [2i32].as_slice() + ); + assert_eq!( + loaded.overlays[1].data_file.fields.as_ref(), + [3i32].as_slice() + ); + } + #[test] fn test_new_fragment() { let path = "foobar.lance"; diff --git a/rust/lance/src/dataset/transaction.rs b/rust/lance/src/dataset/transaction.rs index 32d1833f9bb..900db143b3b 100644 --- a/rust/lance/src/dataset/transaction.rs +++ b/rust/lance/src/dataset/transaction.rs @@ -6288,7 +6288,7 @@ mod tests { bitmap.insert(4); let overlay = DataOverlayFile { data_file: DataFile::new_legacy_from_fields("overlay-0.lance", vec![3], None), - coverage: OverlayCoverage::dense(&bitmap), + coverage: OverlayCoverage::dense(bitmap.clone()), committed_version: 6, }; let pb_overlay = pb::DataOverlayFile::from(&overlay); @@ -6314,7 +6314,10 @@ mod tests { assert_eq!(groups[0].fragment_id, 7); assert_eq!(groups[0].overlays.len(), 1); assert_eq!(groups[0].overlays[0].committed_version, 6); - assert_eq!(groups[0].overlays[0].coverage_for_field(0).unwrap(), bitmap); + assert_eq!( + *groups[0].overlays[0].coverage_for_field(0).unwrap(), + bitmap + ); } other => panic!("expected DataOverlay, got {other:?}"), } From e2f98f51256cc1dcbc7e6c268979dfe0c91ef046 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Wed, 24 Jun 2026 14:23:16 -0700 Subject: [PATCH 09/13] fix(table): correct DataOverlay commit/conflict edge cases Three correctness fixes in the data-overlay write/commit path, each with tests (this path had ~0% coverage): - build_manifest dropped overlays when two DataOverlayGroups targeted the same fragment (HashMap collect kept only the last). Merge groups by fragment_id so all overlays are appended in order. - check_data_overlay_txn did not conflict when a concurrent Update/Delete removed an overlaid fragment, leaving the overlay orphaned and hard-erroring on retry. Mirror check_data_replacement_txn: retryable conflict when removed/deleted_fragment_ids intersect our fragments; fragments merely updated in place stay compatible. - impl PartialEq for Operation had the DataOverlay arm backwards: it reported DataOverlay == Rewrite as true and DataOverlay == DataOverlay as false (conflict semantics copied into the equality impl). Compare groups for the same variant; not equal to other variants. Tests: build_manifest append/stamp, duplicate-group merge, unknown-fragment error; the conflict matrix incl. remove-fragment conflict; Operation equality; OverlayCoverage serde JSON round-trip; coverage_for_field out-of-bounds; apply_feature_flags overlay arm. Co-Authored-By: Claude Opus 4.8 (1M context) --- rust/lance-table/src/feature_flags.rs | 34 ++++ rust/lance-table/src/format/fragment.rs | 35 ++++ rust/lance/src/dataset/transaction.rs | 167 ++++++++++++++++-- rust/lance/src/io/commit/conflict_resolver.rs | 148 +++++++++++++++- 4 files changed, 364 insertions(+), 20 deletions(-) diff --git a/rust/lance-table/src/feature_flags.rs b/rust/lance-table/src/feature_flags.rs index 40cc38aeeda..369c880c062 100644 --- a/rust/lance-table/src/feature_flags.rs +++ b/rust/lance-table/src/feature_flags.rs @@ -172,6 +172,40 @@ mod tests { assert_eq!(FLAG_DATA_OVERLAY_FILES & !supported, 0); } + #[test] + fn test_apply_feature_flags_sets_overlay_flag() { + use crate::format::{ + DataFile, DataOverlayFile, DataStorageFormat, Fragment, OverlayCoverage, + }; + use arrow_schema::{Field as ArrowField, Schema as ArrowSchema}; + use lance_core::datatypes::Schema; + use roaring::RoaringBitmap; + use std::collections::HashMap; + use std::sync::Arc; + + let arrow_schema = ArrowSchema::new(vec![ArrowField::new( + "id", + arrow_schema::DataType::Int64, + false, + )]); + let schema = Schema::try_from(&arrow_schema).unwrap(); + let mut fragment = Fragment::new(0); + fragment.overlays = vec![DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("o.lance", vec![0], None), + coverage: OverlayCoverage::dense(RoaringBitmap::from_iter([0u32])), + committed_version: 1, + }]; + let mut manifest = Manifest::new( + schema, + Arc::new(vec![fragment]), + DataStorageFormat::default(), + HashMap::new(), + ); + apply_feature_flags(&mut manifest, false, false).unwrap(); + assert_ne!(manifest.reader_feature_flags & FLAG_DATA_OVERLAY_FILES, 0); + assert_ne!(manifest.writer_feature_flags & FLAG_DATA_OVERLAY_FILES, 0); + } + #[test] fn test_write_check() { assert!(can_write_dataset(0)); diff --git a/rust/lance-table/src/format/fragment.rs b/rust/lance-table/src/format/fragment.rs index c23f312d3ce..ef5166728ab 100644 --- a/rust/lance-table/src/format/fragment.rs +++ b/rust/lance-table/src/format/fragment.rs @@ -1086,6 +1086,41 @@ mod tests { ); } + #[test] + fn test_overlay_coverage_serde_json_roundtrip() { + // The custom serde impl round-trips through JSON for dense/sparse, + // including empty bitmaps and a zero-bitmap sparse coverage. + for coverage in [ + OverlayCoverage::dense(RoaringBitmap::from_iter([1u32, 5, 100])), + OverlayCoverage::dense(RoaringBitmap::new()), + OverlayCoverage::sparse(vec![ + RoaringBitmap::from_iter([2u32, 3]), + RoaringBitmap::new(), + ]), + OverlayCoverage::sparse(vec![]), + ] { + let json = serde_json::to_string(&coverage).unwrap(); + let back: OverlayCoverage = serde_json::from_str(&json).unwrap(); + assert_eq!(back, coverage); + } + } + + #[test] + fn test_coverage_for_field_out_of_bounds() { + let overlay = DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("o.lance", vec![2, 4], None), + coverage: OverlayCoverage::sparse(vec![ + RoaringBitmap::from_iter([1u32]), + RoaringBitmap::from_iter([2u32]), + ]), + committed_version: 1, + }; + assert!(overlay.coverage_for_field(0).is_ok()); + assert!(overlay.coverage_for_field(1).is_ok()); + let err = overlay.coverage_for_field(5).unwrap_err(); + assert!(err.to_string().contains("field position"), "{err}"); + } + #[test] fn test_new_fragment() { let path = "foobar.lance"; diff --git a/rust/lance/src/dataset/transaction.rs b/rust/lance/src/dataset/transaction.rs index 900db143b3b..33248c2b68d 100644 --- a/rust/lance/src/dataset/transaction.rs +++ b/rust/lance/src/dataset/transaction.rs @@ -1363,15 +1363,7 @@ impl PartialEq for Operation { (Self::Clone { .. }, Self::UpdateBases { .. }) => { std::mem::discriminant(self) == std::mem::discriminant(other) } - // Data overlays are intentionally permissive, like DataReplacement. - // Two overlays stack (the higher committed_version wins each covered - // cell), and overlays are compatible with appends, deletes, column - // rewrites, and concurrent overlays. Only an operation that rewrites - // rows or otherwise invalidates physical offsets (Rewrite, which - // covers compaction and overlay->base folds) conflicts, since the - // overlay's physical offsets would no longer be valid. - (Self::DataOverlay { .. }, Self::Rewrite { .. }) - | (Self::Rewrite { .. }, Self::DataOverlay { .. }) => true, + (Self::DataOverlay { groups: a }, Self::DataOverlay { groups: b }) => compare_vec(a, b), (Self::DataOverlay { .. }, _) | (_, Self::DataOverlay { .. }) => false, } } @@ -2336,10 +2328,16 @@ impl Transaction { let new_version = current_manifest.map_or(1, |m| m.version + 1); let existing_fragments = maybe_existing_fragments?; - let overlays_by_fragment: HashMap> = groups - .iter() - .map(|g| (g.fragment_id, &g.overlays)) - .collect(); + // Multiple groups may target the same fragment; merge them in + // order rather than letting a HashMap collapse drop all but the + // last group's overlays. + let mut overlays_by_fragment: HashMap> = HashMap::new(); + for group in groups { + overlays_by_fragment + .entry(group.fragment_id) + .or_default() + .extend(group.overlays.iter()); + } // Every group must target an existing fragment. for fragment_id in overlays_by_fragment.keys() { @@ -2355,12 +2353,13 @@ impl Transaction { if let Some(new_overlays) = overlays_by_fragment.get(&fragment.id) { // Appended (not replaced) so concurrently-written overlays // survive; later entries are newer. - fragment.overlays.extend(new_overlays.iter().cloned().map( - |mut overlay| { + fragment + .overlays + .extend(new_overlays.iter().map(|&overlay| { + let mut overlay = overlay.clone(); overlay.committed_version = new_version; overlay - }, - )); + })); } final_fragments.push(fragment); } @@ -3947,7 +3946,8 @@ mod tests { use lance_file::version::LanceFileVersion; use lance_io::utils::CachedFileSize; use lance_table::format::{ - RowDatasetVersionMeta, RowDatasetVersionRun, RowDatasetVersionSequence, RowIdMeta, + OverlayCoverage, RowDatasetVersionMeta, RowDatasetVersionRun, RowDatasetVersionSequence, + RowIdMeta, }; use lance_table::rowids::segment::U64Segment; use lance_table::rowids::write_row_ids; @@ -6322,4 +6322,135 @@ mod tests { other => panic!("expected DataOverlay, got {other:?}"), } } + + fn overlay_with_field(field: i32, committed_version: u64) -> DataOverlayFile { + DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("o.lance", vec![field], None), + coverage: OverlayCoverage::dense(roaring::RoaringBitmap::from_iter([0u32])), + committed_version, + } + } + + #[test] + fn test_data_overlay_build_manifest_appends_and_stamps() { + // A fragment already carrying an overlay (committed at v3) gets a new + // overlay appended and stamped to the new dataset version; the existing + // overlay is preserved with its version. + let mut fragment = Fragment::new(0); + fragment.overlays = vec![overlay_with_field(1, 3)]; + let schema = ArrowSchema::new(vec![ArrowField::new("id", DataType::Int32, false)]); + let manifest = Manifest::new( + LanceSchema::try_from(&schema).unwrap(), + Arc::new(vec![fragment]), + lance_table::format::DataStorageFormat::new(LanceFileVersion::V2_0), + HashMap::new(), + ); + + let txn = Transaction::new( + manifest.version, + Operation::DataOverlay { + groups: vec![DataOverlayGroup { + fragment_id: 0, + overlays: vec![overlay_with_field(2, 0)], + }], + }, + None, + ); + + let (result, _) = txn + .build_manifest( + Some(&manifest), + vec![], + "txn", + &ManifestWriteConfig::default(), + ) + .unwrap(); + + let frag = &result.fragments[0]; + assert_eq!(frag.overlays.len(), 2); + assert_eq!(frag.overlays[0].committed_version, 3); + assert_eq!(frag.overlays[1].committed_version, result.version); + assert!(result.version > manifest.version); + } + + #[test] + fn test_data_overlay_build_manifest_merges_duplicate_groups() { + // Two groups targeting the same fragment must both survive (a HashMap + // collapse would have dropped the first). + let manifest = sample_manifest(); + let txn = Transaction::new( + manifest.version, + Operation::DataOverlay { + groups: vec![ + DataOverlayGroup { + fragment_id: 0, + overlays: vec![overlay_with_field(1, 0)], + }, + DataOverlayGroup { + fragment_id: 0, + overlays: vec![overlay_with_field(2, 0)], + }, + ], + }, + None, + ); + + let (result, _) = txn + .build_manifest( + Some(&manifest), + vec![], + "txn", + &ManifestWriteConfig::default(), + ) + .unwrap(); + + let overlays = &result.fragments[0].overlays; + assert_eq!(overlays.len(), 2); + assert_eq!(overlays[0].data_file.fields.as_ref(), [1i32].as_slice()); + assert_eq!(overlays[1].data_file.fields.as_ref(), [2i32].as_slice()); + } + + #[test] + fn test_data_overlay_build_manifest_rejects_unknown_fragment() { + let manifest = sample_manifest(); + let txn = Transaction::new( + manifest.version, + Operation::DataOverlay { + groups: vec![DataOverlayGroup { + fragment_id: 99, + overlays: vec![overlay_with_field(1, 0)], + }], + }, + None, + ); + let err = txn + .build_manifest( + Some(&manifest), + vec![], + "txn", + &ManifestWriteConfig::default(), + ) + .unwrap_err(); + assert!(err.to_string().contains("does not exist"), "{err}"); + } + + #[test] + fn test_data_overlay_operation_eq() { + let overlay = |field: i32| Operation::DataOverlay { + groups: vec![DataOverlayGroup { + fragment_id: 0, + overlays: vec![overlay_with_field(field, 1)], + }], + }; + // Reflexive and value-based (the arm previously returned false for self). + assert_eq!(overlay(1), overlay(1)); + assert_ne!(overlay(1), overlay(2)); + // Not equal to a different operation kind (previously returned true vs Rewrite). + let rewrite = Operation::Rewrite { + groups: vec![], + rewritten_indices: vec![], + frag_reuse_index: None, + }; + assert_ne!(overlay(1), rewrite); + } } diff --git a/rust/lance/src/io/commit/conflict_resolver.rs b/rust/lance/src/io/commit/conflict_resolver.rs index 1af51506f8e..1e46c4acffa 100644 --- a/rust/lance/src/io/commit/conflict_resolver.rs +++ b/rust/lance/src/io/commit/conflict_resolver.rs @@ -1073,8 +1073,6 @@ impl<'a> TransactionRebase<'a> { ) -> Result<()> { match &other_transaction.operation { Operation::Append { .. } - | Operation::Delete { .. } - | Operation::Update { .. } | Operation::CreateIndex { .. } | Operation::ReserveFragments { .. } | Operation::Project { .. } @@ -1084,6 +1082,29 @@ impl<'a> TransactionRebase<'a> { | Operation::UpdateMemWalState { .. } | Operation::DataReplacement { .. } | Operation::DataOverlay { .. } => Ok(()), + // A concurrent Update or Delete that *removes* one of our overlaid + // fragments leaves the overlay orphaned — it is addressed by physical + // offset into a fragment that would no longer exist — so conflict and + // retry against the new fragment list. Fragments merely updated in + // place (deletion vectors, column rewrites) preserve physical offsets, + // so the overlay stays valid and does not conflict. + Operation::Update { + removed_fragment_ids, + .. + } + | Operation::Delete { + deleted_fragment_ids: removed_fragment_ids, + .. + } => { + if removed_fragment_ids + .iter() + .any(|id| self.modified_fragment_ids.contains(id)) + { + Err(self.retryable_conflict_err(other_transaction, other_version)) + } else { + Ok(()) + } + } Operation::Rewrite { groups, .. } => { // A rewrite (compaction / fold) of a fragment we are overlaying // changes its physical row addresses, so our offsets would be @@ -2840,6 +2861,129 @@ mod tests { } } + #[test] + fn test_data_overlay_conflicts() { + use crate::dataset::transaction::DataOverlayGroup; + use ConflictResult::*; + use lance_table::format::{DataOverlayFile, OverlayCoverage}; + use roaring::RoaringBitmap; + + // Our transaction overlays fragment 1. + let overlay_op = |fragment_id: u64| Operation::DataOverlay { + groups: vec![DataOverlayGroup { + fragment_id, + overlays: vec![DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("overlay.lance", vec![0], None), + coverage: OverlayCoverage::dense(RoaringBitmap::from_iter([0u32])), + committed_version: 0, + }], + }], + }; + let update_removing = |removed_fragment_ids: Vec| Operation::Update { + removed_fragment_ids, + updated_fragments: vec![], + new_fragments: vec![], + fields_modified: vec![], + merged_generations: Vec::new(), + fields_for_preserving_frag_bitmap: vec![], + update_mode: None, + inserted_rows_filter: None, + updated_fragment_offsets: None, + }; + let delete = |updated: Vec, deleted: Vec| Operation::Delete { + updated_fragments: updated, + deleted_fragment_ids: deleted, + predicate: "x > 2".to_string(), + }; + let rewrite_of = |old: &Fragment| Operation::Rewrite { + groups: vec![RewriteGroup { + old_fragments: vec![old.clone()], + new_fragments: vec![], + }], + rewritten_indices: vec![], + frag_reuse_index: None, + }; + + let fragment0 = Fragment::new(0); + let fragment1 = Fragment::new(1); + + // Each case is checked against our overlay on fragment 1. + let cases: Vec<(Operation, ConflictResult)> = vec![ + // Permissive: preserves physical offsets / leaves fragment 1 in place. + ( + Operation::Append { + fragments: vec![fragment0.clone()], + }, + Compatible, + ), + ( + Operation::CreateIndex { + new_indices: vec![], + removed_indices: vec![], + }, + Compatible, + ), + ( + Operation::DataReplacement { + replacements: vec![DataReplacementGroup( + 1, + DataFile::new_legacy_from_fields("r.lance", vec![0], None), + )], + }, + Compatible, + ), + // Another overlay on the same fragment stacks rather than conflicts. + (overlay_op(1), Compatible), + // Delete/Update that only updates fragment 1 in place is compatible. + (delete(vec![fragment1.clone()], vec![]), Compatible), + (update_removing(vec![2]), Compatible), + // ...but removing our overlaid fragment 1 orphans the overlay -> conflict. + (delete(vec![], vec![1]), Retryable), + (update_removing(vec![1]), Retryable), + // Rewriting fragment 1 invalidates its physical offsets -> conflict; + // a rewrite of a different fragment does not. + (rewrite_of(&fragment1), Retryable), + (rewrite_of(&fragment0), Compatible), + // Merge rewrites the whole fragment list; Restore replaces the dataset. + ( + Operation::Merge { + fragments: vec![fragment1.clone()], + schema: lance_core::datatypes::Schema::default(), + }, + Retryable, + ), + (Operation::Restore { version: 1 }, NotCompatible), + ]; + + for (other, expected) in cases { + let mut rebase = TransactionRebase { + transaction: Transaction::new(0, overlay_op(1), None), + initial_fragments: HashMap::new(), + modified_fragment_ids: modified_fragment_ids(&overlay_op(1)) + .collect::>(), + affected_rows: None, + conflicting_frag_reuse_indices: Vec::new(), + conflicting_mem_wal_merged_gens: Vec::new(), + }; + let other_txn = Transaction::new(0, other.clone(), None); + let result = rebase.check_txn(&other_txn, 1); + match expected { + Compatible => assert!( + result.is_ok(), + "overlay should be compatible with {other:?}, got {result:?}" + ), + Retryable => assert!( + matches!(result, Err(Error::RetryableCommitConflict { .. })), + "overlay should retryably conflict with {other:?}, got {result:?}" + ), + NotCompatible => assert!( + matches!(result, Err(Error::IncompatibleTransaction { .. })), + "overlay should be incompatible with {other:?}, got {result:?}" + ), + } + } + } + #[test] fn test_create_index_conflicts_only_on_same_name() { let index0 = IndexMetadata { From 8052f87f8e8d29c6903148e4824d01ada7d5eb10 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Tue, 30 Jun 2026 10:18:27 -0700 Subject: [PATCH 10/13] 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 | 707 ++++++++++++++++++++++++++++- rust/lance/src/dataset/overlay.rs | 352 ++++++++++++++ 3 files changed, 1045 insertions(+), 15 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 3e0d77704da..448feb961d7 100644 --- a/rust/lance/src/dataset.rs +++ b/rust/lance/src/dataset.rs @@ -79,6 +79,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 58a3ce0a9b6..113e6d4a649 100644 --- a/rust/lance/src/dataset/fragment.rs +++ b/rust/lance/src/dataset/fragment.rs @@ -15,7 +15,8 @@ use arrow::compute::concat_batches; use arrow_array::cast::as_primitive_array; use arrow_array::types::UInt64Type; use arrow_array::{ - Array, RecordBatch, RecordBatchReader, StructArray, UInt32Array, UInt64Array, new_null_array, + Array, ArrayRef, RecordBatch, RecordBatchReader, StructArray, UInt32Array, UInt64Array, + new_null_array, }; use arrow_schema::Schema as ArrowSchema; use datafusion::logical_expr::Expr; @@ -47,7 +48,7 @@ use lance_table::format::{DataFile, DeletionFile, Fragment}; use lance_table::io::deletion::{deletion_file_path, write_deletion_file}; use lance_table::rowids::RowIdSequence; use lance_table::utils::stream::{ - ReadBatchFutStream, ReadBatchTask, ReadBatchTaskStream, RowIdAndDeletesConfig, + ReadBatchFut, ReadBatchFutStream, ReadBatchTask, ReadBatchTaskStream, RowIdAndDeletesConfig, wrap_with_row_id_and_delete, }; use roaring::RoaringBitmap; @@ -62,6 +63,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::{ + assemble_overlay_column, overlay_indices_newest_first, route_overlays, +}; use crate::io::deletion::read_dataset_deletion_file; /// Result of [`FileFragment::update_columns_with_offsets`]: updated fragment metadata, modified field ids, @@ -900,16 +904,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(); @@ -921,11 +915,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(self.load_overlay_plan(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!( @@ -948,6 +955,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(); } @@ -1123,6 +1134,91 @@ impl FileFragment { Ok(opened_files) } + /// Open the overlay value-column readers for the 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 [`FragmentReader::merge_overlays`]), + /// 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. + async fn load_overlay_plan( + &self, + projection: &Schema, + read_config: &FragReadConfig, + ) -> Result> { + let order = overlay_indices_newest_first(&self.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 = &self.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(self + .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 = &self.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) + } + /// Count the rows in this fragment. pub async fn count_rows(&self, filter: Option) -> Result { match filter { @@ -1999,6 +2095,106 @@ impl From for Fragment { } } +/// 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)] +struct FieldOverlayPlan { + field_name: String, + overlays_newest_first: Vec, +} + +/// 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. +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)?) +} + /// [`FragmentReader`] is an abstract reader for a [`FileFragment`]. /// /// It opens the data files that contains the columns of the projection schema, and @@ -2052,6 +2248,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 @@ -2080,6 +2281,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(), } } } @@ -2147,6 +2349,7 @@ impl FragmentReader { created_at_sequence: None, num_rows, num_physical_rows, + overlay_plans: Arc::new(Vec::new()), }) } @@ -2404,6 +2607,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, @@ -2471,6 +2718,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 { @@ -2641,6 +2890,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 { @@ -2653,7 +2905,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()); @@ -2849,6 +3101,431 @@ 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, 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..b7ff386ea6e --- /dev/null +++ b/rust/lance/src/dataset/overlay.rs @@ -0,0 +1,352 @@ +// 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 arrow_array::{Array, ArrayRef}; +use arrow_select::interleave::interleave; +use lance_core::{Error, Result}; +use roaring::RoaringBitmap; + +use lance_table::format::DataOverlayFile; + +/// 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) +} + +#[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, 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 0a984a06f8e5f5c7e8174a3bcffd7b27eff0a514 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Tue, 30 Jun 2026 12:30:33 -0700 Subject: [PATCH 11/13] 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 113e6d4a649..108c2debcf9 100644 --- a/rust/lance/src/dataset/fragment.rs +++ b/rust/lance/src/dataset/fragment.rs @@ -3107,7 +3107,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; @@ -3524,6 +3526,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 3ba338bb8562ab2ffbf7fc01815587cb10dc19a9 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Tue, 30 Jun 2026 16:11:44 -0700 Subject: [PATCH 12/13] 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 b7ff386ea6e..60f8fe5ff63 100644 --- a/rust/lance/src/dataset/overlay.rs +++ b/rust/lance/src/dataset/overlay.rs @@ -82,10 +82,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 { @@ -333,6 +420,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, OverlayCoverage}; From a2e997fb25020c6738bebff5e30e9ae36a2608b3 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Tue, 23 Jun 2026 09:04:17 -0700 Subject: [PATCH 13/13] feat: overlay-aware compaction scheduler and rewrite MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add overlay-aware compaction (the "Compaction" section of the Data Overlay Files spec). A scheduler reads each fragment's overlay state — overlay count, covered cells, and the version gap between an overlay's committed_version and the base and each covering index's dataset_version — and picks a mode from the staleness signal: - Overlay -> overlay merge: collapse a fragment's overlays into one smaller overlay carrying the per-(offset, field) post-image, stamped with the maximum input committed_version. Base and indexes are untouched. - Overlay -> base fold: materialize the post-image of every covered cell into a fresh base data file, tombstone the folded fields in the old files, and clear the fragment's overlays. Row addresses are preserved (a column rewrite). The same commit drops the folded fragment from any index on a folded field, so an index is never left serving stale values. Both modes read each base/overlay cell at most once: overlay value columns are read only at the ranks that win (a cell a newer overlay supplies is never fetched from an older overlay), via FileFragment::read_overlay_field_winners. Base columns are read whole, the cheaper coalesced read the spec permits. The rewrite is committed via the existing Operation::Update (RewriteColumns) path, so no new persisted operation or conflict-matrix change is needed. Tests cover scheduler mode selection (field-aware, version-gap driven), fold materialization + index reconciliation + NULL override + multi-fragment, merge collapse preserving the max committed_version, and an assert_io check that overlapping coverage reads strictly fewer value bytes than a disjoint control. Co-Authored-By: Claude Opus 4.8 (1M context) --- rust/lance/src/dataset.rs | 1 + rust/lance/src/dataset/fragment.rs | 175 ++- rust/lance/src/dataset/overlay_compaction.rs | 1145 ++++++++++++++++++ 3 files changed, 1320 insertions(+), 1 deletion(-) create mode 100644 rust/lance/src/dataset/overlay_compaction.rs diff --git a/rust/lance/src/dataset.rs b/rust/lance/src/dataset.rs index 448feb961d7..2e96e67f04b 100644 --- a/rust/lance/src/dataset.rs +++ b/rust/lance/src/dataset.rs @@ -80,6 +80,7 @@ pub mod mem_wal; mod metadata; pub mod optimize; pub(crate) mod overlay; +pub mod overlay_compaction; 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 108c2debcf9..787ae8f1de4 100644 --- a/rust/lance/src/dataset/fragment.rs +++ b/rust/lance/src/dataset/fragment.rs @@ -44,7 +44,7 @@ use lance_file::{LanceEncodingsIo, determine_file_version}; use lance_io::ReadBatchParams; use lance_io::scheduler::{FileScheduler, ScanScheduler, SchedulerConfig}; use lance_io::utils::CachedFileSize; -use lance_table::format::{DataFile, DeletionFile, Fragment}; +use lance_table::format::{DataFile, DataOverlayFile, DeletionFile, Fragment}; use lance_table::io::deletion::{deletion_file_path, write_deletion_file}; use lance_table::rowids::RowIdSequence; use lance_table::utils::stream::{ @@ -1219,6 +1219,167 @@ impl FileFragment { Ok(plans) } + /// Read just the values at the given `ranks` (0-based positions within an + /// overlay field's value column) via a `take`, rather than the whole column. + /// A value column has one row per covered offset; position `r` holds the + /// value for the offset whose rank is `r`. Used by overlay compaction so a + /// cell that a newer overlay already supplies is never fetched from an older + /// overlay. + async fn read_overlay_value_column_at( + &self, + overlay: &DataOverlayFile, + field: &lance_core::datatypes::Field, + ranks: &[u32], + read_config: &FragReadConfig, + ) -> Result { + if ranks.is_empty() { + return Ok(arrow_array::new_empty_array(&field.data_type())); + } + let single_field = Schema { + fields: vec![field.clone()], + metadata: Default::default(), + }; + let reader = self + .open_reader(&overlay.data_file, Some(&single_field), read_config) + .await? + .ok_or_else(|| { + Error::internal(format!( + "overlay data file {} does not contain field {} (id {})", + overlay.data_file.path, field.name, field.id + )) + })?; + let mut tasks = reader + .take_all_tasks(ranks, ranks.len() as u32, reader.projection().clone(), 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)?) + } + + /// Resolve, for a single field, the post-image of every offset this + /// fragment's overlays cover, reading each winning cell **at most once**. + /// + /// The overlays are walked newest-first; for each overlay only the offsets + /// not already claimed by a newer overlay are read (via a rank `take`), so + /// overlapping coverage never fetches the same cell from more than one + /// source. The returned [`ResolvedFieldOverlay`]'s `coverage` is the union of + /// every overlay's coverage for the field, and its `values` are in ascending + /// offset order (i.e. rank order within `coverage`). + pub(crate) async fn read_overlay_field_winners( + &self, + field: &lance_core::datatypes::Field, + read_config: &FragReadConfig, + ) -> Result { + let overlays = &self.metadata.overlays; + let order = overlay_indices_newest_first(overlays); + let mut claimed = RoaringBitmap::new(); + let mut taken: Vec = Vec::new(); + // offset -> (index into `taken`, position within that taken array) + let mut sources_by_offset: BTreeMap = BTreeMap::new(); + for overlay_idx in order { + let overlay = &overlays[overlay_idx]; + let Some(field_pos) = overlay + .data_file + .fields + .iter() + .position(|&id| id == field.id) + else { + continue; + }; + let coverage = overlay.coverage_for_field(field_pos)?; + let winners: Vec = (&*coverage - &claimed).iter().collect(); + claimed |= &*coverage; + if winners.is_empty() { + continue; + } + let ranks: Vec = winners + .iter() + .map(|&offset| coverage.rank(offset) as u32 - 1) + .collect(); + let values = self + .read_overlay_value_column_at(overlay, field, &ranks, read_config) + .await?; + let taken_idx = taken.len(); + for (pos, &offset) in winners.iter().enumerate() { + sources_by_offset.insert(offset, (taken_idx, pos)); + } + taken.push(values); + } + if sources_by_offset.is_empty() { + return Ok(ResolvedFieldOverlay { + coverage: RoaringBitmap::new(), + values: arrow_array::new_empty_array(&field.data_type()), + }); + } + // A BTreeMap iterates in ascending key (offset) order, which is exactly + // the rank order of the union coverage. + let coverage: RoaringBitmap = sources_by_offset.keys().copied().collect(); + let indices: Vec<(usize, usize)> = sources_by_offset.values().copied().collect(); + let sources: Vec<&dyn arrow_array::Array> = taken.iter().map(|a| a.as_ref()).collect(); + let values = arrow_select::interleave::interleave(&sources, &indices)?; + Ok(ResolvedFieldOverlay { coverage, values }) + } + + /// The live (non-tombstoned) data file in this fragment that stores + /// `field_id`, if any. A tombstoned column is recorded as field id `-2`, so + /// it never matches a real field id here. + pub(crate) fn base_data_file_for_field(&self, field_id: i32) -> Option<&DataFile> { + self.metadata + .files + .iter() + .find(|f| f.fields.contains(&field_id)) + } + + /// Read a field's full base column for every physical row (counting deleted + /// rows) **without** merging overlays. Returns an all-null column when no + /// base data file holds the field (an overlay over a column the base never + /// materialized). + pub(crate) async fn read_base_field_full( + &self, + field: &lance_core::datatypes::Field, + read_config: &FragReadConfig, + ) -> Result { + let physical_rows = self.physical_rows().await?; + if physical_rows == 0 { + return Ok(arrow_array::new_empty_array(&field.data_type())); + } + let single_field = Schema { + fields: vec![field.clone()], + metadata: Default::default(), + }; + let reader = match self.base_data_file_for_field(field.id) { + Some(data_file) => { + self.open_reader(data_file, Some(&single_field), read_config) + .await? + } + None => None, + }; + let Some(reader) = reader else { + return Ok(arrow_array::new_null_array( + &field.data_type(), + physical_rows, + )); + }; + let mut tasks = reader + .read_range_tasks( + 0..physical_rows as u64, + physical_rows as u32, + reader.projection().clone(), + ) + .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)?) + } + /// Count the rows in this fragment. pub async fn count_rows(&self, filter: Option) -> Result { match filter { @@ -2117,6 +2278,18 @@ struct FieldOverlayPlan { overlays_newest_first: Vec, } +/// The post-image of a single field over every offset this fragment's overlays +/// cover, produced by [`FileFragment::read_overlay_field_winners`] for overlay +/// compaction. `values` is in ascending offset order — i.e. rank order within +/// `coverage` — so a covered offset `o`'s value sits at `coverage`'s rank of `o`. +#[derive(Debug, Clone)] +pub(crate) struct ResolvedFieldOverlay { + /// Union of every overlay's coverage for the field. + pub coverage: RoaringBitmap, + /// The winning value for each covered offset, indexed by rank in `coverage`. + pub values: ArrayRef, +} + /// 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 diff --git a/rust/lance/src/dataset/overlay_compaction.rs b/rust/lance/src/dataset/overlay_compaction.rs new file mode 100644 index 00000000000..1bc919c7ee4 --- /dev/null +++ b/rust/lance/src/dataset/overlay_compaction.rs @@ -0,0 +1,1145 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Overlay-aware compaction (the "Compaction" section of the Data Overlay Files +//! specification). +//! +//! Overlays accumulate read cost — every overlay is a bitmap to test and a file +//! to open — and an overlay newer than an index leaves that index serving stale +//! values until a query re-evaluates the covered rows on the flat path. This +//! module is aware of that per-fragment state and compacts it in one of two +//! modes: +//! +//! - **Overlay → overlay merge.** Collapse a fragment's overlays into a single, +//! smaller overlay carrying the per-`(offset, field)` post-image, stamped with +//! the **maximum** input `committed_version` so the exclusion semantics are +//! preserved. The base and every index are untouched. Cheap; bounds read cost. +//! - **Overlay → base fold.** Materialize the post-image of every covered cell +//! into a fresh base data file, tombstone the folded fields in the old files, +//! and clear the fragment's `overlays`. Row addresses are preserved (a column +//! rewrite, not a row rewrite). Because the fold removes the overlay that was +//! excluding the covered rows from any index on a folded field, the same commit +//! drops those fragments from the index's coverage so they fall to the flat +//! path — the index can never be left serving stale values. +//! +//! The rewrite reads each required base/overlay cell at most once: overlay value +//! columns are read only at the **ranks that win** (a cell a newer overlay already +//! supplies is never fetched from an older overlay), via +//! [`FileFragment::read_overlay_field_winners`]. Base columns are read whole, +//! which is the cheaper coalesced read the specification permits. + +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; + +use arrow_array::ArrayRef; +use object_store::path::Path; +use roaring::RoaringBitmap; +use uuid::Uuid; + +use lance_core::datatypes::Schema; +use lance_core::{Error, Result}; +use lance_file::writer::{FileWriter, FileWriterOptions}; +use lance_io::utils::CachedFileSize; +use lance_table::format::{DataFile, DataOverlayFile, Fragment, IndexMetadata, OverlayCoverage}; + +use crate::dataset::fragment::{FileFragment, FragReadConfig}; +use crate::dataset::overlay::{assemble_overlay_column, route_overlays}; +use crate::dataset::transaction::{Operation, UpdateMode}; +use crate::dataset::{DATA_DIR, Dataset, WriteDestination}; +use crate::index::DatasetIndexExt; + +/// Thresholds the scheduler uses to pick a compaction mode per fragment. +#[derive(Debug, Clone)] +pub struct OverlayCompactionOptions { + /// Fold a fragment to base when an index built on one of its overlaid fields + /// is at least this many dataset versions behind the overlay that made it + /// stale (the `committed_version - index.dataset_version` gap). A fold + /// materializes post-images and reconciles the index, so `1` means "fold as + /// soon as any index is stale with respect to an overlay". `0` disables + /// fold-on-staleness. + pub fold_index_version_gap: u64, + /// Merge a fragment's overlays (overlay → overlay) once it accumulates at + /// least this many, to bound per-read overlay cost. Only applies when the + /// fragment is not already being folded. + pub merge_overlay_count: usize, +} + +impl Default for OverlayCompactionOptions { + fn default() -> Self { + Self { + fold_index_version_gap: 1, + merge_overlay_count: 4, + } + } +} + +/// The action the scheduler selects for a fragment. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum OverlayCompactionMode { + /// Leave the fragment's overlays as they are. + Skip, + /// Collapse the overlays into a single, smaller overlay (base/indexes untouched). + Merge, + /// Materialize overlays into the base and reconcile affected indexes. + Fold, +} + +/// Per-fragment overlay state, the input to the scheduling decision. +#[derive(Debug, Clone)] +pub struct OverlayFragmentState { + pub fragment_id: u64, + /// Number of overlays attached to the fragment. + pub overlay_count: usize, + /// Distinct physical offsets covered by at least one overlay (union over fields). + pub covered_offsets: u64, + /// Sum over fields of each field's coverage popcount — the number of + /// `(offset, field)` cells the overlays supply. + pub covered_cells: u64, + /// Dataset field ids touched by any overlay on the fragment, sorted ascending. + pub covered_fields: Vec, + /// Smallest `committed_version` among the fragment's overlays. + pub min_committed_version: u64, + /// Largest `committed_version` among the fragment's overlays. + pub max_committed_version: u64, + /// Versions between the current dataset version and the newest overlay — how + /// long the newest overlay has gone un-compacted. + pub base_version_gap: u64, + /// Largest staleness gap to any index that covers this fragment and is built + /// on a field the fragment overlays: + /// `max(overlay.committed_version - index.dataset_version)` over such + /// indexes, or `0` when no index is stale with respect to these overlays. + /// This is the version-gap staleness signal that drives the fold decision. + pub max_index_version_gap: u64, +} + +impl OverlayFragmentState { + /// Pick a mode from the version-gap staleness signal and overlay count. + /// + /// Fold takes precedence over merge: a merge keeps the overlays, so a stale + /// index stays stale (queries keep paying the flat re-evaluation cost) — + /// only a fold can reconcile it. + pub fn choose_mode(&self, options: &OverlayCompactionOptions) -> OverlayCompactionMode { + if self.overlay_count == 0 { + return OverlayCompactionMode::Skip; + } + if options.fold_index_version_gap > 0 + && self.max_index_version_gap >= options.fold_index_version_gap + { + return OverlayCompactionMode::Fold; + } + if self.overlay_count >= options.merge_overlay_count { + return OverlayCompactionMode::Merge; + } + OverlayCompactionMode::Skip + } +} + +/// One fragment's scheduled action. +#[derive(Debug, Clone)] +pub struct OverlayCompactionTask { + pub state: OverlayFragmentState, + pub mode: OverlayCompactionMode, +} + +/// The scheduler's output: every overlaid fragment with the mode chosen for it. +#[derive(Debug, Clone)] +pub struct OverlayCompactionPlan { + pub read_version: u64, + pub tasks: Vec, +} + +impl OverlayCompactionPlan { + /// Tasks that actually do something (mode is not [`OverlayCompactionMode::Skip`]). + pub fn actionable_tasks(&self) -> impl Iterator { + self.tasks + .iter() + .filter(|t| t.mode != OverlayCompactionMode::Skip) + } +} + +/// Outcome of a [`compact_overlays`] run. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct OverlayCompactionMetrics { + pub fragments_folded: usize, + pub fragments_merged: usize, + /// Total overlays removed across folded and merged fragments. A fold removes + /// all of a fragment's overlays; a merge removes all but the single overlay + /// it leaves behind. + pub overlays_removed: usize, +} + +/// Compute the overlay state of a single fragment for the scheduler. +fn overlay_fragment_state( + fragment: &Fragment, + indices: &[IndexMetadata], + current_version: u64, +) -> Result { + let mut covered_union = RoaringBitmap::new(); + let mut covered_cells = 0u64; + let mut covered_fields: Vec = Vec::new(); + // Newest overlay version per field, to measure index staleness per field. + let mut field_max_version: HashMap = HashMap::new(); + let mut min_v = u64::MAX; + let mut max_v = 0u64; + + for overlay in &fragment.overlays { + min_v = min_v.min(overlay.committed_version); + max_v = max_v.max(overlay.committed_version); + for (field_pos, &field_id) in overlay.data_file.fields.iter().enumerate() { + if field_id < 0 { + continue; + } + let coverage = overlay.coverage_for_field(field_pos)?; + covered_cells += coverage.len(); + covered_union |= &*coverage; + if !covered_fields.contains(&field_id) { + covered_fields.push(field_id); + } + let entry = field_max_version.entry(field_id).or_insert(0); + *entry = (*entry).max(overlay.committed_version); + } + } + covered_fields.sort_unstable(); + + let mut max_index_version_gap = 0u64; + for index in indices { + let covers_fragment = index + .fragment_bitmap + .as_ref() + .is_some_and(|b| b.contains(fragment.id as u32)); + if !covers_fragment { + continue; + } + for field_id in &index.fields { + if let Some(&overlay_version) = field_max_version.get(field_id) + && overlay_version > index.dataset_version + { + max_index_version_gap = + max_index_version_gap.max(overlay_version - index.dataset_version); + } + } + } + + Ok(OverlayFragmentState { + fragment_id: fragment.id, + overlay_count: fragment.overlays.len(), + covered_offsets: covered_union.len(), + covered_cells, + covered_fields, + min_committed_version: if min_v == u64::MAX { 0 } else { min_v }, + max_committed_version: max_v, + base_version_gap: current_version.saturating_sub(max_v), + max_index_version_gap, + }) +} + +/// Plan overlay compaction: gather each overlaid fragment's state and assign a +/// mode. Reads no data files — only manifest metadata and index descriptors. +pub async fn plan_overlay_compaction( + dataset: &Dataset, + options: &OverlayCompactionOptions, +) -> Result { + let indices = dataset.load_indices().await?; + let current_version = dataset.manifest().version; + + let mut tasks = Vec::new(); + for fragment in dataset.get_fragments() { + let metadata = fragment.metadata(); + if metadata.overlays.is_empty() { + continue; + } + let state = overlay_fragment_state(metadata, &indices, current_version)?; + let mode = state.choose_mode(options); + tasks.push(OverlayCompactionTask { state, mode }); + } + + Ok(OverlayCompactionPlan { + read_version: current_version, + tasks, + }) +} + +/// Plan and execute overlay compaction, committing the result. +/// +/// Merges and folds are committed separately: a merge must leave indexes +/// untouched, so it is committed with no modified fields, while a fold lists its +/// folded fields so the affected fragments drop out of the covering indexes. +pub async fn compact_overlays( + dataset: &mut Dataset, + options: &OverlayCompactionOptions, +) -> Result { + let plan = plan_overlay_compaction(dataset, options).await?; + + let mut merge_fragments: Vec = Vec::new(); + let mut fold_fragments: Vec = Vec::new(); + let mut folded_fields: HashSet = HashSet::new(); + let mut metrics = OverlayCompactionMetrics::default(); + + for task in plan.actionable_tasks() { + let fragment = dataset + .get_fragment(task.state.fragment_id as usize) + .ok_or_else(|| { + Error::internal(format!( + "overlay compaction planned fragment {} which no longer exists", + task.state.fragment_id + )) + })?; + match task.mode { + OverlayCompactionMode::Skip => {} + OverlayCompactionMode::Merge => { + let (updated, removed) = merge_fragment_overlays(dataset, &fragment).await?; + metrics.fragments_merged += 1; + metrics.overlays_removed += removed; + merge_fragments.push(updated); + } + OverlayCompactionMode::Fold => { + let (updated, fields) = fold_fragment_overlays(dataset, &fragment).await?; + metrics.fragments_folded += 1; + metrics.overlays_removed += task.state.overlay_count; + folded_fields.extend(fields); + fold_fragments.push(updated); + } + } + } + + if !merge_fragments.is_empty() { + commit_update(dataset, merge_fragments, Vec::new()).await?; + } + if !fold_fragments.is_empty() { + commit_update(dataset, fold_fragments, folded_fields.into_iter().collect()).await?; + } + + Ok(metrics) +} + +/// Commit an in-place column rewrite via [`Operation::Update`] in +/// [`UpdateMode::RewriteColumns`] mode, replacing the given fragments. The +/// fragments are taken as-is (so cleared/merged `overlays` and tombstoned fields +/// are preserved), and `fields_modified` drops the touched fragments from any +/// index covering those fields. +async fn commit_update( + dataset: &mut Dataset, + updated_fragments: Vec, + fields_modified: Vec, +) -> Result<()> { + let read_version = dataset.manifest().version; + let operation = Operation::Update { + removed_fragment_ids: Vec::new(), + updated_fragments, + new_fragments: Vec::new(), + fields_modified, + merged_generations: Vec::new(), + fields_for_preserving_frag_bitmap: Vec::new(), + update_mode: Some(UpdateMode::RewriteColumns), + inserted_rows_filter: None, + updated_fragment_offsets: None, + }; + let committed = Dataset::commit( + WriteDestination::Dataset(Arc::new(dataset.clone())), + operation, + Some(read_version), + None, + None, + dataset.session.clone(), + false, + ) + .await?; + *dataset = committed; + Ok(()) +} + +/// Merge a fragment's overlays into a single overlay (overlay → overlay). +/// +/// Returns the updated fragment (with `overlays` replaced by the single merged +/// overlay) and the number of overlays removed (`overlay_count - 1`). +async fn merge_fragment_overlays( + dataset: &Dataset, + fragment: &FileFragment, +) -> Result<(Fragment, usize)> { + let metadata = fragment.metadata(); + let original_count = metadata.overlays.len(); + // Preserve the maximum input committed_version so exclusion semantics hold. + let merged_version = metadata + .overlays + .iter() + .map(|o| o.committed_version) + .max() + .ok_or_else(|| { + Error::internal("merge_fragment_overlays called on a fragment with no overlays") + })?; + + let schema = dataset.schema(); + let read_config = FragReadConfig::default(); + + let field_ids = covered_field_ids(metadata)?; + let mut fields = Vec::with_capacity(field_ids.len()); + let mut value_columns: Vec = Vec::with_capacity(field_ids.len()); + let mut per_field_coverage: Vec = Vec::with_capacity(field_ids.len()); + for field_id in &field_ids { + let field = schema + .field_by_id(*field_id) + .ok_or_else(|| Error::internal(format!("overlay field {field_id} not in schema")))?; + let winners = fragment + .read_overlay_field_winners(field, &read_config) + .await?; + fields.push(field.clone()); + value_columns.push(winners.values); + per_field_coverage.push(winners.coverage); + } + + let overlay_schema = Schema { + fields, + metadata: Default::default(), + }; + let data_file = write_data_file(dataset, &overlay_schema, value_columns).await?; + + // Use a single shared bitmap when every field covers the same offsets, else + // store one bitmap per field (a sparse overlay). + let coverage = if per_field_coverage.windows(2).all(|w| w[0] == w[1]) { + OverlayCoverage::dense(per_field_coverage[0].clone()) + } else { + OverlayCoverage::sparse(per_field_coverage) + }; + + let merged_overlay = DataOverlayFile { + data_file, + coverage, + committed_version: merged_version, + }; + + let mut updated = metadata.clone(); + updated.overlays = vec![merged_overlay]; + Ok((updated, original_count - 1)) +} + +/// Fold a fragment's overlays into a fresh base data file (overlay → base). +/// +/// Returns the updated fragment (new file added, folded fields tombstoned in the +/// old files, `overlays` cleared) and the folded field ids. +async fn fold_fragment_overlays( + dataset: &Dataset, + fragment: &FileFragment, +) -> Result<(Fragment, Vec)> { + let metadata = fragment.metadata(); + let schema = dataset.schema(); + let read_config = FragReadConfig::default(); + + let field_ids = covered_field_ids(metadata)?; + let mut fields = Vec::with_capacity(field_ids.len()); + let mut columns: Vec = Vec::with_capacity(field_ids.len()); + for field_id in &field_ids { + let field = schema + .field_by_id(*field_id) + .ok_or_else(|| Error::internal(format!("overlay field {field_id} not in schema")))?; + // Read each winning overlay cell once; read the base column whole. + let winners = fragment + .read_overlay_field_winners(field, &read_config) + .await?; + let base = fragment.read_base_field_full(field, &read_config).await?; + // The base column holds every physical row in order, so row `i` is + // physical offset `i`. `winners.values` is already the field's post-image + // indexed by rank in `winners.coverage`, so routing the full contiguous + // range against that coverage needs exactly those values in order. + let offsets: Vec = (0..base.len() as u32).collect(); + let routing = route_overlays(&offsets, &[&winners.coverage]); + let folded = assemble_overlay_column(&base, &routing, &[winners.values])?; + fields.push(field.clone()); + columns.push(folded); + } + + let fold_schema = Schema { + fields, + metadata: Default::default(), + }; + let new_file = write_data_file(dataset, &fold_schema, columns).await?; + + let folded_set: HashSet = field_ids.iter().copied().collect(); + let mut updated = metadata.clone(); + let new_file_idx = updated.files.len(); + updated.files.push(new_file); + for (idx, file) in updated.files.iter_mut().enumerate() { + if idx == new_file_idx { + continue; + } + let tombstoned: Arc<[i32]> = file + .fields + .iter() + .map(|&id| if folded_set.contains(&id) { -2 } else { id }) + .collect::>() + .into(); + file.fields = tombstoned; + } + updated.overlays.clear(); + + let folded_fields = field_ids.iter().map(|&id| id as u32).collect(); + Ok((updated, folded_fields)) +} + +/// The non-tombstoned dataset field ids that any overlay on the fragment touches, +/// sorted ascending. +fn covered_field_ids(fragment: &Fragment) -> Result> { + let mut ids: Vec = Vec::new(); + for overlay in &fragment.overlays { + for &field_id in overlay.data_file.fields.iter() { + if field_id >= 0 && !ids.contains(&field_id) { + ids.push(field_id); + } + } + } + ids.sort_unstable(); + Ok(ids) +} + +/// Write `columns` (in `schema` field order, lengths may differ for a sparse +/// overlay) to a new data file under the dataset's `data/` directory and return +/// the resulting [`DataFile`] with its field/column index mapping populated. +async fn write_data_file( + dataset: &Dataset, + schema: &Schema, + columns: Vec, +) -> Result { + let version = dataset + .manifest() + .data_storage_format + .lance_file_version()?; + let filename = format!("{}.lance", Uuid::new_v4()); + let path: Path = dataset.base.clone().join(DATA_DIR).join(filename.as_str()); + let object_writer = dataset.object_store.create(&path).await?; + let mut writer = FileWriter::try_new( + object_writer, + schema.clone(), + FileWriterOptions { + format_version: Some(version), + ..Default::default() + }, + )?; + let (major, minor) = writer.version().to_numbers(); + for (column_index, array) in columns.into_iter().enumerate() { + writer.write_column(column_index, array).await?; + } + let summary = writer.finish().await?; + + 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); + Ok(data_file) +} + +#[cfg(test)] +mod tests { + use super::*; + + use std::collections::BTreeMap; + + use arrow_array::{Array, Int32Array, RecordBatch, RecordBatchIterator}; + use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema}; + use futures::TryStreamExt; + use lance_core::utils::tempfile::TempStrDir; + use lance_file::version::LanceFileVersion; + use lance_index::{IndexType, scalar::ScalarIndexParams}; + use lance_io::{assert_io_lt, utils::tracking_store::IoStats}; + use uuid::Uuid; + + use crate::dataset::transaction::{DataOverlayGroup, Operation}; + use crate::dataset::{WriteDestination, WriteParams}; + use crate::index::DatasetIndexExt; + + fn bitmap(offsets: impl IntoIterator) -> RoaringBitmap { + RoaringBitmap::from_iter(offsets) + } + + fn i32_array(values: impl IntoIterator>) -> ArrayRef { + Arc::new(Int32Array::from_iter(values)) + } + + // --------------------------------------------------------------------- + // Scheduler / awareness unit tests (synthetic metadata, no I/O) + // --------------------------------------------------------------------- + + fn synth_overlay(fields: &[i32], offsets: &[u32], committed_version: u64) -> DataOverlayFile { + DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("o.lance", fields.to_vec(), None), + coverage: OverlayCoverage::dense(bitmap(offsets.iter().copied())), + committed_version, + } + } + + fn synth_index( + name: &str, + fields: &[i32], + dataset_version: u64, + frags: &[u32], + ) -> IndexMetadata { + IndexMetadata { + uuid: Uuid::new_v4(), + fields: fields.to_vec(), + name: name.to_string(), + dataset_version, + fragment_bitmap: Some(bitmap(frags.iter().copied())), + index_details: None, + index_version: 0, + created_at: None, + base_id: None, + files: None, + } + } + + #[test] + fn test_state_version_gaps_and_cell_counts() { + let mut fragment = Fragment::new(0).with_physical_rows(6); + // v3 overlays field 1 over {1,2}; v5 overlays fields {1,2} over {3}. + fragment.overlays = vec![ + synth_overlay(&[1], &[1, 2], 3), + synth_overlay(&[1, 2], &[3], 5), + ]; + // Index on field 1 built at version 2, covering fragment 0. + let indices = vec![synth_index("val_idx", &[1], 2, &[0])]; + + let state = overlay_fragment_state(&fragment, &indices, 7).unwrap(); + assert_eq!(state.overlay_count, 2); + assert_eq!(state.covered_fields, vec![1, 2]); + // cells = v3{field1: 2} + v5{field1: 1, field2: 1} = 4. + assert_eq!(state.covered_cells, 4); + // distinct physical offsets = {1,2,3}. + assert_eq!(state.covered_offsets, 3); + assert_eq!(state.min_committed_version, 3); + assert_eq!(state.max_committed_version, 5); + assert_eq!(state.base_version_gap, 2); // current 7 - newest 5. + // field 1's newest overlay (v5) vs index built at v2 -> gap 3. + assert_eq!(state.max_index_version_gap, 3); + } + + #[test] + fn test_mode_fold_when_index_is_stale() { + let mut fragment = Fragment::new(0).with_physical_rows(6); + fragment.overlays = vec![synth_overlay(&[1], &[1], 5)]; + let indices = vec![synth_index("val_idx", &[1], 2, &[0])]; + let state = overlay_fragment_state(&fragment, &indices, 5).unwrap(); + assert_eq!( + state.choose_mode(&OverlayCompactionOptions::default()), + OverlayCompactionMode::Fold + ); + } + + #[test] + fn test_mode_field_aware_no_fold_for_unrelated_index() { + let mut fragment = Fragment::new(0).with_physical_rows(6); + // Overlay touches field 1 only. + fragment.overlays = vec![synth_overlay(&[1], &[1], 5)]; + // Index is on field 2 (unrelated) -> not stale w.r.t. this overlay. + let indices = vec![synth_index("other_idx", &[2], 2, &[0])]; + let state = overlay_fragment_state(&fragment, &indices, 5).unwrap(); + assert_eq!(state.max_index_version_gap, 0); + // No stale index and only one overlay -> nothing to do. + assert_eq!( + state.choose_mode(&OverlayCompactionOptions::default()), + OverlayCompactionMode::Skip + ); + } + + #[test] + fn test_mode_merge_on_overlay_count_without_stale_index() { + let mut fragment = Fragment::new(0).with_physical_rows(6); + fragment.overlays = (1..=4).map(|v| synth_overlay(&[1], &[1], v)).collect(); + // Index already current (built after every overlay) -> not stale. + let indices = vec![synth_index("val_idx", &[1], 10, &[0])]; + let state = overlay_fragment_state(&fragment, &indices, 10).unwrap(); + assert_eq!(state.max_index_version_gap, 0); + let options = OverlayCompactionOptions { + fold_index_version_gap: 1, + merge_overlay_count: 4, + }; + assert_eq!(state.choose_mode(&options), OverlayCompactionMode::Merge); + } + + #[test] + fn test_mode_skip_below_thresholds() { + let mut fragment = Fragment::new(0).with_physical_rows(6); + fragment.overlays = vec![synth_overlay(&[1], &[1], 3), synth_overlay(&[1], &[2], 4)]; + let indices = vec![synth_index("val_idx", &[1], 10, &[0])]; + let state = overlay_fragment_state(&fragment, &indices, 10).unwrap(); + assert_eq!( + state.choose_mode(&OverlayCompactionOptions::default()), + OverlayCompactionMode::Skip + ); + } + + // --------------------------------------------------------------------- + // End-to-end execution tests + // --------------------------------------------------------------------- + + /// Two-fragment Int32 dataset: `id` (field 0) = 0..12 and `val` (field 1) = + /// id * 10, six rows per file (fragments 0 and 1). + async fn create_base_dataset(uri: &str) -> 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(LanceFileVersion::Stable), + ..Default::default() + }; + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + Dataset::write(reader, uri, Some(write_params)) + .await + .unwrap() + } + + /// Write an overlay file covering `fields` of `fragment_id` with the given + /// coverage and per-field value columns, then commit it as a `DataOverlay`. + async fn commit_overlay( + dataset: Dataset, + fragment_id: u64, + fields: &[i32], + coverage: OverlayCoverage, + columns: Vec, + ) -> Dataset { + let read_version = dataset.version().version; + let overlay_schema = dataset.schema().project_by_ids(fields, true); + let filename = format!("{}.lance", Uuid::new_v4()); + let path = dataset.base.clone().join(DATA_DIR).join(filename.as_str()); + let obj_writer = dataset.object_store.create(&path).await.unwrap(); + let mut writer = FileWriter::try_new( + obj_writer, + overlay_schema, + FileWriterOptions { + format_version: Some(LanceFileVersion::Stable), + ..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(|(f, _)| *f as i32) + .collect::>() + .into(); + data_file.column_indices = writer + .field_id_to_column_indices() + .iter() + .map(|(_, c)| *c 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() + } + + /// Scan `id` and `val` and return an `id -> val` map (order-independent). + async fn id_val_map(dataset: &Dataset) -> BTreeMap> { + let mut scanner = dataset.scan(); + scanner.project(&["id", "val"]).unwrap(); + let batches = scanner + .try_into_stream() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let mut out = BTreeMap::new(); + for batch in batches { + let ids = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let vals = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + for i in 0..batch.num_rows() { + let v = if vals.is_null(i) { + None + } else { + Some(vals.value(i)) + }; + out.insert(ids.value(i), v); + } + } + out + } + + #[tokio::test] + async fn test_fold_materializes_overlay_and_clears_overlays() { + let dataset = create_base_dataset("memory://").await; + // Overlay fragment 0: val[1] 10 -> 999, val[4] 40 -> 444. + let dataset = commit_overlay( + dataset, + 0, + &[1], + OverlayCoverage::dense(bitmap([1, 4])), + vec![i32_array([Some(999), Some(444)])], + ) + .await; + + let mut dataset = dataset; + let fragment = dataset.get_fragment(0).unwrap(); + let (updated, fields) = fold_fragment_overlays(&dataset, &fragment).await.unwrap(); + assert_eq!(fields, vec![1]); + assert!(updated.overlays.is_empty()); + commit_update(&mut dataset, vec![updated], fields) + .await + .unwrap(); + + // Overlays gone; values materialized into the base. + assert!( + dataset + .get_fragment(0) + .unwrap() + .metadata() + .overlays + .is_empty() + ); + let map = id_val_map(&dataset).await; + assert_eq!(map[&1], Some(999)); + assert_eq!(map[&4], Some(444)); + assert_eq!(map[&0], Some(0)); + assert_eq!(map[&2], Some(20)); + // Fragment 1 untouched. + assert_eq!(map[&7], Some(70)); + } + + #[tokio::test] + async fn test_fold_null_override() { + let dataset = create_base_dataset("memory://").await; + // Overlay sets val[2] to NULL (covered + null overrides to null). + let dataset = commit_overlay( + dataset, + 0, + &[1], + OverlayCoverage::dense(bitmap([2])), + vec![i32_array([None])], + ) + .await; + + let mut dataset = dataset; + let fragment = dataset.get_fragment(0).unwrap(); + let (updated, fields) = fold_fragment_overlays(&dataset, &fragment).await.unwrap(); + commit_update(&mut dataset, vec![updated], fields) + .await + .unwrap(); + + let map = id_val_map(&dataset).await; + assert_eq!(map[&2], None); // overridden to null + assert_eq!(map[&1], Some(10)); // untouched + } + + #[tokio::test] + async fn test_fold_multi_fragment_multi_overlay() { + let dataset = create_base_dataset("memory://").await; + // Two overlays on fragment 0 (newest wins on the shared offset 1). + let dataset = commit_overlay( + dataset, + 0, + &[1], + OverlayCoverage::dense(bitmap([1, 2])), + vec![i32_array([Some(100), Some(200)])], + ) + .await; + let dataset = commit_overlay( + dataset, + 0, + &[1], + OverlayCoverage::dense(bitmap([1])), + vec![i32_array([Some(111)])], + ) + .await; + // One overlay on fragment 1: offset 0 is id 6 -> val 600. + let dataset = commit_overlay( + dataset, + 1, + &[1], + OverlayCoverage::dense(bitmap([0])), + vec![i32_array([Some(600)])], + ) + .await; + + let mut dataset = dataset; + let frag0 = dataset.get_fragment(0).unwrap(); + let frag1 = dataset.get_fragment(1).unwrap(); + let (u0, f0) = fold_fragment_overlays(&dataset, &frag0).await.unwrap(); + let (u1, f1) = fold_fragment_overlays(&dataset, &frag1).await.unwrap(); + let mut fields = f0; + fields.extend(f1); + fields.sort_unstable(); + fields.dedup(); + commit_update(&mut dataset, vec![u0, u1], fields) + .await + .unwrap(); + + let map = id_val_map(&dataset).await; + assert_eq!(map[&1], Some(111)); // newest overlay won + assert_eq!(map[&2], Some(200)); + assert_eq!(map[&6], Some(600)); + assert_eq!(map[&3], Some(30)); // untouched + for frag in dataset.get_fragments() { + assert!(frag.metadata().overlays.is_empty()); + } + } + + #[tokio::test] + async fn test_fold_reconciles_stale_index() { + let test_dir = TempStrDir::default(); + let mut dataset = create_base_dataset(&test_dir).await; + dataset + .create_index( + &["val"], + IndexType::Scalar, + None, + &ScalarIndexParams::default(), + true, + ) + .await + .unwrap(); + + // Overlay val[1] 10 -> 999 (committed after the index -> index is stale). + dataset = commit_overlay( + dataset, + 0, + &[1], + OverlayCoverage::dense(bitmap([1])), + vec![i32_array([Some(999)])], + ) + .await; + + // The scheduler should fold fragment 0 to reconcile the stale index. + let plan = plan_overlay_compaction(&dataset, &OverlayCompactionOptions::default()) + .await + .unwrap(); + let frag0_task = plan + .tasks + .iter() + .find(|t| t.state.fragment_id == 0) + .unwrap(); + assert_eq!(frag0_task.mode, OverlayCompactionMode::Fold); + assert!(frag0_task.state.max_index_version_gap >= 1); + + let metrics = compact_overlays(&mut dataset, &OverlayCompactionOptions::default()) + .await + .unwrap(); + assert_eq!(metrics.fragments_folded, 1); + + // Index reconciled: fragment 0 dropped from the val index's coverage. + let indices = dataset.load_indices().await.unwrap(); + let val_index = indices + .iter() + .find(|i| i.fields == vec![1]) + .expect("val index present"); + assert!( + !val_index.fragment_bitmap.as_ref().unwrap().contains(0), + "folded fragment must be removed from the stale index's coverage" + ); + + // And the query is correct: val = 999 finds id 1, the stale 10 is gone. + let mut scanner = dataset.scan(); + scanner + .filter("val = 999") + .unwrap() + .project(&["id"]) + .unwrap(); + let batch = scanner.try_into_batch().await.unwrap(); + let ids = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(ids.len(), 1); + assert_eq!(ids.value(0), 1); + assert!( + dataset + .get_fragment(0) + .unwrap() + .metadata() + .overlays + .is_empty() + ); + } + + #[tokio::test] + async fn test_merge_collapses_overlays_preserving_max_version() { + let dataset = create_base_dataset("memory://").await; + // Three overlays on fragment 0; newest (highest committed_version) wins. + let dataset = commit_overlay( + dataset, + 0, + &[1], + OverlayCoverage::dense(bitmap([1, 2])), + vec![i32_array([Some(100), Some(200)])], + ) + .await; + let dataset = commit_overlay( + dataset, + 0, + &[1], + OverlayCoverage::dense(bitmap([1, 3])), + vec![i32_array([Some(111), Some(333)])], + ) + .await; + let dataset = commit_overlay( + dataset, + 0, + &[1], + OverlayCoverage::dense(bitmap([4])), + vec![i32_array([Some(444)])], + ) + .await; + let expected_version = dataset.version().version; // newest overlay's commit version + + let mut dataset = dataset; + let fragment = dataset.get_fragment(0).unwrap(); + let (updated, removed) = merge_fragment_overlays(&dataset, &fragment).await.unwrap(); + assert_eq!(removed, 2); // 3 overlays -> 1 + assert_eq!(updated.overlays.len(), 1); + let merged = &updated.overlays[0]; + assert_eq!(merged.committed_version, expected_version); + // Union coverage over all three overlays for field 1. + assert_eq!(*merged.coverage_for_field(0).unwrap(), bitmap([1, 2, 3, 4])); + + commit_update(&mut dataset, vec![updated], Vec::new()) + .await + .unwrap(); + // The single merged overlay reproduces the newest-wins post-image on read. + let map = id_val_map(&dataset).await; + assert_eq!(map[&1], Some(111)); // second overlay newest for offset 1 + assert_eq!(map[&2], Some(200)); + assert_eq!(map[&3], Some(333)); + assert_eq!(map[&4], Some(444)); + assert_eq!(map[&0], Some(0)); // untouched + } + + /// Build a single-fragment dataset of `n` rows (`val` = field 1) on a + /// tracking local store, attach two overlays over `val` with the given + /// coverage, merge, and return the overlay value bytes/iops the merge read. + async fn merge_value_io(n: i32, cov_a: RoaringBitmap, cov_b: RoaringBitmap) -> IoStats { + let test_dir = TempStrDir::default(); + 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..n)), + Arc::new(Int32Array::from_iter_values((0..n).map(|v| v * 10))), + ], + ) + .unwrap(); + let write_params = WriteParams { + max_rows_per_file: n as usize, + max_rows_per_group: n as usize, + data_storage_version: Some(LanceFileVersion::Stable), + ..Default::default() + }; + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + let dataset = Dataset::write(reader, &test_dir, Some(write_params)) + .await + .unwrap(); + + let vals_a: Vec> = cov_a.iter().map(|o| Some(o as i32)).collect(); + let vals_b: Vec> = cov_b.iter().map(|o| Some(-(o as i32) - 1)).collect(); + let dataset = commit_overlay( + dataset, + 0, + &[1], + OverlayCoverage::dense(cov_a), + vec![i32_array(vals_a)], + ) + .await; + let dataset = commit_overlay( + dataset, + 0, + &[1], + OverlayCoverage::dense(cov_b), + vec![i32_array(vals_b)], + ) + .await; + + let fragment = dataset.get_fragment(0).unwrap(); + // Reset counters, then measure only what the merge reads. + let _ = dataset.object_store.io_stats_incremental(); + let _ = merge_fragment_overlays(&dataset, &fragment).await.unwrap(); + dataset.object_store.io_stats_incremental() + } + + /// Read-each-cell-at-most-once: with equal total coverage, overlapping + /// overlays read strictly fewer value bytes than disjoint ones, because a + /// cell a newer overlay supplies is not also read from the older overlay. + /// The disjoint case is the negative control that makes the bound meaningful. + #[tokio::test] + async fn test_merge_reads_each_winning_cell_at_most_once() { + const N: i32 = 8000; + // Both overlays cover 4000 offsets, so per-file fixed costs match. + // Overlapping: cov_a = [0, 4000), cov_b = [2000, 6000). union = 6000; + // offsets [2000, 4000) are read once, from the newer overlay b only. + let overlap = merge_value_io(N, bitmap(0..4000), bitmap(2000..6000)).await; + // Disjoint negative control: cov_a = [0, 4000), cov_b = [4000, 8000). + // Same per-overlay sizes, but union = 8000 distinct cells. + let disjoint = merge_value_io(N, bitmap(0..4000), bitmap(4000..8000)).await; + + // The only difference between the two is how many distinct cells are + // fetched (6000 vs 8000). If overlapping cells were double-read, the + // overlap case would not read strictly fewer value bytes. + assert!( + overlap.read_bytes > 0 && disjoint.read_bytes > 0, + "both merges must read overlay values" + ); + assert_io_lt!( + overlap, + read_bytes, + disjoint.read_bytes, + "overlapping coverage must not re-read cells a newer overlay already supplies \ + (overlap read {} bytes, disjoint {} bytes)", + overlap.read_bytes, + disjoint.read_bytes + ); + } +}