perf(index): parallelize kmeans centroid recomputation - #8560
perf(index): parallelize kmeans centroid recomputation#8560isaac-dasari wants to merge 7 commits into
Conversation
github-actions
Bot
added
A-index
![lance-gatekeeper[bot]](https://avatars.githubusercontent.com/in/4433141?s=60&v=4){kind=small}
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
lance-gatekeeper
Bot
added
the
K-changes
lance-gatekeeper
Bot
removed
the
K-changes
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
lance-gatekeeper
Bot
added
the
K-changes
lance-gatekeeper
Bot
removed
the
K-changes
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
lance-gatekeeper
Bot
added
the
K-approved
lance-gatekeeper
Bot
removed
the
K-approved
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
lance-gatekeeper
Bot
added
the
K-changes
lance-gatekeeper
Bot
removed
the
K-changes
There was a problem hiding this comment.
❌ Gate recommendation: request changes.
The revision fixes the 512-row fallback and restores owner-local row order, but its row-count-only selector still sends every larger committed workload through an indexed path that remains slower than the live-base rescan.
Keep the owner-rescan path unless selection accounts for worker count, k, dimension, and measured index amortization; none of the committed indexed shapes demonstrates a win.
| // short membership rescans it replaces. Keep the existing owner-rescan | ||
| // path there and reserve indexing for workloads large enough to amortize | ||
| // its two linear setup passes. | ||
| if num_vectors < MIN_OWNER_INDEX_ROWS { |
There was a problem hiding this comment.
This row-count-only cutoff still selects owner indexing when its two serial construction passes, allocation, and indirect row lookup cost more than the membership rescans they replace. Against live base dd08336c, identical 62-thread release-with-debug batched harnesses measured all four committed shapes that take the index path slower on 90894283:
default_ivf_high_dim: 0.329 → 0.605 ms median (+83.9%).default_pq_subvector: 0.327 → 0.743 ms (+127.2%).max_sample_low_dim: 0.668 → 1.269 ms (+90.0%).large_incremental_ivf: 2.956 → 3.212 ms (+8.7%).
The 512-row fallback reaches parity (0.271 → 0.269 ms), confirming that the rescan path fixes the case where it is selected. Two-thread runs also measured PQ 0.744 → 0.946 ms (+27.2%) and IVF 2.089 → 2.169 ms (+3.8%), so worker count must be part of the decision. This is the current successor to the earlier performance finding: owner grouping fixes cluster-major locality, but the selector still needs to retain owner rescans unless the complete workload cost predicts an indexed-path win.
Reproducer
Both exact-head harnesses used the five committed shapes, cyclic membership, allocation outside timing, two warmups, 20 samples, and 50 calls per sample. These commands produced the 62-thread results:
LANCE_CPU_THREADS=62 RAYON_NUM_THREADS=62 BENCH_SAMPLES=20 BENCH_WARMUPS=2 BENCH_BATCH=50 /home/agent/tmp/pr8560-dd0-target-bench/release-with-debug/pr8560-kmeans-bench
LANCE_CPU_THREADS=62 RAYON_NUM_THREADS=62 BENCH_SAMPLES=20 BENCH_WARMUPS=2 BENCH_BATCH=50 /home/agent/tmp/pr8560-908-target-bench/release-with-debug/pr8560-kmeans-benchThe two-thread confirmation used the same binaries with:
LANCE_CPU_THREADS=2 RAYON_NUM_THREADS=2 BENCH_SAMPLES=20 BENCH_WARMUPS=2 BENCH_BATCH=50 BENCH_CASES=default_ivf_high_dim,default_pq_subvector /home/agent/tmp/pr8560-dd0-target-bench/release-with-debug/pr8560-kmeans-bench
LANCE_CPU_THREADS=2 RAYON_NUM_THREADS=2 BENCH_SAMPLES=20 BENCH_WARMUPS=2 BENCH_BATCH=50 BENCH_CASES=default_ivf_high_dim,default_pq_subvector /home/agent/tmp/pr8560-908-target-bench/release-with-debug/pr8560-kmeans-benchExpected: indexed execution should be selected only for workloads where it matches or improves on the owner rescan. Observed: all indexed committed shapes regressed at 62 threads, and PQ also materially regressed at two threads.
lance-gatekeeper
Bot
added
the
K-changes
lance-gatekeeper
Bot
removed
the
K-changes
1 participant
Summary
The previous implementation divided centroids among workers, but every worker scanned the complete input. That made centroid recomputation perform O(num_vectors * workers) input reads. This changes the input scan to O(num_vectors) while avoiding shared-write contention.
Closes #6369.
Validation
Passed locally:
cargo fmt --all -- --checkcargo check -p lance-index --tests --benches --profile bench -j 1cargo clippy -p lance-index --tests --benches --profile bench -j 1 -- -D warningsgit diff --checkThe optimized test/benchmark linker exceeded the available resources on the local 7.6 GiB development machine, so executable test and benchmark runs are intentionally delegated to GitHub Actions. The PR is draft pending those results.