Cosmos: Regression tests for feed-range-scoped query fan-out by simorenoh · Pull Request #4705 · Azure/azure-sdk-for-rust

simorenoh · 2026-07-07T17:57:28Z

Summary

The bug this PR originally fixed — cross-partition queries scoped with FeedScope::range fanning out to partitions outside the requested range — was independently fixed on main by the Gateway 2.0 work (#4319), which landed first. This PR now only adds the regression tests (planner unit tests + an emulator test) for that customer-reported scenario, and backfills the missing changelog entry.

Original problem

Cross-partition queries scoped with FeedScope::range fanned out to partitions outside the requested range. ContainerClient::query_items correctly stores the requested feed range in the CosmosOperation target, but for explicit EPK ranges the query is classified as non-trivial, so the driver fetches a query plan and builds request leaves from the plan's ranges. The query plan's ranges are derived from the query text alone (e.g. [00, FF) for a plain SELECT *) and were never intersected with operation.target().

Example:

Requested scope: [00, 80)
Physical topology: [00, 80), [80, FF)
Query plan: [00, FF)

Expected one request leaf for [00, 80); the bug produced leaves for both [00, 80) and [80, FF) — returning documents outside the requested range and causing repeated scans / RU usage when callers enumerate read_feed_ranges() and query each range.

What changed on `main`

#4319 (Gateway 2.0 thin client) reworked the planner so both fan-out leaf builders — plan_fresh and plan_resume_from_saved_snapshot — intersect each query-plan range with the operation's target, dropping ranges that fall entirely outside it. That is the same fix this PR originally proposed (a clip_to_target helper), so on merge the helper was removed in favor of main's implementation. main's version also correctly handles hierarchical prefix targets (represented as EPK ranges), which the earlier helper deliberately skipped.

What this PR contributes

Planner unit tests (azure_data_cosmos_driver):
- restricts_fanout_to_explicit_target_range — the exact reported scenario (one leaf, not two), against a 2-partition physical topology
- drops_query_ranges_outside_target — disjoint plan range dropped
- clips_query_range_to_partial_target_overlap — partial overlap clipped to target bounds
- drops_query_range_touching_target_boundary — boundary-touch range dropped
- resume_restricts_fanout_to_target_range — resume path honors the target
- rejects_target_disjoint_from_query_ranges — disjoint target yields the empty-ranges error
- A range-aware PhysicalTopologyProvider test double, so an unclipped planner would visibly resolve the out-of-scope partition
Emulator regression test (azure_data_cosmos): feed_range_scoped_query_honors_range provisions a multi-partition container, enumerates read_feed_ranges(), queries each range in isolation, and asserts the union equals the full set exactly once (with the bug, each range returns the whole container, so ids repeat once per range).
Changelog backfill: since this PR adds only tests, it carries no changelog entry of its own. The Bugs Fixed bullet in both crates is attributed to Add Gateway 2.0 (thin client) implementation to Cosmos driver #4319 (which delivered the fix but shipped without a changelog entry for this bug).

Validation

cargo test -p azure_data_cosmos_driver --lib — 2106 pass
cargo fmt / cargo clippy clean on both crates
azure_data_cosmos emulator test target compiles (the emulator test runs in CI)

Follow-up (tracked separately)

Hierarchical prefix FeedScope::partition over-scan is tracked in #4680 (prefix queries not filtered) and #4681 (HPK cross-partition 400); a cross-version continuation-resume edge is tracked in #4714.

Cross-partition queries scoped with `FeedScope::range` were fanning out to partitions outside the requested range. `ContainerClient::query_items` correctly stored the requested feed range in the operation target, but the fan-out planner built request leaves from the query plan's ranges (derived from the query text alone) without intersecting them with that target. With topology `[00,80),[80,FF)` and a query plan range of `[00,FF)`, a request scoped to `[00,80)` produced leaves for both partitions. This could return documents outside the requested feed range and caused extra scans / RU usage when callers enumerated `read_feed_ranges()` and queried each range. Add a `clip_to_target` helper that intersects each query-plan range with the operation's explicit-EPK target and skips ranges that fall entirely outside it. Apply it in both the fresh (`plan_fresh`) and resume (`plan_resume_from_saved_snapshot`) leaf builders. Targets that are absent, full, or logical-partition prefixes are passed through unchanged. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Copilot

Pull request overview

This PR fixes a correctness bug in the Cosmos driver's cross-partition query planner. When a query was scoped to an explicit EPK feed range via FeedScope::range, the driver classified it as non-trivial, fetched a query plan whose ranges are derived from the query text alone (e.g. [00, FF)), and built fan-out request leaves from those plan ranges — without intersecting them against the caller's requested target range. This caused the query to fan out to partitions outside the requested scope, returning out-of-range documents and incurring extra scans/RU.

The fix adds a clip_to_target helper that intersects each query-plan feed range with the operation's explicit-EPK target, dropping ranges entirely outside it. It is applied consistently in both fan-out leaf builders (plan_fresh and plan_resume_from_saved_snapshot). Absent, full(), and logical-partition/prefix targets pass through unchanged, so trivial single-partition queries, full scans, and prefix fan-out (a tracked follow-up) are unaffected.

Changes:

Added clip_to_target and wired it into both the fresh and resume fan-out paths in planner.rs.
Added four planner unit tests covering the reported scenario, disjoint drops, partial-overlap clipping, and the resume path.
Documented the fix in the public azure_data_cosmos CHANGELOG.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/planner.rs	Adds `clip_to_target`, applies it in `plan_fresh` and `plan_resume_from_saved_snapshot`, and adds unit tests for the new clipping behavior.
sdk/cosmos/azure_data_cosmos/CHANGELOG.md	Adds a "Bugs Fixed" entry describing the feed-range scoping fix.

The two headline target-scoping tests passed even with the clip reverted: `MockTopologyProvider` ignores its range argument and replays a fixed queue, so the feed range actually handed to `resolve_ranges` was never observed. Both tests were vacuous. Add a range-aware `PhysicalTopologyProvider` that resolves against a fixed physical layout, returning only partitions overlapping the requested range (mirroring real routing). Rewrite the fresh and resume target tests to use it with a two-partition topology so an unclipped planner would resolve and emit the out-of-scope partition and fail. Also add two edge-case tests: a query-plan range that only touches the target's exclusive upper bound (dropped), and a logical-partition (prefix) target that must not be clipped (pass-through guard). All five range-scoped tests now fail when the clip is neutralized. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

The fix lives in the driver crate (planner.rs) but was only recorded in the public azure_data_cosmos CHANGELOG. Add the matching Bugs Fixed entry to the driver's own release notes, consistent with #4655 which appears in both. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

The fresh fan-out path returns a hard error when a query's target shares no EPKs with any query-plan range (every range clips away, leaving zero request leaves). Add a test that pins this behavior and, via NoopTopologyProvider, asserts no out-of-scope partition is resolved. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

FabianMeiswinkel · 2026-07-08T13:10:39Z


    for query_range in &query_plan.query_ranges {
        let feed_range = query_range_to_feed_range(query_range)?;
+        // Restrict the query-plan range to the operation's requested scope,


Please also add an e2e test - simple SELECT * FROM c query with FeedScope to validate that it indeed honors the rnage (it will with your fix - but having an e2e regression test makes it much simpler to prevent regressions than unit tests which are easier to overlook in larger changes).

I also disagree with the decision to scope out HPK (i would have preferred fixing it in this PR as well) - but whatever - as long as it is also fixed before next release that is fine.

@FabianMeiswinkel Added the e2e test, I've been tracking the HPK scenarios separately through more tests here: #4682

And issues here: #4155, #4680, #4681

FabianMeiswinkel

LGTM except one missing e2e test and the fact that HPK is scoped out (but as long as it is done before next release that is acceptable)

Fabian asked for an emulator-level regression test alongside the unit tests. Add feed_range_scoped_query_honors_range: it provisions a multi-partition container, enumerates read_feed_ranges(), and queries each range in isolation (the reported customer scenario). It asserts the union across ranges equals the full set exactly once -- before the fan-out clip fix each range returned the whole container, so ids would appear once per range. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

The Build Analyze pipeline flagged 'neighbouring' and 'unclipped' from comments added in the feed-range fan-out tests. Add them to the Cosmos-scoped sdk/cosmos/.cspell.json rather than the repo-wide config. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

…range-query-investigation # Conflicts: # sdk/cosmos/.cspell.json # sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/planner.rs

The equivalent planner fix shipped independently with the Gateway 2.0 work (#4319), which landed on main first. Reframe both CHANGELOG bullets to credit #4319 for the code fix and describe #4705 as adding the regression tests (unit + emulator) for the reported feed-range scoping scenario. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

This PR now only adds regression tests, so it needs no changelog entry of its own. Re-point the Bugs Fixed bullet in both crates to #4319 (Tomas' Gateway 2.0 PR), which delivered the planner fix but shipped without a changelog entry for this bug. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

analogrelay

One test fix up, otherwise looks good

analogrelay · 2026-07-09T16:56:39Z

+            let mut union_ids = Vec::new();
+            for range in &ranges {
+                let ids = collect_ids_for_scope(&container_client, FeedScope::range(range.clone()))
+                    .await?;
+                union_ids.extend(ids);
+            }
+            union_ids.sort();
+
+            assert_eq!(
+                union_ids, all_ids,
+                "each feed range must return only its own items: the union across \
+                 ranges should equal the full set exactly once, with no over-scan \
+                 or duplication"
+            );
+
+            Ok(())


I'd prefer this test actually capture each range's results as a separate list and assert those (i.e. a Vec<Vec<String>> and then seen an assert_eq!(vec![vec!["id1", "id2", ...], vec!["idX", "idY", ...]], all_range_ids). If we returned 0 results from one range and all the results in a second range, that would be incorrect but would pass this test.

Good call - updated in 0d29d2e. The test now captures each range's result set into a Vec<Vec<String>> and asserts (1) every range is non-empty, (2) the ranges are pairwise disjoint (no duplicated ids), and (3) flattened they cover the full container exactly once. The empty-range-plus-full-range case you flagged now fails on the non-empty check.

## Summary Fixes two hierarchical-partition-key (HPK / MultiHash) query bugs that were deferred out of #4705: - **Fixes #4680** — a `FeedScope::partition` scope with only a *prefix* of the key hierarchy (e.g. `("USA", "CA")` on a `/country/state/city` container) returned **every** item in the physical partition instead of filtering to the prefix. - **Fixes #4681** — cross-partition queries over an HPK container failed with `400 Bad Request` ("One of the input values is invalid") and an empty `PartitionKeyRangeId`. ## Root cause (unified) Both bugs share a single defect. When the driver scopes a data-plane request to an effective-partition-key (EPK) **range**, it emits `x-ms-start-epk` / `x-ms-end-epk` alongside `x-ms-read-key-type`. The value was hardcoded to the **point** key-type `EffectivePartitionKey`, which the gateway rejects for range-scoped requests. It must be the **range** key-type `EffectivePartitionKeyRange`. This was validated end-to-end against a live Cosmos account (raw HTTP + the Python `azure-cosmos` SDK) and cross-checked against the .NET SDK (`azure-cosmos-dotnet-v3`), whose `RequestInvokerHandler` and Gateway 2.0 thin-client path emit the identical header set. ## Changes - **`azure_data_cosmos_driver/src/driver/pipeline/operation_pipeline.rs`** — emit `x-ms-read-key-type: EffectivePartitionKeyRange` for EPK range-scoped requests; updated the corresponding unit test. - **`azure_data_cosmos_driver/src/driver/dataflow/planner.rs`** — scope a prefix HPK fan-out to `[prefix_epk, prefix_epk + "FF")` (`prefix_epk_range` / `clip_to_prefix`) so a prefix query only fans out over — and is EPK-filtered within — the physical partitions its completions can live in, instead of scanning whole partitions. - **`azure_data_cosmos/src/feed/query.rs`** — documented prefix support on `FeedScope::partition`. - **`azure_data_cosmos/tests/in_memory_emulator_tests/hpk.rs`** (new) — CI-runnable SDK-level tests covering both bugs. - **CHANGELOG** entries in both `azure_data_cosmos` and `azure_data_cosmos_driver`. ## Validation - Driver unit tests pass; new in-memory HPK tests pass; `cargo fmt` + `cargo clippy` clean. - Behavior verified against the .NET reference implementation (header value, prefix `[epk, epk+"FF")` range with the `0x3F`-masked-nibble `FF` sentinel, exact-match pkrangeid collapse, multi-partition prefix fan-out, and change-feed sharing the same path all match). ## Note The canonical `#[ignore]`d emulator tests in `cosmos_hpk.rs` land via the separate #4155 coverage branch; flipping those ignores is deferred to that branch's merge to avoid duplication. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

…range-query-investigation # Conflicts: # sdk/cosmos/azure_data_cosmos/CHANGELOG.md # sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md

Address review feedback: capture each feed range's result set separately instead of only their union. A union-only check would pass even if one range returned every item and another returned nothing. The test now asserts each range is non-empty, the ranges are pairwise disjoint, and together they cover the full container exactly once. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

analogrelay · 2026-07-13T17:42:31Z

+            // Every range must return a non-empty slice. Each physical partition
+            // owns a share of the 10 logical partitions, so a correctly-scoped
+            // query can never come back empty; an empty range would mean the
+            // scope was dropped and the items were served by some other range.
+            for (i, ids) in per_range_ids.iter().enumerate() {
+                assert!(
+                    !ids.is_empty(),
+                    "feed range {i} returned no items; a scoped query must return \
+                     that range's own slice, not defer to another range"
+                );
+            }
+
+            // The ranges must be disjoint: before the fan-out clip fix every
+            // range fanned out across the whole container, so ids were
+            // duplicated across ranges. Flatten and check for duplicates.
+            let mut flattened: Vec<String> = per_range_ids.iter().flatten().cloned().collect();
+            flattened.sort();
+            let mut deduped = flattened.clone();
+            deduped.dedup();
+            assert_eq!(
+                flattened, deduped,
+                "feed ranges must be disjoint: no id may be returned by more than \
+                 one range"
+            );
+
+            // And together the ranges must cover exactly the full container —
+            // no over-scan and nothing missing.
+            assert_eq!(
+                flattened, all_ids,
+                "the union across ranges must equal the full container exactly once"
+            );


Given that we're generating a fixed list of IDs, I feel like it would be better to just have a single assert_eq!(vec![vec![...], ...], per_range_ids) that makes it easy to look at the IDs and say "Yep, that looks correct". This is all correct but it's asserting three separate properties (non-empty, disjoint results, fully spanning) which could also be trivially validated by looking at the actual ranges, and would be easier to read.

Switched to the hardcoded literal you asked for. To de-risk the coupling concern I raised above, I ran the grouping against a real Cosmos DB account and the local emulator: both split into the same two physical ranges at EPK 0x1FFF...FF and produced the identical per-range id grouping. The EPK hash and the split boundary are algorithm-driven, not environment-specific, so the literal is stable. It's now a single assert_eq!(expected, per_range_ids) you can eyeball. Thanks!

…range-query-investigation # Conflicts: # sdk/cosmos/azure_data_cosmos/CHANGELOG.md # sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md

Replace the three separate property checks (non-empty, disjoint, spanning) with a single assert_eq! comparing the actual per-range query results to an independently computed expected grouping. The expected grouping maps each seeded item to its owning physical range via the SDK's partition-key -> EPK range math, so the assertion covers all three properties at once and stays correct regardless of how the emulator splits the container. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Replace the computed grouping oracle with a single assert_eq! against a hardcoded per-range id grouping, per review feedback. The values were captured from a real Cosmos DB account and verified identical on the emulator (both split at EPK 0x1FFF...FF), confirming the grouping is algorithm-driven and environment-independent. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

analogrelay · 2026-07-22T20:33:59Z

@NaluTripician can you rebase, resolve conflicts, and then merge? (I should be able to give a quick approval after the rebase)

Copilot AI review requested due to automatic review settings July 7, 2026 17:57

simorenoh requested a review from a team as a code owner July 7, 2026 17:57

Add PR link to changelog entry

ba795b8

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

github-actions Bot added the Cosmos The azure_cosmos crate label Jul 7, 2026

github-project-automation Bot added this to CosmosDB Rust SDK and Driver Jul 7, 2026

github-project-automation Bot moved this to Todo in CosmosDB Rust SDK and Driver Jul 7, 2026

Copilot started reviewing on behalf of simorenoh July 7, 2026 17:57 View session

Copilot AI reviewed Jul 7, 2026

View reviewed changes

Comment thread sdk/cosmos/azure_data_cosmos/CHANGELOG.md Outdated

simorenoh changed the title ~~Restrict cross-partition query fan-out to feed-range scope~~ Cosmos: Restrict cross-partition query fan-out to feed-range scope Jul 7, 2026

simorenoh and others added 3 commits July 7, 2026 16:05

simorenoh mentioned this pull request Jul 8, 2026

Cosmos: resuming a pre-fix continuation token for a feed-range-scoped cross-partition query can hard-error after upgrade #4714

Closed

Merge branch 'main' into simorenoh/cosmos-feedrange-query-investigation

b97b0cf

FabianMeiswinkel reviewed Jul 8, 2026

View reviewed changes

FabianMeiswinkel approved these changes Jul 8, 2026

View reviewed changes

simorenoh and others added 3 commits July 8, 2026 10:32

Merge branch 'main' into simorenoh/cosmos-feedrange-query-investigation

9f3e986

simorenoh mentioned this pull request Jul 9, 2026

Fix HPK prefix and cross-partition query bugs #4729

Merged

simorenoh and others added 2 commits July 9, 2026 13:13

Merge remote-tracking branch 'origin/main' into simorenoh/cosmos-feed…

6a29dde

…range-query-investigation # Conflicts: # sdk/cosmos/.cspell.json # sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/planner.rs

simorenoh changed the title ~~Cosmos: Restrict cross-partition query fan-out to feed-range scope~~ Cosmos: Regression tests for feed-range-scoped query fan-out Jul 9, 2026

analogrelay requested changes Jul 9, 2026

View reviewed changes

github-project-automation Bot moved this from Todo to Changes Requested in CosmosDB Rust SDK and Driver Jul 9, 2026

simorenoh and others added 2 commits July 9, 2026 18:34

Merge remote-tracking branch 'origin/main' into simorenoh/cosmos-feed…

9d90f1a

…range-query-investigation # Conflicts: # sdk/cosmos/azure_data_cosmos/CHANGELOG.md # sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md

simorenoh closed this Jul 10, 2026

github-project-automation Bot moved this from Changes Requested to Done in CosmosDB Rust SDK and Driver Jul 10, 2026

simorenoh reopened this Jul 10, 2026

simorenoh requested a review from analogrelay July 10, 2026 15:45

Merge branch 'main' into simorenoh/cosmos-feedrange-query-investigation

0d3c8dc

analogrelay requested changes Jul 13, 2026

View reviewed changes

github-project-automation Bot moved this from Done to Changes Requested in CosmosDB Rust SDK and Driver Jul 13, 2026

simorenoh and others added 3 commits July 13, 2026 19:51

Merge remote-tracking branch 'origin/main' into simorenoh/cosmos-feed…

194b406

…range-query-investigation # Conflicts: # sdk/cosmos/azure_data_cosmos/CHANGELOG.md # sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md

simorenoh linked an issue Jul 14, 2026 that may be closed by this pull request

[Cosmos] Testing: Additional sanity and coverage testing #3666

Open

analogrelay approved these changes Jul 14, 2026

View reviewed changes

analogrelay moved this from Needs Attention to Ready To Merge in CosmosDB Rust SDK and Driver Jul 21, 2026

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Cosmos: Regression tests for feed-range-scoped query fan-out#4705

Cosmos: Regression tests for feed-range-scoped query fan-out#4705
simorenoh wants to merge 18 commits into
mainfrom
simorenoh/cosmos-feedrange-query-investigation

simorenoh commented Jul 7, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

FabianMeiswinkel Jul 8, 2026

Uh oh!

simorenoh Jul 8, 2026

Uh oh!

FabianMeiswinkel left a comment

Uh oh!

analogrelay left a comment

Uh oh!

analogrelay Jul 9, 2026

Uh oh!

simorenoh Jul 9, 2026

Uh oh!

analogrelay Jul 13, 2026

Uh oh!

simorenoh Jul 14, 2026

Uh oh!

analogrelay commented Jul 22, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Conversation

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

Uh oh!

Summary

Original problem

What changed on main

What this PR contributes

Validation

Follow-up (tracked separately)

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

FabianMeiswinkel Jul 8, 2026

Choose a reason for hiding this comment

Uh oh!

simorenoh Jul 8, 2026

Choose a reason for hiding this comment

Uh oh!

FabianMeiswinkel left a comment

Choose a reason for hiding this comment

Uh oh!

analogrelay left a comment

Choose a reason for hiding this comment

Uh oh!

analogrelay Jul 9, 2026

Choose a reason for hiding this comment

Uh oh!

simorenoh Jul 9, 2026

Choose a reason for hiding this comment

Uh oh!

analogrelay Jul 13, 2026

Choose a reason for hiding this comment

Uh oh!

simorenoh Jul 14, 2026

Choose a reason for hiding this comment

Uh oh!

analogrelay commented Jul 22, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

simorenoh commented Jul 7, 2026 •

edited

Loading

What changed on `main`