bug: stale rows fill a vector index segment's per-partition top-k and hide rows it still owns

[wombatu-kun]

Problem

An in-place column update (update_columns + LanceOperation.Update) prunes the updated fragment from the old index segment's fragment_bitmap, but that segment's index file still physically holds the pre-update rows. Those rows are ranked inside search_in_partition and can fill the partition's top-k heap before any ownership filter runs, so the segment contributes nothing for the fragments it does still own and the query misses closer rows that are present and current.

Reproduces on main today. It is orthogonal to #7370: #7371 stops the stale rows from being served, but cannot recover the owned rows they displaced, because the displacement happens inside the sub-index.

Reproduction

Run against main at 6fe450bb7 and against #7371's head bc8ab3dd8.

import lance
import numpy as np
import pyarrow as pa


def test_segment_mask_must_precede_partition_topk(tmp_path):
    ndim = 4

    def table(ids, value):
        vectors = np.full((len(ids), ndim), value, dtype=np.float32)
        return pa.table(
            {
                "id": pa.array(ids, type=pa.int64()),
                "vector": pa.FixedSizeListArray.from_arrays(
                    pa.array(vectors.reshape(-1), type=pa.float32()), ndim
                ),
            }
        )

    ds = lance.write_dataset(table(range(20), 1.0), tmp_path, mode="create")
    ds = lance.write_dataset(table(range(100, 120), 0.0), tmp_path, mode="append")
    ds = ds.create_index("vector", index_type="IVF_FLAT", metric="l2", num_partitions=1)

    frag = ds.get_fragment(1)
    rowids = frag.to_table(columns=["id"], with_row_id=True)["_rowid"].to_pylist()
    update_data = pa.table(
        {
            "_rowid": pa.array(rowids, type=pa.uint64()),
            "vector": pa.array(
                [[10.0] * ndim] * len(rowids), type=pa.list_(pa.float32(), ndim)
            ),
        }
    )
    updated, fields = frag.update_columns(update_data)
    ds = lance.LanceDataset.commit(
        ds.uri,
        lance.LanceOperation.Update(updated_fragments=[updated], fields_modified=fields),
        read_version=ds.version,
    )
    ds.optimize.optimize_indices(num_indices_to_merge=0)
    ds = lance.dataset(ds.uri)

    tbl = ds.to_table(
        columns=["id"],
        nearest={"column": "vector", "q": np.zeros(ndim, dtype=np.float32), "k": 5},
    )
    print(f"OBSERVED_IDS={tbl['id'].to_pylist()}")
    print(f"OBSERVED_DISTS={tbl['_distance'].to_pylist()}")
    assert all(row_id < 20 for row_id in tbl["id"].to_pylist())

Observed:

main       OBSERVED_IDS=[100, 101, 102, 103, 104]  OBSERVED_DISTS=[0.0, 0.0, 0.0, 0.0, 0.0]
bc8ab3dd8  OBSERVED_IDS=[100, 101, 102, 103, 104]  OBSERVED_DISTS=[400.0, 400.0, 400.0, 400.0, 400.0]

Expected [0, 1, 2, 3, 4] at distance 4. Fragment 0 holds 20 such rows, the old segment owns them, and it never returns any of them.

On main the segment ranks the stale copies by their pre-update vector and reports distance 0, so stale data is served as current. #7371 fixes that half, and the same ids come back with their true current distance. The missing ids below 20 are unchanged by #7371 and are what this issue tracks.

Root cause

search_in_partition accumulates a local top-k over every row the partition physically contains (push_candidate_local in lance-index/src/vector/flat/index.rs, and the equivalent in pq.rs), bounded by k * refine_factor. Ownership is only consulted afterwards, in ANNIvfSubIndexExec. A row the segment no longer owns therefore occupies a heap slot that an owned row would have taken, and once the heap is truncated that owned row is unrecoverable at any later stage.

Fix

Compose the per-segment ownership mask into the prefilter handed to the sub-index, on both the per-partition and the streaming search_partitions paths, so unowned rows never enter the local heap. The existing post-filter can stay as a boundary check.

The cost has to be measured before this lands: DatasetPreFilter::create_restricted_deletion_mask returns Some for every segment of a multi-delta index, so composing it in makes DatasetPreFilter::is_empty() false for every delta and disables the no-prefilter fast path at rust/lance/src/index/vector/pq.rs:282, rust/lance-index/src/vector/flat/index.rs:151 and :272, and rust/lance-index/src/vector/hnsw/builder.rs:1356. That turns an O(k) post-filter into an O(partition) masked search on every delta query, including the overwhelmingly common case of an index with no stale rows at all.

Found by lance-gatekeeper[bot] while reviewing #7371.