feat: add native DiskANN and adaptive IVF indexes

[JingsongLi]

Summary

• add a native DiskANN implementation with Vamana graph construction, compressed adjacency and row IDs, PQ-guided search, exact reranking, cache-aware multi-range reads, and automatic storage profiles
• replace the unreleased IVF-HNSW variants with IVF-SQ and improve IVF-FLAT, IVF-PQ, and IVF-RQ build/search I/O paths
• expose adaptive tuning, batch search, concurrent reads, and storage-profile controls through Rust, C/C++, Java/JNI, and Python APIs
• add stable storage fixtures and public ANN benchmark support for the new index families

This is the first commit of a seven-commit local series. Architecture notes, user documentation, benchmark-result pages, and release-tool updates remain in follow-up commits and are intentionally excluded from this PR.

Verification

• cargo fmt --all -- --check
• cargo test --workspace (366 passed, 1 ignored in the core crate; all FFI/JNI tests passed)
• cargo clippy --all-targets --workspace -- -D warnings
• mvn -f java/pom.xml test -q
• Python API tests: 18 passed
• ASF license-header and .asf.yaml validation

[leaves12138]

I found one blocking resource regression and one merge-sequencing issue. The IVF list-loading paths need bounded aggregate memory, not only bounded per-pread request size. Also, the root README, root STORAGE_FORMAT.md, and HTML documentation still advertise the removed IVF-HNSW APIs and the removed IVF-RQ query_bits parameter, while the root README identifies the stale root storage document as normative. If the documentation work is intentionally split into a stacked follow-up, the two changes should land atomically; otherwise this PR should update the public entry points before merging.

I reviewed the latest DiskANN inner-product/cosine update and did not find a correctness issue in the metric persistence, cosine normalization/L2 traversal, negative-dot-product ranking, or exact reranking paths. Local Rust, Java/JNI, Python, C, C++, packaging, and additional non-exhaustive recall checks passed.

[leaves12138]

These payloads are allocated for every selected list before pread_batched_slices starts. MAX_IVF_BATCH_READ_BYTES only bounds each pread submission; it does not bound the aggregate live allocation. A single search passes all probed lists here, and batch search passes the union of all queries' lists. The same pattern exists in IVF-PQ and IVF-SQ. For example, IVF-FLAT with 1M 768-dimensional vectors and nprobe == nlist can retain roughly 3 GiB of vector payloads before scanning. This is a regression from the previous per-list read/scan behavior. Please partition selected lists by aggregate payload bytes/range count and perform read -> decode/scan -> release for each chunk, similar to the bounded IVF-RQ path.

[leaves12138]

Thanks for adding aggregate payload batching. One oversized list is still unbounded, though: bounded_ivf_payload_batch_end deliberately returns the first payload even when it is larger than 64 MiB, and each reader allocates that whole list before issuing the positional read. A one-list IVF-FLAT index, or a highly skewed list, can therefore still allocate count * dimension * 4 bytes (potentially multiple GiB), so the original OOM case is only fixed for the sum of multiple lists. Could we stream a single large list in chunks, or otherwise reject/cap it, while keeping the progress guarantee?

[leaves12138]

This sets l_search to the full vector count, so the new IP/cosine tests validate format propagation and exact reranking but do not exercise approximate graph/PQ candidate quality. Please add a deterministic recall regression with l_search << N (ideally covering both PQ-guided/full-precision build and batch/filter paths). I ran an additional 2K x 32 random-data check and saw about 99.5%--99.6% recall@10 at l_search=50, so this is a coverage request rather than a currently observed correctness failure.

[leaves12138]

Could we preserve the selected build-distance mode when the memory budget triggers sharding? build_sharded_with_stats builds each local graph through build_sequential_with_metric, which always uses full-precision distances, so both an explicit diskann.build-distance=product_quantized and the FastBuild/Balanced preset are silently ignored whenever graph_shards > 1. I reproduced this by building the same four-shard graph with the enum toggled; the graphs are identical. This also makes the budget model inconsistent: it accounts for the PQ build-distance table although the sharded execution instead allocates full-precision local vectors. Please add a PQ-guided sharded path, or reject/explicitly resolve this combination before shard selection, with coverage for both build-distance modes.

[leaves12138]

Could we validate max-bytes-per-vector as an actual maximum before returning the plan? It is currently accepted for every index type without a final size check: IVF-FLAT/IVF-SQ ignore it, while DiskANN uses it only to choose PQ/raw-vector encodings even though the format always stores the raw vector, adjacency, row ID, and PQ code. For example, index.type=diskann, dimension=128, and max-bytes-per-vector=32 succeeds and selects F16, but the raw vector alone is already 256 bytes per row; ivf_flat with a one-byte limit also succeeds. This silently violates the requested storage objective. Please estimate the complete persisted per-vector cost and reject unsatisfiable combinations, or narrow/rename the option and reject index types for which it is not enforced.

[leaves12138]

Before making the JNI reader cloneable for batched search, please bound the local references created by one pread call. Dropping the Rust JByteArray wrapper does not call DeleteLocalRef, so every new_byte_array remains live until the native frame returns; copy_java_buffers then creates another local reference per range with get_object_array_element. IVF search can pass one range per probed/unique list, and the Java default maxRangesPerRead() == 0 means unlimited. I reproduced this with 50,000 single-row IVF lists and nprobe=50,000: -Xcheck:jni reported the local-reference count growing to roughly 100,000 during one search. This retains every Java byte array and can overflow/cripple the local-reference table on large probes or batches. Please chunk JNI range batches and/or use local frames / explicit delete_local_ref in both loops.

[leaves12138]

The sharded-fit calculation omits the KMeans peak that runs immediately before the graph allocations. kmeans_train always materializes a train_data copy (up to shard_count * 256 * dimension floats), plus its assignments, centroid buffers, and SGEMM score matrix, but none of those bytes are included here. At the supported limits this can be roughly another 64 MiB just for the copied training vectors (64 * 256 * 1024 * 4), before the other scratch buffers. Consequently graph_build_shard_count can select a shard count whose estimate is within diskann.memory-budget-bytes while the actual build exceeds that budget. Please include the KMeans training peak (and overlap it with the already-live fixed allocations) when deciding whether a sharded build fits.

[leaves12138]

diskann.memory-budget-bytes is not consulted during PQ training, even though this phase can be substantially larger than the later graph estimate. ProductQuantizer::train_hot_start trains all m sub-quantizers with a Rayon par_iter; each active KMeans task copies its sub-vector training data again and may allocate a 16 MiB SGEMM assignment matrix. With the supported 50,000 x 1,024 training sample and 64 concurrent sub-quantizers, the score matrices alone can reach about 1 GiB, on top of the retained ~195 MiB sample and the per-task data copies. A caller can therefore set a 256 MiB build budget, exceed it by several times during finish_training, and only have the budget considered later during graph serialization. Please make PQ-training concurrency/sample scratch honor the build budget (or reject budgets that cannot cover the training peak) rather than treating only Vamana construction as budgeted.

[leaves12138]

Thanks for budgeting the per-task PQ scratch and parallelism. The retained trainer reservoir is still outside this plan when downsampling is required. VectorIndexTrainer always keeps up to 50,000 raw vectors; at d=1024 that is 204,800,000 bytes. With a 128 MiB budget, pq_training_plan succeeds with sample_count=27,892, after which bounded_pq_training_sample allocates another sampled buffer while the original 50,000-vector reservoir remains alive. Cosine training can add a normalized owned copy as well. Since peak_for counts only one sample_bytes buffer, the accepted budget can still be exceeded before KMeans scratch is considered. Could the trainer reservoir limit be derived from the budget, or could the plan include the retained, downsampled, and normalization buffers that coexist?

[leaves12138]

Could we validate the search parameters before converting them to the C ABI? ctypes.c_size_t silently wraps Python integers, so SearchParams.diskann(5, -1).to_ffi().width becomes usize::MAX (and oversized positive values wrap modulo SIZE_MAX). In an actual search this turns a caller error into an effectively exhaustive DiskANN search; SearchParams.ivf(5, -1) has the same behavior for nprobe. The Java/JNI path rejects negative widths. A __post_init__ check for positive top_k and algorithm-specific width, plus the platform upper bound, would keep the Python API consistent and avoid silently changing the requested search.

[leaves12138]

The lazy row-ID order is not reserved from the reader memory budget. resolve_cache_budgets can assign all bytes after the steady resident state to adjacency preload/cache and raw-vector cache, but this check counts only the steady state, current hot adjacency, the order, and its decode scratch. Once the order is retained, both shared caches can still grow to their original limits; if a filtered query loads the order before optimize_for_search, the later adjacency preload also ignores it. I reproduced this with an automatic ObjectStore budget: resident + preload + both cache capacities consumed the configured budget, and retaining the 4 * vector_count order pushed the configured total above it. Could we reserve the order by reducing preload/cache limits when it is loaded, or include current cache usage and refuse/evict accordingly?

[leaves12138]

Should the unified reader validate top_k > 0 before dispatching (and apply the same check to batch/filtered entry points)? All IVF implementations reject k == 0, while DiskANN returns an empty successful result when an explicit l_search is supplied. I reproduced this through the Python/native API: IVF-FLAT, IVF-SQ, IVF-PQ, and IVF-RQ all fail with k must be greater than 0, but DiskANN returns two empty arrays. The same public VectorSearchParams therefore has index-dependent validity across Rust, C/C++, JNI, and Python.

[leaves12138]

The seeded parallel build is nondeterministic even when the Rayon configuration is fixed. I reproduced this by building the same 512-vector graph repeatedly with the same data, parameters, diskann.seed, and the same 8-thread pool; the resulting VamanaGraphs differ. The workers read and mutate the shared adjacency concurrently within a batch, so search/update visibility depends on scheduling (and batch_size additionally changes with the host thread count). This makes the persisted DiskANN graph depend on runtime interleaving rather than only the input and configured seed. Could we make each batch read a stable graph snapshot and apply updates in a deterministic order, or otherwise preserve the seed's reproducibility contract?

[leaves12138]

Calibration mishandles two supported row-ID cases. First, every negative row ID is treated as padding: DiskANN's codec supports the full i64 range (the new golden fixture even stores i64::MIN), but this loop only counts row_id >= 0. Two completely disjoint negative-ID result sets therefore report stability 1.0 because the denominator stays zero; I reproduced this with [-2, -3] versus [-4, -5]. Second, duplicate multiplicity is overcounted because each left occurrence independently uses right_query.contains: [7, 7] versus [7, 8] also reports 1.0, although one Top-K slot changed, and the new filtered-search tests explicitly support duplicate row IDs. calibrate_l_search can therefore select the smallest width while its actual result rows are still changing. Could we distinguish padding from valid negative IDs and compute a multiset-aware overlap?

[leaves12138]

The Python wrapper needs the same native-handle serialization that was added to the Java wrapper. ctypes.CDLL releases the GIL during these calls, so two Python threads can enter paimon_vindex_reader_search on the same handle; the C ABI then creates two simultaneous &mut VectorIndexReader references. I reproduced this with two ThreadPoolExecutor searches on one fresh DiskANN reader (using a barrier in pread_many): by the second iteration one call hit native panic: attempt to multiply with overflow in pq.rs, after concurrent reader initialization/search corrupted the observed PQ state. close() can race the same way and free an in-use handle. Could we guard every reader native call plus close() with a per-instance lock (and do the same for the other mutable handle wrappers)?

[leaves12138]

The cross-thread race is fixed, but threading.RLock leaves a reentrant use-after-free path. search holds the lock while native code invokes input.pread_many; if that callback calls reader.close() (or another native reader method), it runs on the same thread and re-acquires the RLock, freeing/re-entering the handle while the original Rust call still owns &mut VectorIndexReader. I reproduced this with an IVF-FLAT input whose first search-time pread_many calls reader.close(); the Python process segfaults with exit status 139. The Java wrapper already tracks the native-handle owner and rejects reentry. Could the Python wrappers similarly reject all reentrant native-handle operations rather than permitting them? The writer output callbacks have the same pattern.

[leaves12138]

Automatic IVF search stops after the single inferred nprobe on the unfiltered paths, unlike the filtered paths below that call progressive_ivf_search. This can return padded -1 entries even though the index contains enough rows. I reproduced it with 64 lists: the eight closest lists contain one row each and another list contains 1,000 rows; with total_vectors=1008 and top_k=10, auto infers nprobe=8 and returns only eight IDs plus two paddings, while nprobe=64 returns ten IDs. Could we apply progressive expansion for SearchWidth::Auto here as well, including batch search and the other IVF variants?

[leaves12138]

Allowing odd m for 4-bit PQ here breaks the public IVF-PQ path in two ways. write_index now accepts and writes m=3 with a two-byte code, but IVFPQIndexReader::open_with_header still rejects that file with 4-bit PQ requires even m. Before serialization, both distance::scan_4bit_simd and ivfpq::scan_codes_4bit_transposed also still use m / 2 and iterate only complete nibble pairs. I reproduced the latter with m=3: ProductQuantizer::distance_from_table returns 100 and 1 for two codes whose only difference is the third subquantizer, while scan_codes_4bit returns 0 for both and changes the ranking. DiskANN has dedicated packed-4-bit scanners that handle odd m, but generic IVF-PQ does not. Could we either retain the even-m requirement for IVF-PQ 4-bit indexes or update its writer, reader, and both scan layouts consistently?

[leaves12138]

The persistence path now rejects odd m, but the public in-memory API is still incorrect. IVFPQIndex::with_nbits(..., 4, ...) accepts odd m; both scan_4bit_simd and the transposed/FastScan path use m / 2, so the final subquantizer is ignored. I reproduced this with m=3 and two codes that differ only in the third subquantizer: ProductQuantizer::distance_from_table distinguishes them, while IVFPQIndex::search ranks them as equal, both before and after build_search_structures(). Please either reject odd 4-bit m at construction/configuration time or teach all scanners to process the final low nibble.

[leaves12138]

Could we serialize native-handle operations in the C++ Reader as the Java wrapper now does (and cover destruction/move assignment as well)? The new DiskANN reader mutates lazy initialization, caches, workers, calibrated width, and search statistics during search, while each C entry point converts the same raw handle into a Rust &mut VectorIndexReader. Two C++ threads calling search on one Reader therefore create simultaneous mutable Rust references and can corrupt the reader; I reproduced that failure through the Python wrapper, which reaches the same C ABI. The callback thread-safety note only makes concurrent callbacks from one search safe; it does not make concurrent calls on the reader handle safe. If synchronization belongs to callers instead, the C and C++ APIs need an explicit non-thread-safe contract so users do not reasonably share this read-only-looking object across search threads.

[leaves12138]

Could we avoid swallowing call_method failures here? If a custom VectorIndexInput implementation throws from one of these capability methods, jni returns Error::JavaException but the Java exception remains pending. This code converts the failure to 0 and continues opening the native reader; when the native method returns, the pending exception is propagated and the newly allocated reader handle is unreachable/leaked. The subsequent hint calls also run with an exception already pending. Please either read/validate these hints before constructing the reader and propagate the error, or explicitly clear the exception if falling back to zero is intentional. Negative hint values should probably be rejected consistently as well rather than silently becoming zero.

[leaves12138]

Could we preserve an explicitly supplied expected-vector-count here (and in the Java convenience method)? data is the training sample, not necessarily the final indexed corpus. With nlist=auto, a caller can intentionally pass expected-vector-count=1000000 while training on 10,000 sampled vectors, but this assignment silently replaces it with 10000, changing the inferred nlist from 1024 to 128. Please only infer expected-vector-count when the option is absent; an explicit corpus-size hint should take precedence over the convenience fallback.

[leaves12138]

Could we apply the same platform-width validation added for SearchParams here? This only rejects negative values, so on a 64-bit build memory_budget_bytes = ctypes.c_size_t(-1).value + 1 is silently converted to zero by PaimonVindexReaderOptions. The non-negative read-capability hints and the l_search / calibration top_k arguments have the same unchecked c_size_t conversion. Using operator.index and validating against SIZE_MAX before constructing the ctypes values would keep caller errors from changing the requested behavior.

[leaves12138]

The package declares requires-python = ">=3.9", but these int | None annotations require Python 3.10 when annotations are evaluated normally. On Python 3.9 the script fails during function definition before the CLI can run. Please use Optional[int] or add from __future__ import annotations, and ideally include the minimum supported Python version in the tool smoke test.

[leaves12138]

The high-level trainer reservoir is now budgeted, but the public low-level DiskAnnIndex::train path still calls pq_training_plan with one sample buffer. For cosine training with downsampling, this method simultaneously retains the sampled buffer returned by bounded_pq_training_sample and the normalized owned buffer returned by preprocess_vectors. With d=1024, m=256, 50,000 vectors, and a 128 MiB budget, the current plan selects 28,111 samples while the same peak model with two coexisting sample buffers selects only 14,379. Could this path pass the metric-aware sample-buffer count as well, and return an error rather than silently falling back to a one-vector plan when the budget is infeasible?

[leaves12138]

vector_read_plan reports the originally resolved cache budgets, not the effective capacities after lazy state consumes memory. ensure_row_id_order can shrink both shared caches through resize_shared_cache_budgets, but self.options remains unchanged. I extended diskann_row_id_order_reserves_budget_from_shared_caches: after loading the row-ID order, the public plan still reported 32,768 cache bytes while the actual combined cache capacity was 10,060 bytes. Since this API is documented as the concrete read plan, could it report the current hot-adjacency size and shared-cache capacities instead of the initial desired values?

[leaves12138]

Serializing the C++ handle fixes the cross-thread race, but a plain std::mutex turns callback reentry into a permanent deadlock. I reproduced this with an IVF-FLAT InputFile::read_ranges_fn that calls reader.metadata() during reader.search(); the callback blocks reacquiring native_handle_mutex_, and the process timed out after five seconds. Java and Python now explicitly track the owning thread and reject this pattern. Could the C++ wrapper do the same rather than blocking same-thread reentry?

[leaves12138]

Streaming the oversized list bounds the I/O payload, but batch search retains one chunk-sized distance array for every query. I added a temporary unit test that scanned a 65,536-row chunk with 64 SqScanScratch instances; after the loop they retained 16,777,216 bytes (64 * 65,536 * sizeof(f32)) rather than one reusable 256 KiB buffer. Production chunk sizes can be much larger, so a large query batch can still consume GiBs despite the 64 MiB payload bound. Could queries be processed in bounded groups, or use scratch that is released/reused instead of one retained buffer per query? The transposed 8-bit IVF-PQ oversized-list path has the same per-query distance-scratch shape.

[leaves12138]

The input path now bounds local references, but the JNI output path still retains one JByteArray local reference per write_all call until the whole native writeIndex frame returns. I reproduced this with a one-vector IVF-FLAT index and nlist=50,000; under -Xcheck:jni, the local-reference count grew past 69,000 before the run timed out. The offset table alone triggers multiple writes per list. Could write_all use a bounded local frame or explicitly delete_local_ref(jbuf) after call_method?

[leaves12138]

base_read is local to this if block, but the replacement callback captures it by reference through [&] and is invoked only after the block has ended. That leaves a dangling std::function reference in input.read_ranges_fn. The current C++ CI aborts while opening the IVF-FLAT reader (failed to open vector index reader), while my local optimized build happens to pass, which is consistent with this use-after-scope. Could the callback capture base_read by value (for example, base_read = std::move(base_read)) instead?

[leaves12138]

The owner-thread check still deadlocks when the storage callback is invoked by a native worker thread rather than by the thread that entered the Reader. I reproduced this with RAYON_NUM_THREADS=4, a DiskANN search_batch, and an 8 KiB Reader budget so adjacency reads remain cold. A worker-thread read_ranges_fn callback called reader.metadata(); the outer search held operation_mutex_ on the caller thread and waited for the worker, while the worker saw a different owner_ and blocked on the same mutex. The process timed out after five seconds. Since DiskANN explicitly invokes callbacks from separate query workers, could reentry be tracked for the active native operation/callback context rather than only by comparing thread IDs?

[leaves12138]

The same cross-thread callback reentry deadlock exists in the Python wrapper. _NativeHandleLock uses thread-local state, so a Rayon worker callback sees active == false and then blocks on the lock held by the caller whose native batch search is waiting for that worker. I reproduced this with a 5,000-row DiskANN index, RAYON_NUM_THREADS=4, a resident budget only 8 KiB above the required 75,254 bytes, and a 64-query batch. The control run completed with 488 reads across the caller plus four worker threads; when the first worker-thread pread_many callback called reader.metadata(), the process timed out after ten seconds. Could callback reentry be rejected based on the active native operation/callback context rather than per-thread local state?