From fbd230ad9f3e73a54e5d84021606d271fbdc4d82 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Fri, 19 Jun 2026 15:55:38 -0700 Subject: [PATCH 1/9] 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 2/9] 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 3/9] 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 6c3c094d7364e4029be9608b5459a47607a551fb Mon Sep 17 00:00:00 2001 From: Will Jones Date: Mon, 22 Jun 2026 20:06:05 -0700 Subject: [PATCH 4/9] 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 56b93f46c09205f8609beac1dfd3065d079332fe Mon Sep 17 00:00:00 2001 From: Will Jones Date: Mon, 22 Jun 2026 20:09:39 -0700 Subject: [PATCH 5/9] 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 cdcd8e366512018723e1fcd1e58e20238b1111de Mon Sep 17 00:00:00 2001 From: Will Jones Date: Mon, 22 Jun 2026 20:17:41 -0700 Subject: [PATCH 6/9] feat: overlay cell-resolution core (newest-first, rank, NULL override) Adds `dataset::overlay`, the tested heart of reading overlays: given a base column for a physical row range and the overlays covering a field (newest first), `resolve_overlay_column` produces the merged column. An offset is resolved to the newest covering overlay's value at the offset's rank in the coverage bitmap; an uncovered offset falls through to the base; a covered offset whose value is NULL overrides the cell to NULL. `overlay_indices_newest_first` orders a fragment's overlays by `committed_version` then list position. Deletion precedence needs no handling here: the merge runs before the deletion filter, so an overlay value for a deleted offset is computed and dropped with the row. Wiring this into the scan stream and `take` follows on this branch. Unit tests cover rank addressing, multi-overlay precedence, NULL override vs. fall-through, physical-offset base, string columns, and ordering. Part of OSS-1322. Co-Authored-By: Claude Opus 4.8 (1M context) --- rust/lance/src/dataset.rs | 1 + rust/lance/src/dataset/overlay.rs | 242 ++++++++++++++++++++++++++++++ 2 files changed, 243 insertions(+) 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/overlay.rs b/rust/lance/src/dataset/overlay.rs new file mode 100644 index 00000000000..418a2975bd1 --- /dev/null +++ b/rust/lance/src/dataset/overlay.rs @@ -0,0 +1,242 @@ +// 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 physical row range, 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. +//! +//! 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. + +// This module is the tested core of overlay resolution. It is consumed by the +// scan and `take` merge paths landing in follow-up commits on this branch; until +// then its items are exercised only by unit tests. +#![allow(dead_code)] + +use arrow_array::{Array, ArrayRef}; +use arrow_select::interleave::interleave; +use lance_core::{Error, Result}; +use roaring::RoaringBitmap; +use snafu::location; + +use lance_table::format::DataOverlayFile; + +/// One field's contribution from a single overlay: which physical offsets it +/// covers, and the value column holding those offsets' values (indexed by rank). +#[derive(Debug, Clone)] +pub struct ResolvedFieldOverlay { + /// Physical offsets this overlay covers for the field. + pub coverage: RoaringBitmap, + /// The overlay's value column for the field. Its length must equal + /// `coverage.len()`; the value for a covered offset `o` is at `coverage`'s + /// rank of `o`. + pub values: ArrayRef, +} + +/// 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 +} + +/// Resolve a single field's values for the physical row range +/// `[physical_start, physical_start + base.len())`, merging the overlays that +/// cover the field (which must be supplied newest-first). +/// +/// The result has the same length and data type as `base`. A covered offset +/// whose overlay value is NULL resolves **to** NULL (distinct from an offset no +/// overlay covers, which keeps its base value). +pub fn resolve_overlay_column( + base: &ArrayRef, + physical_start: u64, + overlays_newest_first: &[ResolvedFieldOverlay], +) -> Result { + if overlays_newest_first.is_empty() { + return Ok(base.clone()); + } + for (i, overlay) in overlays_newest_first.iter().enumerate() { + if overlay.values.len() as u64 != overlay.coverage.len() { + return Err(Error::invalid_input(format!( + "overlay value column {} has {} values but its coverage has {} offsets", + i, + overlay.values.len(), + overlay.coverage.len() + ))); + } + } + + // Source 0 is the base; source k+1 is overlays_newest_first[k].values. + let mut sources: Vec<&dyn Array> = Vec::with_capacity(overlays_newest_first.len() + 1); + sources.push(base.as_ref()); + for overlay in overlays_newest_first { + sources.push(overlay.values.as_ref()); + } + + let indices: Vec<(usize, usize)> = (0..base.len()) + .map(|i| { + let offset = physical_start + i as u64; + for (k, overlay) in overlays_newest_first.iter().enumerate() { + if overlay.coverage.contains(offset as u32) { + // 0-based rank: number of set bits strictly below `offset`. + let rank = overlay.coverage.rank(offset as u32) as usize - 1; + return (k + 1, rank); + } + } + (0, i) + }) + .collect(); + + interleave(&sources, &indices).map_err(|e| Error::Arrow { + message: format!("failed to merge overlay column: {e}"), + location: location!(), + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow_array::{Int32Array, StringArray}; + 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) + } + + 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_overlay_column(&base, 0, &[]).unwrap(); + 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 = ResolvedFieldOverlay { + coverage: bitmap([1]), + values: i32_array([Some(26)]), + }; + let resolved = resolve_overlay_column(&base, 0, &[overlay]).unwrap(); + 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 = ResolvedFieldOverlay { + coverage: bitmap([0, 2, 3]), + values: i32_array([Some(100), Some(120), Some(130)]), + }; + let resolved = resolve_overlay_column(&base, 0, &[overlay]).unwrap(); + 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 = ResolvedFieldOverlay { + coverage: bitmap([1]), + values: i32_array([Some(999)]), + }; + let older = ResolvedFieldOverlay { + coverage: bitmap([1, 2]), + values: i32_array([Some(111), Some(222)]), + }; + let resolved = resolve_overlay_column(&base, 0, &[newest, older]).unwrap(); + // 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 = ResolvedFieldOverlay { + coverage: bitmap([0]), + values: i32_array([None]), + }; + let resolved = resolve_overlay_column(&base, 0, &[overlay]).unwrap(); + 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 = ResolvedFieldOverlay { + coverage: bitmap([11]), + values: i32_array([Some(7)]), + }; + let resolved = resolve_overlay_column(&base, 10, &[overlay]).unwrap(); + 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 = ResolvedFieldOverlay { + coverage: bitmap([0, 2]), + values: Arc::new(StringArray::from(vec!["A", "C"])), + }; + let resolved = resolve_overlay_column(&base, 0, &[overlay]).unwrap(); + let expected: ArrayRef = Arc::new(StringArray::from(vec!["A", "b", "C"])); + assert_eq!(&resolved, &expected); + } + + #[test] + fn test_value_count_mismatch_errors() { + let base = i32_array([Some(1), Some(2)]); + let overlay = ResolvedFieldOverlay { + coverage: bitmap([0, 1]), + values: i32_array([Some(9)]), // only one value for two covered offsets + }; + assert!(resolve_overlay_column(&base, 0, &[overlay]).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::Shared(vec![]), + 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 1f06d07c93678fee68de9ffd7327c7190256aaa7 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Mon, 22 Jun 2026 19:37:13 -0700 Subject: [PATCH 7/9] 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 3d32d25cf969b434b4fdac5d233dca9a30e71ff9 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Mon, 22 Jun 2026 20:59:24 -0700 Subject: [PATCH 8/9] feat: resolve data overlay files on the take and scan read paths MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wire the overlay cell-resolution core into reads so `take` (and scan) return merged values. `FileFragment::open` loads each projected field's overlay value columns newest-first, and `new_read_impl`/`read_ranges` merge them into base batches by physical offset before deletion filtering — so resolution is identical for take and scan, NULL overrides apply, deletions take precedence (an overlay value for a deleted row is dropped with the row), and fields resolve independently. The merge addresses each row by its physical offset (via `ReadBatchParams::to_offsets_total`), so the resolution core now takes explicit per-row offsets instead of a contiguous start — a single code path for the contiguous scan range and arbitrary take indices. Sparse per-field overlays read each field's value column independently, so unequal-length columns (OSS-1323) are handled without materializing a rectangular batch. Removes the temporary "overlays not supported" guard from OSS-1322. Tests: take of covered/uncovered offsets, multiple overlays (newest-wins), per-field coverage with unequal-length columns, NULL override, overlay on a deleted row (inert), and multi-fragment scan — all over v2.0 and v2.1 files. Co-Authored-By: Claude Opus 4.8 (1M context) --- rust/lance/src/dataset/fragment.rs | 507 ++++++++++++++++++++++++++++- rust/lance/src/dataset/overlay.rs | 183 +++++++++-- 2 files changed, 644 insertions(+), 46 deletions(-) diff --git a/rust/lance/src/dataset/fragment.rs b/rust/lance/src/dataset/fragment.rs index 58a3ce0a9b6..02787863c8b 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; @@ -43,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::{ @@ -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::{ + FieldOverlayPlan, ResolvedFieldOverlay, apply_overlays_to_batch, overlay_indices_newest_first, +}; 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(); @@ -948,6 +942,11 @@ impl FileFragment { Arc::new(self.metadata.clone()), )?; + if !self.metadata.overlays.is_empty() { + reader.overlay_plans = + Arc::new(self.load_overlay_plans(projection, &read_config).await?); + } + if read_config.with_row_id { reader.with_row_id(); } @@ -1123,6 +1122,92 @@ impl FileFragment { Ok(opened_files) } + /// Load the overlay value columns for the projected fields, ordered + /// newest-first, ready to be merged into base batches on read. + /// + /// 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, in which case the field's coverage bitmap and full value + /// column are loaded. Overlays on nested (non-top-level) fields are not yet + /// supported and are simply not matched here. + async fn load_overlay_plans( + &self, + projection: &Schema, + read_config: &FragReadConfig, + ) -> Result> { + let order = overlay_indices_newest_first(&self.metadata.overlays); + let mut plans = Vec::new(); + for field in &projection.fields { + let mut resolved = Vec::new(); + for &overlay_idx in &order { + let overlay = &self.metadata.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 values = self + .read_overlay_value_column(overlay, field, coverage.len(), read_config) + .await?; + resolved.push(ResolvedFieldOverlay { coverage, values }); + } + if !resolved.is_empty() { + plans.push(FieldOverlayPlan { + field_name: field.name.clone(), + overlays_newest_first: resolved, + }); + } + } + Ok(plans) + } + + /// Read a single field's value column fully from an overlay data file. + /// + /// The value column has `num_values` rows (the popcount of the field's + /// coverage); position `r` holds the value for the offset whose rank is `r`. + async fn read_overlay_value_column( + &self, + overlay: &DataOverlayFile, + field: &lance_core::datatypes::Field, + num_values: u64, + read_config: &FragReadConfig, + ) -> Result { + if num_values == 0 { + 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 + .read_range_tasks( + 0..num_values, + num_values 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 { @@ -2052,6 +2137,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 +2170,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 +2238,7 @@ impl FragmentReader { created_at_sequence: None, num_rows, num_physical_rows, + overlay_plans: Arc::new(Vec::new()), }) } @@ -2404,6 +2496,46 @@ 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. + 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 batch = inner.await?; + let slice = &offsets[start..start + batch.num_rows()]; + apply_overlays_to_batch(batch, slice, &plans) + } + .boxed(), + } + }) + .boxed() + } + async fn new_read_impl<'a, F>( &'a self, params: ReadBatchParams, @@ -2471,6 +2603,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 +2775,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 +2790,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 +2986,350 @@ 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(); + let columns: Vec<(usize, ArrayRef)> = columns.into_iter().enumerate().collect(); + writer.write_columns(columns).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(&[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); + } + } + #[rstest] #[tokio::test] async fn test_fragment_scan( diff --git a/rust/lance/src/dataset/overlay.rs b/rust/lance/src/dataset/overlay.rs index 418a2975bd1..907b41f2a14 100644 --- a/rust/lance/src/dataset/overlay.rs +++ b/rust/lance/src/dataset/overlay.rs @@ -4,27 +4,25 @@ //! 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 physical row range, 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. +//! 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. -// This module is the tested core of overlay resolution. It is consumed by the -// scan and `take` merge paths landing in follow-up commits on this branch; until -// then its items are exercised only by unit tests. -#![allow(dead_code)] - -use arrow_array::{Array, ArrayRef}; +use arrow_array::{Array, ArrayRef, RecordBatch}; use arrow_select::interleave::interleave; use lance_core::{Error, Result}; use roaring::RoaringBitmap; -use snafu::location; use lance_table::format::DataOverlayFile; @@ -56,18 +54,25 @@ pub fn overlay_indices_newest_first(overlays: &[DataOverlayFile]) -> Vec indices } -/// Resolve a single field's values for the physical row range -/// `[physical_start, physical_start + base.len())`, merging the overlays that -/// cover the field (which must be supplied newest-first). +/// Resolve a single field's values for the rows whose physical offsets are given +/// by `offsets` (one per base row, in the same order as `base`), merging the +/// overlays that cover the field (which must be supplied newest-first). /// /// The result has the same length and data type as `base`. A covered offset /// whose overlay value is NULL resolves **to** NULL (distinct from an offset no /// overlay covers, which keeps its base value). pub fn resolve_overlay_column( base: &ArrayRef, - physical_start: u64, + offsets: &[u32], overlays_newest_first: &[ResolvedFieldOverlay], ) -> Result { + if offsets.len() != base.len() { + return Err(Error::invalid_input(format!( + "overlay resolution got {} offsets for a base column of {} rows", + offsets.len(), + base.len() + ))); + } if overlays_newest_first.is_empty() { return Ok(base.clone()); } @@ -89,13 +94,14 @@ pub fn resolve_overlay_column( sources.push(overlay.values.as_ref()); } - let indices: Vec<(usize, usize)> = (0..base.len()) - .map(|i| { - let offset = physical_start + i as u64; + let indices: Vec<(usize, usize)> = offsets + .iter() + .enumerate() + .map(|(i, &offset)| { for (k, overlay) in overlays_newest_first.iter().enumerate() { - if overlay.coverage.contains(offset as u32) { + if overlay.coverage.contains(offset) { // 0-based rank: number of set bits strictly below `offset`. - let rank = overlay.coverage.rank(offset as u32) as usize - 1; + let rank = overlay.coverage.rank(offset) as usize - 1; return (k + 1, rank); } } @@ -103,16 +109,53 @@ pub fn resolve_overlay_column( }) .collect(); - interleave(&sources, &indices).map_err(|e| Error::Arrow { - message: format!("failed to merge overlay column: {e}"), - location: location!(), - }) + interleave(&sources, &indices).map_err(Error::from) +} + +/// The overlays that apply to a single projected field, resolved (value columns +/// loaded) and ordered newest-first. `field_name` is the top-level read-batch +/// column name the plan applies to. +#[derive(Debug, Clone)] +pub struct FieldOverlayPlan { + pub field_name: String, + pub overlays_newest_first: Vec, +} + +/// Merge overlay values into a read batch of base values. +/// +/// `offsets[i]` is the physical row offset of `batch` row `i` (as produced by +/// [`lance_io::ReadBatchParams::to_offsets_total`]). Each plan replaces the +/// batch column whose name equals `plan.field_name`; columns with no plan, and +/// the row-id/row-address system columns, pass through unchanged. +/// +/// This runs on physical rows *before* deletion filtering, so an overlay value +/// computed for a deleted row is dropped with the row downstream — giving +/// deletions precedence with no special handling here. +pub fn apply_overlays_to_batch( + batch: RecordBatch, + offsets: &[u32], + plans: &[FieldOverlayPlan], +) -> Result { + if plans.is_empty() { + return Ok(batch); + } + let schema = batch.schema(); + let mut columns = batch.columns().to_vec(); + for plan in plans { + let Some(idx) = schema.index_of(&plan.field_name).ok() else { + // The plan's field is not in this batch's projection; skip it. + continue; + }; + columns[idx] = resolve_overlay_column(&columns[idx], offsets, &plan.overlays_newest_first)?; + } + Ok(RecordBatch::try_new(schema, columns)?) } #[cfg(test)] mod tests { use super::*; use arrow_array::{Int32Array, StringArray}; + use arrow_schema::{DataType, Field, Schema}; use std::sync::Arc; fn i32_array(values: impl IntoIterator>) -> ArrayRef { @@ -123,6 +166,11 @@ mod tests { 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() + } + fn assert_i32_eq(actual: &ArrayRef, expected: impl IntoIterator>) { let actual = actual.as_any().downcast_ref::().unwrap(); assert_eq!(actual, &Int32Array::from_iter(expected)); @@ -131,7 +179,7 @@ mod tests { #[test] fn test_no_overlays_returns_base() { let base = i32_array([Some(1), Some(2), Some(3)]); - let resolved = resolve_overlay_column(&base, 0, &[]).unwrap(); + let resolved = resolve_overlay_column(&base, &offsets(0, 3), &[]).unwrap(); assert_i32_eq(&resolved, [Some(1), Some(2), Some(3)]); } @@ -143,7 +191,7 @@ mod tests { coverage: bitmap([1]), values: i32_array([Some(26)]), }; - let resolved = resolve_overlay_column(&base, 0, &[overlay]).unwrap(); + let resolved = resolve_overlay_column(&base, &offsets(0, 4), &[overlay]).unwrap(); assert_i32_eq(&resolved, [Some(30), Some(26), Some(40), Some(22)]); } @@ -155,7 +203,7 @@ mod tests { coverage: bitmap([0, 2, 3]), values: i32_array([Some(100), Some(120), Some(130)]), }; - let resolved = resolve_overlay_column(&base, 0, &[overlay]).unwrap(); + let resolved = resolve_overlay_column(&base, &offsets(0, 4), &[overlay]).unwrap(); assert_i32_eq(&resolved, [Some(100), Some(11), Some(120), Some(130)]); } @@ -171,7 +219,7 @@ mod tests { coverage: bitmap([1, 2]), values: i32_array([Some(111), Some(222)]), }; - let resolved = resolve_overlay_column(&base, 0, &[newest, older]).unwrap(); + let resolved = resolve_overlay_column(&base, &offsets(0, 3), &[newest, older]).unwrap(); // offset 1 -> newest (999); offset 2 -> only older covers it (222). assert_i32_eq(&resolved, [Some(0), Some(999), Some(222)]); } @@ -185,7 +233,7 @@ mod tests { coverage: bitmap([0]), values: i32_array([None]), }; - let resolved = resolve_overlay_column(&base, 0, &[overlay]).unwrap(); + let resolved = resolve_overlay_column(&base, &offsets(0, 3), &[overlay]).unwrap(); assert_i32_eq(&resolved, [None, Some(2), Some(3)]); } @@ -197,7 +245,7 @@ mod tests { coverage: bitmap([11]), values: i32_array([Some(7)]), }; - let resolved = resolve_overlay_column(&base, 10, &[overlay]).unwrap(); + let resolved = resolve_overlay_column(&base, &offsets(10, 3), &[overlay]).unwrap(); assert_i32_eq(&resolved, [Some(0), Some(7), Some(0)]); } @@ -208,11 +256,37 @@ mod tests { coverage: bitmap([0, 2]), values: Arc::new(StringArray::from(vec!["A", "C"])), }; - let resolved = resolve_overlay_column(&base, 0, &[overlay]).unwrap(); + let resolved = resolve_overlay_column(&base, &offsets(0, 3), &[overlay]).unwrap(); 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 = ResolvedFieldOverlay { + coverage: bitmap([1, 8]), + values: i32_array([Some(11), Some(88)]), + }; + let resolved = resolve_overlay_column(&base, &[5, 1, 8], &[overlay]).unwrap(); + // 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_offset_count_mismatch_errors() { + let base = i32_array([Some(1), Some(2), Some(3)]); + let overlay = ResolvedFieldOverlay { + coverage: bitmap([0]), + values: i32_array([Some(9)]), + }; + // Two offsets for a three-row base column is a caller bug. + assert!(resolve_overlay_column(&base, &[0, 1], &[overlay]).is_err()); + } + #[test] fn test_value_count_mismatch_errors() { let base = i32_array([Some(1), Some(2)]); @@ -220,7 +294,50 @@ mod tests { coverage: bitmap([0, 1]), values: i32_array([Some(9)]), // only one value for two covered offsets }; - assert!(resolve_overlay_column(&base, 0, &[overlay]).is_err()); + assert!(resolve_overlay_column(&base, &offsets(0, 2), &[overlay]).is_err()); + } + + #[test] + fn test_apply_overlays_to_batch_per_field() { + // Two columns; only "age" has an overlay. "name" passes through. + let schema = Arc::new(Schema::new(vec![ + Field::new("age", DataType::Int32, true), + Field::new("name", DataType::Utf8, true), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + i32_array([Some(30), Some(25), Some(40)]), + Arc::new(StringArray::from(vec!["a", "b", "c"])), + ], + ) + .unwrap(); + let plans = vec![FieldOverlayPlan { + field_name: "age".to_string(), + overlays_newest_first: vec![ResolvedFieldOverlay { + coverage: bitmap([2, 5]), + values: i32_array([Some(26), Some(99)]), + }], + }]; + // Batch rows map to physical offsets 4, 5, 6; only offset 5 is covered, + // and offset 5 is at rank 1 in coverage {2, 5}, so its value is 99. + let merged = apply_overlays_to_batch(batch, &[4, 5, 6], &plans).unwrap(); + let ages = merged.column(0); + assert_i32_eq(ages, [Some(30), Some(99), Some(40)]); + let names = merged + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(names, &StringArray::from(vec!["a", "b", "c"])); + } + + #[test] + fn test_apply_overlays_to_batch_empty_plans_is_noop() { + let schema = Arc::new(Schema::new(vec![Field::new("age", DataType::Int32, true)])); + let batch = RecordBatch::try_new(schema, vec![i32_array([Some(1), Some(2)])]).unwrap(); + let merged = apply_overlays_to_batch(batch.clone(), &[0, 1], &[]).unwrap(); + assert_eq!(merged, batch); } #[test] From 7a8cdf25d3ea1ed38d715489e5dab86e778864be Mon Sep 17 00:00:00 2001 From: Will Jones Date: Tue, 23 Jun 2026 09:04:17 -0700 Subject: [PATCH 9/9] 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 | 159 +++ rust/lance/src/dataset/overlay_compaction.rs | 1140 ++++++++++++++++++ 3 files changed, 1300 insertions(+) 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 02787863c8b..d8f46d72ff3 100644 --- a/rust/lance/src/dataset/fragment.rs +++ b/rust/lance/src/dataset/fragment.rs @@ -1208,6 +1208,165 @@ impl FileFragment { Ok(arrow_select::concat::concat(&chunk_refs)?) } + /// Like [`Self::read_overlay_value_column`] but reads only the values at the + /// given `ranks` (0-based positions within the value column) via a `take`, + /// rather than the whole column. 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 { diff --git a/rust/lance/src/dataset/overlay_compaction.rs b/rust/lance/src/dataset/overlay_compaction.rs new file mode 100644 index 00000000000..4951001aa4a --- /dev/null +++ b/rust/lance/src/dataset/overlay_compaction.rs @@ -0,0 +1,1140 @@ +// 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::resolve_overlay_column; +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]) + } 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`. + let offsets: Vec = (0..base.len() as u32).collect(); + let folded = resolve_overlay_column(&base, &offsets, &[winners])?; + 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(); + let indexed_columns: Vec<(usize, ArrayRef)> = columns.into_iter().enumerate().collect(); + writer.write_columns(indexed_columns).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(); + let indexed: Vec<(usize, ArrayRef)> = columns.into_iter().enumerate().collect(); + writer.write_columns(indexed).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 + ); + } +}