Reviewing files that changed from the base of the PR and between dce4a94 and f505987.

• lib/common/common/src/lib.rs (1 hunks)
• lib/common/common/src/maybe_uninit.rs (1 hunks)
• lib/segment/src/vector_storage/chunked_mmap_vectors.rs (5 hunks)
• lib/segment/src/vector_storage/chunked_vector_storage.rs (2 hunks)
• lib/segment/src/vector_storage/dense/appendable_dense_vector_storage.rs (2 hunks)
• lib/segment/src/vector_storage/dense/memmap_dense_vector_storage.rs (2 hunks)
• lib/segment/src/vector_storage/dense/mmap_dense_vectors.rs (2 hunks)
• lib/segment/src/vector_storage/in_ram_persisted_vectors.rs (2 hunks)
• lib/segment/src/vector_storage/multi_dense/appendable_mmap_multi_dense_vector_storage.rs (0 hunks)
• lib/segment/src/vector_storage/multi_dense/simple_multi_dense_vector_storage.rs (1 hunks)
• lib/segment/src/vector_storage/quantized/quantized_custom_query_scorer.rs (0 hunks)
• lib/segment/src/vector_storage/quantized/quantized_query_scorer.rs (0 hunks)
• lib/segment/src/vector_storage/query_scorer/custom_query_scorer.rs (2 hunks)
• lib/segment/src/vector_storage/query_scorer/metric_query_scorer.rs (2 hunks)
• lib/segment/src/vector_storage/query_scorer/mod.rs (2 hunks)
• lib/segment/src/vector_storage/query_scorer/multi_custom_query_scorer.rs (2 hunks)
• lib/segment/src/vector_storage/query_scorer/multi_metric_query_scorer.rs (2 hunks)
• lib/segment/src/vector_storage/query_scorer/sparse_custom_query_scorer.rs (0 hunks)
• lib/segment/src/vector_storage/vector_storage_base.rs (4 hunks)

• lib/segment/src/vector_storage/quantized/quantized_query_scorer.rs
• lib/segment/src/vector_storage/query_scorer/sparse_custom_query_scorer.rs
• lib/segment/src/vector_storage/multi_dense/appendable_mmap_multi_dense_vector_storage.rs
• lib/segment/src/vector_storage/quantized/quantized_custom_query_scorer.rs

• GitHub Check: test-shard-snapshot-api-s3-minio
• GitHub Check: test-snapshot-operations-s3-minio
• GitHub Check: test-consistency
• GitHub Check: test-consensus
• GitHub Check: test-low-resources
• GitHub Check: test (macos-latest)
• GitHub Check: test
• GitHub Check: test (windows-latest)
• GitHub Check: test-consensus-compose
• GitHub Check: test
• GitHub Check: test
• GitHub Check: test (ubuntu-latest)
• GitHub Check: Basic TLS/HTTPS tests

lib/common/common/src/lib.rs (1)

15-15: Looks good: New module for optimizing vector initialization.

The addition of the maybe_uninit module aligns with the PR's objective to optimize performance by eliminating unnecessary vector initialization overhead. This module likely contains the backported version of MaybeUninit::fill_from mentioned in the PR description.

lib/segment/src/vector_storage/query_scorer/custom_query_scorer.rs (2)

3-3: Appropriate import for MaybeUninit.

The import of MaybeUninit is necessary for the optimization changes that follow.

---

96-99: Performance improvement: Using uninitialized memory.

This implementation now uses MaybeUninit to avoid the overhead of initializing the temporary array with default values before it gets populated with actual data. This matches the PR objective of optimizing initialization of temporary arrays used during vector operations.

lib/segment/src/vector_storage/multi_dense/simple_multi_dense_vector_storage.rs (1)

25-25:

✅ Verification successful

Note the deliberate removal of batch method.

Based on the AI summary, the get_batch_multi method has been completely removed from this implementation. This is part of the broader changes to optimize batch operations by using MaybeUninit elsewhere. The import has been updated to only include what's actually used in the file.

---

🏁 Script executed:

#!/bin/bash
# Check for any remaining references to get_batch_multi in this file
rg "get_batch_multi" lib/segment/src/vector_storage/multi_dense/simple_multi_dense_vector_storage.rs

Length of output: 100

---

Removal Verified: The deliberate elimination of get_batch_multi is confirmed

• The search command produced no matches in lib/segment/src/vector_storage/multi_dense/simple_multi_dense_vector_storage.rs, indicating that get_batch_multi is no longer present.
• The updated import now correctly includes only CHUNK_SIZE and StoredRecord, aligning with the optimization strategy utilizing MaybeUninit for batch operations.

lib/segment/src/vector_storage/query_scorer/mod.rs (2)

8-8: Necessary import for the default implementation.

Adding the VECTOR_READ_BATCH_SIZE import is necessary for the new default implementation of score_stored_batch.

---

22-30: Good addition: Default implementation for batch scoring.

Adding a default implementation for score_stored_batch provides a fallback for implementations that don't need specialized batch processing. This is a sensible change that:

1. Ensures all implementations have this method available
2. Maintains the batch size constraints with assertions
3. Provides a simple implementation that calls score_stored for each ID individually

The implementation is correct and aligns with the batch processing pattern seen in the specialized implementors of this trait, while allowing custom optimizations in specific implementations.

lib/segment/src/vector_storage/in_ram_persisted_vectors.rs (1)

1-1: Good addition of necessary import.

The import of MaybeUninit is needed for the changed function signature below.

lib/common/common/src/maybe_uninit.rs (1)

1-33: Well-implemented backport of MaybeUninit::fill_from functionality.

The implementation follows good practices:

• Clear documentation with usage constraints
• Properly annotated safety comments
• Explicit assertion for unsupported types
• Efficient memory handling

The code correctly:

1. Fills uninitialized memory with values from an iterator
2. Tracks the number of initialized elements
3. Returns both the initialized portion and the remainder

One minor observation:

const { assert!(!std::mem::needs_drop::<I::Item>(), "Not supported") };

This compile-time assertion prevents usage with types that need dropping. This is a reasonable limitation for a simplified backport, and it's well-documented.

lib/segment/src/vector_storage/dense/memmap_dense_vector_storage.rs (2)

4-4: Good addition of necessary import.

The import of MaybeUninit is correctly added for the modified function signature.

---

144-151: Method signature improvement for better memory management.

This change enhances performance by removing the need to initialize memory that would immediately be overwritten. The method now follows the pattern of accepting uninitialized memory and returning properly initialized values.

The implementation correctly delegates to the mmap_store.get_vectors method with the updated signature.

lib/segment/src/vector_storage/dense/mmap_dense_vectors.rs (3)

3-3: Good addition of necessary import.

The import of MaybeUninit is correctly added for the modified function signatures.

---

9-9: Good addition of helper function import.

The import of maybe_uninit_fill_from is appropriately added to support the new implementation approach.

---

160-177: Excellent refactoring to improve memory usage efficiency.

This implementation:

1. Properly updates the function signature to use MaybeUninit
2. Uses the new maybe_uninit_fill_from helper function to efficiently populate vectors
3. Maintains the same logic branches (optimized path for sequential reads)
4. Returns the properly initialized slice

The refactoring eliminates unnecessary default initialization of memory that would be immediately overwritten, which should improve performance especially for batch operations with small vectors as mentioned in the PR objectives.

I particularly like how the two code paths (with and without prefetching) both use the same pattern, making the code more consistent and easier to understand.

lib/segment/src/vector_storage/query_scorer/metric_query_scorer.rs (2)

3-3: Added MaybeUninit import for optimized vector initialization.

---

77-81: Improved memory efficiency by using uninitialized memory.

The function now properly uses MaybeUninit to avoid unnecessary initialization of vector values that would immediately be overwritten. This implementation optimizes the performance, especially when working with small vectors, as mentioned in the PR objectives.

lib/segment/src/vector_storage/query_scorer/multi_custom_query_scorer.rs (5)

2-2: Added MaybeUninit import for optimized vector initialization.

---

130-131: Improved assertions for batch processing.

The assertions have been simplified to directly check ids.len() rather than calculating a separate batch size variable, making the code more direct and easier to understand.

---

133-136: Optimized memory initialization for batch vector retrieval.

Using MaybeUninit to allocate uninitialized memory for vectors improves performance by eliminating the unnecessary initialization of default values that would be immediately overwritten.

---

138-138: Improved calculation of total loaded vectors.

The code now directly iterates over the vectors array to calculate the total number of loaded vectors, which is more concise and eliminates the need for slicing.

---

144-144: Simplified batch processing loop.

The loop now directly uses ids.len() instead of a separate batch size variable, making the code more straightforward.

lib/segment/src/vector_storage/query_scorer/multi_metric_query_scorer.rs (5)

2-2: Added MaybeUninit import for optimized vector initialization.

---

112-113: Improved assertions for batch processing.

The assertions have been simplified to directly check ids.len() rather than calculating a separate batch size variable, making the code more direct and easier to understand.

---

115-118: Optimized memory initialization for batch vector retrieval.

Using MaybeUninit to allocate uninitialized memory for vectors improves performance by eliminating the unnecessary initialization of default values that would be immediately overwritten.

---

120-120: Improved calculation of total read vectors.

The code now directly iterates over the vectors array to calculate the total number of read vectors, which is more concise and eliminates the need for slicing.

---

126-126: Simplified batch processing loop.

The loop now directly uses ids.len() instead of a separate batch size variable, making the code more straightforward.

lib/segment/src/vector_storage/dense/appendable_dense_vector_storage.rs (2)

3-3: Added MaybeUninit and helper function import.

The code now imports both MaybeUninit for uninitialized memory and the maybe_uninit_fill_from helper function that was introduced in this PR to efficiently work with uninitialized memory.

Also applies to: 10-10

---

73-85: Optimized batch vector retrieval with MaybeUninit.

The get_dense_batch method has been rewritten to use MaybeUninit for both the input parameter and the internal vector offset handling, which eliminates unnecessary initialization of memory. This implementation is now consistent with the similar code in other storage implementations, as seen in the relevant code snippets from memmap_dense_vector_storage.rs and vector_storage_base.rs.

Key improvements:

1. Changed parameter type to use MaybeUninit<&'a [T]> instead of initialized references
2. Uses uninitialized memory for vector offsets
3. Efficiently fills the offsets using the new maybe_uninit_fill_from helper

This approach aligns with the PR objective of optimizing performance, particularly for operations involving small vectors.

lib/segment/src/vector_storage/chunked_vector_storage.rs (2)

1-1: Good use of MaybeUninit for optimized initialization.

Adding this import supports the new optimized approach for handling uninitialized memory in batch operations.

---

51-55: Well-designed API change for enhanced performance.

The signature change for get_batch is a good optimization that eliminates unnecessary initialization of vectors. By using MaybeUninit, the method now:

1. Takes uninitialized memory and populates it directly
2. Returns a reference to the initialized portion
3. Avoids the overhead of initializing the array with default values

This aligns perfectly with the PR objective of enhancing performance in batch operations.

lib/segment/src/vector_storage/vector_storage_base.rs (3)

1-1: Added imports support the optimization approach.

The added imports are necessary to support the new optimization pattern using MaybeUninit and the helper function maybe_uninit_fill_from.

Also applies to: 8-8, 29-29

---

132-138: Default implementation improved for DenseVectorStorage.

The default implementation of get_dense_batch now uses MaybeUninit and maybe_uninit_fill_from to efficiently initialize the vectors array. The new implementation:

1. Takes uninitialized memory as input
2. Populates it with vector references
3. Returns only the initialized portion

This is more efficient than initializing with default values first and offers a better default implementation for implementors of this trait.

---

154-162: MultiVectorStorage batch implementation optimized.

The get_batch_multi implementation now uses the same optimization pattern with these key improvements:

1. Changed parameter to use MaybeUninit for uninitialized memory
2. Added proper debug assertions to validate inputs
3. Used maybe_uninit_fill_from for safe initialization
4. Returns the initialized portion

The debug assertions ensure that the vectors and keys arrays have matching lengths and that we don't exceed the batch size limit, which is good for catching errors early in development.

lib/segment/src/vector_storage/chunked_mmap_vectors.rs (4)

3-3: Good addition of necessary imports.

These imports support the new optimization pattern using MaybeUninit and the helper function maybe_uninit_fill_from.

Also applies to: 7-7

---

299-316: Optimized batch retrieval implementation.

The get_batch implementation has been improved to use MaybeUninit for more efficient memory handling. The key improvements are:

1. Changed the function signature to use MaybeUninit and return the initialized results
2. Maintained the same debug assertions to verify input correctness
3. Used maybe_uninit_fill_from to populate vectors directly without prior initialization
4. Added proper error handling with panic messages when vectors aren't found

This improves performance by eliminating the unnecessary initialization step while maintaining safety.

---

405-411: Trait implementation aligned with optimization pattern.

The trait implementation for get_batch has been updated to match the new method signature and properly delegates to the ChunkedMmapVectors::get_batch method. This maintains consistency across the codebase.

---

467-478: Test code properly updated to use new pattern.

The test code has been updated to:

1. Use MaybeUninit::uninit() for buffer initialization instead of pre-filled values
2. Capture the return value from get_batch which now returns the initialized vectors
3. Maintain the same validation logic to ensure correctness

This validates that the new implementation works correctly while testing the performance optimization.