Cosmos Driver: Hierarchical Partition Key Support by simorenoh · Pull Request #4087 · Azure/azure-sdk-for-rust

simorenoh · 2026-04-03T19:40:41Z

Adds MultiHash EPK (Effective Partition Key) computation and prefix partition key routing infrastructure to the driver crate, enabling hierarchical partition key (HPK) support for containers with multiple partition key paths.

What changed

MultiHash EPK computation (`effective_partition_key.rs`)

Previously, EffectivePartitionKey::compute() fell through to single-hash V2 for PartitionKeyKind::MultiHash, producing incorrect EPKs. MultiHash requires each component to be hashed independently — this PR adds that per-component hashing.

Added PartitionKeyKind::MultiHash arm to compute() routing to new effective_partition_key_multi_hash_v2()
Each component is independently V2-encoded → MurmurHash3-128 → byte-reversed → top-2-bit masked → hex-encoded, then concatenated (N×32 hex chars)
Extracted shared hash_v2_to_epk() helper used by both single-hash and multi-hash paths
Algorithm verified against cross-SDK baselines (.NET, Go, Java) via existing testdata/*.xml fixtures

Prefix EPK range computation (`effective_partition_key.rs`)

Added compute_range() for partial/prefix partition keys (fewer components than the container definition)
Full key → point range (start == end); prefix key on MultiHash → [prefix_epk, prefix_epk + "FF") range covering all possible suffix completions

Prefix routing in PK range cache (`partition_key_range_cache.rs`)

Added resolve_partition_key_range_ids() that handles both full and prefix partition keys
Full key: point lookup (single range ID); prefix key: EPK range → resolve_overlapping_ranges() → multiple range IDs for fan-out

Tests

9 new unit tests covering:

Single/two/three-component MultiHash EPK computation with expected value verification
MultiHash with Undefined component (partial HPK)
MultiHash vs single-hash divergence for multi-component keys
compute_range() for full keys, prefix keys (1-of-3, 2-of-3), and single-hash (always point)

Follow-up: FeedRange API (PR #3987)

PR #3987 introduces feed_range_from_partition_key(), which currently uses the SDK's hash.rs to compute EPKs. For MultiHash containers, this hits the existing stub and returns incorrect results. Once that method is updated to route through the driver's EPK computation (as part of the SDK-to-driver cutover), it will get correct MultiHash support for free from this PR. Prefix HPK support for the FeedRange API (returning multiple feed ranges for partial keys) will additionally need the compute_range() infrastructure added here.

Follow-up: Query tests and thorough testing

This is the first step of the end-to-end implementation of this feature. The remaining work, piecing together the operations to this logic and ensuring that queries can also use it, rely on the migration to the driver. We need for the migration of requests to the driver to be finalized before we can add these tests.

github-actions · 2026-04-03T19:53:42Z

API Change Check

APIView identified API level changes in this PR and created the following API reviews

azure_data_cosmos_driver

Copilot

Pull request overview

This PR adds hierarchical partition key (HPK) support to the Cosmos driver by implementing correct MultiHash effective partition key (EPK) computation and enabling prefix partition key routing via EPK range fan-out in the partition key range cache.

Changes:

Implement MultiHash V2 EPK computation (per-component hashing) and refactor shared V2 hash→EPK conversion.
Add EffectivePartitionKey::compute_range() to compute point vs prefix EPK ranges for routing.
Add PartitionKeyRangeCache::resolve_partition_key_range_ids() to return one or many PK range IDs depending on full vs prefix PKs.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 5 comments.

File	Description
sdk/cosmos/azure_data_cosmos_driver/src/models/effective_partition_key.rs	Adds MultiHash EPK logic and prefix-range computation, plus new unit tests.
sdk/cosmos/azure_data_cosmos_driver/src/driver/cache/partition_key_range_cache.rs	Adds prefix-aware range ID resolution that can fan out to overlapping ranges.
sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md	Notes HPK/MultiHash support in the unreleased changelog.

tvaron3

besides updating the test coverage, lgtm

analogrelay

The blocking changes are a few asserts and some test fixes. I got to speculating on putting semantic validation in place on the EffectivePartitionKey type, but it might be too much to add for minimal value. The main thing I was going for is the idea that we could normalize EPKs at the time EffectivePartitionKey is created and then comparisons can use mostly standard logic.

FabianMeiswinkel

LGTM - Thanks!

FabianMeiswinkel

LGTM

This PR adds the `FeedRange` type and feed range APIs to `azure_data_cosmos`, enabling consumers to work with physical partition ranges. It uses the driver's EPK (Effective Partition Key) computation from [PR #4087](#4087) for correct MultiHash (hierarchical partition key) support, including prefix partition keys. ### Changes #### `sdk/cosmos/azure_data_cosmos/src/feed_range.rs` (new) New `FeedRange` type representing a contiguous `[min, max)` EPK range: - `FeedRange::full()` — the complete EPK space - `contains()` / `overlaps()` — containment and overlap checks - `Display` / `FromStr` — base64-encoded JSON serialization (cross-SDK compatible) - `Serialize` / `Deserialize` — structured JSON for embedding in documents - Runtime validation of inclusivity flags and min ≤ max ordering - 20 unit tests covering serialization, parsing, containment, overlap, and rejection paths #### `sdk/cosmos/azure_data_cosmos/src/clients/container_client.rs` Two new public methods on `ContainerClient`: | Method | Return type | Description | |---|---|---| | `read_feed_ranges(options)` | `Vec<FeedRange>` | Returns one feed range per physical partition | | `feed_range_from_partition_key(pk, options)` | `Vec<FeedRange>` | Maps a partition key to its covering feed ranges | `feed_range_from_partition_key` handles three cases: | Input | PK Definition | Behavior | |---|---|---| | Full key (N of N components) | Any | `DriverEpk::compute()` → point lookup → 1 feed range | | Prefix key (< N components) | MultiHash | `DriverEpk::compute_range()` → overlapping lookup → 1+ feed ranges | | Prefix key (< N components) | Hash (single) | Error — prefix not supported | Both methods support `ReadFeedRangesOptions::with_force_refresh()` to bypass the routing map cache after partition splits. #### `sdk/cosmos/azure_data_cosmos/src/handler/container_connection.rs` Added three accessor methods: - `partition_key_definition()` — returns the PK definition from the eagerly-resolved `ContainerReference` - `resolve_routing_map(force_refresh)` — resolves the SDK-side routing map with optional cache bypass #### `sdk/cosmos/azure_data_cosmos/src/hash.rs` - Added `PartialOrd`, `Ord`, `Hash` derives to `EffectivePartitionKey` (needed by `FeedRange`) - Added `From<String>` and `From<&str>` impls for `EffectivePartitionKey` - Made `MIN_INCLUSIVE_EFFECTIVE_PARTITION_KEY` / `MAX_EXCLUSIVE_EFFECTIVE_PARTITION_KEY` `pub(crate)` for reuse in `feed_range.rs` #### `sdk/cosmos/azure_data_cosmos_driver/src/models/` (visibility changes) Made the following driver types public so the SDK can use them for EPK computation and routing: - `effective_partition_key` module → `pub` - `partition_key_range` module → `pub` - `EffectivePartitionKey` struct → `pub` - `PartitionKeyRange` struct → `pub` - `PartitionKey::values()` → `pub` #### `sdk/cosmos/azure_data_cosmos/src/options/mod.rs` Added `ReadFeedRangesOptions` with `with_force_refresh(bool)` for cache bypass control. #### `sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_feed_ranges.rs` (new) 5 emulator integration tests: | Test | What it validates | |---|---| | `read_feed_ranges_returns_physical_partitions` | Returns ≥ 2 ranges with 11000 RU/s, non-overlapping, serializable | | `feed_range_from_partition_key_maps_correctly` | Full key → single range matching a physical partition, deterministic | | `feed_range_from_full_hpk_returns_single_range` | Full 3-component HPK → single range within full EPK space | | `feed_range_from_prefix_hpk_returns_ranges` | Prefix HPK (1-of-3, 2-of-3) → 1+ non-overlapping ranges | | `feed_range_from_partition_key_single_hash_full_key` | Regression: full key on single-hash container succeeds | ### Architecture Note The feed range methods use a hybrid approach: - **EPK computation** → driver's `EffectivePartitionKey::compute()` / `compute_range()` (correct MultiHash support from [PR #4087](#4087)) - **Routing map lookup** → SDK-side `PartitionKeyRangeCache` (via `ContainerConnection::resolve_routing_map()`) This is because `CosmosDriver` does not yet expose a `resolve_routing_map()` method — the driver's `PartitionKeyRangeCache` exists but isn't wired into the driver struct. Once that infrastructure is added (tracked separately), the feed range methods will switch to the driver's routing map, eliminating the dual-cache situation. --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Addresses three findings from the PR Deep Reviewer second pass: - F-A: replace the non-existent epk_length_aware_cmp citation with EffectivePartitionKey's Ord/cmp impl, cite the actual epk_cmp_* tests in container_routing_map.rs and the binary_search_by consumer site. Point PR Azure#4087 at the correct claim. - F-B: fix the numerically wrong UUID worked example. The previous example for 12345678-1234-5678-1234-567812345678 wrote MSB bytes 78 56 34 12 34 12 78 56, conflating writeLongLE with byte-reversal of the hyphen groups. Replace with 0a1b2c3d-4e5f-6789-abcd- ef0123456789 so MSB and LSB give visually distinct LE sequences. - F-C: add a "Proxy unreachable definition" subsection enumerating transport-level (TCP refuse/timeout, TLS handshake, HTTP/2 GOAWAY, reqwest::Error connect/timeout/request before any status) and HTTP-infrastructure classes (502, 504, 503-without-Cosmos- substatus). Explicitly exclude responses carrying a Cosmos sub-status. Defer to TRANSPORT_PIPELINE_SPEC for broader classification. Cross-reference from the Retry Decision Table. Also add a "Java parity" subsection to Phase 4 documenting that ThinClientStoreModel extends RxGatewayStoreModel, that none of the Java retry policies have thin-client-specific code, and that the Rust failure-fallback counter is more thin-client-aware than Java's. Flag a Java behavioral nuance worth NOT replicating: Java marks the gateway endpoint (not the thin-client endpoint) unavailable on a thin-client 503. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

implementation changes

6143a84

github-actions Bot added the Cosmos The azure_cosmos crate label Apr 3, 2026

github-project-automation Bot added this to CosmosDB Go/Rust Crew Apr 3, 2026

github-project-automation Bot moved this to Todo in CosmosDB Go/Rust Crew Apr 3, 2026

changelog

14ddfcb

simorenoh added 2 commits April 3, 2026 16:41

some updates

d746cbd

debugs

e18213c

simorenoh marked this pull request as ready for review April 6, 2026 16:37

simorenoh requested a review from a team as a code owner April 6, 2026 16:37

Copilot AI review requested due to automatic review settings April 6, 2026 16:37

Copilot started reviewing on behalf of simorenoh April 6, 2026 16:37 View session

Copilot AI reviewed Apr 6, 2026

View reviewed changes

tvaron3 reviewed Apr 6, 2026

View reviewed changes

Comment thread sdk/cosmos/azure_data_cosmos_driver/src/models/effective_partition_key.rs Outdated

simorenoh added 2 commits April 7, 2026 12:40

update with hpk logic

c1bbec0

doc updates from review

89826c8

simorenoh linked an issue Apr 7, 2026 that may be closed by this pull request

Cosmos: Add HPK support #4016

Closed

length aware epk comparison/ nomalization

86ac4b4

analogrelay requested changes Apr 7, 2026

View reviewed changes

github-project-automation Bot moved this from Todo to Changes Requested in CosmosDB Go/Rust Crew Apr 7, 2026

simorenoh added 2 commits April 7, 2026 17:59

address comments

79d8f85

add more test guidelines

7fb84da

analogrelay approved these changes Apr 7, 2026

View reviewed changes

FabianMeiswinkel reviewed Apr 8, 2026

View reviewed changes

Comment thread sdk/cosmos/azure_data_cosmos_driver/src/driver/cache/container_routing_map.rs

FabianMeiswinkel approved these changes Apr 8, 2026

View reviewed changes

tvaron3 approved these changes Apr 9, 2026

View reviewed changes

simorenoh added 4 commits April 9, 2026 11:42

add zero padding test

6f4b1c0

Merge branch 'release/azure_data_cosmos-previews' into cosmos-hpk

cedd41a

cspell post merging in branch

410e92c

Merge branch 'release/azure_data_cosmos-previews' into cosmos-hpk

a6750ca

FabianMeiswinkel approved these changes Apr 9, 2026

View reviewed changes

tvaron3 approved these changes Apr 9, 2026

View reviewed changes

simorenoh merged commit 25233c9 into release/azure_data_cosmos-previews Apr 9, 2026
13 checks passed

simorenoh deleted the cosmos-hpk branch April 9, 2026 19:20

github-project-automation Bot moved this from Changes Requested to Done in CosmosDB Go/Rust Crew Apr 9, 2026

This was referenced Apr 10, 2026

Cosmos Driver: Wire PartitionKeyRangeCache into CosmosDriver #4137

Closed

[Cosmos] Feed Range API #4149

Merged

Conversation

simorenoh commented Apr 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What changed

MultiHash EPK computation (effective_partition_key.rs)

Prefix EPK range computation (effective_partition_key.rs)

Prefix routing in PK range cache (partition_key_range_cache.rs)

Tests

Follow-up: FeedRange API (PR #3987)

Follow-up: Query tests and thorough testing

Uh oh!

github-actions Bot commented Apr 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

API Change Check

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

tvaron3 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

analogrelay left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

FabianMeiswinkel left a comment

Choose a reason for hiding this comment

Uh oh!

FabianMeiswinkel left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

simorenoh commented Apr 3, 2026 •

edited

Loading

MultiHash EPK computation (`effective_partition_key.rs`)

Prefix EPK range computation (`effective_partition_key.rs`)

Prefix routing in PK range cache (`partition_key_range_cache.rs`)

github-actions Bot commented Apr 3, 2026 •

edited

Loading