Fix HPK prefix and cross-partition query bugs

[simorenoh]

Summary

Fixes two hierarchical-partition-key (HPK / MultiHash) query bugs that were deferred out of #4705:

• Fixes Cosmos: HPK prefix queries are not filtered (return all items) #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 Cosmos: cross-partition queries over HPK containers fail with 400 #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.

[Copilot]

Pull request overview

This PR fixes two hierarchical-partition-key (HPK / MultiHash) query bugs that were deferred out of #4705. Both share a single root cause: when the driver scopes a data-plane request to an effective-partition-key (EPK) range (emitting x-ms-start-epk/x-ms-end-epk), it was sending the point key-type x-ms-read-key-type: EffectivePartitionKey, which the gateway rejects with 400. The fix emits EffectivePartitionKeyRange and adds prefix→EPK-range clipping so a partial HPK key (FeedScope::partition prefix) fans out over — and is EPK-filtered within — only the physical partitions its completions can occupy.

Changes:

• Emit x-ms-read-key-type: EffectivePartitionKeyRange for EPK range-scoped requests (fixes both #4680 filtered-prefix queries and #4681 cross-partition 400s).
• Add prefix_epk_range / clip_to_prefix in the planner to scope prefix fan-out to [prefix_epk, prefix_epk + "FF"), applied in both fresh and resume leaf builders.
• Add CI-runnable SDK-level HPK tests plus driver unit tests, document prefix support on FeedScope::partition, and add CHANGELOG entries.

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
azure_data_cosmos_driver/src/driver/pipeline/operation_pipeline.rs	Sends EffectivePartitionKeyRange key type for EPK range-scoped requests; updates the corresponding unit test.
azure_data_cosmos_driver/src/driver/dataflow/planner.rs	Adds prefix EPK-range computation and clipping into both fan-out leaf builders, with three new unit tests.
azure_data_cosmos/src/feed/query.rs	Documents that FeedScope::partition now supports HPK prefixes as filtered cross-partition queries.
azure_data_cosmos/tests/in_memory_emulator_tests/hpk.rs	New SDK-level tests covering prefix filtering, no-match, cross-partition full-container, and full-key cases.
azure_data_cosmos/tests/in_memory_emulator_tests/mod.rs	Registers the new hpk test module.
azure_data_cosmos/CHANGELOG.md	Adds a Bugs Fixed entry for #4680/#4681.
azure_data_cosmos_driver/CHANGELOG.md	Adds a Bugs Fixed entry describing the x-ms-read-key-type fix.

[FabianMeiswinkel]

LGTM except for few NITs

[ananth7592]

Hi @simorenoh — nice work here. The EffectivePartitionKeyRange header correction in this PR is exactly the fix I'd landed independently in #4726, and your version carries the HPK context plus .NET validation, so I'm closing #4726 in favor of this PR and @tvaron3's #4319.

One artifact from #4726 that I don't think is covered elsewhere: a standard-gateway (non-GW2.0, single-hash /pk) end-to-end regression test. It seeds one doc per partition key, computes each doc's EPK locally, and asserts that FeedScope::range(interior_window) returns exactly the windowed subset — with wide, tight, and empty-gap windows. #4319's query test is GW2.0-specific and your new tests are in-memory HPK, so this exercises a different path: it fails on pre-fix main with SCOPE BUG (full-scan fallback) and passes once both the header fix and the planner scope-clip are in.

It's #[ignore]-gated on test_category="emulator" and already adapted to current main's API (to_hex, byte-hex From, EffectivePartitionKey: Ord). Would it be cleaner to fold it into this PR, or would you prefer I bring it as a small standalone follow-up? Whatever's least noise for you. Dropping the full file here for convenience:

// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

//! End-to-end regression test for the cross-partition query **scope bug**: the
//! planner used to ignore the caller's `FeedScope::range(..)` feed range.
//!
//! Before the fix, `plan_fresh` built its request ranges purely from the query
//! plan (`query_plan.query_ranges`) and never intersected them with the
//! operation's scope feed range (`operation.target()`). For a plain `SELECT *`
//! the query plan reports the whole container, so a `FeedScope::range([X, Y))`
//! window was silently dropped and the query did a **full scan** — returning
//! documents outside the requested window.
//!
//! This test seeds one document per partition key, computes each document's
//! effective partition key locally, and issues cross-partition queries scoped
//! to interior EPK windows. It asserts the result set is exactly the windowed
//! subset. Without the planner fix these assertions fail with
//! `SCOPE BUG` (too many documents returned); with the fix — plus the
//! `x-ms-read-key-type: EffectivePartitionKeyRange` header correction that lets
//! the gateway accept the now-emitted interior EPK window — they pass.
//!
//! Run against the emulator:
//!   RUSTFLAGS='--cfg test_category="emulator"' \
//!     cargo test -p azure_data_cosmos --test emulator \
//!     --features "key_auth fault_injection" feed_range_scope

use super::framework;

use std::borrow::Cow;
use std::collections::BTreeSet;
use std::error::Error;

use azure_data_cosmos::clients::ContainerClient;
use azure_data_cosmos::feed::{FeedRange, FeedScope};
use azure_data_cosmos::models::{
    ContainerProperties, EffectivePartitionKey, PartitionKeyDefinition,
};
use azure_data_cosmos::{PartitionKey, Query};
use futures::TryStreamExt;
use serde::{Deserialize, Serialize};

use framework::{TestClient, TestOptions};

const DOC_COUNT: usize = 40;

#[derive(Clone, Debug, Serialize, Deserialize)]
struct Doc {
    id: String,
    pk: String,
}

fn pk_definition() -> PartitionKeyDefinition {
    PartitionKeyDefinition::new(vec![Cow::Borrowed("/pk")])
}

/// Runs a feed-range-scoped `SELECT *` and returns the set of returned ids.
async fn drain_ids(
    container: &ContainerClient,
    scope: FeedScope,
) -> Result<BTreeSet<String>, Box<dyn Error>> {
    let mut iter = container
        .query_items::<Doc>(Query::from("SELECT * FROM c"), scope, None)
        .await?;
    let mut got = BTreeSet::new();
    while let Some(doc) = iter.try_next().await? {
        got.insert(doc.id);
    }
    Ok(got)
}

#[tokio::test]
#[cfg_attr(
    not(any(test_category = "emulator", test_category = "emulator_vnext")),
    ignore = "requires test_category 'emulator' or 'emulator_vnext'"
)]
pub async fn feed_range_scope_restricts_cross_partition_query() -> Result<(), Box<dyn Error>> {
    TestClient::run_with_unique_db(
        async |run_context, db_client| {
            let properties = ContainerProperties::new("ScopeBugContainer", pk_definition());
            let container = run_context
                .create_container(db_client, properties, None)
                .await?;

            // Seed DOC_COUNT documents, each on its own partition key.
            let mut docs: Vec<Doc> = Vec::with_capacity(DOC_COUNT);
            for i in 0..DOC_COUNT {
                let id = format!("doc-{i:03}");
                let doc = Doc {
                    id: id.clone(),
                    pk: id.clone(),
                };
                container
                    .create_item(PartitionKey::from(doc.pk.clone()), &id, doc.clone(), None)
                    .await?;
                docs.push(doc);
            }

            // Compute each document's effective partition key locally (same hash
            // the service uses), then sort ascending by EPK.
            let pk_def = pk_definition();
            let mut points: Vec<(String, FeedRange)> = docs
                .iter()
                .map(|d| {
                    let fr = FeedRange::for_partition(PartitionKey::from(d.pk.clone()), &pk_def);
                    (d.id.clone(), fr)
                })
                .collect();
            points.sort_by(|a, b| a.1.min_inclusive().cmp(b.1.min_inclusive()));
            let k = points.len();

            // Control: a full-container scan must return every seeded document.
            let all: BTreeSet<String> = points.iter().map(|(id, _)| id.clone()).collect();
            let control = drain_ids(&container, FeedScope::full_container()).await?;
            assert_eq!(
                control, all,
                "control full-container scan should return all {DOC_COUNT} docs"
            );

            // Test B: a WIDE interior window [X1, X_{k-1}) must exclude the
            // globally smallest and largest EPKs => expect k-2 documents.
            let window_b = FeedRange::new(
                points[1].1.min_inclusive().clone(),
                points[k - 1].1.min_inclusive().clone(),
            )?;
            let expected_b: BTreeSet<String> =
                points[1..k - 1].iter().map(|(id, _)| id.clone()).collect();
            let got_b = drain_ids(&container, FeedScope::range(window_b)).await?;
            assert_eq!(
                got_b,
                expected_b,
                "SCOPE BUG (wide window): FeedScope::range was ignored and the \
                 query fell back to a full scan. got {} docs, expected {}. \
                 unexpected(outside window)={:?}",
                got_b.len(),
                expected_b.len(),
                got_b.difference(&expected_b).collect::<Vec<_>>()
            );

            // Test C: the TIGHTEST interior window [X_mid, X_{mid+1}) must
            // return exactly one document.
            let mid = k / 2;
            let window_c = FeedRange::new(
                points[mid].1.min_inclusive().clone(),
                points[mid + 1].1.min_inclusive().clone(),
            )?;
            let expected_c: BTreeSet<String> = std::iter::once(points[mid].0.clone()).collect();
            let got_c = drain_ids(&container, FeedScope::range(window_c)).await?;
            assert_eq!(
                got_c,
                expected_c,
                "SCOPE BUG (tight window): FeedScope::range was ignored. \
                 got {} docs, expected 1. unexpected(outside window)={:?}",
                got_c.len(),
                got_c.difference(&expected_c).collect::<Vec<_>>()
            );

            // Test D: a window that lies ENTIRELY in the gap between two adjacent
            // EPKs must return zero documents. Appending a whole byte (`80`) to
            // `X_mid`'s hex keeps `X_mid` as a byte prefix, so the result is
            // strictly greater than `X_mid` and — because adjacent single-hash
            // EPKs differ within their leading bytes — strictly less than
            // `X_{mid+1}`. So `[X_mid || 0x80, X_{mid+1})` contains no doc's EPK.
            // (A single hex nibble would be dropped by the byte-wise hex parser,
            // so a full byte is required.) This guards against an off-by-one
            // where the lower bound is treated as inclusive of `X_mid`.
            let gap_start =
                EffectivePartitionKey::from(format!("{}80", points[mid].1.min_inclusive().to_hex()));
            let window_d = FeedRange::new(gap_start, points[mid + 1].1.min_inclusive().clone())?;
            let got_d = drain_ids(&container, FeedScope::range(window_d)).await?;
            assert!(
                got_d.is_empty(),
                "SCOPE BUG (empty gap window): a window between adjacent EPKs \
                 returned {} docs, expected 0. unexpected={:?}",
                got_d.len(),
                got_d.iter().collect::<Vec<_>>()
            );

            Ok(())
        },
        Some(TestOptions::for_emulator()),
    )
    .await
}

Happy to discuss any of this.

[simorenoh]

Thanks @ananth7592 — appreciate you closing #4726 in favor of this, and the regression test is a great addition. It exercises a path neither the in-memory HPK tests nor the GW2.0 query tests cover, so I've folded it straight in rather than taking it as a follow-up (least noise for everyone).

It now lives in sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_feed_ranges.rs as feed_range_scope_restricts_cross_partition_query, with the standard #[cfg_attr(not(any(test_category = "emulator", test_category = "emulator_vnext")), ignore)] gate. One small tweak: drain_ids now drains page-by-page via .into_pages() instead of item-level try_next, which keeps the query future under CI's clippy::large-futures threshold. Wording and assertions are otherwise yours.