Add cross-partition DISTINCT queries by tvaron3 · Pull Request #5026 · Azure/azure-sdk-for-rust

Tomas Varon (tvaron3) · 2026-08-07T19:51:25Z

Adds ordered and unordered cross-partition DISTINCT execution to the Cosmos DB Rust driver. Before this, planner::validate_query_info and validate_query_plan_for_streaming_order_by rejected any plan with distinctType != None, so every SELECT DISTINCT cross-partition query failed as an unsupported feature.

Fixes #4754.

Design

Distinct is a composition stage above the fan-out root — Distinct -> SequentialDrain when unordered, Distinct -> StreamingOrderedMerge when ordered. It keys on a structural, type-aware 128-bit hash of the whole projected row rather than the ORDER BY items, so one node serves both modes and only the retained state differs. The hash module is standalone so GROUP BY can reuse it.

	Ordered	Unordered
Retained state	last hash, O(1)	full `HashSet<Hash128>`
Continuation	resumable — 16 bytes	refused at mint time

Ordered DISTINCT deduplicates by adjacency and resumes from the 16 bytes PipelineNodeState::Distinct persists: a value the stage has moved past can never reappear. This complements rather than duplicates the merge's own resume trim, which is positional (_rid + skipCount) — last_hash catches a different _rid carrying the same projected value, i.e. two documents that are one DISTINCT row but two ORDER BY rows.

Unordered DISTINCT can't be resumed: the set is the state, serializing it would produce an unbounded token, and truncating it would silently re-emit duplicates. snapshot_state fails with the new CLIENT_DISTINCT_CONTINUATION_UNSUPPORTED so to_continuation_token errors while the caller still holds a live plan and can rewrite with a matching ORDER BY. In-process paging is fully supported, and once drained there's no state left to lose, so it snapshots as Drained like any other finished node.

The `Ordered` boundary is measured, not assumed

Against a live account with production's SUPPORTED_QUERY_FEATURES:

Query	`distinctType`
`SELECT DISTINCT VALUE c.name … ORDER BY c.name`	`Ordered` (with a 172-char `rewrittenQuery`)
`SELECT DISTINCT c.name … ORDER BY c.name` (list form)	`Unordered`
`SELECT DISTINCT VALUE c.name … ORDER BY c.name, c.other`	`Unordered`
`SELECT DISTINCT c.name, c.city … ORDER BY c.name`	`Unordered`
`SELECT DISTINCT VALUE c.name FROM c`	`Unordered`

Only the VALUE form whose ORDER BY is exactly the projected path stays Ordered — notably not a multi-column ORDER BY that leads with it, where adjacency would in fact still hold. Advertising Distinct is mandatory: without it the service rejects the query with 400/1004.

Adds ordered and unordered cross-partition `DISTINCT` execution to the Cosmos DB Rust driver. Before this, `planner::validate_query_info` and `validate_query_plan_for_streaming_order_by` rejected any plan with `distinctType != None`, so every `SELECT DISTINCT` cross-partition query failed as an unsupported feature. `Distinct` is a composition stage above the fan-out root — `Distinct -> SequentialDrain` when unordered, `Distinct -> StreamingOrderedMerge` when ordered. It keys on a structural, type-aware 128-bit hash of the whole projected row rather than the ORDER BY items, so one node serves both modes and only the retained state differs. The hash module is standalone so GROUP BY can reuse it. Ordered DISTINCT deduplicates by adjacency, keeps a single hash, and resumes from the 16 bytes `PipelineNodeState::Distinct` persists — a value the stage has moved past can never reappear. Unordered DISTINCT retains every hash seen and is not resumable: the set is the state, so `snapshot_state` fails with the new `CLIENT_DISTINCT_CONTINUATION_UNSUPPORTED` at token-mint time, while the caller still holds a live plan and can rewrite with a matching ORDER BY. Once drained there is no state left to lose, so it snapshots as `Drained` like any other finished node. The `Ordered` boundary is measured, not assumed. Against a live account with production's `SUPPORTED_QUERY_FEATURES`, the service reports `Ordered` only for `SELECT DISTINCT VALUE <path> ... ORDER BY <same path>`; every other shape is `Unordered`, including a multi-column ORDER BY that leads with the projected path. Advertising `Distinct` is mandatory — without it the service rejects the query with 400/1004. Two pre-existing bugs surfaced while validating and are fixed here: - The local plan generator collapsed every constant DISTINCT to `None`. The service only does that without a FROM clause; with one the query yields a row per document, so deduplication is real work. The old behavior made the driver skip the stage and return one row per document. - `PartitionKey::as_headers` emitted `x-ms-documentdb-query-enablecrosspartition` for an empty key, contradicting `PartitionKey::EMPTY`'s own contract. That header is never emitted by any other path, and cross-partition queries route per pkrange instead. The empty-list guard is kept — without it the trailing-comma pop strips the opening bracket. The `distinctType` carve-out in the Gateway parity test is removed: its stated rule (downgrade whenever a `rewrittenQuery` is emitted) is disproved by the measurement above, and with the local generator matching the service the tolerance is unreachable — verified by instrumenting the branch and running the suite live. Tests are driven by a source-attributed scenario catalog reusing the streaming ORDER BY fixture structure, replayed across map, mock-pipeline, and in-memory-emulator layers, with a strict schema test. Because the emulator layer partly simulates behavior this change also implements, every scenario shape has a live counterpart. Four cases neither the .NET nor the Java SDK covers get explicit scenarios: -0.0 versus 0.0, object key-order commutativity, structural array ordering, and partition splits.

The entries linked the tracking issue; repo convention is to link the PR that introduced the change.

Tomas Varon (tvaron3) · 2026-08-07T19:53:06Z

/azp run rust - cosmos - weekly

azure-pipelines · 2026-08-07T19:53:13Z

Azure Pipelines will not run the associated pipelines, because the pull request was updated after the run command was issued. Review the pull request again and issue a new run command.

Tomas Varon (tvaron3) · 2026-08-07T19:54:29Z

/azp run rust - cosmos - weekly

azure-pipelines · 2026-08-07T19:54:35Z

Azure Pipelines will not run the associated pipelines, because the pull request was updated after the run command was issued. Review the pull request again and issue a new run command.

…artition-distinct # Conflicts: # sdk/cosmos/azure_data_cosmos/CHANGELOG.md # sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md # sdk/cosmos/azure_data_cosmos_driver/src/error/cosmos_status.rs

Tomas Varon (tvaron3) · 2026-08-07T20:52:07Z

/azp run rust - cosmos - weekly

azure-pipelines · 2026-08-07T20:52:30Z

Azure Pipelines: Successfully started running 1 pipeline(s).

…artition-distinct

Tomas Varon (tvaron3) · 2026-08-10T17:09:19Z

/azp run rust - cosmos - weekly

azure-pipelines · 2026-08-10T17:09:46Z

Azure Pipelines: Successfully started running 1 pipeline(s).

…s-partition-distinct # Conflicts: # sdk/cosmos/.cspell.json

Tomas Varon (tvaron3) · 2026-08-11T19:10:19Z

/azp run rust - cosmos - weekly

azure-pipelines · 2026-08-11T19:10:42Z

Azure Pipelines: Successfully started running 1 pipeline(s).

Ashley Stanton-Nurse (analogrelay)

Looks good to me, on the right track.

Tomas Varon (tvaron3) · 2026-08-11T22:05:06Z

/azp run rust - cosmos - weekly

azure-pipelines · 2026-08-11T22:05:30Z

Azure Pipelines: Successfully started running 1 pipeline(s).

`SplitRequired` replaces the node that emits it, so a `Distinct` node that forwarded one would be spliced out along with its deduplication map, letting already-suppressed values reappear. Refuse the split at the node instead, with a dedicated `CLIENT_DISTINCT_CANNOT_FORWARD_SPLIT` status. The wrapped fan-out node absorbs splits internally, so this is unreachable today; it becomes reachable as soon as another node is layered above `Distinct`, which is why it fails loudly rather than silently corrupting results. Also fixes two adjacent defects this exposed: - `SubStatusCode::name()` had arms for every neighbouring client code but omitted 20122/20123, so both `DISTINCT` statuses returned `None` and lost their searchable names in diagnostics. - The `Pipeline::next_page` comment still claimed the root is always a `Request`, `SequentialDrain`, or `DrainedLeaf`, which `wrap_root` made untrue once `Distinct` could be the root. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: eab410d0-1b95-4c38-9c03-2e6bb4a9a758

Resolves five conflicts plus two semantic integrations with the `OFFSET`/`LIMIT`/`TOP` support that landed upstream in Azure#4870. Conflicts: - `SUPPORTED_QUERY_FEATURES` — union both sides so the gateway may return plans carrying `Distinct` *and* `OffsetAndLimit`/`Top`. - Client sub-status codes — upstream took 20122 for `CLIENT_QUERY_REWRITE_BODY_INVALID`, so the three `DISTINCT` codes shift to 20123/20124/20125. - The emulator's rewritten-query builder now skips a leading `DISTINCT` *and* a leading `TOP <n>`; a query can carry both. - `response_with_charge` was added by both sides; kept one. - HPK doc comment — `DISTINCT` and `OFFSET`/`LIMIT`/`TOP` are both servable now; aggregates and `GROUP BY` remain rejected. Composition order (the load-bearing part): Auto-merge nested the stages as `Distinct(SkipTake(fan-out))`, which applies the row window to raw rows and then deduplicates, so `SELECT DISTINCT TOP 2 ...` over a stream whose first value repeats returned one row instead of two. SQL applies `DISTINCT` first, so `SkipTake` now wraps `Distinct`, and the resume peels `SkipTake` before `Distinct` to match the token nesting. Both are reachable because the merged feature list lets the service return a plan with both set. `Distinct` also had to learn the page shapes it can now be handed: upstream made `StreamingOrderedMerge` emit pre-split `Items`, so it normalizes `Items`/`Bytes`/`NoPayload` on input and emits `Items`, mirroring `SkipTake`. That retires `retain_documents`, `parse_document_page`, and `Pipeline::wrap_root`. The `unsupported_combination_top` catalog scenario becomes `combination_top_applies_after_distinct`: no longer an unsupported-feature error, it now pins the stage order and fails with the pre-merge nesting. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: eab410d0-1b95-4c38-9c03-2e6bb4a9a758

azure-pipelines · 2026-08-13T18:13:14Z

Azure Pipelines: Successfully started running 1 pipeline(s). 2 pipeline(s) were filtered out due to trigger conditions. There may be pipelines that require an authorized user to comment /azp run to run.

Copilot

Pull request overview

Adds cross-partition ordered and unordered DISTINCT execution to the Cosmos DB query pipeline.

Changes:

Adds structural hashing, deduplication, continuation handling, and split support.
Public API: adds three CosmosStatus constants for DISTINCT failures.
Expands planner, emulator, integration, catalog, and live-split coverage.

Reviewed changes

Copilot reviewed 32 out of 32 changed files in this pull request and generated 9 comments.

Show a summary per file

File	Description
`azure_data_cosmos/tests/split_tests/mod.rs`	Registers DISTINCT split tests.
`azure_data_cosmos/tests/split_tests/cosmos_query_distinct_split.rs`	Tests DISTINCT across live splits.
`azure_data_cosmos/tests/emulator_tests/cosmos_query.rs`	Adds SDK-level DISTINCT scenarios.
`azure_data_cosmos/tests/emulator_tests/cosmos_hpk.rs`	Verifies DISTINCT over HPK containers.
`azure_data_cosmos/CHANGELOG.md`	Documents SDK support.
`azure_data_cosmos_driver/tests/in_memory_emulator_tests/mod.rs`	Registers emulator tests.
`azure_data_cosmos_driver/tests/in_memory_emulator_tests/distinct.rs`	Exercises end-to-end DISTINCT behavior.
`azure_data_cosmos_driver/tests/gateway_query_plan_comparison.rs`	Expands Gateway plan parity checks.
`azure_data_cosmos_driver/tests/distinct_scenario_catalog.rs`	Validates the scenario catalog.
`query/plan/tests/query_plan_comparison.rs`	Tests local DISTINCT classification.
`query/plan/mod.rs`	Classifies ordered and constant DISTINCT plans.
`query/mod.rs`	Advertises DISTINCT support.
`query/eval/mod.rs`	Adds emulator-side partition deduplication.
`query_plan_native/native_dll_tests.rs`	Updates Copy-based test handling.
`models/partition_key.rs`	Stops emitting headers for empty keys.
`models/distributed_transaction.rs`	Explicitly rejects empty transaction keys.
`in_memory_emulator/operations.rs`	Rewrites ordered DISTINCT queries.
`error/cosmos_status.rs`	Adds DISTINCT status codes.
`dataflow/snapshot.rs`	Persists ordered DISTINCT state.
`dataflow/query_plan.rs`	Makes `DistinctType` copyable.
`dataflow/planner.rs`	Composes DISTINCT into query pipelines.
`dataflow/pipeline.rs`	Updates root-node documentation.
`dataflow/order_by.rs`	Reuses numeric classification.
`dataflow/mod.rs`	Registers DISTINCT modules.
`dataflow/mocks.rs`	Relocates charged-response helper.
`dataflow/distinct.rs`	Implements the DISTINCT stage.
`dataflow/distinct_hash.rs`	Implements structural JSON hashing.
`driver/cosmos_driver.rs`	Updates query-feature documentation.
`docs/FEED_OPERATIONS_REQS.md`	Documents DISTINCT semantics.
`azure_data_cosmos_driver/CHANGELOG.md`	Documents driver support.
`sdk/cosmos/.cspell.json`	Adds new terminology.

💡 Add a code-review agent skill for context-aware, tailored reviews. Learn more in the docs.

**Bound the suppressed diagnostics.** A long run of all-duplicate pages retained one `DiagnosticsContext` per page, growing independently of `max_request_diagnostics` and breaking the bounded-diagnostics contract. Fold each suppressed page into a single accumulator instead: `aggregate_sub_operations` re-bounds the concatenated records to the cap, so the artifact is O(cap) rather than O(pages). The newest page stays last, so the aggregate still inherits its operation-level fields and sums durations. **Stop asserting `DISTINCT TOP` / `OFFSET`-`LIMIT` are rejected.** The `SkipTake` composition that arrived with the upstream merge makes both servable, but the emulator test still expected a 400. It is gated behind `test_category`, so no local run reached it and CI would have failed. Those shapes are now asserted positively — the window must apply to deduplicated values — leaving only GROUP BY and aggregates in the rejected set. **Move `CLIENT_DISTINCT_CANNOT_FORWARD_SPLIT` to the invariant range.** It is an "unreachable by construction" guard, so 2012x (caller input) was wrong; its sibling `CLIENT_ROOT_NODE_CANNOT_REQUEST_SPLIT` is 20208. Uses 20217 rather than the next free 20216, which a stacked follow-up branch already claims. **Repair the split-test fallback.** It used a no-op hook and returned before `force_split_and_wait`, so the half it claimed still ran performed no split at all, and a catch-all `Err` let any continuation regression take that path and go green. Only the unsupported-continuation sub-status is tolerated now, and it fails loudly rather than skipping. **Correct the docs against measured behaviour.** `Ordered` applies only to `SELECT DISTINCT VALUE <path> ... ORDER BY <same path>`; "prefixes its projection" implied list and multi-column forms were resumable. Also fixes a stale sub-status number, and a claim that Java refuses an unordered continuation — it does not, it emits a token and silently drops the dedup state on resume. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: eab410d0-1b95-4c38-9c03-2e6bb4a9a758

The stage interaction that arrived with the upstream `OFFSET`/`LIMIT`/`TOP` merge had one execution test — an ordered `TOP` scenario. Everything else was plan-shape assertions, which pass regardless of how the pipeline composes the stages, so a wrong nesting produced silently short results rather than a failure. Adds four scenarios and one test: - Ordered `OFFSET`/`LIMIT`, which had no execution coverage at all (only `TOP` did). - Continuation round-trips for `DISTINCT` under both `TOP` and `OFFSET`/`LIMIT`. A token here nests `SkipTake { child: Distinct { .. } }`, and nothing previously replayed that shape: lose the window's remaining budget and the query over-returns, lose the dedup hash and it repeats the boundary value. Drains page-by-page and compares against a single drain. - A partition split mid-drain while the window is active, extending the existing split test, which used no window. - The same windowed shape across a real split in the live suite, where the fan-out is genuinely rebuilt rather than simulated. Verified by mutation: restoring the pre-merge `Distinct(SkipTake(..))` nesting fails three test functions, where it previously failed one. `TOP 3` returns two rows and `OFFSET 1 LIMIT 2` returns one, because the window then counts raw rows instead of deduplicated values. Not added: `DISTINCT` scenarios in `skip_take_scenarios.json`. Its runner reads `item["id"]` from each row, so a `VALUE` projection panics, and a list projection over unique ids makes deduplication a no-op — a test that cannot fail. The same paths are covered from the DISTINCT catalog with real value assertions. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: eab410d0-1b95-4c38-9c03-2e6bb4a9a758

Tomas Varon (tvaron3) · 2026-08-14T17:43:58Z

/azp run rust - cosmos - weekly

azure-pipelines · 2026-08-14T17:44:23Z

Azure Pipelines: Successfully started running 1 pipeline(s).

Tomas Varon (tvaron3) added 2 commits August 7, 2026 12:46

Point DISTINCT changelog entries at the PR

66259ed

The entries linked the tracking issue; repo convention is to link the PR that introduced the change.

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

Merge remote-tracking branch 'azure/main' into tvaron3-cosmos-cross-p…

7a34a66

…artition-distinct # Conflicts: # sdk/cosmos/azure_data_cosmos/CHANGELOG.md # sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md # sdk/cosmos/azure_data_cosmos_driver/src/error/cosmos_status.rs

Merge remote-tracking branch 'azure/main' into tvaron3-cosmos-cross-p…

15e1705

…artition-distinct

Merge remote-tracking branch 'upstream/main' into tvaron3-cosmos-cros…

3abf9b0

…s-partition-distinct # Conflicts: # sdk/cosmos/.cspell.json

Ashley Stanton-Nurse (analogrelay) reviewed Aug 11, 2026

View reviewed changes

Comment thread sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/distinct.rs Outdated

Tomas Varon (tvaron3) and others added 2 commits August 12, 2026 10:52

Tomas Varon (tvaron3) mentioned this pull request Aug 13, 2026

Add binary encoding for Cosmos query pages tvaron3/azure-sdk-for-rust#17

Closed

Tomas Varon (tvaron3) marked this pull request as ready for review August 13, 2026 18:13

Tomas Varon (tvaron3) requested a review from a team as a code owner August 13, 2026 18:13

Copilot AI balanced review requested due to automatic review settings August 13, 2026 18:13

Copilot started reviewing on behalf of Tomas Varon (tvaron3) August 13, 2026 18:14 View session

Copilot AI reviewed Aug 13, 2026

View reviewed changes

Tomas Varon (tvaron3) and others added 2 commits August 13, 2026 13:21

Tomas Varon (tvaron3) mentioned this pull request Aug 14, 2026

Add binary encoding for Cosmos query pages #5070

Draft

Uh oh!

Conversation

Tomas Varon (tvaron3) commented Aug 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Design

The Ordered boundary is measured, not assumed

Uh oh!

Tomas Varon (tvaron3) commented Aug 7, 2026

Uh oh!

azure-pipelines Bot commented Aug 7, 2026

Uh oh!

Tomas Varon (tvaron3) commented Aug 7, 2026

Uh oh!

azure-pipelines Bot commented Aug 7, 2026

Uh oh!

Tomas Varon (tvaron3) commented Aug 7, 2026

Uh oh!

azure-pipelines Bot commented Aug 7, 2026

Uh oh!

Tomas Varon (tvaron3) commented Aug 10, 2026

Uh oh!

azure-pipelines Bot commented Aug 10, 2026

Uh oh!

Tomas Varon (tvaron3) commented Aug 11, 2026

Uh oh!

azure-pipelines Bot commented Aug 11, 2026

Uh oh!

Ashley Stanton-Nurse (analogrelay) left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Tomas Varon (tvaron3) commented Aug 11, 2026

Uh oh!

azure-pipelines Bot commented Aug 11, 2026

Uh oh!

azure-pipelines Bot commented Aug 13, 2026

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!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Tomas Varon (tvaron3) commented Aug 14, 2026

Uh oh!

azure-pipelines Bot commented Aug 14, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Tomas Varon (tvaron3) commented Aug 7, 2026 •

edited

Loading

The `Ordered` boundary is measured, not assumed