Remove AsThreadPool generic and take &RayonThreadPool directly

[JordanMaples]

Remove AsThreadPool trait, forward_threadpool! macro, and execute_with_rayon.

Addresses and Closes #905

Copilot summary of changes and usage:

Removed:

• AsThreadPool trait + sealed module
• forward_threadpool! macro
• execute_with_rayon free function

Added:

• RayonThreadPool — owned pool, created via create_thread_pool(n)
• as_ref() → RayonThreadPoolRef<'_> for borrowing
• RayonThreadPoolRef<'a> — Copy borrowed handle accepted by all pool-taking APIs
• RayonThreadPoolRef::new(&rayon::ThreadPool) — bring-your-own-pool support

Updated:

• All ParallelIteratorInPool methods take RayonThreadPoolRef<'_> instead of &RayonThreadPool
• All pool-accepting functions across 6 crates de-genericized to take RayonThreadPoolRef<'_> directly
• PQGenerationContext pool field changed from type parameter to borrowed RayonThreadPoolRef<'_>
• Test and bench callers updated to create explicit pools at call site

Usage

 // DiskANN-owned pool
 let pool = create_thread_pool(num_threads)?;
 generate_pq_pivots(pool.as_ref())?;

 // Bring-your-own pool (e.g. Cosmos DB, Azure SQL embedding scenarios)
 let external = rayon::ThreadPoolBuilder::new().num_threads(4).build()?;
 generate_pq_pivots(RayonThreadPoolRef::new(&external))?;

[Copilot]

Pull request overview

This PR simplifies the Rayon thread-pool plumbing by removing the AsThreadPool generic + forward_threadpool! macro approach and standardizing pool-accepting APIs to take a borrowed &RayonThreadPool instead. This reduces monomorphization higher in the call stack and makes pool usage more explicit across providers/tools/disk build paths.

Changes:

• Removed AsThreadPool, forward_threadpool!, and execute_with_rayon; updated exports accordingly.
• Updated PQ/kmeans/partition/quantization APIs to accept &RayonThreadPool and adjusted call sites (including tests) to construct explicit pools.
• Updated disk quantization generator and PQ generation context types to borrow the pool instead of carrying a generic pool type parameter.

Reviewed changes

Copilot reviewed 16 out of 16 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
diskann-tools/src/utils/build_pq.rs	Creates a thread pool once and passes &RayonThreadPool into PQ generation calls.
diskann-providers/src/utils/rayon_util.rs	Removes execute_with_rayon/AsThreadPool infra; retains concrete RayonThreadPool and in-pool iterator helpers.
diskann-providers/src/utils/mod.rs	Stops re-exporting removed thread-pool abstractions/functions.
diskann-providers/src/storage/index_storage.rs	Updates tests to pass &RayonThreadPool instead of a thread-count integer.
diskann-providers/src/model/pq/pq_construction.rs	Degenericizes PQ construction APIs to take &RayonThreadPool.
diskann-providers/src/index/wrapped_async.rs	Updates tests to pass &RayonThreadPool.
diskann-providers/src/index/diskann_async.rs	Degenericizes train_pq and updates test call sites to pass &RayonThreadPool.
diskann-disk/src/utils/partition.rs	Degenericizes partitioning helper to take &RayonThreadPool.
diskann-disk/src/utils/math_util.rs	Degenericizes math utilities to take &RayonThreadPool.
diskann-disk/src/utils/kmeans.rs	Degenericizes kmeans utilities to take &RayonThreadPool.
diskann-disk/src/storage/quant/pq/pq_generation.rs	Removes pool type parameter from PQ generation context; now borrows &RayonThreadPool.
diskann-disk/src/storage/quant/generator.rs	Degenericizes quant data generator to take &RayonThreadPool; updates test call site accordingly.
diskann-disk/src/build/builder/quantizer.rs	Trains PQ quantizer by constructing a thread pool at call site and passing &RayonThreadPool.
diskann-disk/src/build/builder/core.rs	Updates merged index workflow partitioning call to the new signature.
diskann-disk/src/build/builder/build.rs	Updates PQ generation usage and generator invocation to the new borrowed-pool signatures.

Comments suppressed due to low confidence (1)

diskann-providers/src/utils/rayon_util.rs:15

• create_thread_pool is documented to treat num_threads == 0 as “use logical CPUs”, but the implementation always calls ThreadPoolBuilder::num_threads(num_threads). If callers pass 0 (e.g., as an “auto” sentinel), this will rely on Rayon accepting 0; if Rayon treats 0 as invalid, this will return an error at runtime. Consider explicitly handling num_threads == 0 by omitting .num_threads(...) (or mapping 0 to available_parallelism()) so behavior matches the docs and is stable across Rayon versions.

/// Creates a new thread pool with the specified number of threads.
/// If `num_threads` is 0, it defaults to the number of logical CPUs.
pub fn create_thread_pool(num_threads: usize) -> ANNResult<RayonThreadPool> {
    let pool = rayon::ThreadPoolBuilder::new()
        .num_threads(num_threads)
        .build()
        .map_err(|err| ANNError::log_thread_pool_error(err.to_string()))?;
    Ok(RayonThreadPool(pool))

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 77.46479% with 48 lines in your changes missing coverage. Please review.
✅ Project coverage is 89.49%. Comparing base (e4cd2f3) to head (d75fa17).

Files with missing lines	Patch %	Lines
diskann-tools/src/utils/ground_truth.rs	25.45%	41 Missing ⚠️
diskann-tools/src/utils/build_pq.rs	0.00%	4 Missing ⚠️
diskann-disk/src/utils/kmeans.rs	85.71%	2 Missing ⚠️
diskann-tools/src/utils/search_disk_index.rs	0.00%	1 Missing ⚠️

❌ Your patch status has failed because the patch coverage (77.46%) is below the target coverage (90.00%). You can increase the patch coverage or adjust the target coverage.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #967      +/-   ##
==========================================
- Coverage   89.50%   89.49%   -0.01%     
==========================================
  Files         448      448              
  Lines       84167    84118      -49     
==========================================
- Hits        75333    75285      -48     
+ Misses       8834     8833       -1     

Flag	Coverage Δ
miri	89.49% <77.46%> (-0.01%)	⬇️
unittests	89.34% <77.46%> (-0.01%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
diskann-benchmark/src/backend/index/product.rs	100.00% <ø> (ø)
diskann-disk/src/build/builder/build.rs	94.15% <100.00%> (ø)
diskann-disk/src/build/builder/core.rs	95.23% <100.00%> (ø)
diskann-disk/src/build/builder/quantizer.rs	90.52% <100.00%> (ø)
diskann-disk/src/search/provider/disk_provider.rs	90.89% <100.00%> (ø)
diskann-disk/src/storage/quant/generator.rs	92.67% <100.00%> (-0.06%)	⬇️
diskann-disk/src/storage/quant/pq/pq_generation.rs	93.22% <100.00%> (-0.12%)	⬇️
diskann-disk/src/utils/math_util.rs	98.81% <100.00%> (-0.01%)	⬇️
diskann-disk/src/utils/partition.rs	92.51% <100.00%> (-0.04%)	⬇️
diskann-providers/src/index/diskann_async.rs	95.99% <100.00%> (+0.02%)	⬆️
... and 8 more

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[JordanMaples]

This closes #905

[hildebrandmw]

Thanks Jordan. One quick comment: can you please update the PR description to reflect that bring-your-own-threadpool is now supported! Cheers!