bug: test_ann flakes because its IVF fixture leaves partitions with fewer than k rows

Open Beginner friendly
  • Is something wrong?
#10,782 0 comments 0 reactions 0 assignees View on GitHub

Maintainers usually reply within 2 days

@lance-gatefixer is already working on this.

Since Nov 6, 2026.

  • #10785 by @lance-gatefixer — open

Assessment

Difficulty
2/5
Estimated time
1-3 hours
Newbie friendliness
75/100
Issue type
Bug
Clarity
Clearly specified
Activity status
Active
Tech stack
python, rust
Domain
machine-learning, testing-qa

Research direction

Start with python/tests/test_vector_index.py::test_ann and the indexed_dataset fixture, then check how its IVF_PQ parameters leave some partitions with fewer than 15 rows. Adjust the fixture so k-means trains on enough rows to keep every partition comfortably above k, while preserving the behavior needed by test_rowid_order and the indexed_dataset cases of test_distributed_vector. Done when these tests pass reliably and the nprobes=1 assertion remains unchanged.

Written by the indexing model from the issue text.

Description

bug

python/tests/test_vector_index.py::test_ann fails now and then with assert 11 == 15: most recently in this job of an unrelated PR, and #8789 recorded the same assertion on main.

Cause

The indexed_dataset fixture builds IVF_PQ over 1000 random vectors with num_partitions=4, max_iters=2 and sample_rate=2, so k-means trains on num_partitions * sample_rate = 8 rows (sample_size_hint in rust/lance/src/index/vector/ivf.rs). That almost always leaves a tiny partition: in 28 of 30 builds of this fixture the smallest partition held fewer than 15 rows, usually a single row. run() searches with "nprobes": 1, which ProbeParams::resolve (rust/lance-index/src/vector.rs) makes both the minimum and the maximum number of partitions to probe. When the random query's nearest centroid is the tiny partition's, the search returns only that partition's rows and assert len(rs) == 15 fails. A random query rarely lands there, which is why the test usually passes.

Reproduction

Querying the tiny partition's own centroid makes the failure deterministic, and probing all four partitions is the control. With pylance 15.0.0b11 the script fails on every run:

import os
import tempfile

# Keeps the IVF centroids in index_stats() and silences its warning about them.
os.environ.setdefault("LANCE_INCLUDE_VECTOR_CENTROIDS", "true")

import lance
import numpy as np
import pyarrow as pa

# The index of test_ann's `indexed_dataset` fixture: 1000 random 128-d vectors,
# IVF_PQ with 4 partitions trained on num_partitions * sample_rate = 8 rows.
for _ in range(10):
    with tempfile.TemporaryDirectory() as tmp:
        mat = np.random.randn(1000, 128).astype(np.float32)
        vectors = pa.FixedSizeListArray.from_arrays(pa.array(mat.ravel()), 128)
        ds = lance.write_dataset(pa.table({"vector": vectors}), tmp)
        ds = ds.create_index(
            "vector",
            index_type="IVF_PQ",
            num_partitions=4,
            num_sub_vectors=16,
            max_iters=2,
            sample_rate=2,
        )
        ivf = ds.stats.index_stats("vector_idx")["indices"][0]
        sizes = [p["size"] for p in ivf["partitions"]]
        small = int(np.argmin(sizes))
        if sizes[small] >= 15:
            continue
        # A query whose nearest centroid is the small partition's, which is
        # what a random query in test_ann occasionally gets.
        q = np.array(ivf["centroids"][small], dtype=np.float32)
        nearest = {"column": "vector", "q": q, "k": 15}
        all_probes = ds.to_table(nearest={**nearest, "nprobes": 4})
        one_probe = ds.to_table(nearest={**nearest, "nprobes": 1})
        print(f"partition sizes: {sizes}")
        print(f"rows with nprobes=4: {len(all_probes)}, with nprobes=1: {len(one_probe)}")
        assert len(all_probes) == 15
        assert len(one_probe) == 15  # test_ann's assertion
        break
else:
    raise SystemExit("every build had partitions of 15 or more rows")

partition sizes: [1, 220, 544, 235]
rows with nprobes=4: 15, with nprobes=1: 1
Traceback (most recent call last):
  File "test_ann_repro.py", line 40, in <module>
    assert len(one_probe) == 15  # test_ann's assertion
AssertionError
Expected

test_ann passes on every run: every partition of the fixture's index should hold well over k rows, for example by letting k-means train on all rows instead of 8.

What not to change

A fixed nprobes is meant to cap the search, so fewer than k rows from a small partition is correct library behavior: do not change ProbeParams, and do not weaken the assertion to len(rs) <= 15. The fixture is also used by test_rowid_order and the indexed_dataset cases of test_distributed_vector, which should keep passing.

Dominant language
Rust
Stars
7.3k
Forks
924
Avg merge
3d 18h
Merged PRs (30d)
575

Getting set up

Open in Codespaces

Starts the project's dev container in your browser, under your own GitHub account.

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from lance-format/lance

All issues in lance-format/lance

Similar issues

More Rust issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.