Reviewing files that changed from the base of the PR and between f50292e and 24f51b8.

• lib/collection/src/collection/mod.rs (3 hunks)
• lib/collection/src/collection/snapshots.rs (2 hunks)
• lib/collection/src/collection/state_management.rs (2 hunks)
• lib/collection/src/collection_state.rs (2 hunks)
• lib/collection/src/shards/shard_holder/mod.rs (6 hunks)
• lib/collection/src/shards/shard_holder/shard_mapping.rs (2 hunks)
• lib/storage/src/content_manager/toc/collection_meta_ops.rs (1 hunks)
• src/migrations/single_to_cluster.rs (1 hunks)

lib/collection/src/collection_state.rs (1)

26-27: Implementation matches TODO comment - good work!

This change successfully implements the TODO comment by replacing the wrapper type with the new ShardKeyMapping type.

src/migrations/single_to_cluster.rs (2)

100-100: Simplified accessing of shard key mapping entries

The change from using .to_map() to directly using .iter() on the shards_key_mapping object aligns with the overall refactoring to simplify the ShardKeyMapping implementation.

---

104-104: Removed reference operator when accessing shard ID

Removed the reference operator (&) from shard_id when calling get(), which matches with how the new ShardKeyMapping type handles shard IDs.

lib/collection/src/collection/state_management.rs (2)

179-179: Signature successfully updated to reflect the new type.

This change aligns well with the overall refactor to use ShardKeyMapping.

---

12-12:

Details

❓ Verification inconclusive

Confirm old references are fully removed.

Swapping from ShardKeyMappingWrapper to ShardKeyMapping is consistent with the new design. Please ensure no references to ShardKeyMappingWrapper remain in the codebase.

---

🏁 Script executed:

#!/bin/bash
rg 'ShardKeyMappingWrapper'

Length of output: 27

---

Let's perform a slightly more targeted search before finalizing this review comment. I'll re-run the search specifically in Rust files to be sure there are no stale references:

Once you verify that this command produces no output, we can confidently conclude that all references to ShardKeyMappingWrapper have been removed.

Rewritten Review Comment:

Ensure complete removal of ShardKeyMappingWrapper.
The change from ShardKeyMappingWrapper to ShardKeyMapping in lib/collection/src/collection/state_management.rs (line 12) aligns with the new design. Please confirm by re-running a thorough search (e.g., using rg -t rust 'ShardKeyMappingWrapper') that no lingering references to ShardKeyMappingWrapper exist anywhere in the codebase.

---

🏁 Script executed:

#!/bin/bash
rg -t rust 'ShardKeyMappingWrapper'

Length of output: 35

---

Action: Confirm removal of ShardKeyMappingWrapper from the entire codebase.
The change in lib/collection/src/collection/state_management.rs (line 12) replacing ShardKeyMappingWrapper with ShardKeyMapping follows the new design. Automated searches using both generic and Rust-specific queries have not returned any occurrences of ShardKeyMappingWrapper. However, the output was minimal, so please perform a manual review (or run an alternate search) across the repository to fully ensure no residual references remain.

lib/collection/src/collection/mod.rs (3)

50-50: Import change for the refactored type.

Switching the import to ShardKeyMapping reflects the new approach. Looks good.

---

108-108: Parameter type refactor.

Replacing Option<ShardKeyMappingWrapper> with Option<ShardKeyMapping> ensures consistency with the updated mapping structure.

---

352-353: Correct approach to return owned keys.

Using .cloned().collect() to produce a vector of owned ShardKeys is concise and clear.

lib/collection/src/shards/shard_holder/shard_mapping.rs (7)

2-2: Additional import for ops.

Importing std::ops to implement Deref and DerefMut is appropriate. No issues found.

---

12-16: New struct definition for simplified shard mapping.

Defining ShardKeyMapping with serde attributes and derived traits appears solid. Good job incorporating the new format handling.

---

58-63: Iterator logic is straightforward.

Collecting all shard IDs through iter_shard_ids is clear and follows well from using HashSet for uniqueness.

---

68-74: Check for multiple keys referencing the same ID.

This returns the first key found if a single shard ID appears under multiple shard keys. Please verify that such overlap cannot occur or else handle conflicts explicitly.

---

77-91: Backward compatibility for old vs. new storage formats.

Parsing the helper’s Old or New variants supports existing and new workflows effectively. Nicely done.

---

94-99: Untagged enum for flexible deserialization.

Untagged serde approach provides a seamless fallback for legacy data. Good choice here.

---

117-139: Simple intermediate struct for new format.

KeyIdsPair is straightforward and helps cleanly separate serialization logic. Looks good.

lib/collection/src/shards/shard_holder/mod.rs (6)

17-17: LGTM! Updated import for ShardKeyMapping.

This import change is part of the refactoring to use the direct ShardKeyMapping type instead of a wrapper.

---

59-59: LGTM! Updated field type from wrapper to direct type.

Changed the field type from SaveOnDiskShardKeyMappingWrapper to SaveOnDisk<ShardKeyMapping>, removing an unnecessary wrapper layer.

---

79-80: LGTM! Updated initialization to use direct ShardKeyMapping type.

Simplified initialization now loads a SaveOnDisk<ShardKeyMapping> directly instead of wrapping it.

---

119-120: LGTM! Simplified return type.

Method now returns ShardKeyMapping directly instead of wrapping it, consistent with the type simplification.

---

275-275: LGTM! Updated parameter type.

Changed parameter type from ShardKeyMappingWrapper to ShardKeyMapping for consistency.

---

132-139:

Details

❓ Verification inconclusive

LGTM! Streamlined implementation with direct mapping.

The implementation has been simplified to:

1. Extract the shard ID to key mapping
2. Write the mapping directly to storage
3. Update the local cache

This is cleaner than the previous approach that required conversion to a map first.

---

🏁 Script executed:

#!/bin/bash
# Verify that the ShardKeyMapping type implements the necessary shards() method
ast-grep --pattern 'impl $_ for ShardKeyMapping {
  $$$
  fn shards($$$) -> $_ {
    $$$
  }
  $$$
}'

Length of output: 103

---

Action Required: Verify ShardKeyMapping Implementation

The updated streamlined implementation looks clean and correct. However, the initial automated check for the shards() method on ShardKeyMapping did not return output. Please manually verify that the ShardKeyMapping type indeed implements a shards() method with the correct behavior, ensuring that it returns the expected mapping.

Bullet points for review:

• Confirm that ShardKeyMapping has an implementation of fn shards(&self) -> _ (with the correct return type) that extracts the shard ID to key mapping.
• Ensure that the behavior of this method matches the expectations in the streamlined code (i.e., direct mapping extraction, proper updating of storage, and cache).

Reviewing files that changed from the base of the PR and between 24f51b8 and 2f8a986.

• lib/collection/src/shards/shard_holder/shard_mapping.rs (2 hunks)

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

lib/collection/src/shards/shard_holder/shard_mapping.rs (8)

2-2: Well-structured encapsulation with Deref pattern

The refactoring encapsulates the shard key mapping inside a struct with proper derive traits for serialization/deserialization, which is a better design approach than the previous type alias. Using the Deref pattern (importing std::ops) allows for backward compatibility while providing proper encapsulation.

Also applies to: 12-16

---

18-30: Good implementation of Deref and DerefMut

The Deref and DerefMut traits allow the struct to be used as if it were a HashMap directly, providing backward compatibility with existing code while maintaining proper encapsulation. This is a standard Rust idiom for wrapping collections.

---

35-44: Method implementation looks correct

The implementation of shards() correctly uses flat_map to transform the mapping from shard key to shard IDs into a mapping from shard ID to shard key. The approach is efficient and clear.

---

59-63: Good iterator implementation

Creating a dedicated iterator method is a good practice in Rust, allowing for more flexible and potentially more efficient code. Using copied() instead of cloning each shard ID is also more efficient.

---

77-92: Clean implementation of conversion from SerdeHelper

The conversion from the serialization helper to the actual struct is clean and handles both the new and old formats correctly. The matching is straightforward and maps directly to the struct field.

---

94-99: Good serialization strategy

Using an untagged enum for serialization is a good approach to handle multiple formats without complicating the main struct. This separates serialization concerns from the core data structure.

---

101-117: Format selection aligns with previous discussions

The logic for selecting between new and old formats correctly follows the requirements discussed in the previous reviews. It uses the new format if either the feature flag is enabled or if any numeric shard key is used, which preserves type information.

This implementation matches the discussion in the past review comments where using numeric shard keys should opt-in to the new format to preserve type information, even if it breaks backward compatibility.

---

119-141: Well-designed helper struct

The KeyIdsPair struct and its implementations are well-designed for serialization purposes. The into_parts method and From implementation make the conversion between the pair and tuple formats clean and explicit.

Reviewing files that changed from the base of the PR and between 2f8a986 and 3ec1bd0.

• lib/collection/src/collection/mod.rs (4 hunks)
• lib/collection/src/collection/state_management.rs (3 hunks)
• lib/collection/src/collection_state.rs (2 hunks)
• lib/collection/src/shards/shard_holder/mod.rs (6 hunks)
• lib/collection/src/shards/shard_holder/shard_mapping.rs (2 hunks)

• lib/collection/src/collection_state.rs
• lib/collection/src/collection/state_management.rs
• lib/collection/src/collection/mod.rs

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

lib/collection/src/shards/shard_holder/shard_mapping.rs (8)

2-2: No issues with adding std::ops.
This import is clearly needed for implementing Deref and DerefMut.

---

12-16: Structured and well-derived type definition.
The use of #[derive(Default)] along with Deserialize and Serialize streamlines the creation and persistence of ShardKeyMapping. Overall, this approach is clean.

---

32-43: Efficient method for reversing the shard mapping.
Generating an inverse map on-the-fly is straightforward and maintains clarity. Complexity is reasonable for most use cases.

---

57-62: Iterator method looks good.
The fluent approach of returning an iterator (rather than a collected list) is memory-friendly.

---

75-91: Deserialization logic aligns well with migrating formats.
The fallback strategy for the old format ensures backward compatibility while enabling updates to the new format.

---

92-105: Preserving older format under a feature flag is a sensible approach.
The plan to remove this in a future version (as noted by TODO) is clear. No issues.

---

107-123: Correctly switching to the new format for numeric shard keys.
This addresses previous concerns about losing numeric type information, matching user and reviewer feedback.

---

126-148: Helper struct and conversions are straightforward.
Round-trip conversions for (ShardKey, HashSet<ShardId>) look clean. No further recommendations.

lib/collection/src/shards/shard_holder/mod.rs (9)

17-17: Import looks appropriate.
No concerns regarding referencing ShardKeyMapping here.

---

59-59: Storing the mapping in SaveOnDisk is consistent.
Persisting ShardKeyMapping with the same pattern as other data structures is logical and maintains uniformity.

---

79-80: Initialization strategy makes sense.
Using load_or_init_default is straightforward for ensuring the mapping is present after restarts.

---

119-120: Simple getter returning a cloned mapping.
This is fine if returning a full copy is expected. Monitor performance if the mapping grows large.

---

132-133: Straightforward signature update.
Accepting a plain ShardKeyMapping parameter aligns with the new design.

---

134-135: Good use of shard_id_to_shard_key().
Generating the inverse map once prevents repeated computation, improving readability and performance.

---

137-138: Overwriting the stored mapping is handled cleanly.
Using .write_optional with Some(shard_key_mapping) is clear and consistent with the rest of the codebase.

---

139-140: Properly caching the inverse mapping in memory.
No issues with synchronizing it right after the write operation.

---

275-276: apply_shards_state parameter update is consistent.
Integrating ShardKeyMapping aligns the method with the rest of the refactor.