`merge_insert` on deleted data fails with `Execution error: The input to a take operation specified fragment id 0 but this fragment does not exist in the dataset`

[ilya-zlobintsev]

Description

Performing the sequence of operations described in the repro script fails with this error:

OSError: LanceError(IO): Execution error: The input to a take operation specified fragment id 0 but this fragment does not exist in the dataset (uri=/tmp/lance-stale-take-repro-ai7gnf30/stale-take.lance, version=6, manifest=tmp/lance-stale-take-repro-ai7gnf30/stale-take.lance/_versions/18446744073709551609.manifest, branch=main), /rustc/4a4ef493e3a1488c6e321570238084b38948f6db/library/core/src/task/poll.rs:291:44

This seems to be a similar bug to what was fixed in #6563, but a slightly different case.

Steps to reproduce

import tempfile
from pathlib import Path

import pyarrow as pa

import lance

with tempfile.TemporaryDirectory(prefix="lance-stale-take-repro-") as tmpdir:
    uri = str(Path(tmpdir) / "stale-take.lance")
    print(f"lance version: {lance.__version__}")

    # 1. Fragment 0: ids 0..4, fragment 1: ids 5..9; BTREE index on `id`.
    def table(ids, vals):
        return pa.table(
            {
                "id": pa.array(ids, pa.uint64()),
                "val": pa.array(vals, pa.uint64()),
            }
        )

    ds = lance.write_dataset(table(range(0, 5), [0] * 5), uri)
    ds = lance.write_dataset(table(range(5, 10), [0] * 5), uri, mode="append")
    ds.create_scalar_index("id", index_type="BTREE")

    def show(label):
        idx = ds.list_indices()[0]
        frags = {
            f.fragment_id: f.metadata.deletion_file is not None for f in ds.get_fragments()
        }
        print(
            f"{label}: fragments (id: has_deletions) = {frags}, "
            f"index bitmap = {sorted(idx.get('fragment_ids') or [])}"
        )

    show("initial")

    # 2. Partial update of one row of fragment 0 (source columns reordered
    #    w.r.t. the dataset schema, which routes merge_insert through the
    #    in-place partial-update path): drops fragment 0 from the index
    #    bitmap, leaves the btree data untouched.
    ds.merge_insert("id").when_matched_update_all().execute(
        pa.table({"val": pa.array([1], pa.uint64()), "id": pa.array([0], pa.uint64())})
    )
    show("after partial patch (fragment 0 pruned from bitmap)")

    # 3. Full-schema upsert matching every row of fragment 0: the fragment is
    #    fully drained and removed; its rows move to a new fragment. The
    #    btree data still holds fragment 0's addresses.
    ds.merge_insert("id").when_matched_update_all().execute(table(range(0, 5), [9] * 5))
    show("after full drain (fragment 0 removed)")

    # 4. One deletion in fragment 1, which IS in the index bitmap. This forces
    #    the prefilter's block-list branch. Comment this line out and the
    #    final merge_insert succeeds.
    ds.delete("id = 9")
    show("after delete (block-list branch armed)")

    # 5. Probe a key whose stale index entry points at dead fragment 0.
    print("running merge_insert probing id=0; affected builds fail here:")
    ds.merge_insert("id").when_matched_update_all().execute(table([0], [7]))
    print("merge_insert succeeded (bug not present)")

Expected behavior

The merge insert succeeds

Lance version

8.0.0

Language binding

Python

Environment

Local disk, Linux host

Logs / traceback

[ilya-zlobintsev]

Additional note: the provided example performs a merge_insert on a single indexed field, which has had the bug for many lance versions.
If you do a merge_insert on a compound key, where only one of the columns in indexed, it regressed from lance 7 to 8 (previously worked, now hitting this bug) as with #6878 it now also uses the indexed path.