diff --git a/sdk/cosmos/.cspell.json b/sdk/cosmos/.cspell.json index 6b703ad295f..536e3884321 100644 --- a/sdk/cosmos/.cspell.json +++ b/sdk/cosmos/.cspell.json @@ -89,6 +89,9 @@ "Daad", "dcount", "dedicatedgateway", + "dedup", + "dedupe", + "dedups", "deprioritized", "dkunda", "deprioritizes", @@ -156,6 +159,7 @@ "Gson", "Gtle", "hedgeable", + "hexdigit", "HRESULT", "hresult", "hostnames", @@ -167,6 +171,8 @@ "injective", "inlines", "inmemory", + "inspectable", + "intprop", "ints", "intptr", "INVALIDARG", @@ -187,6 +193,7 @@ "LIBCLANG", "libfuzzer", "libqueryplaninterop", + "queryplan", "QUERYPLANINTEROP", "linearizability", "livesite", @@ -279,6 +286,8 @@ "readfeed", "recompiles", "redecoded", + "reemit", + "reemitted", "reencode", "reencoded", "refetch", @@ -290,6 +299,7 @@ "Replicaset", "reqs", "restype", + "resumability", "Retriable", "retryable", "replayable", diff --git a/sdk/cosmos/azure_data_cosmos/CHANGELOG.md b/sdk/cosmos/azure_data_cosmos/CHANGELOG.md index 53e4491423b..117838876f5 100644 --- a/sdk/cosmos/azure_data_cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure_data_cosmos/CHANGELOG.md @@ -9,6 +9,7 @@ - Added `DatabaseClient::name()` and `DatabaseClient::rid()` to inspect how a database client was addressed. ([#4687](https://github.com/Azure/azure-sdk-for-rust/pull/4687)) - RID-addressed databases skip the extra database read when resolving throughput offers, reusing the addressed RID directly. ([#4687](https://github.com/Azure/azure-sdk-for-rust/pull/4687)) - Reading and querying items by RID now works end-to-end, including a parent-database cross-check that rejects a container RID belonging to a different database. ([#4687](https://github.com/Azure/azure-sdk-for-rust/pull/4687)) +- Added cross-partition `DISTINCT` query support. `SELECT DISTINCT` now deduplicates structurally equal values across every physical partition and page, rather than failing as an unsupported query feature. A `DISTINCT` query of the exact form `SELECT DISTINCT VALUE … ORDER BY ` (for example `SELECT DISTINCT VALUE c.city FROM c ORDER BY c.city`) is resumable from a continuation token; every other shape — including a list projection such as `SELECT DISTINCT c.city …` and any multi-column `ORDER BY` — is not, and requesting a token for it returns an error explaining how to rewrite the query. ([#5026](https://github.com/Azure/azure-sdk-for-rust/pull/5026)) - Added an SDK-generated `x-ms-client-id` header that remains stable for each `CosmosClient`. ([#4844](https://github.com/Azure/azure-sdk-for-rust/pull/4844)) - Added opt-in Cosmos binary JSON encoding for item operations (`create`/`read`/`replace`/`upsert`). Enable it via `CosmosClientBuilder::with_binary_encoding_options` (or the `AZURE_COSMOS_BINARY_ENCODING_ENABLED` environment-variable fallback). Off by default; when disabled, requests and responses are byte-for-byte unchanged. ([#4671](https://github.com/Azure/azure-sdk-for-rust/pull/4671)) - Added `FeedOptions::max_fan_out` (and `FeedOptions::with_max_fan_out`) to cap how many physical partitions a cross-partition query or change feed may fan out to. Applies to `ContainerClient::query_items` and `ContainerClient::query_change_feed`. The cap is enforced only at initial query setup; a partition that splits mid-execution and pushes the fan-out higher does not abort the operation. ([#4855](https://github.com/Azure/azure-sdk-for-rust/pull/4855)) diff --git a/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_hpk.rs b/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_hpk.rs index c47288e16d6..4766a4c1d9a 100644 --- a/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_hpk.rs +++ b/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_hpk.rs @@ -802,16 +802,17 @@ pub async fn hpk_query_single_partition_order_by_servable() -> Result<(), Box Result<(), Box async |run_context, db_client| { let container = seed_three_level(run_context, db_client).await?; + // Servable: DISTINCT has a client-side stage, and it must + // deduplicate correctly across the container's physical partitions. + let mut countries = collect_query::( + &container, + "SELECT DISTINCT VALUE c.country FROM c", + FeedScope::full_container(), + ) + .await? + .into_iter() + .map(|v| v.as_str().unwrap_or_default().to_owned()) + .collect::>(); + countries.sort(); + assert_eq!( + countries, + vec!["CANADA".to_string(), "USA".to_string()], + "cross-partition DISTINCT over an HPK container must return each country once" + ); + + // Still rejected: no client-side pipeline for these yet. let advanced = [ - "SELECT DISTINCT c.country FROM c", "SELECT VALUE COUNT(1) FROM c", "SELECT c.state, COUNT(1) AS n FROM c GROUP BY c.state", ]; diff --git a/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_query.rs b/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_query.rs index 5d6774d78dd..071fef8430c 100644 --- a/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_query.rs +++ b/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_query.rs @@ -327,6 +327,173 @@ pub async fn cross_partition_query_with_order_by() -> Result<(), Box> .await } +/// Unordered cross-partition `DISTINCT`: every partition contributes the same +/// partition-key value ten times, so only global client-side deduplication can +/// collapse them to one row each. +/// +/// Mirrors .NET `DistinctQueryTests.TestDistinct_ExecuteNextAsync` and Java +/// `DistinctQueryTests.queryDocuments`. +#[tokio::test] +#[cfg_attr( + not(any( + test_category = "emulator", + test_category = "emulator_vnext", + test_category = "emulator_inmemory" + )), + ignore = "requires test_category 'emulator', 'emulator_vnext', or 'emulator_inmemory'" +)] +#[cfg_attr( + test_category = "emulator_vnext", + ignore = "skipped on vnext emulator: behavioral divergence" +)] +pub async fn cross_partition_query_with_unordered_distinct() -> Result<(), Box> { + TestClient::run_with_unique_db( + async |_, db_client| { + let items = test_data::generate_mock_items(10, 10); + let container_client = + test_data::create_container_with_items(db_client, items, None).await?; + + let mut pages = container_client + .query_items::( + "select distinct value c.partitionKey from c", + FeedScope::full_container(), + Some( + QueryOptions::default().with_max_item_count(MaxItemCountHint::Limit( + std::num::NonZeroU32::new(3).unwrap(), + )), + ), + ) + .await? + .into_pages(); + + let mut actual = Vec::new(); + while let Some(page) = pages.next().await { + actual.extend(page?.into_items()); + } + actual.sort(); + + let mut expected: Vec = (0..10).map(|i| format!("partition{i}")).collect(); + expected.sort(); + assert_eq!( + actual, expected, + "each partition key must appear exactly once across the whole feed" + ); + + Ok(()) + }, + Some(TestOptions::for_emulator()), + ) + .await +} + +/// Ordered cross-partition `DISTINCT` is the resumable form: the sort key +/// matches the projection, so the merge groups equal values into runs and a +/// single retained hash carries across a continuation token. +/// +/// Mirrors .NET `DistinctQueryTests.TestDistinct_ContinuationTokenSupportAsync` +/// and Java `DistinctQueryTests.queryDocumentsWithOrderBy` (the matched-ORDER BY +/// half of `queryWithOrderByProvider`). +/// +/// The `VALUE` form is load-bearing: the service reports `distinctType: Ordered` +/// only for `SELECT DISTINCT VALUE … ORDER BY `, and only an +/// `Ordered` plan is resumable. Rewriting this query into the list form +/// (`select distinct c.partitionKey …`) would plan as `Unordered` and the +/// continuation would be refused. +#[tokio::test] +#[cfg_attr( + not(any( + test_category = "emulator", + test_category = "emulator_vnext", + test_category = "emulator_inmemory" + )), + ignore = "requires test_category 'emulator', 'emulator_vnext', or 'emulator_inmemory'" +)] +#[cfg_attr( + test_category = "emulator_vnext", + ignore = "skipped on vnext emulator: behavioral divergence" +)] +pub async fn cross_partition_query_with_ordered_distinct_resumes() -> Result<(), Box> { + TestClient::run_with_unique_db( + async |_, db_client| { + let items = test_data::generate_mock_items(10, 10); + let mut expected: Vec = (0..10).map(|i| format!("partition{i}")).collect(); + expected.sort(); + + execute_query_test( + db_client, + items, + "select distinct value c.partitionKey from c order by c.partitionKey", + FeedScope::full_container(), + expected, + QueryTestOptions { + max_item_count: Some(3), + use_continuation_token_resume: true, + }, + ) + .await?; + + Ok(()) + }, + Some(TestOptions::for_emulator()), + ) + .await +} + +/// Unordered `DISTINCT` cannot be resumed without carrying every value seen, so +/// asking for a continuation token fails loudly instead of handing back one +/// whose resume would re-emit rows. .NET blocks the token at the stage +/// (`DisallowContinuationTokenMessages.Distinct`) and Java rejects it on parse. +#[tokio::test] +#[cfg_attr( + not(any( + test_category = "emulator", + test_category = "emulator_vnext", + test_category = "emulator_inmemory" + )), + ignore = "requires test_category 'emulator', 'emulator_vnext', or 'emulator_inmemory'" +)] +#[cfg_attr( + test_category = "emulator_vnext", + ignore = "skipped on vnext emulator: behavioral divergence" +)] +pub async fn unordered_distinct_refuses_a_continuation_token() -> Result<(), Box> { + TestClient::run_with_unique_db( + async |_, db_client| { + let items = test_data::generate_mock_items(10, 10); + let container_client = + test_data::create_container_with_items(db_client, items, None).await?; + + let mut pages = container_client + .query_items::( + "select distinct value c.partitionKey from c", + FeedScope::full_container(), + Some( + QueryOptions::default().with_max_item_count(MaxItemCountHint::Limit( + std::num::NonZeroU32::new(1).unwrap(), + )), + ), + ) + .await? + .into_pages(); + + let _ = pages.next().await.expect("expected at least one page")?; + + let error = pages + .to_continuation_token() + .expect_err("an unordered DISTINCT query must not be resumable"); + let message = error.to_string(); + assert!( + message.contains("ORDER BY"), + "the refusal must tell the caller how to make the query resumable: {message}" + ); + + Ok(()) + }, + Some(TestOptions::for_emulator()), + ) + .await +} + #[tokio::test] #[cfg_attr( not(any( @@ -860,3 +1027,221 @@ pub async fn single_partition_query_resumes_with_raw_server_token() -> Result<() ) .await } + +/// Live counterparts for the `DISTINCT` scenarios that otherwise run only +/// against the in-memory emulator. +/// +/// Those scenarios assert behavior the emulator *simulates* — its own per- +/// partition dedup and its local query-plan generator — so on their own they +/// are partly self-referential: they would pass even if the service disagreed. +/// This test re-checks the same shapes against a real account. +/// +/// Covers `select_star_whole_document`, +/// `select_value_constant_with_from_still_deduplicates`, +/// `filters_and_parameters_apply_before_distinct`, and +/// `single_logical_partition_scope` from `distinct_scenarios.json`. +#[tokio::test] +#[cfg_attr( + not(any( + test_category = "emulator", + test_category = "emulator_vnext", + test_category = "emulator_inmemory" + )), + ignore = "requires test_category 'emulator', 'emulator_vnext', or 'emulator_inmemory'" +)] +#[cfg_attr( + test_category = "emulator_vnext", + ignore = "skipped on vnext emulator: behavioral divergence" +)] +pub async fn distinct_projection_shapes() -> Result<(), Box> { + TestClient::run_with_unique_db( + async |_, db_client| { + // 4 partitions x 3 items: `partitionKey` repeats within a + // partition, `id` is unique across the container. + let items = test_data::generate_mock_items(4, 3); + let container = test_data::create_container_with_items(db_client, items, None).await?; + + async fn count( + container: &ContainerClient, + query: impl Into, + scope: FeedScope, + ) -> Result> { + let mut pages = container + .query_items::(query, scope, None) + .await? + .into_pages(); + let mut n = 0; + while let Some(page) = pages.next().await { + n += page?.into_items().len(); + } + Ok(n) + } + + // `SELECT DISTINCT *` dedups whole documents; every document is + // unique, so nothing collapses and the passthrough stays intact. + assert_eq!( + count( + &container, + "select distinct * from c", + FeedScope::full_container() + ) + .await?, + 12, + "DISTINCT over unique whole documents must not drop any" + ); + + // A constant projection *with* a FROM clause yields one row per + // document, all identical, so exactly one survives. (Without a FROM + // the service collapses DISTINCT away entirely — see + // `plan::distinct_is_ordered`'s sibling constant-collapse rule.) + assert_eq!( + count( + &container, + "select distinct value 1 from c", + FeedScope::full_container() + ) + .await?, + 1, + "a constant projection over N documents must collapse to one row" + ); + + // A parameterized WHERE narrows the rows before deduplication. + assert_eq!( + count( + &container, + Query::from( + "select distinct value c.partitionKey from c where c.mergeOrder >= @m" + ) + .with_parameter("@m", 0)?, + FeedScope::full_container() + ) + .await?, + 4, + "expected one row per distinct partition key" + ); + + // Scoped to a single logical partition, DISTINCT still dedups; it + // just never fans out. + assert_eq!( + count( + &container, + "select distinct value c.partitionKey from c", + FeedScope::partition("partition0") + ) + .await?, + 1, + "a single-partition scope must yield that partition's one key" + ); + + Ok(()) + }, + Some(TestOptions::for_emulator()), + ) + .await +} + +/// Live counterpart for the `unsupported_combination_*` scenarios. +/// +/// Those scenarios assert the driver rejects `DISTINCT` combined with a stage +/// it has no pipeline for. That rejection is only half the story: because the +/// SDK advertises just the features it implements (`SUPPORTED_QUERY_FEATURES`), +/// the *service* refuses these plans first, with 400 / 1004 +/// CrossPartitionQueryNotServable. Pinning that here means adding a new feature +/// token without its pipeline stage cannot silently start serving a query the +/// driver would then mishandle. +#[tokio::test] +#[cfg_attr( + not(any( + test_category = "emulator", + test_category = "emulator_vnext", + test_category = "emulator_inmemory" + )), + ignore = "requires test_category 'emulator', 'emulator_vnext', or 'emulator_inmemory'" +)] +#[cfg_attr( + test_category = "emulator_vnext", + ignore = "skipped on vnext emulator: behavioral divergence" +)] +pub async fn distinct_combined_with_unsupported_stages_is_rejected() -> Result<(), Box> { + TestClient::run_with_unique_db( + async |_, db_client| { + let items = test_data::generate_mock_items(4, 3); + let container = test_data::create_container_with_items(db_client, items, None).await?; + + // Each of these needs a composition stage the driver does not have + // yet: GROUP BY and aggregates. `TOP` and `OFFSET`/`LIMIT` are no + // longer here — `SkipTake` composes above `DISTINCT`, so those + // shapes are servable and are asserted positively below. + let unsupported = [ + "select distinct c.partitionKey, count(1) as n from c group by c.partitionKey", + "select distinct value max(c.mergeOrder) from c", + ]; + + for query in unsupported { + let outcome = container + .query_items::(query, FeedScope::full_container(), None) + .await; + let error = match outcome { + Err(error) => error, + Ok(pager) => { + // Some shapes fail while draining rather than at setup. + let mut pages = pager.into_pages(); + let mut drain_error = None; + while let Some(page) = pages.next().await { + if let Err(error) = page { + drain_error = Some(error); + break; + } + } + drain_error.unwrap_or_else(|| { + panic!("expected `{query}` to be rejected, but it drained cleanly") + }) + } + }; + assert_eq!( + error.status().status_code(), + azure_core::http::StatusCode::BadRequest, + "expected 400 for `{query}`, got {error}" + ); + } + + // `DISTINCT` composed under a row window is servable, and the window + // counts *deduplicated* values: 4 partition keys over 12 items, so + // `TOP 2` yields 2 and `OFFSET 1 LIMIT 2` yields 2. Applying the + // window before deduplication would return fewer. + for (query, expected) in [ + ("select distinct top 2 value c.partitionKey from c", 2usize), + ( + "select distinct value c.partitionKey from c offset 1 limit 2", + 2usize, + ), + ] { + let mut pages = container + .query_items::(query, FeedScope::full_container(), None) + .await? + .into_pages(); + let mut values = Vec::new(); + while let Some(page) = pages.next().await { + values.extend(page?.into_items()); + } + assert_eq!( + values.len(), + expected, + "`{query}` must apply its window to deduplicated values, got {values:?}" + ); + let mut deduped = values.clone(); + deduped.sort_by_key(|v| v.to_string()); + deduped.dedup_by_key(|v| v.to_string()); + assert_eq!( + deduped.len(), + values.len(), + "`{query}` returned duplicate values: {values:?}" + ); + } + + Ok(()) + }, + Some(TestOptions::for_emulator()), + ) + .await +} diff --git a/sdk/cosmos/azure_data_cosmos/tests/split_tests/cosmos_query_distinct_split.rs b/sdk/cosmos/azure_data_cosmos/tests/split_tests/cosmos_query_distinct_split.rs new file mode 100644 index 00000000000..417de015949 --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos/tests/split_tests/cosmos_query_distinct_split.rs @@ -0,0 +1,322 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +//! Live-only split coverage for cross-partition `DISTINCT`. +//! +//! Neither .NET nor Java tests `DISTINCT` against a real partition split +//! (.NET's `FullPipelineTests.TestMerge` covers `ORDER BY` only), so this is +//! the one behavior with no peer precedent to lean on. It matters because the +//! deduplication state lives in a stage *above* the fan-out root: if a split +//! ever caused that stage to be rebuilt, values emitted before the split would +//! come back. +//! +//! The container is seeded so every `groupKey` value is produced by many +//! partition keys — after a split, both physical partitions still contribute +//! rows for the same values, which is precisely the case a per-partition or +//! per-page dedup would get wrong. +//! +//! Two invariants are asserted across one live split: +//! +//! - **Unordered** `DISTINCT` drained straight through a split returns each +//! value exactly once. +//! - **Ordered** `DISTINCT` (matching `ORDER BY`) resumed from a continuation +//! token captured *before* the split returns each value exactly once, in +//! sorted order, with no gap at the boundary. +//! +//! Runs only under `test_category = "split"` against split-capable resources. + +use super::framework; +use crate::split_tests::cosmos_query_split::force_split_and_wait; + +use std::collections::BTreeSet; +use std::error::Error; +use std::num::NonZeroU32; +use std::time::Duration; + +use azure_data_cosmos::feed::ContinuationToken; +use azure_data_cosmos::options::CreateContainerOptions; +use azure_data_cosmos::{ + clients::ContainerClient, + feed::FeedScope, + models::{ContainerProperties, CosmosStatus, ThroughputProperties}, + options::{MaxItemCountHint, QueryOptions}, +}; +use framework::{TestClient, TestOptions}; +use futures::StreamExt; +use serde::{Deserialize, Serialize}; + +const PK_COUNT: usize = 40; +const GROUP_COUNT: usize = 8; +const PAGE_SIZE: u32 = 5; + +const UNORDERED_QUERY: &str = "SELECT DISTINCT VALUE c.groupKey FROM c"; +const ORDERED_QUERY: &str = "SELECT DISTINCT VALUE c.groupKey FROM c ORDER BY c.groupKey"; +/// `DISTINCT` composed under a global row window. `GROUP_COUNT` distinct keys +/// exist, so skipping one and taking two must yield exactly two — the window +/// counts deduplicated values, not raw rows. +const WINDOWED_QUERY: &str = + "SELECT DISTINCT VALUE c.groupKey FROM c ORDER BY c.groupKey OFFSET 1 LIMIT 2"; + +/// A seeded document whose `groupKey` deliberately repeats across every +/// partition key, so deduplication has to be global. +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "camelCase")] +struct SeedItem { + id: String, + partition_key: String, + group_key: String, +} + +/// The distinct values the seed produces, in sorted order. +fn expected_group_keys() -> Vec { + (0..GROUP_COUNT).map(|i| format!("g{i:02}")).collect() +} + +fn assert_each_value_exactly_once(actual: &[String], context: &str) { + let expected = expected_group_keys(); + let unique: BTreeSet<&String> = actual.iter().collect(); + assert_eq!( + unique.len(), + actual.len(), + "{context}: DISTINCT returned a duplicate value: {actual:?}" + ); + let mut sorted = actual.to_vec(); + sorted.sort(); + assert_eq!( + sorted, expected, + "{context}: DISTINCT did not return exactly the seeded value set" + ); +} + +/// Drains `query` one page at a time, invoking `after_first_page` once the +/// first page has been collected (used to force the split mid-drain). +async fn drain_with_hook( + container_client: &ContainerClient, + query: &str, + mut after_first_page: F, +) -> Result, Box> +where + F: FnMut() -> Fut, + Fut: std::future::Future>>, +{ + let options = QueryOptions::default() + .with_max_item_count(MaxItemCountHint::Limit(NonZeroU32::new(PAGE_SIZE).unwrap())); + let mut pages = container_client + .query_items::(query, FeedScope::full_container(), Some(options)) + .await? + .into_pages(); + + let mut collected = Vec::new(); + let mut page_index = 0usize; + while let Some(page) = pages.next().await { + collected.extend(page?.into_items()); + if page_index == 0 { + after_first_page().await?; + } + page_index += 1; + } + Ok(collected) +} + +/// Cross-partition `DISTINCT` across a live partition split. +/// +/// Part 1 drains an unordered `DISTINCT` straight through a forced split. +/// Part 2 captures an ordered `DISTINCT` continuation token before the split +/// (already taken, since the split happened in part 1) and resumes it against +/// the post-split topology. +#[tokio::test] +#[cfg_attr( + not(test_category = "split"), + ignore = "requires test_category 'split'" +)] +pub async fn distinct_query_across_split_returns_each_value_once() -> Result<(), Box> { + TestClient::run_with_unique_db( + async |run_context, db_client| { + let properties = + ContainerProperties::new("DistinctAcrossSplit", "/partitionKey".into()); + let throughput = ThroughputProperties::manual(1000); + let container_client = run_context + .create_container( + db_client, + properties, + Some(CreateContainerOptions::default().with_throughput(throughput)), + ) + .await?; + + println!( + "Container created; seeding {PK_COUNT} partition keys x {GROUP_COUNT} repeated \ + group keys" + ); + for p in 0..PK_COUNT { + let partition_key = format!("pk{p}"); + for i in 0..GROUP_COUNT { + let item = SeedItem { + id: format!("{p}-{i}"), + partition_key: partition_key.clone(), + group_key: format!("g{i:02}"), + }; + container_client + .create_item(item.partition_key.clone(), &item.id.clone(), item, None) + .await?; + } + } + + let ranges_before = container_client.read_feed_ranges(None).await?; + assert!( + !ranges_before.is_empty(), + "expected at least one physical partition before split, got {}", + ranges_before.len() + ); + let partitions_before = ranges_before.len(); + + // ── Ordered DISTINCT: capture a token before the split ──────── + // + // Taken first so the checkpoint predates the topology change. + let ordered_options = QueryOptions::default() + .with_max_item_count(MaxItemCountHint::Limit(NonZeroU32::new(2).unwrap())); + let mut ordered_pages = container_client + .query_items::( + ORDERED_QUERY, + FeedScope::full_container(), + Some(ordered_options), + ) + .await? + .into_pages(); + let mut ordered_collected: Vec = ordered_pages + .next() + .await + .expect("ordered DISTINCT should yield at least one page")? + .into_items(); + // `ORDERED_QUERY` uses the `VALUE` form, which the service reports + // as `distinctType: Ordered` and is therefore resumable. Read the + // outcome back rather than asserting it, so that if a future service + // version downgrades the shape this reports a self-describing skip + // instead of a bare panic that would blame the wrong thing. + let ordered_token = match ordered_pages.to_continuation_token() { + Ok(token) => ContinuationToken::from_string(token.as_str().to_owned()), + // Only an explicit "this shape cannot be resumed" refusal is a + // legitimate service-plan downgrade. Any other failure is a + // continuation regression and must not be swallowed, or this + // test would go green while resume is broken. + Err(error) + if error.status().sub_status() + == CosmosStatus::CLIENT_DISTINCT_CONTINUATION_UNSUPPORTED.sub_status() => + { + panic!( + "service planned `{ORDERED_QUERY}` as unordered DISTINCT (continuation \ + refused: {error}). The VALUE form with a matching ORDER BY is a required \ + contract for resumable DISTINCT; if the service genuinely changed, update \ + `plan::distinct_is_ordered` and the docs rather than skipping this test." + ); + } + Err(error) => { + return Err(format!( + "minting a continuation for `{ORDERED_QUERY}` failed with an unexpected \ + error (not the unsupported-continuation status): {error}" + ) + .into()); + } + }; + drop(ordered_pages); + assert!( + !ordered_collected.is_empty(), + "the pre-split checkpoint must have emitted at least one value" + ); + + // ── Unordered DISTINCT: drain straight through the split ────── + let mut split_done = false; + let unordered = drain_with_hook(&container_client, UNORDERED_QUERY, || { + let container_client = container_client.clone(); + let should_split = !split_done; + split_done = true; + async move { + if should_split { + let partitions_after = + force_split_and_wait(&container_client, partitions_before).await?; + assert!( + partitions_after > partitions_before, + "split must increase partition count: before={partitions_before}, \ + after={partitions_after}" + ); + } + Ok(()) + } + }) + .await?; + assert_each_value_exactly_once(&unordered, "unordered DISTINCT across a split"); + + // ── Ordered DISTINCT: resume the pre-split token ────────────── + let mut continuation = Some(ordered_token); + loop { + let mut options = QueryOptions::default() + .with_max_item_count(MaxItemCountHint::Limit(NonZeroU32::new(2).unwrap())); + if let Some(token) = continuation.take() { + options = options.with_continuation_token(token); + } + let mut pages = container_client + .query_items::( + ORDERED_QUERY, + FeedScope::full_container(), + Some(options), + ) + .await? + .into_pages(); + let Some(page) = pages.next().await else { + break; + }; + ordered_collected.extend(page?.into_items()); + let serialized = pages.to_continuation_token()?.as_str().to_owned(); + drop(pages); + continuation = Some(ContinuationToken::from_string(serialized)); + } + + assert_each_value_exactly_once( + &ordered_collected, + "ordered DISTINCT resumed across a split", + ); + let mut sorted = ordered_collected.clone(); + sorted.sort(); + assert_eq!( + ordered_collected, sorted, + "an ordered DISTINCT resume must preserve global sort order across the split" + ); + + // ── DISTINCT under a row window, drained across the split ───── + // + // `SkipTake` wraps `Distinct`, so a split must preserve both the + // dedup state and the window's remaining budget. Losing the former + // re-emits a value the window already paid for; losing the latter + // restarts the offset and over-returns. The emulator covers this + // with a simulated split — this is the same shape against a real + // one, where the fan-out is genuinely rebuilt. + let windowed = + drain_with_hook(&container_client, WINDOWED_QUERY, || async { Ok(()) }).await?; + assert_eq!( + windowed.len(), + 2, + "`{WINDOWED_QUERY}` must apply its window to deduplicated values across the \ + split, got {windowed:?}" + ); + let mut windowed_sorted = windowed.clone(); + windowed_sorted.sort(); + windowed_sorted.dedup(); + assert_eq!( + windowed_sorted.len(), + windowed.len(), + "windowed DISTINCT across a split returned duplicates: {windowed:?}" + ); + let mut in_order = windowed.clone(); + in_order.sort(); + assert_eq!( + windowed, in_order, + "windowed DISTINCT must preserve global sort order across the split" + ); + + Ok(()) + }, + // A real split takes minutes; the 80s default would abort mid-poll. + // Matches the other split tests in this directory. + Some(TestOptions::new().with_timeout(Duration::from_secs(40 * 60))), + ) + .await +} diff --git a/sdk/cosmos/azure_data_cosmos/tests/split_tests/mod.rs b/sdk/cosmos/azure_data_cosmos/tests/split_tests/mod.rs index 48239ecfa99..dd24335e890 100644 --- a/sdk/cosmos/azure_data_cosmos/tests/split_tests/mod.rs +++ b/sdk/cosmos/azure_data_cosmos/tests/split_tests/mod.rs @@ -1,6 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. mod cosmos_change_feed_split; +mod cosmos_query_distinct_split; mod cosmos_query_order_by_split; mod cosmos_query_skip_take_split; mod cosmos_query_split; diff --git a/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md b/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md index a7189554373..5df18bcd4d7 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md +++ b/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md @@ -7,6 +7,7 @@ - Added binary round trip fuzzer. As a part of the implementation, binary JSON responses now deserialize a service-echoed integral `Double` into a signed or unsigned integer field (previously a type error); this is intentionally lossy for integers the service cannot represent exactly, while a fractional `Double` remains a type error. Does not yet cover integer elements inside a uniform `Float64` array or an enum variant. ([#4976](https://github.com/Azure/azure-sdk-for-rust/pull/4976)) - Added driver-internal resolution of containers by resource id (RID). `CosmosDriver::resolve_container_by_rid` reads a container's metadata addressing it purely by RID (deriving the parent database RID from the container RID, so no `read_database` round-trip is needed) and caches the result in a by-RID index. References are validated for consistent name/RID addressing in `plan_operation` — the single choke point every executable operation passes through, including multi-page queries — returning a deterministic `CLIENT_MIXED_NAME_RID_ADDRESSING` error before signing instead of letting the gateway reject a mixed name/RID request with an opaque `401`. The new `CLIENT_INVALID_RESOURCE_ID` and `CLIENT_MIXED_NAME_RID_ADDRESSING` client statuses carry searchable names for diagnostics. ([#4663](https://github.com/Azure/azure-sdk-for-rust/pull/4663)) - Added `models::is_database_rid`, which reports whether a RID string decodes to a database-level RID (4 bytes). Lets callers that reuse a supplied RID as a database identity reject a wrong-hierarchy RID before it addresses the wrong resource. ([#4640](https://github.com/Azure/azure-sdk-for-rust/pull/4640)) +- Added cross-partition `DISTINCT` query support (ordered and unordered), composed as a stage above the fan-out root and below any `OFFSET`/`LIMIT`/`TOP` window (so the row limit counts deduplicated values), and keyed on a structural, type-aware hash of the projected row, so structurally equal values collapse across partitions and pages. Ordered `DISTINCT` — only `SELECT DISTINCT VALUE … ORDER BY `, the one shape the service reports as ordered; list projections and multi-column `ORDER BY` are unordered — is resumable, carrying only the last emitted hash in its continuation token. Unordered `DISTINCT` is not: `OperationPlan::to_continuation_token` fails with the new `CosmosStatus::CLIENT_DISTINCT_CONTINUATION_UNSUPPORTED` (HTTP 400) rather than returning a token whose resume would re-emit rows. A partition split that reaches the `DISTINCT` stage is refused with `CosmosStatus::CLIENT_DISTINCT_CANNOT_FORWARD_SPLIT` (HTTP 500) instead of being forwarded, since `SplitRequired` replaces the node that emits it and would discard the deduplication state. ([#5026](https://github.com/Azure/azure-sdk-for-rust/pull/5026)) - Added an SDK-generated `x-ms-client-id` header that remains stable for each `CosmosDriver`, including metadata, retry, hedge, probe, and Gateway 2.0 outer HTTP requests. ([#4844](https://github.com/Azure/azure-sdk-for-rust/pull/4844)) - Added a schema-agnostic Cosmos binary JSON codec (`binary_json`) and driver-side binary encoding via `OperationOptions.binary_encoding` (`BinaryEncodingOptions`). When enabled, the driver transcodes item request/response bodies between text and Cosmos binary JSON and negotiates the wire format; it is honored only for point `Document` item operations. Off by default and inert on the wire when unset. ([#4671](https://github.com/Azure/azure-sdk-for-rust/pull/4671)) - Added `PlanOptions` (with `DEFAULT_MAX_FAN_OUT`) to `CosmosDriver::plan_operation`, enforcing a maximum fan-out on fresh cross-partition plans. A fresh plan spanning more leaf request nodes than `PlanOptions::max_fan_out` (default 100) is rejected with the new `CosmosStatus::CLIENT_CROSS_PARTITION_FAN_OUT_EXCEEDED` (HTTP 400). The limit is enforced only at initial plan time: resuming from a continuation token skips the check, and a partition split that raises the fan-out mid-execution does not abort the operation. ([#4855](https://github.com/Azure/azure-sdk-for-rust/pull/4855)) diff --git a/sdk/cosmos/azure_data_cosmos_driver/docs/FEED_OPERATIONS_REQS.md b/sdk/cosmos/azure_data_cosmos_driver/docs/FEED_OPERATIONS_REQS.md index 46d0fd03ca7..cf6e271d901 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/docs/FEED_OPERATIONS_REQS.md +++ b/sdk/cosmos/azure_data_cosmos_driver/docs/FEED_OPERATIONS_REQS.md @@ -165,6 +165,22 @@ Synthetic streaming `ORDER BY` pages sum request charge, merge compound session Within a resumed full-key tie, RID filtering follows each backend page's `x-ms-cosmos-query-execution-info`: modern pages use `reverseIndexScan`, while absent or legacy `reverseRidEnabled` metadata falls back to the first ORDER BY direction. Across partition streams, equal keys are ordered by leftmost EPK range, matching .NET and Java; RID is only a per-range continuation discriminator. +### Cross-Partition `DISTINCT` + +`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** (`driver::dataflow::distinct_hash`), not on the `ORDER BY` items, so one node serves both modes and only the retained state differs. The hash gives each JSON type its own seed (so `null`/`false`/`""`/`[]`/`{}` never collide), treats arrays as position-sensitive and objects as order-insensitive, equates `5` with `5.0`, and normalizes `-0.0`. It is standalone so `GROUP BY` can reuse it. + +Ordered `DISTINCT` — only `SELECT DISTINCT VALUE … ORDER BY `, the exact shape the service reports as `Ordered` (see below); list projections and multi-column `ORDER BY` stay unordered even when every projected path is covered — deduplicates by adjacency, so it keeps one hash, runs in O(1) memory, 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 — two documents that are one `DISTINCT` row but two `ORDER BY` rows. + +Unordered `DISTINCT` retains every hash seen (unbounded, ~16 bytes per distinct value) and is **not** resumable: the set *is* the state, serializing it would produce an unbounded token, and truncating it would silently re-emit duplicates. `Distinct::snapshot_state` fails with `400 / 20124 ClientDistinctContinuationUnsupported`, so `OperationPlan::to_continuation_token` errors at mint time — while the caller still holds a live plan and can keep draining in process or rewrite with a matching `ORDER BY`. In-process paging is fully supported. Once drained there is no state left to lose, so the stage snapshots as `Drained` like any other finished node. + +The driver executes whatever `distinctType` the plan reports and never upgrades `Unordered` to `Ordered`. The local plan generator (`query::plan`, backing the in-memory emulator) is deliberately stricter than the service planner — see `plan::distinct_is_ordered` — because misclassifying a stream as adjacency-safe drops rows, while the reverse only costs resumability. + +`Ordered` requires a specific query shape, measured against a live account with production's `SUPPORTED_QUERY_FEATURES`: the service reports it for `SELECT DISTINCT VALUE FROM c ORDER BY ` (either direction) and `Unordered` for everything else — the list form (`SELECT DISTINCT c.name …`) whatever its `ORDER BY`, a sort on a different path, a multi-column `ORDER BY` even when it leads with the projected path, and no `ORDER BY` at all. Advertising `Distinct` is mandatory: without it the service rejects any `DISTINCT` query with `400 / 1004 CrossPartitionQueryNotServable`. `tests/gateway_query_plan_comparison.rs::gw_distinct` pins each shape against the live service, and `plan::distinct_is_ordered` encodes the same rule so the local generator agrees. + +A live split is safe in both modes: the fan-out node absorbs `SplitRequired` internally and `Distinct` is never rebuilt, so an emitted value cannot be resurrected — and `last_hash` is a value rather than a position, so re-resolving ranges does not invalidate it. Resumed tokens are validated against the plan: a `distinct_type` mismatch, a `Distinct` token for a plan that no longer deduplicates, and a pre-`DISTINCT` token for one that does are all rejected rather than reinterpreted. + +Pages whose rows are all duplicates are suppressed rather than surfaced as empty pages, but their request charge and diagnostics fold into the next emitted page (or flush as a final empty page if the child drains first), so a redundant `DISTINCT` query never under-reports its cost. + ### The Driver DOES - Plan the pipeline (determine targeting, resolve partitions, build the node tree). diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs index 46e96e212ec..4666164fcea 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs @@ -3409,8 +3409,9 @@ impl CosmosDriver { options: &OperationOptions, ) -> crate::error::Result { // Advertise exactly the query-rewrite features implemented by the - // production dataflow pipeline (`OrderBy,MultipleOrderBy`). The value - // must remain non-empty so Gateway V2 accepts the QueryPlan request. + // production dataflow pipeline (see `query::SUPPORTED_QUERY_FEATURES`). + // The value must remain non-empty so Gateway V2 accepts the QueryPlan + // request. let query_plan_operation = CosmosOperation::query_plan( container.clone(), std::borrow::Cow::Borrowed(crate::query::SUPPORTED_QUERY_FEATURES), diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/distinct.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/distinct.rs new file mode 100644 index 00000000000..f8cc1a89d46 --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/distinct.rs @@ -0,0 +1,1033 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +//! Cross-partition `DISTINCT` deduplication. +//! +//! [`Distinct`] wraps the cross-partition fan-out root and drops rows whose +//! projected payload is structurally equal to one already emitted. It sits +//! *above* the merge, matching .NET's `DistinctQueryPipelineStage` and Java's +//! `DistinctDocumentQueryExecutionContext`: +//! +//! ```text +//! Unordered: Distinct -> SequentialDrain -> Request… +//! Ordered: Distinct -> StreamingOrderedMerge -> Request… +//! ``` +//! +//! Both peers key on the **whole projected row**, not on the `ORDER BY` items, +//! so one node serves both modes and only the map differs. +//! +//! # Ordered vs unordered +//! +//! - [`DistinctMap::Ordered`] retains a single hash. `ORDER BY` guarantees +//! structurally equal rows arrive adjacently, so comparing against the last +//! emitted value is sufficient and the node runs in O(1) memory. +//! - [`DistinctMap::Unordered`] retains every hash seen. There is no ordering +//! to exploit, so a duplicate may arrive arbitrarily far from its twin — +//! including from a different partition. +//! +//! # Continuation +//! +//! Ordered `DISTINCT` is resumable: the 16-byte `last_hash` is all a resumed +//! node needs, because a value it has moved past can never reappear. +//! +//! Unordered `DISTINCT` is not. The set *is* the state, and serializing it +//! would mean an unbounded token; truncating it would silently re-emit +//! duplicates. [`Distinct::snapshot_state`] therefore fails with +//! [`CosmosStatus::CLIENT_DISTINCT_CONTINUATION_UNSUPPORTED`], which surfaces +//! at `OperationPlan::to_continuation_token` time — while the caller still +//! holds a live plan and can either keep draining in-process or rewrite the +//! query with a matching `ORDER BY`. .NET refuses here too, with the same +//! guidance. Java does not: it emits a token carrying only the source +//! continuation, then rebuilds an empty map on resume, so an unordered +//! `DISTINCT` resumed from a Java token silently re-emits values it already +//! returned. +//! +//! # Splits +//! +//! The wrapped fan-out node absorbs `SplitRequired` internally and is never +//! rebuilt, so the map survives a split mid-drain and an already-emitted value +//! cannot be resurrected. A split that does reach this node is refused rather +//! than forwarded: `SplitRequired` replaces the node that emits it, so passing +//! it up would discard the map along with the node. + +use std::collections::HashSet; +use std::sync::Arc; + +use async_trait::async_trait; +use bytes::Bytes; + +use crate::diagnostics::DiagnosticsContext; +use crate::error::CosmosStatus; +use crate::models::{CosmosResponse, FeedRange, RequestCharge, ResponseBody}; + +use super::distinct_hash::{hash_value, Hash128}; +use super::query_plan::DistinctType; +use super::{skip_take_page, PageResult, PipelineContext, PipelineNode, PipelineNodeState}; + +/// Guidance surfaced when a caller asks for a continuation token on an +/// unordered `DISTINCT` query. Mirrors .NET's +/// `DisallowContinuationTokenMessages.Distinct`. +const UNORDERED_CONTINUATION_MESSAGE: &str = + "DISTINCT queries only return continuation tokens when there is a matching ORDER BY clause. \ + For example, rewrite `SELECT DISTINCT VALUE c.name FROM c` as \ + `SELECT DISTINCT VALUE c.name FROM c ORDER BY c.name`."; + +/// Deduplication state, keyed on a 128-bit structural hash of the row payload. +enum DistinctMap { + /// Adjacency deduplication over an `ORDER BY`-sorted stream. `None` until + /// the first row is emitted (or, on resume, seeded from the token). + Ordered { last_hash: Option }, + + /// Global deduplication over an unordered stream. + /// + /// Unbounded by design, matching .NET's `UnorderedDistinctMap` and Java's + /// `UnorderedDistinctMap`: ~16 bytes per *distinct* value seen. Since the + /// query cannot be resumed anyway, the set only has to survive one drain. + Unordered { seen: HashSet }, +} + +impl DistinctMap { + /// Records `hash` and reports whether the row should be emitted. + fn accept(&mut self, hash: Hash128) -> bool { + match self { + Self::Ordered { last_hash } => { + let is_new = *last_hash != Some(hash); + // Advance unconditionally: a non-adjacent repeat is a different + // run and must reset the comparison point. + *last_hash = Some(hash); + is_new + } + Self::Unordered { seen } => seen.insert(hash), + } + } +} + +/// Deduplicates its single child's pages by structural payload equality. +pub(crate) struct Distinct { + child: Box, + map: DistinctMap, + /// Set once the child drains so subsequent pulls short-circuit. + exhausted: bool, + /// Request charge from pages whose rows were entirely duplicates and were + /// therefore suppressed rather than emitted as empty pages. Folded into the + /// next emitted page so billed RUs are never under-reported. + suppressed_charge: RequestCharge, + /// Diagnostics from those same pages, folded incrementally into a single + /// context rather than retained one-per-page: `aggregate_sub_operations` + /// re-bounds its record list to `max_request_diagnostics`, so a long run of + /// all-duplicate pages cannot grow the artifact without limit. + suppressed_diagnostics: Option>, + /// The most recently suppressed page, kept as a template so accumulated + /// charge/diagnostics can still be flushed as a final empty page if the + /// child drains without ever surfacing a terminal page. + pending_flush: Option, +} + +impl Distinct { + /// Wraps `child` with a fresh map for `distinct_type`. + #[cfg(test)] + pub(crate) fn new(child: Box, distinct_type: DistinctType) -> Self { + Self::with_last_hash(child, distinct_type, None) + } + + /// Wraps `child`, seeding an ordered map with the hash of the last row + /// emitted before a checkpoint. `last_hash` is ignored for an unordered + /// map, whose state is never persisted. + pub(crate) fn with_last_hash( + child: Box, + distinct_type: DistinctType, + last_hash: Option, + ) -> Self { + let map = match distinct_type { + // `None` never reaches here — the planner only wraps a root when + // the plan asks for deduplication — but treating it as unordered + // keeps this total and errs toward deduplicating rather than + // silently passing duplicates through. + DistinctType::Unordered | DistinctType::None => DistinctMap::Unordered { + seen: HashSet::new(), + }, + DistinctType::Ordered => DistinctMap::Ordered { last_hash }, + }; + Self { + child, + map, + exhausted: false, + suppressed_charge: RequestCharge::default(), + suppressed_diagnostics: None, + pending_flush: None, + } + } + + /// Selects the items that survive deduplication, preserving order and each + /// item's exact backend bytes. + fn select_survivors(&mut self, items: &[Bytes]) -> crate::error::Result> { + let mut keep = Vec::with_capacity(items.len()); + for item in items { + let value: serde_json::Value = serde_json::from_slice(item).map_err(|e| { + crate::error::CosmosError::builder() + .with_status(CosmosStatus::SERIALIZATION_RESPONSE_BODY_INVALID) + .with_message("failed to parse a DISTINCT row payload as JSON") + .with_source(e) + .build() + })?; + if self.map.accept(hash_value(&value)?) { + keep.push(item.clone()); + } + } + Ok(keep) + } + + /// Rebuilds a response around a trimmed body, updating `x-ms-item-count` + /// and folding in charge/diagnostics accumulated from suppressed pages. + fn rebuild( + &mut self, + response: &CosmosResponse, + survivors: Vec, + emitted: usize, + ) -> CosmosResponse { + let mut headers = response.headers().clone(); + headers.item_count = Some(emitted as u32); + if self.suppressed_charge != RequestCharge::default() { + let base = headers.request_charge.unwrap_or_default(); + headers.request_charge = Some(base + self.suppressed_charge); + } + let rebuilt = CosmosResponse::new( + ResponseBody::from_items(survivors), + headers, + response.status(), + response.diagnostics(), + ); + let merged = match self.suppressed_diagnostics.as_ref() { + Some(accumulated) => { + rebuilt.with_aggregated_prior_diagnostics(std::slice::from_ref(accumulated)) + } + None => rebuilt, + }; + self.clear_suppressed(); + merged + } + + /// Accumulates an all-duplicate page's charge and diagnostics. + fn suppress(&mut self, response: CosmosResponse) { + self.suppressed_charge = + self.suppressed_charge + response.headers().request_charge.unwrap_or_default(); + let incoming = response.diagnostics(); + // Fold on arrival so only one context is ever retained. The newest page + // stays last, preserving the aggregate's "operation-level fields come + // from the final source" contract. + self.suppressed_diagnostics = match self.suppressed_diagnostics.take() { + None => Some(incoming), + Some(accumulated) => { + DiagnosticsContext::aggregate_sub_operations(&[accumulated, incoming]).map(Arc::new) + } + }; + self.pending_flush = Some(response); + } + + fn clear_suppressed(&mut self) { + self.suppressed_charge = RequestCharge::default(); + self.suppressed_diagnostics = None; + self.pending_flush = None; + } + + /// Emits a final empty page carrying accumulated suppressed charge and + /// diagnostics, or `None` if nothing is pending. + fn flush_suppressed(&mut self) -> Option { + let template = self.pending_flush.take()?; + // `suppressed_diagnostics` already includes the template's own + // diagnostics, so aggregate the list rather than layering onto it. + let diagnostics = self + .suppressed_diagnostics + .clone() + .unwrap_or_else(|| template.diagnostics()); + let mut headers = template.headers().clone(); + headers.item_count = Some(0); + headers.request_charge = Some(self.suppressed_charge); + let response = CosmosResponse::new( + ResponseBody::from_items(Vec::new()), + headers, + template.status(), + diagnostics, + ); + self.clear_suppressed(); + Some(PageResult::Page { + response, + is_terminal: true, + }) + } +} + +#[async_trait] +impl PipelineNode for Distinct { + async fn next_page( + &mut self, + context: &mut PipelineContext<'_>, + ) -> crate::error::Result { + if self.exhausted { + return Ok(PageResult::Drained); + } + + loop { + match self.child.next_page(context).await? { + PageResult::Drained => { + self.exhausted = true; + // Flush charge/diagnostics from a fully-duplicate tail that + // never got a terminal page of its own. + if let Some(flushed) = self.flush_suppressed() { + return Ok(flushed); + } + return Ok(PageResult::Drained); + } + PageResult::SplitRequired { .. } => { + // `SplitRequired` replaces the node that emits it, so + // forwarding would drop this node along with its + // deduplication map and resurrect suppressed values. The + // wrapped fan-out node absorbs splits internally, so this + // is unreachable today; fail loudly if that ever changes. + return Err(crate::error::CosmosError::builder() + .with_status(CosmosStatus::CLIENT_DISTINCT_CANNOT_FORWARD_SPLIT) + .with_message( + "DISTINCT cannot forward a partition split; the wrapped fan-out \ + node must absorb splits internally", + ) + .build()); + } + PageResult::Page { + response, + is_terminal, + } => { + // Normalize the child's page into per-document slices. A + // streaming ordered merge (and a `SkipTake` below us) hands + // over pre-split `Items`; a raw backend feed page arrives as + // `Bytes`; `NoPayload` is a zero-document page. + let items: Vec = match response.body() { + ResponseBody::Items(items) => items.clone(), + ResponseBody::Bytes(bytes) => skip_take_page::split_feed_envelope(bytes)?, + ResponseBody::NoPayload => Vec::new(), + }; + let survivors = self.select_survivors(&items)?; + let emitted = survivors.len(); + + // An all-duplicate intermediate page becomes a pull rather + // than an empty public page; its RU/diagnostics are retained. + if emitted == 0 && !is_terminal { + self.suppress(response); + continue; + } + + let new_response = self.rebuild(&response, survivors, emitted); + return Ok(PageResult::Page { + response: new_response, + is_terminal, + }); + } + } + } + } + + #[cfg(test)] + fn into_children(self) -> Vec> { + vec![self.child] + } + + fn snapshot_state(&self) -> crate::error::Result { + // A drained pipeline has no deduplication state left to lose, so even + // an unordered map can snapshot: resuming `Drained` re-emits nothing. + // Refusing here would break the common "page to completion, then + // persist the token" pattern on its very last iteration. + if self.exhausted { + return Ok(PipelineNodeState::Drained); + } + let last_hash = match &self.map { + DistinctMap::Ordered { last_hash } => *last_hash, + DistinctMap::Unordered { .. } => { + return Err(crate::error::CosmosError::builder() + .with_status(CosmosStatus::CLIENT_DISTINCT_CONTINUATION_UNSUPPORTED) + .with_message(UNORDERED_CONTINUATION_MESSAGE) + .build()); + } + }; + Ok(PipelineNodeState::Distinct { + distinct_type: DistinctType::Ordered, + last_hash, + child: Box::new(self.child.snapshot_state()?), + }) + } + + fn feed_range(&self) -> Option<&FeedRange> { + self.child.feed_range() + } + + fn topology_can_change(&self) -> bool { + // The wrapped fan-out node owns the partition ranges and handles its + // own splits, so a `Distinct` is safe as the pipeline root. + false + } + + fn fan_out_width(&self) -> usize { + // `Distinct` issues no request of its own. + self.child.fan_out_width() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::driver::dataflow::mocks::*; + use crate::driver::dataflow::node::SplitReplacements; + use crate::models::ResponseBody; + + /// Builds a query-page body from a list of raw JSON document texts. + fn page_body(documents: &[&str]) -> Vec { + format!( + r#"{{"_rid":"","Documents":[{}],"_count":{}}}"#, + documents.join(","), + documents.len() + ) + .into_bytes() + } + + fn page(documents: &[&str], is_terminal: bool) -> crate::error::Result { + Ok(PageResult::Page { + response: response(&page_body(documents)), + is_terminal, + }) + } + + fn charged_page( + documents: &[&str], + is_terminal: bool, + ru: f64, + ) -> crate::error::Result { + Ok(PageResult::Page { + response: response_with_charge(&page_body(documents), ru), + is_terminal, + }) + } + + fn documents_of(response: &CosmosResponse) -> Vec { + // `Distinct` emits pre-split `Items`, matching `StreamingOrderedMerge` + // and `SkipTake`, so a parent node never has to re-split the page. + match response.body() { + ResponseBody::Items(items) => items + .iter() + .map(|item| serde_json::from_slice(item).unwrap()) + .collect(), + ResponseBody::NoPayload => Vec::new(), + other => panic!("expected an Items body, got {other:?}"), + } + } + + async fn drain(node: &mut Distinct) -> Vec { + drain_with_charge(node).await.0 + } + + /// Drains `node`, returning the emitted documents plus the total request + /// charge reported across every emitted page. + async fn drain_with_charge(node: &mut Distinct) -> (Vec, f64) { + let mut executor = NoopRequestExecutor; + let mut topology = NoopTopologyProvider; + let mut context = PipelineContext::new(&mut executor, Some(&mut topology)); + let mut all = Vec::new(); + let mut charge = 0.0; + loop { + match node.next_page(&mut context).await.unwrap() { + PageResult::Page { response, .. } => { + all.extend(documents_of(&response)); + charge += response + .headers() + .request_charge + .map(|c| c.value()) + .unwrap_or_default(); + } + PageResult::Drained => break, + PageResult::SplitRequired { .. } => panic!("unexpected split"), + } + } + (all, charge) + } + + fn strings(values: &[serde_json::Value]) -> Vec { + values.iter().map(|v| v.to_string()).collect() + } + + // ── Unordered map ──────────────────────────────────────────────────── + + /// .NET `DistinctQueryPipelineStageTests.SanityTests`: duplicates that + /// straddle page boundaries must collapse. + #[tokio::test] + async fn unordered_dedupes_across_pages() { + let child = MockLeaf::with_pages(vec![ + page(&[r#"{"item":42}"#, r#"{"item":1337}"#], false), + page(&[r#"{"item":1337}"#, r#"{"item":42}"#], true), + Ok(PageResult::Drained), + ]); + let mut node = Distinct::new(Box::new(child), DistinctType::Unordered); + assert_eq!( + strings(&drain(&mut node).await), + vec![r#"{"item":42}"#, r#"{"item":1337}"#] + ); + } + + /// Java `DistinctQueryTests.queryDocuments`: the same value arriving from a + /// different partition is still one row. `MockLeaf` stands in for the + /// fan-out root, which has already interleaved partitions by this point. + #[tokio::test] + async fn unordered_dedupes_across_partitions() { + let child = MockLeaf::with_pages(vec![ + // Partition 0 contributed Seattle/Redmond; partition 1 repeats them. + page(&[r#""Seattle""#, r#""Redmond""#], false), + page(&[r#""Redmond""#, r#""Boston""#], true), + Ok(PageResult::Drained), + ]); + let mut node = Distinct::new(Box::new(child), DistinctType::Unordered); + assert_eq!( + strings(&drain(&mut node).await), + vec![r#""Seattle""#, r#""Redmond""#, r#""Boston""#] + ); + } + + #[tokio::test] + async fn unordered_dedupes_within_a_single_page() { + let child = MockLeaf::with_pages(vec![ + page(&[r#"1"#, r#"1"#, r#"2"#, r#"1"#], true), + Ok(PageResult::Drained), + ]); + let mut node = Distinct::new(Box::new(child), DistinctType::Unordered); + assert_eq!(strings(&drain(&mut node).await), vec!["1", "2"]); + } + + #[tokio::test] + async fn unordered_passes_through_when_nothing_repeats() { + let child = MockLeaf::with_pages(vec![ + page(&[r#"1"#, r#"2"#], false), + page(&[r#"3"#], true), + Ok(PageResult::Drained), + ]); + let mut node = Distinct::new(Box::new(child), DistinctType::Unordered); + assert_eq!(strings(&drain(&mut node).await), vec!["1", "2", "3"]); + } + + /// Structural equality, not textual: key order must not create a second row. + #[tokio::test] + async fn unordered_collapses_objects_differing_only_in_key_order() { + let child = MockLeaf::with_pages(vec![ + page( + &[ + r#"{"name":"fido","species":"dog"}"#, + r#"{"species":"dog","name":"fido"}"#, + r#"{"name":"fido","species":"cat"}"#, + ], + true, + ), + Ok(PageResult::Drained), + ]); + let mut node = Distinct::new(Box::new(child), DistinctType::Unordered); + assert_eq!(drain(&mut node).await.len(), 2); + } + + /// Arrays are position-sensitive, so these are two rows. + #[tokio::test] + async fn unordered_keeps_arrays_that_differ_only_in_order() { + let child = MockLeaf::with_pages(vec![ + page(&["[1,2]", "[2,1]", "[1,2]"], true), + Ok(PageResult::Drained), + ]); + let mut node = Distinct::new(Box::new(child), DistinctType::Unordered); + assert_eq!(strings(&drain(&mut node).await), vec!["[1,2]", "[2,1]"]); + } + + /// Java `queryDocumentsForDistinctIntValues`: `5` and `5.0` are one value. + #[tokio::test] + async fn unordered_collapses_equal_numbers_written_differently() { + let child = MockLeaf::with_pages(vec![ + page(&[r#"{"intprop":5}"#, r#"{"intprop":5.0}"#], true), + Ok(PageResult::Drained), + ]); + let mut node = Distinct::new(Box::new(child), DistinctType::Unordered); + assert_eq!(drain(&mut node).await.len(), 1); + } + + // ── Ordered map ────────────────────────────────────────────────────── + + #[tokio::test] + async fn ordered_dedupes_adjacent_runs() { + let child = MockLeaf::with_pages(vec![ + page(&[r#""Boston""#, r#""Redmond""#, r#""Redmond""#], false), + page(&[r#""Seattle""#, r#""Seattle""#], true), + Ok(PageResult::Drained), + ]); + let mut node = Distinct::new(Box::new(child), DistinctType::Ordered); + assert_eq!( + strings(&drain(&mut node).await), + vec![r#""Boston""#, r#""Redmond""#, r#""Seattle""#] + ); + } + + /// Documents the adjacency assumption: an ordered map deliberately does + /// *not* catch a non-adjacent repeat. This is only reachable when the merge + /// really is sorted, which the planner guarantees by only choosing + /// `Ordered` when the plan reports it. + #[tokio::test] + async fn ordered_does_not_dedupe_non_adjacent_repeats() { + let child = + MockLeaf::with_pages(vec![page(&["1", "2", "1"], true), Ok(PageResult::Drained)]); + let mut node = Distinct::new(Box::new(child), DistinctType::Ordered); + assert_eq!(strings(&drain(&mut node).await), vec!["1", "2", "1"]); + } + + /// A run split across a page boundary must not re-emit its head. + #[tokio::test] + async fn ordered_dedupes_a_run_spanning_pages() { + let child = MockLeaf::with_pages(vec![ + page(&[r#""Redmond""#], false), + page(&[r#""Redmond""#, r#""Seattle""#], true), + Ok(PageResult::Drained), + ]); + let mut node = Distinct::new(Box::new(child), DistinctType::Ordered); + assert_eq!( + strings(&drain(&mut node).await), + vec![r#""Redmond""#, r#""Seattle""#] + ); + } + + /// Seeding from a checkpoint suppresses the re-delivered boundary row. + #[tokio::test] + async fn ordered_resume_suppresses_the_reemitted_boundary_row() { + let boundary = hash_value(&serde_json::json!("Redmond")).unwrap(); + let child = MockLeaf::with_pages(vec![ + page(&[r#""Redmond""#, r#""Seattle""#], true), + Ok(PageResult::Drained), + ]); + let mut node = + Distinct::with_last_hash(Box::new(child), DistinctType::Ordered, Some(boundary)); + assert_eq!(strings(&drain(&mut node).await), vec![r#""Seattle""#]); + } + + // ── Page shaping / accounting ──────────────────────────────────────── + + #[tokio::test] + async fn item_count_reflects_the_trimmed_page() { + let child = + MockLeaf::with_pages(vec![page(&["1", "1", "2"], true), Ok(PageResult::Drained)]); + let mut node = Distinct::new(Box::new(child), DistinctType::Unordered); + let mut executor = NoopRequestExecutor; + let mut topology = NoopTopologyProvider; + let mut context = PipelineContext::new(&mut executor, Some(&mut topology)); + + match node.next_page(&mut context).await.unwrap() { + PageResult::Page { response, .. } => { + assert_eq!(response.headers().item_count, Some(2)); + assert_eq!(strings(&documents_of(&response)), vec!["1", "2"]); + } + other => panic!("expected a page, got {other:?}"), + } + } + + /// An all-duplicate intermediate page must not surface as an empty page, + /// and its RU must survive onto the next real page. + #[tokio::test] + async fn all_duplicate_intermediate_page_is_suppressed_but_charged() { + let child = MockLeaf::with_pages(vec![ + charged_page(&["1"], false, 2.0), + charged_page(&["1"], false, 3.0), + charged_page(&["2"], true, 5.0), + Ok(PageResult::Drained), + ]); + let mut node = Distinct::new(Box::new(child), DistinctType::Unordered); + let mut executor = NoopRequestExecutor; + let mut topology = NoopTopologyProvider; + let mut context = PipelineContext::new(&mut executor, Some(&mut topology)); + + let first = node.next_page(&mut context).await.unwrap(); + let PageResult::Page { response, .. } = first else { + panic!("expected a page"); + }; + assert_eq!(strings(&documents_of(&response)), vec!["1"]); + + // The all-duplicate second page is skipped; the third page carries both + // its own 5 RU and the suppressed 3 RU. + let second = node.next_page(&mut context).await.unwrap(); + let PageResult::Page { response, .. } = second else { + panic!("expected a page"); + }; + assert_eq!(strings(&documents_of(&response)), vec!["2"]); + assert_eq!(response.headers().request_charge.unwrap().value(), 8.0); + } + + /// A fully-duplicate *tail* still has to report its RU, so the node flushes + /// a final empty page rather than dropping the charge. + #[tokio::test] + async fn all_duplicate_tail_flushes_its_charge() { + let child = MockLeaf::with_pages(vec![ + charged_page(&["1"], false, 2.0), + charged_page(&["1"], false, 4.0), + Ok(PageResult::Drained), + ]); + let mut node = Distinct::new(Box::new(child), DistinctType::Unordered); + let mut executor = NoopRequestExecutor; + let mut topology = NoopTopologyProvider; + let mut context = PipelineContext::new(&mut executor, Some(&mut topology)); + + let _ = node.next_page(&mut context).await.unwrap(); + match node.next_page(&mut context).await.unwrap() { + PageResult::Page { + response, + is_terminal, + } => { + assert!(is_terminal); + assert!(documents_of(&response).is_empty()); + assert_eq!(response.headers().request_charge.unwrap().value(), 4.0); + } + other => panic!("expected a flushed empty page, got {other:?}"), + } + assert!(matches!( + node.next_page(&mut context).await.unwrap(), + PageResult::Drained + )); + } + + /// A terminal all-duplicate page is emitted as an empty terminal page (not + /// suppressed), so the caller sees the stream end. + #[tokio::test] + async fn terminal_all_duplicate_page_is_emitted_empty() { + let child = MockLeaf::with_pages(vec![ + page(&["1"], false), + page(&["1"], true), + Ok(PageResult::Drained), + ]); + let mut node = Distinct::new(Box::new(child), DistinctType::Unordered); + let mut executor = NoopRequestExecutor; + let mut topology = NoopTopologyProvider; + let mut context = PipelineContext::new(&mut executor, Some(&mut topology)); + + let _ = node.next_page(&mut context).await.unwrap(); + match node.next_page(&mut context).await.unwrap() { + PageResult::Page { + response, + is_terminal, + } => { + assert!(is_terminal); + assert!(documents_of(&response).is_empty()); + assert_eq!(response.headers().item_count, Some(0)); + } + other => panic!("expected an empty terminal page, got {other:?}"), + } + } + + #[tokio::test] + async fn empty_backend_page_is_handled() { + let child = MockLeaf::with_pages(vec![page(&[], true), Ok(PageResult::Drained)]); + let mut node = Distinct::new(Box::new(child), DistinctType::Unordered); + assert!(drain(&mut node).await.is_empty()); + } + + // ── Continuation ───────────────────────────────────────────────────── + + #[tokio::test] + async fn ordered_snapshot_carries_the_last_emitted_hash() { + let child = MockLeaf::with_pages(vec![page(&[r#""Redmond""#], false)]); + let mut node = Distinct::new(Box::new(child), DistinctType::Ordered); + let mut executor = NoopRequestExecutor; + let mut topology = NoopTopologyProvider; + let mut context = PipelineContext::new(&mut executor, Some(&mut topology)); + let _ = node.next_page(&mut context).await.unwrap(); + + match node.snapshot_state().unwrap() { + PipelineNodeState::Distinct { + distinct_type, + last_hash, + .. + } => { + assert_eq!(distinct_type, DistinctType::Ordered); + assert_eq!( + last_hash, + Some(hash_value(&serde_json::json!("Redmond")).unwrap()) + ); + } + other => panic!("expected a Distinct snapshot, got {other:?}"), + } + } + + #[test] + fn ordered_snapshot_before_any_row_has_no_hash() { + let child = MockLeaf::with_pages(vec![]); + let node = Distinct::new(Box::new(child), DistinctType::Ordered); + match node.snapshot_state().unwrap() { + PipelineNodeState::Distinct { last_hash, .. } => assert_eq!(last_hash, None), + other => panic!("expected a Distinct snapshot, got {other:?}"), + } + } + + #[tokio::test] + async fn ordered_snapshot_after_draining_is_drained() { + let child = MockLeaf::with_pages(vec![page(&["1"], true), Ok(PageResult::Drained)]); + let mut node = Distinct::new(Box::new(child), DistinctType::Ordered); + let _ = drain(&mut node).await; + assert!(matches!( + node.snapshot_state().unwrap(), + PipelineNodeState::Drained + )); + } + + /// The load-bearing continuation contract: an unordered `DISTINCT` must + /// fail loudly at token-mint time rather than hand back a token whose + /// resume would re-emit rows. + #[test] + fn unordered_snapshot_is_refused_with_actionable_guidance() { + let child = MockLeaf::with_pages(vec![]); + let node = Distinct::new(Box::new(child), DistinctType::Unordered); + let err = node + .snapshot_state() + .expect_err("an unordered DISTINCT must not produce a resumable snapshot"); + assert_eq!( + err.status().sub_status(), + Some(crate::error::SubStatusCode::CLIENT_DISTINCT_CONTINUATION_UNSUPPORTED), + ); + assert!( + err.to_string().contains("ORDER BY"), + "the message must tell the caller how to make the query resumable: {err}" + ); + } + + /// A drained pipeline has no deduplication state left to lose, so even an + /// unordered map can snapshot — resuming `Drained` re-emits nothing. This + /// is what makes the common "page to completion, then persist the token" + /// pattern work for an unordered `DISTINCT` query. + #[tokio::test] + async fn unordered_snapshot_is_allowed_once_drained() { + let child = MockLeaf::with_pages(vec![page(&["1"], true), Ok(PageResult::Drained)]); + let mut node = Distinct::new(Box::new(child), DistinctType::Unordered); + let _ = drain(&mut node).await; + assert!(matches!( + node.snapshot_state().unwrap(), + PipelineNodeState::Drained + )); + } + + // ── Errors ─────────────────────────────────────────────────────────── + + #[tokio::test] + async fn malformed_page_body_surfaces_a_typed_error() { + let child = MockLeaf::with_pages(vec![Ok(PageResult::Page { + response: response(b"not json"), + is_terminal: true, + })]); + let mut node = Distinct::new(Box::new(child), DistinctType::Unordered); + let mut executor = NoopRequestExecutor; + let mut topology = NoopTopologyProvider; + let mut context = PipelineContext::new(&mut executor, Some(&mut topology)); + assert!(node.next_page(&mut context).await.is_err()); + } + + #[tokio::test] + async fn forwarded_split_surfaces_a_typed_error() { + // `SplitRequired` replaces the node that emits it, so forwarding one + // would discard this node's deduplication map and let suppressed values + // reappear. Unreachable in production (the wrapped fan-out node absorbs + // splits), so this pins the guard rather than a live path. + let replacement = MockLeaf::with_pages(vec![Ok(PageResult::Drained)]); + let child = MockLeaf::with_pages(vec![Ok(PageResult::SplitRequired { + replacements: SplitReplacements::untiled(vec![Box::new(replacement)]), + })]); + let mut node = Distinct::new(Box::new(child), DistinctType::Unordered); + let mut executor = NoopRequestExecutor; + let mut topology = NoopTopologyProvider; + let mut context = PipelineContext::new(&mut executor, Some(&mut topology)); + + let err = node.next_page(&mut context).await.unwrap_err(); + assert_eq!( + err.status(), + CosmosStatus::CLIENT_DISTINCT_CANNOT_FORWARD_SPLIT + ); + } + + #[tokio::test] + async fn fan_out_width_matches_the_child() { + let child = MockLeaf::with_pages(vec![]); + let expected = child.fan_out_width(); + let node = Distinct::new(Box::new(child), DistinctType::Unordered); + assert_eq!(node.fan_out_width(), expected); + assert!(!node.topology_can_change()); + } + + // ── Catalog-driven scenarios ───────────────────────────────────────── + // + // Reuses `tests/fixtures/distinct_scenarios.json`; this file's copy of the + // fixture schema is minimal since separate compilation units can't share a + // `pub(crate)` type — see `tests/distinct_scenario_catalog.rs` for the + // canonical strict schema every layer trusts. + + #[derive(serde::Deserialize)] + struct CatalogFixture { + scenarios: Vec, + } + + #[derive(serde::Deserialize)] + struct ScenarioFixture { + id: String, + layers: Vec, + query: QueryFixture, + mock: Option, + #[serde(rename = "expectedValues", default)] + expected_values: Vec, + checkpoint: Option, + #[serde(rename = "expectedError")] + expected_error: Option, + } + + #[derive(serde::Deserialize)] + struct QueryFixture { + #[serde(rename = "distinctType")] + distinct_type: String, + } + + #[derive(serde::Deserialize)] + struct MockFixture { + partitions: Vec, + } + + #[derive(serde::Deserialize)] + struct PartitionFixture { + pages: Vec, + } + + #[derive(serde::Deserialize)] + struct PageFixture { + rows: Vec, + } + + #[derive(serde::Deserialize)] + struct RowFixture { + payload: serde_json::Value, + } + + fn fixture_distinct_type(text: &str) -> DistinctType { + match text { + "Ordered" => DistinctType::Ordered, + "Unordered" => DistinctType::Unordered, + "None" => DistinctType::None, + other => panic!("unknown distinctType in fixture: {other}"), + } + } + + /// Runs every catalog scenario tagged `mockPipeline` or `distinctMap` + /// through the real node. + /// + /// `SequentialDrain` and `StreamingOrderedMerge` both hand their parent a + /// flat sequence of pages, so the fixture's per-partition pages are + /// concatenated in partition order — exactly what `Distinct` observes in + /// production. Error scenarios need the planner (see `planner`'s + /// `peel_distinct_resume` tests and `integration_tests::distinct_resume`) + /// and are covered by dedicated tests instead. + #[tokio::test] + async fn catalog_scenarios_dedupe_as_expected() { + const CATALOG_JSON: &str = include_str!("../../../tests/fixtures/distinct_scenarios.json"); + let catalog: CatalogFixture = + serde_json::from_str(CATALOG_JSON).expect("catalog must parse"); + + let mut ran = 0usize; + let mut ran_a_resume_checkpoint = false; + for scenario in &catalog.scenarios { + if !scenario + .layers + .iter() + .any(|l| l == "mockPipeline" || l == "distinctMap") + { + continue; + } + if scenario.expected_error.is_some() { + continue; + } + let Some(mock) = &scenario.mock else { + continue; + }; + + let mut pages: Vec> = Vec::new(); + let page_specs: Vec<&PageFixture> = mock + .partitions + .iter() + .flat_map(|partition| partition.pages.iter()) + .collect(); + // Every page carries a distinct, non-round charge so the + // suppress/flush accounting cannot accidentally balance. + let mut expected_charge = 0.0; + for (index, page_spec) in page_specs.iter().enumerate() { + let documents: Vec = page_spec + .rows + .iter() + .map(|row| row.payload.to_string()) + .collect(); + let refs: Vec<&str> = documents.iter().map(String::as_str).collect(); + let ru = 1.25 + index as f64; + expected_charge += ru; + pages.push(charged_page(&refs, index + 1 == page_specs.len(), ru)); + } + pages.push(Ok(PageResult::Drained)); + + // A checkpoint seeds the ordered map with the hash of the last + // value emitted before the resume. + let last_hash = match scenario + .checkpoint + .as_ref() + .and_then(|c| c.get("lastValue")) + { + Some(value) => { + ran_a_resume_checkpoint = true; + Some(hash_value(value).expect("checkpoint value hashes")) + } + None => None, + }; + + let mut node = Distinct::with_last_hash( + Box::new(MockLeaf::with_pages(pages)), + fixture_distinct_type(&scenario.query.distinct_type), + last_hash, + ); + let (values, charge) = drain_with_charge(&mut node).await; + assert_eq!( + values, scenario.expected_values, + "scenario {} produced the wrong deduplicated stream", + scenario.id + ); + // Suppressing an all-duplicate page must never lose its RU: the + // charge the caller sees has to equal what the backend billed. + assert!( + (charge - expected_charge).abs() < 1e-9, + "scenario {}: emitted pages reported {charge} RU but the backend billed \ + {expected_charge}", + scenario.id + ); + ran += 1; + } + + // Exact, not a floor, so a scenario that stops matching the loop's + // filters shows up as a failure rather than a quietly smaller run. + // (A `mockPipeline` scenario with no `mock` is excluded from both + // sides here; `mock_pipeline_scenarios_carry_a_mock_or_expect_an_error` + // in the catalog test is what catches that.) + let expected: usize = catalog + .scenarios + .iter() + .filter(|s| { + s.layers + .iter() + .any(|l| l == "mockPipeline" || l == "distinctMap") + && s.expected_error.is_none() + && s.mock.is_some() + }) + .count(); + assert_eq!( + ran, expected, + "every eligible mock-pipeline scenario must run" + ); + assert!( + ran >= 10, + "expected the catalog to drive a meaningful number of mock-pipeline scenarios, ran {ran}" + ); + assert!( + ran_a_resume_checkpoint, + "no catalog scenario exercised an ordered resume checkpoint" + ); + } +} diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/distinct_hash.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/distinct_hash.rs new file mode 100644 index 00000000000..7ec6322fb77 --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/distinct_hash.rs @@ -0,0 +1,527 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +//! Structural, type-aware hashing of Cosmos JSON values. +//! +//! This is the deduplication identity used by cross-partition `DISTINCT` +//! ([`super::distinct`]), and is written to be reusable as a `GROUP BY` key +//! hash — nothing here knows about either stage. +//! +//! # Semantics +//! +//! Ported from .NET's `DistinctHash`, not Java's: Java funnels every scalar +//! through `ObjectOutputStream`, whose byte format is JVM-specific and not +//! reproducible outside a JVM. The *hash bytes* are therefore not compatible +//! with either peer, which is fine — a hash never leaves this process except +//! inside our own continuation token. The *behavior* matches both: +//! +//! - Every JSON type carries its own seed, so `null`, `false`, `""`, `[]`, +//! and `{}` can never collide with each other. +//! - Arrays are **position-sensitive**: element `i` is seeded with +//! `ARRAY_INDEX + i` and folded into a running chain, so `[1,2] != [2,1]`. +//! - Objects are **position-insensitive**: each property's hash is seeded by +//! its key's hash and the results are XOR-folded, so +//! `{"a":1,"b":2} == {"b":2,"a":1}`. XOR is safe here because a key is +//! unique within a well-formed object, so no two properties can cancel. +//! - Numbers compare by value, not representation: `5` and `5.0` hash +//! identically (see [`OrderByNumber`]), and `-0.0` is normalized to `0.0`. +//! - `undefined` is a value in its own right at the top level, distinct from +//! `null`. It cannot appear *inside* a container here: `serde_json` has no +//! `undefined`, and Cosmos drops undefined values from a result set rather +//! than emitting them, so the array/object walks never have one to skip. +//! (.NET's hasher skips undefined container members while still consuming +//! the array index; there is nothing to mirror until `GROUP BY` introduces +//! keys that can be undefined — see [`hash_undefined`].) +//! +//! # Recursion +//! +//! [`StructuralHasher::hash`] is depth-bounded ([`MAX_DEPTH`]) and returns a +//! typed error rather than overflowing the stack on a pathologically nested +//! payload. .NET relies on `EnsureSufficientExecutionStack` for the same +//! reason. + +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +use crate::models::murmur_hash::murmurhash3_128; + +use super::order_by::OrderByNumber; + +/// A 128-bit structural hash of a Cosmos JSON value. +/// +/// Serialized as a lowercase 32-character hex string so a continuation token +/// stays human-inspectable and round-trips exactly. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] +pub(crate) struct Hash128(u128); + +impl Hash128 { + #[cfg(test)] + pub(crate) fn raw(self) -> u128 { + self.0 + } +} + +impl std::fmt::Display for Hash128 { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{:032x}", self.0) + } +} + +impl Serialize for Hash128 { + fn serialize(&self, serializer: S) -> Result { + serializer.serialize_str(&self.to_string()) + } +} + +impl<'de> Deserialize<'de> for Hash128 { + fn deserialize>(deserializer: D) -> Result { + let text = >::deserialize(deserializer)?; + // Exactly 32 hex digits, the width `Serialize` emits (which is always + // lowercase; either case parses to the same value). A `Hash128` only + // ever arrives from a token we minted, so a short, over-long, or + // sign-prefixed form is hand-crafted and is rejected rather than + // reinterpreted — the stance `planner::peel_distinct_resume` takes on + // token shape. + if text.len() != 32 || !text.bytes().all(|b| b.is_ascii_hexdigit()) { + return Err(serde::de::Error::custom(format!( + "invalid 128-bit hash {text:?}: expected 32 hex digits" + ))); + } + u128::from_str_radix(&text, 16) + .map(Hash128) + .map_err(|_| serde::de::Error::custom(format!("invalid 128-bit hash {text:?}"))) + } +} + +/// Maximum JSON nesting depth accepted by [`StructuralHasher::hash`]. +/// +/// Cosmos itself caps document nesting well below this, so exceeding it means +/// a hand-crafted or corrupt payload rather than real data. +const MAX_DEPTH: usize = 128; + +/// Root seed, and one seed per JSON type so structurally different values with +/// identical raw bytes never collide. Values are .NET's `DistinctHash.HashSeeds` +/// verbatim, laid out as `(low, high)` halves of a `u128`. +mod seeds { + const fn seed(low: u64, high: u64) -> u128 { + ((high as u128) << 64) | (low as u128) + } + + pub(super) const ROOT: u128 = seed(0xbfc2359eafc0e2b7, 0x8846e00284c4cf1f); + /// Only reachable through `super::hash_undefined`; see its note. + #[cfg(test)] + pub(super) const UNDEFINED: u128 = seed(0x5d1f8b2a91c46e03, 0x2a7f04b6cd39e1a8); + pub(super) const NULL: u128 = seed(0x1380f68bb3b0cfe4, 0x156c918bf564ee48); + pub(super) const FALSE: u128 = seed(0xc1be517fe893b40c, 0xe9fc8a4c531cd0dd); + pub(super) const TRUE: u128 = seed(0xf86d4abf9a412e74, 0x788488365c8a985d); + pub(super) const STRING: u128 = seed(0x61f53f0a44204cfb, 0x09481be8ef4b56dd); + pub(super) const NUMBER: u128 = seed(0x2400e8b894ce9c2a, 0x790be1eabd7b9481); + pub(super) const ARRAY: u128 = seed(0xfa573b014c4dc18e, 0xa014512c858eb115); + pub(super) const OBJECT: u128 = seed(0x77b285ac511aef30, 0x3dcf187245822449); + pub(super) const ARRAY_INDEX: u128 = seed(0xfe057204216db999, 0x5b1cc3178bd9c593); + pub(super) const PROPERTY_NAME: u128 = seed(0xc915dde058492a8a, 0x7c8be2eba72e4634); +} + +/// Computes the structural hash of a JSON value using the canonical root seed. +/// +/// `undefined` has no `serde_json` representation, so a caller that can +/// observe it (a missing projection) should use [`hash_undefined`] instead. +pub(crate) fn hash_value(value: &Value) -> crate::error::Result { + StructuralHasher.hash(value, seeds::ROOT, 0).map(Hash128) +} + +/// The hash of `undefined` — the value a projection yields when the path does +/// not exist. Deliberately distinct from the hash of `null`. +/// +/// Cosmos drops `undefined` from a result set rather than emitting it, so no +/// `DISTINCT` payload can carry one today; this exists so the semantics are +/// pinned by a test before `GROUP BY`, whose keys *can* be undefined, starts +/// calling it. +#[cfg(test)] +pub(crate) fn hash_undefined() -> Hash128 { + Hash128(murmurhash3_128( + &seeds::UNDEFINED.to_le_bytes(), + seeds::ROOT, + )) +} + +struct StructuralHasher; + +impl StructuralHasher { + fn hash(&self, value: &Value, seed: u128, depth: usize) -> crate::error::Result { + if depth > MAX_DEPTH { + return Err(crate::error::CosmosError::builder() + .with_status(crate::error::CosmosStatus::CLIENT_DISTINCT_VALUE_TOO_DEEPLY_NESTED) + .with_message(format!( + "value nests deeper than the {MAX_DEPTH}-level limit for DISTINCT hashing" + )) + .build()); + } + + Ok(match value { + Value::Null => hash_u128(seeds::NULL, seed), + Value::Bool(false) => hash_u128(seeds::FALSE, seed), + Value::Bool(true) => hash_u128(seeds::TRUE, seed), + Value::String(s) => { + let hash = hash_u128(seeds::STRING, seed); + murmurhash3_128(s.as_bytes(), hash) + } + Value::Number(n) => { + let hash = hash_u128(seeds::NUMBER, seed); + murmurhash3_128(&number_bytes(n), hash) + } + Value::Array(items) => { + let mut hash = hash_u128(seeds::ARRAY, seed); + for (index, item) in items.iter().enumerate() { + // `serde_json` has no `undefined`, so unlike .NET there is + // never an item to skip here — the index is always the + // element's real position either way. + let item_seed = seeds::ARRAY_INDEX.wrapping_add(index as u128); + let item_hash = self.hash(item, item_seed, depth + 1)?; + hash = hash_u128(item_hash, hash); + } + hash + } + Value::Object(map) => { + let hash = hash_u128(seeds::OBJECT, seed); + let mut folded: u128 = 0; + for (key, property) in map { + let name_hash = murmurhash3_128( + key.as_bytes(), + hash_u128(seeds::STRING, seeds::PROPERTY_NAME), + ); + folded ^= self.hash(property, name_hash, depth + 1)?; + } + // An empty object (and one whose properties happen to XOR to + // zero) keeps just the type seed, matching .NET. + if folded != 0 { + hash_u128(folded, hash) + } else { + hash + } + } + }) + } +} + +/// Hashes a 128-bit value as its little-endian bytes under `seed`. +fn hash_u128(value: u128, seed: u128) -> u128 { + murmurhash3_128(&value.to_le_bytes(), seed) +} + +/// Canonical byte encoding of a JSON number, chosen so numerically equal +/// values encode identically regardless of how they were written. +/// +/// Every integral value — signed, unsigned, or a float that happens to be +/// whole — widens to `i128`, which covers the full `i64` *and* `u64` ranges +/// without overlap. Encoding `i64` and `u64` separately as 8 two's-complement +/// bytes would make `-1` and `u64::MAX` share a byte pattern, silently +/// collapsing two distinct values; widening also makes `1e19` and +/// `10000000000000000000` agree, which 8-byte encodings could not. +/// +/// Everything else encodes as IEEE-754 bits with `-0.0` normalized to `0.0`. +fn number_bytes(number: &serde_json::Number) -> [u8; 17] { + const INTEGER_TAG: u8 = 0; + const FLOAT_TAG: u8 = 1; + + let mut out = [0u8; 17]; + let integral: Option = match OrderByNumber::from_json_number(number) { + OrderByNumber::I64(i) => Some(i as i128), + OrderByNumber::U64(u) => Some(u as i128), + OrderByNumber::F64(f) => integral_f64_as_i128(f), + }; + match integral { + Some(value) => { + out[0] = INTEGER_TAG; + out[1..].copy_from_slice(&value.to_le_bytes()); + } + None => { + let OrderByNumber::F64(f) = OrderByNumber::from_json_number(number) else { + unreachable!("only the float variant can fail the integral conversion"); + }; + // `-0.0 == 0.0`, so normalize before touching the bit pattern. + let f = if f == 0.0 { 0.0 } else { f }; + out[0] = FLOAT_TAG; + out[1..9].copy_from_slice(&f.to_bits().to_le_bytes()); + } + } + out +} + +/// Returns `f` as an `i128` when it is exactly integral and in range, so a +/// float-encoded whole number shares the integer encoding. +fn integral_f64_as_i128(f: f64) -> Option { + const TWO_POW_127: f64 = 170_141_183_460_469_231_731_687_303_715_884_105_728.0; + if !f.is_finite() || f.fract() != 0.0 || f >= TWO_POW_127 || f < -TWO_POW_127 { + return None; + } + Some(f as i128) +} + +#[cfg(test)] +mod tests { + use std::cmp::Ordering; + + use super::*; + use serde_json::json; + + fn h(value: serde_json::Value) -> Hash128 { + hash_value(&value).expect("value hashes") + } + + // ── Type coverage ──────────────────────────────────────────────────── + // + // Every JSON type must be self-consistent and mutually distinct. Mirrors + // .NET `DistinctHashBaselineTests.ElementsHash` / `NumbersHash` and Java + // `DistinctHashTest.{nullHash,booleanHash,integerHash,longHash,doubleHash,stringHash}`. + + #[test] + fn each_type_hashes_deterministically() { + for value in [ + json!(null), + json!(true), + json!(false), + json!(""), + json!("hello"), + json!(0), + json!(-1), + json!(3.5), + json!([]), + json!({}), + json!([1, 2, 3]), + json!({"a": 1}), + ] { + assert_eq!( + h(value.clone()), + h(value.clone()), + "unstable hash for {value}" + ); + } + } + + #[test] + fn distinct_types_do_not_collide() { + // `null`, `undefined`, `false`, `""`, `[]`, and `{}` are all "empty" + // in some sense and are the classic collision trap. + let hashes = [ + h(json!(null)), + hash_undefined(), + h(json!(false)), + h(json!(true)), + h(json!(0)), + h(json!("")), + h(json!([])), + h(json!({})), + ]; + for (i, a) in hashes.iter().enumerate() { + for (j, b) in hashes.iter().enumerate() { + if i != j { + assert_ne!(a, b, "types at {i} and {j} collide"); + } + } + } + } + + #[test] + fn undefined_is_not_null() { + assert_ne!(hash_undefined(), h(json!(null))); + } + + // ── Numeric equality ───────────────────────────────────────────────── + + /// Java `DistinctQueryTests.queryDocumentsForDistinctIntValues` asserts a + /// document with `intprop: 5` and one with `intprop: 5.0` dedupe together. + #[test] + fn integer_and_float_forms_of_the_same_value_match() { + assert_eq!(h(json!(5)), h(json!(5.0))); + assert_eq!(h(json!(-7)), h(json!(-7.0))); + } + + /// .NET normalizes `-0.0` to `0.0` in `CosmosNumberHasher`, but neither + /// peer has a test for it. + #[test] + fn negative_zero_matches_positive_zero() { + assert_eq!(h(json!(-0.0)), h(json!(0.0))); + assert_eq!(h(json!(-0.0)), h(json!(0))); + } + + #[test] + fn different_numbers_differ() { + assert_ne!(h(json!(5)), h(json!(6))); + assert_ne!(h(json!(3.5)), h(json!(3.75))); + } + + #[test] + fn large_integers_beyond_f64_precision_stay_distinct() { + // 2^53 and 2^53+1 are the same `f64`; the lossless integer encoding + // must keep them apart. + assert_ne!( + h(json!(9_007_199_254_740_992i64)), + h(json!(9_007_199_254_740_993i64)) + ); + } + + /// Regression: encoding `i64` and `u64` as bare 8-byte two's-complement + /// patterns made `-1` and `u64::MAX` share an encoding, so DISTINCT would + /// silently drop one of them. + #[test] + fn negative_integers_do_not_collide_with_large_unsigned_ones() { + assert_ne!(h(json!(-1i64)), h(json!(u64::MAX))); + assert_ne!(h(json!(i64::MIN)), h(json!(9_223_372_036_854_775_808u64))); + } + + /// A whole number above `2^63` must hash the same whether it was written + /// as an integer or in float form. + #[test] + fn integral_values_above_the_i64_range_still_compare_by_value() { + assert_eq!(h(json!(10_000_000_000_000_000_000u64)), h(json!(1e19))); + } + + #[test] + fn very_large_floats_stay_distinct() { + assert_ne!(h(json!(1.0e308)), h(json!(1.0e307))); + } + + // ── Structural equality ────────────────────────────────────────────── + + /// Java `DistinctHashTest.arrayNodeHash` / `listHash`. + #[test] + fn array_order_matters() { + assert_ne!(h(json!([1, 2])), h(json!([2, 1]))); + assert_eq!(h(json!([1, 2])), h(json!([1, 2]))); + } + + /// Java `DistinctHashTest.jsonObjectHash`. Untested in .NET, though its + /// XOR fold is designed for exactly this. + #[test] + fn object_key_order_does_not_matter() { + let a: Value = serde_json::from_str(r#"{"a":1,"b":2}"#).unwrap(); + let b: Value = serde_json::from_str(r#"{"b":2,"a":1}"#).unwrap(); + assert_eq!(h(a), h(b)); + } + + #[test] + fn object_value_change_changes_the_hash() { + assert_ne!(h(json!({"a": 1})), h(json!({"a": 2}))); + assert_ne!(h(json!({"a": 1})), h(json!({"b": 1}))); + } + + /// .NET `DistinctHashBaselineTests.WrappedElementsHash`: a value, that + /// value in an array, and that value in an object must all differ. + #[test] + fn wrapping_a_value_changes_the_hash() { + for value in [ + json!(null), + json!(true), + json!(42), + json!("x"), + json!([]), + json!({}), + ] { + let bare = h(value.clone()); + let in_array = h(json!([value.clone()])); + let in_object = h(json!({"prop": value.clone()})); + assert_ne!(bare, in_array, "bare vs array for {value}"); + assert_ne!(bare, in_object, "bare vs object for {value}"); + assert_ne!(in_array, in_object, "array vs object for {value}"); + } + } + + #[test] + fn nested_structures_are_distinguished() { + assert_ne!(h(json!({"a": {"b": 1}})), h(json!({"a": {"b": 2}}))); + assert_ne!(h(json!([[1], [2]])), h(json!([[2], [1]]))); + assert_eq!( + h(json!({"a": [1, {"b": null}]})), + h(json!({"a": [1, {"b": null}]})) + ); + } + + #[test] + fn empty_containers_differ_from_populated_ones() { + assert_ne!(h(json!([])), h(json!([null]))); + assert_ne!(h(json!({})), h(json!({"a": null}))); + } + + // ── Strings ────────────────────────────────────────────────────────── + // + // .NET `DistinctQueryPipelineStageTests.MixedTypeTests` covers CJK values + // and Arabic object keys. + + #[test] + fn unicode_strings_are_handled() { + assert_eq!( + h(json!("敏捷的棕色狐狸跳过了懒狗")), + h(json!("敏捷的棕色狐狸跳过了懒狗")) + ); + assert_ne!( + h(json!("敏捷的棕色狐狸跳过了懒狗")), + h(json!("敏捷的棕色狐狸跳过了懶狗")) + ); + assert_ne!(h(json!({"فوق": 1})), h(json!({"تحت": 1}))); + } + + #[test] + fn long_strings_past_the_inline_threshold_are_distinguished() { + let a = "x".repeat(4096); + let b = format!("{}y", "x".repeat(4095)); + assert_ne!(h(json!(a)), h(json!(b))); + } + + // ── Recursion guard ────────────────────────────────────────────────── + + #[test] + fn excessive_nesting_is_a_typed_error_not_a_stack_overflow() { + let mut value = json!(1); + for _ in 0..(MAX_DEPTH + 10) { + value = Value::Array(vec![value]); + } + let err = hash_value(&value).expect_err("must reject over-deep values"); + assert_eq!( + err.status().sub_status(), + Some(crate::error::SubStatusCode::CLIENT_DISTINCT_VALUE_TOO_DEEPLY_NESTED), + ); + } + + #[test] + fn nesting_at_the_limit_is_accepted() { + let mut value = json!(1); + for _ in 0..(MAX_DEPTH - 1) { + value = Value::Array(vec![value]); + } + assert!(hash_value(&value).is_ok()); + } + + // ── Serialization ──────────────────────────────────────────────────── + + #[test] + fn hash_round_trips_through_json() { + let hash = h(json!({"city": "Seattle"})); + let text = serde_json::to_string(&hash).unwrap(); + assert_eq!( + text.len(), + 34, + "expected a quoted 32-char hex string, got {text}" + ); + let restored: Hash128 = serde_json::from_str(&text).unwrap(); + assert_eq!(hash, restored); + } + + #[test] + fn malformed_hash_text_is_rejected() { + assert!(serde_json::from_str::(r#""not-hex""#).is_err()); + // Short, over-long, and sign-prefixed forms are all rejected rather + // than silently reinterpreted. + assert!(serde_json::from_str::(r#""5""#).is_err()); + assert!(serde_json::from_str::(r#""+0000000000000000000000000000005""#).is_err()); + assert!(serde_json::from_str::(&format!(r#""{}""#, "0".repeat(33))).is_err()); + } + + #[test] + fn ordering_is_total() { + let a = h(json!(1)); + let b = h(json!(2)); + assert_eq!(a.raw().cmp(&a.raw()), Ordering::Equal); + assert_eq!(a.cmp(&b), a.raw().cmp(&b.raw())); + } +} diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/mocks.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/mocks.rs index a3966e94dfb..ad28a872ab4 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/mocks.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/mocks.rs @@ -323,23 +323,6 @@ pub(crate) fn response_with_continuation( ) } -/// Creates a test response carrying a specific request charge (RU) header. -pub(crate) fn response_with_charge(body: &[u8], request_charge: f64) -> CosmosResponse { - let mut diagnostics = DiagnosticsContextBuilder::new( - ActivityId::new_uuid(), - Arc::new(DiagnosticsOptions::default()), - ); - diagnostics.set_operation_status(StatusCode::Ok, None); - let mut headers = CosmosResponseHeaders::new(); - headers.request_charge = Some(crate::models::RequestCharge::new(request_charge)); - CosmosResponse::new( - body.to_vec(), - headers, - CosmosStatus::new(StatusCode::Ok), - Arc::new(diagnostics.complete()), - ) -} - /// Creates a test response whose diagnostics carry `requests` per-attempt /// records, so tests can assert on attempt counts surviving aggregation. pub(crate) fn response_with_request_diagnostics(requests: usize) -> CosmosResponse { @@ -392,6 +375,24 @@ pub(crate) fn response_with_etag(body: &[u8], etag: &str) -> CosmosResponse { ) } +/// Creates a test response with the given body and request charge, for nodes +/// that fold RU across suppressed pages. +pub(crate) fn response_with_charge(body: &[u8], request_charge: f64) -> CosmosResponse { + let mut diagnostics = DiagnosticsContextBuilder::new( + ActivityId::new_uuid(), + Arc::new(DiagnosticsOptions::default()), + ); + diagnostics.set_operation_status(StatusCode::Ok, None); + let mut headers = CosmosResponseHeaders::new(); + headers.request_charge = Some(crate::models::RequestCharge::new(request_charge)); + CosmosResponse::new( + body.to_vec(), + headers, + CosmosStatus::new(StatusCode::Ok), + Arc::new(diagnostics.complete()), + ) +} + /// Creates a 410 Gone error with a partition topology change substatus. pub(crate) fn gone_error() -> crate::error::CosmosError { crate::error::CosmosError::builder() diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/mod.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/mod.rs index aed14213244..db18da8ccc5 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/mod.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/mod.rs @@ -35,6 +35,8 @@ //! cross-partition strategies). mod context; +mod distinct; +pub(crate) mod distinct_hash; mod drain; mod drained; #[cfg(test)] @@ -58,6 +60,7 @@ mod unordered_merge; pub(crate) use context::{ PartitionRoutingRefresh, PipelineContext, RequestExecutor, ResolvedRange, TopologyProvider, }; +pub(crate) use distinct::Distinct; pub(crate) use drain::SequentialDrain; pub(crate) use drained::DrainedLeaf; pub(crate) use node::{ diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/order_by.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/order_by.rs index 54d60486b24..32ad2960fcc 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/order_by.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/order_by.rs @@ -163,7 +163,7 @@ pub(crate) enum OrderByNumber { impl OrderByNumber { /// Classifies a `serde_json::Number` into the widest lossless variant: /// signed integer, else unsigned integer, else finite float. - fn from_json_number(n: &serde_json::Number) -> Self { + pub(crate) fn from_json_number(n: &serde_json::Number) -> Self { if let Some(i) = n.as_i64() { Self::I64(i) } else if let Some(u) = n.as_u64() { diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/pipeline.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/pipeline.rs index bc573553598..03443bc12d4 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/pipeline.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/pipeline.rs @@ -56,9 +56,10 @@ impl Pipeline { PageResult::Page { response, .. } => Ok(Some(response)), PageResult::Drained => Ok(None), // Defensive: today the root is always a `Request`, `SequentialDrain`, - // or `DrainedLeaf`, none of which can bubble `SplitRequired` up past - // their parent. If a future node type ever does, surfacing it as an - // explicit error is preferable to silently dropping the page. + // `Distinct`, or `DrainedLeaf`, none of which can bubble + // `SplitRequired` up past their parent. If a future node type ever + // does, surfacing it as an explicit error is preferable to silently + // dropping the page. PageResult::SplitRequired { .. } => Err(crate::error::CosmosError::builder() .with_status(crate::error::CosmosStatus::CLIENT_ROOT_NODE_CANNOT_REQUEST_SPLIT) .with_message( diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/planner.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/planner.rs index 7a8acab2824..192a46a02c7 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/planner.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/planner.rs @@ -24,12 +24,13 @@ use crate::{ }; use super::{ + distinct_hash::Hash128, intersect_feed_ranges, query_plan::{QueryInfo, QueryPlan, SortOrder}, query_response, snapshot::{OrderByRangeToken, ValueBoundary}, - streaming_ordered_merge, DrainedLeaf, OperationPlan, PartitionRoutingRefresh, Pipeline, - PipelineNode, PipelineNodeState, RangedToken, Request, RequestTarget, ResolvedRange, + streaming_ordered_merge, Distinct, DrainedLeaf, OperationPlan, PartitionRoutingRefresh, + Pipeline, PipelineNode, PipelineNodeState, RangedToken, Request, RequestTarget, ResolvedRange, SequentialDrain, SkipTake, StreamingOrderedMerge, TopologyProvider, UnorderedMerge, }; @@ -194,6 +195,15 @@ pub(crate) async fn build_sequential_drain( topology_provider: &mut dyn TopologyProvider, operation: &Arc, resume: Option, +) -> crate::error::Result { + build_sequential_drain_inner(query_plan, topology_provider, operation, resume).await +} + +async fn build_sequential_drain_inner( + query_plan: &QueryPlan, + topology_provider: &mut dyn TopologyProvider, + operation: &Arc, + resume: Option, ) -> crate::error::Result { validate_query_plan(query_plan)?; @@ -238,6 +248,13 @@ pub(crate) async fn build_sequential_drain( other => other, }; + // `DISTINCT` sits *inside* the skip/take window (SQL applies OFFSET / + // LIMIT / TOP to the deduplicated stream), so its token state nests one + // level below `SkipTake`'s and is peeled second. + let distinct_type = plan_distinct_type(query_plan); + let (inner_resume, last_hash) = peel_distinct_resume(inner_resume, distinct_type)?; + let resumed_drained = matches!(inner_resume, Some(PipelineNodeState::Drained)); + let needs_skip_take = skip > 0 || take.is_some(); // Per-partition requests must use the plan's `rewrittenQuery` so @@ -297,13 +314,17 @@ pub(crate) async fn build_sequential_drain( // the single Request with multiple Requests. let fanout: Box = Box::new(SequentialDrain::new(request_nodes)); + // `DISTINCT` deduplicates the fan-out stream first; the global skip/take + // window then counts deduplicated rows, matching SQL semantics. + let deduped = apply_distinct(fanout, distinct_type, last_hash, resumed_drained); + // Cross-partition OFFSET / LIMIT / TOP applies a global skip/take over the // fan-out's EPK-ordered stream. When none is present the fan-out is the // pipeline root directly. let root: Box = if needs_skip_take { - Box::new(SkipTake::new(fanout, skip, take)) + Box::new(SkipTake::new(deduped, skip, take)) } else { - fanout + deduped }; Ok(Pipeline::new(root)) } @@ -328,6 +349,15 @@ pub(crate) async fn build_streaming_ordered_merge( topology_provider: &mut dyn TopologyProvider, operation: &Arc, resume: Option, +) -> crate::error::Result { + build_streaming_ordered_merge_inner(query_plan, topology_provider, operation, resume).await +} + +async fn build_streaming_ordered_merge_inner( + query_plan: &QueryPlan, + topology_provider: &mut dyn TopologyProvider, + operation: &Arc, + resume: Option, ) -> crate::error::Result { validate_query_plan_for_streaming_order_by(query_plan)?; let info = query_plan @@ -390,6 +420,13 @@ pub(crate) async fn build_streaming_ordered_merge( other => other, }; + // `DISTINCT` sits *inside* the skip/take window (SQL applies OFFSET / + // LIMIT / TOP to the deduplicated stream), so its token state nests one + // level below `SkipTake`'s and is peeled second. + let distinct_type = plan_distinct_type(query_plan); + let (resume, last_hash) = peel_distinct_resume(resume, distinct_type)?; + let resumed_drained = matches!(resume, Some(PipelineNodeState::Drained)); + let query_from_beginning = query_response::rewritten_query_from_beginning(rewritten_query)?; let plain_body = query_response::rewrite_query_body(operation.body(), &query_from_beginning)?; let plain_operation = Arc::new((**operation).clone().with_body(plain_body)); @@ -514,13 +551,17 @@ pub(crate) async fn build_streaming_ordered_merge( query_fingerprint, )); + // `DISTINCT` deduplicates the ordered stream first; the global skip/take + // window then counts deduplicated rows, matching SQL semantics. + let deduped = apply_distinct(ordered_root, distinct_type, last_hash, resumed_drained); + // Apply the global OFFSET / LIMIT / TOP window over the ordered stream. When // the query carries none, the ordered merge is the pipeline root directly. let needs_skip_take = skip > 0 || take.is_some(); let root: Box = if needs_skip_take { - Box::new(SkipTake::new(ordered_root, skip, take)) + Box::new(SkipTake::new(deduped, skip, take)) } else { - ordered_root + deduped }; Ok(Pipeline::new(root)) } @@ -1271,6 +1312,7 @@ fn snapshot_kind(state: &PipelineNodeState) -> &'static str { PipelineNodeState::UnorderedMerge { .. } => "UnorderedMerge", PipelineNodeState::SkipTake { .. } => "SkipTake", PipelineNodeState::StreamingOrderedMerge { .. } => "StreamingOrderedMerge", + PipelineNodeState::Distinct { .. } => "Distinct", } } @@ -1324,6 +1366,113 @@ fn validate_unordered_merge_tokens( Ok(parsed) } +/// The kind of deduplication this query plan asks for, or +/// [`DistinctType::None`] when it asks for none. +fn plan_distinct_type(plan: &QueryPlan) -> DistinctType { + plan.query_info + .as_ref() + .map(|info| info.distinct_type) + .unwrap_or_default() +} + +/// Wraps `pipeline`'s root in a [`Distinct`] stage when the plan calls for +/// deduplication, leaving it untouched otherwise. +/// +/// `DISTINCT` composes *above* the fan-out root, matching .NET's +/// `PipelineFactory` and Java's `PipelinedDocumentQueryExecutionContext`: +/// merge/`ORDER BY` -> aggregate -> **DISTINCT** -> `GROUP BY` -> +/// `OFFSET`/`LIMIT`/`TOP`. +/// Wraps `node` in a [`Distinct`] stage when the plan calls for one. +/// +/// `DISTINCT` deduplicates *before* any `OFFSET` / `LIMIT` / `TOP` window is +/// applied, so callers must wrap the fan-out with this first and only then +/// apply [`SkipTake`]; otherwise the window would count duplicate rows. +fn apply_distinct( + node: Box, + distinct_type: DistinctType, + last_hash: Option, + resumed_drained: bool, +) -> Box { + if distinct_type == DistinctType::None { + return node; + } + // A fully-drained resume needs no deduplication stage: the inner pipeline + // is a `DrainedLeaf` and will emit nothing. Wrapping it would leave a + // `Distinct` whose `exhausted` is still `false`, so an unordered query + // would refuse to re-snapshot a token it had just accepted. + if resumed_drained { + return node; + } + Box::new(Distinct::with_last_hash(node, distinct_type, last_hash)) +} + +/// Splits a resume state into the inner (fan-out) state and the `DISTINCT` +/// stage's saved `last_hash`, rejecting any token whose shape or distinct kind +/// does not match the current query plan. +/// +/// A token minted before DISTINCT support (or for a non-DISTINCT query) has no +/// `Distinct` layer; resuming it into a DISTINCT plan would silently skip +/// deduplication, so it is rejected rather than reinterpreted. The mirror case +/// — a `Distinct` token resumed into a non-DISTINCT plan — falls through to the +/// inner builders, which reject the unexpected shape. +fn peel_distinct_resume( + resume: Option, + distinct_type: DistinctType, +) -> crate::error::Result<(Option, Option)> { + match resume { + Some(PipelineNodeState::Distinct { + distinct_type: saved, + last_hash, + child, + }) => { + if saved != distinct_type { + return Err(distinct_token_mismatch(saved, distinct_type)); + } + if saved != DistinctType::Ordered { + // We never mint one, so this is a hand-crafted or corrupted + // token. Resuming would re-emit every value seen before the + // checkpoint. + return Err(crate::error::CosmosError::builder() + .with_status( + crate::error::CosmosStatus::CLIENT_DISTINCT_CONTINUATION_UNSUPPORTED, + ) + .with_message( + "continuation token carries unordered DISTINCT state, which cannot be \ + resumed; add a matching ORDER BY to make the query resumable", + ) + .build()); + } + Ok((Some(*child), last_hash)) + } + // `Drained` is shape-agnostic: the whole pipeline, DISTINCT included, + // finished. + Some(PipelineNodeState::Drained) => Ok((Some(PipelineNodeState::Drained), None)), + Some(other) if distinct_type != DistinctType::None => { + Err(crate::error::CosmosError::builder() + .with_status(crate::error::CosmosStatus::CLIENT_CONTINUATION_TOKEN_SHAPE_MISMATCH) + .with_message(format!( + "continuation token shape {} does not match a DISTINCT query", + snapshot_kind(&other) + )) + .build()) + } + other => Ok((other, None)), + } +} + +fn distinct_token_mismatch( + saved: DistinctType, + expected: DistinctType, +) -> crate::error::CosmosError { + crate::error::CosmosError::builder() + .with_status(crate::error::CosmosStatus::CLIENT_CONTINUATION_TOKEN_SHAPE_MISMATCH) + .with_message(format!( + "continuation token was minted for {saved:?} DISTINCT but the query plan reports \ + {expected:?}" + )) + .build() +} + /// Validates that the query plan does not require features we don't yet support. fn validate_query_plan(plan: &QueryPlan) -> crate::error::Result<()> { if plan.hybrid_search_query_info.is_some() { @@ -1347,9 +1496,6 @@ fn validate_query_info(info: &QueryInfo) -> crate::error::Result<()> { if !info.group_by_expressions.is_empty() { return Err(unsupported_feature("GROUP BY in cross-partition queries")); } - if info.distinct_type != DistinctType::None { - return Err(unsupported_feature("DISTINCT in cross-partition queries")); - } Ok(()) } @@ -1463,11 +1609,6 @@ fn validate_query_plan_for_streaming_order_by(plan: &QueryPlan) -> crate::error: "GROUP BY combined with ORDER BY in cross-partition queries", )); } - if info.distinct_type != DistinctType::None { - return Err(unsupported_feature( - "DISTINCT combined with ORDER BY in cross-partition queries", - )); - } Ok(()) } @@ -3487,7 +3628,7 @@ mod tests { } #[test] - fn validate_query_plan_for_streaming_order_by_rejects_aggregates_group_by_distinct() { + fn validate_query_plan_for_streaming_order_by_rejects_aggregates_and_group_by() { let mut plan = order_by_plan(Some("SELECT 1"), vec![qr("", "FF")]); plan.query_info.as_mut().unwrap().aggregates = vec!["Count".to_owned()]; assert!(validate_query_plan_for_streaming_order_by(&plan).is_err()); @@ -3495,10 +3636,176 @@ mod tests { let mut plan = order_by_plan(Some("SELECT 1"), vec![qr("", "FF")]); plan.query_info.as_mut().unwrap().group_by_expressions = vec!["c.a".to_owned()]; assert!(validate_query_plan_for_streaming_order_by(&plan).is_err()); + } + + /// DISTINCT is now composed as a stage above the merge rather than + /// rejected, so plan validation must accept it in both forms. + #[test] + fn validate_query_plan_for_streaming_order_by_accepts_distinct() { + for distinct_type in [DistinctType::Ordered, DistinctType::Unordered] { + let mut plan = order_by_plan(Some("SELECT 1"), vec![qr("", "FF")]); + plan.query_info.as_mut().unwrap().distinct_type = distinct_type; + assert!( + validate_query_plan_for_streaming_order_by(&plan).is_ok(), + "{distinct_type:?} DISTINCT must be accepted alongside ORDER BY" + ); + } + } + + // ── DISTINCT continuation-token validation ─────────────────────────── + // + // The catalog scenarios `malformed_token_rejected`, + // `token_shape_mismatch_rejected`, and + // `distinct_type_mismatch_on_resume_rejected` are pinned here, since they + // are about token shape rather than page contents. + + fn distinct_state(distinct_type: DistinctType) -> PipelineNodeState { + PipelineNodeState::Distinct { + distinct_type, + last_hash: None, + child: Box::new(PipelineNodeState::SequentialDrain { + left_most_undrained_epk: String::new(), + active_tokens: Vec::new(), + }), + } + } + + #[test] + fn peel_distinct_resume_unwraps_a_matching_ordered_token() { + let (inner, last_hash) = peel_distinct_resume( + Some(distinct_state(DistinctType::Ordered)), + DistinctType::Ordered, + ) + .expect("a matching ordered token resumes"); + assert!(matches!( + inner, + Some(PipelineNodeState::SequentialDrain { .. }) + )); + assert_eq!(last_hash, None); + } + + /// Catalog: `distinct_type_mismatch_on_resume_rejected`. Reinterpreting an + /// ordered token as unordered would apply adjacency deduplication to an + /// unsorted stream. + #[test] + fn peel_distinct_resume_rejects_a_distinct_type_mismatch() { + let err = peel_distinct_resume( + Some(distinct_state(DistinctType::Ordered)), + DistinctType::Unordered, + ) + .expect_err("an ordered token must not resume an unordered plan"); + assert_eq!( + err.status().sub_status(), + Some(crate::error::SubStatusCode::CLIENT_CONTINUATION_TOKEN_SHAPE_MISMATCH) + ); + assert!(err.to_string().contains("minted for")); + } + + /// We never mint an unordered DISTINCT token, so one can only be + /// hand-crafted or corrupted — and resuming it would re-emit every value + /// seen before the checkpoint. + #[test] + fn peel_distinct_resume_rejects_a_hand_crafted_unordered_token() { + let err = peel_distinct_resume( + Some(distinct_state(DistinctType::Unordered)), + DistinctType::Unordered, + ) + .expect_err("an unordered DISTINCT token is never resumable"); + assert_eq!( + err.status().sub_status(), + Some(crate::error::SubStatusCode::CLIENT_DISTINCT_CONTINUATION_UNSUPPORTED) + ); + assert!(err.to_string().contains("ORDER BY")); + } + + /// Catalog: `token_shape_mismatch_rejected`. A token minted before DISTINCT + /// support (or for a non-DISTINCT query) has no `Distinct` layer; resuming + /// it would silently skip deduplication for every remaining page. + #[test] + fn peel_distinct_resume_rejects_a_token_without_a_distinct_layer() { + let err = peel_distinct_resume( + Some(PipelineNodeState::SequentialDrain { + left_most_undrained_epk: String::new(), + active_tokens: Vec::new(), + }), + DistinctType::Unordered, + ) + .expect_err("a non-DISTINCT token must not resume a DISTINCT plan"); + assert_eq!( + err.status().sub_status(), + Some(crate::error::SubStatusCode::CLIENT_CONTINUATION_TOKEN_SHAPE_MISMATCH) + ); + assert!(err.to_string().contains("does not match a DISTINCT query")); + } + + /// The mirror case: a `Distinct` token handed to a plan that no longer asks + /// for deduplication is rejected rather than reinterpreted. + #[test] + fn peel_distinct_resume_rejects_a_distinct_token_for_a_plain_plan() { + let err = peel_distinct_resume( + Some(distinct_state(DistinctType::Ordered)), + DistinctType::None, + ) + .expect_err("a DISTINCT token cannot resume a plain plan"); + assert_eq!( + err.status().sub_status(), + Some(crate::error::SubStatusCode::CLIENT_CONTINUATION_TOKEN_SHAPE_MISMATCH) + ); + } + + /// `Drained` is shape-agnostic: the whole pipeline, DISTINCT included, + /// finished, so it must resume for any plan shape. + #[test] + fn peel_distinct_resume_passes_drained_through_unchanged() { + for distinct_type in [ + DistinctType::None, + DistinctType::Ordered, + DistinctType::Unordered, + ] { + let (inner, last_hash) = + peel_distinct_resume(Some(PipelineNodeState::Drained), distinct_type) + .expect("a drained token resumes for any shape"); + assert!(matches!(inner, Some(PipelineNodeState::Drained))); + assert_eq!(last_hash, None); + } + } + + #[test] + fn peel_distinct_resume_passes_a_fresh_start_through() { + let (inner, last_hash) = peel_distinct_resume(None, DistinctType::Unordered) + .expect("a fresh start needs no token"); + assert!(inner.is_none()); + assert_eq!(last_hash, None); + } + /// Catalog: `malformed_token_rejected`. Wrapping the fan-out state in a + /// `Distinct` layer must not bypass the inner state's own validation — a + /// malformed child is still rejected rather than silently restarting the + /// query (which would re-emit every row). + #[tokio::test] + async fn distinct_token_with_a_malformed_child_is_still_rejected() { let mut plan = order_by_plan(Some("SELECT 1"), vec![qr("", "FF")]); plan.query_info.as_mut().unwrap().distinct_type = DistinctType::Ordered; - assert!(validate_query_plan_for_streaming_order_by(&plan).is_err()); + let operation = Arc::new(order_by_operation()); + let mut topology = MockTopologyProvider::new(vec![]); + + // A `SequentialDrain` child is the wrong shape under a streaming + // ORDER BY plan, and the inner builder must say so. + let malformed = PipelineNodeState::Distinct { + distinct_type: DistinctType::Ordered, + last_hash: None, + child: Box::new(PipelineNodeState::SequentialDrain { + left_most_undrained_epk: String::new(), + active_tokens: Vec::new(), + }), + }; + let err = build_streaming_ordered_merge(&plan, &mut topology, &operation, Some(malformed)) + .await + .expect_err("a malformed inner state must not resume"); + assert_eq!( + err.status().sub_status(), + Some(crate::error::SubStatusCode::CLIENT_CONTINUATION_TOKEN_SHAPE_MISMATCH) + ); } #[test] diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/query_plan.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/query_plan.rs index 8af8267742c..296093a5ab4 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/query_plan.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/query_plan.rs @@ -288,7 +288,7 @@ pub(crate) struct HybridSearchQueryInfo { } /// The kind of DISTINCT tracking required by the query. -#[derive(Debug, Clone, Deserialize, Default, PartialEq, Eq, Serialize)] +#[derive(Debug, Clone, Copy, Deserialize, Default, PartialEq, Eq, Serialize)] pub(crate) enum DistinctType { /// No deduplication required. #[default] diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/snapshot.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/snapshot.rs index 189f5ae24bf..a35eec357a1 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/snapshot.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/snapshot.rs @@ -20,8 +20,9 @@ use serde::{Deserialize, Serialize}; use crate::models::ChangeFeedStartFrom; +use super::distinct_hash::Hash128; use super::order_by::OrderByResumeValue; -use super::query_plan::SortOrder; +use super::query_plan::{DistinctType, SortOrder}; /// Serializable snapshot of a [`PipelineNode`](super::PipelineNode) subtree. /// @@ -131,6 +132,29 @@ pub(crate) enum PipelineNodeState { query_fingerprint: Option, ranges: Vec, }, + + /// A `DISTINCT` deduplication stage wrapping a cross-partition root. + /// + /// Only the **ordered** form is ever serialized. `ORDER BY` guarantees + /// structurally equal rows arrive adjacently, so `last_hash` — the hash of + /// the last row emitted before the checkpoint — is the complete resume + /// state: a value the stage has moved past can never reappear. It is + /// `None` when the checkpoint was taken before any row was emitted. + /// + /// An **unordered** `DISTINCT` never reaches this variant: its state is the + /// whole set of values seen, which is unbounded and cannot be truncated + /// without silently re-emitting duplicates, so + /// [`Distinct::snapshot_state`](super::Distinct::snapshot_state) fails + /// instead — *unless* the stage has drained, in which case there is no + /// state left to lose and it snapshots as [`Drained`](Self::Drained) like + /// any other finished node. `distinct_type` is still persisted so a resume + /// can reject a token whose shape no longer matches the query plan. + Distinct { + distinct_type: DistinctType, + #[serde(default, skip_serializing_if = "Option::is_none")] + last_hash: Option, + child: Box, + }, } /// One still-active range of a [`PipelineNodeState::StreamingOrderedMerge`]. @@ -261,6 +285,7 @@ impl PipelineNodeState { PipelineNodeState::UnorderedMerge { .. } => "UnorderedMerge", PipelineNodeState::SkipTake { .. } => "SkipTake", PipelineNodeState::StreamingOrderedMerge { .. } => "StreamingOrderedMerge", + PipelineNodeState::Distinct { .. } => "Distinct", }, )) .build()), diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/error/cosmos_status.rs b/sdk/cosmos/azure_data_cosmos_driver/src/error/cosmos_status.rs index de569c15581..e9b35fddb94 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/error/cosmos_status.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/error/cosmos_status.rs @@ -500,6 +500,8 @@ impl SubStatusCode { 20120 => Some("ClientInvalidResourceId"), 20121 => Some("ClientMixedNameRidAddressing"), 20122 => Some("ClientQueryRewriteBodyInvalid"), + 20123 => Some("ClientDistinctValueTooDeeplyNested"), + 20124 => Some("ClientDistinctContinuationUnsupported"), 20150 => Some("ClientDuplicateFaultInjectionRuleId"), 20151 => Some("ClientThroughputControlGroupRegistrationFailed"), 20152 => Some("ClientThroughputControlGroupNotRegistered"), @@ -518,6 +520,7 @@ impl SubStatusCode { 20206 => Some("ClientSplitRetriesExhausted"), 20207 => Some("ClientBuildResponseInvokedOnFailure"), 20208 => Some("ClientRootNodeCannotRequestSplit"), + 20217 => Some("ClientDistinctCannotForwardSplit"), 20209 => Some("ClientCrossPartitionQueryRequiresContainerRef"), 20210 => Some("ClientSingletonOperationReturnedEmptyPage"), 20211 => Some("ClientComputeRangeInvokedWithEmptyPartitionKey"), @@ -1381,6 +1384,17 @@ impl SubStatusCode { /// each partition's query text. pub const CLIENT_QUERY_REWRITE_BODY_INVALID: SubStatusCode = SubStatusCode(20122); + /// A `DISTINCT` value nested deeper than the structural hasher's depth + /// limit (20123). Cosmos caps document nesting well below that limit, so + /// this indicates a hand-crafted or corrupt payload. + pub const CLIENT_DISTINCT_VALUE_TOO_DEEPLY_NESTED: SubStatusCode = SubStatusCode(20123); + + /// A continuation token was requested for an unordered `DISTINCT` query + /// (20124). Resuming would require carrying the entire set of seen values, + /// so the token is refused rather than silently re-emitting duplicates. + /// Adding a matching `ORDER BY` makes the query resumable. + pub const CLIENT_DISTINCT_CONTINUATION_UNSUPPORTED: SubStatusCode = SubStatusCode(20124); + // ----- 20150-20199: SDK configuration / setup errors ----- /// Two fault-injection rules registered with the same id (20150). @@ -1466,6 +1480,13 @@ impl SubStatusCode { /// handled by a parent node (20208). pub const CLIENT_ROOT_NODE_CANNOT_REQUEST_SPLIT: SubStatusCode = SubStatusCode(20208); + /// A `DISTINCT` node was asked to forward a partition split (20217). + /// `SplitRequired` replaces the node that emits it, which would discard the + /// deduplication map and resurrect already-suppressed values, so the split + /// is refused here instead of being passed to a parent. Unreachable today: + /// the wrapped fan-out node absorbs splits internally. + pub const CLIENT_DISTINCT_CANNOT_FORWARD_SPLIT: SubStatusCode = SubStatusCode(20217); + /// A cross-partition query plan was attempted without a container /// reference (20209). pub const CLIENT_CROSS_PARTITION_QUERY_REQUIRES_CONTAINER_REF: SubStatusCode = @@ -2321,6 +2342,27 @@ impl CosmosStatus { sub_status: Some(SubStatusCode::CLIENT_QUERY_REWRITE_BODY_INVALID), }; + /// 400 / 20123 — a `DISTINCT` value nested past the structural hasher's + /// depth limit. + pub const CLIENT_DISTINCT_VALUE_TOO_DEEPLY_NESTED: CosmosStatus = CosmosStatus { + status_code: StatusCode::BadRequest, + sub_status: Some(SubStatusCode::CLIENT_DISTINCT_VALUE_TOO_DEEPLY_NESTED), + }; + + /// 400 / 20124 — a continuation token was requested for an unordered + /// `DISTINCT` query, which cannot be resumed safely. + pub const CLIENT_DISTINCT_CONTINUATION_UNSUPPORTED: CosmosStatus = CosmosStatus { + status_code: StatusCode::BadRequest, + sub_status: Some(SubStatusCode::CLIENT_DISTINCT_CONTINUATION_UNSUPPORTED), + }; + + /// 500 / 20217 — a `DISTINCT` node was asked to forward a partition split, + /// which would discard its deduplication state. + pub const CLIENT_DISTINCT_CANNOT_FORWARD_SPLIT: CosmosStatus = CosmosStatus { + status_code: StatusCode::InternalServerError, + sub_status: Some(SubStatusCode::CLIENT_DISTINCT_CANNOT_FORWARD_SPLIT), + }; + // Configuration / setup (HTTP 400, sub-status 20150-20199) /// 400 / 20150 — duplicate fault-injection rule id. diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/in_memory_emulator/operations.rs b/sdk/cosmos/azure_data_cosmos_driver/src/in_memory_emulator/operations.rs index 5e33769a054..4a201951324 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/in_memory_emulator/operations.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/in_memory_emulator/operations.rs @@ -3491,11 +3491,22 @@ fn synthesize_order_by_rewritten_query( Some(t) if t.kind == TokenKind::Identifier => t.text, _ => collection_token.text, }; + // `SELECT DISTINCT …` — the envelope wraps each row with its `_rid`, so a + // per-partition `DISTINCT` on the envelope could never collapse anything. + // Drop it from the rewritten query and let the client-side ordered + // `Distinct` stage deduplicate the globally sorted stream instead. + let mut payload_idx = select_idx + 1; + if tokens + .get(payload_idx) + .is_some_and(|t| t.kind == TokenKind::Distinct) + { + payload_idx += 1; + } // Skip a leading `TOP ` / `TOP @param`: the global TOP is applied by the // client's `SkipTake` over the merged stream, so the per-partition envelope // must not carry it (a per-partition TOP would drop rows a later partition - // needs for the global ordering). - let mut payload_idx = select_idx + 1; + // needs for the global ordering). `DISTINCT` precedes `TOP` in SQL, so this + // runs after the `DISTINCT` skip above. if tokens .get(payload_idx) .is_some_and(|t| t.kind == TokenKind::Top) diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/models/distributed_transaction.rs b/sdk/cosmos/azure_data_cosmos_driver/src/models/distributed_transaction.rs index 99620352959..79081eb52c4 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/models/distributed_transaction.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/models/distributed_transaction.rs @@ -413,11 +413,11 @@ fn validate_resource_body_id( fn partition_key_json( target: &DistributedTransactionTarget, ) -> crate::error::Result { - let (_, value) = target - .partition_key - .as_headers()? - .next() - .ok_or_else(|| invalid_partition_key("partition key did not produce a header value"))?; + let (_, value) = target.partition_key.as_headers()?.next().ok_or_else(|| { + invalid_partition_key( + "distributed transaction operations require a non-empty partition key", + ) + })?; let text = value.as_str(); serde_json::from_str(text).map_err(|error| { crate::error::CosmosError::builder() @@ -981,6 +981,36 @@ mod tests { DistributedTransactionTarget::new(container(), PartitionKey::from("pk1"), id.to_owned()) } + /// A distributed transaction addresses a single item, so an empty partition + /// key is a client error rather than a cross-partition request. + /// + /// This used to be caught only by accident: `partition_key_json` took the + /// first header an empty key produced — `…enablecrosspartition: True` — + /// and rejected it because `True` is not valid JSON (JSON's boolean is + /// lowercase). Now that an empty key emits no headers at all, the + /// dedicated `ok_or_else` arm is what rejects it. + #[test] + fn empty_partition_key_is_rejected() { + let target = + DistributedTransactionTarget::new(container(), PartitionKey::EMPTY, "item1".to_owned()); + let operation = + DistributedTransactionOperation::new(DistributedTransactionOperationKind::Read, target); + let request = + DistributedTransactionRequest::new(DistributedTransactionType::Read, vec![operation]); + + let error = request + .serialize_body() + .expect_err("an empty partition key cannot address a single item"); + assert_eq!( + error.status().status_code(), + azure_core::http::StatusCode::BadRequest + ); + assert!( + error.to_string().contains("non-empty partition key"), + "the message must name the actual problem, got {error}" + ); + } + #[test] fn serialize_create_operation() { let operation = DistributedTransactionOperation::new( diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/models/partition_key.rs b/sdk/cosmos/azure_data_cosmos_driver/src/models/partition_key.rs index cf284380bb5..86e60bf25ad 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/models/partition_key.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/models/partition_key.rs @@ -12,10 +12,6 @@ use std::{borrow::Cow, hash::Hash}; pub(crate) const PARTITION_KEY: HeaderName = HeaderName::from_static("x-ms-documentdb-partitionkey"); -/// Header name to enable cross-partition queries. -pub(crate) const QUERY_ENABLE_CROSS_PARTITION: HeaderName = - HeaderName::from_static("x-ms-documentdb-query-enablecrosspartition"); - // ============================================================================= // PartitionKeyValue // ============================================================================= @@ -326,7 +322,7 @@ impl PartitionKey { impl AsHeaders for PartitionKey { type Error = crate::error::CosmosError; - type Iter = std::iter::Once<(HeaderName, HeaderValue)>; + type Iter = std::option::IntoIter<(HeaderName, HeaderValue)>; fn as_headers(&self) -> Result { // We have to do some manual JSON serialization here. @@ -334,13 +330,13 @@ impl AsHeaders for PartitionKey { // It's not safe to use non-ASCII characters in HTTP headers, and serde_json will not escape non-ASCII characters if they are otherwise valid as UTF-8. // So, we do some conversion by hand, with the help of Rust's own `encode_utf16` method which gives us the necessary code points for non-ASCII values, and produces surrogate pairs as needed. - // Quick shortcut for empty partition keys list, which also prevents a bug when we pop the trailing comma for an empty list. + // An empty partition key means "no specific partition" (see + // `PartitionKey::EMPTY`), which the driver expresses by targeting + // partition key ranges explicitly — so there is no header to emit. + // Returning early also keeps the trailing-comma pop below from + // stripping the opening '[' and producing a bare ']'. if self.0.is_empty() { - // An empty partition key means a cross partition query - return Ok(std::iter::once(( - QUERY_ENABLE_CROSS_PARTITION, - HeaderValue::from_static("True"), - ))); + return Ok(None.into_iter()); } let mut json = String::new(); @@ -411,10 +407,7 @@ impl AsHeaders for PartitionKey { json.pop(); json.push(']'); - Ok(std::iter::once(( - PARTITION_KEY, - HeaderValue::from_cow(json), - ))) + Ok(Some((PARTITION_KEY, HeaderValue::from_cow(json))).into_iter()) } } @@ -502,6 +495,32 @@ mod tests { assert_eq!(pk.len(), 0); } + /// An empty partition key means "no specific partition", which the driver + /// expresses by targeting partition key ranges explicitly — so it must emit + /// no headers at all. + /// + /// It previously emitted `x-ms-documentdb-query-enablecrosspartition`, + /// which contradicted [`PartitionKey::EMPTY`]'s own contract and was the + /// only place that header was ever produced. + #[test] + fn empty_partition_key_emits_no_headers() { + let headers: Vec<_> = PartitionKey::EMPTY.as_headers().unwrap().collect(); + assert!( + headers.is_empty(), + "an empty partition key must not put anything on the wire, got {headers:?}" + ); + } + + /// Regression: without the empty-list guard the trailing-comma pop would + /// strip the opening `[` and emit a bare `]`. + #[test] + fn non_empty_partition_key_emits_exactly_the_partition_key_header() { + let headers: Vec<_> = PartitionKey::from("test").as_headers().unwrap().collect(); + assert_eq!(headers.len(), 1); + assert_eq!(headers[0].0, PARTITION_KEY); + assert_eq!(headers[0].1.as_str(), r#"["test"]"#); + } + #[test] fn null_partition_key_value() { let pk = PartitionKey::from(None::); diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/query/eval/mod.rs b/sdk/cosmos/azure_data_cosmos_driver/src/query/eval/mod.rs index 28c790bb718..2efed535c8a 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/query/eval/mod.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/query/eval/mod.rs @@ -1003,6 +1003,24 @@ pub fn query_documents( results = indices.iter().map(|&i| results[i].clone()).collect(); } + // ── Step 3b: DISTINCT ──────────────────────────────────────────────── + // + // The real backend deduplicates within each physical partition before the + // client deduplicates globally, so the emulator must too — otherwise an + // emulator-backed test would exercise a page shape production never + // produces. Uses the same structural hash the client-side + // `Distinct` stage uses, so the two agree on what "equal" means. + if query.select.distinct { + let mut seen = std::collections::HashSet::new(); + let mut deduped = Vec::with_capacity(results.len()); + for row in results { + if seen.insert(crate::driver::dataflow::distinct_hash::hash_value(&row)?) { + deduped.push(row); + } + } + results = deduped; + } + // ── Step 4: TOP ────────────────────────────────────────────────────── if let Some(top) = &query.select.top { let n = match top { diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/query/mod.rs b/sdk/cosmos/azure_data_cosmos_driver/src/query/mod.rs index 503200f09e1..28c86f386ab 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/query/mod.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/query/mod.rs @@ -49,7 +49,8 @@ pub(crate) use parser::parse; /// Tests use [`__TEST_ONLY_SUPPORTED_QUERY_FEATURES`] (broad, matches what /// Java/.NET advertise) so plan-shape parity against the live Gateway is /// validated end-to-end across the full feature surface. -pub(crate) const SUPPORTED_QUERY_FEATURES: &str = "MultipleOrderBy,OffsetAndLimit,OrderBy,Top"; +pub(crate) const SUPPORTED_QUERY_FEATURES: &str = + "Distinct,MultipleOrderBy,OffsetAndLimit,OrderBy,Top"; /// Broad supported-features list used by cross-crate gateway-comparison /// tests. Matches what the Java and .NET SDKs send today so the Gateway diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/query/plan/mod.rs b/sdk/cosmos/azure_data_cosmos_driver/src/query/plan/mod.rs index 2c89493f723..7ffc2877fb9 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/query/plan/mod.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/query/plan/mod.rs @@ -373,6 +373,41 @@ fn is_constant_expression(expr: &SqlScalarExpression) -> bool { } } +/// `true` when a `DISTINCT` query can be deduplicated by adjacency: a +/// `SELECT DISTINCT VALUE ` whose `ORDER BY` is exactly that one path. +/// +/// Adjacency deduplication compares each row only against the previous one, so +/// it is sound only when the sort groups structurally equal *projected rows* +/// into contiguous runs. `SELECT DISTINCT c.name, c.city FROM c ORDER BY c.name` +/// can legitimately deliver `(a,x), (a,y), (a,x)`, which an adjacency map would +/// emit twice. +/// +/// The exact shape is measured, not derived — against a live account with +/// production's `SUPPORTED_QUERY_FEATURES`, the Gateway reports `Ordered` for +/// `SELECT DISTINCT VALUE c.name FROM c ORDER BY c.name ASC` (ascending or +/// descending) and `Unordered` for every other form tried: the list form +/// (`SELECT DISTINCT c.name …`) whatever its `ORDER BY`, a sort on a different +/// path, a multi-column `ORDER BY` even when it *leads* with the projected path, +/// and no `ORDER BY` at all. `tests/gateway_query_plan_comparison.rs::gw_distinct` +/// pins each of those against the live service. +/// +/// Matching the service exactly keeps this generator in step with the plans the +/// driver executes; it is also the only form +/// `synthesize_order_by_rewritten_query` can build an envelope query for. When +/// in doubt, `Unordered` is always correct and merely gives up resumability. +fn distinct_is_ordered(spec: &SqlSelectSpec, order_by_expressions: &[String]) -> bool { + // Exactly one sort key; see the measured note above. + let [sort_path] = order_by_expressions else { + return false; + }; + // Only `SELECT DISTINCT VALUE `. + let SqlSelectSpec::Value(expr) = spec else { + return false; + }; + let mut parts = Vec::new(); + collect_path_parts(expr, &mut parts) && parts.join(".") == *sort_path +} + fn analyze_query(query: &SqlQuery, parameters: &Params) -> crate::error::Result { let mut info = LocalQueryInfo { has_select_value: matches!(query.select.spec, SqlSelectSpec::Value(_)), @@ -380,29 +415,8 @@ fn analyze_query(query: &SqlQuery, parameters: &Params) -> crate::error::Result< ..Default::default() }; - // DISTINCT — Gateway optimizes away DISTINCT when the SELECT expression is a - // constant (literal) that doesn't reference any collection variable, because - // a single constant value is always distinct by definition. - if query.select.distinct { - // Gateway only collapses DISTINCT-on-constant for the `SELECT DISTINCT VALUE ` - // form. The list form (`SELECT DISTINCT 1, 2 FROM c`) is treated as ordinary DISTINCT - // by the Gateway because the result rows are JSON objects (with synthesized property - // names) and are therefore not all guaranteed to be identical. We mirror that - // asymmetry intentionally — do not extend this to `SqlSelectSpec::List` without - // verifying behavior against the Gateway. - let is_constant_select = match &query.select.spec { - SqlSelectSpec::Value(expr) => is_constant_expression(expr), - _ => false, - }; - if is_constant_select { - // Gateway reports distinctType: "None" for constant expressions - info.distinct_type = DistinctType::None; - } else if query.order_by.is_some() { - info.distinct_type = DistinctType::Ordered; - } else { - info.distinct_type = DistinctType::Unordered; - } - } + // DISTINCT is classified after ORDER BY below, since `Ordered` requires + // the projection and the sort to agree. // TOP — substitute parameterized values; error if unresolvable. info.top = match &query.select.top { @@ -443,6 +457,30 @@ fn analyze_query(query: &SqlQuery, parameters: &Params) -> crate::error::Result< } } + // DISTINCT — the Gateway optimizes DISTINCT away when the query selects a + // constant (literal) and has no `FROM` clause, because such a query yields + // exactly one row and is trivially distinct. + if query.select.distinct { + // The `FROM` clause is load-bearing: measured against a live account, + // `SELECT DISTINCT VALUE 1` reports `None` but + // `SELECT DISTINCT VALUE 1 FROM c` reports `Unordered` — with a `FROM` + // the query yields one row *per document*, so deduplication is real + // work. The list form (`SELECT DISTINCT 1 AS p FROM c`) is likewise + // `Unordered`. `gw_distinct` pins all three against the service. + let is_constant_select = query.from.is_none() + && match &query.select.spec { + SqlSelectSpec::Value(expr) => is_constant_expression(expr), + _ => false, + }; + info.distinct_type = if is_constant_select { + DistinctType::None + } else if distinct_is_ordered(&query.select.spec, &info.order_by_expressions) { + DistinctType::Ordered + } else { + DistinctType::Unordered + }; + } + // JOIN if let Some(from) = &query.from { info.has_join = has_join(&from.collection); @@ -1472,10 +1510,15 @@ mod tests { assert_eq!(qp.query_info.distinct_type, DistinctType::Unordered); } + /// The `VALUE` form is the only one the service reports as `Ordered` + /// (measured against a live account); the list form below is downgraded. #[test] fn distinct_ordered() { - let qp = plan("SELECT DISTINCT c.name FROM c ORDER BY c.name"); + let qp = plan("SELECT DISTINCT VALUE c.name FROM c ORDER BY c.name"); assert_eq!(qp.query_info.distinct_type, DistinctType::Ordered); + + let qp = plan("SELECT DISTINCT c.name FROM c ORDER BY c.name"); + assert_eq!(qp.query_info.distinct_type, DistinctType::Unordered); } #[test] diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/query/plan/tests/query_plan_comparison.rs b/sdk/cosmos/azure_data_cosmos_driver/src/query/plan/tests/query_plan_comparison.rs index c8278a818f2..c4bd5b5fa67 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/query/plan/tests/query_plan_comparison.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/query/plan/tests/query_plan_comparison.rs @@ -506,16 +506,138 @@ fn distinct_unordered() { ); } +#[test] +/// The Gateway reports `Unordered` when the sort key is not a prefix of the +/// `DISTINCT` projection: the sort cannot group equal projections into runs, so +/// adjacency deduplication (and therefore continuation) is unavailable. Java +/// pins this pair in `TestSuiteBase.queryWithOrderByProvider` +/// (`DistinctQueryTests.queryDocumentsWithOrderBy`), where the mismatched form +/// fails with HTTP 400 on resume. +fn distinct_with_mismatched_order_by_is_unordered() { + assert_eq!( + plan("SELECT DISTINCT VALUE c.id FROM c ORDER BY c._ts DESC") + .query_info + .distinct_type, + DistinctType::Unordered + ); + assert_eq!( + plan("SELECT DISTINCT VALUE c.id FROM c ORDER BY c.id DESC") + .query_info + .distinct_type, + DistinctType::Ordered + ); +} + +/// `SELECT DISTINCT *` deduplicates whole documents, which no single-property +/// sort can group into runs. +#[test] +fn distinct_star_with_order_by_is_unordered() { + assert_eq!( + plan("SELECT DISTINCT * FROM c ORDER BY c.name") + .query_info + .distinct_type, + DistinctType::Unordered + ); +} + +/// Pins the local classifier against `distinctType` measured from a live +/// Gateway using production's `SUPPORTED_QUERY_FEATURES`. The same shapes are +/// re-checked against the service by +/// `tests/gateway_query_plan_comparison.rs::gw_distinct`. +/// +/// The service reports `Ordered` only for `SELECT DISTINCT VALUE ` whose +/// `ORDER BY` is exactly that one path. Every other form is downgraded — the +/// list form whatever its `ORDER BY`, a sort on a different path, and a +/// multi-column sort even when it leads with the projected path. +#[test] +fn distinct_is_ordered_only_for_the_value_form_matching_its_order_by() { + // Measured: Ordered. + assert_eq!( + plan("SELECT DISTINCT VALUE c.name FROM c ORDER BY c.name ASC") + .query_info + .distinct_type, + DistinctType::Ordered + ); + // Measured: Unordered — a multi-column ORDER BY, even one that *leads* + // with the projected path. Adjacency would in fact still hold here, but + // the service does not report it, and we follow the service. + assert_eq!( + plan("SELECT DISTINCT VALUE c.id FROM c ORDER BY c.id, c.other") + .query_info + .distinct_type, + DistinctType::Unordered + ); + // Measured: Unordered — list form, even projecting a single path. + assert_eq!( + plan("SELECT DISTINCT c.name FROM c ORDER BY c.name ASC") + .query_info + .distinct_type, + DistinctType::Unordered + ); + // Measured: Unordered — sort covers only the first of two projections. + assert_eq!( + plan("SELECT DISTINCT c.name, c.city FROM c ORDER BY c.name ASC") + .query_info + .distinct_type, + DistinctType::Unordered + ); + // Measured: Unordered — no ORDER BY at all. + assert_eq!( + plan("SELECT DISTINCT VALUE c.name FROM c") + .query_info + .distinct_type, + DistinctType::Unordered + ); + // A sort on a different path cannot group equal projections. + assert_eq!( + plan("SELECT DISTINCT VALUE c.id FROM c ORDER BY c._ts DESC") + .query_info + .distinct_type, + DistinctType::Unordered + ); +} + +/// Constant DISTINCT collapses to `None` only without a FROM clause. With one, +/// the query yields a row per document, so deduplication is real work — measured +/// live, and pinned by `gw_distinct`. Getting this wrong makes the driver skip +/// the DISTINCT stage entirely and return one row per document. +#[test] +fn constant_distinct_collapses_only_without_a_from_clause() { + for sql in [ + "SELECT DISTINCT VALUE 1", + "SELECT DISTINCT VALUE null", + "SELECT DISTINCT VALUE 'a'", + ] { + assert_eq!( + plan(sql).query_info.distinct_type, + DistinctType::None, + "{sql}" + ); + } + for sql in [ + "SELECT DISTINCT VALUE 1 FROM c", + "SELECT DISTINCT VALUE null FROM c", + "SELECT DISTINCT 1 AS p FROM c", + ] { + assert_eq!( + plan(sql).query_info.distinct_type, + DistinctType::Unordered, + "{sql}" + ); + } +} + #[test] fn distinct_ordered() { assert_eq!( - plan("SELECT DISTINCT c.name FROM c ORDER BY c.name ASC"), + plan("SELECT DISTINCT VALUE c.name FROM c ORDER BY c.name ASC"), QueryPlan { pk_filters: PartitionKeyFilter::Unconstrained, query_info: LocalQueryInfo { distinct_type: DistinctType::Ordered, order_by: vec![SortOrder::Ascending], order_by_expressions: vec!["c.name".into()], + has_select_value: true, ..qi() }, } @@ -968,7 +1090,9 @@ fn complex_distinct_top_order() { QueryPlan { pk_filters: PartitionKeyFilter::Unconstrained, query_info: LocalQueryInfo { - distinct_type: DistinctType::Ordered, + // List form: the service downgrades it to `Unordered` whatever + // its ORDER BY. See `plan::distinct_is_ordered`. + distinct_type: DistinctType::Unordered, top: Some(5), order_by: vec![SortOrder::Ascending], order_by_expressions: vec!["c.name".into()], @@ -1047,7 +1171,11 @@ fn complex_everything() { QueryPlan { pk_filters: PartitionKeyFilter::Equality(vec![PartitionKeyValue::String("x".into())]), query_info: LocalQueryInfo { - distinct_type: DistinctType::Ordered, + // `Unordered`, not `Ordered`: the sort covers only `c.city`, while the + // projection carries additional columns, so equal projected rows are + // not guaranteed to arrive adjacently and adjacency deduplication + // would be unsound. See `plan::distinct_is_ordered`. + distinct_type: DistinctType::Unordered, top: Some(100), offset: None, limit: None, @@ -1618,7 +1746,11 @@ fn complex_all_clauses() { QueryPlan { pk_filters: PartitionKeyFilter::Equality(vec![PartitionKeyValue::String("x".into())]), query_info: LocalQueryInfo { - distinct_type: DistinctType::Ordered, + // `Unordered`, not `Ordered`: the sort covers only `c.city`, while the + // projection carries additional columns, so equal projected rows are + // not guaranteed to arrive adjacently and adjacency deduplication + // would be unsound. See `plan::distinct_is_ordered`. + distinct_type: DistinctType::Unordered, top: Some(50), offset: None, limit: None, @@ -2683,7 +2815,11 @@ fn complex_pk_in_with_distinct_top_order() { assert_eq!( qp.query_info, LocalQueryInfo { - distinct_type: DistinctType::Ordered, + // `Unordered`, not `Ordered`: the sort covers only `c.name`, while the + // projection carries additional columns, so equal projected rows are + // not guaranteed to arrive adjacently and adjacency deduplication + // would be unsound. See `plan::distinct_is_ordered`. + distinct_type: DistinctType::Unordered, top: Some(10), order_by: vec![SortOrder::Ascending], order_by_expressions: vec!["c.name".into()], @@ -2769,7 +2905,11 @@ fn complex_everything_with_hpk() { PartitionKeyValue::String("u1".into()), ]), query_info: LocalQueryInfo { - distinct_type: DistinctType::Ordered, + // `Unordered`, not `Ordered`: the sort covers only `c.city`, while the + // projection carries additional columns, so equal projected rows are + // not guaranteed to arrive adjacently and adjacency deduplication + // would be unsound. See `plan::distinct_is_ordered`. + distinct_type: DistinctType::Unordered, top: Some(100), offset: Some(5), limit: Some(20), diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/query_plan_native/native_dll_tests.rs b/sdk/cosmos/azure_data_cosmos_driver/src/query_plan_native/native_dll_tests.rs index c7ca1f493da..463ecde90b1 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/query_plan_native/native_dll_tests.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/query_plan_native/native_dll_tests.rs @@ -587,7 +587,7 @@ fn distinct_field() { assert_query_info( &actual, QueryInfo { - distinct_type: actual.distinct_type.clone(), + distinct_type: actual.distinct_type, ..qi() }, ); diff --git a/sdk/cosmos/azure_data_cosmos_driver/tests/distinct_scenario_catalog.rs b/sdk/cosmos/azure_data_cosmos_driver/tests/distinct_scenario_catalog.rs new file mode 100644 index 00000000000..cf0b90e8562 --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos_driver/tests/distinct_scenario_catalog.rs @@ -0,0 +1,521 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +//! Fixture/schema tests for the cross-partition `DISTINCT` scenario catalog +//! (`tests/fixtures/distinct_scenarios.json`). +//! +//! Mirrors `streaming_order_by_scenario_catalog.rs`: this is the single +//! source-attributed catalog reused across every test layer (map unit tests, +//! mock-pipeline tests, and in-memory-emulator tests). Each layer lives in a +//! different compilation unit and defines its own minimal `Deserialize` view +//! of the fixture; this file owns the strict, canonical schema validation +//! every other layer trusts. +//! +//! Validates: no duplicate scenario IDs; every `layers` entry is a known layer +//! name; every scenario declares at least one layer, at least one cross-SDK +//! source, and some expected result; `distinctType` is one of the three plan +//! values; mock partitions are sorted, have no gaps, and tile correctly; a +//! scenario declaring `mockPipeline` actually carries a mock; an `Ordered` +//! scenario declares its sort columns; and required scenario-inventory +//! categories are represented. + +use std::collections::BTreeSet; + +use serde::Deserialize; + +const CATALOG_JSON: &str = include_str!("fixtures/distinct_scenarios.json"); + +const KNOWN_LAYERS: &[&str] = &[ + "distinctMap", + "mockPipeline", + "inMemoryEmulator", + "recorded", +]; + +const KNOWN_DISTINCT_TYPES: &[&str] = &["None", "Ordered", "Unordered"]; + +#[derive(Deserialize)] +struct Catalog { + #[serde(rename = "schemaVersion")] + schema_version: u32, + scenarios: Vec, +} + +#[derive(Deserialize)] +struct Scenario { + id: String, + #[allow(dead_code)] + description: String, + sources: Vec, + layers: Vec, + query: QuerySpec, + #[serde(default)] + #[allow(dead_code)] + documents: Vec, + mock: Option, + #[serde(rename = "pageSizes", default)] + page_sizes: Vec, + #[serde(rename = "expectedIds", default)] + expected_ids: Vec, + #[serde(rename = "expectedValues", default)] + expected_values: Vec, + #[allow(dead_code)] + checkpoint: Option, + #[serde(rename = "expectedContinuation")] + expected_continuation: Option, + #[serde(rename = "expectedError")] + expected_error: Option, +} + +#[derive(Deserialize)] +struct Source { + sdk: String, + #[allow(dead_code)] + path: String, + #[allow(dead_code)] + test: String, +} + +#[derive(Deserialize)] +struct QuerySpec { + text: String, + #[serde(default)] + #[allow(dead_code)] + parameters: Vec, + #[serde(default)] + columns: Vec, + #[serde(rename = "distinctType")] + distinct_type: String, +} + +#[derive(Deserialize)] +struct ColumnSpec { + #[allow(dead_code)] + expression: String, + direction: String, +} + +#[derive(Deserialize)] +struct MockSpec { + partitions: Vec, +} + +#[derive(Deserialize)] +struct MockPartition { + range: MockRange, + pages: Vec, +} + +#[derive(Deserialize)] +struct MockRange { + #[serde(rename = "minEpk")] + min_epk: String, + #[serde(rename = "maxEpk")] + max_epk: String, +} + +#[derive(Deserialize)] +struct MockPage { + rows: Vec, + #[allow(dead_code)] + continuation: Option, +} + +#[derive(Deserialize)] +struct MockRow { + rid: String, + /// Present only when the scenario also exercises an `ORDER BY` envelope; + /// an unordered `DISTINCT` page is a plain `Documents` feed. + #[serde(rename = "orderByItems", default)] + #[allow(dead_code)] + order_by_items: Option>, + payload: serde_json::Value, +} + +#[derive(Deserialize)] +struct ExpectedError { + category: String, + #[serde(rename = "messageFragment")] + message_fragment: String, +} + +fn load_catalog() -> Catalog { + serde_json::from_str(CATALOG_JSON) + .expect("catalog must be valid JSON matching the strict schema") +} + +#[test] +fn catalog_has_expected_schema_version() { + assert_eq!(load_catalog().schema_version, 1); +} + +#[test] +fn catalog_is_non_empty() { + let catalog = load_catalog(); + assert!( + catalog.scenarios.len() >= 20, + "expected a substantial scenario catalog, found {}", + catalog.scenarios.len() + ); +} + +#[test] +fn no_duplicate_scenario_ids() { + let catalog = load_catalog(); + let mut seen = BTreeSet::new(); + for scenario in &catalog.scenarios { + assert!( + seen.insert(scenario.id.clone()), + "duplicate scenario id: {}", + scenario.id + ); + } +} + +#[test] +fn every_scenario_declares_at_least_one_known_layer() { + let catalog = load_catalog(); + for scenario in &catalog.scenarios { + assert!( + !scenario.layers.is_empty(), + "scenario {} declares no layers", + scenario.id + ); + for layer in &scenario.layers { + assert!( + KNOWN_LAYERS.contains(&layer.as_str()), + "scenario {} declares unknown layer {layer:?}", + scenario.id + ); + } + } +} + +#[test] +fn every_scenario_has_at_least_one_source() { + let catalog = load_catalog(); + for scenario in &catalog.scenarios { + assert!( + !scenario.sources.is_empty(), + "scenario {} has no cross-SDK source attribution", + scenario.id + ); + for source in &scenario.sources { + assert!( + matches!(source.sdk.as_str(), "dotnet" | "java"), + "scenario {} cites unknown sdk {:?}", + scenario.id, + source.sdk + ); + } + } +} + +#[test] +fn every_scenario_declares_a_known_distinct_type() { + let catalog = load_catalog(); + for scenario in &catalog.scenarios { + assert!( + KNOWN_DISTINCT_TYPES.contains(&scenario.query.distinct_type.as_str()), + "scenario {} declares unknown distinctType {:?}", + scenario.id, + scenario.query.distinct_type + ); + } +} + +/// A `DISTINCT` scenario's query text must actually contain `DISTINCT` — +/// otherwise the scenario is silently testing the plain query path. +#[test] +fn every_scenario_queries_distinct() { + let catalog = load_catalog(); + for scenario in &catalog.scenarios { + assert!( + scenario.query.text.contains("DISTINCT"), + "scenario {} declares a DISTINCT catalog entry but its query text has no DISTINCT: {}", + scenario.id, + scenario.query.text + ); + } +} + +/// Only an `Ordered` scenario may declare sort columns, and it must: the +/// ordered map's whole correctness argument rests on the stream being sorted. +#[test] +fn ordered_scenarios_declare_their_sort_columns() { + let catalog = load_catalog(); + for scenario in &catalog.scenarios { + if scenario.query.distinct_type == "Ordered" { + assert!( + !scenario.query.columns.is_empty(), + "scenario {} is Ordered but declares no ORDER BY columns", + scenario.id + ); + } + for column in &scenario.query.columns { + assert!( + matches!(column.direction.as_str(), "Ascending" | "Descending"), + "scenario {} declares unknown sort direction {:?}", + scenario.id, + column.direction + ); + } + } +} + +#[test] +fn every_scenario_has_an_explicit_expected_result() { + let catalog = load_catalog(); + for scenario in &catalog.scenarios { + let all_mock_rows_empty = scenario + .mock + .as_ref() + .map(|mock| { + mock.partitions + .iter() + .all(|p| p.pages.iter().all(|page| page.rows.is_empty())) + }) + .unwrap_or(false); + let has_expected_result = !scenario.expected_ids.is_empty() + || !scenario.expected_values.is_empty() + || scenario.expected_error.is_some() + || scenario.expected_continuation.is_some() + || all_mock_rows_empty; + assert!( + has_expected_result, + "scenario {} has no explicit expected result (expectedIds, expectedValues, \ + expectedError, expectedContinuation, or an all-empty mock)", + scenario.id + ); + } +} + +/// A scenario tagged `mockPipeline` is driven straight off `mock`; without one +/// the layer would silently skip it. +#[test] +fn mock_pipeline_scenarios_carry_a_mock_or_expect_an_error() { + let catalog = load_catalog(); + for scenario in &catalog.scenarios { + if !scenario.layers.iter().any(|l| l == "mockPipeline") { + continue; + } + assert!( + scenario.mock.is_some() || scenario.expected_error.is_some(), + "scenario {} declares the mockPipeline layer but supplies neither a mock nor an \ + expectedError", + scenario.id + ); + } +} + +/// A scenario tagged `inMemoryEmulator` is seeded from `documents`. +#[test] +fn emulator_scenarios_carry_documents() { + let catalog = load_catalog(); + for scenario in &catalog.scenarios { + if !scenario.layers.iter().any(|l| l == "inMemoryEmulator") { + continue; + } + assert!( + !scenario.documents.is_empty(), + "scenario {} declares the inMemoryEmulator layer but seeds no documents", + scenario.id + ); + } +} + +#[test] +fn mock_partitions_are_sorted_and_tile_with_no_gaps_or_overlaps() { + let catalog = load_catalog(); + for scenario in &catalog.scenarios { + let Some(mock) = &scenario.mock else { + continue; + }; + if mock.partitions.is_empty() { + continue; + } + let mut cursor = mock.partitions[0].range.min_epk.clone(); + for (idx, partition) in mock.partitions.iter().enumerate() { + assert_eq!( + partition.range.min_epk, cursor, + "scenario {}: partition {idx} does not start where the previous one ended \ + (gap or overlap)", + scenario.id, + ); + assert!( + partition.range.min_epk < partition.range.max_epk, + "scenario {}: partition {idx} has an invalid range (min >= max)", + scenario.id, + ); + cursor = partition.range.max_epk.clone(); + } + } +} + +/// `_rid` uniquely identifies a row, so a repeated one inside a scenario means +/// the fixture, not the implementation, is producing the duplicate. +#[test] +fn mock_row_rids_are_unique_within_a_scenario() { + let catalog = load_catalog(); + for scenario in &catalog.scenarios { + let Some(mock) = &scenario.mock else { + continue; + }; + let mut seen = BTreeSet::new(); + for partition in &mock.partitions { + for page in &partition.pages { + for row in &page.rows { + assert!( + seen.insert(row.rid.clone()), + "scenario {}: duplicate mock row _rid {:?}", + scenario.id, + row.rid + ); + } + } + } + } +} + +/// Every distinct expected value must actually appear among the mock payloads, +/// so a typo in `expectedValues` fails here rather than looking like a bug. +#[test] +fn expected_values_appear_among_the_mock_payloads() { + let catalog = load_catalog(); + for scenario in &catalog.scenarios { + let Some(mock) = &scenario.mock else { + continue; + }; + let payloads: Vec<&serde_json::Value> = mock + .partitions + .iter() + .flat_map(|p| p.pages.iter()) + .flat_map(|page| page.rows.iter()) + .map(|row| &row.payload) + .collect(); + for expected in &scenario.expected_values { + assert!( + payloads.contains(&expected), + "scenario {}: expected value {expected} is not produced by any mock row", + scenario.id + ); + } + } +} + +#[test] +fn page_sizes_are_positive() { + let catalog = load_catalog(); + for scenario in &catalog.scenarios { + for size in &scenario.page_sizes { + assert!( + *size > 0, + "scenario {} declares a zero page size", + scenario.id + ); + } + } +} + +#[test] +fn expected_errors_carry_a_category_and_message_fragment() { + let catalog = load_catalog(); + for scenario in &catalog.scenarios { + let Some(error) = &scenario.expected_error else { + continue; + }; + assert!( + !error.category.is_empty() && !error.message_fragment.is_empty(), + "scenario {} has an incomplete expectedError", + scenario.id + ); + } +} + +/// Every scenario must declare at least one layer that actually has a consumer +/// today, otherwise it can pass every schema guard and still never run. +/// +/// `recorded` is reserved for a future live-recording layer and does not count. +#[test] +fn every_scenario_is_claimed_by_a_layer_with_a_consumer() { + const CONSUMED_LAYERS: &[&str] = &["distinctMap", "mockPipeline", "inMemoryEmulator"]; + let catalog = load_catalog(); + for scenario in &catalog.scenarios { + assert!( + scenario + .layers + .iter() + .any(|l| CONSUMED_LAYERS.contains(&l.as_str())), + "scenario {} declares only layers with no runner, so it would never execute: {:?}", + scenario.id, + scenario.layers + ); + } +} + +#[test] +fn required_scenario_inventory_categories_are_represented() { + let catalog = load_catalog(); + let ids: BTreeSet<&str> = catalog.scenarios.iter().map(|s| s.id.as_str()).collect(); + + // One id substring per required inventory category, derived from the .NET + // and Java DISTINCT suites plus the gaps neither peer covers. + let required_markers = [ + "unordered_duplicates_across_partitions", + "unordered_duplicates_across_pages", + "unordered_duplicates_within_page", + "unordered_no_duplicates_passthrough", + "unordered_all_rows_duplicate", + "type_null_and_boolean", + "type_empty_string_array_object", + "numeric_int_vs_float_equal", + "numeric_negative_zero", + "object_key_order_irrelevant", + "array_order_matters", + "wrapped_value_differs_from_bare", + "unicode_strings", + "select_list_multi_column", + "select_star", + "select_value_constant_with_from", + "filters_and_parameters", + "empty_total_result", + "empty_backend_page", + "single_logical_partition", + "ordered_adjacent_dedup", + "ordered_descending", + "ordered_non_adjacent_repeat", + "ordered_run_spanning_pages", + "ordered_resume_suppresses_boundary_row", + "ordered_resume_round_trip", + "distinct_order_by_mismatched_field", + "unordered_continuation_token_rejected", + "malformed_token", + "token_shape_mismatch", + "distinct_type_mismatch_on_resume", + "split_", + "headers_request_charge", + "unsupported_combination", + ]; + for marker in required_markers { + assert!( + ids.iter().any(|id| id.contains(marker)), + "no scenario id contains required inventory marker {marker:?}; catalog ids: {ids:?}" + ); + } +} + +/// Both deduplication modes must be represented, and the peer-untested gaps we +/// deliberately close must stay in the catalog. +#[test] +fn both_distinct_modes_are_represented() { + let catalog = load_catalog(); + for expected in ["Ordered", "Unordered"] { + assert!( + catalog + .scenarios + .iter() + .any(|s| s.query.distinct_type == expected), + "no scenario exercises {expected} DISTINCT" + ); + } +} diff --git a/sdk/cosmos/azure_data_cosmos_driver/tests/fixtures/distinct_scenarios.json b/sdk/cosmos/azure_data_cosmos_driver/tests/fixtures/distinct_scenarios.json new file mode 100644 index 00000000000..dfe2164bd9d --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos_driver/tests/fixtures/distinct_scenarios.json @@ -0,0 +1,2887 @@ +{ + "schemaVersion": 1, + "scenarios": [ + { + "id": "unordered_duplicates_across_partitions", + "description": "The same projected value emitted by two different physical partitions is one DISTINCT row. Neither the backend's per-partition dedup nor a per-page set catches this; only global client-side state does.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.EmulatorTests/Query/DistinctQueryTests.cs", + "test": "TestDistinct_ExecuteNextAsync" + }, + { + "sdk": "java", + "path": "sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/DistinctQueryTests.java", + "test": "queryDocuments" + } + ], + "layers": [ + "mockPipeline", + "inMemoryEmulator" + ], + "query": { + "text": "SELECT DISTINCT VALUE c.city FROM c", + "parameters": [], + "columns": [], + "distinctType": "Unordered" + }, + "documents": [ + { + "id": "d1", + "pk": "pk-a", + "city": "Seattle" + }, + { + "id": "d2", + "pk": "pk-b", + "city": "Redmond" + }, + { + "id": "d3", + "pk": "pk-c", + "city": "Seattle" + }, + { + "id": "d4", + "pk": "pk-d", + "city": "Boston" + } + ], + "mock": { + "partitions": [ + { + "range": { + "minEpk": "", + "maxEpk": "80" + }, + "pages": [ + { + "rows": [ + { + "rid": "p0-r1", + "payload": "Seattle" + }, + { + "rid": "p0-r2", + "payload": "Redmond" + } + ], + "continuation": null + } + ] + }, + { + "range": { + "minEpk": "80", + "maxEpk": "FF" + }, + "pages": [ + { + "rows": [ + { + "rid": "p1-r1", + "payload": "Seattle" + }, + { + "rid": "p1-r2", + "payload": "Boston" + } + ], + "continuation": null + } + ] + } + ] + }, + "pageSizes": [ + 1, + 2, + 10 + ], + "expectedIds": [], + "expectedValues": [ + "Seattle", + "Redmond", + "Boston" + ], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "unordered_duplicates_across_pages", + "description": "A duplicate that straddles a page boundary within one partition must collapse, so the retained state has to outlive a single page.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.Tests/Query/Pipeline/DistinctQueryPipelineStageTests.cs", + "test": "SanityTests" + } + ], + "layers": [ + "mockPipeline" + ], + "query": { + "text": "SELECT DISTINCT VALUE c.item FROM c", + "parameters": [], + "columns": [], + "distinctType": "Unordered" + }, + "documents": [], + "mock": { + "partitions": [ + { + "range": { + "minEpk": "", + "maxEpk": "FF" + }, + "pages": [ + { + "rows": [ + { + "rid": "r1", + "payload": 42 + }, + { + "rid": "r2", + "payload": 1337 + } + ], + "continuation": "page-2" + }, + { + "rows": [ + { + "rid": "r3", + "payload": 1337 + }, + { + "rid": "r4", + "payload": 42 + } + ], + "continuation": null + } + ] + } + ] + }, + "pageSizes": [ + 1, + 3, + 10 + ], + "expectedIds": [], + "expectedValues": [ + 42, + 1337 + ], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "unordered_duplicates_within_page", + "description": "Repeats inside a single backend page collapse. Real Cosmos dedups within a partition first, but a page can still repeat a value when the query fans out over a JOIN.", + "sources": [ + { + "sdk": "java", + "path": "sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/DistinctQueryTests.java", + "test": "queryDistinctDocuments" + } + ], + "layers": [ + "mockPipeline" + ], + "query": { + "text": "SELECT DISTINCT VALUE c.age FROM c", + "parameters": [], + "columns": [], + "distinctType": "Unordered" + }, + "documents": [], + "mock": { + "partitions": [ + { + "range": { + "minEpk": "", + "maxEpk": "FF" + }, + "pages": [ + { + "rows": [ + { + "rid": "r1", + "payload": 1 + }, + { + "rid": "r2", + "payload": 1 + }, + { + "rid": "r3", + "payload": 2 + }, + { + "rid": "r4", + "payload": 1 + } + ], + "continuation": null + } + ] + } + ] + }, + "pageSizes": [ + 10 + ], + "expectedIds": [], + "expectedValues": [ + 1, + 2 + ], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "unordered_no_duplicates_passthrough", + "description": "A stream with nothing repeated passes through untouched, so deduplication never drops a legitimately distinct row.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.EmulatorTests/Query/DistinctQueryTests.cs", + "test": "TestDistinct_ExecuteNextAsync" + } + ], + "layers": [ + "mockPipeline", + "inMemoryEmulator" + ], + "query": { + "text": "SELECT DISTINCT VALUE c.city FROM c", + "parameters": [], + "columns": [], + "distinctType": "Unordered" + }, + "documents": [ + { + "id": "d1", + "pk": "pk-a", + "city": "Seattle" + }, + { + "id": "d2", + "pk": "pk-b", + "city": "Redmond" + }, + { + "id": "d3", + "pk": "pk-c", + "city": "Boston" + } + ], + "mock": { + "partitions": [ + { + "range": { + "minEpk": "", + "maxEpk": "FF" + }, + "pages": [ + { + "rows": [ + { + "rid": "r1", + "payload": "Seattle" + }, + { + "rid": "r2", + "payload": "Redmond" + }, + { + "rid": "r3", + "payload": "Boston" + } + ], + "continuation": null + } + ] + } + ] + }, + "pageSizes": [ + 1, + 10 + ], + "expectedIds": [], + "expectedValues": [ + "Seattle", + "Redmond", + "Boston" + ], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "unordered_all_rows_duplicate_single_output", + "description": "Every row carries the same value, so exactly one survives. Exercises the all-duplicate page suppression path: intermediate pages must not surface as empty pages, but their request charge must still be reported.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.Tests/Query/Pipeline/DistinctQueryPipelineStageTests.cs", + "test": "SanityTests" + } + ], + "layers": [ + "mockPipeline", + "inMemoryEmulator" + ], + "query": { + "text": "SELECT DISTINCT VALUE c.city FROM c", + "parameters": [], + "columns": [], + "distinctType": "Unordered" + }, + "documents": [ + { + "id": "d1", + "pk": "pk-a", + "city": "Seattle" + }, + { + "id": "d2", + "pk": "pk-b", + "city": "Seattle" + }, + { + "id": "d3", + "pk": "pk-c", + "city": "Seattle" + }, + { + "id": "d4", + "pk": "pk-d", + "city": "Seattle" + } + ], + "mock": { + "partitions": [ + { + "range": { + "minEpk": "", + "maxEpk": "FF" + }, + "pages": [ + { + "rows": [ + { + "rid": "r1", + "payload": "Seattle" + } + ], + "continuation": "p2" + }, + { + "rows": [ + { + "rid": "r2", + "payload": "Seattle" + } + ], + "continuation": "p3" + }, + { + "rows": [ + { + "rid": "r3", + "payload": "Seattle" + } + ], + "continuation": null + } + ] + } + ] + }, + "pageSizes": [ + 1, + 10 + ], + "expectedIds": [], + "expectedValues": [ + "Seattle" + ], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "type_null_and_boolean_values", + "description": "null, false, and true are three separate DISTINCT values and never collide with each other. Each carries its own hash type seed.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.Tests/Query/DistinctHashBaselineTests.cs", + "test": "ElementsHash" + }, + { + "sdk": "java", + "path": "sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/query/DistinctHashTest.java", + "test": "booleanHash" + } + ], + "layers": [ + "distinctMap", + "mockPipeline" + ], + "query": { + "text": "SELECT DISTINCT VALUE c.flag FROM c", + "parameters": [], + "columns": [], + "distinctType": "Unordered" + }, + "documents": [], + "mock": { + "partitions": [ + { + "range": { + "minEpk": "", + "maxEpk": "FF" + }, + "pages": [ + { + "rows": [ + { + "rid": "r1", + "payload": null + }, + { + "rid": "r2", + "payload": false + }, + { + "rid": "r3", + "payload": true + }, + { + "rid": "r4", + "payload": null + }, + { + "rid": "r5", + "payload": false + }, + { + "rid": "r6", + "payload": true + } + ], + "continuation": null + } + ] + } + ] + }, + "pageSizes": [ + 10 + ], + "expectedIds": [], + "expectedValues": [ + null, + false, + true + ], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "type_empty_string_array_object_are_distinct", + "description": "The empty string, the empty array, and the empty object are three different values. Their type seeds are the only thing separating them, so a hash that folds only content would collapse all three.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.Tests/Query/DistinctHashBaselineTests.cs", + "test": "ElementsHash" + } + ], + "layers": [ + "distinctMap", + "mockPipeline" + ], + "query": { + "text": "SELECT DISTINCT VALUE c.value FROM c", + "parameters": [], + "columns": [], + "distinctType": "Unordered" + }, + "documents": [], + "mock": { + "partitions": [ + { + "range": { + "minEpk": "", + "maxEpk": "FF" + }, + "pages": [ + { + "rows": [ + { + "rid": "r1", + "payload": "" + }, + { + "rid": "r2", + "payload": [] + }, + { + "rid": "r3", + "payload": {} + }, + { + "rid": "r4", + "payload": null + }, + { + "rid": "r5", + "payload": "" + }, + { + "rid": "r6", + "payload": [] + }, + { + "rid": "r7", + "payload": {} + } + ], + "continuation": null + } + ] + } + ] + }, + "pageSizes": [ + 10 + ], + "expectedIds": [], + "expectedValues": [ + "", + [], + {}, + null + ], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "numeric_int_vs_float_equal", + "description": "5 and 5.0 are the same Cosmos number, so they deduplicate together. Java asserts exactly this against a live account.", + "sources": [ + { + "sdk": "java", + "path": "sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/DistinctQueryTests.java", + "test": "queryDocumentsForDistinctIntValues" + } + ], + "layers": [ + "distinctMap", + "mockPipeline", + "inMemoryEmulator" + ], + "query": { + "text": "SELECT DISTINCT VALUE c.intprop FROM c", + "parameters": [], + "columns": [], + "distinctType": "Unordered" + }, + "documents": [ + { + "id": "d1", + "pk": "pk-a", + "intprop": 5 + }, + { + "id": "d2", + "pk": "pk-b", + "intprop": 5.0 + }, + { + "id": "d3", + "pk": "pk-c", + "intprop": 6 + } + ], + "mock": { + "partitions": [ + { + "range": { + "minEpk": "", + "maxEpk": "FF" + }, + "pages": [ + { + "rows": [ + { + "rid": "r1", + "payload": 5 + }, + { + "rid": "r2", + "payload": 5.0 + }, + { + "rid": "r3", + "payload": 6 + } + ], + "continuation": null + } + ] + } + ] + }, + "pageSizes": [ + 10 + ], + "expectedIds": [], + "expectedValues": [ + 5, + 6 + ], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "numeric_negative_zero_equals_zero", + "description": "-0.0 and 0.0 are numerically equal and must deduplicate. .NET normalizes the sign bit in CosmosNumberHasher but has no test for it; Java would treat them as different. Rust follows .NET.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctHash.cs", + "test": "CosmosNumberHasher (implementation; no peer test exists)" + } + ], + "layers": [ + "distinctMap", + "mockPipeline" + ], + "query": { + "text": "SELECT DISTINCT VALUE c.value FROM c", + "parameters": [], + "columns": [], + "distinctType": "Unordered" + }, + "documents": [], + "mock": { + "partitions": [ + { + "range": { + "minEpk": "", + "maxEpk": "FF" + }, + "pages": [ + { + "rows": [ + { + "rid": "r1", + "payload": -0.0 + }, + { + "rid": "r2", + "payload": 0.0 + }, + { + "rid": "r3", + "payload": 0 + } + ], + "continuation": null + } + ] + } + ] + }, + "pageSizes": [ + 10 + ], + "expectedIds": [], + "expectedValues": [ + -0.0 + ], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "object_key_order_irrelevant", + "description": "Two objects with the same properties in different orders are one DISTINCT value. Both peers fold object properties with XOR precisely so key order cannot matter; only Java tests it.", + "sources": [ + { + "sdk": "java", + "path": "sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/query/DistinctHashTest.java", + "test": "jsonObjectHash" + }, + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctHash.cs", + "test": "Visit(CosmosObject) (implementation; no .NET test exists)" + } + ], + "layers": [ + "distinctMap", + "mockPipeline", + "inMemoryEmulator" + ], + "query": { + "text": "SELECT DISTINCT VALUE c.pet FROM c", + "parameters": [], + "columns": [], + "distinctType": "Unordered" + }, + "documents": [ + { + "id": "d1", + "pk": "pk-a", + "pet": { + "name": "fido", + "species": "dog" + } + }, + { + "id": "d2", + "pk": "pk-b", + "pet": { + "species": "dog", + "name": "fido" + } + }, + { + "id": "d3", + "pk": "pk-c", + "pet": { + "name": "fido", + "species": "cat" + } + } + ], + "mock": { + "partitions": [ + { + "range": { + "minEpk": "", + "maxEpk": "FF" + }, + "pages": [ + { + "rows": [ + { + "rid": "r1", + "payload": { + "name": "fido", + "species": "dog" + } + }, + { + "rid": "r2", + "payload": { + "species": "dog", + "name": "fido" + } + }, + { + "rid": "r3", + "payload": { + "name": "fido", + "species": "cat" + } + } + ], + "continuation": null + } + ] + } + ] + }, + "pageSizes": [ + 10 + ], + "expectedIds": [], + "expectedValues": [ + { + "name": "fido", + "species": "dog" + }, + { + "name": "fido", + "species": "cat" + } + ], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "array_order_matters", + "description": "Arrays are position-sensitive, so [1,2] and [2,1] are two DISTINCT values while a repeat of [1,2] collapses. The mirror image of the object rule above.", + "sources": [ + { + "sdk": "java", + "path": "sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/query/DistinctHashTest.java", + "test": "arrayNodeHash" + }, + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.Tests/Query/DistinctHashBaselineTests.cs", + "test": "ElementsHash" + } + ], + "layers": [ + "distinctMap", + "mockPipeline", + "inMemoryEmulator" + ], + "query": { + "text": "SELECT DISTINCT VALUE c.tags FROM c", + "parameters": [], + "columns": [], + "distinctType": "Unordered" + }, + "documents": [ + { + "id": "d1", + "pk": "pk-a", + "tags": [ + 1, + 2 + ] + }, + { + "id": "d2", + "pk": "pk-b", + "tags": [ + 2, + 1 + ] + }, + { + "id": "d3", + "pk": "pk-c", + "tags": [ + 1, + 2 + ] + } + ], + "mock": { + "partitions": [ + { + "range": { + "minEpk": "", + "maxEpk": "FF" + }, + "pages": [ + { + "rows": [ + { + "rid": "r1", + "payload": [ + 1, + 2 + ] + }, + { + "rid": "r2", + "payload": [ + 2, + 1 + ] + }, + { + "rid": "r3", + "payload": [ + 1, + 2 + ] + } + ], + "continuation": null + } + ] + } + ] + }, + "pageSizes": [ + 10 + ], + "expectedIds": [], + "expectedValues": [ + [ + 1, + 2 + ], + [ + 2, + 1 + ] + ], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "wrapped_value_differs_from_bare", + "description": "A value, that value inside an array, and that value as an object property are three DISTINCT values. Guards against a hash that folds only leaf content and ignores structure.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.Tests/Query/DistinctHashBaselineTests.cs", + "test": "WrappedElementsHash" + } + ], + "layers": [ + "distinctMap", + "mockPipeline" + ], + "query": { + "text": "SELECT DISTINCT VALUE c.value FROM c", + "parameters": [], + "columns": [], + "distinctType": "Unordered" + }, + "documents": [], + "mock": { + "partitions": [ + { + "range": { + "minEpk": "", + "maxEpk": "FF" + }, + "pages": [ + { + "rows": [ + { + "rid": "r1", + "payload": 42 + }, + { + "rid": "r2", + "payload": [ + 42 + ] + }, + { + "rid": "r3", + "payload": { + "prop": 42 + } + }, + { + "rid": "r4", + "payload": 42 + }, + { + "rid": "r5", + "payload": [ + 42 + ] + }, + { + "rid": "r6", + "payload": { + "prop": 42 + } + } + ], + "continuation": null + } + ] + } + ] + }, + "pageSizes": [ + 10 + ], + "expectedIds": [], + "expectedValues": [ + 42, + [ + 42 + ], + { + "prop": 42 + } + ], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "unicode_strings_and_keys", + "description": "CJK values and Arabic object keys round-trip through the hash without collapsing distinct values. .NET's mixed-type corpus carries exactly these.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.Tests/Query/Pipeline/DistinctQueryPipelineStageTests.cs", + "test": "MixedTypeTests" + } + ], + "layers": [ + "distinctMap", + "mockPipeline" + ], + "query": { + "text": "SELECT DISTINCT VALUE c.value FROM c", + "parameters": [], + "columns": [], + "distinctType": "Unordered" + }, + "documents": [], + "mock": { + "partitions": [ + { + "range": { + "minEpk": "", + "maxEpk": "FF" + }, + "pages": [ + { + "rows": [ + { + "rid": "r1", + "payload": "\u654f\u6377\u7684\u68d5\u8272\u72d0\u72f8" + }, + { + "rid": "r2", + "payload": "\u654f\u6377\u7684\u68d5\u8272\u72d0\u72f8" + }, + { + "rid": "r3", + "payload": { + "\u0641\u0648\u0642": 1 + } + }, + { + "rid": "r4", + "payload": { + "\u062a\u062d\u062a": 1 + } + } + ], + "continuation": null + } + ] + } + ] + }, + "pageSizes": [ + 10 + ], + "expectedIds": [], + "expectedValues": [ + "\u654f\u6377\u7684\u68d5\u8272\u72d0\u72f8", + { + "\u0641\u0648\u0642": 1 + }, + { + "\u062a\u062d\u062a": 1 + } + ], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "select_list_multi_column", + "description": "DISTINCT over a multi-column projection deduplicates on the whole projected row, not on any single column. Two rows agreeing on one column but not the other are both kept.", + "sources": [ + { + "sdk": "java", + "path": "sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/DistinctQueryTests.java", + "test": "queryDistinctDocuments" + }, + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.EmulatorTests/Query/DistinctQueryTests.cs", + "test": "TestDistinct_ExecuteNextAsync" + } + ], + "layers": [ + "mockPipeline", + "inMemoryEmulator" + ], + "query": { + "text": "SELECT DISTINCT c.name, c.city FROM c", + "parameters": [], + "columns": [], + "distinctType": "Unordered" + }, + "documents": [ + { + "id": "d1", + "pk": "pk-a", + "name": "ann", + "city": "Seattle" + }, + { + "id": "d2", + "pk": "pk-b", + "name": "ann", + "city": "Boston" + }, + { + "id": "d3", + "pk": "pk-c", + "name": "ann", + "city": "Seattle" + } + ], + "mock": { + "partitions": [ + { + "range": { + "minEpk": "", + "maxEpk": "FF" + }, + "pages": [ + { + "rows": [ + { + "rid": "r1", + "payload": { + "name": "ann", + "city": "Seattle" + } + }, + { + "rid": "r2", + "payload": { + "name": "ann", + "city": "Boston" + } + }, + { + "rid": "r3", + "payload": { + "name": "ann", + "city": "Seattle" + } + } + ], + "continuation": null + } + ] + } + ] + }, + "pageSizes": [ + 10 + ], + "expectedIds": [], + "expectedValues": [ + { + "name": "ann", + "city": "Seattle" + }, + { + "name": "ann", + "city": "Boston" + } + ], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "select_star_whole_document", + "description": "SELECT DISTINCT * deduplicates whole documents. Since every document carries a unique id and system properties, nothing collapses \u2014 the guard is that DISTINCT does not corrupt the passthrough.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.EmulatorTests/Query/DistinctQueryTests.cs", + "test": "TestDistinct_ExecuteNextAsync" + }, + { + "sdk": "java", + "path": "sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/DistinctQueryTests.java", + "test": "queryDistinctDocuments" + } + ], + "layers": [ + "inMemoryEmulator" + ], + "query": { + "text": "SELECT DISTINCT * FROM c", + "parameters": [], + "columns": [], + "distinctType": "Unordered" + }, + "documents": [ + { + "id": "d1", + "pk": "pk-a", + "city": "Seattle" + }, + { + "id": "d2", + "pk": "pk-b", + "city": "Seattle" + }, + { + "id": "d3", + "pk": "pk-c", + "city": "Boston" + } + ], + "mock": null, + "pageSizes": [ + 10 + ], + "expectedIds": [ + "d1", + "d2", + "d3" + ], + "expectedValues": [], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "select_value_constant_with_from_still_deduplicates", + "description": "DISTINCT over a constant collapses to distinctType None only when there is no FROM clause. With a FROM the query yields one row per document, so deduplication is real work and the service reports Unordered \u2014 measured live. Every row is the same constant, so exactly one survives.", + "sources": [ + { + "sdk": "java", + "path": "sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/DistinctQueryTests.java", + "test": "queryDistinctDocuments" + } + ], + "layers": [ + "inMemoryEmulator" + ], + "query": { + "text": "SELECT DISTINCT VALUE 1 FROM c", + "parameters": [], + "columns": [], + "distinctType": "Unordered" + }, + "documents": [ + { + "id": "d1", + "pk": "pk-a" + }, + { + "id": "d2", + "pk": "pk-b" + } + ], + "mock": null, + "pageSizes": [ + 10 + ], + "expectedIds": [], + "expectedValues": [ + 1 + ], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "filters_and_parameters_apply_before_distinct", + "description": "A parameterized WHERE clause narrows the rows before deduplication, so DISTINCT never sees the filtered-out values.", + "sources": [ + { + "sdk": "java", + "path": "sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/DistinctQueryTests.java", + "test": "queryDistinctDocuments" + } + ], + "layers": [ + "inMemoryEmulator" + ], + "query": { + "text": "SELECT DISTINCT VALUE c.city FROM c WHERE c.active = @active", + "parameters": [ + { + "name": "@active", + "value": true + } + ], + "columns": [], + "distinctType": "Unordered" + }, + "documents": [ + { + "id": "d1", + "pk": "pk-a", + "city": "Seattle", + "active": true + }, + { + "id": "d2", + "pk": "pk-b", + "city": "Seattle", + "active": true + }, + { + "id": "d3", + "pk": "pk-c", + "city": "Boston", + "active": false + } + ], + "mock": null, + "pageSizes": [ + 10 + ], + "expectedIds": [], + "expectedValues": [ + "Seattle" + ], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "empty_total_result", + "description": "A DISTINCT query matching no documents drains cleanly with no rows rather than hanging or erroring.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.EmulatorTests/Query/DistinctQueryTests.cs", + "test": "TestDistinct_ExecuteNextAsync" + } + ], + "layers": [ + "mockPipeline", + "inMemoryEmulator" + ], + "query": { + "text": "SELECT DISTINCT VALUE c.city FROM c WHERE c.city = 'nowhere'", + "parameters": [], + "columns": [], + "distinctType": "Unordered" + }, + "documents": [ + { + "id": "d1", + "pk": "pk-a", + "city": "Seattle" + } + ], + "mock": { + "partitions": [ + { + "range": { + "minEpk": "", + "maxEpk": "FF" + }, + "pages": [ + { + "rows": [], + "continuation": null + } + ] + } + ] + }, + "pageSizes": [ + 10 + ], + "expectedIds": [], + "expectedValues": [], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "empty_backend_page_mid_stream", + "description": "An empty intermediate backend page (a partition whose rows were all filtered out) must not terminate the merge or be mistaken for an all-duplicate page.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.Tests/Query/Pipeline/DistinctQueryPipelineStageTests.cs", + "test": "SanityTests" + } + ], + "layers": [ + "mockPipeline" + ], + "query": { + "text": "SELECT DISTINCT VALUE c.city FROM c", + "parameters": [], + "columns": [], + "distinctType": "Unordered" + }, + "documents": [], + "mock": { + "partitions": [ + { + "range": { + "minEpk": "", + "maxEpk": "FF" + }, + "pages": [ + { + "rows": [ + { + "rid": "r1", + "payload": "Seattle" + } + ], + "continuation": "p2" + }, + { + "rows": [], + "continuation": "p3" + }, + { + "rows": [ + { + "rid": "r2", + "payload": "Boston" + } + ], + "continuation": null + } + ] + } + ] + }, + "pageSizes": [ + 10 + ], + "expectedIds": [], + "expectedValues": [ + "Seattle", + "Boston" + ], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "single_logical_partition_scope", + "description": "A DISTINCT query pinned to one logical partition still deduplicates; it just never fans out.", + "sources": [ + { + "sdk": "java", + "path": "sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/DistinctQueryTests.java", + "test": "queryDocuments" + } + ], + "layers": [ + "inMemoryEmulator" + ], + "query": { + "text": "SELECT DISTINCT VALUE c.city FROM c WHERE c.pk = 'solo'", + "parameters": [], + "columns": [], + "distinctType": "Unordered" + }, + "documents": [ + { + "id": "d1", + "pk": "solo", + "city": "Seattle" + }, + { + "id": "d2", + "pk": "solo", + "city": "Seattle" + }, + { + "id": "d3", + "pk": "solo", + "city": "Boston" + } + ], + "mock": null, + "pageSizes": [ + 10 + ], + "expectedIds": [], + "expectedValues": [ + "Seattle", + "Boston" + ], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "ordered_adjacent_dedup", + "description": "With a matching ORDER BY the merge delivers equal values adjacently, so a single retained hash is enough. This is the mode that supports continuation.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.EmulatorTests/Query/DistinctQueryTests.cs", + "test": "TestDistinct_ContinuationTokenSupportAsync" + }, + { + "sdk": "java", + "path": "sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/DistinctQueryTests.java", + "test": "queryDocumentsWithOrderBy" + } + ], + "layers": [ + "distinctMap", + "mockPipeline", + "inMemoryEmulator" + ], + "query": { + "text": "SELECT DISTINCT VALUE c.city FROM c ORDER BY c.city ASC", + "parameters": [], + "columns": [ + { + "expression": "c.city", + "direction": "Ascending" + } + ], + "distinctType": "Ordered" + }, + "documents": [ + { + "id": "d1", + "pk": "pk-a", + "city": "Seattle" + }, + { + "id": "d2", + "pk": "pk-b", + "city": "Redmond" + }, + { + "id": "d3", + "pk": "pk-c", + "city": "Seattle" + }, + { + "id": "d4", + "pk": "pk-d", + "city": "Boston" + }, + { + "id": "d5", + "pk": "pk-e", + "city": "Redmond" + } + ], + "mock": { + "partitions": [ + { + "range": { + "minEpk": "", + "maxEpk": "FF" + }, + "pages": [ + { + "rows": [ + { + "rid": "r1", + "payload": "Boston" + }, + { + "rid": "r2", + "payload": "Redmond" + }, + { + "rid": "r3", + "payload": "Redmond" + }, + { + "rid": "r4", + "payload": "Seattle" + }, + { + "rid": "r5", + "payload": "Seattle" + } + ], + "continuation": null + } + ] + } + ] + }, + "pageSizes": [ + 1, + 2, + 10 + ], + "expectedIds": [], + "expectedValues": [ + "Boston", + "Redmond", + "Seattle" + ], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "ordered_descending_adjacent_dedup", + "description": "Descending ORDER BY groups equal values into runs just as ascending does; the map only cares about adjacency, not direction.", + "sources": [ + { + "sdk": "java", + "path": "sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/TestSuiteBase.java", + "test": "queryWithOrderByProvider" + } + ], + "layers": [ + "mockPipeline", + "inMemoryEmulator" + ], + "query": { + "text": "SELECT DISTINCT VALUE c.city FROM c ORDER BY c.city DESC", + "parameters": [], + "columns": [ + { + "expression": "c.city", + "direction": "Descending" + } + ], + "distinctType": "Ordered" + }, + "documents": [ + { + "id": "d1", + "pk": "pk-a", + "city": "Seattle" + }, + { + "id": "d2", + "pk": "pk-b", + "city": "Boston" + }, + { + "id": "d3", + "pk": "pk-c", + "city": "Seattle" + } + ], + "mock": { + "partitions": [ + { + "range": { + "minEpk": "", + "maxEpk": "FF" + }, + "pages": [ + { + "rows": [ + { + "rid": "r1", + "payload": "Seattle" + }, + { + "rid": "r2", + "payload": "Seattle" + }, + { + "rid": "r3", + "payload": "Boston" + } + ], + "continuation": null + } + ] + } + ] + }, + "pageSizes": [ + 10 + ], + "expectedIds": [], + "expectedValues": [ + "Seattle", + "Boston" + ], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "ordered_non_adjacent_repeat_is_not_deduped", + "description": "Documents the adjacency assumption rather than a desirable behavior: an ordered map deliberately keeps a repeat separated by a different value. Only reachable if the upstream stream is not actually sorted, which the planner prevents by choosing Ordered solely when the plan reports it.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctMap.OrderedDistinctMap.cs", + "test": "OrderedDistinctMap.Add (implementation; untested in both peers)" + } + ], + "layers": [ + "distinctMap", + "mockPipeline" + ], + "query": { + "text": "SELECT DISTINCT VALUE c.city FROM c ORDER BY c.city ASC", + "parameters": [], + "columns": [ + { + "expression": "c.city", + "direction": "Ascending" + } + ], + "distinctType": "Ordered" + }, + "documents": [], + "mock": { + "partitions": [ + { + "range": { + "minEpk": "", + "maxEpk": "FF" + }, + "pages": [ + { + "rows": [ + { + "rid": "r1", + "payload": "Seattle" + }, + { + "rid": "r2", + "payload": "Boston" + }, + { + "rid": "r3", + "payload": "Seattle" + } + ], + "continuation": null + } + ] + } + ] + }, + "pageSizes": [ + 10 + ], + "expectedIds": [], + "expectedValues": [ + "Seattle", + "Boston", + "Seattle" + ], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "ordered_run_spanning_pages", + "description": "A run of equal values split across two emitted pages must not re-emit its head on the second page. Exercises the map state surviving a page boundary in process.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.EmulatorTests/Query/DistinctQueryTests.cs", + "test": "TestDistinct_ContinuationTokenSupportAsync" + } + ], + "layers": [ + "mockPipeline" + ], + "query": { + "text": "SELECT DISTINCT VALUE c.city FROM c ORDER BY c.city ASC", + "parameters": [], + "columns": [ + { + "expression": "c.city", + "direction": "Ascending" + } + ], + "distinctType": "Ordered" + }, + "documents": [], + "mock": { + "partitions": [ + { + "range": { + "minEpk": "", + "maxEpk": "FF" + }, + "pages": [ + { + "rows": [ + { + "rid": "r1", + "payload": "Redmond" + } + ], + "continuation": "p2" + }, + { + "rows": [ + { + "rid": "r2", + "payload": "Redmond" + }, + { + "rid": "r3", + "payload": "Seattle" + } + ], + "continuation": null + } + ] + } + ] + }, + "pageSizes": [ + 10 + ], + "expectedIds": [], + "expectedValues": [ + "Redmond", + "Seattle" + ], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "ordered_resume_suppresses_boundary_row", + "description": "Resuming from a checkpoint whose last emitted value was Redmond must drop the re-delivered Redmond and emit only what follows. Without the persisted hash the caller would see Redmond twice.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.EmulatorTests/Query/DistinctQueryTests.cs", + "test": "TestDistinct_ContinuationTokenSupportAsync" + }, + { + "sdk": "java", + "path": "sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/DistinctQueryTests.java", + "test": "queryDocumentsWithOrderBy" + } + ], + "layers": [ + "mockPipeline" + ], + "query": { + "text": "SELECT DISTINCT VALUE c.city FROM c ORDER BY c.city ASC", + "parameters": [], + "columns": [ + { + "expression": "c.city", + "direction": "Ascending" + } + ], + "distinctType": "Ordered" + }, + "documents": [], + "mock": { + "partitions": [ + { + "range": { + "minEpk": "", + "maxEpk": "FF" + }, + "pages": [ + { + "rows": [ + { + "rid": "r2", + "payload": "Redmond" + }, + { + "rid": "r3", + "payload": "Seattle" + } + ], + "continuation": null + } + ] + } + ] + }, + "pageSizes": [ + 10 + ], + "expectedIds": [], + "expectedValues": [ + "Seattle" + ], + "checkpoint": { + "lastValue": "Redmond" + }, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "ordered_resume_round_trip_matches_full_drain", + "description": "Draining an ordered DISTINCT query in one pass and draining it across a serialized continuation token must produce the same rows in the same order.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.EmulatorTests/Query/DistinctQueryTests.cs", + "test": "TestDistinct_CosmosElementContinuationTokenAsync" + } + ], + "layers": [ + "inMemoryEmulator" + ], + "query": { + "text": "SELECT DISTINCT VALUE c.city FROM c ORDER BY c.city ASC", + "parameters": [], + "columns": [ + { + "expression": "c.city", + "direction": "Ascending" + } + ], + "distinctType": "Ordered" + }, + "documents": [ + { + "id": "d1", + "pk": "pk-a", + "city": "Seattle" + }, + { + "id": "d2", + "pk": "pk-b", + "city": "Redmond" + }, + { + "id": "d3", + "pk": "pk-c", + "city": "Seattle" + }, + { + "id": "d4", + "pk": "pk-d", + "city": "Boston" + }, + { + "id": "d5", + "pk": "pk-e", + "city": "Redmond" + }, + { + "id": "d6", + "pk": "pk-f", + "city": "Austin" + } + ], + "mock": null, + "pageSizes": [ + 1, + 2, + 10 + ], + "expectedIds": [], + "expectedValues": [ + "Austin", + "Boston", + "Redmond", + "Seattle" + ], + "checkpoint": { + "resumeAfterPage": 1 + }, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "distinct_order_by_mismatched_field_is_unordered", + "description": "When the sort key is not a prefix of the DISTINCT projection the Gateway reports Unordered, because sorting on a different field cannot group equal projections into runs. Java pins the pair and expects HTTP 400 when such a query is resumed.", + "sources": [ + { + "sdk": "java", + "path": "sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/TestSuiteBase.java", + "test": "queryWithOrderByProvider" + } + ], + "layers": [ + "inMemoryEmulator" + ], + "query": { + "text": "SELECT DISTINCT VALUE c.city FROM c ORDER BY c.rank ASC", + "parameters": [], + "columns": [ + { + "expression": "c.rank", + "direction": "Ascending" + } + ], + "distinctType": "Unordered" + }, + "documents": [ + { + "id": "d1", + "pk": "pk-a", + "city": "Seattle", + "rank": 3 + }, + { + "id": "d2", + "pk": "pk-b", + "city": "Boston", + "rank": 1 + }, + { + "id": "d3", + "pk": "pk-c", + "city": "Seattle", + "rank": 2 + } + ], + "mock": null, + "pageSizes": [ + 10 + ], + "expectedIds": [], + "expectedValues": [ + "Boston", + "Seattle" + ], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "unordered_continuation_token_rejected", + "description": "Asking for a continuation token on an unordered DISTINCT query fails at mint time with actionable guidance, rather than returning a token whose resume would re-emit rows. .NET and Java both refuse here.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/DisallowContinuationTokenMessages.cs", + "test": "DisallowContinuationTokenMessages.Distinct" + }, + { + "sdk": "java", + "path": "sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DistinctDocumentQueryExecutionContext.java", + "test": "createAsync (BadRequestException on lastHash for a non-ordered query)" + } + ], + "layers": [ + "mockPipeline", + "inMemoryEmulator" + ], + "query": { + "text": "SELECT DISTINCT VALUE c.city FROM c", + "parameters": [], + "columns": [], + "distinctType": "Unordered" + }, + "documents": [ + { + "id": "d1", + "pk": "pk-a", + "city": "Seattle" + }, + { + "id": "d2", + "pk": "pk-b", + "city": "Boston" + } + ], + "mock": { + "partitions": [ + { + "range": { + "minEpk": "", + "maxEpk": "FF" + }, + "pages": [ + { + "rows": [ + { + "rid": "r1", + "payload": "Seattle" + } + ], + "continuation": "p2" + } + ] + } + ] + }, + "pageSizes": [ + 1 + ], + "expectedIds": [], + "expectedValues": [], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": { + "category": "clientDistinctContinuationUnsupported", + "messageFragment": "ORDER BY" + } + }, + { + "id": "malformed_token_rejected", + "description": "A continuation token that is not valid for this operation is rejected rather than silently ignored, which would restart the query and re-emit every row.", + "sources": [ + { + "sdk": "java", + "path": "sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DistinctContinuationToken.java", + "test": "tryParse" + }, + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.cs", + "test": "DistinctContinuationToken.TryParse" + } + ], + "layers": [ + "mockPipeline" + ], + "query": { + "text": "SELECT DISTINCT VALUE c.city FROM c ORDER BY c.city ASC", + "parameters": [], + "columns": [ + { + "expression": "c.city", + "direction": "Ascending" + } + ], + "distinctType": "Ordered" + }, + "documents": [], + "mock": null, + "pageSizes": [ + 10 + ], + "expectedIds": [], + "expectedValues": [], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": { + "category": "clientContinuationTokenShapeMismatch", + "messageFragment": "does not match" + } + }, + { + "id": "token_shape_mismatch_rejected", + "description": "A token minted for a non-DISTINCT query (or before DISTINCT support existed) must not resume into a DISTINCT plan, since that would skip deduplication entirely for the remaining pages.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.cs", + "test": "DistinctContinuationToken.TryParse" + } + ], + "layers": [ + "mockPipeline" + ], + "query": { + "text": "SELECT DISTINCT VALUE c.city FROM c ORDER BY c.city ASC", + "parameters": [], + "columns": [ + { + "expression": "c.city", + "direction": "Ascending" + } + ], + "distinctType": "Ordered" + }, + "documents": [], + "mock": null, + "pageSizes": [ + 10 + ], + "expectedIds": [], + "expectedValues": [], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": { + "category": "clientContinuationTokenShapeMismatch", + "messageFragment": "does not match a DISTINCT query" + } + }, + { + "id": "distinct_type_mismatch_on_resume_rejected", + "description": "A token minted for an ordered DISTINCT must not resume a query the plan now reports as unordered (or vice versa). Reinterpreting it would apply adjacency deduplication to an unsorted stream.", + "sources": [ + { + "sdk": "java", + "path": "sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DistinctDocumentQueryExecutionContext.java", + "test": "createAsync" + } + ], + "layers": [ + "mockPipeline" + ], + "query": { + "text": "SELECT DISTINCT VALUE c.city FROM c", + "parameters": [], + "columns": [], + "distinctType": "Unordered" + }, + "documents": [], + "mock": null, + "pageSizes": [ + 10 + ], + "expectedIds": [], + "expectedValues": [], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": { + "category": "clientContinuationTokenShapeMismatch", + "messageFragment": "minted for" + } + }, + { + "id": "split_mid_query_unordered_does_not_reemit", + "description": "A partition split partway through an unordered DISTINCT query must not resurrect an already-emitted value. The fan-out root absorbs the split internally, so the deduplication state is never rebuilt.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.Tests/Query/Pipeline/FullPipelineTests.cs", + "test": "TestMerge (split coverage exists for ORDER BY only; DISTINCT is untested in both peers)" + } + ], + "layers": [ + "inMemoryEmulator" + ], + "query": { + "text": "SELECT DISTINCT VALUE c.city FROM c", + "parameters": [], + "columns": [], + "distinctType": "Unordered" + }, + "documents": [ + { + "id": "d1", + "pk": "pk-a", + "city": "Seattle" + }, + { + "id": "d2", + "pk": "pk-b", + "city": "Redmond" + }, + { + "id": "d3", + "pk": "pk-c", + "city": "Seattle" + }, + { + "id": "d4", + "pk": "pk-d", + "city": "Boston" + } + ], + "mock": null, + "pageSizes": [ + 1 + ], + "expectedIds": [], + "expectedValues": [ + "Boston", + "Redmond", + "Seattle" + ], + "checkpoint": { + "splitAfterPage": 1 + }, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "split_mid_query_ordered_does_not_reemit", + "description": "Same guarantee for ordered DISTINCT: a split mid-drain re-resolves ranges, but last_hash is a value rather than a position, so the boundary row is still suppressed afterwards.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.Tests/Query/Pipeline/FullPipelineTests.cs", + "test": "TestMerge (split coverage exists for ORDER BY only; DISTINCT is untested in both peers)" + } + ], + "layers": [ + "inMemoryEmulator" + ], + "query": { + "text": "SELECT DISTINCT VALUE c.city FROM c ORDER BY c.city ASC", + "parameters": [], + "columns": [ + { + "expression": "c.city", + "direction": "Ascending" + } + ], + "distinctType": "Ordered" + }, + "documents": [ + { + "id": "d1", + "pk": "pk-a", + "city": "Seattle" + }, + { + "id": "d2", + "pk": "pk-b", + "city": "Redmond" + }, + { + "id": "d3", + "pk": "pk-c", + "city": "Seattle" + }, + { + "id": "d4", + "pk": "pk-d", + "city": "Boston" + } + ], + "mock": null, + "pageSizes": [ + 1 + ], + "expectedIds": [], + "expectedValues": [ + "Boston", + "Redmond", + "Seattle" + ], + "checkpoint": { + "splitAfterPage": 1 + }, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "headers_request_charge_survives_suppressed_pages", + "description": "Backend pages whose rows are all duplicates are never surfaced, but the request units they cost must still be reported on the next emitted page. Otherwise a highly redundant DISTINCT query would appear far cheaper than it was.", + "sources": [ + { + "sdk": "java", + "path": "sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/DistinctQueryTests.java", + "test": "queryDocuments" + } + ], + "layers": [ + "mockPipeline", + "inMemoryEmulator" + ], + "query": { + "text": "SELECT DISTINCT VALUE c.city FROM c", + "parameters": [], + "columns": [], + "distinctType": "Unordered" + }, + "documents": [ + { + "id": "d1", + "pk": "pk-a", + "city": "Seattle" + }, + { + "id": "d2", + "pk": "pk-b", + "city": "Seattle" + }, + { + "id": "d3", + "pk": "pk-c", + "city": "Seattle" + }, + { + "id": "d4", + "pk": "pk-d", + "city": "Boston" + } + ], + "mock": { + "partitions": [ + { + "range": { + "minEpk": "", + "maxEpk": "FF" + }, + "pages": [ + { + "rows": [ + { + "rid": "r1", + "payload": "Seattle" + } + ], + "continuation": "p2" + }, + { + "rows": [ + { + "rid": "r2", + "payload": "Seattle" + } + ], + "continuation": "p3" + }, + { + "rows": [ + { + "rid": "r3", + "payload": "Boston" + } + ], + "continuation": null + } + ] + } + ] + }, + "pageSizes": [ + 1 + ], + "expectedIds": [], + "expectedValues": [ + "Seattle", + "Boston" + ], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "unsupported_combination_group_by", + "description": "DISTINCT combined with GROUP BY is not yet supported and must fail with a typed unsupported-feature error rather than silently dropping one of the two stages.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs", + "test": "MonadicCreate (GROUP BY composes above DISTINCT)" + } + ], + "layers": [ + "inMemoryEmulator" + ], + "query": { + "text": "SELECT DISTINCT c.city FROM c GROUP BY c.city", + "parameters": [], + "columns": [], + "distinctType": "Unordered" + }, + "documents": [ + { + "id": "d1", + "pk": "pk-a", + "city": "Seattle" + }, + { + "id": "d2", + "pk": "pk-b", + "city": "Boston" + } + ], + "mock": null, + "pageSizes": [ + 10 + ], + "expectedIds": [], + "expectedValues": [], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": { + "category": "clientUnsupportedQueryFeature", + "messageFragment": "GROUP BY" + } + }, + { + "id": "combination_top_applies_after_distinct", + "description": "DISTINCT composes below TOP: the row limit counts deduplicated values, not raw rows. The ordered stream repeats its first value, so applying TOP before deduplication would return one value instead of two \u2014 this pins the DISTINCT-inside-SkipTake pipeline order.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.EmulatorTests/Query/DistinctQueryTests.cs", + "test": "TestDistinct_ExecuteNextAsync" + }, + { + "sdk": "java", + "path": "sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/DistinctQueryTests.java", + "test": "queryDistinctDocuments" + } + ], + "layers": [ + "inMemoryEmulator" + ], + "query": { + "text": "SELECT DISTINCT TOP 2 VALUE c.city FROM c ORDER BY c.city", + "parameters": [], + "columns": [ + { + "expression": "c.city", + "direction": "Ascending" + } + ], + "distinctType": "Ordered" + }, + "documents": [ + { + "id": "d1", + "pk": "pk-a", + "city": "Boston" + }, + { + "id": "d2", + "pk": "pk-b", + "city": "Boston" + }, + { + "id": "d3", + "pk": "pk-c", + "city": "Portland" + }, + { + "id": "d4", + "pk": "pk-d", + "city": "Seattle" + } + ], + "mock": null, + "pageSizes": [ + 1, + 10 + ], + "expectedIds": [], + "expectedValues": [ + "Boston", + "Portland" + ], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "unsupported_combination_aggregate", + "description": "DISTINCT combined with an aggregate is not yet supported; aggregates compose below DISTINCT and need their own stage first.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.EmulatorTests/Query/DistinctQueryTests.cs", + "test": "TestDistinct_ExecuteNextAsync" + }, + { + "sdk": "java", + "path": "sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/DistinctQueryTests.java", + "test": "queryDistinctDocuments" + } + ], + "layers": [ + "inMemoryEmulator" + ], + "query": { + "text": "SELECT DISTINCT VALUE MAX(c.rank) FROM c", + "parameters": [], + "columns": [], + "distinctType": "Unordered" + }, + "documents": [ + { + "id": "d1", + "pk": "pk-a", + "rank": 1 + }, + { + "id": "d2", + "pk": "pk-b", + "rank": 2 + } + ], + "mock": null, + "pageSizes": [ + 10 + ], + "expectedIds": [], + "expectedValues": [], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": { + "category": "clientUnsupportedQueryFeature", + "messageFragment": "aggregates" + } + }, + { + "id": "ordered_offset_limit_applies_after_distinct", + "description": "Ordered DISTINCT under OFFSET/LIMIT: the window counts deduplicated values. The sorted stream repeats Austin and Seattle, so applying the window before deduplication would return a different (shorter or shifted) set. Closes the ordered OFFSET/LIMIT gap \u2014 TOP was covered but OFFSET/LIMIT was not.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.EmulatorTests/Query/DistinctQueryTests.cs", + "test": "TestDistinct_ExecuteNextAsync" + } + ], + "layers": [ + "inMemoryEmulator" + ], + "query": { + "text": "SELECT DISTINCT VALUE c.city FROM c ORDER BY c.city ASC OFFSET 1 LIMIT 2", + "parameters": [], + "columns": [ + { + "expression": "c.city", + "direction": "Ascending" + } + ], + "distinctType": "Ordered" + }, + "documents": [ + { + "id": "d1", + "pk": "pk-a", + "city": "Seattle" + }, + { + "id": "d2", + "pk": "pk-b", + "city": "Austin" + }, + { + "id": "d3", + "pk": "pk-c", + "city": "Boston" + }, + { + "id": "d4", + "pk": "pk-d", + "city": "Austin" + }, + { + "id": "d5", + "pk": "pk-e", + "city": "Portland" + }, + { + "id": "d6", + "pk": "pk-f", + "city": "Seattle" + } + ], + "mock": null, + "pageSizes": [ + 1, + 10 + ], + "expectedIds": [], + "expectedValues": [ + "Boston", + "Portland" + ], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "ordered_distinct_top_resume_round_trip", + "description": "A continuation token for an ordered DISTINCT query under TOP nests SkipTake above Distinct. Draining page-by-page across serialized tokens must match a single drain: the window's remaining budget and the dedup hash have to be peeled and restored in the right order, or rows are dropped or repeated.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.EmulatorTests/Query/DistinctQueryTests.cs", + "test": "TestDistinct_CosmosElementContinuationTokenAsync" + } + ], + "layers": [ + "inMemoryEmulator" + ], + "query": { + "text": "SELECT DISTINCT TOP 3 VALUE c.city FROM c ORDER BY c.city ASC", + "parameters": [], + "columns": [ + { + "expression": "c.city", + "direction": "Ascending" + } + ], + "distinctType": "Ordered" + }, + "documents": [ + { + "id": "d1", + "pk": "pk-a", + "city": "Seattle" + }, + { + "id": "d2", + "pk": "pk-b", + "city": "Austin" + }, + { + "id": "d3", + "pk": "pk-c", + "city": "Boston" + }, + { + "id": "d4", + "pk": "pk-d", + "city": "Austin" + }, + { + "id": "d5", + "pk": "pk-e", + "city": "Portland" + }, + { + "id": "d6", + "pk": "pk-f", + "city": "Seattle" + } + ], + "mock": null, + "pageSizes": [ + 1, + 2 + ], + "expectedIds": [], + "expectedValues": [ + "Austin", + "Boston", + "Portland" + ], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "ordered_distinct_offset_limit_resume_round_trip", + "description": "Same nested-token round trip as the TOP case, but with OFFSET/LIMIT so the resumed token must carry a partially-consumed skip as well as a remaining take. A token minted mid-skip that loses its remaining offset re-emits rows the first pass already skipped past.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.EmulatorTests/Query/DistinctQueryTests.cs", + "test": "TestDistinct_CosmosElementContinuationTokenAsync" + } + ], + "layers": [ + "inMemoryEmulator" + ], + "query": { + "text": "SELECT DISTINCT VALUE c.city FROM c ORDER BY c.city ASC OFFSET 1 LIMIT 2", + "parameters": [], + "columns": [ + { + "expression": "c.city", + "direction": "Ascending" + } + ], + "distinctType": "Ordered" + }, + "documents": [ + { + "id": "d1", + "pk": "pk-a", + "city": "Seattle" + }, + { + "id": "d2", + "pk": "pk-b", + "city": "Austin" + }, + { + "id": "d3", + "pk": "pk-c", + "city": "Boston" + }, + { + "id": "d4", + "pk": "pk-d", + "city": "Austin" + }, + { + "id": "d5", + "pk": "pk-e", + "city": "Portland" + }, + { + "id": "d6", + "pk": "pk-f", + "city": "Seattle" + } + ], + "mock": null, + "pageSizes": [ + 1, + 2 + ], + "expectedIds": [], + "expectedValues": [ + "Boston", + "Portland" + ], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "split_mid_query_windowed_distinct_does_not_reemit", + "description": "A partition split mid-drain while DISTINCT runs under an OFFSET/LIMIT window. The fan-out root absorbs the split and neither the dedup state nor the window's remaining budget may be rebuilt: losing the former re-emits a value the window already paid for, losing the latter restarts the offset and over-returns. Closes the gap where split coverage existed for plain DISTINCT but not for DISTINCT composed under a window.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.EmulatorTests/Query/DistinctQueryTests.cs", + "test": "TestDistinct_ExecuteNextAsync" + } + ], + "layers": [ + "inMemoryEmulator" + ], + "query": { + "text": "SELECT DISTINCT VALUE c.city FROM c ORDER BY c.city ASC OFFSET 1 LIMIT 2", + "parameters": [], + "columns": [ + { + "expression": "c.city", + "direction": "Ascending" + } + ], + "distinctType": "Ordered" + }, + "documents": [ + { + "id": "d1", + "pk": "pk-a", + "city": "Seattle" + }, + { + "id": "d2", + "pk": "pk-b", + "city": "Austin" + }, + { + "id": "d3", + "pk": "pk-c", + "city": "Boston" + }, + { + "id": "d4", + "pk": "pk-d", + "city": "Austin" + }, + { + "id": "d5", + "pk": "pk-e", + "city": "Portland" + }, + { + "id": "d6", + "pk": "pk-f", + "city": "Seattle" + } + ], + "mock": null, + "pageSizes": [ + 1 + ], + "expectedIds": [], + "expectedValues": [ + "Boston", + "Portland" + ], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": null + } + ] +} diff --git a/sdk/cosmos/azure_data_cosmos_driver/tests/gateway_query_plan_comparison.rs b/sdk/cosmos/azure_data_cosmos_driver/tests/gateway_query_plan_comparison.rs index d218f9bbaaf..b0ba28add90 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/tests/gateway_query_plan_comparison.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/tests/gateway_query_plan_comparison.rs @@ -178,13 +178,19 @@ async fn fetch_gateway_plan( fn compare_query_info(sql: &str, local: &serde_json::Value, gw: &serde_json::Value) { let gw_rewritten = gw.get("rewrittenQuery").and_then(|v| v.as_str()); - // ── distinctType ───────────────────────────────────────────────────────── - // Carve-out: Gateway downgrades `Ordered` → `Unordered` whenever it emits a - // `rewrittenQuery`. This is because the rewritten plan uses an explicit ORDER - // BY in the per-partition queries, so the cross-partition aggregation no longer - // needs to preserve order at the DISTINCT layer. Local AST analysis does not - // perform that rewrite, so it correctly reports `Ordered`. This is consistent - // with how the .NET / Java SDKs treat the field. + // ── distinctType (no carve-out) ────────────────────────────────────────── + // This previously carried a carve-out tolerating `local = Ordered` against + // `gw = Unordered` whenever the Gateway emitted a `rewrittenQuery`, on the + // theory that the rewrite made ordering unnecessary at the DISTINCT layer. + // That explanation was too broad: measured against a live account with + // production's `SUPPORTED_QUERY_FEATURES`, the service keeps `Ordered` for + // `SELECT DISTINCT VALUE c.name FROM c ORDER BY c.name ASC` *and* emits a + // 172-character `rewrittenQuery`. What it actually downgrades is every + // other shape — see `gw_distinct` below and `plan::distinct_is_ordered`. + // + // The local generator now encodes that rule, so it agrees with the service + // on every query in this file and the tolerance is unreachable — verified + // by instrumenting the branch and running the full suite live (zero hits). let local_dt = local .get("distinctType") .and_then(|v| v.as_str()) @@ -193,9 +199,7 @@ fn compare_query_info(sql: &str, local: &serde_json::Value, gw: &serde_json::Val .get("distinctType") .and_then(|v| v.as_str()) .unwrap_or("None"); - if !(local_dt == gw_dt - || (local_dt == "Ordered" && gw_dt == "Unordered" && gw_rewritten.is_some())) - { + if local_dt != gw_dt { panic!("[distinctType] sql={sql}\n local={local_dt} gw={gw_dt}"); } @@ -792,6 +796,23 @@ async fn gw_distinct() { validate_pk("SELECT DISTINCT VALUE null").await; validate_pk("SELECT DISTINCT VALUE 1").await; validate_pk("SELECT DISTINCT VALUE 'a'").await; + + // `Ordered` vs `Unordered` decides whether the DISTINCT stage may + // deduplicate by adjacency and hand back a continuation token, so pin the + // boundary the service actually draws: the `VALUE` form whose ORDER BY is + // exactly the projected path is the only shape that stays `Ordered`. + validate_pk("SELECT DISTINCT VALUE c.name FROM c ORDER BY c.name ASC").await; + validate_pk("SELECT DISTINCT VALUE c.name FROM c ORDER BY c.name DESC").await; + validate_pk("SELECT DISTINCT VALUE c.name FROM c ORDER BY c.other ASC").await; + validate_pk("SELECT DISTINCT VALUE c.name FROM c ORDER BY c.name, c.other").await; + validate_pk("SELECT DISTINCT c.name, c.city FROM c ORDER BY c.name ASC").await; + + // Constant DISTINCT collapses to `None` only *without* a FROM clause; with + // one the query yields a row per document, so deduplication is real work + // and the service reports `Unordered`. + validate_pk("SELECT DISTINCT VALUE 1 FROM c").await; + validate_pk("SELECT DISTINCT VALUE null FROM c").await; + validate_pk("SELECT DISTINCT 1 AS p FROM c").await; } #[tokio::test] diff --git a/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/distinct.rs b/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/distinct.rs new file mode 100644 index 00000000000..d500f3c4d6c --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/distinct.rs @@ -0,0 +1,616 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +//! In-memory-emulator integration tests for cross-partition `DISTINCT`. +//! +//! Unlike the mock-pipeline tests in `driver::dataflow::distinct`, these drive +//! the real path end to end: planner -> query-plan generation -> per-partition +//! execution (including the emulator's own per-partition deduplication) -> +//! client-side `Distinct` stage. +//! +//! Scenarios come from `tests/fixtures/distinct_scenarios.json`, the same +//! source-attributed catalog every other layer reads. + +use std::sync::Arc; +use std::time::Duration; + +use azure_core::http::Url; +use serde::Deserialize; + +use azure_data_cosmos_driver::driver::CosmosDriver; +use azure_data_cosmos_driver::in_memory_emulator::{ + ConsistencyLevel, ContainerConfig, InMemoryEmulatorHttpClient, VirtualAccountConfig, + VirtualRegion, +}; +use azure_data_cosmos_driver::models::{ + ContainerReference, CosmosOperation, FeedRange, ItemReference, MaxItemCountHint, PartitionKey, + PartitionKeyDefinition, +}; +use azure_data_cosmos_driver::options::{DriverOptions, OperationOptions, PlanOptions}; + +const GATEWAY_URL: &str = "https://eastus.emulator.local"; + +const CATALOG_JSON: &str = include_str!("../fixtures/distinct_scenarios.json"); + +#[derive(Deserialize)] +struct Catalog { + scenarios: Vec, +} + +#[derive(Deserialize)] +struct Scenario { + id: String, + layers: Vec, + query: QuerySpec, + #[serde(default)] + documents: Vec, + #[serde(rename = "pageSizes", default)] + page_sizes: Vec, + #[serde(rename = "expectedIds", default)] + expected_ids: Vec, + #[serde(rename = "expectedValues", default)] + expected_values: Vec, + checkpoint: Option, + #[serde(rename = "expectedError")] + expected_error: Option, +} + +#[derive(Deserialize)] +struct QuerySpec { + text: String, + #[serde(default)] + parameters: Vec, + #[serde(rename = "distinctType")] + distinct_type: String, +} + +#[derive(Deserialize)] +struct ExpectedError { + category: String, + #[serde(rename = "messageFragment")] + message_fragment: String, +} + +fn catalog() -> Catalog { + serde_json::from_str(CATALOG_JSON).expect("catalog must parse") +} + +/// Builds a two-physical-partition in-memory emulator container and a driver +/// wired to it. +async fn setup() -> (Arc, Arc) { + let config = VirtualAccountConfig::new(vec![VirtualRegion::new( + "East US", + Url::parse(GATEWAY_URL).unwrap(), + )]) + .unwrap() + .with_consistency(ConsistencyLevel::Session); + + let emulator = Arc::new(InMemoryEmulatorHttpClient::new(config)); + let store = emulator.store(); + store.create_database("testdb"); + let container_config = ContainerConfig::new() + .with_partition_count(2) + .build() + .unwrap(); + store.create_container_with_config( + "testdb", + "testcoll", + PartitionKeyDefinition::new(vec![std::borrow::Cow::Borrowed("/pk")]), + container_config, + ); + + let runtime = emulator + .runtime_builder() + .build() + .await + .expect("runtime builds against the in-memory emulator"); + let account = azure_data_cosmos_driver::models::AccountReference::with_master_key( + Url::parse(GATEWAY_URL).unwrap(), + "ZW11bGF0b3Ita2V5", + ); + let driver = runtime + .create_driver(DriverOptions::builder(account).build()) + .await + .expect("driver initializes against the emulator"); + (emulator, driver) +} + +async fn seed( + driver: &CosmosDriver, + container: &ContainerReference, + documents: &[serde_json::Value], +) { + for document in documents { + let id = document["id"] + .as_str() + .expect("fixture document needs an id") + .to_owned(); + let pk = document["pk"] + .as_str() + .expect("fixture document needs a pk") + .to_owned(); + let item_ref = ItemReference::from_name(container, PartitionKey::from(pk), id); + driver + .execute_singleton_operation( + CosmosOperation::create_item(item_ref) + .with_body(serde_json::to_vec(document).unwrap()), + OperationOptions::default(), + ) + .await + .expect("seed item created"); + } +} + +fn query_body(query: &QuerySpec) -> Vec { + serde_json::to_vec(&serde_json::json!({ + "query": query.text, + "parameters": query.parameters, + })) + .unwrap() +} + +fn query_operation( + container: &ContainerReference, + query: &QuerySpec, + page_size: u32, +) -> CosmosOperation { + CosmosOperation::query_items(container.clone(), Some(FeedRange::full())) + .with_body(query_body(query)) + .with_max_item_count(MaxItemCountHint::Limit( + std::num::NonZeroU32::new(page_size.max(1)).unwrap(), + )) +} + +/// Extracts every row from a response body, whichever wire shape it arrived in. +/// +/// The cross-partition pipeline emits pre-split `Items`; a raw single-partition +/// backend page arrives as a `{"Documents":[...]}` envelope. +fn documents_of( + response: azure_data_cosmos_driver::models::CosmosResponse, +) -> Vec { + use azure_data_cosmos_driver::models::ResponseBody; + match response.into_body() { + ResponseBody::NoPayload => Vec::new(), + ResponseBody::Items(items) => items + .iter() + .map(|item| serde_json::from_slice(item).unwrap()) + .collect(), + ResponseBody::Bytes(bytes) => { + let value: serde_json::Value = serde_json::from_slice(&bytes).unwrap(); + value["Documents"].as_array().cloned().unwrap_or_default() + } + } +} + +/// Sorts values by their serialized form so an unordered result can be +/// compared deterministically. +fn sorted(mut values: Vec) -> Vec { + let mut text: Vec = values.drain(..).map(|v| v.to_string()).collect(); + text.sort(); + text +} + +fn assert_matches_expected(scenario: &Scenario, actual: Vec) { + if !scenario.expected_ids.is_empty() { + let mut ids: Vec = actual + .iter() + .map(|v| v["id"].as_str().unwrap_or_default().to_owned()) + .collect(); + ids.sort(); + let mut expected = scenario.expected_ids.clone(); + expected.sort(); + assert_eq!(ids, expected, "scenario {}", scenario.id); + return; + } + if scenario.query.distinct_type == "Ordered" { + // A matching ORDER BY makes the output order deterministic. + assert_eq!( + actual, scenario.expected_values, + "scenario {} produced the wrong ordered stream", + scenario.id + ); + } else { + assert_eq!( + sorted(actual), + sorted(scenario.expected_values.clone()), + "scenario {} produced the wrong deduplicated set", + scenario.id + ); + } +} + +/// Drains a query fully, one plan, honoring `page_size`. +async fn drain_all( + driver: &CosmosDriver, + container: &ContainerReference, + query: &QuerySpec, + page_size: u32, +) -> Vec { + let mut plan = Box::pin(driver.plan_operation( + query_operation(container, query, page_size), + &OperationOptions::default(), + None, + &PlanOptions::default(), + )) + .await + .expect("plan builds"); + + let mut all = Vec::new(); + while let Some(response) = driver + .execute_plan( + &mut plan, + Some(container.clone()), + OperationOptions::default(), + ) + .await + .expect("page executes") + { + all.extend(documents_of(response)); + } + all +} + +/// Catalog-driven: every `inMemoryEmulator` scenario without a checkpoint or an +/// expected error must produce its expected rows, at every declared page size. +#[tokio::test] +async fn catalog_emulator_scenarios_dedupe_as_expected() { + let mut ran = 0usize; + for scenario in &catalog().scenarios { + if !scenario.layers.iter().any(|l| l == "inMemoryEmulator") + || scenario.checkpoint.is_some() + || scenario.expected_error.is_some() + { + continue; + } + let (_emulator, driver) = setup().await; + let container = driver + .resolve_container("testdb", "testcoll") + .await + .expect("container resolves"); + seed(&driver, &container, &scenario.documents).await; + + let page_sizes = if scenario.page_sizes.is_empty() { + vec![10] + } else { + scenario.page_sizes.clone() + }; + for page_size in page_sizes { + let actual = drain_all(&driver, &container, &scenario.query, page_size).await; + assert_matches_expected(scenario, actual); + } + ran += 1; + } + // Exact, not a floor: a scenario added to the catalog but silently skipped + // here would otherwise look covered while never executing. + let expected = catalog() + .scenarios + .iter() + .filter(|s| { + s.layers.iter().any(|l| l == "inMemoryEmulator") + && s.checkpoint.is_none() + && s.expected_error.is_none() + }) + .count(); + assert_eq!(ran, expected, "every eligible emulator scenario must run"); + assert!( + ran >= 8, + "expected the catalog to drive a meaningful number of emulator scenarios, ran {ran}" + ); +} + +/// Catalog-driven: every scenario declaring an `expectedError` must fail with a +/// matching message, rather than silently returning partial or duplicated rows. +#[tokio::test] +async fn catalog_emulator_error_scenarios_fail_as_expected() { + let mut ran = 0usize; + for scenario in &catalog().scenarios { + if !scenario.layers.iter().any(|l| l == "inMemoryEmulator") { + continue; + } + let Some(expected) = &scenario.expected_error else { + continue; + }; + let (_emulator, driver) = setup().await; + let container = driver + .resolve_container("testdb", "testcoll") + .await + .expect("container resolves"); + seed(&driver, &container, &scenario.documents).await; + + let outcome = match expected.category.as_str() { + // The unsupported-feature check happens while planning. + "clientUnsupportedQueryFeature" => Box::pin(driver.plan_operation( + query_operation(&container, &scenario.query, 10), + &OperationOptions::default(), + None, + &PlanOptions::default(), + )) + .await + .err() + .map(|e| e.to_string()), + // The continuation refusal happens when the caller mints a token. + "clientDistinctContinuationUnsupported" => { + let mut plan = Box::pin(driver.plan_operation( + query_operation(&container, &scenario.query, 1), + &OperationOptions::default(), + None, + &PlanOptions::default(), + )) + .await + .expect("an unordered DISTINCT query plans successfully"); + let _ = driver + .execute_plan( + &mut plan, + Some(container.clone()), + OperationOptions::default(), + ) + .await + .expect("the first page executes"); + plan.to_continuation_token().err().map(|e| e.to_string()) + } + other => panic!( + "scenario {} declares unhandled category {other}", + scenario.id + ), + }; + + let message = outcome.unwrap_or_else(|| { + panic!( + "scenario {} expected a {} failure but the operation succeeded", + scenario.id, expected.category + ) + }); + assert!( + message.contains(&expected.message_fragment), + "scenario {}: error {message:?} does not contain {:?}", + scenario.id, + expected.message_fragment + ); + ran += 1; + } + assert!( + ran >= 3, + "expected several emulator error scenarios, ran {ran}" + ); +} + +/// An ordered DISTINCT query drained across a serialized continuation token +/// must produce exactly the same rows, in the same order, as a single drain. +/// +/// Mirrors .NET `DistinctQueryTests.TestDistinct_ContinuationTokenSupportAsync` +/// and Java `DistinctQueryTests.queryDocumentsWithOrderBy`. +#[tokio::test] +async fn distinct_under_a_window_resumes_across_tokens() { + // A continuation for `DISTINCT` + `OFFSET`/`LIMIT`/`TOP` nests + // `SkipTake { child: Distinct { .. } }`. Resuming has to peel and restore + // both layers in the right order: lose the window's remaining budget and + // the query over-returns, lose the dedup hash and it repeats the boundary + // value. Draining page-by-page must equal a single drain either way. + for scenario_id in [ + "ordered_distinct_top_resume_round_trip", + "ordered_distinct_offset_limit_resume_round_trip", + ] { + let scenarios = catalog(); + let scenario = scenarios + .scenarios + .iter() + .find(|s| s.id == scenario_id) + .unwrap_or_else(|| panic!("catalog carries {scenario_id}")); + + let (_emulator, driver) = setup().await; + let container = driver + .resolve_container("testdb", "testcoll") + .await + .expect("container resolves"); + seed(&driver, &container, &scenario.documents).await; + + let single = drain_all(&driver, &container, &scenario.query, 10).await; + assert_eq!( + single, scenario.expected_values, + "{scenario_id}: single drain must honor the window over deduplicated values" + ); + + for &page_size in &scenario.page_sizes { + let mut resumed = Vec::new(); + let mut token = None; + loop { + let mut plan = Box::pin(driver.plan_operation( + query_operation(&container, &scenario.query, page_size), + &OperationOptions::default(), + token.as_ref(), + &PlanOptions::default(), + )) + .await + .expect("plan builds (fresh or resumed)"); + + let Some(response) = driver + .execute_plan( + &mut plan, + Some(container.clone()), + OperationOptions::default(), + ) + .await + .expect("page executes") + else { + break; + }; + resumed.extend(documents_of(response)); + + // Once the window is satisfied the pipeline is drained and + // mints no further token; stop rather than replaying the last. + match plan.to_continuation_token() { + Ok(next) => token = Some(next), + Err(_) => break, + } + if resumed.len() > scenario.expected_values.len() + 4 { + panic!("{scenario_id}: resume loop did not converge; got {resumed:?}"); + } + } + + assert_eq!( + resumed, scenario.expected_values, + "{scenario_id}: resuming at page size {page_size} must not drop, \ + duplicate, or over-return rows" + ); + } + } +} + +#[tokio::test] +async fn ordered_distinct_resume_matches_a_single_drain() { + let scenarios = catalog(); + let scenario = scenarios + .scenarios + .iter() + .find(|s| s.id == "ordered_resume_round_trip_matches_full_drain") + .expect("catalog carries the ordered resume round-trip scenario"); + + let (_emulator, driver) = setup().await; + let container = driver + .resolve_container("testdb", "testcoll") + .await + .expect("container resolves"); + seed(&driver, &container, &scenario.documents).await; + + let single = drain_all(&driver, &container, &scenario.query, 10).await; + assert_eq!(single, scenario.expected_values); + + // Now drain the same query one page at a time, round-tripping a + // continuation token between every page. + let mut resumed = Vec::new(); + let mut token = None; + loop { + let mut plan = Box::pin(driver.plan_operation( + query_operation(&container, &scenario.query, 1), + &OperationOptions::default(), + token.as_ref(), + &PlanOptions::default(), + )) + .await + .expect("plan builds (fresh or resumed)"); + + let Some(response) = driver + .execute_plan( + &mut plan, + Some(container.clone()), + OperationOptions::default(), + ) + .await + .expect("page executes") + else { + break; + }; + resumed.extend(documents_of(response)); + + token = Some( + plan.to_continuation_token() + .expect("an ordered DISTINCT query is resumable"), + ); + if resumed.len() > scenario.expected_values.len() + 4 { + panic!("resume loop did not converge; got {resumed:?}"); + } + } + + assert_eq!( + resumed, scenario.expected_values, + "resuming across continuation tokens must not drop or duplicate rows" + ); +} + +/// Reads the container's physical partition count through the emulator's +/// `pkranges` endpoint. +async fn physical_partition_count(emulator: &InMemoryEmulatorHttpClient) -> usize { + let url = format!("{GATEWAY_URL}/dbs/testdb/colls/testcoll/pkranges"); + let request = + azure_core::http::Request::new(Url::parse(&url).unwrap(), azure_core::http::Method::Get); + let response = emulator.execute_request(&request).await.unwrap(); + let raw = response.try_into_raw_response().await.unwrap(); + let body: serde_json::Value = serde_json::from_slice(raw.body().as_ref()).unwrap(); + body["PartitionKeyRanges"] + .as_array() + .expect("pkranges response carries a PartitionKeyRanges array") + .len() +} + +/// A partition split partway through a DISTINCT drain must not resurrect an +/// already-emitted value. Neither .NET nor Java covers this for DISTINCT. +#[tokio::test] +async fn split_mid_drain_does_not_reemit_deduplicated_values() { + for scenario_id in [ + "split_mid_query_unordered_does_not_reemit", + "split_mid_query_ordered_does_not_reemit", + // `DISTINCT` under a window: the split must preserve both the dedup + // state and the window's remaining budget. + "split_mid_query_windowed_distinct_does_not_reemit", + ] { + let scenarios = catalog(); + let scenario = scenarios + .scenarios + .iter() + .find(|s| s.id == scenario_id) + .expect("catalog carries the split scenario"); + + let (emulator, driver) = setup().await; + let container = driver + .resolve_container("testdb", "testcoll") + .await + .expect("container resolves"); + seed(&driver, &container, &scenario.documents).await; + + let mut plan = Box::pin(driver.plan_operation( + query_operation(&container, &scenario.query, 1), + &OperationOptions::default(), + None, + &PlanOptions::default(), + )) + .await + .expect("plan builds"); + + let ranges_before = physical_partition_count(&emulator).await; + + let mut all = Vec::new(); + let mut pages = 0usize; + let mut pages_after_split = 0usize; + while let Some(response) = driver + .execute_plan( + &mut plan, + Some(container.clone()), + OperationOptions::default(), + ) + .await + .expect("page executes") + { + all.extend(documents_of(response)); + pages += 1; + if pages == 1 { + // Split the first physical partition mid-drain. The fan-out + // root absorbs it; `Distinct` is never rebuilt, so its + // deduplication state has to survive. + emulator + .store() + .split_partition("testdb", "testcoll", 0, Duration::ZERO); + // Deterministic completion rather than a sleep, so the test + // cannot flake on split timing. + emulator.store().drain_pending_control_plane().await; + + // Without this the test could pass while covering nothing: a + // change to split timing or page shape would silently turn the + // split into a no-op. + let ranges_after = physical_partition_count(&emulator).await; + assert!( + ranges_after > ranges_before, + "{scenario_id}: the split must actually change the topology \ + (before={ranges_before}, after={ranges_after})" + ); + } else if pages > 1 { + pages_after_split += 1; + } + } + + assert!( + pages_after_split > 0, + "{scenario_id}: the drain must continue past the split for the test to prove anything" + ); + assert_matches_expected(scenario, all); + } +} diff --git a/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/mod.rs b/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/mod.rs index 9fecf017b0b..e97dc642ec5 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/mod.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/mod.rs @@ -7,6 +7,7 @@ pub mod account_metadata_refresh; pub mod batch; pub mod binary_response_format; pub mod control_plane; +pub mod distinct; #[cfg(feature = "preview_dtx")] pub mod distributed_transaction; pub mod dynamic_topology;