perf: use ShardMergeReader for non-PQ index types

[hushengquan]

Problem Description

When processing the merge_partial_vector_auxiliary_files stage for the ivf-sq index on a dataset of 5 million records with 8 workers and 2236 partitions, the operation takes an excessively long time (3780 seconds), which significantly impacts the overall efficiency of vector index merging.(#6288 )

Root Cause

The current implementation of merging partial vector auxiliary files lacks shard-based parallelism optimization, leading to serial processing bottlenecks when handling a large number of partitions and datasets.

Proposed Solution

Use ShardMergeReader to leverage shard-level parallelism for merging partial vector auxiliary files. This optimization distributes the merge workload across multiple shards, reducing the overall processing time.

Performance Metrics (Before/After)

Scenario (8 workers / 2236 partitions / 5M dataset)	Before Optimization	After Optimization	Improvement
merge_partial_vector_auxiliary_files (ivf_sq index)	3780 seconds	208 seconds	~18.1x faster

Additional Context

• Index type: ivf_sq
• Dataset size: 5 million records
• Worker count: 8
• Partition count: 2236
• Optimization core: ShardMergeReader (shard-based parallel merge logic)

[github-actions]

PR Review

Clean change — reusing ShardMergeReader for non-PQ paths is the obvious fix and the 8.8x improvement speaks for itself.

Two items to address:

1. Dead code: write_partition_rows (P1)

After this PR, pub async fn write_partition_rows(...) (line ~290) has zero callers — the only call site was the code you replaced. It should be removed in this PR to avoid a dead-code warning (it's pub but not used outside this file/crate).

2. Missing empty-batches guard (P1 — consistency)

The PQ path (lines 1231-1236) returns an error when batches.is_empty() for a partition with accumulated_lengths[pid] > 0. The new non-PQ path silently skips this case. Consider adding the same guard for consistency, or document why it's intentionally omitted.

---

Otherwise LGTM. Existing test_merge_ivf_flat_success_basic covers this path. Nice improvement.

[codecov]

Codecov Report

❌ Patch coverage is 70.83333% with 7 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...lance-index/src/vector/distributed/index_merger.rs	70.83%	4 Missing and 3 partials ⚠️

📢 Thoughts on this report? Let us know!

[Xuanwo]

Thank you very much for this PR! I think this change is great, but I want to raise one point: we are also limited by memory usage. A non-PQ index uses more memory than a PQ index, that’s why we didn’t have this feature in the first place.

IVF_SQ/IVF_FLAT can have large payloads and large partitions. Therefore, I believe we should ensure our new logic works well in multi‑batch and large‑partition scenarios, and make sure we do not store an entire partition in memory before writing it out.

[hushengquan]

Thank you very much for this PR! I think this change is great, but I want to raise one point: we are also limited by memory usage. A non-PQ index uses more memory than a PQ index, that’s why we didn’t have this feature in the first place.

IVF_SQ/IVF_FLAT can have large payloads and large partitions. Therefore, I believe we should ensure our new logic works well in multi‑batch and large‑partition scenarios, and make sure we do not store an entire partition in memory before writing it out.

Got it, thank you!