[dshepelev15]

Summary

Four hot spots found by profiling IVF_PQ search on an f16 vector column (4M rows, fsl<f16,768>, 128 partitions × 48 sub-vectors), each fixed where the profile pointed:

• f16 codebooks use the prepared L2 targets. PQStorage::with_prepared_targets and the Float16 branch of dist_calculator now take the quantizer's widened l2_targets, so each partition's distance table no longer goes through the generic f16 kernel (8% of a query).
• Prefilters are applied to a partition in bulk. RowAddrMask groups rows by fragment, looks the fragment's selection up once per run and intersects ascending runs of 64+ rows with the fragment bitmap. selected_indices keeps its Vec<u64> signature and enumerates positions as usize; selected_offsets_into writes the u32 offsets the flat index uses at its three prefilter sites through the existing QueryScratch::u32 buffer, whose documented bound is one index partition (fewer than u32::MAX rows). PQDistCalculator overrides accumulate_filtered_topk_with_scratch to score accepted rows with the selected-rows kernels in 256-row chunks: the existing 4-bit kernel, and a new compute_pq_distance_8bit_rows that interleaves eight rows so the 48 dependent adds of one row no longer bound the loop. Per-row RowAddrMask::selected plus per-row distance() were 62% + 15% of a masked query.
• One query scores a large chunk on several cores. IVFIndex::score_chunk splits a global-top-k chunk over concurrent spawn_cpu tasks when it holds at least twice LANCE_IVF_SCORING_TASK_BYTES (1 MiB) of prepared partitions, each task with its own heap, merged back to the k * refine best. The split honors the query's query_parallelism: an explicit 1 keeps the single-worker scoring of feat(vector): add partition search parallelism #6475, an explicit count caps the tasks, and automatic (0) or -1 use the CPU pool when it has at least two workers; a one-worker pool keeps one task. Small partitions keep the single dispatch. Which of several rows tied exactly at the k-th distance survives can differ from sequential scoring.
• Less per-partition churn. PQ partition scans write distances into the query scratch instead of allocating a vector per partition (compute_pq_distance_into), and accumulate_distances_into_heap rejects rows worse than the current k-th with one plain float compare before the total-order compares.

No public struct or trait changes its shape: QueryScratch keeps its fields and selected_indices its signature.

Branch state

This branch is the author's head plus an ordinary merge of main at d27fd01 and one repair commit. Versus current main it touches 9 files, +840/−81.

• selected_indices no longer routes its positions through the internal u32 buffer, so the public path has no 4-billion-row precondition. selected_offsets_into keeps that buffer for the three partition hot paths and documents the partition bound.
• The scoring-split test derives its expectation from get_num_compute_intensive_cpus() >= 2, so a single-worker pool exercises the production fallback instead of failing. It was run in separate processes with the default 32-CPU pool and with LANCE_CPU_THREADS=1; the explicit query_parallelism = 1 case stays single-task in both.
• The masked PQ partition test now writes two fragments, so a partition offset is not a row address; it proves the masked scorers map offsets through storage.row_id.
• Limits: there is no >u32::MAX-row fixture (the public offset width follows from the usize enumeration, not from a large fixture); the k-th tie winner is still not guaranteed; no new cache, protocol, or scheduler.

Measurements

Author-measured on the original feature head and not re-run for this branch update. Same dataset on an EPYC 9354, k=10, nprobes=20, p50 latency of a single client, index cache warm:

mode	before	after
plain	13.3 ms	5.1 ms
allowlist mask (5% of rows)	25.2 ms	6.4 ms
refine_factor 10	20.3 ms	11.4 ms
plain, 16 concurrent clients	756 qps	1400 qps

recall@10 against a flat scan is unchanged (0.53–0.57 over 20 queries, within sampling noise).

Tests

• test_selected_indices_bulk_matches_per_row: bulk, per-row, unsorted and duplicate runs, full and absent fragments, allow and block lists, buffer reuse.
• test_f16_pq8_uses_prepared_l2_targets: f16 PQ8 partition storage gets the targets and scores within rounding of the generic kernel.
• test_filtered_accumulate_topk_matches_per_row: filtered scan equals per-row distance() bit for bit for PQ4/PQ8 × L2/Dot, from an empty and a prefilled heap, across a chunk boundary.
• test_compute_pq_distance_8bit_rows_partial_group: a row count that is not a multiple of eight and an output longer than the ids.
• test_balanced_byte_ranges, test_merge_topk_heaps_keeps_the_best: the split and merge helpers, including the below-two-tasks, one-partition and capped cases.
• test_batch_partition_prefilter_preserves_bulk_scoring: flat/PQ4/PQ8 partition search under allow, block and empty masks over a two-fragment dataset; masked results equal the unfiltered result filtered per row.
• test_knn_scores_large_chunk_on_concurrent_tasks: an IVF_FLAT index whose two partitions exceed the split threshold is scored by concurrent tasks under automatic parallelism when the pool has two or more workers, and on one task under query_parallelism = 1 or a one-worker pool; both match a flat scan.

cargo fmt --all and cargo clippy --all --tests --benches --profile ci --locked -- -D warnings are clean, as are cargo check --workspace --tests --benches --profile ci --locked and the targeted test runs above.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[dshepelev15]

Addressed both notes: score_chunk now takes its task cap from the query, so an explicit query_parallelism = 1 keeps the single-worker scoring of #6475 (an explicit count caps the split, automatic and -1 use the pool); test_knn_scores_large_chunk_on_concurrent_tasks covers both the automatic and the single-worker case. QueryScratch no longer gains a field: the selection goes through the existing u32 scratch via RowAddrMask::selected_offsets_into, so no public struct changes shape and the breaking-change label should not be needed.

[lance-gatekeeper]

✅ Gate recommendation: approve.

The public mask path keeps full-width positions, and masked Flat/PQ searches keep fragment row addresses distinct from partition offsets. The explicit query_parallelism=1 limit remains intact. No additional changes are needed.

[Xuanwo]

Approving at head 13f0315.

All prior findings are resolved: query_parallelism = 1 still pins scoring to one worker (reproducer added as a test), and QueryScratch keeps its original fields, so the withdrawn breaking-change note does not apply. All applicable checks completed at this head; the four skipped jobs are path-filtered for Python/bindings and Arrow-dependency changes that this Rust-only diff does not make.

A squash of this head onto current main (9c0de40) merges cleanly to tree 3503b89e3f8a25cd29235c37e3ffbc311fd65176, which is the candidate tree dfd02e71 plus main's #10826 delta, with every changed-file blob identical on both sides.