diff --git a/sdk/cosmos/.cspell.json b/sdk/cosmos/.cspell.json index ab818543976..788a98c1c21 100644 --- a/sdk/cosmos/.cspell.json +++ b/sdk/cosmos/.cspell.json @@ -6,6 +6,7 @@ "accountname", "activityid", "ALPN", + "antisymmetry", "apacsoutheast", "Appleby's", "argtypes", @@ -61,6 +62,7 @@ "consumingly", "Conv", "cooldown", + "copysign", "cosmosclient", "cosmosdriver", "correlator", @@ -118,6 +120,7 @@ "fixdate", "flamegraph", "fmix", + "formattableorderbyquery", "fract", "francecentral", "francesouth", @@ -308,6 +311,7 @@ "uncollapsed", "uncontended", "undrained", + "unemitted", "unfaulted", "ungoverned", "unparseable", diff --git a/sdk/cosmos/azure_data_cosmos/CHANGELOG.md b/sdk/cosmos/azure_data_cosmos/CHANGELOG.md index 250d218b5a9..e43ca5720c1 100644 --- a/sdk/cosmos/azure_data_cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure_data_cosmos/CHANGELOG.md @@ -13,6 +13,7 @@ - Added change feed pull support via `ContainerClient::query_change_feed()`, which takes a required `ChangeFeedStartFrom` start position (`Beginning`, `Now`, `PointInTime`) and returns a `ChangeFeedPageIterator` that streams `FeedPage` results. New `feed` types `ChangeFeedPageIterator`, `FeedScope`, and `ContinuationToken`, plus `options` types `ChangeFeedOptions` and `ChangeFeedMode` (currently `LatestVersion`); supports single-partition, per-partition-key, and full-container (cross-partition fan-out) reads with continuation-token resumption that persists the original start position so never-polled partitions don't replay history on resume. ([#4621](https://github.com/Azure/azure-sdk-for-rust/pull/4621)) - Change feed items are now surfaced as an envelope. `ContainerClient::query_change_feed::()` yields `ChangeFeedItem`, binding the envelope into the return type so the post-change document is read via `ChangeFeedItem::current()` and cannot be silently deserialized away. The envelope also exposes the pre-change document (`previous()`) and per-change `metadata()` (populated by full-fidelity reads; absent for `LatestVersion`). A full-fidelity delete returns an empty `current` object, which maps to `None` so callers with strict document types still deserialize the delete; the deleted item's identity is available via `ChangeFeedMetadata::id()` and `ChangeFeedMetadata::partition_key()`. `ChangeFeedOperationType` includes an `Unknown` catch-all so a future operation type cannot fail a page and stall the feed. A backend that does not envelope change feed items (such as the Cosmos emulator) returns the bare document, which is mapped onto `current()` so no data is lost. Added the `models` types `ChangeFeedItem`, `ChangeFeedMetadata`, `ChangeFeedOperationType`, and `LogicalSequenceNumber`. ([#4723](https://github.com/Azure/azure-sdk-for-rust/pull/4723)) - Added `TlsBackend` (re-exported) and a `tls_backend` option on `ConnectionPoolOptions` (`ConnectionPoolOptionsBuilder::with_tls_backend`), defaulting to `TlsBackend::Rustls`, available under the `rustls` feature, to pin the TLS backend used by the transport. This is additive and changes no behavior for the default (rustls) build; it only has an effect in builds that compile in multiple reqwest TLS backends, where reqwest would otherwise default to native-tls and the driver now pins rustls instead. ([#4649](https://github.com/Azure/azure-sdk-for-rust/pull/4649)) +- Added resumable cross-partition streaming `ORDER BY` query support. ([#4800](https://github.com/Azure/azure-sdk-for-rust/pull/4800)) ### Breaking Changes 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 ceda74a5bf5..bd82b34c39c 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 @@ -10,7 +10,6 @@ use azure_data_cosmos::feed::ContinuationToken; use azure_data_cosmos::{ clients::DatabaseClient, feed::FeedScope, - models::CosmosStatus, options::{MaxItemCountHint, QueryOptions}, Query, }; @@ -263,56 +262,27 @@ pub async fn cross_partition_query_with_projection_and_filter() -> Result<(), Bo test_category = "emulator_vnext", ignore = "skipped on vnext emulator: behavioral divergence" )] -pub async fn cross_partition_query_with_order_by_fails() -> Result<(), Box> { +pub async fn cross_partition_query_with_order_by() -> 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.clone(), None).await?; + let mut expected = items.clone(); + expected.sort_by_key(|item| item.merge_order); + let expected_ids = expected.into_iter().map(|item| item.id).collect(); - let Err(err) = container_client - .query_items::( - "select value c.id from c order by c.mergeOrder", - FeedScope::full_container(), - None, - ) - .await - else { - panic!("Expected query to fail due to cross-partition ORDER BY"); - }; - assert_eq!( - err.status(), - CosmosStatus::CROSS_PARTITION_QUERY_NOT_SERVABLE, - "Expected 400 / 1004 (CrossPartitionQueryNotServable) for cross-partition ORDER BY" - ); + execute_query_test( + db_client, + items, + "select value c.id from c order by c.mergeOrder", + FeedScope::full_container(), + expected_ids, + QueryTestOptions { + max_item_count: Some(7), + use_continuation_token_resume: true, + }, + ) + .await?; - let body = err - .response() - .and_then(|r| match r.body() { - azure_data_cosmos_driver::models::ResponseBody::Bytes(b) => Some(b.as_ref()), - _ => None, - }) - .expect("service error should carry a response body"); - #[derive(serde::Deserialize)] - struct ErrorDetail { - code: String, - message: String, - } - let error_detail: ErrorDetail = - serde_json::from_slice(body).expect("response body must be JSON"); - assert_eq!(error_detail.code, "BadRequest"); - - // Take only the first two lines of the message for comparison, since the full message may contain additional details that could change over time - let clean_message = error_detail - .message - .lines() - .take(2) - .collect::>() - .join("\n"); - assert_eq!( - clean_message, - "Query contains 1 or more unsupported features. Upgrade your SDK to a version that does support the requested features:\nQuery contained OrderBy, which the calling client does not support." - ); Ok(()) }, Some(TestOptions::for_emulator()), diff --git a/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator.rs b/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator.rs index 4fb690f8528..d56739cefd4 100644 --- a/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator.rs +++ b/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator.rs @@ -6,4 +6,8 @@ // the entire translation unit (and its dev-only deps such as `uuid`) out of // the build when the feature is off. +// Test futures compose many operations and near clippy's `large_futures` +// threshold on tokio's large test stacks; allowed here (mirrors `tests/split.rs`). +#![allow(clippy::large_futures)] + mod in_memory_emulator_tests; diff --git a/sdk/cosmos/azure_data_cosmos/tests/split_tests/cosmos_query_order_by_split.rs b/sdk/cosmos/azure_data_cosmos/tests/split_tests/cosmos_query_order_by_split.rs new file mode 100644 index 00000000000..7da6a2b386b --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos/tests/split_tests/cosmos_query_order_by_split.rs @@ -0,0 +1,624 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +//! Live-only split-resume coverage for the cross-partition streaming +//! `ORDER BY` pipeline (`StreamingOrderedMerge`), asserting the stronger +//! invariant an `ORDER BY` query must uphold: exact global sort order is +//! preserved across a real partition split and resume, not just "the same +//! items eventually come back". +//! +//! One container is seeded once with `mergeOrder` values that *repeat* +//! across every partition key (ties spread across partitions, not a +//! globally-unique key per item). Both an ASC and a DESC query are started +//! and each mid-tie-group continuation token is captured before a single +//! forced split, then both are resumed against the post-split topology. +//! Assertions cover every seeded id appearing exactly once, monotonic keys +//! in each direction, and — since ties can't be resolved by `mergeOrder` +//! alone — that DESC's tied-group order is the exact reverse of ASC's, +//! proving both share the same underlying document-`_rid` total order +//! (see `driver::dataflow::order_by::compare_rids`) rather than arbitrary +//! or storage-order noise that happens to differ per direction. +//! +//! Exhaustive split/merge permutations are covered in-process by +//! `azure_data_cosmos_driver`'s mock-pipeline and driver-level tests; this +//! single live test proves the same behavior against a real service split. +//! 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::error::Error; +use std::num::NonZeroU32; +use std::time::Duration; + +use azure_core::http::StatusCode; +use azure_data_cosmos::feed::ContinuationToken; +use azure_data_cosmos::options::CreateContainerOptions; +use azure_data_cosmos::{ + clients::ContainerClient, + feed::FeedScope, + models::{ + CompositeIndex, CompositeIndexOrder, CompositeIndexProperty, ContainerProperties, + IndexingPolicy, ThroughputProperties, + }, + options::{MaxItemCountHint, QueryOptions}, +}; +use framework::{MockItem, TestClient, TestOptions}; +use futures::StreamExt; +use serde::{de::DeserializeOwned, Deserialize, Serialize}; + +/// A seeded document carrying both a numeric (`mergeOrder`) and a string +/// (`sortText`) sort key, so one container/one split covers ORDER BY +/// resume for both key kinds. Queries deserialize only the fields they need +/// ([`MockItem`] for the numeric key, [`StringKeyItem`] for the string key). +#[derive(Serialize, Clone, Debug)] +#[serde(rename_all = "camelCase")] +struct SeedItem { + id: String, + partition_key: String, + merge_order: usize, + sort_text: String, +} + +/// Query projection for the string-keyed ORDER BY (reads only `id` and +/// `sortText`; serde ignores the seeded `partitionKey`/`mergeOrder`). +#[derive(Deserialize, Clone, Debug)] +#[serde(rename_all = "camelCase")] +struct StringKeyItem { + id: String, + sort_text: String, +} + +/// A seeded document carrying an array-valued sort key, for the dedicated +/// array-keyed `ORDER BY` split-resume test. +#[derive(Serialize, Clone, Debug)] +#[serde(rename_all = "camelCase")] +struct ArraySeedItem { + id: String, + partition_key: String, + array_key: Vec, +} + +/// Projection for the array-keyed `ORDER BY` (reads only `id`). +#[derive(Deserialize, Clone, Debug)] +#[serde(rename_all = "camelCase")] +struct ArrayKeyItem { + id: String, +} + +/// A string sort key for tie group `i`: ordered by the zero-padded index +/// prefix, with a constant suffix carrying a quote, backslash, and tab so +/// the resume boundary contains SQL special characters the parameterized filter +/// must round-trip without inlining them into query text. +fn string_key(i: usize) -> String { + format!("g{i:02}-x'\\\t") +} + +/// Runs `query`'s first page only (before any split), returning the items +/// it yielded plus a re-serialized continuation token — mirroring how a +/// real caller would capture and persist a mid-fan-out checkpoint. +async fn capture_first_page( + container_client: &ContainerClient, + query: &str, + page_size: u32, +) -> Result<(Vec, ContinuationToken), Box> { + 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 first_page = pages + .next() + .await + .expect("query should yield at least one page before split")?; + let items = first_page.into_items(); + let serialized = pages.to_continuation_token()?.as_str().to_owned(); + Ok((items, ContinuationToken::from_string(serialized))) +} + +/// Resumes `query` from `continuation` and drains it to completion, +/// appending to `collected`. Round-trips the continuation through string +/// serialization each page, matching how a real caller persists it across +/// processes. +async fn drain_resumed( + container_client: &ContainerClient, + query: &str, + page_size: u32, + mut collected: Vec, + mut continuation: Option, +) -> Result, Box> { + loop { + let mut resume_options = QueryOptions::default() + .with_max_item_count(MaxItemCountHint::Limit(NonZeroU32::new(page_size).unwrap())); + if let Some(t) = continuation.take() { + resume_options = resume_options.with_continuation_token(t); + } + let mut pages = container_client + .query_items::(query, FeedScope::full_container(), Some(resume_options)) + .await? + .into_pages(); + + let Some(page) = pages.next().await else { + break; + }; + collected.extend(page?.into_items()); + let serialized = pages.to_continuation_token()?.as_str().to_owned(); + drop(pages); + continuation = Some(ContinuationToken::from_string(serialized)); + } + Ok(collected) +} + +/// Cross-partition streaming `ORDER BY` resume after a live split, for both +/// directions, with repeated (tied) sort-key values, and for both a numeric +/// (`mergeOrder`) and a string (`sortText`) sort key. The string key carries +/// SQL special characters (quote, backslash, tab), so its resume boundary +/// exercises the parameterized service-safe seek filter across a real split. +/// +/// Seeds `PK_COUNT` partition keys, each contributing exactly one item per +/// group `0..TIE_GROUP_COUNT` — every sort-key value therefore ties across +/// all `PK_COUNT` partition keys (and, after the split, across both physical +/// partitions) instead of each item having a globally unique key. Captures +/// an ASC and a DESC continuation token mid-tie-group for each key kind, +/// forces one live split, then resumes all four against the post-split +/// topology. +/// +/// Reusing the same container and single split, also drains a mixed-direction +/// multi-column query (`mergeOrder ASC, sortText DESC`) — whose entire tuple +/// ties within every group and is therefore broken only by `_rid` — to +/// completion pre-split for an authoritative baseline, then asserts its +/// post-split resume reproduces that exact sequence, not merely the same ids. +#[tokio::test] +#[cfg_attr( + not(test_category = "split"), + ignore = "requires test_category 'split'" +)] +pub async fn order_by_query_resume_across_split_preserves_global_order( +) -> Result<(), Box> { + const PK_COUNT: usize = 50; + const TIE_GROUP_COUNT: usize = 10; + const PAGE_SIZE: u32 = 10; + const ASC_QUERY: &str = "SELECT * FROM c ORDER BY c.mergeOrder ASC"; + const DESC_QUERY: &str = "SELECT * FROM c ORDER BY c.mergeOrder DESC"; + const STRING_ASC_QUERY: &str = "SELECT * FROM c ORDER BY c.sortText ASC"; + const STRING_DESC_QUERY: &str = "SELECT * FROM c ORDER BY c.sortText DESC"; + // Mixed-direction, multi-column: both columns are constant within a tie + // group, so the whole tuple ties across every partition key and only the + // document `_rid` order (first column ASC) breaks it. + const MULTI_QUERY: &str = "SELECT * FROM c ORDER BY c.mergeOrder ASC, c.sortText DESC"; + + TestClient::run_with_unique_db( + async |run_context, db_client| { + let composite_index = CompositeIndex::default() + .with_property(CompositeIndexProperty::new( + "/mergeOrder", + CompositeIndexOrder::Ascending, + )) + .with_property(CompositeIndexProperty::new( + "/sortText", + CompositeIndexOrder::Descending, + )); + let properties = + ContainerProperties::new("OrderByResumeAcrossSplit", "/partitionKey".into()) + .with_indexing_policy( + IndexingPolicy::default().with_composite_index(composite_index), + ); + let throughput = ThroughputProperties::manual(1000); + let container_client = std::sync::Arc::new( + run_context + .create_container( + db_client, + properties, + Some(CreateContainerOptions::default().with_throughput(throughput)), + ) + .await?, + ); + + println!( + "Container created; seeding {PK_COUNT} partition keys x {TIE_GROUP_COUNT} \ + tied groups (repeated numeric + string sort values spread across partitions)" + ); + for p in 0..PK_COUNT { + let partition_key = format!("pk{p}"); + for i in 0..TIE_GROUP_COUNT { + let item = SeedItem { + id: format!("{p}-{i}"), + partition_key: partition_key.clone(), + merge_order: i, + sort_text: string_key(i), + }; + match container_client + .create_item(item.partition_key.clone(), &item.id.clone(), item, None) + .await + { + Ok(_) => {} + // A retried transport 503 can find the item already + // committed; unique ids make Conflict here benign. + Err(error) if error.status().status_code() == StatusCode::Conflict => {} + Err(error) => return Err(error.into()), + } + } + } + + 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(); + + // Capture ASC + DESC tokens mid-tie-group (each group has more + // members than one page) for BOTH the numeric and string keys, + // before the single forced split below. + let (asc_first, asc_token) = + capture_first_page::(&container_client, ASC_QUERY, PAGE_SIZE).await?; + let (desc_first, desc_token) = + capture_first_page::(&container_client, DESC_QUERY, PAGE_SIZE).await?; + let (str_asc_first, str_asc_token) = + capture_first_page::(&container_client, STRING_ASC_QUERY, PAGE_SIZE) + .await?; + let (str_desc_first, str_desc_token) = capture_first_page::( + &container_client, + STRING_DESC_QUERY, + PAGE_SIZE, + ) + .await?; + + // Fully drain the multi-column query on the pre-split topology for + // the authoritative expected order, then capture a mid-tie + // continuation to resume across the split and compare against it. + let multi_baseline = drain_resumed::( + &container_client, + MULTI_QUERY, + PAGE_SIZE, + Vec::new(), + None, + ) + .await?; + let multi_baseline_ids: Vec = + multi_baseline.iter().map(|item| item.id.clone()).collect(); + let (multi_first, multi_token) = + capture_first_page::(&container_client, MULTI_QUERY, PAGE_SIZE).await?; + + println!( + "Captured numeric ASC/DESC ({}/{}) and string ASC/DESC ({}/{}) tokens; forcing \ + the single split all four will resume across", + asc_first.len(), + desc_first.len(), + str_asc_first.len(), + str_desc_first.len(), + ); + 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}" + ); + + // Resume each direction's continuation independently against + // the now-split topology. + let asc_all = drain_resumed::( + &container_client, + ASC_QUERY, + PAGE_SIZE, + asc_first, + Some(asc_token), + ) + .await?; + let desc_all = drain_resumed::( + &container_client, + DESC_QUERY, + PAGE_SIZE, + desc_first, + Some(desc_token), + ) + .await?; + let str_asc_all = drain_resumed::( + &container_client, + STRING_ASC_QUERY, + PAGE_SIZE, + str_asc_first, + Some(str_asc_token), + ) + .await?; + let str_desc_all = drain_resumed::( + &container_client, + STRING_DESC_QUERY, + PAGE_SIZE, + str_desc_first, + Some(str_desc_token), + ) + .await?; + let multi_all = drain_resumed::( + &container_client, + MULTI_QUERY, + PAGE_SIZE, + multi_first, + Some(multi_token), + ) + .await?; + + // ── Every seeded id exactly once, in each direction ────────── + let mut expected_ids: Vec = (0..PK_COUNT) + .flat_map(|p| (0..TIE_GROUP_COUNT).map(move |i| format!("{p}-{i}"))) + .collect(); + expected_ids.sort(); + + let mut asc_ids: Vec = asc_all.iter().map(|item| item.id.clone()).collect(); + asc_ids.sort(); + assert_eq!( + asc_ids, expected_ids, + "ASC drain must return every seeded id exactly once (no duplicates or losses \ + across the split)" + ); + + let mut desc_ids: Vec = desc_all.iter().map(|item| item.id.clone()).collect(); + desc_ids.sort(); + assert_eq!( + desc_ids, expected_ids, + "DESC drain must return every seeded id exactly once (no duplicates or losses \ + across the split)" + ); + + // ── Monotonic keys per direction ───────────────────────────── + let asc_keys: Vec = asc_all.iter().map(|item| item.merge_order).collect(); + assert!( + asc_keys.windows(2).all(|w| w[0] <= w[1]), + "ASC mergeOrder must be non-decreasing across the split/resume boundary: \ + {asc_keys:?}" + ); + let desc_keys: Vec = desc_all.iter().map(|item| item.merge_order).collect(); + assert!( + desc_keys.windows(2).all(|w| w[0] >= w[1]), + "DESC mergeOrder must be non-increasing across the split/resume boundary: \ + {desc_keys:?}" + ); + + // ── Tie order is a real, direction-symmetric rid order ─────── + // For every tied `mergeOrder` group, DESC's item sequence must + // be the *exact reverse* of ASC's: both are broken by the same + // document-`_rid` total order, just applied in opposite + // directions. Anything else (arbitrary order, storage-order + // noise, or a direction-dependent tie-break bug) would make the + // two sequences disagree once reversed. + for group in 0..TIE_GROUP_COUNT { + let asc_group: Vec<&str> = asc_all + .iter() + .filter(|item| item.merge_order == group) + .map(|item| item.id.as_str()) + .collect(); + let mut desc_group: Vec<&str> = desc_all + .iter() + .filter(|item| item.merge_order == group) + .map(|item| item.id.as_str()) + .collect(); + desc_group.reverse(); + assert_eq!( + asc_group, desc_group, + "tie group mergeOrder={group}: DESC order reversed must equal ASC order \ + (both broken by the same rid total order)" + ); + assert_eq!( + asc_group.len(), + PK_COUNT, + "tie group mergeOrder={group} must contain exactly one item per partition key" + ); + } + + // ── String-keyed ORDER BY resume (service-safe boundaries) ──── + // The same invariants for the special-character-laden string key, + // proving the parameterized resume filter round-trips a boundary + // containing a quote, backslash, and tab across a real split. + let mut str_asc_ids: Vec = + str_asc_all.iter().map(|item| item.id.clone()).collect(); + str_asc_ids.sort(); + assert_eq!( + str_asc_ids, expected_ids, + "string ASC drain must return every seeded id exactly once across the split" + ); + let mut str_desc_ids: Vec = + str_desc_all.iter().map(|item| item.id.clone()).collect(); + str_desc_ids.sort(); + assert_eq!( + str_desc_ids, expected_ids, + "string DESC drain must return every seeded id exactly once across the split" + ); + + let str_asc_keys: Vec<&str> = str_asc_all + .iter() + .map(|item| item.sort_text.as_str()) + .collect(); + assert!( + str_asc_keys.windows(2).all(|w| w[0] <= w[1]), + "string ASC sortText must be non-decreasing across the split/resume boundary: \ + {str_asc_keys:?}" + ); + let str_desc_keys: Vec<&str> = str_desc_all + .iter() + .map(|item| item.sort_text.as_str()) + .collect(); + assert!( + str_desc_keys.windows(2).all(|w| w[0] >= w[1]), + "string DESC sortText must be non-increasing across the split/resume boundary: \ + {str_desc_keys:?}" + ); + + for group in 0..TIE_GROUP_COUNT { + let key = string_key(group); + let asc_group: Vec<&str> = str_asc_all + .iter() + .filter(|item| item.sort_text == key) + .map(|item| item.id.as_str()) + .collect(); + let mut desc_group: Vec<&str> = str_desc_all + .iter() + .filter(|item| item.sort_text == key) + .map(|item| item.id.as_str()) + .collect(); + desc_group.reverse(); + assert_eq!( + asc_group, desc_group, + "string tie group {group}: DESC order reversed must equal ASC order \ + (both broken by the same rid total order)" + ); + assert_eq!( + asc_group.len(), + PK_COUNT, + "string tie group {group} must contain exactly one item per partition key" + ); + } + + // ── Multi-column resume equals the pre-split baseline exactly ─ + // The strongest cross-split invariant: a full tuple tie is resolved + // only by `_rid`, and `_rid`s are stable across a split, so the exact + // pre-split drain order must survive the split/resume — not merely the + // same id set. + assert_eq!( + multi_baseline_ids.len(), + PK_COUNT * TIE_GROUP_COUNT, + "multi-column baseline must observe every seeded id exactly once pre-split" + ); + let multi_first_col: Vec = + multi_baseline.iter().map(|item| item.merge_order).collect(); + assert!( + multi_first_col.windows(2).all(|w| w[0] <= w[1]), + "multi-column baseline's first column (mergeOrder ASC) must be non-decreasing: \ + {multi_first_col:?}" + ); + let multi_ids: Vec = multi_all.iter().map(|item| item.id.clone()).collect(); + assert_eq!( + multi_ids, multi_baseline_ids, + "multi-column mixed-direction resume across the split must reproduce the exact \ + pre-split drain order (full tuple ties broken only by `_rid`)" + ); + + Ok(()) + }, + Some(TestOptions::new().with_timeout(Duration::from_secs(40 * 60))), + ) + .await +} + +/// A dedicated live check that a cross-partition streaming `ORDER BY` on an +/// **array-valued** sort key resumes across a real split with no omissions +/// or duplicates. +/// +/// Array/object sort values are ordered by the backend's bounded hash (there +/// is no documented structural order), and Rust sends that boundary to the +/// backend as the structured `resumeFilter` (its complex value serialized as +/// `{"type":"array","low":..,"high":..}`). This asserts the resume-correctness +/// invariant — every seeded id returned exactly once across the split — rather +/// than an exact global sequence (which is undefined for complex keys). +/// +/// By design this exercises the **saved-token** resume path: the first page is +/// drained and its continuation token captured, the split is then forced while +/// *no query is active*, and the query is reissued from the saved token. That +/// token's complex boundary replays through the structured `resumeFilter`, +/// whose backend `DistinctHash` seek excludes already-emitted rows across the +/// new topology — so the test is expected to pass. Live in-process splits — +/// where the split child forwards its own backend continuation into each +/// replacement, so no client-side discard is needed — are covered exhaustively +/// by the `streaming_ordered_merge` unit tests. +/// +/// Kept separate from `order_by_query_resume_across_split_preserves_global_order` +/// so its default (index-everything) container can serve an array-path +/// `ORDER BY` without a hand-tuned composite index, and so a service that does +/// not support ordering by an array value surfaces here without destabilizing +/// the scalar coverage. Gated on `test_category = "split"`. +#[tokio::test] +#[cfg_attr( + not(test_category = "split"), + ignore = "requires test_category 'split'" +)] +pub async fn order_by_array_key_resume_across_split_has_no_omissions() -> Result<(), Box> +{ + const PK_COUNT: usize = 30; + const TIE_GROUP_COUNT: usize = 5; + const PAGE_SIZE: u32 = 10; + const ARRAY_QUERY: &str = "SELECT * FROM c ORDER BY c.arrayKey ASC"; + + TestClient::run_with_unique_db( + async |run_context, db_client| { + // Default (index-everything) indexing so the array sort path is + // servable without a hand-tuned composite index. + let properties = ContainerProperties::new( + "OrderByArrayKeyResumeAcrossSplit", + "/partitionKey".into(), + ); + let throughput = ThroughputProperties::manual(1000); + let container_client = std::sync::Arc::new( + run_context + .create_container( + db_client, + properties, + Some(CreateContainerOptions::default().with_throughput(throughput)), + ) + .await?, + ); + + // Each group's array value `[i]` ties across every partition key, + // so a tie run spreads across both post-split physical partitions. + for p in 0..PK_COUNT { + let partition_key = format!("pk{p}"); + for i in 0..TIE_GROUP_COUNT { + let item = ArraySeedItem { + id: format!("{p}-{i}"), + partition_key: partition_key.clone(), + array_key: vec![i as i64], + }; + match container_client + .create_item(item.partition_key.clone(), &item.id.clone(), item, None) + .await + { + Ok(_) => {} + Err(error) if error.status().status_code() == StatusCode::Conflict => {} + Err(error) => return Err(error.into()), + } + } + } + + let ranges_before = container_client.read_feed_ranges(None).await?; + let partitions_before = ranges_before.len(); + + let (array_first, array_token) = + capture_first_page::(&container_client, ARRAY_QUERY, PAGE_SIZE) + .await?; + + 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}" + ); + + let array_all = drain_resumed::( + &container_client, + ARRAY_QUERY, + PAGE_SIZE, + array_first, + Some(array_token), + ) + .await?; + + let mut expected_ids: Vec = (0..PK_COUNT) + .flat_map(|p| (0..TIE_GROUP_COUNT).map(move |i| format!("{p}-{i}"))) + .collect(); + expected_ids.sort(); + let mut array_ids: Vec = array_all.iter().map(|item| item.id.clone()).collect(); + array_ids.sort(); + assert_eq!( + array_ids, expected_ids, + "array-keyed ORDER BY resume across the split must return every seeded id \ + exactly once (no omissions or duplicates)" + ); + Ok(()) + }, + 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 e3b2d3b70eb..515571a648d 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_order_by_split; mod cosmos_query_split; mod cosmos_split_offers; diff --git a/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md b/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md index 0affa5f2fb0..c5cbeb8b5cb 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md +++ b/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md @@ -12,6 +12,7 @@ - Added preview distributed transaction driver models, request serialization, response parsing/reordering, strict session-token merge, DTX retry handling, diagnostics, and a baseline in-memory emulator `/operations/dtc` path behind the disabled-by-default `preview_dtx` feature. ([#4702](https://github.com/Azure/azure-sdk-for-rust/pull/4702)) - Added change feed support in the dataflow pipeline: a new `UnorderedMerge` node fans a change feed read out across physical partitions and round-robins their pages, and `CosmosOperation::change_feed` builds incremental-feed operations with the appropriate wire headers. A new public `ChangeFeedStartFrom` enum (`Beginning`, `Now`, `PointInTime`) records the feed's original start position and is persisted in the continuation token so partitions never polled before a checkpoint re-apply it on resume instead of replaying history; set it via `CosmosOperation::with_change_feed_start`. ([#4621](https://github.com/Azure/azure-sdk-for-rust/pull/4621)) - Added `TlsBackend` (currently `TlsBackend::Rustls`, the default) and a `tls_backend` option on `ConnectionPoolOptions` (`ConnectionPoolOptionsBuilder::with_tls_backend` / `ConnectionPoolOptions::tls_backend`), available under the `rustls` feature. The driver asserts the selected backend on the `reqwest` transport, giving a supported way to pin the TLS backend without direct transport access. This is additive and changes no behavior for the default (rustls-only) build, where reqwest already negotiates rustls; it only has an effect in builds that compile in multiple reqwest TLS backends (e.g. `rustls` plus `native_tls`, absent reqwest's `http3` feature), where reqwest would otherwise default to native-tls and the driver now pins rustls instead. ([#4649](https://github.com/Azure/azure-sdk-for-rust/pull/4649)) +- Added resumable cross-partition streaming `ORDER BY` query support. ([#4800](https://github.com/Azure/azure-sdk-for-rust/pull/4800)) ### Breaking Changes diff --git a/sdk/cosmos/azure_data_cosmos_driver/Cargo.toml b/sdk/cosmos/azure_data_cosmos_driver/Cargo.toml index f28c783ad95..daa72575a38 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/Cargo.toml +++ b/sdk/cosmos/azure_data_cosmos_driver/Cargo.toml @@ -32,7 +32,11 @@ percent-encoding = { workspace = true, optional = true } rand.workspace = true reqwest = { workspace = true, optional = true } serde.workspace = true -serde_json.workspace = true +# `raw_value` lets the cross-partition streaming ORDER BY pipeline +# (`driver::dataflow::query_response`) retain each item's `payload` as its +# exact original JSON bytes while parsing the rewritten-query envelope, +# instead of re-serializing (and potentially reformatting numbers). +serde_json = { workspace = true, features = ["raw_value"] } time.workspace = true tokio = { workspace = true, optional = true, features = ["rt", "time"] } tracing.workspace = true 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 e89309dc2db..a2eca991686 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 @@ -1,8 +1,9 @@ + # Feed Operations — Requirements & Design Primer **Crate:** `azure_data_cosmos_driver` **Scope:** Driver-internal architecture for feed operations (queries, future read-many, change feed) -**Current focus:** `SELECT * [WHERE ]` using natural order +**Current focus:** `SELECT * [WHERE ]` using natural order, cross-partition streaming `ORDER BY`, and change feed --- @@ -22,7 +23,7 @@ All operations — point and feed — are expressed as a **Dataflow Pipeline**: - The pipeline is a **tree**. Nodes own their children. Fan-out creates branching. - **Leaf nodes** issue a single Cosmos DB request via the existing operation pipeline (retry, failover, auth, transport). -- **Intermediate nodes** orchestrate their children. The first intermediate node type is `SequentialDrain`, which iterates children in EPK order, fully draining one before advancing to the next. +- **Intermediate nodes** orchestrate their children. `SequentialDrain` iterates children in EPK order, fully draining one before advancing to the next (natural-order queries). `UnorderedMerge` polls children round-robin without evicting them (change feed). `StreamingOrderedMerge` k-way merges globally-ordered results across children, each executing a Gateway-rewritten per-partition query (cross-partition `ORDER BY`) — see §8 and `driver::dataflow::streaming_ordered_merge` for the implementation. - **Trivial pipelines** (point operations, single-partition feeds) are a single leaf node with no intermediate parent. These must add near-zero overhead compared to today's direct execution path. ### Pipeline Lifecycle @@ -34,12 +35,12 @@ All operations — point and feed — are expressed as a **Dataflow Pipeline**: ### Future Node Types (Design For, Don't Implement Yet) -- **UnorderedMerge**: concurrent fan-out, results returned in arrival order (Read Many). -- **StreamingOrderedMerge**: k-way merge of pre-sorted partition streams (streaming ORDER BY). -- **BufferedOrderedMerge**: collect all results, then sort (non-streaming ORDER BY). +- **BufferedOrderedMerge**: collect all results, then sort (non-streaming `ORDER BY` — issue #4755). - **HybridSearch**: issues multiple distinct sub-queries (e.g., vector similarity + full-text keyword) against different child pipelines, then combines/re-ranks their results. Demonstrates that an intermediate node may have heterogeneous children with different semantics. - **Aggregate**: client-side aggregation across partitions. +Implemented since this primer was written: **UnorderedMerge** (change feed fan-out) and **StreamingOrderedMerge** (streaming cross-partition `ORDER BY`; see §8). + --- ## 3. Key Invariants @@ -139,6 +140,8 @@ The initial implementation targets `SELECT [...] [WHERE ]` queries: - Integration with the existing operation pipeline for each sub-request. - Pipeline repair on partition splits/merges. +Since this initial milestone, change feed (`UnorderedMerge`, §2) and cross-partition streaming `ORDER BY` (`StreamingOrderedMerge`, §2, §7, §8) have landed as additional cross-partition strategies over the same pipeline model. + --- ## 7. Design Boundaries @@ -154,7 +157,9 @@ The initial implementation targets `SELECT [...] [WHERE ]` queries: For `SequentialDrain`, item bodies are fully opaque binary payloads. The pipeline does not inspect them — ordering is already established by the backend. -Future node types (e.g., streaming ORDER BY, hybrid search) may require partial parsing of item bodies. The backend query plan can rewrite the query to use a standardized envelope (promoting ordering keys to top-level fields and demoting the raw user document to a `payload` field). This varied-shape pattern must be considered in the overall design direction, but does not need to be accommodated in the current implementation. +`StreamingOrderedMerge` (cross-partition `ORDER BY`) is the first node type that parses item bodies: the backend query plan's `rewrittenQuery` reshapes each per-partition result into a standardized envelope (`{"_rid": ..., "orderByItems": [...], "payload": ...}` — sort keys promoted to top-level fields, the raw user document demoted to `payload`). The pipeline parses just enough of this envelope to compare rows and emit `payload` — see `driver::dataflow::order_by` (the value/comparator/resume-value model) and `driver::dataflow::query_response` (envelope parsing and response reassembly). Hybrid search, if implemented, would need the same pattern for its own envelope shape. + +Resume for streaming `ORDER BY` is per-range and topology-change-safe. Each still-active range persists its last-emitted key tuple, envelope `_rid`, and a `_rid`-tie `skipCount`. On resume — including after a partition split or merge — the boundary is sent to the backend as the .NET-compatible structured `resumeFilter` field of the query body (`{"value": [], "rid": , "exclude": false}` — see `driver::dataflow::query_response::with_resume_filter`), not an SDK-rewritten SQL predicate. The query text is unchanged (its Gateway resume-filter placeholder is replaced with `true`, exactly as for a fresh start), and the caller's `parameters` are preserved verbatim; sort direction lives in the query's `ORDER BY` clause. Rust persists one boundary per range and resumes each range with its own boundary in the .NET "target" partition style (`rid` present, `exclude: false`); the already-emitted prefix of the boundary tie run is discarded client-side by a three-phase `(sort key, _rid, skipCount)` seek matching .NET's `FilterNextAsync`: keys before the boundary and `_rid`s before the boundary `_rid` (numeric document-ordinal order, see `models::resource_id::compare_document_rids`) are dropped, then exactly `skipCount` rows sharing the boundary's exact `(sort key, _rid)` are dropped. `skipCount` (mirroring .NET's `OrderByContinuationToken.SkipCount`) is needed only because a JOIN/array-unwind query can emit several result rows from one document that share both the sort key and `_rid`; for ordinary single-row-per-document queries it is always `1` (just the boundary row). The seek is per row, so resume stays exact and clips to sub-ranges across a split — only the sub-range that re-returns the boundary `_rid` consumes `skipCount`. Scalar boundaries round-trip exactly; **array/object** (complex) boundaries travel as a bounded 128-bit hash (`{"type": "array"|"object", "low": , "high": }`) that is an exact port of the backend / .NET SDK structural `DistinctHash` (see `driver::dataflow::distinct_hash`), so it is byte-identical to the backend's own hash. A **saved-token** resume across a split or merge is therefore safe for scalar and complex keys alike (there is no positional-rescan or unchanged-topology restriction): the persisted boundary is replayed through the structured `resumeFilter`, and the backend's own `DistinctHash` seek excludes already-emitted complex rows. A **live** split — one hit mid-fan-out — is equally safe: `Request::split_for_topology_change` forwards the split child's own backend continuation into every replacement, so (by the Cosmos contract also relied on by `SequentialDrain` — a parent partition's continuation stays valid on each post-split child under EPK scoping) each replacement resumes *after* every already-emitted row. `StreamingOrderedMerge::handle_split` therefore installs no client-side discard on a live split's first replacement page — reapplying one would wrongly drop later JOIN rows that share the boundary `(sort key, _rid)` — and forwards only the boundary's `skipCount` bookkeeping for future snapshots. Scalar and complex live splits, and every saved-token resume, are all fully supported. ### The Driver DOES: @@ -171,10 +176,13 @@ Future node types (e.g., streaming ORDER BY, hybrid search) may require partial These capabilities must be achievable without redesigning the pipeline model: -- **Streaming ORDER BY**: k-way merge of partition streams. Requires fetching a backend query plan to determine sort keys. New intermediate node type. -- **Buffered ORDER BY**: collect all partition results, sort client-side. Same query plan requirement. Different intermediate node. +- **Buffered ORDER BY**: collect all partition results, sort client-side (issue #4755, non-streaming `ORDER BY`). Different intermediate node than streaming `ORDER BY` (below); same query-plan requirement. - **Vector / Hybrid Search**: may require preliminary requests to fetch full-text statistics before issuing the main query. Multi-phase pipeline execution. - **Read Many Items**: fan-out by (ID, PK) pairs grouped by partition. Concurrent leaf execution with an unordered merge intermediate node. -- **Change Feed**: per-range continuation tokens. Different resumption semantics. + +Implemented since this primer was written: + +- **Streaming ORDER BY**: `StreamingOrderedMerge` k-way merges partition streams using the backend query plan's rewritten per-partition query and sort keys. See §2, §7, and `driver::dataflow::streaming_ordered_merge`. +- **Change Feed**: `UnorderedMerge` polls per-range continuation tokens round-robin; see §2. The pipeline's tree structure, typed node hierarchy, and separation of planning from execution accommodate all of these as new node types and planning strategies without changing the core execution loop. 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 4bb66f98f18..dda5e22c387 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 @@ -2844,6 +2844,23 @@ impl CosmosDriver { |container, continuation| self.fetch_pk_ranges_from_service(container, continuation), ); + // Route streaming ORDER BY queries to the k-way merge instead of + // the natural-order sequential drain. + if query_plan + .query_info + .as_ref() + .is_some_and(planner::is_streaming_order_by) + { + let pipeline = planner::build_streaming_ordered_merge( + &query_plan, + &mut topology, + &operation, + resume_state, + ) + .await?; + return Ok(OperationPlan::new(pipeline, operation)); + } + let pipeline = planner::build_sequential_drain(&query_plan, &mut topology, &operation, resume_state) .await?; 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..40b399b8d0a --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/distinct_hash.rs @@ -0,0 +1,416 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +//! Structural 128-bit hash of a JSON value, an exact port of the Cosmos DB +//! backend / .NET SDK `DistinctHash` +//! (`Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct.DistinctHash`). +//! +//! The backend represents a complex (array/object) `ORDER BY` resume value as +//! this hash, so producing a byte-identical hash is what lets the structured +//! `resumeFilter` seek correctly from an array/object boundary — including +//! across a partition split or merge. +//! +//! Structurally-equal values hash equal regardless of object property order +//! (property hashes are XOR-combined). The 128-bit output is the raw +//! `UInt128` whose low/high 64-bit halves match .NET's `UInt128.GetLow()` / +//! `GetHigh()`; hashing a `UInt128` or `DoubleEx` feeds its little-endian +//! memory bytes to [`murmurhash3_128`], exactly as .NET's +//! `MemoryMarshal.AsBytes` does on a little-endian host. +//! +//! This operates on [`OrderByItem`] (the driver's typed JSON key model), so +//! it is reusable by future `DISTINCT` / `GROUP BY` work: any +//! `serde_json::Value` bridges in via [`OrderByItem::from_json`]. + +use super::order_by::{OrderByItem, OrderByNumber}; +use crate::models::murmur_hash::murmurhash3_128; + +/// Combines a `UInt128`'s low and high 64-bit words into a `u128`, matching +/// .NET's `UInt128.Create(low, high)`. +const fn seed(low: u64, high: u64) -> u128 { + ((high as u128) << 64) | (low as u128) +} + +/// Root and per-type hash seeds, verbatim from .NET `DistinctHash` +/// (`HashSeeds` / `RootHashSeed` and `CosmosNumberHasher.HashSeeds`). +mod seeds { + use super::seed; + + pub(super) const ROOT: u128 = seed(0xbfc2359eafc0e2b7, 0x8846e00284c4cf1f); + 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 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); + pub(super) const NUMBER64: u128 = seed(0x2400e8b894ce9c2a, 0x790be1eabd7b9481); +} + +/// Structural 128-bit `DistinctHash` of a JSON value, seeded from the root — +/// the canonical entry point used to hash a complex `ORDER BY` resume value. +pub(crate) fn distinct_hash(item: &OrderByItem) -> u128 { + structural_hash(item, seeds::ROOT) +} + +/// Hashes an element under `seed`. The type seed is always mixed in first +/// (`Murmur(type_seed_bytes, seed)`); .NET's `RootCache` is just memoization +/// of this exact expression for `seed == ROOT`, so no root special-case is +/// needed here. +fn structural_hash(item: &OrderByItem, seed: u128) -> u128 { + match item { + // Undefined is ignored while hashing: it returns the seed untouched. + OrderByItem::Undefined => seed, + OrderByItem::Null => type_hash(seeds::NULL, seed), + OrderByItem::Boolean(false) => type_hash(seeds::FALSE, seed), + OrderByItem::Boolean(true) => type_hash(seeds::TRUE, seed), + OrderByItem::Number(number) => hash_number(*number, seed), + OrderByItem::String(value) => { + let hash = type_hash(seeds::STRING, seed); + murmurhash3_128(value.as_bytes(), hash) + } + OrderByItem::Array(items) => hash_array(items, seed), + OrderByItem::Object(properties) => hash_object(properties, seed), + } +} + +/// `Murmur(type_seed's little-endian bytes, seed)` — the first mix applied to +/// every element, distinguishing e.g. an empty array from an empty object. +fn type_hash(type_seed: u128, seed: u128) -> u128 { + murmurhash3_128(&type_seed.to_le_bytes(), seed) +} + +fn hash_array(items: &[OrderByItem], seed: u128) -> u128 { + let mut hash = type_hash(seeds::ARRAY, seed); + for (index, item) in items.iter().enumerate() { + // Undefined items are skipped, but the position index still advances. + if !matches!(item, OrderByItem::Undefined) { + let item_seed = seeds::ARRAY_INDEX.wrapping_add(index as u128); + let item_hash = structural_hash(item, item_seed); + hash = murmurhash3_128(&item_hash.to_le_bytes(), hash); + } + } + hash +} + +fn hash_object(properties: &[(String, OrderByItem)], seed: u128) -> u128 { + let hash = type_hash(seeds::OBJECT, seed); + // Each property value is seeded with a hash of its name, so property order + // can be XOR-folded (order-independent) without duplicate values cancelling. + let name_seed = murmurhash3_128(&seeds::STRING.to_le_bytes(), seeds::PROPERTY_NAME); + let mut intermediate: u128 = 0; + for (name, value) in properties { + if !matches!(value, OrderByItem::Undefined) { + let name_hash = murmurhash3_128(name.as_bytes(), name_seed); + intermediate ^= structural_hash(value, name_hash); + } + } + // Only fold in the properties for a non-empty object, so `{}` keeps a + // distinct hash. `intermediate` is unsigned, so `!= 0` matches .NET's + // `if (intermediateHash > 0)` exactly. + if intermediate != 0 { + murmurhash3_128(&intermediate.to_le_bytes(), hash) + } else { + hash + } +} + +fn hash_number(number: OrderByNumber, seed: u128) -> u128 { + let hash = type_hash(seeds::NUMBER64, seed); + let (double_value, extra_bits) = number_to_double_ex(number); + // .NET hashes the `DoubleEx` struct's raw bytes: `[double LE][u16 LE]`, + // 10 bytes total (`[StructLayout(Sequential, Pack = 2)]`). + let mut bytes = [0u8; 10]; + bytes[..8].copy_from_slice(&double_value.to_le_bytes()); + bytes[8..].copy_from_slice(&extra_bits.to_le_bytes()); + murmurhash3_128(&bytes, hash) +} + +/// Converts an `ORDER BY` number to the backend's `Number64.DoubleEx` +/// `(double, extraBits)` pair. Cosmos `Number64` is either a signed `long` or +/// a `double`, so an integer within `i64` range takes the exact long path and +/// anything larger (only reachable via a `u64` above `i64::MAX`) or a float +/// takes the double path. +fn number_to_double_ex(number: OrderByNumber) -> (f64, u16) { + match number { + OrderByNumber::I64(value) => long_to_double_ex(value), + OrderByNumber::U64(value) if value <= i64::MAX as u64 => long_to_double_ex(value as i64), + // Above i64::MAX, Number64 can only carry it as a double. + OrderByNumber::U64(value) => (value as f64, 0), + OrderByNumber::F64(value) => (value, 0), + } +} + +/// Exact port of `Number64.DoubleEx`'s `implicit operator DoubleEx(long)`: +/// small integers become a plain `double`, while an integer whose significant +/// bits span more than 52 positions is encoded losslessly across the double's +/// mantissa plus 16 `extraBits`. +fn long_to_double_ex(value: i64) -> (f64, u16) { + // long.MinValue is special-cased: Math.Abs would overflow. + if value == i64::MIN { + return (value as f64, 0); + } + + let abs_value = value.unsigned_abs(); + if abs_value != 0 { + let msb_index = 63 - abs_value.leading_zeros() as i32; + let lsb_index = abs_value.trailing_zeros() as i32; + // Only when the value needs more than a double's 52-bit mantissa + // (both spanning past bit 52 and covering more than 52 bit positions) + // do we take the extended-precision path. + if msb_index > 52 && (msb_index - lsb_index) > 52 { + let exponent_value = msb_index; + let exponent_bits: i64 = ((exponent_value as i64) + 1023) << 52; + // Shift in u64 to match .NET's unchecked `long` shift (the mask + // then clears the implicit leading 1, keeping 62 mantissa bits). + let mantissa = (abs_value << (62 - exponent_value)) & 0x3FFF_FFFF_FFFF_FFFF; + let extra_bits = ((mantissa & 0x3FF) << 6) as u16; + let mantissa = (mantissa >> 10) as i64; + let mut value_bits = exponent_bits | mantissa; + if value < 0 { + value_bits = (value_bits as u64 | 0x8000_0000_0000_0000) as i64; + } + return (f64::from_bits(value_bits as u64), extra_bits); + } + } + + (value as f64, 0) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Parses a `.NET`-style `"89-8E-..."` 16-byte hash dump into a `u128`. + /// The dump is the `UInt128` in big-endian order (most-significant byte + /// first), so `from_be_bytes` recovers the value. + fn expected(dump: &str) -> u128 { + let mut bytes = [0u8; 16]; + for (i, part) in dump.split('-').enumerate() { + bytes[i] = u8::from_str_radix(part, 16).expect("valid hex byte"); + } + u128::from_be_bytes(bytes) + } + + fn repeated_string_key() -> String { + // Kept as a repeat to avoid a long unknown word in the spell checker. + "asdf".repeat(6) + } + + // ── Mandatory .NET baseline vectors ────────────────────────────────── + + #[test] + fn baseline_null() { + assert_eq!( + distinct_hash(&OrderByItem::Null), + expected("89-8E-FB-F2-68-0D-AE-A3-9E-24-D6-AE-66-53-41-7D") + ); + } + + #[test] + fn baseline_false() { + assert_eq!( + distinct_hash(&OrderByItem::Boolean(false)), + expected("6E-88-AA-42-42-F0-17-B7-70-5F-12-58-1A-61-11-8F") + ); + } + + #[test] + fn baseline_true() { + assert_eq!( + distinct_hash(&OrderByItem::Boolean(true)), + expected("CD-4D-51-AB-86-9F-07-1F-F0-19-C3-E4-A5-5C-01-25") + ); + } + + #[test] + fn baseline_empty_string() { + assert_eq!( + distinct_hash(&OrderByItem::String(String::new())), + expected("A7-84-46-CF-28-70-29-A5-85-B3-95-A1-61-A2-A9-A1") + ); + } + + #[test] + fn baseline_string() { + assert_eq!( + distinct_hash(&OrderByItem::String(repeated_string_key())), + expected("8F-12-43-68-40-1E-0B-91-AD-58-0C-A5-42-E0-35-DB") + ); + } + + #[test] + fn baseline_empty_array() { + assert_eq!( + distinct_hash(&OrderByItem::Array(Vec::new())), + expected("7D-54-F7-29-FF-4A-63-FF-A2-EF-3C-A3-59-53-80-82") + ); + } + + #[test] + fn baseline_mixed_array() { + let array = OrderByItem::Array(vec![ + OrderByItem::Null, + OrderByItem::Boolean(false), + OrderByItem::Boolean(true), + OrderByItem::Number(1234_i64.into()), + OrderByItem::String(repeated_string_key()), + ]); + assert_eq!( + distinct_hash(&array), + expected("A9-40-79-16-DE-27-F3-16-CC-5A-8E-4E-B9-AF-D5-E6") + ); + } + + #[test] + fn baseline_empty_object() { + assert_eq!( + distinct_hash(&OrderByItem::Object(Vec::new())), + expected("14-B8-6A-60-AA-01-A7-38-14-38-74-3E-80-B9-65-D8") + ); + } + + #[test] + fn baseline_object() { + let object = OrderByItem::Object(vec![ + ("null".to_owned(), OrderByItem::Null), + ("false".to_owned(), OrderByItem::Boolean(false)), + ("true".to_owned(), OrderByItem::Boolean(true)), + ( + "cosmosNumber".to_owned(), + OrderByItem::Number(1234_i64.into()), + ), + ( + "cosmosString".to_owned(), + OrderByItem::String(repeated_string_key()), + ), + ]); + assert_eq!( + distinct_hash(&object), + expected("3B-6A-A9-F4-F4-3A-AE-C0-E4-8E-BE-2B-C0-20-D5-5C") + ); + } + + #[test] + fn baseline_number_1234() { + assert_eq!( + distinct_hash(&OrderByItem::Number(1234_i64.into())), + expected("A9-4B-F6-13-35-C9-FB-A4-2C-28-D7-D9-89-5D-14-34") + ); + } + + // ── Structural properties ──────────────────────────────────────────── + + #[test] + fn object_hash_is_property_order_independent() { + let forward = OrderByItem::Object(vec![ + ("a".to_owned(), OrderByItem::Number(1_i64.into())), + ("b".to_owned(), OrderByItem::String("x".to_owned())), + ("c".to_owned(), OrderByItem::Null), + ]); + let reversed = OrderByItem::Object(vec![ + ("c".to_owned(), OrderByItem::Null), + ("b".to_owned(), OrderByItem::String("x".to_owned())), + ("a".to_owned(), OrderByItem::Number(1_i64.into())), + ]); + assert_eq!(distinct_hash(&forward), distinct_hash(&reversed)); + } + + #[test] + fn array_hash_is_order_sensitive() { + let forward = OrderByItem::Array(vec![ + OrderByItem::Boolean(true), + OrderByItem::Boolean(false), + ]); + let swapped = OrderByItem::Array(vec![ + OrderByItem::Boolean(false), + OrderByItem::Boolean(true), + ]); + assert_ne!(distinct_hash(&forward), distinct_hash(&swapped)); + } + + #[test] + fn distinct_scalars_hash_differently() { + assert_ne!( + distinct_hash(&OrderByItem::Null), + distinct_hash(&OrderByItem::Boolean(false)) + ); + assert_ne!( + distinct_hash(&OrderByItem::Number(1_i64.into())), + distinct_hash(&OrderByItem::Number(2_i64.into())) + ); + } + + // ── DoubleEx conversion (hand-computed from the .NET operator) ──────── + // + // These assert the `long -> DoubleEx` conversion (double bits + extra + // bits) hand-derived from .NET's operator. There is deliberately no + // full end-to-end `distinct_hash` baseline vector for an *extended* + // DoubleEx value (extra bits != 0): no published .NET hash dump was + // available for one, and deriving an "expected" value from this crate's + // own implementation would be a circular self-test, not an independent + // baseline. `adjacent_large_integers_hash_differently` still guards the + // lossless path structurally. TODO: add extended-DoubleEx baseline + // vectors here once a published .NET dump is available. + + #[test] + fn double_ex_small_integer_is_plain_double() { + assert_eq!(long_to_double_ex(1234), (1234.0_f64, 0)); + assert_eq!(long_to_double_ex(-1234), (-1234.0_f64, 0)); + assert_eq!(long_to_double_ex(0), (0.0_f64, 0)); + } + + #[test] + fn double_ex_min_value_special_case() { + assert_eq!(long_to_double_ex(i64::MIN), (i64::MIN as f64, 0)); + } + + #[test] + fn double_ex_two_pow_53_plus_one() { + // 2^53 + 1: msb 53, lsb 0 -> extended path. Hand-computed: + // doubleValue = 2^53 (0x4340_0000_0000_0000), extraBits = 0x8000. + let (double_value, extra_bits) = long_to_double_ex(9_007_199_254_740_993); + assert_eq!(double_value.to_bits(), 0x4340_0000_0000_0000); + assert_eq!(extra_bits, 0x8000); + } + + #[test] + fn double_ex_i64_max() { + // i64::MAX = 2^63 - 1: msb 62, lsb 0 -> extended path. Hand-computed: + // valueBits = 0x43DF_FFFF_FFFF_FFFF, extraBits = 0xFFC0. + let (double_value, extra_bits) = long_to_double_ex(i64::MAX); + assert_eq!(double_value.to_bits(), 0x43DF_FFFF_FFFF_FFFF); + assert_eq!(extra_bits, 0xFFC0); + } + + #[test] + fn double_ex_negative_extended_sets_sign_bit() { + // The negative of the 2^53+1 case: identical magnitude, sign bit set. + let (positive, positive_extra) = long_to_double_ex(9_007_199_254_740_993); + let (negative, negative_extra) = long_to_double_ex(-9_007_199_254_740_993); + assert_eq!(negative_extra, positive_extra); + assert_eq!( + negative.to_bits(), + positive.to_bits() | 0x8000_0000_0000_0000 + ); + } + + #[test] + fn adjacent_large_integers_hash_differently() { + // Would collide only if the value were routed through a lossy f64. + assert_ne!( + distinct_hash(&OrderByItem::Number(9_007_199_254_740_992_i64.into())), + distinct_hash(&OrderByItem::Number(9_007_199_254_740_993_i64.into())) + ); + } + + #[test] + fn u64_above_i64_max_uses_double_path() { + // A u64 beyond i64::MAX can only be a Number64 double. + let big = u64::MAX; + assert_eq!( + number_to_double_ex(OrderByNumber::U64(big)), + (big as f64, 0) + ); + } +} diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/integration_tests/mod.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/integration_tests/mod.rs index e8c842f551a..4fb7f635a58 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/integration_tests/mod.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/integration_tests/mod.rs @@ -7,4 +7,5 @@ //! they cover; cross-layer scenarios live here. mod change_feed_resume; +mod order_by_resume; mod query_resume; diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/integration_tests/order_by_resume.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/integration_tests/order_by_resume.rs new file mode 100644 index 00000000000..e4a5c1e911b --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/integration_tests/order_by_resume.rs @@ -0,0 +1,1691 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +// cspell:ignore repolled + +//! Driver-level integration tests for the cross-partition streaming +//! `ORDER BY` pipeline (`StreamingOrderedMerge`). +//! +//! Composes the real planner, merge node, snapshot, and continuation-token +//! layers against the `MockRequestExecutor` / `MockTopologyProvider` from +//! `dataflow::mocks`, mirroring `query_resume.rs`'s pattern for +//! `SequentialDrain`. Each mocked backend page uses the rewritten-envelope +//! shape (`{"_rid": ..., "Documents": [{"_rid", "orderByItems", "payload"}]}`) +//! `StreamingOrderedMerge` expects. + +use std::sync::Arc; + +use super::super::{ + mocks::{MockRequestExecutor, MockTopologyProvider}, + order_by::{OrderByItem, OrderByResumeValue}, + planner::build_streaming_ordered_merge, + query_plan::{QueryInfo, QueryPlan, QueryRange, SortOrder}, + snapshot::{OrderByRangeToken, ValueBoundary}, + Pipeline, PipelineContext, PipelineNodeState, ResolvedRange, +}; +use crate::{ + diagnostics::DiagnosticsContextBuilder, + models::{ + effective_partition_key::EffectivePartitionKey, AccountReference, ActivityId, + ContainerProperties, ContainerReference, ContinuationToken, CosmosOperation, + CosmosResponse, CosmosResponseHeaders, CosmosStatus, FeedRange, MaxItemCountHint, + ResolvedToken, SystemProperties, + }, + options::DiagnosticsOptions, +}; + +// ── Test fixtures ─────────────────────────────────────────────────────────── + +fn test_account() -> AccountReference { + AccountReference::with_master_key( + url::Url::parse("https://test.documents.azure.com:443/").unwrap(), + "dGVzdA==", + ) +} + +fn test_container_props() -> ContainerProperties { + use std::borrow::Cow; + ContainerProperties { + id: Cow::Owned("coll".into()), + partition_key: serde_json::from_str(r#"{"paths":["/pk"]}"#).unwrap(), + system_properties: SystemProperties::default(), + } +} + +fn test_container() -> ContainerReference { + ContainerReference::new( + test_account(), + "db", + "db_rid", + "coll", + "coll_rid", + &test_container_props(), + ) +} + +fn order_by_operation() -> Arc { + Arc::new( + CosmosOperation::query_items(test_container(), Some(FeedRange::full())) + .with_body(br#"{"query":"SELECT * FROM c ORDER BY c.rank","parameters":[]}"#.to_vec()), + ) +} + +/// Like [`order_by_operation`] but with an explicit `max_item_count`, to +/// force a precise page boundary. +fn order_by_operation_with_page_size(n: u32) -> Arc { + Arc::new( + CosmosOperation::query_items(test_container(), Some(FeedRange::full())) + .with_body(br#"{"query":"SELECT * FROM c ORDER BY c.rank","parameters":[]}"#.to_vec()) + .with_max_item_count(MaxItemCountHint::Limit( + std::num::NonZeroU32::new(n).unwrap(), + )), + ) +} + +/// A single-ascending-column ORDER BY plan spanning the full container. +fn order_by_plan() -> QueryPlan { + QueryPlan { + partitioned_query_execution_info_version: 2, + query_info: Some(QueryInfo { + order_by: vec![SortOrder::Ascending], + order_by_expressions: vec!["c.rank".to_owned()], + rewritten_query: Some( + "SELECT c._rid, [{\"item\":c.rank}] AS orderByItems, c AS payload FROM c \ + WHERE {documentdb-formattableorderbyquery-filter} ORDER BY c.rank ASC" + .to_owned(), + ), + ..Default::default() + }), + query_ranges: vec![QueryRange { + min: String::new(), + max: "FF".to_string(), + is_min_inclusive: true, + is_max_inclusive: false, + }], + hybrid_search_query_info: None, + } +} + +/// Like [`order_by_plan`] but with `c.rank`'s sort `direction` configurable, +/// for DESC-specific coverage (`order_by_plan` stays ASC-only since many +/// existing tests depend on its exact shape). +fn order_by_plan_with_direction(direction: SortOrder) -> QueryPlan { + let keyword = match direction { + SortOrder::Ascending => "ASC", + SortOrder::Descending => "DESC", + }; + QueryPlan { + partitioned_query_execution_info_version: 2, + query_info: Some(QueryInfo { + order_by: vec![direction], + order_by_expressions: vec!["c.rank".to_owned()], + rewritten_query: Some(format!( + "SELECT c._rid, [{{\"item\":c.rank}}] AS orderByItems, c AS payload FROM c \ + WHERE {{documentdb-formattableorderbyquery-filter}} ORDER BY c.rank {keyword}" + )), + ..Default::default() + }), + query_ranges: vec![QueryRange { + min: String::new(), + max: "FF".to_string(), + is_min_inclusive: true, + is_max_inclusive: false, + }], + hybrid_search_query_info: None, + } +} + +/// A realistic, 16-byte-encoded document `_rid` (matching real Cosmos DB's +/// hierarchical RID layout) whose little-endian document-ordinal segment is +/// `doc_id` — exercises `models::resource_id::compare_document_rids`'s +/// numeric decode path, unlike the short synthetic rids used elsewhere in +/// this file (e.g. `"a"`, `"tied-1"`), which take its raw-string fallback. +fn real_rid(doc_id: u64) -> String { + let mut bytes = [0u8; 16]; + bytes[0..4].copy_from_slice(&[0x0A, 0x0B, 0x0C, 0x0D]); + bytes[4..8].copy_from_slice(&[0x80, 0x01, 0x02, 0x03]); + bytes[8..16].copy_from_slice(&doc_id.to_le_bytes()); + crate::models::resource_id::encode_rid(&bytes) +} + +fn resolved(min: &str, max: &str, pk_range_id: &str) -> ResolvedRange { + ResolvedRange { + partition_key_range_id: pk_range_id.to_string(), + range: FeedRange::new( + EffectivePartitionKey::from(min), + EffectivePartitionKey::from(max), + ) + .unwrap(), + } +} + +/// Builds a rewritten-envelope page with one row per `(rid, rank)` pair. +fn envelope_page(rows: &[(&str, i64)], continuation: Option<&str>) -> CosmosResponse { + let documents: Vec = rows + .iter() + .map(|(rid, rank)| { + serde_json::json!({ + "_rid": rid, + "orderByItems": [{"item": rank}], + "payload": {"id": rid, "rank": rank}, + }) + }) + .collect(); + let body = serde_json::json!({ + "_rid": "", + "Documents": documents, + "_count": documents.len(), + }); + let mut diagnostics = DiagnosticsContextBuilder::new( + ActivityId::new_uuid(), + Arc::new(DiagnosticsOptions::default()), + ); + diagnostics.set_operation_status(azure_core::http::StatusCode::Ok, None); + let mut headers = CosmosResponseHeaders::new(); + headers.continuation = continuation.map(str::to_owned); + headers.request_charge = Some(crate::models::RequestCharge::new(1.5)); + CosmosResponse::new( + serde_json::to_vec(&body).unwrap(), + headers, + CosmosStatus::new(azure_core::http::StatusCode::Ok), + Arc::new(diagnostics.complete()), + ) +} + +/// Like [`envelope_page`] but JOIN-shaped: `(rid, rank, id)` rows where +/// several rows may share one `_rid` (a single document expanded by a JOIN) +/// while carrying distinct payload `id`s. Drives the `skip_count` resume path. +fn join_envelope_page(rows: &[(&str, i64, &str)], continuation: Option<&str>) -> CosmosResponse { + let documents: Vec = rows + .iter() + .map(|(rid, rank, id)| { + serde_json::json!({ + "_rid": rid, + "orderByItems": [{"item": rank}], + "payload": {"id": id}, + }) + }) + .collect(); + let body = serde_json::json!({ + "_rid": "", + "Documents": documents, + "_count": documents.len(), + }); + let mut diagnostics = DiagnosticsContextBuilder::new( + ActivityId::new_uuid(), + Arc::new(DiagnosticsOptions::default()), + ); + diagnostics.set_operation_status(azure_core::http::StatusCode::Ok, None); + let mut headers = CosmosResponseHeaders::new(); + headers.continuation = continuation.map(str::to_owned); + headers.request_charge = Some(crate::models::RequestCharge::new(1.5)); + CosmosResponse::new( + serde_json::to_vec(&body).unwrap(), + headers, + CosmosStatus::new(azure_core::http::StatusCode::Ok), + Arc::new(diagnostics.complete()), + ) +} + +/// Like [`envelope_page`] but with a string sort key per row, for +/// string-boundary resume coverage. +fn string_envelope_page(rows: &[(&str, &str)], continuation: Option<&str>) -> CosmosResponse { + let documents: Vec = rows + .iter() + .map(|(rid, key)| { + serde_json::json!({ + "_rid": rid, + "orderByItems": [{"item": key}], + "payload": {"id": rid, "key": key}, + }) + }) + .collect(); + let body = serde_json::json!({ + "_rid": "", + "Documents": documents, + "_count": documents.len(), + }); + let mut diagnostics = DiagnosticsContextBuilder::new( + ActivityId::new_uuid(), + Arc::new(DiagnosticsOptions::default()), + ); + diagnostics.set_operation_status(azure_core::http::StatusCode::Ok, None); + let mut headers = CosmosResponseHeaders::new(); + headers.continuation = continuation.map(str::to_owned); + headers.request_charge = Some(crate::models::RequestCharge::new(1.5)); + CosmosResponse::new( + serde_json::to_vec(&body).unwrap(), + headers, + CosmosStatus::new(azure_core::http::StatusCode::Ok), + Arc::new(diagnostics.complete()), + ) +} + +/// Like [`envelope_page`] but with a complex array sort key `[value]` per +/// row, for complex-boundary resume coverage. +fn array_envelope_page(rows: &[(&str, i64)], continuation: Option<&str>) -> CosmosResponse { + let documents: Vec = rows + .iter() + .map(|(rid, value)| { + serde_json::json!({ + "_rid": rid, + "orderByItems": [{"item": [value]}], + "payload": {"id": rid}, + }) + }) + .collect(); + let body = serde_json::json!({ + "_rid": "", + "Documents": documents, + "_count": documents.len(), + }); + let mut diagnostics = DiagnosticsContextBuilder::new( + ActivityId::new_uuid(), + Arc::new(DiagnosticsOptions::default()), + ); + diagnostics.set_operation_status(azure_core::http::StatusCode::Ok, None); + let mut headers = CosmosResponseHeaders::new(); + headers.continuation = continuation.map(str::to_owned); + headers.request_charge = Some(crate::models::RequestCharge::new(1.5)); + CosmosResponse::new( + serde_json::to_vec(&body).unwrap(), + headers, + CosmosStatus::new(azure_core::http::StatusCode::Ok), + Arc::new(diagnostics.complete()), + ) +} + +/// Extracts every item's `id` across `page`'s `Documents`, in wire order. +fn ids_in_page(page: &CosmosResponse) -> Vec { + let value: serde_json::Value = serde_json::from_slice(page.body_bytes()).unwrap(); + value["Documents"] + .as_array() + .unwrap() + .iter() + .map(|item| item["id"].as_str().unwrap().to_owned()) + .collect() +} + +async fn drain_all(pipeline: &mut Pipeline, executor: &mut MockRequestExecutor) -> Vec { + let mut ids = Vec::new(); + let mut topology = super::super::mocks::NoopTopologyProvider; + loop { + let mut context = PipelineContext::new(executor, Some(&mut topology)); + match pipeline.next_page(&mut context).await.unwrap() { + Some(response) => ids.extend(ids_in_page(&response)), + None => break, + } + } + ids +} + +/// Like [`drain_all`] but supplies a real topology provider, needed when a +/// split happens during iteration. +async fn drain_all_with_topology( + pipeline: &mut Pipeline, + executor: &mut MockRequestExecutor, + topology: &mut MockTopologyProvider, +) -> Vec { + let mut ids = Vec::new(); + loop { + let mut context = PipelineContext::new(executor, Some(topology)); + match pipeline.next_page(&mut context).await.unwrap() { + Some(response) => ids.extend(ids_in_page(&response)), + None => break, + } + } + ids +} + +async fn drain_one(pipeline: &mut Pipeline, executor: &mut MockRequestExecutor) -> Vec { + let mut topology = super::super::mocks::NoopTopologyProvider; + let mut context = PipelineContext::new(executor, Some(&mut topology)); + let response = pipeline + .next_page(&mut context) + .await + .unwrap() + .expect("expected a page, not drained"); + ids_in_page(&response) +} + +fn round_trip_state(state: PipelineNodeState, op: &CosmosOperation) -> PipelineNodeState { + let token = ContinuationToken::encode_v1(op, &state).expect("encode succeeds"); + let resolved = token.resolve().expect("decode succeeds"); + match resolved { + ResolvedToken::ClientV1(token_state) => { + token_state + .is_valid_for_operation(op) + .expect("operation compatible"); + token_state.into_root_node_state() + } + ResolvedToken::ServerOpaque(_) => panic!("expected ClientV1 token"), + } +} + +// ── Tests ──────────────────────────────────────────────────────────────── + +/// Baseline: two partitions each locally sorted ascending by `rank`; the +/// merge must interleave them into one globally sorted stream. +#[tokio::test] +async fn merges_two_partitions_into_global_order() { + let op = order_by_operation(); + let plan = order_by_plan(); + + let mut topology = MockTopologyProvider::new(vec![Ok(vec![ + resolved("", "80", "pk-left"), + resolved("80", "FF", "pk-right"), + ])]); + let mut executor = MockRequestExecutor::new(vec![ + Ok(envelope_page(&[("l1", 1), ("l2", 3), ("l3", 5)], None)), + Ok(envelope_page(&[("r1", 2), ("r2", 4), ("r3", 6)], None)), + ]); + + let mut pipeline = build_streaming_ordered_merge(&plan, &mut topology, &op, None) + .await + .unwrap(); + let ids = drain_all(&mut pipeline, &mut executor).await; + + assert_eq!( + ids, + vec!["l1", "r1", "l2", "r2", "l3", "r3"] + .into_iter() + .map(str::to_owned) + .collect::>(), + "rows must interleave in ascending rank order across both partitions" + ); +} + +/// A single partition, single page: the trivial case must still flow +/// through the merge machinery correctly (no fan-out needed). +#[tokio::test] +async fn single_partition_passthrough() { + let op = order_by_operation(); + let plan = order_by_plan(); + + let mut topology = MockTopologyProvider::new(vec![Ok(vec![resolved("", "FF", "pk-0")])]); + let mut executor = + MockRequestExecutor::new(vec![Ok(envelope_page(&[("a", 1), ("b", 2)], None))]); + + let mut pipeline = build_streaming_ordered_merge(&plan, &mut topology, &op, None) + .await + .unwrap(); + let ids = drain_all(&mut pipeline, &mut executor).await; + assert_eq!(ids, vec!["a".to_owned(), "b".to_owned()]); +} + +/// An empty backend page carrying a continuation must be transparently +/// re-polled, not mistaken for "drained" or surfaced as an empty result. +#[tokio::test] +async fn empty_page_with_continuation_is_repolled() { + let op = order_by_operation(); + let plan = order_by_plan(); + + let mut topology = MockTopologyProvider::new(vec![Ok(vec![resolved("", "FF", "pk-0")])]); + let mut executor = MockRequestExecutor::new(vec![ + Ok(envelope_page(&[], Some("ct-empty"))), + Ok(envelope_page(&[("a", 1)], None)), + ]); + + let mut pipeline = build_streaming_ordered_merge(&plan, &mut topology, &op, None) + .await + .unwrap(); + let ids = drain_all(&mut pipeline, &mut executor).await; + assert_eq!(ids, vec!["a".to_owned()]); + assert_eq!( + executor.continuation_calls, + vec![None, Some("ct-empty".to_owned())] + ); +} + +/// Empty total result must surface as a single empty, terminal page +/// (matching a plain `Request`), not an error. +#[tokio::test] +async fn empty_total_result_surfaces_as_terminal_empty_page() { + let op = order_by_operation(); + let plan = order_by_plan(); + + let mut topology = MockTopologyProvider::new(vec![Ok(vec![resolved("", "FF", "pk-0")])]); + let mut executor = MockRequestExecutor::new(vec![Ok(envelope_page(&[], None))]); + + let mut pipeline = build_streaming_ordered_merge(&plan, &mut topology, &op, None) + .await + .unwrap(); + let ids = drain_all(&mut pipeline, &mut executor).await; + assert!(ids.is_empty()); +} + +/// Baseline resume: session 1 drains one page (continuation "ct-1" +/// pending); session 2 must forward it, not fresh-start. +#[tokio::test] +async fn resume_with_unchanged_topology_forwards_continuation() { + let op = order_by_operation_with_page_size(1); + let plan = order_by_plan(); + + let mut topology1 = MockTopologyProvider::new(vec![Ok(vec![resolved("", "FF", "pk-0")])]); + let mut executor1 = + MockRequestExecutor::new(vec![Ok(envelope_page(&[("a", 1)], Some("ct-1")))]); + let mut pipeline1 = build_streaming_ordered_merge(&plan, &mut topology1, &op, None) + .await + .unwrap(); + let ids1 = drain_one(&mut pipeline1, &mut executor1).await; + assert_eq!(ids1, vec!["a".to_owned()]); + let state = pipeline1.snapshot_state().unwrap(); + match &state { + PipelineNodeState::StreamingOrderedMerge { ranges, .. } => { + assert_eq!(ranges.len(), 1); + assert_eq!(ranges[0].server_continuation, Some("ct-1".to_owned())); + assert!( + ranges[0].boundary.is_some(), + "boundary must be recorded even though a plain continuation is also available" + ); + } + other => panic!("expected StreamingOrderedMerge snapshot, got {other:?}"), + } + drop(pipeline1); + + let resumed_state = round_trip_state(state, &op); + let mut topology2 = MockTopologyProvider::new(vec![Ok(vec![resolved("", "FF", "pk-0")])]); + let mut executor2 = MockRequestExecutor::new(vec![Ok(envelope_page(&[("b", 2)], None))]); + let mut pipeline2 = + build_streaming_ordered_merge(&plan, &mut topology2, &op, Some(resumed_state)) + .await + .unwrap(); + let ids2 = drain_all(&mut pipeline2, &mut executor2).await; + assert_eq!(ids2, vec!["b".to_owned()]); + assert_eq!( + executor2.continuation_calls, + vec![Some("ct-1".to_owned())], + "resume must reuse the saved plain continuation when topology is unchanged" + ); +} + +/// Regression: a mid-page checkpoint (no safe `server_continuation`) +/// resumed against unchanged topology must apply the value-boundary +/// resume path, not a fresh restart that re-emits delivered rows. +#[tokio::test] +async fn resume_with_unchanged_topology_and_no_saved_continuation_does_not_duplicate_rows() { + let op = order_by_operation_with_page_size(1); + let plan = order_by_plan(); + + // `max_item_count = 1` surfaces one row per call, leaving two rows + // buffered — the "mid-page" case forcing `server_continuation = None`. + let mut topology1 = MockTopologyProvider::new(vec![Ok(vec![resolved("", "FF", "pk-0")])]); + let mut executor1 = MockRequestExecutor::new(vec![Ok(envelope_page( + &[("a", 1), ("b", 2), ("c", 3)], + None, + ))]); + let mut pipeline1 = build_streaming_ordered_merge(&plan, &mut topology1, &op, None) + .await + .unwrap(); + let ids1 = drain_one(&mut pipeline1, &mut executor1).await; + assert_eq!(ids1, vec!["a".to_owned()]); + + let state = pipeline1.snapshot_state().unwrap(); + match &state { + PipelineNodeState::StreamingOrderedMerge { ranges, .. } => { + assert_eq!(ranges.len(), 1); + assert_eq!( + ranges[0].server_continuation, None, + "buffer held unread rows at checkpoint time, so no plain \ + continuation is safe to save" + ); + assert!( + ranges[0].boundary.is_some(), + "a row was already emitted, so a resume boundary must be recorded" + ); + } + other => panic!("expected StreamingOrderedMerge snapshot, got {other:?}"), + } + drop(pipeline1); + + let resumed_state = round_trip_state(state, &op); + // Mock re-returns the full unfiltered page (mock can't evaluate SQL + // filters); the client-side discard must strip the "a" row back out. + let mut topology2 = MockTopologyProvider::new(vec![Ok(vec![resolved("", "FF", "pk-0")])]); + let mut executor2 = MockRequestExecutor::new(vec![Ok(envelope_page( + &[("a", 1), ("b", 2), ("c", 3)], + None, + ))]); + let mut pipeline2 = + build_streaming_ordered_merge(&plan, &mut topology2, &op, Some(resumed_state)) + .await + .unwrap(); + let ids2 = drain_all(&mut pipeline2, &mut executor2).await; + assert_eq!( + ids2, + vec!["b".to_owned(), "c".to_owned()], + "row \"a\" was already emitted before the checkpoint and must not \ + be re-emitted on resume" + ); +} + +/// Catalog-sourced regression for `equal_key_resume_requiring_skip_count`: +/// resuming a tied-row value boundary on unchanged topology must apply +/// the `_rid`-aware discard, not a fresh restart. +#[tokio::test] +async fn catalog_equal_key_resume_requiring_skip_count_replays_correctly() { + const CATALOG_JSON: &str = + include_str!("../../../../tests/fixtures/streaming_order_by_scenarios.json"); + + // Parse generically and pick out only the one scenario this test needs. + let catalog: serde_json::Value = + serde_json::from_str(CATALOG_JSON).expect("catalog must parse as JSON"); + let scenario = catalog["scenarios"] + .as_array() + .expect("catalog must have a scenarios array") + .iter() + .find(|s| s["id"] == "equal_key_resume_requiring_skip_count") + .expect("catalog must contain scenario equal_key_resume_requiring_skip_count"); + + let expected_ids: Vec = scenario["expectedIds"] + .as_array() + .expect("scenario must declare expectedIds") + .iter() + .map(|v| { + v.as_str() + .expect("expectedIds entries must be strings") + .to_owned() + }) + .collect(); + + let rows: Vec<(String, i64, String)> = scenario["mock"]["partitions"][0]["pages"][0]["rows"] + .as_array() + .expect("scenario must declare mock rows") + .iter() + .map(|row| { + let rid = row["_rid"] + .as_str() + .or_else(|| row["rid"].as_str()) + .expect("row must have a rid") + .to_owned(); + let item = row["orderByItems"][0]["item"] + .as_i64() + .expect("this scenario's sort key is a plain integer"); + // A JOIN row's payload `id` is distinct from its `_rid` (several + // rows share one document's `_rid`), so read it explicitly. + let id = row["payload"]["id"] + .as_str() + .expect("row payload must have an id") + .to_owned(); + (rid, item, id) + }) + .collect(); + let row_refs: Vec<(&str, i64, &str)> = rows + .iter() + .map(|(rid, rank, id)| (rid.as_str(), *rank, id.as_str())) + .collect(); + + let checkpoint = &scenario["checkpoint"]; + let resume_values: Vec = + serde_json::from_value(checkpoint["resumeValues"].clone()) + .expect("checkpoint.resumeValues must parse as OrderByResumeValue"); + let last_rid = checkpoint["lastRid"] + .as_str() + .expect("checkpoint must declare lastRid") + .to_owned(); + // A real boundary always emitted at least the boundary row, so a fixture + // that omits `skipCount` resumes as 1. + let skip_count = checkpoint + .get("skipCount") + .and_then(|v| v.as_u64()) + .map_or(1, |n| n as u32); + + let op = order_by_operation(); + let plan = order_by_plan(); + let resumed_state = PipelineNodeState::StreamingOrderedMerge { + directions: vec![SortOrder::Ascending], + ranges: vec![OrderByRangeToken { + min_epk: String::new(), + max_epk: "FF".to_owned(), + server_continuation: None, + boundary: Some(ValueBoundary { + resume_values, + last_rid, + skip_count, + }), + }], + }; + + let mut topology = MockTopologyProvider::new(vec![Ok(vec![resolved("", "FF", "pk-0")])]); + let mut executor = MockRequestExecutor::new(vec![Ok(join_envelope_page(&row_refs, None))]); + let mut pipeline = + build_streaming_ordered_merge(&plan, &mut topology, &op, Some(resumed_state)) + .await + .unwrap(); + let ids = drain_all(&mut pipeline, &mut executor).await; + assert_eq!( + ids, expected_ids, + "scenario {} drained ids do not match the catalog's expectedIds", + scenario["id"] + ); +} + +/// Split during live iteration: the merge must splice in replacement +/// ranges and keep global ordering across the remaining rows. +#[tokio::test] +async fn split_mid_merge_splices_replacements_and_preserves_order() { + let op = order_by_operation(); + let plan = order_by_plan(); + + // The split child (a `Request`) resolves the new topology once when it + // produces its replacement nodes; the merge consumes those directly, so no + // second resolution is queued. + let mut topology = MockTopologyProvider::new(vec![ + Ok(vec![resolved("", "FF", "pk-0")]), + Ok(vec![ + resolved("", "80", "pk-left"), + resolved("80", "FF", "pk-right"), + ]), + ]); + let mut executor = MockRequestExecutor::new(vec![ + Err(super::super::mocks::gone_error()), + Ok(envelope_page(&[("l1", 1), ("l2", 3)], None)), + Ok(envelope_page(&[("r1", 2), ("r2", 4)], None)), + ]); + + let mut pipeline = build_streaming_ordered_merge(&plan, &mut topology, &op, None) + .await + .unwrap(); + let ids = drain_all_with_topology(&mut pipeline, &mut executor, &mut topology).await; + assert_eq!( + ids, + vec!["l1", "r1", "l2", "r2"] + .into_iter() + .map(str::to_owned) + .collect::>() + ); +} + +/// Regression for the in-flight split ordering defect: the split happens +/// while *replenishing the popped winner mid-pop-loop* (not during the +/// initial prime), fanning P0 into two sub-ranges. Both replacements must be +/// primed before the next selection, or P0b's `10, 20` would be skipped and +/// P1's `50` emitted ahead of them. A default (large) page cap keeps popping +/// within a single page so the mis-ordering would surface immediately. +#[tokio::test] +async fn split_during_replenish_primes_all_replacements_preserving_order() { + let op = order_by_operation(); + let plan = order_by_plan(); + + // Initial resolve yields two live partitions; P0 then splits, resolving + // the post-split topology once when it produces its replacement nodes + // (the merge consumes them directly, without a second resolution). + let mut topology = MockTopologyProvider::new(vec![ + Ok(vec![ + resolved("", "80", "pk-p0"), + resolved("80", "FF", "pk-p1"), + ]), + Ok(vec![ + resolved("", "40", "pk-a"), + resolved("40", "80", "pk-b"), + ]), + ]); + // P0's first page delivers 1, 2 (continuation pending); replenishing it + // 410s into a split. Each replacement carries P0's forwarded continuation + // (Cosmos contract), so it resumes past the `2` boundary with no client + // discard: P0a resumes with 3, P0b with 10, 20. + let mut executor = MockRequestExecutor::new(vec![ + Ok(envelope_page(&[("d1", 1), ("d2", 2)], Some("p0-ct"))), + Ok(envelope_page(&[("d50", 50)], None)), + Err(super::super::mocks::gone_error()), + Ok(envelope_page(&[("d3", 3)], None)), + Ok(envelope_page(&[("d10", 10), ("d20", 20)], None)), + ]); + + let mut pipeline = build_streaming_ordered_merge(&plan, &mut topology, &op, None) + .await + .unwrap(); + let ids = drain_all_with_topology(&mut pipeline, &mut executor, &mut topology).await; + assert_eq!( + ids, + vec!["d1", "d2", "d3", "d10", "d20", "d50"] + .into_iter() + .map(str::to_owned) + .collect::>(), + "the second split replacement's rows (10, 20) must precede 50" + ); +} + +/// Companion to the large-cap regression: a page cap that fills exactly as a +/// mid-pop-loop split completes must checkpoint immediately after the split +/// (snapshotting both replacements), then resume in global order. Page 1 +/// yields 1, 2, 3; the persisted token resumes 10, 20, 50 across the split. +#[tokio::test] +async fn split_during_replenish_checkpoint_resumes_in_global_order() { + let op = order_by_operation_with_page_size(3); + let plan = order_by_plan(); + + let mut topology1 = MockTopologyProvider::new(vec![ + Ok(vec![ + resolved("", "80", "pk-p0"), + resolved("80", "FF", "pk-p1"), + ]), + Ok(vec![ + resolved("", "40", "pk-a"), + resolved("40", "80", "pk-b"), + ]), + ]); + let mut executor1 = MockRequestExecutor::new(vec![ + Ok(envelope_page(&[("d1", 1), ("d2", 2)], Some("p0-ct"))), + Ok(envelope_page(&[("d50", 50)], None)), + Err(super::super::mocks::gone_error()), + Ok(envelope_page(&[("d3", 3)], None)), + Ok(envelope_page(&[("d10", 10), ("d20", 20)], None)), + ]); + + let mut pipeline1 = build_streaming_ordered_merge(&plan, &mut topology1, &op, None) + .await + .unwrap(); + let page1 = { + let mut context = PipelineContext::new(&mut executor1, Some(&mut topology1)); + pipeline1 + .next_page(&mut context) + .await + .unwrap() + .expect("expected a first page") + }; + assert_eq!( + ids_in_page(&page1), + vec!["d1".to_owned(), "d2".to_owned(), "d3".to_owned()], + "page 1 stops at the cap right after the split" + ); + + let state = pipeline1.snapshot_state().unwrap(); + match &state { + PipelineNodeState::StreamingOrderedMerge { ranges, .. } => { + assert_eq!( + ranges.len(), + 2, + "both post-split children (P0b + P1) must survive into the checkpoint" + ); + } + other => panic!("expected StreamingOrderedMerge snapshot, got {other:?}"), + } + drop(pipeline1); + + let resumed_state = round_trip_state(state, &op); + // On resume both saved ranges are unchanged; P0b re-seeks past its value + // boundary (mock returns 10, 20 unfiltered), P1 restarts fresh (50). + let mut topology2 = MockTopologyProvider::new(vec![ + Ok(vec![resolved("40", "80", "pk-b")]), + Ok(vec![resolved("80", "FF", "pk-p1")]), + ]); + let mut executor2 = MockRequestExecutor::new(vec![ + Ok(envelope_page(&[("d10", 10), ("d20", 20)], None)), + Ok(envelope_page(&[("d50", 50)], None)), + ]); + let mut pipeline2 = + build_streaming_ordered_merge(&plan, &mut topology2, &op, Some(resumed_state)) + .await + .unwrap(); + let resumed = drain_all(&mut pipeline2, &mut executor2).await; + assert_eq!( + resumed, + vec!["d10", "d20", "d50"] + .into_iter() + .map(str::to_owned) + .collect::>(), + "resume continues the global order across the split checkpoint" + ); +} + +/// Two full snapshot/resume cycles over a JOIN tie run: one document (`docA`) +/// expands into four result rows sharing its `_rid`, so the boundary's +/// `skip_count` must accumulate across cycles — each cycle emits two more +/// docA rows, and the next resume skips exactly the ones already delivered. +/// Every result id must appear exactly once across all three pages. +#[tokio::test] +async fn two_snapshot_resume_cycles_accumulate_join_skip_count() { + // Page size 2 forces a checkpoint mid-tie-run. The mock can't honor the + // structured resumeFilter, so it re-returns the whole page each cycle and + // the client-side `skip_count` discard trims the already-emitted prefix. + let op = order_by_operation_with_page_size(2); + let plan = order_by_plan(); + let full_page = || { + join_envelope_page( + &[ + ("docA", 5, "a1"), + ("docA", 5, "a2"), + ("docA", 5, "a3"), + ("docA", 5, "a4"), + ("docB", 6, "b1"), + ], + None, + ) + }; + + let mut collected: Vec = Vec::new(); + let mut resume: Option = None; + // Cycle 0 emits a1,a2; cycle 1 emits a3,a4; cycle 2 emits b1 (terminal). + for cycle in 0..3 { + let mut topology = MockTopologyProvider::new(vec![Ok(vec![resolved("", "FF", "pk-0")])]); + let mut executor = MockRequestExecutor::new(vec![Ok(full_page())]); + let mut pipeline = build_streaming_ordered_merge(&plan, &mut topology, &op, resume.clone()) + .await + .unwrap(); + collected.extend(drain_one(&mut pipeline, &mut executor).await); + if cycle < 2 { + resume = Some(round_trip_state(pipeline.snapshot_state().unwrap(), &op)); + } + } + + assert_eq!( + collected, + vec!["a1", "a2", "a3", "a4", "b1"] + .into_iter() + .map(str::to_owned) + .collect::>(), + "each JOIN result row is emitted exactly once across two resume cycles" + ); +} + +/// Regression: after a split, both sub-ranges resume from the same +/// boundary; the `_rid`-aware discard must avoid dropping/duplicating rows. +#[tokio::test] +async fn resume_after_split_with_emitted_ties_has_no_omissions_or_duplicates() { + let op = order_by_operation(); + let plan = order_by_plan(); + let resumed_state = PipelineNodeState::StreamingOrderedMerge { + directions: vec![SortOrder::Ascending], + ranges: vec![OrderByRangeToken { + min_epk: String::new(), + max_epk: "FF".to_owned(), + server_continuation: None, + boundary: Some(ValueBoundary { + resume_values: vec![OrderByResumeValue::Number { value: 5.0.into() }], + last_rid: "c".to_owned(), + skip_count: 1, + }), + }], + }; + + // The saved range resolves to two post-split sub-ranges. + let mut topology = MockTopologyProvider::new(vec![Ok(vec![ + resolved("", "80", "pk-left"), + resolved("80", "FF", "pk-right"), + ])]); + // Mock can't honor the server-side filter, so each sub-range returns + // rows unfiltered. `[,80)` holds a, c (emitted), e (unemitted tie), + // m (rank 7); `[80,FF)` holds b (emitted), z (rank 6). + let mut executor = MockRequestExecutor::new(vec![ + Ok(envelope_page( + &[("a", 5), ("c", 5), ("e", 5), ("m", 7)], + None, + )), + Ok(envelope_page(&[("b", 5), ("z", 6)], None)), + ]); + + let mut pipeline = + build_streaming_ordered_merge(&plan, &mut topology, &op, Some(resumed_state)) + .await + .unwrap(); + let ids = drain_all(&mut pipeline, &mut executor).await; + assert_eq!( + ids, + vec!["e", "z", "m"] + .into_iter() + .map(str::to_owned) + .collect::>(), + "already-emitted tied rows (a, b, c) are dropped by `_rid`; the \ + unemitted tied row `e` survives with no duplicates" + ); +} + +/// A complex (array/object) boundary now also resumes across a split via +/// the structured `resumeFilter`, discarding already-emitted ties by hash + +/// `_rid` client-side — no more topology-change rejection. +#[tokio::test] +async fn resume_complex_boundary_across_split_has_no_omissions_or_duplicates() { + let op = order_by_operation(); + let plan = order_by_plan(); + // Boundary is the array `[5]`, last emitted at rid "c". Built from the + // same integer representation the envelope rows carry, so the hashes match. + let complex = OrderByItem::Array(vec![OrderByItem::Number(5_i64.into())]).to_resume_value(); + let resumed_state = PipelineNodeState::StreamingOrderedMerge { + directions: vec![SortOrder::Ascending], + ranges: vec![OrderByRangeToken { + min_epk: String::new(), + max_epk: "FF".to_owned(), + server_continuation: None, + boundary: Some(ValueBoundary { + resume_values: vec![complex], + last_rid: "c".to_owned(), + skip_count: 1, + }), + }], + }; + + // The saved range resolves to two post-split sub-ranges. + let mut topology = MockTopologyProvider::new(vec![Ok(vec![ + resolved("", "80", "pk-left"), + resolved("80", "FF", "pk-right"), + ])]); + // Every row ties on the same array `[5]` (a whole tie run split across + // the two sub-ranges), so the client-side hash + `_rid` discard is + // deterministic: rids at or before "c" (a, b, c) were already emitted; + // "d" and "e" survive. (A *differently*-valued complex row's ordering is + // the backend's hash order, which a `MockLeaf` can't reproduce, so the + // synthetic split test exercises the tie run.) + let mut executor = MockRequestExecutor::new(vec![ + Ok(array_envelope_page(&[("a", 5), ("c", 5), ("e", 5)], None)), + Ok(array_envelope_page(&[("b", 5), ("d", 5)], None)), + ]); + + let mut pipeline = + build_streaming_ordered_merge(&plan, &mut topology, &op, Some(resumed_state)) + .await + .expect("a complex boundary now resumes across a split"); + // Assert the resume-filtered request bodies carry the structured + // complex `resumeFilter` (reaching the executor), then check ordering. + let ids = drain_all(&mut pipeline, &mut executor).await; + for i in 0..2 { + assert_is_complex_resume_filtered_query(&executor.body_text(i)); + } + assert_eq!( + ids, + vec!["d", "e"] + .into_iter() + .map(str::to_owned) + .collect::>(), + "already-emitted tied rows (a, b, c) are dropped by hash+`_rid`; the \ + unemitted tied rows (d, e) survive with no duplicates" + ); +} + +/// Regression: a saved sub-range resolving to a wider merged partition +/// must be clipped to scope, not rejected as "over-covering". +#[tokio::test] +async fn resume_after_merge_clips_widened_range_and_drains() { + let op = order_by_operation(); + let plan = order_by_plan(); + let resumed_state = PipelineNodeState::StreamingOrderedMerge { + directions: vec![SortOrder::Ascending], + ranges: vec![OrderByRangeToken { + min_epk: String::new(), + max_epk: "80".to_owned(), + server_continuation: None, + boundary: Some(ValueBoundary { + resume_values: vec![OrderByResumeValue::Number { value: 5.0.into() }], + last_rid: "c".to_owned(), + skip_count: 1, + }), + }], + }; + + // Post-merge: the saved [,80) sub-range is now served by a wider + // physical partition [,FF). + let mut topology = MockTopologyProvider::new(vec![Ok(vec![resolved("", "FF", "pk-merged")])]); + let mut executor = MockRequestExecutor::new(vec![Ok(envelope_page( + &[("a", 5), ("c", 5), ("e", 5), ("n", 8)], + None, + ))]); + + let mut pipeline = + build_streaming_ordered_merge(&plan, &mut topology, &op, Some(resumed_state)) + .await + .expect("resume across a merge (widened physical range) must not be rejected"); + let ids = drain_all(&mut pipeline, &mut executor).await; + assert_eq!( + ids, + vec!["e", "n"] + .into_iter() + .map(str::to_owned) + .collect::>(), + "already-emitted rows (a, c) are dropped; the range's unemitted rows drain" + ); +} + +/// A continuation token shaped for a different node type (e.g. +/// `SequentialDrain`) must be rejected with a clear shape-mismatch error. +#[tokio::test] +async fn resume_rejects_wrong_node_shape() { + let op = order_by_operation(); + let plan = order_by_plan(); + + let wrong_shape = PipelineNodeState::SequentialDrain { + left_most_undrained_epk: String::new(), + active_tokens: vec![], + }; + + let mut topology = MockTopologyProvider::new(vec![]); + let err = build_streaming_ordered_merge(&plan, &mut topology, &op, Some(wrong_shape)) + .await + .unwrap_err(); + assert_eq!( + err.status(), + CosmosStatus::CLIENT_CONTINUATION_TOKEN_SHAPE_MISMATCH + ); +} + +/// A `Drained` continuation token must short-circuit to a drained pipeline +/// without issuing any requests. +#[tokio::test] +async fn resume_from_drained_short_circuits() { + let op = order_by_operation(); + let plan = order_by_plan(); + + let mut topology = MockTopologyProvider::new(vec![]); + let mut executor = MockRequestExecutor::new(vec![]); + let mut pipeline = + build_streaming_ordered_merge(&plan, &mut topology, &op, Some(PipelineNodeState::Drained)) + .await + .unwrap(); + let ids = drain_all(&mut pipeline, &mut executor).await; + assert!(ids.is_empty()); +} + +/// Request charge is summed across every backend page; the emitted page +/// never carries a raw backend continuation header. +#[tokio::test] +async fn aggregates_request_charge_and_omits_backend_continuation_header() { + let op = order_by_operation(); + let plan = order_by_plan(); + + let mut topology = MockTopologyProvider::new(vec![Ok(vec![ + resolved("", "80", "pk-left"), + resolved("80", "FF", "pk-right"), + ])]); + let mut executor = MockRequestExecutor::new(vec![ + Ok(envelope_page(&[("l1", 1)], None)), + Ok(envelope_page(&[("r1", 2)], None)), + ]); + + let mut pipeline = build_streaming_ordered_merge(&plan, &mut topology, &op, None) + .await + .unwrap(); + let mut noop_topology = super::super::mocks::NoopTopologyProvider; + let mut context = PipelineContext::new(&mut executor, Some(&mut noop_topology)); + let response = pipeline + .next_page(&mut context) + .await + .unwrap() + .expect("expected a page"); + assert_eq!( + response.headers().request_charge, + Some(crate::models::RequestCharge::new(3.0)), + "charge from both partitions' pages must be summed" + ); + assert!(response.headers().continuation.is_none()); +} + +// ── Query-shape / continuation-binding regression ─────────────────────────── + +/// Asserts a recorded request body carries the .NET-compatible structured +/// `resumeFilter` (Rust's per-range "target" style: `rid` present, +/// `exclude:false`) — not a fresh query, an SDK-generated `@cosmos...` +/// parameter, or an inline SQL seek predicate. +fn assert_is_resume_filtered_query(body: &str) { + let value: serde_json::Value = + serde_json::from_str(body).expect("recorded request body must be valid JSON"); + assert!( + !body.contains("{documentdb-formattableorderbyquery-filter}"), + "the resume filter placeholder must be substituted, not sent verbatim: {body}" + ); + assert!( + !body.contains("@cosmosResumeFilter"), + "resume must not synthesize query parameters: {body}" + ); + let query = value["query"].as_str().expect("body has a query string"); + assert!( + !query.contains("IS_NUMBER(") && !query.contains(" > @") && !query.contains(" < @"), + "resume must not rewrite the SQL with a seek predicate: {query}" + ); + let filter = &value["resumeFilter"]; + assert!( + filter.is_object(), + "a value-boundary resume must carry a structured resumeFilter: {body}" + ); + assert!( + filter["value"].is_array() && !filter["value"].as_array().unwrap().is_empty(), + "resumeFilter.value must be a non-empty array: {body}" + ); + assert!( + filter["rid"].is_string(), + "Rust resumes each range target-style with its own rid present: {body}" + ); + assert_eq!( + filter["exclude"], + serde_json::Value::Bool(false), + "Rust resumes target-style with exclude:false: {body}" + ); +} + +/// Like [`assert_is_resume_filtered_query`], but for a complex (array) +/// boundary: the single resume value is the typed hash object with signed +/// `low`/`high` halves. +fn assert_is_complex_resume_filtered_query(body: &str) { + let value: serde_json::Value = + serde_json::from_str(body).expect("recorded request body must be valid JSON"); + let entry = &value["resumeFilter"]["value"][0]; + assert_eq!( + entry["type"], "array", + "a complex array boundary serializes as a typed hash: {body}" + ); + assert!( + entry["low"].is_i64() && entry["high"].is_i64(), + "the hash halves are signed i64 (matching .NET's (long)ulong): {body}" + ); + assert_eq!( + value["resumeFilter"]["exclude"], + serde_json::Value::Bool(false) + ); +} + +/// Shared driver for a fresh → mid-page resume-filtered → second-resume +/// cycle, asserting the resume-filtered query's backend continuation is +/// never snapshotted into the plain `server_continuation` nor replayed +/// against the plain query on the next resume. +/// +/// This is the end-to-end regression for the query-shape binding bug: +/// before the fix, once a resume-filtered child reached an empty buffer +/// with a live backend continuation, `snapshot_state` stored that opaque, +/// filtered-query-bound token in the plain field, and the next resume fed +/// it to the plain query — an opaque-token/query-shape mismatch (400 / +/// duplication / omission risk). +/// +/// Each session emits one row (`max_item_count = 1`): session 1 emits the +/// first row and buffers the rest (mid-page checkpoint); session 2 resumes +/// via the resume-filtered query, emits one row, and checkpoints with an +/// empty buffer while its continuation is still live; session 3 resumes +/// again and drains the last row. The mock can't evaluate the server-side +/// filter, so pages return unfiltered — the client-side discard strips +/// already-emitted rows. Since the mock replies FIFO regardless of the +/// token handed to it, drained order alone wouldn't prove token binding, +/// so recorded request bodies and continuation calls are inspected directly. +async fn run_resume_filtered_binding_cycle( + session1_page: &[(&str, i64)], + session2_page: &[(&str, i64)], + session3_page: &[(&str, i64)], + resume_filtered_ct: &str, + expected_order: &[&str], +) { + let op = order_by_operation_with_page_size(1); + let plan = order_by_plan(); + + // ── Session 1: fresh, mid-page checkpoint ─────────────────────────── + let mut topology1 = MockTopologyProvider::new(vec![Ok(vec![resolved("", "FF", "pk-0")])]); + let mut executor1 = + MockRequestExecutor::new(vec![Ok(envelope_page(session1_page, Some("s1-more")))]); + let mut pipeline1 = build_streaming_ordered_merge(&plan, &mut topology1, &op, None) + .await + .unwrap(); + let ids1 = drain_one(&mut pipeline1, &mut executor1).await; + assert_eq!(ids1.len(), 1, "page size 1 must surface exactly one row"); + // The fresh child ran the plain rewritten query: the Gateway + // resume-filter placeholder was replaced with `true` (never wrapped as + // an outer subquery, and never sent verbatim). + let fresh_body = executor1.body_text(0); + assert!( + !fresh_body.contains("SELECT VALUE r FROM ("), + "session 1 must issue the plain rewritten query, not a wrapped subquery: {fresh_body}" + ); + assert!( + !fresh_body.contains("{documentdb-formattableorderbyquery-filter}"), + "session 1 must substitute the fresh placeholder, not send it verbatim: {fresh_body}" + ); + assert!( + fresh_body.contains("WHERE true"), + "a fresh streaming ORDER BY query replaces the placeholder with `true`: {fresh_body}" + ); + + let state1 = pipeline1.snapshot_state().unwrap(); + match &state1 { + PipelineNodeState::StreamingOrderedMerge { ranges, .. } => { + assert_eq!(ranges.len(), 1); + assert_eq!( + ranges[0].server_continuation, None, + "mid-page checkpoint holds unread rows, so no plain continuation is safe" + ); + assert!( + ranges[0].boundary.is_some(), + "a row was emitted, so a scalar resume boundary must be recorded" + ); + } + other => panic!("expected StreamingOrderedMerge snapshot, got {other:?}"), + } + drop(pipeline1); + + // ── Session 2: scalar value-boundary resume-filtered query ────────── + let resumed1 = round_trip_state(state1, &op); + let mut topology2 = MockTopologyProvider::new(vec![Ok(vec![resolved("", "FF", "pk-0")])]); + let mut executor2 = MockRequestExecutor::new(vec![Ok(envelope_page( + session2_page, + Some(resume_filtered_ct), + ))]); + let mut pipeline2 = build_streaming_ordered_merge(&plan, &mut topology2, &op, Some(resumed1)) + .await + .unwrap(); + let ids2 = drain_one(&mut pipeline2, &mut executor2).await; + assert_eq!(ids2.len(), 1); + + // The resume must run the resume-filtered `_rid`-aware query with a + // *fresh* start (no forwarded backend continuation). + assert_eq!( + executor2.continuation_calls, + vec![None], + "a value-boundary resume starts the resume-filtered query fresh" + ); + assert_is_resume_filtered_query(&executor2.body_text(0)); + + let state2 = pipeline2.snapshot_state().unwrap(); + let (state2_continuation, state2_boundary_rid) = match &state2 { + PipelineNodeState::StreamingOrderedMerge { ranges, .. } => { + assert_eq!(ranges.len(), 1); + ( + ranges[0].server_continuation.clone(), + ranges[0].boundary.as_ref().map(|b| b.last_rid.clone()), + ) + } + other => panic!("expected StreamingOrderedMerge snapshot, got {other:?}"), + }; + // The core assertion: the resume-filtered query's live backend + // continuation is NOT captured into the plain `server_continuation` + // field. Before the fix this was `Some(resume_filtered_ct)`. + assert_eq!( + state2_continuation, None, + "a resume-filtered value-boundary child must never snapshot its backend \ + continuation into the plain server_continuation field" + ); + assert_eq!( + state2_boundary_rid.as_deref(), + Some(ids2[0].as_str()), + "the scalar boundary must advance to the row just emitted" + ); + drop(pipeline2); + + // ── Session 3: second resume ──────────────────────────────────────── + let resumed2 = round_trip_state(state2, &op); + let mut topology3 = MockTopologyProvider::new(vec![Ok(vec![resolved("", "FF", "pk-0")])]); + let mut executor3 = MockRequestExecutor::new(vec![Ok(envelope_page(session3_page, None))]); + let mut pipeline3 = build_streaming_ordered_merge(&plan, &mut topology3, &op, Some(resumed2)) + .await + .unwrap(); + let ids3 = drain_all(&mut pipeline3, &mut executor3).await; + + // The decisive proof: the second resume rebuilds the resume-filtered + // query from the scalar boundary and starts it fresh. It must NOT replay + // the resume-filtered token against the plain query. + assert!( + !executor3 + .continuation_calls + .contains(&Some(resume_filtered_ct.to_owned())), + "the resume-filtered backend continuation must never be replayed on resume; \ + continuation_calls = {:?}", + executor3.continuation_calls + ); + assert_eq!( + executor3.continuation_calls.first(), + Some(&None), + "the second resume must start the rebuilt resume-filtered query fresh, not forward a token" + ); + assert_is_resume_filtered_query(&executor3.body_text(0)); + + // End-to-end correctness: every row is delivered exactly once, in order. + let mut all_ids = ids1; + all_ids.extend(ids2); + all_ids.extend(ids3); + assert_eq!( + all_ids, + expected_order + .iter() + .map(|s| s.to_string()) + .collect::>(), + "across two snapshot/resume cycles every row must be delivered exactly once, in order" + ); +} + +/// Regression (distinct keys): a resume-filtered checkpoint with a live +/// continuation must not persist it; next resume re-derives from boundary. +#[tokio::test] +async fn resume_filtered_query_never_replays_backend_continuation_against_plain_query() { + run_resume_filtered_binding_cycle( + &[("a", 1), ("b", 2)], + &[("a", 1), ("b", 2)], + &[("b", 2), ("c", 3)], + "resume-filtered-ct-1", + &["a", "b", "c"], + ) + .await; +} + +/// Regression (tied keys): correctness rests on the `_rid` tiebreak +/// surviving the resume-filtered snapshot without mis-binding. +#[tokio::test] +async fn resume_filtered_query_with_tied_keys_never_replays_backend_continuation() { + run_resume_filtered_binding_cycle( + &[("a", 5), ("b", 5)], + &[("a", 5), ("b", 5)], + &[("b", 5), ("c", 5)], + "resume-filtered-ct-tied", + &["a", "b", "c"], + ) + .await; +} + +/// A resume whose boundary value is a string full of SQL special characters +/// (quote, backslash, newline, tab, non-ASCII) must travel in the structured +/// `resumeFilter.value`, never inlined as SQL text. Asserts the recorded +/// request body: the raw string never appears in the query SQL, no +/// `@cosmos...` parameter is synthesized, and the `resumeFilter` carries the +/// exact string verbatim as a plain value. +#[tokio::test] +async fn string_boundary_resume_travels_in_resume_filter_not_inline_sql() { + let nasty = "a' OR 1=1 -- \\ \n\t\u{2713}"; + let op = order_by_operation(); + let plan = order_by_plan(); + let resumed_state = PipelineNodeState::StreamingOrderedMerge { + directions: vec![SortOrder::Ascending], + ranges: vec![OrderByRangeToken { + min_epk: String::new(), + max_epk: "FF".to_owned(), + server_continuation: None, + boundary: Some(ValueBoundary { + resume_values: vec![OrderByResumeValue::String { + value: nasty.to_owned(), + }], + last_rid: "rid-1".to_owned(), + skip_count: 1, + }), + }], + }; + + let mut topology = MockTopologyProvider::new(vec![Ok(vec![resolved("", "FF", "pk-0")])]); + let mut executor = MockRequestExecutor::new(vec![Ok(envelope_page(&[], None))]); + let mut pipeline = + build_streaming_ordered_merge(&plan, &mut topology, &op, Some(resumed_state)) + .await + .unwrap(); + let _ = drain_all(&mut pipeline, &mut executor).await; + + let body = executor.body_text(0); + let value: serde_json::Value = + serde_json::from_str(&body).expect("recorded request body must be valid JSON"); + let query = value["query"].as_str().expect("body has a query string"); + assert!( + !body.contains("@cosmosResumeFilter"), + "no query parameter may be synthesized: {body}" + ); + // The adversarial text must never leak into the SQL — only into the filter. + for needle in ["OR 1=1", "'", "\\"] { + assert!( + !query.contains(needle), + "boundary text {needle:?} must not appear in the query SQL: {query}" + ); + } + assert_eq!( + value["resumeFilter"]["value"][0], + serde_json::Value::String(nasty.to_owned()), + "the exact boundary string round-trips as the resumeFilter value" + ); +} + +/// Asserts a recorded request body's resume filter carries `expected` as its +/// first `resumeFilter.value` entry and never inlines it into the SQL. +fn assert_string_resume_filter_boundary(body: &str, expected: &str) { + let value: serde_json::Value = + serde_json::from_str(body).expect("recorded request body must be valid JSON"); + let query = value["query"].as_str().expect("body has a query string"); + assert!( + !query.contains(expected), + "boundary text {expected:?} must not appear in the query SQL: {query}" + ); + assert_eq!( + value["resumeFilter"]["value"][0], + serde_json::Value::String(expected.to_owned()), + "the boundary string travels verbatim in the resumeFilter value" + ); +} + +/// Repeated resume with string sort keys carrying SQL special characters: +/// fresh → checkpoint → resume → checkpoint → resume. Every resumed request +/// body must bind the boundary string as a parameter (never inline it), the +/// binding must survive each token serialize/deserialize round-trip, and all +/// rows must be delivered exactly once in order. +#[tokio::test] +async fn repeated_string_boundary_resume_binds_parameter_each_cycle() { + // Keys sort a < b < c and each carries a distinct special character. + let (ka, kb, kc) = ("k1' ", "k2\\", "k3\t"); + let op = order_by_operation_with_page_size(1); + let plan = order_by_plan(); + + // ── Session 1: fresh, emit "a", checkpoint mid-page ───────────────── + let mut topology1 = MockTopologyProvider::new(vec![Ok(vec![resolved("", "FF", "pk-0")])]); + let mut executor1 = MockRequestExecutor::new(vec![Ok(string_envelope_page( + &[("a", ka), ("b", kb), ("c", kc)], + Some("s1-more"), + ))]); + let mut pipeline1 = build_streaming_ordered_merge(&plan, &mut topology1, &op, None) + .await + .unwrap(); + let ids1 = drain_one(&mut pipeline1, &mut executor1).await; + assert_eq!(ids1, vec!["a".to_owned()]); + let state1 = pipeline1.snapshot_state().unwrap(); + drop(pipeline1); + + // ── Session 2: resume from the "a" (ka) boundary ──────────────────── + let resumed1 = round_trip_state(state1, &op); + let mut topology2 = MockTopologyProvider::new(vec![Ok(vec![resolved("", "FF", "pk-0")])]); + let mut executor2 = MockRequestExecutor::new(vec![Ok(string_envelope_page( + &[("b", kb), ("c", kc)], + Some("s2-more"), + ))]); + let mut pipeline2 = build_streaming_ordered_merge(&plan, &mut topology2, &op, Some(resumed1)) + .await + .unwrap(); + let ids2 = drain_one(&mut pipeline2, &mut executor2).await; + assert_eq!(ids2, vec!["b".to_owned()]); + assert_string_resume_filter_boundary(&executor2.body_text(0), ka); + let state2 = pipeline2.snapshot_state().unwrap(); + drop(pipeline2); + + // ── Session 3: second resume from the "b" (kb) boundary ───────────── + let resumed2 = round_trip_state(state2, &op); + let mut topology3 = MockTopologyProvider::new(vec![Ok(vec![resolved("", "FF", "pk-0")])]); + let mut executor3 = + MockRequestExecutor::new(vec![Ok(string_envelope_page(&[("c", kc)], None))]); + let mut pipeline3 = build_streaming_ordered_merge(&plan, &mut topology3, &op, Some(resumed2)) + .await + .unwrap(); + let ids3 = drain_all(&mut pipeline3, &mut executor3).await; + assert_eq!(ids3, vec!["c".to_owned()]); + assert_string_resume_filter_boundary(&executor3.body_text(0), kb); + + let mut all = ids1; + all.extend(ids2); + all.extend(ids3); + assert_eq!( + all, + vec!["a", "b", "c"] + .into_iter() + .map(str::to_owned) + .collect::>(), + "across two string-boundary resume cycles every row is delivered once, in order" + ); +} + +// ── Full-key tie, direction-aware, multi-page, two-resume regression ─────── + +/// Drives one direction's full-key-tie resume regression end to end: six +/// documents, all tied on `rank`, with real (16-byte-encoded) rids in +/// `direction`'s document-ordinal order — matching how the real backend +/// returns a tied run within a partition (see `order_by::compare_rids`). +/// +/// Exercises, together: +/// - a tie run spanning more than one raw backend page fetch, both before +/// (session A) and after (session C, where an entire page is discarded +/// and the fetch loop must continue to the next one) a checkpoint; +/// - two full snapshot/resume cycles (A -> B -> C); +/// - the resume request carrying the structured `resumeFilter` (the +/// direction lives in the query's `ORDER BY` clause), with the numeric +/// `_rid` tie-break applied client-side (a base64 `_rid` string does not +/// sort in document order). +/// +/// Every one of the six rids must be delivered exactly once, in strict +/// `direction` order — proving no duplicates and no omissions. +async fn tied_full_key_resume_spans_pages_and_two_cycles(direction: SortOrder) { + const RANK: i64 = 5; + // Ordinals in `direction`'s document order, matching real backend + // behavior for a full-key tie within one partition. + let ordinals: [u64; 6] = match direction { + SortOrder::Ascending => [10, 20, 30, 40, 50, 60], + SortOrder::Descending => [60, 50, 40, 30, 20, 10], + }; + let rid = |i: usize| real_rid(ordinals[i]); + let rows: Vec<(String, i64)> = (0..6).map(|i| (rid(i), RANK)).collect(); + let refs = |range: std::ops::Range| -> Vec<(&str, i64)> { + rows[range].iter().map(|(r, v)| (r.as_str(), *v)).collect() + }; + let order_keyword = match direction { + SortOrder::Ascending => "ASC", + SortOrder::Descending => "DESC", + }; + + let op = order_by_operation_with_page_size(1); + let plan = order_by_plan_with_direction(direction); + + // ── Session A: fresh. The tie run spans two raw backend page fetches + // (rid(0),rid(1) then rid(2),rid(3)) before any checkpoint. ── + let mut topology_a = MockTopologyProvider::new(vec![Ok(vec![resolved("", "FF", "pk-0")])]); + let mut executor_a = MockRequestExecutor::new(vec![ + Ok(envelope_page(&refs(0..2), Some("bp-1"))), + Ok(envelope_page(&refs(2..4), None)), + ]); + let mut pipeline_a = build_streaming_ordered_merge(&plan, &mut topology_a, &op, None) + .await + .unwrap(); + let mut ids_a = Vec::new(); + ids_a.extend(drain_one(&mut pipeline_a, &mut executor_a).await); // rid(0): backend page 1 + ids_a.extend(drain_one(&mut pipeline_a, &mut executor_a).await); // rid(1): still buffered + ids_a.extend(drain_one(&mut pipeline_a, &mut executor_a).await); // rid(2): fetches backend page 2 + assert_eq!(ids_a, vec![rid(0), rid(1), rid(2)]); + assert_eq!( + executor_a.query_bodies.len(), + 2, + "the tie run must have spanned exactly two raw backend page fetches" + ); + + let state_a = pipeline_a.snapshot_state().unwrap(); + match &state_a { + PipelineNodeState::StreamingOrderedMerge { ranges, .. } => { + assert_eq!(ranges.len(), 1); + assert_eq!( + ranges[0].server_continuation, None, + "rid(3) is still buffered (unread) from backend page 2, so no plain \ + continuation is safe to save" + ); + } + other => panic!("expected StreamingOrderedMerge snapshot, got {other:?}"), + } + drop(pipeline_a); + + // ── Session B: first resume. The mock can't evaluate the SQL filter, so + // it replays the whole unfiltered tie run; the client-side discard + // strips rid(0)..=rid(2) (already emitted), surfacing only rid(3). ── + let resumed_a = round_trip_state(state_a, &op); + let mut topology_b = MockTopologyProvider::new(vec![Ok(vec![resolved("", "FF", "pk-0")])]); + let mut executor_b = MockRequestExecutor::new(vec![ + Ok(envelope_page(&refs(0..4), Some("bp-2"))), + Ok(envelope_page(&refs(4..6), None)), + ]); + let mut pipeline_b = + build_streaming_ordered_merge(&plan, &mut topology_b, &op, Some(resumed_a)) + .await + .unwrap(); + let ids_b = drain_one(&mut pipeline_b, &mut executor_b).await; + assert_eq!( + ids_b, + vec![rid(3)], + "already-emitted rid(0)..=rid(2) must be discarded client-side" + ); + assert_eq!( + executor_b.continuation_calls, + vec![None], + "a value-boundary resume starts the resume-filtered query fresh" + ); + assert_eq!( + executor_b.query_bodies.len(), + 1, + "rid(3) surfaced from the first backend page alone; the second must not be fetched yet" + ); + let body_b = executor_b.body_text(0); + assert_is_resume_filtered_query(&body_b); + let body_b_value: serde_json::Value = serde_json::from_str(&body_b).unwrap(); + assert!( + body_b_value["query"] + .as_str() + .unwrap() + .contains(&format!("ORDER BY c.rank {order_keyword}")), + "{body_b}: the query's ORDER BY direction must match {direction:?}" + ); + assert_eq!( + body_b_value["resumeFilter"]["value"][0], + serde_json::json!(RANK), + "the boundary value travels in the structured resumeFilter, not inline SQL" + ); + + let state_b = pipeline_b.snapshot_state().unwrap(); + match &state_b { + PipelineNodeState::StreamingOrderedMerge { ranges, .. } => { + assert_eq!( + ranges[0].server_continuation, None, + "a resume-filtered child's continuation must never be snapshotted, even \ + though its buffer is now empty and \"bp-2\" is technically still live" + ); + } + other => panic!("expected StreamingOrderedMerge snapshot, got {other:?}"), + } + drop(pipeline_b); + + // ── Session C: second resume. Backend page 1 here (rid(0)..=rid(3)) is + // *entirely* discarded (every row ties at-or-before the rid(3) + // boundary) — the discard must stay active across that whole page and + // into backend page 2, where rid(4) finally survives. ── + let resumed_b = round_trip_state(state_b, &op); + let mut topology_c = MockTopologyProvider::new(vec![Ok(vec![resolved("", "FF", "pk-0")])]); + let mut executor_c = MockRequestExecutor::new(vec![ + Ok(envelope_page(&refs(0..4), Some("bp-3"))), + Ok(envelope_page(&refs(4..6), None)), + ]); + let mut pipeline_c = + build_streaming_ordered_merge(&plan, &mut topology_c, &op, Some(resumed_b)) + .await + .unwrap(); + let ids_c = drain_all(&mut pipeline_c, &mut executor_c).await; + assert_eq!( + ids_c, + vec![rid(4), rid(5)], + "already-emitted rid(0)..=rid(3) must be discarded again on this second resume" + ); + assert_eq!( + executor_c.query_bodies.len(), + 2, + "the fully-discarded first page must not stop the fetch loop before the second" + ); + assert_eq!( + executor_c.continuation_calls.first(), + Some(&None), + "the second resume must start the rebuilt resume-filtered query fresh — never \ + forwarding session B's live \"bp-2\" continuation, which its snapshot never kept" + ); + + // End-to-end: every row delivered exactly once, in strict `direction` + // rid order — proving no duplicates and no omissions. + let mut all_ids = ids_a; + all_ids.extend(ids_b); + all_ids.extend(ids_c); + assert_eq!( + all_ids, + (0..6).map(rid).collect::>(), + "every tied row must be delivered exactly once, in strict {direction:?} rid order" + ); +} + +#[tokio::test] +async fn ascending_full_key_ties_span_backend_pages_and_survive_two_resumes() { + tied_full_key_resume_spans_pages_and_two_cycles(SortOrder::Ascending).await; +} + +#[tokio::test] +async fn descending_full_key_ties_span_backend_pages_and_survive_two_resumes() { + tied_full_key_resume_spans_pages_and_two_cycles(SortOrder::Descending).await; +} 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 029665dd9a1..b0e66053cc0 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 @@ -115,6 +115,9 @@ pub(crate) struct MockRequestExecutor { pub refresh_calls: Vec, pub continuation_calls: Vec>, pub target_calls: Vec, + /// Request body of every executed operation, in call order — lets + /// tests assert which query text was sent with which continuation. + pub query_bodies: Vec>>, } impl MockRequestExecutor { @@ -124,14 +127,26 @@ impl MockRequestExecutor { refresh_calls: Vec::new(), continuation_calls: Vec::new(), target_calls: Vec::new(), + query_bodies: Vec::new(), } } + + /// Returns the recorded request body at `index` as UTF-8 (lossy). + pub fn body_text(&self, index: usize) -> String { + let body = self + .query_bodies + .get(index) + .unwrap_or_else(|| panic!("no recorded request at index {index}")) + .as_deref() + .unwrap_or_default(); + String::from_utf8_lossy(body).into_owned() + } } impl RequestExecutor for MockRequestExecutor { fn execute_request<'a>( &'a mut self, - _operation: &'a CosmosOperation, + operation: &'a CosmosOperation, target: RequestTarget, partition_routing_refresh: PartitionRoutingRefresh, continuation: Option, @@ -139,6 +154,7 @@ impl RequestExecutor for MockRequestExecutor { self.refresh_calls.push(partition_routing_refresh); self.continuation_calls.push(continuation); self.target_calls.push(target); + self.query_bodies.push(operation.body().map(<[u8]>::to_vec)); let response = self.responses.pop_front().expect("mock request response"); Box::pin(async move { response }) } 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 f7eec4c3d6c..d860171e980 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 @@ -16,10 +16,13 @@ //! - Intermediate nodes: [`SequentialDrain`] iterates EPK-ordered children //! left-to-right, draining each before advancing. [`UnorderedMerge`] polls //! children round-robin without evicting them, suitable for change feed. +//! [`StreamingOrderedMerge`] k-way merges globally-ordered `ORDER BY` +//! results across children, each executing a Gateway-rewritten query. //! - Planner: [`planner::build_trivial_pipeline`] handles point reads and //! single-partition operations; [`planner::build_sequential_drain`] handles -//! cross-partition queries by consuming a backend query plan and resolving -//! it against the current topology. +//! natural-order cross-partition queries; [`planner::build_streaming_ordered_merge`] +//! handles cross-partition `ORDER BY` queries — all by consuming a backend +//! query plan and resolving it against the current topology. //! - Serializable state: [`PipelineNodeState`] (see [`snapshot`]) is the //! in-memory shape of a continuation snapshot; the wire-format token lives //! in [`crate::models::ContinuationToken`]. @@ -32,6 +35,7 @@ //! cross-partition strategies). mod context; +mod distinct_hash; mod drain; mod drained; #[cfg(test)] @@ -39,11 +43,14 @@ mod integration_tests; #[cfg(test)] pub(crate) mod mocks; mod node; +mod order_by; mod pipeline; pub(crate) mod planner; pub(crate) mod query_plan; +mod query_response; mod request; mod snapshot; +mod streaming_ordered_merge; mod topology; mod unordered_merge; @@ -57,6 +64,7 @@ pub use pipeline::OperationPlan; pub(crate) use pipeline::Pipeline; pub(crate) use request::{intersect_feed_ranges, Request, RequestTarget}; pub(crate) use snapshot::{PipelineNodeState, RangedToken}; +pub(crate) use streaming_ordered_merge::StreamingOrderedMerge; pub(crate) use topology::CachedTopologyProvider; pub(crate) use unordered_merge::UnorderedMerge; 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 new file mode 100644 index 00000000000..de736d7b6cf --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/order_by.rs @@ -0,0 +1,1481 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +//! Cross-partition streaming `ORDER BY` value model. +//! +//! Shared vocabulary between [`StreamingOrderedMerge`] and the +//! continuation-token snapshot (`super::snapshot`): [`OrderByItem`] parses +//! one column's envelope value (distinguishing `Undefined` from `null`); +//! [`compare_key_tuples`] is the shared Cosmos-ordered comparator; +//! [`OrderByResumeValue`] is the bounded "last emitted" value persisted in +//! a token (arrays/objects hashed to 128 bits so a token never grows with +//! document size). On resume, a range's boundary is sent to the backend as +//! a structured [`resume_filter_json`] — the .NET-compatible `resumeFilter` +//! field of the query body — and the already-emitted prefix of the boundary +//! tie run is trimmed client-side via [`classify_row_vs_boundary`]. +//! +//! # Type order +//! +//! Canonical Cosmos ascending type order: +//! `Undefined < Null < Boolean < Number < String < Array < Object`. Must +//! stay in sync with `crate::query::eval`'s `sort_type_order`. + +use std::cmp::Ordering; + +use serde::{Deserialize, Serialize}; + +use super::distinct_hash::distinct_hash; +use super::query_plan::SortOrder; + +#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +enum CosmosType { + Undefined, + Null, + Boolean, + Number, + String, + Array, + Object, +} + +/// A single `ORDER BY` key value parsed from a rewritten envelope's +/// `orderByItems[i]`. Distinguishes `Undefined` (no `item` field) from +/// every JSON value, including `null` — load-bearing, since Cosmos ranks +/// `Undefined` before `Null`. +#[derive(Debug, Clone)] +pub(crate) enum OrderByItem { + Undefined, + Null, + Boolean(bool), + Number(OrderByNumber), + String(String), + Array(Vec), + /// Key/value pairs in wire order; comparison sorts by key internally + /// so differently-ordered wire keys still compare equal. + Object(Vec<(String, OrderByItem)>), +} + +/// Total-ordering comparison for `f64` (`-0.0 == 0.0`, matching IEEE/Cosmos +/// numeric equality); falls back to `total_cmp` only for `NaN` so this +/// never panics (required for a strict-total-order comparator). +fn cmp_f64(a: f64, b: f64) -> Ordering { + a.partial_cmp(&b).unwrap_or_else(|| a.total_cmp(&b)) +} + +/// Compares `i64` with `u64` exactly, without a lossy cast: a negative +/// signed value is always smaller than any unsigned value, otherwise the +/// non-negative signed value widens to `u64` losslessly. +fn cmp_i64_u64(a: i64, b: u64) -> Ordering { + if a < 0 { + Ordering::Less + } else { + (a as u64).cmp(&b) + } +} + +/// Compares an `i64` with an `f64` for a *total* order. A finite `f64` is +/// compared exactly (see below); a `NaN` is ordered by `f64::total_cmp` +/// convention so the comparator stays a strict total order even if a test or +/// internal construction slips a `NaN` in — negative `NaN` sorts before every +/// finite value, positive `NaN` after (production deserialization rejects +/// non-finite values, so this never arises from real data). `2^63` is exactly +/// representable, so it bounds the range where `floor` fits an `i64`; the +/// fractional part then breaks a floor-tie without any lossy cast. +fn cmp_i64_f64(a: i64, b: f64) -> Ordering { + const TWO_POW_63: f64 = 9_223_372_036_854_775_808.0; + if b.is_nan() { + // `a` is finite: below +NaN, above -NaN (matches `f64::total_cmp`). + return if b.is_sign_negative() { + Ordering::Greater + } else { + Ordering::Less + }; + } + if b >= TWO_POW_63 { + return Ordering::Less; // a <= i64::MAX < 2^63 <= b + } + if b < -TWO_POW_63 { + return Ordering::Greater; // a >= i64::MIN = -2^63 > b + } + let floor = b.floor(); // -2^63 <= floor < 2^63, fits i64 exactly + match a.cmp(&(floor as i64)) { + // a == floor(b) but b has a fractional part, so a < b. + Ordering::Equal if b > floor => Ordering::Less, + other => other, + } +} + +/// Compares a `u64` with an `f64` for a *total* order (see [`cmp_i64_f64`] for +/// the `NaN` convention; `2^64` is exactly representable and bounds the `u64` +/// range). +fn cmp_u64_f64(a: u64, b: f64) -> Ordering { + const TWO_POW_64: f64 = 18_446_744_073_709_551_616.0; + if b.is_nan() { + // `a` is finite: below +NaN, above -NaN (matches `f64::total_cmp`). + return if b.is_sign_negative() { + Ordering::Greater + } else { + Ordering::Less + }; + } + if b < 0.0 { + return Ordering::Greater; // a >= 0 > b + } + if b >= TWO_POW_64 { + return Ordering::Less; // a <= u64::MAX < 2^64 <= b + } + let floor = b.floor(); // 0 <= floor < 2^64, fits u64 exactly + match a.cmp(&(floor as u64)) { + Ordering::Equal if b > floor => Ordering::Less, + other => other, + } +} + +/// A lossless `ORDER BY` numeric value. Cosmos/JSON numbers may be signed +/// integers, unsigned integers, or finite floats, and no single `f64` +/// represents every `i64`/`u64` exactly (e.g. `2^53 + 1`). Preserving the +/// original variant keeps comparison, resume-token round-trips, and SQL +/// literals mathematically exact. +/// +/// Serialization is lossless: each variant emits its native JSON number +/// and re-parses back to the same variant. +#[derive(Debug, Clone, Copy)] +pub(crate) enum OrderByNumber { + I64(i64), + U64(u64), + F64(f64), +} + +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 { + if let Some(i) = n.as_i64() { + Self::I64(i) + } else if let Some(u) = n.as_u64() { + Self::U64(u) + } else { + Self::F64(n.as_f64().unwrap_or(0.0)) + } + } + + /// Renders the exact JSON value (integers stay integers) for the wire + /// form of a scalar resume boundary (its `resumeFilter.value` entry). + fn to_json_value(self) -> serde_json::Value { + match self { + Self::I64(i) => serde_json::Value::Number(i.into()), + Self::U64(u) => serde_json::Value::Number(u.into()), + Self::F64(f) => serde_json::Number::from_f64(f) + .map(serde_json::Value::Number) + .unwrap_or(serde_json::Value::Null), + } + } + + #[cfg(test)] + fn cosmos_cmp(&self, other: &Self) -> Ordering { + self.cmp(other) + } +} + +/// Numeric equality is value-based across variants (`I64(5) == F64(5.0)`), +/// matching Cosmos, so representation never changes an equality result. +impl PartialEq for OrderByNumber { + fn eq(&self, other: &Self) -> bool { + self.cmp(other) == Ordering::Equal + } +} + +// `Eq` soundness: `cmp` is a strict total order for *every* `OrderByNumber`, +// including `NaN` (the cross-variant helpers order `NaN` by `f64::total_cmp` +// convention, and `F64`/`F64` already does), so the `cmp`-derived `PartialEq` +// is a proper equivalence relation — reflexive, symmetric, and transitive — +// even if a test or internal construction introduces a `NaN`. Production never +// does: the only production constructor is `from_json_number` (a +// `serde_json::Number` is always finite), `deserialize` rejects non-finite +// floats below, and the `From` convenience is `#[cfg(test)]`-only. +impl Eq for OrderByNumber {} + +impl PartialOrd for OrderByNumber { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for OrderByNumber { + fn cmp(&self, other: &Self) -> Ordering { + use OrderByNumber::{F64, I64, U64}; + match (self, other) { + (I64(a), I64(b)) => a.cmp(b), + (U64(a), U64(b)) => a.cmp(b), + (F64(a), F64(b)) => cmp_f64(*a, *b), + (I64(a), U64(b)) => cmp_i64_u64(*a, *b), + (U64(a), I64(b)) => cmp_i64_u64(*b, *a).reverse(), + (I64(a), F64(b)) => cmp_i64_f64(*a, *b), + (F64(a), I64(b)) => cmp_i64_f64(*b, *a).reverse(), + (U64(a), F64(b)) => cmp_u64_f64(*a, *b), + (F64(a), U64(b)) => cmp_u64_f64(*b, *a).reverse(), + } + } +} + +impl From for OrderByNumber { + fn from(value: i64) -> Self { + Self::I64(value) + } +} + +impl From for OrderByNumber { + fn from(value: u64) -> Self { + Self::U64(value) + } +} + +// Test-only: production never builds an `OrderByNumber` from an arbitrary +// `f64` (that path is `from_json_number`, which is always finite), so +// keeping this out of production builds guarantees a non-finite value can +// never enter and undermine the `Eq` impl above. +#[cfg(test)] +impl From for OrderByNumber { + fn from(value: f64) -> Self { + Self::F64(value) + } +} + +impl serde::Serialize for OrderByNumber { + fn serialize(&self, serializer: S) -> Result { + match self { + Self::I64(i) => serializer.serialize_i64(*i), + Self::U64(u) => serializer.serialize_u64(*u), + Self::F64(f) => serializer.serialize_f64(*f), + } + } +} + +impl<'de> serde::Deserialize<'de> for OrderByNumber { + fn deserialize>(deserializer: D) -> Result { + let number = serde_json::Number::deserialize(deserializer)?; + let value = Self::from_json_number(&number); + // A persisted resume value is always a finite JSON number. Reject + // non-finite floats so the `Eq` equivalence stays airtight — they + // can't arise from valid JSON but would otherwise be a silent hole. + if let Self::F64(f) = value { + if !f.is_finite() { + return Err(serde::de::Error::custom( + "ORDER BY numeric resume value must be finite", + )); + } + } + Ok(value) + } +} + +impl OrderByItem { + fn cosmos_type(&self) -> CosmosType { + match self { + Self::Undefined => CosmosType::Undefined, + Self::Null => CosmosType::Null, + Self::Boolean(_) => CosmosType::Boolean, + Self::Number(_) => CosmosType::Number, + Self::String(_) => CosmosType::String, + Self::Array(_) => CosmosType::Array, + Self::Object(_) => CosmosType::Object, + } + } + + /// Converts a JSON value into an `OrderByItem`. Never produces + /// `Undefined` — only the wire-envelope parser does, for a missing + /// `item` key. + pub(crate) fn from_json(value: &serde_json::Value) -> Self { + match value { + serde_json::Value::Null => Self::Null, + serde_json::Value::Bool(b) => Self::Boolean(*b), + serde_json::Value::Number(n) => Self::Number(OrderByNumber::from_json_number(n)), + serde_json::Value::String(s) => Self::String(s.clone()), + serde_json::Value::Array(items) => { + Self::Array(items.iter().map(Self::from_json).collect()) + } + serde_json::Value::Object(map) => Self::Object( + map.iter() + .map(|(k, v)| (k.clone(), Self::from_json(v))) + .collect(), + ), + } + } +} + +impl PartialEq for OrderByItem { + fn eq(&self, other: &Self) -> bool { + self.cmp(other) == Ordering::Equal + } +} + +impl Eq for OrderByItem {} + +impl PartialOrd for OrderByItem { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for OrderByItem { + fn cmp(&self, other: &Self) -> Ordering { + let rank_cmp = self.cosmos_type().cmp(&other.cosmos_type()); + if rank_cmp != Ordering::Equal { + return rank_cmp; + } + match (self, other) { + (Self::Undefined, Self::Undefined) => Ordering::Equal, + (Self::Null, Self::Null) => Ordering::Equal, + (Self::Boolean(a), Self::Boolean(b)) => a.cmp(b), + (Self::Number(a), Self::Number(b)) => a.cmp(b), + (Self::String(a), Self::String(b)) => a.cmp(b), + (Self::Array(a), Self::Array(b)) => { + for (x, y) in a.iter().zip(b.iter()) { + let c = x.cmp(y); + if c != Ordering::Equal { + return c; + } + } + a.len().cmp(&b.len()) + } + (Self::Object(a), Self::Object(b)) => { + let mut a_sorted: Vec<&(String, OrderByItem)> = a.iter().collect(); + let mut b_sorted: Vec<&(String, OrderByItem)> = b.iter().collect(); + a_sorted.sort_by(|x, y| x.0.cmp(&y.0)); + b_sorted.sort_by(|x, y| x.0.cmp(&y.0)); + for (x, y) in a_sorted.iter().zip(b_sorted.iter()) { + let key_cmp = x.0.cmp(&y.0); + if key_cmp != Ordering::Equal { + return key_cmp; + } + let val_cmp = x.1.cmp(&y.1); + if val_cmp != Ordering::Equal { + return val_cmp; + } + } + a_sorted.len().cmp(&b_sorted.len()) + } + _ => unreachable!("rank_cmp already distinguished differing variants"), + } + } +} + +impl OrderByItem { + /// Converts this item to its bounded, serializable resume-value form. + pub(crate) fn to_resume_value(&self) -> OrderByResumeValue { + match self { + Self::Undefined => OrderByResumeValue::Undefined, + Self::Null => OrderByResumeValue::Null, + Self::Boolean(b) => OrderByResumeValue::Boolean { value: *b }, + Self::Number(n) => OrderByResumeValue::Number { value: *n }, + Self::String(s) => OrderByResumeValue::String { value: s.clone() }, + Self::Array(_) => OrderByResumeValue::Complex { + complex_type: ComplexTypeTag::Array, + hash: ComplexHash::of(self), + }, + Self::Object(_) => OrderByResumeValue::Complex { + complex_type: ComplexTypeTag::Object, + hash: ComplexHash::of(self), + }, + } + } + + #[cfg(test)] + fn cosmos_cmp(&self, other: &Self) -> Ordering { + self.cmp(other) + } +} + +/// Parses a rewritten envelope's `orderByItems` array into one +/// [`OrderByItem`] per column, validating the wire shape: must be a JSON +/// array of length `expected_len`, each element a JSON object; a missing +/// `item` key parses as [`OrderByItem::Undefined`]. +/// +/// Returns a typed [`crate::error::CosmosError`] (not a panic) for any +/// shape violation. +pub(crate) fn parse_order_by_items( + value: &serde_json::Value, + expected_len: usize, +) -> crate::error::Result> { + let elements = value.as_array().ok_or_else(|| { + envelope_error(format!( + "rewritten envelope `orderByItems` must be a JSON array, found {}", + json_type_name(value) + )) + })?; + if elements.len() != expected_len { + return Err(envelope_error(format!( + "rewritten envelope `orderByItems` has {} entries but the query defines {expected_len} ORDER BY column(s)", + elements.len() + ))); + } + elements + .iter() + .map(|element| { + let obj = element.as_object().ok_or_else(|| { + envelope_error(format!( + "rewritten envelope `orderByItems` entry must be a JSON object, found {}", + json_type_name(element) + )) + })?; + Ok(match obj.get("item") { + Some(item) => OrderByItem::from_json(item), + None => OrderByItem::Undefined, + }) + }) + .collect() +} + +fn json_type_name(v: &serde_json::Value) -> &'static str { + match v { + serde_json::Value::Null => "null", + serde_json::Value::Bool(_) => "bool", + serde_json::Value::Number(_) => "number", + serde_json::Value::String(_) => "string", + serde_json::Value::Array(_) => "array", + serde_json::Value::Object(_) => "object", + } +} + +fn envelope_error(message: impl Into>) -> crate::error::CosmosError { + crate::error::CosmosError::builder() + .with_status(crate::error::CosmosStatus::SERVICE_ORDER_BY_ENVELOPE_INVALID) + .with_message(message) + .build() +} + +/// Compares two same-length key tuples column-by-column, applying each +/// column's direction, stopping at the first non-equal column. Panics on +/// a length mismatch — callers must validate column-count agreement first. +pub(crate) fn compare_key_tuples( + a: &[OrderByItem], + b: &[OrderByItem], + directions: &[SortOrder], +) -> Ordering { + debug_assert_eq!(a.len(), directions.len()); + debug_assert_eq!(b.len(), directions.len()); + for (i, direction) in directions.iter().enumerate() { + let c = a[i].cmp(&b[i]); + let c = match direction { + SortOrder::Ascending => c, + SortOrder::Descending => c.reverse(), + }; + if c != Ordering::Equal { + return c; + } + } + Ordering::Equal +} + +/// Compares two document `_rid`s in `direction`'s Cosmos document order — +/// the order the backend breaks full-key `ORDER BY` ties in within a +/// partition. Ascending uses numeric document-ordinal order (see +/// [`crate::models::resource_id::compare_document_rids`]); descending +/// reverses it. The direction is the query's first sort column, matching +/// .NET's `ReverseRidEnabled` fallback (no `reverseIndexScan` signal is +/// available in the Gateway query plan the driver consumes). +pub(crate) fn compare_rids(a: &str, b: &str, direction: SortOrder) -> Ordering { + let ascending = crate::models::resource_id::compare_document_rids(a, b); + match direction { + SortOrder::Ascending => ascending, + SortOrder::Descending => ascending.reverse(), + } +} + +/// Discriminates which complex JSON shape a hashed resume value came from, +/// so two colliding hashes from different shapes are never a tie. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub(crate) enum ComplexTypeTag { + Array, + Object, +} + +/// A bounded 128-bit hash of a complex (array/object) `ORDER BY` value, +/// split into low/high halves for JSON round-tripping. +/// +/// This is the backend / .NET SDK structural `DistinctHash` (see +/// [`super::distinct_hash`]), so it is byte-identical to the hash the +/// backend derives for the same array/object value. That is what makes a +/// structured `resumeFilter` seek correctly from a complex boundary, +/// including across a partition split or merge. Structurally-equal values +/// hash equal regardless of object property order. +/// +/// Its *ordering* of distinct values is still not Cosmos sort order — +/// MurmurHash output is not monotonic — so the client-side discard treats +/// only an exact-hash tie as already-emitted and never infers less/greater +/// between two distinct complex hashes (see [`classify_row_vs_boundary`]). +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct ComplexHash { + pub(crate) low64: u64, + pub(crate) high64: u64, +} + +impl ComplexHash { + fn of(item: &OrderByItem) -> Self { + let hash = distinct_hash(item); + Self { + low64: hash as u64, + high64: (hash >> 64) as u64, + } + } +} + +/// The bounded, serializable representation of one column's "last emitted" +/// value, persisted in a continuation token. Scalars round-trip exactly; +/// arrays/objects are represented only by their [`ComplexHash`], so a +/// token never grows with document size. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub(crate) enum OrderByResumeValue { + Undefined, + Null, + Boolean { + value: bool, + }, + Number { + value: OrderByNumber, + }, + String { + value: String, + }, + Complex { + complex_type: ComplexTypeTag, + hash: ComplexHash, + }, +} + +impl OrderByResumeValue { + /// Whether this boundary column is a complex (array/object) value, + /// represented only by its bounded [`ComplexHash`]. + #[cfg(test)] + pub(crate) fn is_complex(&self) -> bool { + matches!(self, Self::Complex { .. }) + } + + /// Reconstructs the [`OrderByItem`] this came from, or `None` for + /// [`Self::Complex`] (bytes aren't recoverable from the hash). + #[cfg(test)] + pub(crate) fn to_scalar_order_by_item(&self) -> Option { + match self { + Self::Undefined => Some(OrderByItem::Undefined), + Self::Null => Some(OrderByItem::Null), + Self::Boolean { value } => Some(OrderByItem::Boolean(*value)), + Self::Number { value } => Some(OrderByItem::Number(*value)), + Self::String { value } => Some(OrderByItem::String(value.clone())), + Self::Complex { .. } => None, + } + } + + /// The .NET-compatible wire form of this resume value for a query + /// body's `resumeFilter.value` array (distinct from the token/snapshot + /// serde form): `Undefined` -> `[]`; `Null`/`Boolean`/`Number`/`String` + /// -> the raw JSON value (integers keep exact i64/u64 precision); a + /// complex (array/object) value -> `{"type":..,"low":,"high":}` + /// where `low`/`high` reinterpret the 64-bit hash halves' bit patterns + /// as signed integers (matching .NET's `(long)ulong` cast). + fn to_wire_value(&self) -> serde_json::Value { + match self { + Self::Undefined => serde_json::Value::Array(Vec::new()), + Self::Null => serde_json::Value::Null, + Self::Boolean { value } => serde_json::Value::Bool(*value), + Self::Number { value } => value.to_json_value(), + Self::String { value } => serde_json::Value::String(value.clone()), + Self::Complex { complex_type, hash } => { + let type_name = match complex_type { + ComplexTypeTag::Array => "array", + ComplexTypeTag::Object => "object", + }; + serde_json::json!({ + "type": type_name, + "low": hash.low64 as i64, + "high": hash.high64 as i64, + }) + } + } + } + + fn cosmos_type(&self) -> CosmosType { + match self { + Self::Undefined => CosmosType::Undefined, + Self::Null => CosmosType::Null, + Self::Boolean { .. } => CosmosType::Boolean, + Self::Number { .. } => CosmosType::Number, + Self::String { .. } => CosmosType::String, + Self::Complex { complex_type, .. } => match complex_type { + ComplexTypeTag::Array => CosmosType::Array, + ComplexTypeTag::Object => CosmosType::Object, + }, + } + } +} + +/// Builds the .NET-compatible `resumeFilter` object inserted into a +/// resumed query body, mirroring +/// `Microsoft.Azure.Cosmos.Query.Core.SqlQueryResumeFilter`: +/// `{"value":[],"rid":,"exclude":}`. +/// +/// `rid` is omitted when `None`. Rust persists a last-emitted boundary per +/// range and always resumes a range with its own boundary as the .NET +/// "target" partition does — `rid` present and `exclude:false` — then trims +/// the already-emitted prefix of the boundary tie run client-side (see +/// [`classify_row_vs_boundary`]). +pub(crate) fn resume_filter_json( + resume_values: &[OrderByResumeValue], + rid: Option<&str>, + exclude: bool, +) -> serde_json::Value { + let mut obj = serde_json::Map::new(); + obj.insert( + "value".to_owned(), + serde_json::Value::Array( + resume_values + .iter() + .map(OrderByResumeValue::to_wire_value) + .collect(), + ), + ); + if let Some(rid) = rid { + obj.insert("rid".to_owned(), serde_json::Value::String(rid.to_owned())); + } + obj.insert("exclude".to_owned(), serde_json::Value::Bool(exclude)); + serde_json::Value::Object(obj) +} + +/// One column's comparison of a returned row's value against a persisted +/// resume value. `Ordered` is a reliable Cosmos-sort-order comparison +/// (scalars by value, cross-type by Cosmos type rank); `Indeterminate` +/// means both sides are the same complex (array/object) type but their +/// bounded [`ComplexHash`]es differ. That hash is now the backend's exact +/// structural `DistinctHash`, yet its *ordering* is still not Cosmos sort +/// order (MurmurHash output is not monotonic), so a differing hash can't +/// place the row relative to the boundary — only an exact-hash tie is +/// meaningful. +enum ColumnCmp { + Ordered(Ordering), + Indeterminate, +} + +/// Ascending Cosmos comparison of a returned row's `item` against a +/// persisted resume-boundary `value` for one column. Cross-type pairs order +/// by Cosmos type rank; same-type scalars by value; same-type complex +/// values compare **only** for hash-equality (a differing hash is +/// [`ColumnCmp::Indeterminate`], never a spurious `Less`/`Greater`). +fn column_cmp(item: &OrderByItem, value: &OrderByResumeValue) -> ColumnCmp { + let rank_cmp = item.cosmos_type().cmp(&value.cosmos_type()); + if rank_cmp != Ordering::Equal { + return ColumnCmp::Ordered(rank_cmp); + } + match (item, value) { + (OrderByItem::Undefined, OrderByResumeValue::Undefined) => { + ColumnCmp::Ordered(Ordering::Equal) + } + (OrderByItem::Null, OrderByResumeValue::Null) => ColumnCmp::Ordered(Ordering::Equal), + (OrderByItem::Boolean(a), OrderByResumeValue::Boolean { value }) => { + ColumnCmp::Ordered(a.cmp(value)) + } + (OrderByItem::Number(a), OrderByResumeValue::Number { value }) => { + ColumnCmp::Ordered(a.cmp(value)) + } + (OrderByItem::String(a), OrderByResumeValue::String { value }) => { + ColumnCmp::Ordered(a.as_str().cmp(value.as_str())) + } + ( + OrderByItem::Array(_) | OrderByItem::Object(_), + OrderByResumeValue::Complex { hash, .. }, + ) => { + if ComplexHash::of(item) == *hash { + ColumnCmp::Ordered(Ordering::Equal) + } else { + ColumnCmp::Indeterminate + } + } + _ => unreachable!("type rank already distinguished differing variants"), + } +} + +/// Where a returned row's key tuple sits relative to a persisted resume +/// boundary, for the client-side discard. +pub(crate) enum RowVsBoundary { + /// The row sorts reliably strictly before the boundary — already emitted. + Before, + /// Exact full-key tie; the caller breaks it by `_rid`. + Tie, + /// The row sorts at or after the boundary, or its position cannot be + /// determined (a differing complex column, whose hash order is not sort + /// order). Never treated as already-emitted, so the discard can never + /// drop an un-emitted row. + AfterOrIndeterminate, +} + +/// Classifies a returned row's key tuple against a persisted resume boundary +/// (bounded [`OrderByResumeValue`]s), applying each column's direction and +/// stopping at the first non-equal column — the mixed-scalar/complex +/// counterpart of [`compare_key_tuples`], used by the client-side discard. +/// +/// A row is reported [`RowVsBoundary::Before`] (droppable) only when a +/// scalar (or cross-type) prefix column proves it; a differing complex +/// column yields [`RowVsBoundary::AfterOrIndeterminate`] so the discard +/// keeps it rather than risk dropping an un-emitted row. Panics on a length +/// mismatch; callers validate column-count agreement first. +pub(crate) fn classify_row_vs_boundary( + keys: &[OrderByItem], + resume_values: &[OrderByResumeValue], + directions: &[SortOrder], +) -> RowVsBoundary { + debug_assert_eq!(keys.len(), directions.len()); + debug_assert_eq!(resume_values.len(), directions.len()); + for (i, direction) in directions.iter().enumerate() { + match column_cmp(&keys[i], &resume_values[i]) { + ColumnCmp::Ordered(ord) => { + let ord = match direction { + SortOrder::Ascending => ord, + SortOrder::Descending => ord.reverse(), + }; + match ord { + Ordering::Less => return RowVsBoundary::Before, + Ordering::Greater => return RowVsBoundary::AfterOrIndeterminate, + Ordering::Equal => {} + } + } + // A complex column that differs from the boundary by hash: its + // true sort position is unknown, so never infer "before". + ColumnCmp::Indeterminate => return RowVsBoundary::AfterOrIndeterminate, + } + } + RowVsBoundary::Tie +} + +#[cfg(test)] +mod tests { + use super::*; + + // ── Type order ─────────────────────────────────────────────────────── + + #[test] + fn ascending_type_order_matches_cosmos_canonical_order() { + let items = [ + OrderByItem::Undefined, + OrderByItem::Null, + OrderByItem::Boolean(false), + OrderByItem::Number(0.0.into()), + OrderByItem::String(String::new()), + OrderByItem::Array(vec![]), + OrderByItem::Object(vec![]), + ]; + for i in 0..items.len() { + for j in 0..items.len() { + let expected = i.cmp(&j); + assert_eq!(items[i].cosmos_cmp(&items[j]), expected, "index {i} vs {j}"); + } + } + } + + #[test] + fn undefined_sorts_before_null() { + assert_eq!( + OrderByItem::Undefined.cosmos_cmp(&OrderByItem::Null), + Ordering::Less + ); + assert_eq!( + OrderByItem::Null.cosmos_cmp(&OrderByItem::Undefined), + Ordering::Greater + ); + } + + #[test] + fn numbers_compare_numerically() { + assert_eq!( + OrderByItem::Number(1.0.into()).cosmos_cmp(&OrderByItem::Number(2.0.into())), + Ordering::Less + ); + assert_eq!( + OrderByItem::Number((-0.0).into()).cosmos_cmp(&OrderByItem::Number(0.0.into())), + Ordering::Equal + ); + } + + #[test] + fn strings_compare_lexicographically() { + assert_eq!( + OrderByItem::String("a".into()).cosmos_cmp(&OrderByItem::String("b".into())), + Ordering::Less + ); + } + + #[test] + fn arrays_compare_lexicographically_with_prefix_shorter_first() { + let a = OrderByItem::Array(vec![OrderByItem::Number(1.0.into())]); + let b = OrderByItem::Array(vec![ + OrderByItem::Number(1.0.into()), + OrderByItem::Number(2.0.into()), + ]); + assert_eq!(a.cosmos_cmp(&b), Ordering::Less); + + let c = OrderByItem::Array(vec![ + OrderByItem::Number(1.0.into()), + OrderByItem::Number(1.0.into()), + ]); + let d = OrderByItem::Array(vec![ + OrderByItem::Number(1.0.into()), + OrderByItem::Number(2.0.into()), + ]); + assert_eq!(c.cosmos_cmp(&d), Ordering::Less); + } + + #[test] + fn objects_compare_by_sorted_key_regardless_of_wire_order() { + let a = OrderByItem::Object(vec![ + ("b".to_owned(), OrderByItem::Number(1.0.into())), + ("a".to_owned(), OrderByItem::Number(2.0.into())), + ]); + let b = OrderByItem::Object(vec![ + ("a".to_owned(), OrderByItem::Number(2.0.into())), + ("b".to_owned(), OrderByItem::Number(1.0.into())), + ]); + assert_eq!( + a.cosmos_cmp(&b), + Ordering::Equal, + "same content, different wire key order" + ); + + let c = OrderByItem::Object(vec![("a".to_owned(), OrderByItem::Number(3.0.into()))]); + assert_eq!( + a.cosmos_cmp(&c), + Ordering::Less, + "sorted-key comparison first differs on key 'a': 2 < 3" + ); + } + + // ── OrderByNumber: lossless cross-variant comparison ──────────────── + + #[test] + fn adjacent_two_pow_53_integers_remain_distinct_and_ordered() { + // 2^53 is the largest integer exactly representable in `f64`; + // 2^53 + 1 is not (it would round to 2^53 or 2^53 + 2). A + // comparator that ever cast through `f64` would collapse these two + // values; `OrderByNumber` must keep them distinct and ordered. + let a = OrderByNumber::from(9_007_199_254_740_992_i64); // 2^53 + let b = OrderByNumber::from(9_007_199_254_740_993_i64); // 2^53 + 1 + assert_ne!(a, b); + assert_eq!(a.cosmos_cmp(&b), Ordering::Less); + assert_eq!(b.cosmos_cmp(&a), Ordering::Greater); + + let ua = OrderByNumber::from(9_007_199_254_740_992_u64); + let ub = OrderByNumber::from(9_007_199_254_740_993_u64); + assert_ne!(ua, ub); + assert_eq!(ua.cosmos_cmp(&ub), Ordering::Less); + } + + #[test] + fn i64_and_u64_extremes_compare_correctly() { + let i_min = OrderByNumber::from(i64::MIN); + let i_max = OrderByNumber::from(i64::MAX); + let u_max = OrderByNumber::from(u64::MAX); + let u_zero = OrderByNumber::from(0_u64); + + assert_eq!( + i_min.cosmos_cmp(&u_zero), + Ordering::Less, + "any negative i64 is less than any u64" + ); + assert_eq!(i_min.cosmos_cmp(&i_max), Ordering::Less); + assert_eq!( + i_max.cosmos_cmp(&u_max), + Ordering::Less, + "i64::MAX < u64::MAX" + ); + assert_eq!(u_max.cosmos_cmp(&i_max), Ordering::Greater); + assert_eq!(u_max.cosmos_cmp(&u_max), Ordering::Equal); + assert_eq!(i_min.cosmos_cmp(&i_min), Ordering::Equal); + // Same numeric value across variants widens losslessly to equal. + assert_eq!( + OrderByNumber::from(i64::MAX).cosmos_cmp(&OrderByNumber::from(i64::MAX as u64)), + Ordering::Equal + ); + } + + #[test] + fn cross_int_float_equality_and_order() { + assert_eq!(OrderByNumber::from(5_i64), OrderByNumber::from(5.0_f64)); + assert_eq!(OrderByNumber::from(5_u64), OrderByNumber::from(5.0_f64)); + assert_eq!( + OrderByNumber::from(5_i64).cosmos_cmp(&OrderByNumber::from(5.5_f64)), + Ordering::Less + ); + assert_eq!( + OrderByNumber::from(5_i64).cosmos_cmp(&OrderByNumber::from(4.5_f64)), + Ordering::Greater + ); + + // `2^63`/`2^64` are the exact boundaries `cmp_i64_f64`/`cmp_u64_f64` + // branch on; `1e300` is far enough past either to be unambiguous + // (unlike `-2^63 - 1.0`, which rounds right back to `-2^63` at that + // magnitude's `f64` precision). + const TWO_POW_63: f64 = 9_223_372_036_854_775_808.0; + const TWO_POW_64: f64 = 18_446_744_073_709_551_616.0; + assert_eq!( + OrderByNumber::from(i64::MAX).cosmos_cmp(&OrderByNumber::from(TWO_POW_63)), + Ordering::Less + ); + assert_eq!( + OrderByNumber::from(i64::MIN).cosmos_cmp(&OrderByNumber::from(-TWO_POW_63)), + Ordering::Equal + ); + assert_eq!( + OrderByNumber::from(i64::MIN).cosmos_cmp(&OrderByNumber::from(-1e300_f64)), + Ordering::Greater + ); + assert_eq!( + OrderByNumber::from(u64::MAX).cosmos_cmp(&OrderByNumber::from(TWO_POW_64)), + Ordering::Less + ); + assert_eq!( + OrderByNumber::from(u64::MAX).cosmos_cmp(&OrderByNumber::from(1e300_f64)), + Ordering::Less + ); + assert_eq!( + OrderByNumber::from(0_u64).cosmos_cmp(&OrderByNumber::from(-1.0_f64)), + Ordering::Greater + ); + } + + #[test] + fn nan_orders_totally_against_finite_ints_and_floats() { + // `cmp` must stay a strict total order even if a `NaN` slips in via a + // test/internal construction. `f64::total_cmp` convention: `-NaN` + // sorts before every finite value, `+NaN` after. (Production rejects + // non-finite — see `deserialize_rejects_non_finite_numbers`.) + let neg_nan = OrderByNumber::from(f64::NAN.copysign(-1.0)); + let pos_nan = OrderByNumber::from(f64::NAN); + for finite in [ + OrderByNumber::from(i64::MIN), + OrderByNumber::from(-1_i64), + OrderByNumber::from(0_u64), + OrderByNumber::from(3.5_f64), + OrderByNumber::from(u64::MAX), + OrderByNumber::from(f64::INFINITY), + ] { + assert_eq!( + neg_nan.cosmos_cmp(&finite), + Ordering::Less, + "-NaN < {finite:?}" + ); + assert_eq!( + finite.cosmos_cmp(&neg_nan), + Ordering::Greater, + "{finite:?} > -NaN" + ); + assert_eq!( + pos_nan.cosmos_cmp(&finite), + Ordering::Greater, + "+NaN > {finite:?}" + ); + assert_eq!( + finite.cosmos_cmp(&pos_nan), + Ordering::Less, + "{finite:?} < +NaN" + ); + } + assert_eq!(neg_nan.cosmos_cmp(&pos_nan), Ordering::Less); + assert_eq!(pos_nan.cosmos_cmp(&neg_nan), Ordering::Greater); + // Reflexive on identical bit patterns. + assert_eq!(pos_nan.cosmos_cmp(&pos_nan), Ordering::Equal); + assert_eq!(neg_nan.cosmos_cmp(&neg_nan), Ordering::Equal); + } + + #[test] + fn cmp_is_transitive_and_antisymmetric_with_nan_and_finite_values() { + // A single intransitive triple would make sorting / `BTreeMap` + // undefined. Exhaustively verify antisymmetry and `a <= b <= c => + // a <= c` over a set spanning -NaN, finite ints/floats across every + // variant, and +NaN. + let values = [ + OrderByNumber::from(f64::NAN.copysign(-1.0)), + OrderByNumber::from(i64::MIN), + OrderByNumber::from(-2.5_f64), + OrderByNumber::from(-1_i64), + OrderByNumber::from(0_u64), + OrderByNumber::from(0.0_f64), + OrderByNumber::from(1_i64), + OrderByNumber::from(9_007_199_254_740_993_i64), // 2^53 + 1 + OrderByNumber::from(u64::MAX), + OrderByNumber::from(f64::INFINITY), + OrderByNumber::from(f64::NAN), + ]; + for a in &values { + for b in &values { + assert_eq!( + a.cosmos_cmp(b), + b.cosmos_cmp(a).reverse(), + "antisymmetry: {a:?} vs {b:?}" + ); + for c in &values { + if a.cosmos_cmp(b) != Ordering::Greater && b.cosmos_cmp(c) != Ordering::Greater + { + assert_ne!( + a.cosmos_cmp(c), + Ordering::Greater, + "transitivity: {a:?} <= {b:?} <= {c:?} but a > c" + ); + } + } + } + } + } + + #[test] + fn scalar_number_wire_value_preserves_exact_integer_precision() { + // The resume boundary sends the number as the raw JSON `resumeFilter` + // value; integers must keep full precision (never degrade through f64). + let wire = |n: OrderByNumber| OrderByResumeValue::Number { value: n }.to_wire_value(); + assert_eq!( + wire(OrderByNumber::from(i64::MIN)), + serde_json::json!(i64::MIN) + ); + assert_eq!( + wire(OrderByNumber::from(u64::MAX)), + serde_json::json!(u64::MAX) + ); + assert_eq!( + serde_json::to_string(&wire(OrderByNumber::from(9_007_199_254_740_993_i64))).unwrap(), + "9007199254740993", + "must not degrade to float notation and lose precision" + ); + assert_eq!( + wire(OrderByNumber::from(5.5_f64)), + serde_json::json!(5.5_f64) + ); + } + + #[test] + fn deserialize_rejects_non_finite_numbers() { + // JSON has no NaN/Inf, so a non-finite float can never round-trip + // into an `OrderByNumber` — the `Eq` soundness contract. Overflowing + // exponents that a parser might coerce to infinity are rejected too. + for json in [ + r#"{"type":"number","value":1e400}"#, + r#"{"type":"number","value":-1e400}"#, + ] { + assert!( + serde_json::from_str::(json).is_err(), + "non-finite resume value {json} must be rejected on deserialize" + ); + } + // Finite values still round-trip, and `Eq` holds on the result. + let value: OrderByResumeValue = + serde_json::from_str(r#"{"type":"number","value":5.5}"#).unwrap(); + assert_eq!(value, OrderByResumeValue::Number { value: 5.5.into() }); + } + + #[test] + fn large_integer_resume_values_round_trip_through_json_exactly() { + for number in [ + OrderByNumber::from(i64::MIN), + OrderByNumber::from(i64::MAX), + OrderByNumber::from(u64::MAX), + OrderByNumber::from(9_007_199_254_740_993_i64), // 2^53 + 1 + ] { + let item = OrderByItem::Number(number); + let resume = item.to_resume_value(); + let json = serde_json::to_string(&resume).unwrap(); + let back: OrderByResumeValue = serde_json::from_str(&json).unwrap(); + assert_eq!( + back, resume, + "token round-trip must preserve the exact value for {number:?}" + ); + assert_eq!( + back.to_scalar_order_by_item(), + Some(OrderByItem::Number(number)), + "must reconstruct the exact numeric variant, not a lossy float, for {number:?}" + ); + // `OrderByNumber`'s `PartialEq` is intentionally value-based + // (`I64(5) == F64(5.0)`), so for a value like `i64::MIN` that + // also has an exact `f64` representation, the assertions above + // alone can't catch a regression that silently degrades the + // variant to `F64`. Comparing `Debug` output (which prints the + // variant tag) closes that gap. + assert_eq!( + format!("{back:?}"), + format!("{resume:?}"), + "round-trip must preserve the exact variant (I64/U64/F64), not just the \ + numeric value, for {number:?}" + ); + } + } + + #[test] + fn complex_hash_distinguishes_adjacent_large_integers() { + // Integers within `i64` range take the exact `Number64` long path + // (mantissa + `extraBits`), so adjacent values beyond 2^53 never + // collide the way a lossy `f64` cast would. + let a = OrderByItem::Array(vec![OrderByItem::Number(9_007_199_254_740_992_i64.into())]); + let b = OrderByItem::Array(vec![OrderByItem::Number(9_007_199_254_740_993_i64.into())]); + assert_ne!(a.to_resume_value(), b.to_resume_value()); + + // Even at the very top of the `i64` range the `extraBits` keep + // adjacent values distinct. + let c = OrderByItem::Object(vec![( + "n".to_owned(), + OrderByItem::Number((i64::MAX - 1).into()), + )]); + let d = OrderByItem::Object(vec![("n".to_owned(), OrderByItem::Number(i64::MAX.into()))]); + assert_ne!(c.to_resume_value(), d.to_resume_value()); + + // Above `i64::MAX`, `Number64` can only carry the value as a `double` + // (matching the backend), so `u64::MAX` and `u64::MAX - 1` both round + // to 2^64 and intentionally share a hash — the client must agree with + // the backend's lossy representation, not out-precision it. + let e = OrderByItem::Object(vec![( + "n".to_owned(), + OrderByItem::Number((u64::MAX - 1).into()), + )]); + let f = OrderByItem::Object(vec![("n".to_owned(), OrderByItem::Number(u64::MAX.into()))]); + assert_eq!(e.to_resume_value(), f.to_resume_value()); + } + + // ── Envelope parsing ───────────────────────────────────────────────── + + #[test] + fn parses_defined_and_undefined_items() { + let value = serde_json::json!([{"item": 1}, {}, {"item": null}]); + let items = parse_order_by_items(&value, 3).unwrap(); + assert_eq!( + items, + vec![ + OrderByItem::Number(1.0.into()), + OrderByItem::Undefined, + OrderByItem::Null, + ] + ); + } + + #[test] + fn rejects_wrong_length() { + let value = serde_json::json!([{"item": 1}]); + let err = parse_order_by_items(&value, 2).unwrap_err(); + assert_eq!( + err.status(), + crate::error::CosmosStatus::SERVICE_ORDER_BY_ENVELOPE_INVALID + ); + } + + #[test] + fn rejects_non_array() { + let value = serde_json::json!({"item": 1}); + let err = parse_order_by_items(&value, 1).unwrap_err(); + assert_eq!( + err.status(), + crate::error::CosmosStatus::SERVICE_ORDER_BY_ENVELOPE_INVALID + ); + } + + #[test] + fn rejects_non_object_element() { + let value = serde_json::json!([1]); + let err = parse_order_by_items(&value, 1).unwrap_err(); + assert_eq!( + err.status(), + crate::error::CosmosStatus::SERVICE_ORDER_BY_ENVELOPE_INVALID + ); + } + + // ── Multi-column comparator ────────────────────────────────────────── + + #[test] + fn multi_column_mixed_direction() { + let directions = vec![SortOrder::Ascending, SortOrder::Descending]; + let a = vec![ + OrderByItem::Number(1.0.into()), + OrderByItem::String("b".into()), + ]; + let b = vec![ + OrderByItem::Number(1.0.into()), + OrderByItem::String("a".into()), + ]; + // Column 0 ties; column 1 DESC means "b" sorts before "a". + assert_eq!(compare_key_tuples(&a, &b, &directions), Ordering::Less); + } + + // ── Resume value round trip ────────────────────────────────────────── + + #[test] + fn scalar_resume_values_round_trip_through_json() { + for item in [ + OrderByItem::Undefined, + OrderByItem::Null, + OrderByItem::Boolean(true), + OrderByItem::Number(5.5.into()), + OrderByItem::String("s".into()), + ] { + let resume = item.to_resume_value(); + let json = serde_json::to_string(&resume).unwrap(); + let back: OrderByResumeValue = serde_json::from_str(&json).unwrap(); + assert_eq!(back, resume); + // A scalar resume value reconstructs to the exact item it came + // from (the basis of the `_rid`-aware client-side discard). + assert_eq!(resume.to_scalar_order_by_item(), Some(item)); + } + } + + #[test] + fn complex_resume_values_hash_deterministically_and_distinguish_shape() { + let array = OrderByItem::Array(vec![ + OrderByItem::Number(1.0.into()), + OrderByItem::Number(2.0.into()), + ]); + let object = OrderByItem::Object(vec![("x".to_owned(), OrderByItem::Number(1.0.into()))]); + + let array_resume = array.to_resume_value(); + assert!(array_resume.is_complex()); + // An array and an object never share a resume value even if their + // hashes were to collide: the shape tag differs. + assert_ne!(array_resume, object.to_resume_value()); + + // Same content, different wire key order -> identical resume value + // (structurally-equal objects hash identically). + let reordered_object_source = serde_json::json!({"x": 1.0}); + let reordered = OrderByItem::from_json(&reordered_object_source); + assert_eq!(reordered.to_resume_value(), object.to_resume_value()); + } + + #[test] + fn distinct_arrays_do_not_share_a_resume_value() { + let a = OrderByItem::Array(vec![OrderByItem::Number(1.0.into())]); + let b = OrderByItem::Array(vec![OrderByItem::Number(2.0.into())]); + assert_ne!(a.to_resume_value(), b.to_resume_value()); + } + + // ── Resume filter wire model ───────────────────────────────────────── + + #[test] + fn wire_value_matches_dotnet_serialization_contract() { + // Undefined -> `[]`; Null/Bool/Number/String -> raw JSON value. + assert_eq!( + OrderByResumeValue::Undefined.to_wire_value(), + serde_json::json!([]) + ); + assert_eq!( + OrderByResumeValue::Null.to_wire_value(), + serde_json::Value::Null + ); + assert_eq!( + OrderByResumeValue::Boolean { value: true }.to_wire_value(), + serde_json::json!(true) + ); + assert_eq!( + OrderByResumeValue::Number { + value: 5_i64.into() + } + .to_wire_value(), + serde_json::json!(5) + ); + assert_eq!( + OrderByResumeValue::String { + value: "mid".to_owned(), + } + .to_wire_value(), + serde_json::json!("mid") + ); + } + + #[test] + fn wire_value_number_keeps_exact_i64_u64_f64() { + for (n, expected) in [ + (OrderByNumber::from(i64::MIN), serde_json::json!(i64::MIN)), + (OrderByNumber::from(u64::MAX), serde_json::json!(u64::MAX)), + // 2^53 + 1 is not exactly representable as f64; must stay integer. + ( + OrderByNumber::from(9_007_199_254_740_993_i64), + serde_json::json!(9_007_199_254_740_993_i64), + ), + (OrderByNumber::from(1.5_f64), serde_json::json!(1.5_f64)), + ] { + assert_eq!( + OrderByResumeValue::Number { value: n }.to_wire_value(), + expected + ); + } + assert_eq!( + serde_json::to_string( + &OrderByResumeValue::Number { + value: OrderByNumber::from(9_007_199_254_740_993_i64), + } + .to_wire_value() + ) + .unwrap(), + "9007199254740993", + "an exact integer beyond 2^53 must not degrade to float notation" + ); + } + + #[test] + fn complex_wire_value_reinterprets_hash_halves_as_signed_i64() { + // The `low`/`high` halves are the 64-bit hash bits cast to signed + // i64 (matching .NET's `(long)ulong`): the top-bit-set half is + // negative, the other stays positive. + let value = OrderByResumeValue::Complex { + complex_type: ComplexTypeTag::Array, + hash: ComplexHash { + low64: u64::MAX, + high64: 1, + }, + }; + assert_eq!( + value.to_wire_value(), + serde_json::json!({"type": "array", "low": -1_i64, "high": 1_i64}) + ); + + let object = OrderByResumeValue::Complex { + complex_type: ComplexTypeTag::Object, + hash: ComplexHash { + low64: 0x8000_0000_0000_0000, + high64: 0x7fff_ffff_ffff_ffff, + }, + }; + assert_eq!( + object.to_wire_value(), + serde_json::json!({"type": "object", "low": i64::MIN, "high": i64::MAX}) + ); + } + + #[test] + fn complex_wire_value_round_trips_the_real_hash_bits() { + let array = OrderByItem::Array(vec![OrderByItem::Number(1.0.into())]).to_resume_value(); + let OrderByResumeValue::Complex { hash, .. } = array else { + panic!("array resume value must be complex"); + }; + assert_eq!( + array.to_wire_value(), + serde_json::json!({ + "type": "array", + "low": hash.low64 as i64, + "high": hash.high64 as i64, + }) + ); + } + + #[test] + fn resume_filter_json_is_target_style_with_rid_and_exclude_false() { + let filter = resume_filter_json( + &[ + OrderByResumeValue::Number { value: 5.0.into() }, + OrderByResumeValue::String { + value: "mid".to_owned(), + }, + ], + Some("rid-1"), + false, + ); + assert_eq!( + filter, + serde_json::json!({ + "value": [5.0, "mid"], + "rid": "rid-1", + "exclude": false, + }) + ); + } + + #[test] + fn resume_filter_json_omits_rid_when_absent() { + let filter = resume_filter_json(&[OrderByResumeValue::Null], None, true); + assert_eq!( + filter, + serde_json::json!({"value": [null], "exclude": true}) + ); + assert!( + filter.get("rid").is_none(), + "an absent rid must be omitted, not serialized as null" + ); + } + + #[test] + fn classify_row_vs_boundary_scalar_ascending() { + let directions = [SortOrder::Ascending]; + let boundary = [OrderByResumeValue::Number { value: 5.0.into() }]; + let classify = |v: f64| { + classify_row_vs_boundary(&[OrderByItem::Number(v.into())], &boundary, &directions) + }; + assert!(matches!(classify(4.0), RowVsBoundary::Before)); + assert!(matches!(classify(5.0), RowVsBoundary::Tie)); + assert!(matches!(classify(6.0), RowVsBoundary::AfterOrIndeterminate)); + } + + #[test] + fn classify_row_vs_boundary_applies_descending_direction() { + let directions = [SortOrder::Descending]; + let boundary = [OrderByResumeValue::Number { value: 5.0.into() }]; + // Under DESC, a larger value sorts *before* the boundary. + assert!(matches!( + classify_row_vs_boundary(&[OrderByItem::Number(6.0.into())], &boundary, &directions,), + RowVsBoundary::Before + )); + } + + #[test] + fn classify_row_vs_boundary_complex_ties_by_hash_never_infers_order() { + let array = OrderByItem::Array(vec![ + OrderByItem::Number(1.0.into()), + OrderByItem::Number(2.0.into()), + ]); + let boundary = [array.to_resume_value()]; + let directions = [SortOrder::Ascending]; + // The same structural array ties the complex boundary (equal hash). + assert!(matches!( + classify_row_vs_boundary(&[array], &boundary, &directions), + RowVsBoundary::Tie + )); + // A *different* array is never inferred as "before" (its hash order + // is not sort order); it is kept, not dropped. + assert!(matches!( + classify_row_vs_boundary( + &[OrderByItem::Array(vec![OrderByItem::Number(9.0.into())])], + &boundary, + &directions, + ), + RowVsBoundary::AfterOrIndeterminate + )); + } + + #[test] + fn classify_row_vs_boundary_orders_across_types() { + // A String row is always after a Number boundary (type rank). + assert!(matches!( + classify_row_vs_boundary( + &[OrderByItem::String("x".into())], + &[OrderByResumeValue::Number { value: 5.0.into() }], + &[SortOrder::Ascending], + ), + RowVsBoundary::AfterOrIndeterminate + )); + // A Number row is reliably *before* a String boundary (lower rank), + // even though the boundary column is not the row's own type. + assert!(matches!( + classify_row_vs_boundary( + &[OrderByItem::Number(5.0.into())], + &[OrderByResumeValue::String { + value: "x".to_owned() + }], + &[SortOrder::Ascending], + ), + RowVsBoundary::Before + )); + } + + #[test] + fn scalar_resume_values_reconstruct_their_order_by_item() { + for item in [ + OrderByItem::Undefined, + OrderByItem::Null, + OrderByItem::Boolean(true), + OrderByItem::Number(3.5.into()), + OrderByItem::String("x".into()), + ] { + let resume = item.to_resume_value(); + assert_eq!(resume.to_scalar_order_by_item(), Some(item)); + } + let complex = OrderByItem::Array(vec![OrderByItem::Number(1.0.into())]).to_resume_value(); + assert_eq!(complex.to_scalar_order_by_item(), None); + } +} 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 716476653c4..b892d476608 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 @@ -1,6 +1,8 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +// cspell:ignore unemitted checkpointed + //! Pipeline planner for Cosmos DB operations. //! //! The planner validates an operation's target against its resource type and @@ -22,9 +24,12 @@ use crate::{ use super::{ intersect_feed_ranges, - query_plan::{QueryInfo, QueryPlan}, - DrainedLeaf, PartitionRoutingRefresh, Pipeline, PipelineNode, PipelineNodeState, RangedToken, - Request, RequestTarget, ResolvedRange, SequentialDrain, TopologyProvider, UnorderedMerge, + query_plan::{QueryInfo, QueryPlan, SortOrder}, + query_response, + snapshot::{OrderByRangeToken, ValueBoundary}, + streaming_ordered_merge, DrainedLeaf, PartitionRoutingRefresh, Pipeline, PipelineNode, + PipelineNodeState, RangedToken, Request, RequestTarget, ResolvedRange, SequentialDrain, + StreamingOrderedMerge, TopologyProvider, UnorderedMerge, }; /// Builds a single-node [`Pipeline`] for a trivial operation. @@ -204,6 +209,240 @@ pub(crate) async fn build_sequential_drain( Ok(Pipeline::new(root)) } +/// `true` if `query_info` selects the streaming `ORDER BY` pipeline (one or +/// more `ORDER BY` columns not requiring the non-streaming buffered sort). +pub(crate) fn is_streaming_order_by(info: &QueryInfo) -> bool { + !info.order_by.is_empty() && !info.has_non_streaming_order_by +} + +/// Builds a [`streaming_ordered_merge::StreamingOrderedMerge`] pipeline +/// from a backend query plan whose `queryInfo.orderBy` is non-empty. +/// Mirrors [`build_sequential_drain`]'s shape, but snapshots every +/// still-active range explicitly since global ordering means any range +/// may still have unemitted rows. +/// +/// `resume` re-resolves each saved range against current topology and +/// rebuilds it via [`streaming_ordered_merge::build_children`], the same +/// path a live split uses. +pub(crate) async fn build_streaming_ordered_merge( + 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 + .query_info + .as_ref() + .expect("is_streaming_order_by requires query_info to be Some"); + let rewritten_query = info + .rewritten_query + .as_deref() + .filter(|s| !s.is_empty()) + .ok_or_else(|| { + crate::error::CosmosError::builder() + .with_status( + crate::error::CosmosStatus::SERVICE_QUERY_PLAN_ORDER_BY_MISSING_REWRITTEN_QUERY, + ) + .with_message( + "query plan reported one or more ORDER BY columns but did not supply a \ + non-empty rewrittenQuery", + ) + .build() + })?; + let directions = info.order_by.clone(); + + 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)); + + let is_resume = resume.is_some(); + let saved_ranges = match resume { + None => None, + Some(PipelineNodeState::Drained) => { + return Ok(Pipeline::new(Box::new(DrainedLeaf))); + } + Some(PipelineNodeState::StreamingOrderedMerge { + directions: saved_directions, + ranges, + }) => Some(validate_streaming_order_by_snapshot( + &directions, + &saved_directions, + ranges, + )?), + Some(other) => { + return 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 streaming ORDER BY operation", + snapshot_kind(&other) + )) + .build()); + } + }; + + let mut children = Vec::new(); + + if let Some(saved_ranges) = saved_ranges { + for saved in saved_ranges { + let resolved = topology_provider + .resolve_ranges(&saved.range, PartitionRoutingRefresh::UseCached) + .await?; + let mut range_children = streaming_ordered_merge::build_children( + &resolved, + &saved.range, + &plain_operation, + &directions, + saved.server_continuation, + saved.boundary.as_ref(), + )?; + children.append(&mut range_children); + } + } else { + // See `plan_fresh` for rationale on intersecting with the operation scope. + let scope_range = operation.target(); + let normalized_len = operation + .container() + .and_then(|c| normalized_epk_len(c.partition_key_definition())); + for query_range in &query_plan.query_ranges { + let plan_range = query_range_to_feed_range(query_range, normalized_len)?; + let feed_range = match scope_range { + Some(scope) => match intersect_feed_ranges(scope, &plan_range) { + Some(r) => r, + None => continue, + }, + None => plan_range, + }; + let resolved = topology_provider + .resolve_ranges(&feed_range, PartitionRoutingRefresh::UseCached) + .await?; + let mut range_children = streaming_ordered_merge::build_children( + &resolved, + &feed_range, + &plain_operation, + &directions, + None, + None, + )?; + children.append(&mut range_children); + } + } + + if children.is_empty() { + if is_resume { + return Ok(Pipeline::new(Box::new(DrainedLeaf))); + } + return Err(crate::error::CosmosError::builder() + .with_status(crate::error::CosmosStatus::CLIENT_QUERY_PLAN_PRODUCED_EMPTY_RANGES) + .with_message("query plan produced no partition ranges to query") + .build()); + } + + let root = Box::new(StreamingOrderedMerge::new( + plain_operation, + directions, + children, + )); + Ok(Pipeline::new(root)) +} + +/// A saved [`OrderByRangeToken`], parsed and validated into planner-ready +/// types. +struct ParsedOrderByRange { + range: FeedRange, + server_continuation: Option, + boundary: Option, +} + +/// Validates a resumed `StreamingOrderedMerge` continuation's `ORDER BY` +/// direction discriminator against the current query, then validates and +/// parses every saved range: well-formed, sorted, non-overlapping bounds, +/// and a boundary whose resume-value count matches the columns. +fn validate_streaming_order_by_snapshot( + directions: &[SortOrder], + saved_directions: &[SortOrder], + ranges: Vec, +) -> crate::error::Result> { + if saved_directions != directions { + return Err(order_by_state_invalid(format!( + "continuation token has {} ORDER BY column(s)/direction(s) but the current query \ + has {}", + saved_directions.len(), + directions.len(), + ))); + } + if ranges.is_empty() { + return Err(order_by_state_invalid( + "continuation token has an empty StreamingOrderedMerge range list; a fully-drained \ + operation must use the Drained shape instead", + )); + } + + let mut parsed = Vec::with_capacity(ranges.len()); + let mut prev_max: Option = None; + for entry in ranges { + let min = EffectivePartitionKey::from(entry.min_epk); + let max = EffectivePartitionKey::from(entry.max_epk); + if min >= max { + return Err(order_by_state_invalid(format!( + "continuation token has an invalid range (min `{}` >= max `{}`)", + min.to_hex(), + max.to_hex(), + ))); + } + if let Some(prev) = &prev_max { + if &min < prev { + return Err(order_by_state_invalid( + "continuation token ranges must be sorted ascending and non-overlapping", + )); + } + } + prev_max = Some(max.clone()); + + if let Some(boundary) = &entry.boundary { + if boundary.resume_values.len() != directions.len() { + return Err(order_by_state_invalid(format!( + "continuation token range boundary has {} resume value(s) but the query has \ + {} ORDER BY column(s)", + boundary.resume_values.len(), + directions.len(), + ))); + } + if boundary.last_rid.is_empty() { + return Err(order_by_state_invalid( + "continuation token range boundary has an empty RID", + )); + } + // A well-formed boundary counts at least its own boundary row, so + // `skip_count` is always >= 1. An explicit 0 is corrupt (a legacy + // token that omits the field is read back as 1 by serde default, + // not 0, so this only rejects a genuinely malformed value). + if boundary.skip_count == 0 { + return Err(order_by_state_invalid( + "continuation token range boundary has a skip count of 0 (must be >= 1)", + )); + } + } + + parsed.push(ParsedOrderByRange { + range: FeedRange::new(min, max)?, + server_continuation: entry.server_continuation, + boundary: entry.boundary, + }); + } + + Ok(parsed) +} + +fn order_by_state_invalid( + message: impl Into>, +) -> crate::error::CosmosError { + crate::error::CosmosError::builder() + .with_status(crate::error::CosmosStatus::CLIENT_CONTINUATION_TOKEN_ORDER_BY_STATE_INVALID) + .with_message(message) + .build() +} + /// Builds an [`UnorderedMerge`] pipeline for change feed operations. /// /// Unlike [`build_sequential_drain`], this does not require a query plan. @@ -809,6 +1048,7 @@ fn snapshot_kind(state: &PipelineNodeState) -> &'static str { PipelineNodeState::Request { .. } => "Request", PipelineNodeState::SequentialDrain { .. } => "SequentialDrain", PipelineNodeState::UnorderedMerge { .. } => "UnorderedMerge", + PipelineNodeState::StreamingOrderedMerge { .. } => "StreamingOrderedMerge", } } @@ -899,6 +1139,70 @@ fn validate_query_info(info: &QueryInfo) -> crate::error::Result<()> { Ok(()) } +/// Validates a query plan for [`build_streaming_ordered_merge`]: `ORDER BY` +/// is expected, but every other unsupported feature (TOP, non-streaming +/// `ORDER BY`, DISTINCT/GROUP BY/aggregates/OFFSET/LIMIT/hybrid-search) is +/// still rejected the same way [`validate_query_info`] rejects it. +fn validate_query_plan_for_streaming_order_by(plan: &QueryPlan) -> crate::error::Result<()> { + if plan.hybrid_search_query_info.is_some() { + return Err(unsupported_feature("hybrid search queries")); + } + let info = plan.query_info.as_ref().ok_or_else(|| { + // Precondition of `is_streaming_order_by`; an internal planner bug if violated. + crate::error::CosmosError::builder() + .with_status(crate::error::CosmosStatus::CLIENT_UNSUPPORTED_QUERY_FEATURE) + .with_message( + "internal error: streaming ORDER BY path selected with no queryInfo present", + ) + .build() + })?; + if info.has_non_streaming_order_by { + return Err(unsupported_feature( + "non-streaming ORDER BY in cross-partition queries", + )); + } + if info.top.is_some() { + return Err(unsupported_feature( + "TOP combined with streaming ORDER BY (requires the TOP composition stage)", + )); + } + if info.offset.is_some() || info.limit.is_some() { + return Err(unsupported_feature( + "OFFSET/LIMIT combined with ORDER BY in cross-partition queries", + )); + } + if !info.aggregates.is_empty() { + return Err(unsupported_feature( + "aggregates combined with ORDER BY in cross-partition queries", + )); + } + if !info.group_by_expressions.is_empty() { + return Err(unsupported_feature( + "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", + )); + } + // Parallel arrays paired by the resume filter; must match in length. + if info.order_by_expressions.len() != info.order_by.len() { + return Err(crate::error::CosmosError::builder() + .with_status( + crate::error::CosmosStatus::SERVICE_QUERY_PLAN_ORDER_BY_EXPRESSIONS_MISMATCH, + ) + .with_message(format!( + "query plan's ORDER BY metadata is inconsistent: {} sort direction(s) but {} \ + sort-key expression(s)", + info.order_by.len(), + info.order_by_expressions.len(), + )) + .build()); + } + Ok(()) +} + fn unsupported_feature(feature: &str) -> crate::error::CosmosError { crate::error::CosmosError::builder() .with_status(crate::error::CosmosStatus::CLIENT_UNSUPPORTED_QUERY_FEATURE) @@ -2473,4 +2777,309 @@ mod tests { ); } } + + // ── Streaming ORDER BY selection and validation ─────────────────────── + + fn order_by_query_info(rewritten_query: Option<&str>) -> QueryInfo { + QueryInfo { + order_by: vec![SortOrder::Ascending], + order_by_expressions: vec!["c.rank".to_owned()], + rewritten_query: rewritten_query.map(str::to_owned), + ..Default::default() + } + } + + fn order_by_plan(rewritten_query: Option<&str>, ranges: Vec) -> QueryPlan { + QueryPlan { + partitioned_query_execution_info_version: 2, + query_info: Some(order_by_query_info(rewritten_query)), + query_ranges: ranges, + hybrid_search_query_info: None, + } + } + + fn order_by_operation() -> CosmosOperation { + CosmosOperation::query_items(test_container(), Some(FeedRange::full())) + .with_body(br#"{"query":"SELECT * FROM c ORDER BY c.rank","parameters":[]}"#.to_vec()) + } + + #[test] + fn is_streaming_order_by_true_only_for_non_empty_streaming_order_by() { + assert!(!is_streaming_order_by(&QueryInfo::default())); + assert!(is_streaming_order_by(&order_by_query_info(Some( + "SELECT 1" + )))); + + let mut non_streaming = order_by_query_info(Some("SELECT 1")); + non_streaming.has_non_streaming_order_by = true; + assert!(!is_streaming_order_by(&non_streaming)); + } + + #[test] + fn validate_query_plan_for_streaming_order_by_accepts_plain_order_by() { + let plan = order_by_plan(Some("SELECT 1"), vec![qr("", "FF")]); + assert!(validate_query_plan_for_streaming_order_by(&plan).is_ok()); + } + + #[test] + fn validate_query_plan_for_streaming_order_by_rejects_non_streaming_order_by() { + let mut plan = order_by_plan(Some("SELECT 1"), vec![qr("", "FF")]); + plan.query_info.as_mut().unwrap().has_non_streaming_order_by = true; + let err = validate_query_plan_for_streaming_order_by(&plan).unwrap_err(); + assert_eq!( + err.status(), + crate::error::CosmosStatus::CLIENT_UNSUPPORTED_QUERY_FEATURE + ); + } + + #[test] + fn validate_query_plan_for_streaming_order_by_rejects_top() { + let mut plan = order_by_plan(Some("SELECT 1"), vec![qr("", "FF")]); + plan.query_info.as_mut().unwrap().top = Some(5); + assert!(validate_query_plan_for_streaming_order_by(&plan).is_err()); + } + + #[test] + fn validate_query_plan_for_streaming_order_by_rejects_offset_limit() { + let mut plan = order_by_plan(Some("SELECT 1"), vec![qr("", "FF")]); + plan.query_info.as_mut().unwrap().offset = Some(1); + assert!(validate_query_plan_for_streaming_order_by(&plan).is_err()); + + let mut plan = order_by_plan(Some("SELECT 1"), vec![qr("", "FF")]); + plan.query_info.as_mut().unwrap().limit = Some(1); + assert!(validate_query_plan_for_streaming_order_by(&plan).is_err()); + } + + #[test] + fn validate_query_plan_for_streaming_order_by_rejects_aggregates_group_by_distinct() { + 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()); + + 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()); + + 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()); + } + + #[test] + fn validate_query_plan_for_streaming_order_by_rejects_hybrid_search() { + let mut plan = order_by_plan(Some("SELECT 1"), vec![qr("", "FF")]); + plan.hybrid_search_query_info = + Some(crate::driver::dataflow::query_plan::HybridSearchQueryInfo { + global_statistics_query: String::new(), + component_query_infos: vec![], + component_weights: vec![], + skip: None, + take: None, + requires_global_statistics: false, + }); + assert!(validate_query_plan_for_streaming_order_by(&plan).is_err()); + } + + #[tokio::test] + async fn build_streaming_ordered_merge_rejects_missing_rewritten_query() { + let op = Arc::new(order_by_operation()); + let plan = order_by_plan(None, vec![qr("", "FF")]); + let mut topology = MockTopologyProvider::new(vec![]); + let err = build_streaming_ordered_merge(&plan, &mut topology, &op, None) + .await + .unwrap_err(); + assert_eq!( + err.status(), + crate::error::CosmosStatus::SERVICE_QUERY_PLAN_ORDER_BY_MISSING_REWRITTEN_QUERY + ); + } + + #[tokio::test] + async fn build_streaming_ordered_merge_rejects_empty_rewritten_query() { + let op = Arc::new(order_by_operation()); + let plan = order_by_plan(Some(""), vec![qr("", "FF")]); + let mut topology = MockTopologyProvider::new(vec![]); + let err = build_streaming_ordered_merge(&plan, &mut topology, &op, None) + .await + .unwrap_err(); + assert_eq!( + err.status(), + crate::error::CosmosStatus::SERVICE_QUERY_PLAN_ORDER_BY_MISSING_REWRITTEN_QUERY + ); + } + + #[tokio::test] + async fn build_streaming_ordered_merge_builds_one_child_per_resolved_range() { + let op = Arc::new(order_by_operation()); + let plan = order_by_plan(Some("SELECT c._rid, [{\"item\":c.rank}] AS orderByItems, c AS payload FROM c ORDER BY c.rank ASC"), vec![qr("", "FF")]); + let mut topology = MockTopologyProvider::new(vec![Ok(vec![ + rr("", "80", "pk-left"), + rr("80", "FF", "pk-right"), + ])]); + + let pipeline = build_streaming_ordered_merge(&plan, &mut topology, &op, None) + .await + .unwrap(); + let root = pipeline.into_root(); + let merge = root + .downcast::() + .expect("root must be a StreamingOrderedMerge"); + let children = merge.into_children(); + assert_eq!( + children.len(), + 2, + "one child per resolved physical partition" + ); + } + + #[tokio::test] + async fn build_streaming_ordered_merge_empty_ranges_errors_on_fresh_start() { + let op = Arc::new(order_by_operation()); + // Empty `query_ranges` exercises the "produced nothing" guard. + let plan = order_by_plan(Some("SELECT 1"), vec![]); + let mut topology = MockTopologyProvider::new(vec![]); + let err = build_streaming_ordered_merge(&plan, &mut topology, &op, None) + .await + .unwrap_err(); + assert_eq!( + err.status(), + crate::error::CosmosStatus::CLIENT_QUERY_PLAN_PRODUCED_EMPTY_RANGES + ); + } + + /// An operation whose SQL body carries `sql` as the `query` field. + fn order_by_operation_with_query(sql: &str) -> CosmosOperation { + let body = serde_json::json!({ "query": sql, "parameters": [] }); + CosmosOperation::query_items(test_container(), Some(FeedRange::full())) + .with_body(serde_json::to_vec(&body).unwrap()) + } + + #[tokio::test] + async fn build_streaming_ordered_merge_accepts_join_query() { + // A JOIN can emit multiple rows per document `_rid`; the resume cursor + // now tracks that with a `_rid`-tie skip count (mirroring .NET), so a + // JOIN-shaped ORDER BY builds instead of being rejected up front. + let op = Arc::new(order_by_operation_with_query( + "SELECT * FROM c JOIN t IN c.tags ORDER BY c.rank", + )); + let plan = order_by_plan( + Some("SELECT c._rid, [{\"item\":c.rank}] AS orderByItems, c AS payload FROM c JOIN t IN c.tags ORDER BY c.rank ASC"), + vec![qr("", "FF")], + ); + let mut topology = MockTopologyProvider::new(vec![Ok(vec![rr("", "FF", "pk-0")])]); + let pipeline = build_streaming_ordered_merge(&plan, &mut topology, &op, None) + .await + .expect("a JOIN-shaped ORDER BY must build a pipeline"); + assert!( + pipeline + .into_root() + .downcast::() + .is_some(), + "a JOIN-shaped ORDER BY builds a StreamingOrderedMerge" + ); + } + + #[tokio::test] + async fn build_streaming_ordered_merge_accepts_ordinary_parameterized_order_by() { + // An ordinary single-source, parameterized ORDER BY builds normally + // (no local SQL parsing gate stands in the way). + let op = Arc::new(order_by_operation_with_query( + "SELECT * FROM c WHERE c.tenant = @t ORDER BY c.rank", + )); + let plan = order_by_plan( + Some("SELECT c._rid, [{\"item\":c.rank}] AS orderByItems, c AS payload FROM c ORDER BY c.rank ASC"), + vec![qr("", "FF")], + ); + let mut topology = MockTopologyProvider::new(vec![Ok(vec![rr("", "FF", "pk-0")])]); + let pipeline = build_streaming_ordered_merge(&plan, &mut topology, &op, None) + .await + .expect("ordinary parameterized ORDER BY must build a pipeline"); + assert!( + pipeline + .into_root() + .downcast::() + .is_some(), + "ordinary parameterized ORDER BY builds a StreamingOrderedMerge" + ); + } + + /// A boundary in a resumed `StreamingOrderedMerge` snapshot always counts + /// at least its own boundary row, so an explicit `skip_count` of 0 is a + /// corrupt token and must be rejected. + #[test] + fn streaming_order_by_snapshot_rejects_zero_skip_count() { + let ranges = vec![OrderByRangeToken { + min_epk: String::new(), + max_epk: "FF".to_owned(), + server_continuation: None, + boundary: Some(ValueBoundary { + resume_values: vec![ + crate::driver::dataflow::order_by::OrderByResumeValue::Number { + value: 5.0.into(), + }, + ], + last_rid: "rid-1".to_owned(), + skip_count: 0, + }), + }]; + let err = validate_streaming_order_by_snapshot( + &[SortOrder::Ascending], + &[SortOrder::Ascending], + ranges, + ) + .err() + .expect("a boundary skip_count of 0 must be rejected"); + assert_eq!( + err.status().sub_status(), + Some(crate::error::SubStatusCode::CLIENT_CONTINUATION_TOKEN_ORDER_BY_STATE_INVALID), + "a boundary skip_count of 0 must be rejected, got: {err}" + ); + } + + /// The smallest well-formed boundary carries `skip_count == 1` (just the + /// boundary row) and validates. + #[test] + fn streaming_order_by_snapshot_accepts_skip_count_one() { + let ranges = vec![OrderByRangeToken { + min_epk: String::new(), + max_epk: "FF".to_owned(), + server_continuation: None, + boundary: Some(ValueBoundary { + resume_values: vec![ + crate::driver::dataflow::order_by::OrderByResumeValue::Number { + value: 5.0.into(), + }, + ], + last_rid: "rid-1".to_owned(), + skip_count: 1, + }), + }]; + let parsed = validate_streaming_order_by_snapshot( + &[SortOrder::Ascending], + &[SortOrder::Ascending], + ranges, + ) + .expect("a well-formed boundary with skip_count 1 is valid"); + assert_eq!(parsed.len(), 1); + assert_eq!(parsed[0].boundary.as_ref().unwrap().skip_count, 1); + } + + /// A legacy token that omits `skip_count` deserializes as 1 (not 0), so it + /// still validates — its boundary row is discarded on resume, never + /// re-emitted. + #[test] + fn streaming_order_by_snapshot_defaults_missing_skip_count_to_one() { + let token: OrderByRangeToken = serde_json::from_str( + r#"{"min_epk":"","max_epk":"FF","boundary":{"resume_values":[{"type":"number","value":5.0}],"last_rid":"rid-1"}}"#, + ) + .unwrap(); + assert_eq!(token.boundary.as_ref().unwrap().skip_count, 1); + let parsed = validate_streaming_order_by_snapshot( + &[SortOrder::Ascending], + &[SortOrder::Ascending], + vec![token], + ) + .expect("a legacy boundary missing skip_count defaults to 1 and is valid"); + assert_eq!(parsed[0].boundary.as_ref().unwrap().skip_count, 1); + } } diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/query_response.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/query_response.rs new file mode 100644 index 00000000000..7ad0cf8e8be --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/query_response.rs @@ -0,0 +1,532 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +//! Rewritten-query construction and response-envelope plumbing for the +//! cross-partition streaming `ORDER BY` pipeline. +//! +//! - **Request**: [`rewrite_query_body`] swaps only the `"query"` text for +//! the Gateway's `QueryInfo::rewritten_query` (its resume-filter +//! placeholder replaced with `true`), preserving `"parameters"`. On +//! resume, [`with_resume_filter`] additionally inserts the .NET-compatible +//! structured `"resumeFilter"` field — never rewriting the SQL or the +//! caller's parameters. +//! - **Response**: [`parse_envelope_page`] parses a backend page into +//! strict [`EnvelopeRow`]s (envelope shape `{"_rid", "orderByItems", +//! "payload"}`), retaining `payload` as raw JSON bytes. +//! [`PageAggregator`] accumulates charge/diagnostics across pages, and +//! [`PageAggregator::build_page`] reconstructs a synthetic +//! `Documents`-array [`CosmosResponse`] — the same wire shape every +//! other feed node returns. +//! +//! Backend continuations are never copied onto the emitted page's headers: +//! `OperationPlan::to_continuation_token` owns the client-issued token, and +//! a raw per-partition backend token would be meaningless to the caller. + +use std::sync::Arc; + +use serde::Deserialize; +use serde_json::value::RawValue; + +use crate::diagnostics::{DiagnosticsContext, DiagnosticsContextBuilder}; +use crate::models::{ + ActivityId, CosmosResponse, CosmosResponseHeaders, CosmosStatus, ResponseBody, +}; +use crate::options::DiagnosticsOptions; + +use super::order_by::{self, OrderByItem, OrderByResumeValue}; + +/// Placeholder emitted by the Gateway in rewritten streaming `ORDER BY` +/// queries so SDKs can inject a continuation resume filter. +const ORDER_BY_FILTER_PLACEHOLDER: &str = "{documentdb-formattableorderbyquery-filter}"; + +/// Produces the executable rewritten query used for a fresh range: the +/// Gateway's filter placeholder is replaced with `true` (matching +/// .NET/Java), or left as-is if the placeholder is absent. +pub(crate) fn rewritten_query_from_beginning(rewritten_query: &str) -> String { + rewritten_query.replace(ORDER_BY_FILTER_PLACEHOLDER, "true") +} + +/// Replaces only the `"query"` field of a query operation's JSON body with +/// `rewritten_query`, preserving `"parameters"` (and any other field) +/// verbatim. +/// +/// # Errors +/// +/// Returns a typed error if `original_body` is missing, is not valid JSON, +/// or is not a JSON object (the query operation body is always expected to +/// be `{"query": ..., "parameters": [...]}`). +pub(crate) fn rewrite_query_body( + original_body: Option<&[u8]>, + rewritten_query: &str, +) -> crate::error::Result> { + let mut value = parse_query_body(original_body)?; + let obj = value + .as_object_mut() + .ok_or_else(|| body_error_msg("original query body is not a JSON object"))?; + obj.insert( + "query".to_owned(), + serde_json::Value::String(rewritten_query.to_owned()), + ); + serde_json::to_vec(&value) + .map_err(|e| body_error("failed to serialize rewritten query body", e)) +} + +/// Inserts the .NET-compatible structured `"resumeFilter"` field into an +/// already-rewritten query `body` (its `"query"` text already +/// placeholder-substituted with `true` by [`rewrite_query_body`]), +/// preserving `"query"`, the caller's `"parameters"`, and every other field +/// verbatim. Mirrors `SqlQuerySpec.resumeFilter`: the backend seeks to the +/// resume point with no SDK-side SQL rewriting and no generated parameters. +/// +/// Rust persists one boundary per range and resumes each with `rid` present +/// and `exclude:false` (the .NET "target" partition style); the +/// already-emitted prefix of the boundary tie run is trimmed client-side +/// (see [`super::order_by::classify_row_vs_boundary`]). +/// +/// # Errors +/// +/// Returns a typed error if `body` is missing, is not valid JSON, or is not +/// a JSON object. +pub(crate) fn with_resume_filter( + body: Option<&[u8]>, + resume_values: &[OrderByResumeValue], + rid: Option<&str>, + exclude: bool, +) -> crate::error::Result> { + let mut value = parse_query_body(body)?; + let obj = value + .as_object_mut() + .ok_or_else(|| body_error_msg("query body is not a JSON object"))?; + obj.insert( + "resumeFilter".to_owned(), + order_by::resume_filter_json(resume_values, rid, exclude), + ); + serde_json::to_vec(&value) + .map_err(|e| body_error("failed to serialize query body with resume filter", e)) +} + +/// Parses a query operation's JSON body, erroring on a missing body or +/// invalid JSON. +fn parse_query_body(body: Option<&[u8]>) -> crate::error::Result { + let body = body.ok_or_else(|| { + body_error_msg("cannot rewrite an ORDER BY query operation with no request body") + })?; + serde_json::from_slice(body) + .map_err(|e| body_error("failed to parse original query body as JSON", e)) +} + +/// One row parsed from a rewritten-envelope backend page, ready for the +/// merge heap. `payload` retains the item's exact original JSON bytes +/// (via [`RawValue`]) rather than a re-serialized value, so the emitted +/// item is byte-identical to what the backend returned. +#[derive(Debug)] +pub(crate) struct EnvelopeRow { + pub(crate) keys: Vec, + pub(crate) rid: String, + pub(crate) payload: Box, +} + +/// A single rewritten-envelope item, as deserialized off the wire. +#[derive(Deserialize)] +struct EnvelopeItem { + #[serde(rename = "_rid")] + rid: Option, + #[serde(rename = "orderByItems")] + order_by_items: Option, + payload: Option>, +} + +/// The standard Cosmos feed-response wire shape: +/// `{"_rid": ..., "Documents": [...], "_count": N}`. Every item is kept as +/// an owned [`RawValue`] so its exact bytes are available for +/// [`parse_envelope_item`] without a second parse-and-copy pass. +#[derive(Deserialize)] +struct RawFeedBody { + #[serde(alias = "Documents")] + documents: Vec>, +} + +/// Parses a child backend page's response body into envelope rows, in wire +/// order (the backend already returned them in this partition's local +/// sort order). +/// +/// # Errors +/// +/// Returns a typed [`crate::error::CosmosError`] with +/// [`CosmosStatus::SERVICE_ORDER_BY_ENVELOPE_INVALID`] for any malformed +/// envelope: a non-feed body shape, an item missing `payload` or a +/// non-empty `_rid`, or an `orderByItems` array whose length does not +/// match `order_by_column_count`. +pub(crate) fn parse_envelope_page( + body: &ResponseBody, + order_by_column_count: usize, +) -> crate::error::Result> { + let bytes = match body { + ResponseBody::NoPayload => return Ok(Vec::new()), + ResponseBody::Bytes(b) => b.clone(), + ResponseBody::Items(_) => { + return Err(envelope_error( + "rewritten-query backend page returned an already-split `Items` body; \ + expected a raw `Documents`-array feed body", + )); + } + }; + let feed: RawFeedBody = serde_json::from_slice(&bytes).map_err(|e| { + body_error( + "failed to parse rewritten-query backend page as a feed body", + e, + ) + })?; + feed.documents + .into_iter() + .map(|raw| parse_envelope_item(&raw, order_by_column_count)) + .collect() +} + +fn parse_envelope_item( + raw: &RawValue, + order_by_column_count: usize, +) -> crate::error::Result { + let item: EnvelopeItem = serde_json::from_str(raw.get()) + .map_err(|e| body_error("failed to parse rewritten envelope item", e))?; + let rid = item + .rid + .filter(|s| !s.is_empty()) + .ok_or_else(|| envelope_error("rewritten envelope item is missing a non-empty `_rid`"))?; + let order_by_items = item + .order_by_items + .ok_or_else(|| envelope_error("rewritten envelope item is missing `orderByItems`"))?; + let keys = order_by::parse_order_by_items(&order_by_items, order_by_column_count)?; + let payload = item + .payload + .ok_or_else(|| envelope_error("rewritten envelope item is missing `payload`"))?; + Ok(EnvelopeRow { keys, rid, payload }) +} + +/// Accumulates request charge and diagnostics across every backend page +/// consumed while assembling one emitted output page, then reconstructs a +/// single [`CosmosResponse`] carrying only the raw payload items. +pub(crate) struct PageAggregator { + request_charge: crate::models::RequestCharge, + diagnostics_sources: Vec>, + activity_id: Option, + status: CosmosStatus, +} + +impl Default for PageAggregator { + fn default() -> Self { + Self { + request_charge: crate::models::RequestCharge::default(), + diagnostics_sources: Vec::new(), + activity_id: None, + status: CosmosStatus::new(azure_core::http::StatusCode::Ok), + } + } +} + +impl PageAggregator { + pub(crate) fn new() -> Self { + Self::default() + } + + /// Folds one consumed backend page's headers/diagnostics/status into + /// the running aggregate. The *last* absorbed response's activity ID + /// and status win (matching how a single multi-page fetch already + /// surfaces "the most recent" status to callers); request charge is + /// summed across every absorbed response. + pub(crate) fn absorb(&mut self, response: &CosmosResponse) { + let charge = response.headers().request_charge.unwrap_or_default(); + self.request_charge = self.request_charge + charge; + self.diagnostics_sources.push(response.diagnostics()); + if let Some(id) = &response.headers().activity_id { + self.activity_id = Some(id.clone()); + } + self.status = response.status(); + } + + /// Builds the emitted page from the accumulated aggregate plus the + /// final ordered list of raw item payloads. + /// + /// Body is a synthetic `{"_rid": "", "Documents": [...], "_count": N}` + /// envelope of `payloads`' bytes unmodified — the wire shape every + /// other feed node returns. `_rid` is left empty (per-item `_rid` + /// identifies each item, not the feed-level one). + /// + /// It's valid for no backend page to have been absorbed (page + /// assembled entirely from previously-buffered rows); it then reports + /// zero charge and a fresh, empty [`DiagnosticsContext`]. + pub(crate) fn build_page( + self, + payloads: &[Box], + ) -> crate::error::Result { + let mut body = + Vec::with_capacity(64 + payloads.iter().map(|p| p.get().len() + 1).sum::()); + body.extend_from_slice(br#"{"_rid":"","Documents":["#); + for (i, payload) in payloads.iter().enumerate() { + if i > 0 { + body.push(b','); + } + body.extend_from_slice(payload.get().as_bytes()); + } + body.extend_from_slice(br#"],"_count":"#); + body.extend_from_slice(payloads.len().to_string().as_bytes()); + body.push(b'}'); + + let diagnostics = DiagnosticsContext::aggregate_sub_operations(&self.diagnostics_sources) + .map(Arc::new) + .unwrap_or_else(empty_diagnostics); + + let headers = CosmosResponseHeaders { + activity_id: self.activity_id, + request_charge: Some(self.request_charge), + item_count: Some(payloads.len() as u32), + // Omitted: `continuation` (owned by + // `OperationPlan::to_continuation_token`), `session_token`/`etag` + // (not meaningful once pages are interleaved). + ..Default::default() + }; + + Ok(CosmosResponse::new( + ResponseBody::from_bytes(body), + headers, + self.status, + diagnostics, + )) + } +} + +/// A fresh, empty [`DiagnosticsContext`] for a page needing no new backend +/// fetch. Uses a new activity ID since it corresponds to no real request. +fn empty_diagnostics() -> Arc { + let mut builder = DiagnosticsContextBuilder::new( + ActivityId::new_uuid(), + Arc::new(DiagnosticsOptions::default()), + ); + builder.set_operation_status(azure_core::http::StatusCode::Ok, None); + Arc::new(builder.complete()) +} + +fn envelope_error(message: impl Into>) -> crate::error::CosmosError { + crate::error::CosmosError::builder() + .with_status(CosmosStatus::SERVICE_ORDER_BY_ENVELOPE_INVALID) + .with_message(message) + .build() +} + +fn body_error(message: &'static str, source: serde_json::Error) -> crate::error::CosmosError { + crate::error::CosmosError::builder() + .with_status(CosmosStatus::SERVICE_ORDER_BY_ENVELOPE_INVALID) + .with_message(message) + .with_source(source) + .build() +} + +fn body_error_msg(message: &'static str) -> crate::error::CosmosError { + crate::error::CosmosError::builder() + .with_status(CosmosStatus::SERVICE_ORDER_BY_ENVELOPE_INVALID) + .with_message(message) + .build() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::driver::dataflow::mocks::response; + + #[test] + fn rewrite_query_body_replaces_query_and_preserves_parameters() { + let original = br#"{"query":"SELECT * FROM c","parameters":[{"name":"@p","value":1}]}"#; + let rewritten = rewrite_query_body(Some(original), "SELECT c._rid FROM c").unwrap(); + let value: serde_json::Value = serde_json::from_slice(&rewritten).unwrap(); + assert_eq!(value["query"], "SELECT c._rid FROM c"); + assert_eq!(value["parameters"][0]["name"], "@p"); + assert_eq!(value["parameters"][0]["value"], 1); + } + + #[test] + fn with_resume_filter_inserts_target_style_filter_and_preserves_caller_parameters() { + // The caller's query text and parameters are untouched; the resume + // point is a top-level structured `resumeFilter` (rid present, + // exclude false), never SDK-generated SQL or parameters. + let plain = br#"{"query":"SELECT c._rid FROM c WHERE true","parameters":[{"name":"@p","value":1}]}"#; + let resume = [ + OrderByResumeValue::Number { + value: 9_007_199_254_740_993_i64.into(), + }, + OrderByResumeValue::String { + value: "a' OR 1=1".to_owned(), + }, + ]; + let body = with_resume_filter(Some(plain), &resume, Some("rid-1"), false).unwrap(); + let value: serde_json::Value = serde_json::from_slice(&body).unwrap(); + // Caller query and parameters are byte-for-byte preserved. + assert_eq!(value["query"], "SELECT c._rid FROM c WHERE true"); + assert_eq!( + value["parameters"], + serde_json::json!([{"name": "@p", "value": 1}]), + "caller parameters must be preserved with nothing appended" + ); + // The resume filter is the .NET structured shape; the exact integer + // survives the JSON round-trip and the adversarial string is a plain + // value (never interpolated into SQL). + assert_eq!( + value["resumeFilter"], + serde_json::json!({ + "value": [9_007_199_254_740_993_i64, "a' OR 1=1"], + "rid": "rid-1", + "exclude": false, + }) + ); + } + + #[test] + fn with_resume_filter_creates_no_parameters_array_and_omits_absent_rid() { + let plain = br#"{"query":"SELECT 1"}"#; + let body = + with_resume_filter(Some(plain), &[OrderByResumeValue::Null], None, true).unwrap(); + let value: serde_json::Value = serde_json::from_slice(&body).unwrap(); + assert!( + value.get("parameters").is_none(), + "no parameters must be synthesized when the caller had none" + ); + assert_eq!( + value["resumeFilter"], + serde_json::json!({"value": [null], "exclude": true}) + ); + assert!( + value["resumeFilter"].get("rid").is_none(), + "an absent rid must be omitted from the wire filter" + ); + } + + #[test] + fn rewritten_query_from_beginning_replaces_order_by_filter_placeholder() { + let rewritten = "SELECT * FROM c WHERE {documentdb-formattableorderbyquery-filter}"; + assert_eq!( + rewritten_query_from_beginning(rewritten), + "SELECT * FROM c WHERE true" + ); + } + + #[test] + fn rewritten_query_from_beginning_preserves_query_without_placeholder() { + let rewritten = "SELECT * FROM c"; + assert_eq!(rewritten_query_from_beginning(rewritten), rewritten); + } + + #[test] + fn rewrite_query_body_rejects_missing_body() { + let err = rewrite_query_body(None, "SELECT 1").unwrap_err(); + assert_eq!( + err.status(), + CosmosStatus::SERVICE_ORDER_BY_ENVELOPE_INVALID + ); + } + + #[test] + fn rewrite_query_body_rejects_non_object_body() { + let err = rewrite_query_body(Some(b"[1,2,3]"), "SELECT 1").unwrap_err(); + assert_eq!( + err.status(), + CosmosStatus::SERVICE_ORDER_BY_ENVELOPE_INVALID + ); + } + + #[test] + fn parse_envelope_page_parses_documents_array() { + let body = ResponseBody::from_bytes( + br#"{"_rid":"abc","Documents":[{"_rid":"r1","orderByItems":[{"item":1}],"payload":{"id":"d1"}},{"_rid":"r2","orderByItems":[{}],"payload":{"id":"d2"}}],"_count":2}"#.to_vec(), + ); + let rows = parse_envelope_page(&body, 1).unwrap(); + assert_eq!(rows.len(), 2); + assert_eq!(rows[0].rid, "r1"); + assert_eq!(rows[0].keys, vec![OrderByItem::Number(1.0.into())]); + assert_eq!(rows[0].payload.get(), r#"{"id":"d1"}"#); + assert_eq!(rows[1].rid, "r2"); + assert_eq!(rows[1].keys, vec![OrderByItem::Undefined]); + } + + #[test] + fn parse_envelope_page_empty_body_yields_no_rows() { + let rows = parse_envelope_page(&ResponseBody::NoPayload, 1).unwrap(); + assert!(rows.is_empty()); + } + + #[test] + fn parse_envelope_page_rejects_missing_rid() { + let body = ResponseBody::from_bytes( + br#"{"Documents":[{"orderByItems":[{"item":1}],"payload":{"id":"d1"}}]}"#.to_vec(), + ); + let err = parse_envelope_page(&body, 1).unwrap_err(); + assert_eq!( + err.status(), + CosmosStatus::SERVICE_ORDER_BY_ENVELOPE_INVALID + ); + } + + #[test] + fn parse_envelope_page_rejects_missing_payload() { + let body = ResponseBody::from_bytes( + br#"{"Documents":[{"_rid":"r1","orderByItems":[{"item":1}]}]}"#.to_vec(), + ); + let err = parse_envelope_page(&body, 1).unwrap_err(); + assert_eq!( + err.status(), + CosmosStatus::SERVICE_ORDER_BY_ENVELOPE_INVALID + ); + } + + #[test] + fn parse_envelope_page_rejects_items_body_shape() { + let body = ResponseBody::from_items(vec![]); + let err = parse_envelope_page(&body, 1).unwrap_err(); + assert_eq!( + err.status(), + CosmosStatus::SERVICE_ORDER_BY_ENVELOPE_INVALID + ); + } + + #[test] + fn page_aggregator_sums_charge_and_builds_documents_envelope() { + let mut aggregator = PageAggregator::new(); + aggregator.absorb(&response(b"{}")); + aggregator.absorb(&response(b"{}")); + let payloads: Vec> = vec![ + RawValue::from_string(r#"{"id":"a"}"#.to_owned()).unwrap(), + RawValue::from_string(r#"{"id":"b"}"#.to_owned()).unwrap(), + ]; + let page = aggregator.build_page(&payloads).unwrap(); + let body = page.body_bytes(); + let value: serde_json::Value = serde_json::from_slice(body).unwrap(); + assert_eq!(value["Documents"].as_array().unwrap().len(), 2); + assert_eq!(value["Documents"][0]["id"], "a"); + assert_eq!(value["_count"], 2); + assert!(page.headers().continuation.is_none()); + } + + #[test] + fn page_aggregator_builds_empty_page_when_polled_children_yielded_no_rows() { + // At least one page absorbed, but zero rows contributed. + let mut aggregator = PageAggregator::new(); + aggregator.absorb(&response(b"{}")); + let page = aggregator.build_page(&[]).unwrap(); + let value: serde_json::Value = serde_json::from_slice(page.body_bytes()).unwrap(); + assert_eq!(value["Documents"].as_array().unwrap().len(), 0); + assert_eq!(value["_count"], 0); + } + + #[test] + fn page_aggregator_with_no_absorbed_response_builds_zero_charge_page() { + // No new fetch needed; must still build a valid zero-charge page. + let aggregator = PageAggregator::new(); + let payloads: Vec> = + vec![RawValue::from_string(r#"{"id":"a"}"#.to_owned()).unwrap()]; + let page = aggregator.build_page(&payloads).unwrap(); + assert_eq!( + page.headers().request_charge, + Some(crate::models::RequestCharge::new(0.0)) + ); + let value: serde_json::Value = serde_json::from_slice(page.body_bytes()).unwrap(); + assert_eq!(value["Documents"][0]["id"], "a"); + } +} 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 05c56c648a9..135f3111cda 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 @@ -1,6 +1,8 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +// cspell:ignore unemitted rescan + //! Pipeline node snapshot state used to serialize / deserialize continuation //! tokens. //! @@ -18,6 +20,9 @@ use serde::{Deserialize, Serialize}; use crate::models::ChangeFeedStartFrom; +use super::order_by::OrderByResumeValue; +use super::query_plan::SortOrder; + /// Serializable snapshot of a [`PipelineNode`](super::PipelineNode) subtree. /// /// The shape is intentionally open to future intermediate node kinds so a @@ -78,6 +83,77 @@ pub(crate) enum PipelineNodeState { #[serde(default, skip_serializing_if = "Option::is_none")] start_from: Option, }, + + /// A streaming cross-partition `ORDER BY` k-way merge. + /// + /// `directions` is a query-shape discriminator validated on resume + /// alongside range/boundary compatibility, so a token can never resume + /// a structurally incompatible query. `ranges` lists every still-active + /// EPK range explicitly (unlike `SequentialDrain`'s sparse cursor, since + /// any range may still have unemitted rows), sorted ascending by + /// `min_epk`; a fully-drained range is omitted. + StreamingOrderedMerge { + directions: Vec, + ranges: Vec, + }, +} + +/// One still-active range of a [`PipelineNodeState::StreamingOrderedMerge`]. +/// `server_continuation` and `boundary` are independent: +/// +/// - `server_continuation` is present only when the buffer was empty and +/// the range ran the plain query — resumable like +/// [`PipelineNodeState::Request`]. A resume-filtered range's token is +/// opaque and bound to that filtered request body, so it's never recorded +/// here; such a range carries only its `boundary`. +/// - `boundary` is present once a row has been emitted; it's the required +/// resume path whenever `server_continuation` is absent or untrusted: +/// reissue with the structured `resumeFilter` (scalar or complex keys +/// alike — see `super::order_by::resume_filter_json`). +/// +/// Both absent means this range has never been touched. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct OrderByRangeToken { + pub(crate) min_epk: String, + pub(crate) max_epk: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub(crate) server_continuation: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub(crate) boundary: Option, +} + +/// The last-emitted row's key tuple + RID for one range, plus a `_rid`-tie +/// `skip_count`. +/// +/// Resumes by sending the boundary to the backend as a structured +/// `resumeFilter` (see `super::order_by::resume_filter_json`), with the +/// already-emitted prefix of the boundary tie run trimmed client-side. +/// +/// `skip_count` mirrors .NET's `OrderByContinuationToken.SkipCount`: a JOIN +/// (or array-unwind) query can emit several result rows from a single +/// document, all sharing the same `_rid` and the same sort key, so the +/// `(sort key, _rid)` cursor alone cannot say how many of that group were +/// already emitted. `skip_count` records exactly that — the number of +/// already-emitted rows sharing the boundary's `(resume_values, last_rid)` — +/// so resume skips precisely those duplicates and no more. A well-formed +/// boundary always carries `skip_count >= 1` (it counts at least the boundary +/// row itself); a token predating the field is read back as `1` +/// (`default_boundary_skip_count`) so its boundary row is still discarded and +/// never re-emitted. Tokens minted by earlier builds may also carry a +/// now-ignored `rows_emitted` field; serde skips it. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct ValueBoundary { + pub(crate) resume_values: Vec, + pub(crate) last_rid: String, + #[serde(default = "default_boundary_skip_count")] + pub(crate) skip_count: u32, +} + +/// A well-formed value boundary always emitted at least its boundary row, so a +/// token predating the `skip_count` field resumes as if that one row was +/// counted — the discard drops it rather than re-emitting it. +fn default_boundary_skip_count() -> u32 { + 1 } /// One entry in a [`PipelineNodeState::SequentialDrain`] `active_tokens` @@ -145,6 +221,7 @@ impl PipelineNodeState { PipelineNodeState::Request { .. } => "Request", PipelineNodeState::SequentialDrain { .. } => "SequentialDrain", PipelineNodeState::UnorderedMerge { .. } => "UnorderedMerge", + PipelineNodeState::StreamingOrderedMerge { .. } => "StreamingOrderedMerge", }, )) .build()), @@ -328,6 +405,225 @@ mod tests { ); } + #[test] + fn into_child_contribution_rejects_nested_streaming_ordered_merge() { + let err = PipelineNodeState::StreamingOrderedMerge { + directions: vec![SortOrder::Ascending], + ranges: vec![], + } + .into_child_contribution("Parent", 0, 1) + .expect_err("nested StreamingOrderedMerge is not a supported child shape"); + let msg = format!("{err:?}"); + assert!( + msg.contains("StreamingOrderedMerge"), + "error should name the offending variant: {msg}" + ); + } + + #[test] + fn streaming_ordered_merge_round_trips_untouched_range() { + let state = PipelineNodeState::StreamingOrderedMerge { + directions: vec![SortOrder::Ascending, SortOrder::Descending], + ranges: vec![OrderByRangeToken { + min_epk: "".to_owned(), + max_epk: "FF".to_owned(), + server_continuation: None, + boundary: None, + }], + }; + let json = serde_json::to_string(&state).unwrap(); + let parsed: PipelineNodeState = serde_json::from_str(&json).unwrap(); + assert_eq!(parsed, state); + } + + /// An older token carrying the retired `query_fingerprint` field must + /// still deserialize (serde ignores unknown fields), and a freshly + /// serialized token must omit it. + #[test] + fn streaming_ordered_merge_ignores_legacy_query_fingerprint_field() { + let json = r#"{ + "kind": "streaming_ordered_merge", + "directions": ["Ascending"], + "query_fingerprint": "deadbeef", + "ranges": [{ + "min_epk": "", + "max_epk": "FF" + }] + }"#; + let parsed: PipelineNodeState = serde_json::from_str(json).unwrap(); + assert_eq!( + parsed, + PipelineNodeState::StreamingOrderedMerge { + directions: vec![SortOrder::Ascending], + ranges: vec![OrderByRangeToken { + min_epk: "".to_owned(), + max_epk: "FF".to_owned(), + server_continuation: None, + boundary: None, + }], + } + ); + + let new_json = serde_json::to_string(&parsed).unwrap(); + assert!( + !new_json.contains("query_fingerprint"), + "new tokens must not emit the retired field: {new_json}" + ); + } + + #[test] + fn streaming_ordered_merge_round_trips_clean_page_boundary() { + let state = PipelineNodeState::StreamingOrderedMerge { + directions: vec![SortOrder::Ascending], + ranges: vec![OrderByRangeToken { + min_epk: "".to_owned(), + max_epk: "FF".to_owned(), + server_continuation: Some("backend-token".to_owned()), + boundary: Some(ValueBoundary { + resume_values: vec![OrderByResumeValue::Number { value: 5.0.into() }], + last_rid: "rid-1".to_owned(), + skip_count: 1, + }), + }], + }; + let json = serde_json::to_string(&state).unwrap(); + let parsed: PipelineNodeState = serde_json::from_str(&json).unwrap(); + assert_eq!(parsed, state); + } + + #[test] + fn streaming_ordered_merge_round_trips_complex_value_boundary() { + let state = PipelineNodeState::StreamingOrderedMerge { + directions: vec![SortOrder::Ascending], + ranges: vec![OrderByRangeToken { + min_epk: "40".to_owned(), + max_epk: "80".to_owned(), + server_continuation: None, + boundary: Some(ValueBoundary { + resume_values: vec![OrderByResumeValue::Complex { + complex_type: super::super::order_by::ComplexTypeTag::Array, + hash: super::super::order_by::ComplexHash { + low64: 1, + high64: 2, + }, + }], + last_rid: "rid-2".to_owned(), + skip_count: 1, + }), + }], + }; + let json = serde_json::to_string(&state).unwrap(); + let parsed: PipelineNodeState = serde_json::from_str(&json).unwrap(); + assert_eq!(parsed, state); + } + + /// A boundary's `skip_count` round-trips, and a token that predates the + /// field parses as `skip_count: 1` (via `default_boundary_skip_count`), so + /// an older continuation token never fails to deserialize and still drops + /// its boundary row on resume. + #[test] + fn streaming_ordered_merge_boundary_skip_count_round_trips_and_defaults() { + let with_skip = r#"{ + "kind": "streaming_ordered_merge", + "directions": ["Ascending"], + "ranges": [{ + "min_epk": "", + "max_epk": "FF", + "boundary": { + "resume_values": [{"type": "number", "value": 5.0}], + "last_rid": "rid-1", + "skip_count": 3 + } + }] + }"#; + let parsed: PipelineNodeState = serde_json::from_str(with_skip).unwrap(); + let PipelineNodeState::StreamingOrderedMerge { ranges, .. } = &parsed else { + panic!("expected StreamingOrderedMerge"); + }; + assert_eq!(ranges[0].boundary.as_ref().unwrap().skip_count, 3); + // Re-serializing preserves the value. + let round_tripped: PipelineNodeState = + serde_json::from_str(&serde_json::to_string(&parsed).unwrap()).unwrap(); + assert_eq!(round_tripped, parsed); + + // A token missing `skip_count` parses as 1 — a real boundary always + // emitted its boundary row, so resume still discards it. + let without_skip = r#"{ + "kind": "streaming_ordered_merge", + "directions": ["Ascending"], + "ranges": [{ + "min_epk": "", + "max_epk": "FF", + "boundary": { + "resume_values": [{"type": "number", "value": 5.0}], + "last_rid": "rid-1" + } + }] + }"#; + let parsed: PipelineNodeState = serde_json::from_str(without_skip).unwrap(); + let PipelineNodeState::StreamingOrderedMerge { ranges, .. } = &parsed else { + panic!("expected StreamingOrderedMerge"); + }; + assert_eq!(ranges[0].boundary.as_ref().unwrap().skip_count, 1); + } + + #[test] + fn streaming_ordered_merge_ignores_legacy_rows_emitted_field() { + let json = r#"{ + "kind": "streaming_ordered_merge", + "directions": ["Ascending"], + "ranges": [{ + "min_epk": "", + "max_epk": "80", + "boundary": { + "resume_values": [{"type": "number", "value": 5.0}], + "last_rid": "rid-3", + "rows_emitted": 40 + } + }] + }"#; + let parsed: PipelineNodeState = serde_json::from_str(json).unwrap(); + assert_eq!( + parsed, + PipelineNodeState::StreamingOrderedMerge { + directions: vec![SortOrder::Ascending], + ranges: vec![OrderByRangeToken { + min_epk: String::new(), + max_epk: "80".to_owned(), + server_continuation: None, + boundary: Some(ValueBoundary { + resume_values: vec![OrderByResumeValue::Number { value: 5.0.into() }], + last_rid: "rid-3".to_owned(), + skip_count: 1, + }), + }], + } + ); + let new_json = serde_json::to_string(&parsed).unwrap(); + assert!( + !new_json.contains("rows_emitted"), + "new tokens must not emit the retired field: {new_json}" + ); + } + + #[test] + fn streaming_ordered_merge_omits_absent_continuation_and_boundary() { + let state = PipelineNodeState::StreamingOrderedMerge { + directions: vec![SortOrder::Ascending], + ranges: vec![OrderByRangeToken { + min_epk: String::new(), + max_epk: "FF".to_owned(), + server_continuation: None, + boundary: None, + }], + }; + let json = serde_json::to_string(&state).unwrap(); + assert!( + !json.contains("server_continuation") && !json.contains("boundary"), + "absent fields must be omitted from the wire form: {json}" + ); + } + #[test] fn unordered_merge_round_trips_empty() { let state = PipelineNodeState::UnorderedMerge { diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/streaming_ordered_merge.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/streaming_ordered_merge.rs new file mode 100644 index 00000000000..f4cbd8497d0 --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/streaming_ordered_merge.rs @@ -0,0 +1,2398 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +// cspell:ignore unemitted unfetched rescan + +//! Streaming cross-partition `ORDER BY` k-way merge node. +//! +//! [`StreamingOrderedMerge`] polls one leaf [`Request`] per active EPK range, +//! buffers one locally-sorted row per range, and repeatedly emits the +//! globally smallest buffered row (per [`compare_key_tuples`]). With the +//! documented fan-out cap (`FEED_OPERATIONS_REQS.md` §3, default 100), a +//! linear per-pop scan over active children is simpler than a real heap and +//! not a meaningful cost next to per-page network I/O. +//! +//! # Resume model +//! +//! Unlike [`super::SequentialDrain`]'s sparse cursor, global ordering means +//! *any* range may still have unemitted rows, so every active range is +//! snapshotted explicitly (see +//! [`super::snapshot::PipelineNodeState::StreamingOrderedMerge`]). A range +//! resumes at a **clean page boundary** (buffer empty, topology unchanged, +//! plain query) by replaying its saved backend continuation, or at a +//! **value boundary** (mid-page, topology changed, or already +//! resume-filtered) by reissuing the query with the .NET-compatible +//! structured `resumeFilter` field (see +//! [`super::query_response::with_resume_filter`]) and discarding +//! already-emitted rows client-side (see +//! [`super::order_by::classify_row_vs_boundary`]). +//! +//! A backend continuation is opaque and bound to the exact request body +//! that minted it, so one captured for a resume-filtered request is never +//! persisted or reused (see [`ChildQueryShape`]). On a live split, +//! [`Request::split_for_topology_change`] forwards the split child's backend +//! continuation into every replacement leaf. The Cosmos contract (also relied +//! on by [`super::SequentialDrain`]) is that a parent partition's continuation +//! stays valid on each post-split child under EPK scoping, so a replacement +//! resumes *after* every row the split child already emitted. +//! [`StreamingOrderedMerge::handle_split`] therefore forwards only the split +//! child's `last_emitted` bookkeeping (for `skip_count` accumulation and +//! future snapshots) and installs *no* client-side discard on the first +//! replacement page — reapplying one would wrongly drop later JOIN rows that +//! share the boundary `(key, _rid)`. This holds for scalar and complex +//! boundaries alike. A **saved-token** resume across a split instead rebuilds +//! each range through the structured `resumeFilter` (see +//! [`build_value_boundary_child`]), whose backend `DistinctHash` seek handles +//! scalar and complex boundaries alike. A replacement that carries no usable +//! continuation yet inherits an emitted boundary (a resume-filtered range that +//! split before its first page, or a generic non-`Request` node) is rebuilt +//! via that boundary discard when its shape allows, else rejected with a typed +//! `CLIENT_STREAMING_MERGE_SPLIT_REPLACEMENT_INVALID` rather than guessing at +//! an unknown stream's position. + +use std::cmp::Ordering; +use std::collections::VecDeque; +use std::sync::Arc; + +use async_trait::async_trait; +use serde_json::value::RawValue; + +use crate::models::{CosmosOperation, FeedRange, MaxItemCountHint}; + +use super::order_by::{ + classify_row_vs_boundary, compare_key_tuples, compare_rids, OrderByItem, OrderByResumeValue, + RowVsBoundary, +}; +use super::query_plan::SortOrder; +use super::query_response::{self, PageAggregator}; +use super::snapshot::{OrderByRangeToken, ValueBoundary}; +use super::{ + intersect_feed_ranges, PageResult, PipelineContext, PipelineNode, PipelineNodeState, Request, + RequestTarget, ResolvedRange, +}; + +/// Default emitted-page size when no `max_item_count` hint is set +/// (matches Cosmos's default; independent of backend per-child page size). +const DEFAULT_MAX_ITEM_COUNT: usize = 100; + +/// Maximum consecutive split retries per child before giving up (mirrors +/// `SequentialDrain`/`UnorderedMerge`). +const MAX_SPLIT_RETRIES: usize = 10; + +/// A range's last-emitted key tuple (bounded resume-value form), RID, and a +/// `_rid`-tie `skip_count`. `skip_count` is the number of consecutive rows +/// emitted so far that share this exact `(resume_values, rid)` — a JOIN (or +/// array unwind) can emit several result rows from one document with the same +/// `_rid` and sort key, so the count is needed to skip precisely those +/// already-emitted duplicates on resume (mirrors .NET's +/// `OrderByContinuationToken.SkipCount`). It is always `>= 1` once any row is +/// recorded. +#[derive(Clone)] +struct LastEmitted { + resume_values: Vec, + rid: String, + skip_count: u32, +} + +/// What a child's next fetched page must discard before its rows become +/// visible to the merge, right after resuming a value boundary. +enum PendingDiscard { + /// Nothing to discard — a fresh start or a plain continuation resume. + None, + /// Discard every row at or before the resume boundary, applying the same + /// three-phase `(sort key, _rid, skip_count)` seek .NET's + /// `FilterNextAsync` uses: + /// + /// 1. **sort key** — keys strictly before the boundary (per + /// [`classify_row_vs_boundary`]) are dropped; a key strictly after + /// stops the discard. + /// 2. **`_rid`** — within a full-key tie, a row whose `_rid` is strictly + /// before `last_rid` in the first column's direction (numeric + /// document-ordinal order via [`compare_rids`], matching the backend) + /// is dropped; a `_rid` strictly after stops the discard. + /// 3. **`skip_count`** — for the exact `(boundary key, last_rid)` group, + /// drop exactly `skip_count` rows (the already-emitted JOIN + /// duplicates), then stop so the remaining rows of that group emit. + /// + /// A per-row predicate, so it stays correct across pages and split + /// sub-ranges even when the backend's structured `resumeFilter` already + /// trims most of the prefix. `skip_count` is decremented as matching rows + /// are dropped and persists across empty/partial pages until consumed. + ResumeBoundary { + resume_values: Vec, + last_rid: String, + skip_count: u32, + directions: Vec, + }, +} + +impl PendingDiscard { + fn apply(&mut self, rows: &mut VecDeque) { + match self { + PendingDiscard::None => {} + PendingDiscard::ResumeBoundary { + resume_values, + last_rid, + skip_count, + directions, + } => { + // The first sort column governs the `_rid` tie direction + // (matching the backend within a full-key tie run). + let rid_direction = directions.first().copied().unwrap_or(SortOrder::Ascending); + while let Some(front) = rows.front() { + let discard = + match classify_row_vs_boundary(&front.keys, resume_values, directions) { + // Phase 1: sorts strictly before the boundary key. + RowVsBoundary::Before => true, + // Phase 1: at/after the boundary key (or an + // indeterminate complex column) — nothing left to + // discard. + RowVsBoundary::AfterOrIndeterminate => false, + RowVsBoundary::Tie => { + match compare_rids(&front.rid, last_rid, rid_direction) { + // Phase 2: `_rid` strictly before the boundary. + Ordering::Less => true, + // Phase 2: `_rid` strictly after the boundary. + Ordering::Greater => false, + // Phase 3: exact `(key, _rid)` group — drop + // exactly `skip_count` already-emitted + // duplicates, then keep the rest. + Ordering::Equal => { + if *skip_count > 0 { + *skip_count -= 1; + true + } else { + false + } + } + } + } + }; + if !discard { + // Rows are sorted, so every later row is past the boundary too. + *self = PendingDiscard::None; + return; + } + rows.pop_front(); + } + // Boundary not reached this page (empty page, or a tie/skip run + // spanning pages): keep the discard active — with any remaining + // `skip_count` — for the next page. + } + } + } +} + +/// Which request body a child's [`Request`] is executing, so +/// [`StreamingOrderedMerge::snapshot_state`] knows whether its backend +/// continuation (opaque, bound to the exact request body) is portable to +/// the plain-continuation resume path. +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +enum ChildQueryShape { + /// Runs the plain query body (no `resumeFilter`); continuation may be + /// snapshotted. + Plain, + /// Runs the query body with a structured `resumeFilter` injected; + /// continuation is bound to that body and must never be stored — resumes + /// from its value boundary instead. + ResumeFilterInjected, +} + +/// One still-active child stream: an owned EPK range, a leaf [`Request`] +/// executing the query for that range, and buffered rows from its current +/// backend page. `pub(super)` only so `planner` can hold the opaque +/// `Vec` [`build_children`] returns. +pub(super) struct ChildStream { + range: FeedRange, + node: Box, + buffered: VecDeque, + /// `true` once a terminal (no-continuation) page was observed. + drained: bool, + last_emitted: Option, + pending_discard: PendingDiscard, + /// Which request body `node` is executing (see [`ChildQueryShape`]). + query_shape: ChildQueryShape, +} + +impl ChildStream { + fn fresh(range: FeedRange, node: Box) -> Self { + Self { + range, + node, + buffered: VecDeque::new(), + drained: false, + last_emitted: None, + pending_discard: PendingDiscard::None, + // Overridden to `ResumeFilterInjected` only by + // `build_value_boundary_child`. + query_shape: ChildQueryShape::Plain, + } + } + + /// Records the just-emitted row as the new resume boundary, tracking the + /// `_rid`-tie `skip_count`: consecutive rows sharing the exact + /// `(resume_values, rid)` increment the count (a JOIN duplicate of the + /// same document), and any change to the key tuple or `_rid` resets it to + /// `1`. Returns a typed error only on the (practically unreachable) u32 + /// overflow, so a boundary is never silently truncated. + fn record_emission(&mut self, keys: &[OrderByItem], rid: &str) -> crate::error::Result<()> { + let resume_values: Vec = + keys.iter().map(OrderByItem::to_resume_value).collect(); + let skip_count = match &self.last_emitted { + Some(last) if last.rid == rid && last.resume_values == resume_values => { + last.skip_count.checked_add(1).ok_or_else(|| { + crate::error::CosmosError::builder() + .with_status( + crate::error::CosmosStatus::CLIENT_CONTINUATION_TOKEN_ORDER_BY_STATE_INVALID, + ) + .with_message( + "ORDER BY resume skip_count overflowed u32 (a single document \ + produced more than u32::MAX result rows sharing one _rid)", + ) + .build() + })? + } + // First row overall, or a new key/_rid group: this row is the + // first (and so far only) emission of its `(key, _rid)` group. + _ => 1, + }; + self.last_emitted = Some(LastEmitted { + resume_values, + rid: rid.to_owned(), + skip_count, + }); + Ok(()) + } + + fn boundary(&self) -> Option { + self.last_emitted.as_ref().map(|le| ValueBoundary { + resume_values: le.resume_values.clone(), + last_rid: le.rid.clone(), + skip_count: le.skip_count, + }) + } +} + +/// Streams cross-partition `ORDER BY` results in globally sorted order +/// (see module docs for the merge algorithm and resume model). +pub(crate) struct StreamingOrderedMerge { + /// The per-range operation with no `resumeFilter` (its `query` text + /// already placeholder-substituted with `true`) — the base every + /// child's `Request` starts from. A resumed range injects a structured + /// `resumeFilter` into a clone of this body. + plain_operation: Arc, + directions: Vec, + children: Vec, +} + +impl StreamingOrderedMerge { + pub(super) fn new( + plain_operation: Arc, + directions: Vec, + children: Vec, + ) -> Self { + Self { + plain_operation, + directions, + children, + } + } + + fn max_item_count(&self) -> usize { + match self.plain_operation.request_headers().max_item_count { + Some(MaxItemCountHint::Limit(n)) => n.get() as usize, + Some(MaxItemCountHint::ServerDecides) | None => DEFAULT_MAX_ITEM_COUNT, + } + } + + /// Ensures the child at `idx` has a buffered row or is drained, + /// fetching/splitting as needed. Absorbs pages into `aggregator`. + async fn prime_child( + &mut self, + idx: usize, + context: &mut PipelineContext<'_>, + aggregator: &mut PageAggregator, + ) -> crate::error::Result<()> { + let mut split_retries = 0; + loop { + if !self.children[idx].buffered.is_empty() || self.children[idx].drained { + return Ok(()); + } + + match self.children[idx].node.next_page(context).await? { + PageResult::Page { + response, + is_terminal, + } => { + aggregator.absorb(&response); + let mut rows: VecDeque = + query_response::parse_envelope_page( + response.body(), + self.directions.len(), + )? + .into(); + self.children[idx].pending_discard.apply(&mut rows); + self.children[idx].buffered = rows; + if is_terminal { + self.children[idx].drained = true; + } + if !self.children[idx].buffered.is_empty() || self.children[idx].drained { + return Ok(()); + } + // Empty page with a continuation pending is not drained; re-poll. + } + PageResult::Drained => { + self.children[idx].drained = true; + return Ok(()); + } + PageResult::SplitRequired { replacement_nodes } => { + split_retries += 1; + if split_retries > MAX_SPLIT_RETRIES { + return Err(crate::error::CosmosError::builder() + .with_status(crate::error::CosmosStatus::CLIENT_SPLIT_RETRIES_EXHAUSTED) + .with_message(format!( + "exceeded maximum split retries ({MAX_SPLIT_RETRIES}) \ + in StreamingOrderedMerge" + )) + .build()); + } + self.handle_split(idx, replacement_nodes)?; + // Loop: index `idx` now refers to the first replacement. + } + } + } + } + + /// Primes every currently-active child, restoring the merge invariant + /// that before each [`select_min_child_index`] every non-drained child + /// has a buffered head row (or is proven drained) — so no child is ever + /// skipped for lacking a head. Re-reads `len()` each step because + /// [`prime_child`]'s split handling can splice several replacements in at + /// once (only the first is primed inline); already-buffered/drained + /// children short-circuit, so no page is re-fetched. + async fn prime_all_active_children( + &mut self, + context: &mut PipelineContext<'_>, + aggregator: &mut PageAggregator, + ) -> crate::error::Result<()> { + let mut idx = 0; + while idx < self.children.len() { + self.prime_child(idx, context, aggregator).await?; + idx += 1; + } + Ok(()) + } + + /// Handles a child's `SplitRequired` by consuming the `replacement_nodes` + /// the child produced, keeping the merge ignorant of their concrete type. + /// + /// The split child provides one replacement leaf per post-split sub-range; + /// this validates each carries a [`feed_range`](PipelineNode::feed_range), + /// sorts them, and confirms they exactly tile the split child's scope with + /// no gaps or overlaps, then wraps each in a [`ChildStream`] that inherits + /// the split child's resume state (see [`wrap_split_replacement`]). No + /// topology is re-resolved here — the split child already did that when + /// producing the nodes, forwarding its backend continuation into each + /// replacement so they resume past every already-emitted row. + fn handle_split( + &mut self, + idx: usize, + replacement_nodes: Vec>, + ) -> crate::error::Result<()> { + let scope = self.children[idx].range.clone(); + let prior_boundary = self.children[idx].boundary(); + let query_shape = self.children[idx].query_shape; + + // Every replacement must own an EPK sub-range; a missing one would + // make coverage unverifiable, so reject rather than risk a gap. + let mut replacements: Vec<(FeedRange, Box)> = + Vec::with_capacity(replacement_nodes.len()); + for node in replacement_nodes { + let range = node + .feed_range() + .ok_or_else(|| { + split_replacement_invalid( + "StreamingOrderedMerge split replacement node has no feed_range", + ) + })? + .clone(); + replacements.push((range, node)); + } + replacements.sort_by(|a, b| a.0.min_inclusive().cmp(b.0.min_inclusive())); + validate_exact_coverage(&scope, replacements.iter().map(|(range, _)| range))?; + + // Wrap each replacement before mutating `self.children`, so a rejected + // replacement leaves the merge unchanged rather than half-spliced. + let mut wrapped = Vec::with_capacity(replacements.len()); + for (range, node) in replacements { + wrapped.push(wrap_split_replacement( + range, + node, + prior_boundary.as_ref(), + query_shape, + &self.directions, + )?); + } + + self.children.remove(idx); + for (i, child) in wrapped.into_iter().enumerate() { + self.children.insert(idx + i, child); + } + Ok(()) + } + + /// Index of the active child whose buffered head row compares smallest + /// (per [`compare_key_tuples`], tie-broken by the direction-aware + /// numeric `_rid` then range identity), or `None` if no child has a + /// buffered row. + fn select_min_child_index(&self) -> Option { + let mut best: Option = None; + for idx in 0..self.children.len() { + if self.children[idx].buffered.front().is_none() { + continue; + } + best = Some(match best { + None => idx, + Some(best_idx) => { + if self.row_less_than(idx, best_idx) { + idx + } else { + best_idx + } + } + }); + } + best + } + + fn row_less_than(&self, a_idx: usize, b_idx: usize) -> bool { + let a = self.children[a_idx] + .buffered + .front() + .expect("caller only compares children with a buffered row"); + let b = self.children[b_idx] + .buffered + .front() + .expect("caller only compares children with a buffered row"); + // The full-key tie-break must match the backend's per-partition + // order — numeric `_rid` in the first sort column's direction — so a + // resumed range's scalar boundary is a clean cut point across a split. + let rid_direction = self + .directions + .first() + .copied() + .unwrap_or(SortOrder::Ascending); + let ordering = compare_key_tuples(&a.keys, &b.keys, &self.directions) + .then_with(|| compare_rids(&a.rid, &b.rid, rid_direction)) + .then_with(|| { + self.children[a_idx] + .range + .min_inclusive() + .cmp(self.children[b_idx].range.min_inclusive()) + }); + ordering == std::cmp::Ordering::Less + } +} + +#[async_trait] +impl PipelineNode for StreamingOrderedMerge { + async fn next_page( + &mut self, + context: &mut PipelineContext<'_>, + ) -> crate::error::Result { + if self.children.is_empty() { + return Ok(PageResult::Drained); + } + + let mut aggregator = PageAggregator::new(); + + // Prime every child up front so the first `select_min_child_index` + // sees a head row for each non-drained child (see + // `prime_all_active_children`). + self.prime_all_active_children(context, &mut aggregator) + .await?; + + let cap = self.max_item_count(); + let mut payloads: Vec> = Vec::new(); + + while payloads.len() < cap { + let Some(winner) = self.select_min_child_index() else { + break; + }; + let row = self.children[winner] + .buffered + .pop_front() + .expect("select_min_child_index only returns indices with a buffered row"); + self.children[winner].record_emission(&row.keys, &row.rid)?; + payloads.push(row.payload); + // Re-prime before the next selection only if there's still room, + // so unread rows stay unfetched until the next `next_page` call. + // Priming *all* active children (not just `winner`) is required: + // replenishing `winner` may split it into several replacements, + // and every one must have a head row before the next selection or + // a later replacement's smaller rows would be skipped. + if payloads.len() < cap { + self.prime_all_active_children(context, &mut aggregator) + .await?; + } + } + + // Evict fully-drained empty children, mirroring `SequentialDrain`, + // so a later snapshot never references them. + self.children + .retain(|child| !(child.drained && child.buffered.is_empty())); + let is_terminal = self.children.is_empty(); + + let response = aggregator.build_page(&payloads)?; + Ok(PageResult::Page { + response, + is_terminal, + }) + } + + #[cfg(test)] + fn into_children(self) -> Vec> { + self.children.into_iter().map(|c| c.node).collect() + } + + fn snapshot_state(&self) -> crate::error::Result { + if self.children.is_empty() { + return Ok(PipelineNodeState::Drained); + } + + let mut ranges = Vec::with_capacity(self.children.len()); + for (idx, child) in self.children.iter().enumerate() { + // Safe to snapshot the backend continuation into the plain + // `server_continuation` field only when nothing is buffered + // (else it points past unemitted rows) and the child runs the + // plain query (`ChildQueryShape::Plain`) — a resume-filtered + // child's continuation is bound to that filtered text and would + // mismatch the plain query on resume; it resumes from its + // scalar `boundary` instead. + let server_continuation = if child.buffered.is_empty() + && child.query_shape == ChildQueryShape::Plain + { + match child.node.snapshot_state()? { + PipelineNodeState::Request { + server_continuation, + } => server_continuation, + PipelineNodeState::Drained => None, + other => { + return Err(crate::error::CosmosError::builder() + .with_status( + crate::error::CosmosStatus::CLIENT_CONTINUATION_TOKEN_UNEXPECTED_NESTED_SHAPE, + ) + .with_message(format!( + "StreamingOrderedMerge child {idx} of {total} produced an \ + unsupported snapshot shape: {other:?}", + total = self.children.len(), + )) + .build()); + } + } + } else { + None + }; + + ranges.push(OrderByRangeToken { + min_epk: child.range.min_inclusive().to_hex(), + max_epk: child.range.max_exclusive().to_hex(), + server_continuation, + boundary: child.boundary(), + }); + } + + Ok(PipelineNodeState::StreamingOrderedMerge { + directions: self.directions.clone(), + ranges, + }) + } + + fn topology_can_change(&self) -> bool { + // Splits are handled internally (`handle_split`); no parent needed. + false + } +} + +/// Builds the child streams needed to cover `scope`, given its topology +/// resolution and prior resume state (`None` for a fresh range). Shared +/// by planner construction and the live split handler. `directions` seed +/// the client-side discard's boundary comparison. +/// +/// - Unchanged topology + `prior_continuation`: replays it as-is (safe — +/// see `snapshot_state` / [`ChildQueryShape`]). +/// - Else + `prior_boundary`: rebuilds via `build_value_boundary_child`, +/// which sends the boundary as a structured `resumeFilter`. This is a +/// per-row seek, so it stays correct across a split for scalar and +/// complex boundaries alike. +/// - Else: a fresh unfiltered start. On a topology change, resolved +/// ranges are clipped to `scope` and must exactly tile it. +pub(super) fn build_children( + resolved: &[ResolvedRange], + scope: &FeedRange, + plain_operation: &Arc, + directions: &[SortOrder], + prior_continuation: Option, + prior_boundary: Option<&ValueBoundary>, +) -> crate::error::Result> { + let unchanged = resolved.len() == 1 + && resolved[0].range.min_inclusive() == scope.min_inclusive() + && resolved[0].range.max_exclusive() == scope.max_exclusive(); + + if unchanged { + let resolved_range = &resolved[0]; + let target = RequestTarget::effective_partition_key_range( + scope.clone(), + resolved_range.partition_key_range_id.clone(), + resolved_range.range.clone(), + ); + let child = if let Some(continuation) = prior_continuation { + // Safe to replay as-is (see doc comment above); carry the + // last-emitted boundary forward for the next snapshot. + let mut child = ChildStream::fresh( + scope.clone(), + Box::new(Request::new( + Arc::clone(plain_operation), + target, + Some(continuation), + )), + ); + if let Some(boundary) = prior_boundary { + child.last_emitted = Some(LastEmitted { + resume_values: boundary.resume_values.clone(), + rid: boundary.last_rid.clone(), + skip_count: boundary.skip_count, + }); + } + child + } else if let Some(boundary) = prior_boundary { + // Rows were emitted but no safe continuation was saved (a + // mid-page checkpoint); a plain restart would re-emit them. + build_value_boundary_child( + scope.clone(), + target, + plain_operation, + directions, + boundary, + )? + } else { + // Genuinely fresh: no continuation, no rows ever emitted. + ChildStream::fresh( + scope.clone(), + Box::new(Request::new(Arc::clone(plain_operation), target, None)), + ) + }; + return Ok(vec![child]); + } + + // Topology changed (split or merge). Clip resolved ranges to `scope` + // before validating coverage — a merge yields a range wider than + // `scope`, and the raw bounds would spuriously fail. Mirrors + // `SequentialDrain`'s resume path via `intersect_feed_ranges`. + let mut clipped: Vec<(FeedRange, &ResolvedRange)> = Vec::with_capacity(resolved.len()); + for resolved_range in resolved { + if let Some(owned) = intersect_feed_ranges(&resolved_range.range, scope) { + clipped.push((owned, resolved_range)); + } + // Non-overlapping ranges contribute nothing; a coverage gap fails below. + } + clipped.sort_by(|a, b| a.0.min_inclusive().cmp(b.0.min_inclusive())); + validate_exact_coverage(scope, clipped.iter().map(|(range, _)| range))?; + + // Every sub-range resumes from the same last-emitted boundary via a + // structured `resumeFilter` (a per-row seek), so a split needs no + // per-range row-count attribution — scalar and complex boundaries alike + // are split-safe. + let mut children = Vec::with_capacity(clipped.len()); + for (owned, resolved_range) in clipped { + let target = RequestTarget::effective_partition_key_range( + owned.clone(), + resolved_range.partition_key_range_id.clone(), + resolved_range.range.clone(), + ); + let child = match prior_boundary { + None => ChildStream::fresh( + owned, + Box::new(Request::new(Arc::clone(plain_operation), target, None)), + ), + Some(boundary) => { + build_value_boundary_child(owned, target, plain_operation, directions, boundary)? + } + }; + children.push(child); + } + Ok(children) +} + +/// Returns `Ok(())` if `ranges` (yielded in ascending `min_inclusive` order, +/// each already clipped to `scope`) exactly tiles `scope` end-to-end with no +/// gaps or overlaps. Shared by the planner's split/merge resume path +/// ([`build_children`]) and the live split handler +/// ([`StreamingOrderedMerge::handle_split`]). +fn validate_exact_coverage<'a>( + scope: &FeedRange, + ranges: impl Iterator, +) -> crate::error::Result<()> { + let mut cursor = scope.min_inclusive().clone(); + let mut saw_any = false; + for owned in ranges { + saw_any = true; + if owned.min_inclusive() != &cursor { + return Err(split_replacement_invalid(format!( + "replacement range [{}, {}) does not start at the expected cursor {}", + owned.min_inclusive().to_hex(), + owned.max_exclusive().to_hex(), + cursor.to_hex(), + ))); + } + cursor = owned.max_exclusive().clone(); + } + if !saw_any { + return Err(split_replacement_invalid( + "topology resolution produced no replacement ranges", + )); + } + if &cursor != scope.max_exclusive() { + return Err(split_replacement_invalid(format!( + "replacement ranges cover up to {} but the prior range extended to {}", + cursor.to_hex(), + scope.max_exclusive().to_hex(), + ))); + } + Ok(()) +} + +/// Wraps one split-replacement leaf into a [`ChildStream`] carrying the split +/// child's resume state. `query_shape` is copied from the split child so a +/// resume-filtered range's continuation is never persisted or reused. +/// +/// `prior_boundary` is the split child's current boundary (`None` if it never +/// emitted a row): +/// +/// - `None`: a fresh start — no rows emitted before the split, nothing to +/// discard. +/// - `Some` + the leaf reports a `Request` with a forwarded backend +/// continuation: trust it. [`Request::split_for_topology_change`] carried the +/// split child's continuation into the replacement, so it already resumes +/// *after* every emitted row. Only `last_emitted` is forwarded (for +/// `skip_count` accumulation and future snapshots) — installing a discard +/// here would wrongly drop later JOIN rows sharing the boundary +/// `(key, _rid)`. +/// - `Some` + no forwarded continuation + a resume-filtered range: rebuild the +/// [`PendingDiscard::ResumeBoundary`] from the boundary, exactly like a +/// saved-token resume. Such a leaf re-runs the structured `resumeFilter` +/// (the backend seeks to the boundary), so the discard only trims the +/// already-emitted prefix — safe for scalar and complex boundaries. +/// - `Some` + no forwarded continuation + a plain range (or generic node): +/// the leaf's resume position is unknown (a plain replay would re-fetch from +/// the start, which the client discard can't order for a complex boundary), +/// so reject with a typed error rather than guess. For a real plain +/// `Request` this is unreachable — a plain child that had emitted a row is +/// always `Continuing` when it splits, so its replacement always carries a +/// forwarded continuation. +fn wrap_split_replacement( + range: FeedRange, + node: Box, + prior_boundary: Option<&ValueBoundary>, + query_shape: ChildQueryShape, + directions: &[SortOrder], +) -> crate::error::Result { + let mut child = ChildStream::fresh(range, node); + child.query_shape = query_shape; + + let Some(boundary) = prior_boundary else { + // Fresh start: no rows emitted before the split, nothing to discard. + return Ok(child); + }; + + let forwarded_continuation = matches!( + child.node.snapshot_state()?, + PipelineNodeState::Request { + server_continuation: Some(_), + } + ); + if forwarded_continuation { + // The forwarded continuation positions the replacement past every + // emitted row; carry `last_emitted` only (no client discard). + child.last_emitted = Some(LastEmitted { + resume_values: boundary.resume_values.clone(), + rid: boundary.last_rid.clone(), + skip_count: boundary.skip_count, + }); + return Ok(child); + } + + // No forwarded continuation but the split child had emitted rows. A + // resume-filtered replacement re-seeks the backend to the boundary, so + // rebuild the discard to trim the already-emitted prefix (scalar or + // complex). A plain (or generic) replacement's position is unknown; reject. + if query_shape != ChildQueryShape::ResumeFilterInjected { + return Err(split_replacement_invalid( + "StreamingOrderedMerge split replacement carries no continuation to \ + resume a mid-group boundary and cannot be repositioned", + )); + } + child.last_emitted = Some(LastEmitted { + resume_values: boundary.resume_values.clone(), + rid: boundary.last_rid.clone(), + skip_count: boundary.skip_count, + }); + child.pending_discard = PendingDiscard::ResumeBoundary { + resume_values: boundary.resume_values.clone(), + last_rid: boundary.last_rid.clone(), + skip_count: boundary.skip_count, + directions: directions.to_vec(), + }; + Ok(child) +} + +fn split_replacement_invalid( + message: impl Into>, +) -> crate::error::CosmosError { + crate::error::CosmosError::builder() + .with_status(crate::error::CosmosStatus::CLIENT_STREAMING_MERGE_SPLIT_REPLACEMENT_INVALID) + .with_message(message) + .build() +} + +/// Builds one child via the value-boundary resume path: sends the +/// last-emitted boundary to the backend as a structured `resumeFilter` +/// (`rid` present, `exclude:false`) injected into a clone of the plain +/// query body, and installs a matching [`PendingDiscard::ResumeBoundary`] +/// guard so the already-emitted prefix of the boundary tie run is trimmed +/// client-side. Works for scalar and complex (array/object) boundaries +/// alike — the backend seek is a per-row predicate, so it stays correct +/// across a split. +fn build_value_boundary_child( + owned_range: FeedRange, + target: RequestTarget, + plain_operation: &Arc, + directions: &[SortOrder], + boundary: &ValueBoundary, +) -> crate::error::Result { + let body = query_response::with_resume_filter( + plain_operation.body(), + &boundary.resume_values, + Some(&boundary.last_rid), + false, + )?; + let operation = Arc::new((*Arc::clone(plain_operation)).clone().with_body(body)); + + let mut child = + ChildStream::fresh(owned_range, Box::new(Request::new(operation, target, None))); + child.last_emitted = Some(LastEmitted { + resume_values: boundary.resume_values.clone(), + rid: boundary.last_rid.clone(), + skip_count: boundary.skip_count, + }); + child.pending_discard = PendingDiscard::ResumeBoundary { + resume_values: boundary.resume_values.clone(), + last_rid: boundary.last_rid.clone(), + skip_count: boundary.skip_count, + directions: directions.to_vec(), + }; + // The continuation is bound to the resume-filtered body; never snapshot it. + child.query_shape = ChildQueryShape::ResumeFilterInjected; + Ok(child) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::driver::dataflow::mocks::{self, MockLeaf}; + use crate::models::effective_partition_key::EffectivePartitionKey; + + fn range(min: &str, max: &str) -> FeedRange { + FeedRange::new( + EffectivePartitionKey::from(min), + EffectivePartitionKey::from(max), + ) + .unwrap() + } + + /// Builds a rewritten-envelope backend `CosmosResponse` with one row per + /// `(rid, rank)` pair (for a `MockRequestExecutor` reply). + fn envelope_response( + rows: &[(&str, i64)], + continuation: Option<&str>, + ) -> crate::models::CosmosResponse { + let documents: Vec = rows + .iter() + .map(|(rid, rank)| { + serde_json::json!({ + "_rid": rid, + "orderByItems": [{"item": rank}], + "payload": {"id": rid, "rank": rank}, + }) + }) + .collect(); + let body = + serde_json::json!({"_rid": "", "Documents": documents, "_count": documents.len()}); + mocks::response_with_continuation(&serde_json::to_vec(&body).unwrap(), continuation) + } + + /// Wraps [`envelope_response`] in a `PageResult` for a `MockLeaf` page. + fn envelope_page( + rows: &[(&str, i64)], + continuation: Option<&str>, + ) -> crate::error::Result { + Ok(PageResult::Page { + response: envelope_response(rows, continuation), + is_terminal: continuation.is_none(), + }) + } + + /// A single-column complex (array) envelope backend `CosmosResponse`: + /// each row's sort key is the array `[value]`. Shared by + /// [`array_envelope_page`] (MockLeaf pages) and `MockRequestExecutor` + /// replies for the resume-filtered split test. + fn array_envelope_response( + rows: &[(&str, i64)], + continuation: Option<&str>, + ) -> crate::models::CosmosResponse { + let documents: Vec = rows + .iter() + .map(|(rid, value)| { + serde_json::json!({ + "_rid": rid, + "orderByItems": [{"item": [value]}], + "payload": {"id": rid}, + }) + }) + .collect(); + let body = + serde_json::json!({"_rid": "", "Documents": documents, "_count": documents.len()}); + mocks::response_with_continuation(&serde_json::to_vec(&body).unwrap(), continuation) + } + + /// A single-column complex (array) envelope page: each row's sort key is + /// the array `[value]`, exercising the hash-based boundary discard. + fn array_envelope_page( + rows: &[(&str, i64)], + continuation: Option<&str>, + ) -> crate::error::Result { + Ok(PageResult::Page { + response: array_envelope_response(rows, continuation), + is_terminal: continuation.is_none(), + }) + } + + /// A JOIN-shaped envelope backend `CosmosResponse`: `(rid, rank, id)` rows + /// where several rows can share one `_rid` (a single document expanded by a + /// JOIN) while carrying distinct payload `id`s. Shared by + /// [`join_envelope_page`] (MockLeaf pages) and `MockRequestExecutor` replies + /// for the live-split forwarded-continuation test. + fn join_envelope_response( + rows: &[(&str, i64, &str)], + continuation: Option<&str>, + ) -> crate::models::CosmosResponse { + let documents: Vec = rows + .iter() + .map(|(rid, rank, id)| { + serde_json::json!({ + "_rid": rid, + "orderByItems": [{"item": rank}], + "payload": {"id": id}, + }) + }) + .collect(); + let body = + serde_json::json!({"_rid": "", "Documents": documents, "_count": documents.len()}); + mocks::response_with_continuation(&serde_json::to_vec(&body).unwrap(), continuation) + } + + /// A JOIN-shaped envelope page: `(rid, rank, id)` rows where several rows + /// can share one `_rid` (a single document expanded by a JOIN) while + /// carrying distinct payload `id`s. Exercises the `skip_count` phase of the + /// resume discard, which a local emulator cannot produce since it does not + /// execute JOINs. + fn join_envelope_page( + rows: &[(&str, i64, &str)], + continuation: Option<&str>, + ) -> crate::error::Result { + Ok(PageResult::Page { + response: join_envelope_response(rows, continuation), + is_terminal: continuation.is_none(), + }) + } + + fn ids(response: &crate::models::CosmosResponse) -> Vec { + let value: serde_json::Value = serde_json::from_slice(response.body_bytes()).unwrap(); + value["Documents"] + .as_array() + .unwrap() + .iter() + .map(|item| item["id"].as_str().unwrap().to_owned()) + .collect() + } + + fn merge(children: Vec, directions: Vec) -> StreamingOrderedMerge { + StreamingOrderedMerge::new(Arc::new(mocks::operation()), directions, children) + } + + async fn next_page(node: &mut StreamingOrderedMerge) -> PageResult { + let mut executor = mocks::NoopRequestExecutor; + let mut topology = mocks::NoopTopologyProvider; + let mut context = PipelineContext::new(&mut executor, Some(&mut topology)); + node.next_page(&mut context).await.unwrap() + } + + /// A `MockLeaf` split page carrying `replacement_nodes` — the leaves the + /// merge splices in and wraps with the split child's resume state. Using + /// non-`Request` (`MockLeaf`) replacements proves the merge consumes the + /// supplied nodes directly, staying ignorant of their concrete type. + fn split_page( + replacement_nodes: Vec>, + ) -> crate::error::Result { + Ok(PageResult::SplitRequired { replacement_nodes }) + } + + /// A `MockLeaf` split-replacement leaf covering `[min, max)` with pre-set + /// pages, reported through [`PipelineNode::feed_range`] so the merge can + /// verify the replacements tile the split child's scope. Its default + /// snapshot is `Drained` (no forwarded continuation), so it models a + /// *generic* replacement whose resume position the merge cannot trust — + /// see [`positioned_replacement_leaf`] for the common live-split case. + fn replacement_leaf( + min: &str, + max: &str, + pages: Vec>, + ) -> Box { + Box::new(MockLeaf::with_pages(pages).with_feed_range(range(min, max))) + } + + /// A split-replacement leaf that reports a `Request` snapshot carrying a + /// forwarded backend continuation, modeling what + /// [`Request::split_for_topology_change`] hands the merge on a real live + /// split. `handle_split` reads this snapshot to confirm the replacement + /// already resumes past every emitted row, so it installs no client-side + /// discard. + fn positioned_replacement_leaf( + min: &str, + max: &str, + pages: Vec>, + ) -> Box { + Box::new( + MockLeaf::with_pages(pages) + .with_feed_range(range(min, max)) + .with_snapshot(PipelineNodeState::Request { + server_continuation: Some("split-forwarded-ct".to_owned()), + }), + ) + } + + /// Drains `node` to completion against a no-op executor/topology (a merge + /// handles splits internally by consuming replacement nodes, so no live + /// topology resolution is needed), returning every emitted id in order. + async fn drain_all_ids(node: &mut StreamingOrderedMerge) -> Vec { + let mut all = Vec::new(); + loop { + match next_page(node).await { + PageResult::Page { + response, + is_terminal, + } => { + all.extend(ids(&response)); + if is_terminal { + break; + } + } + PageResult::Drained => break, + PageResult::SplitRequired { .. } => { + panic!("merge must handle splits internally, never surface SplitRequired") + } + } + } + all + } + + /// Regression for the in-flight split ordering defect: when replenishing + /// the popped winner fans it out into several sub-ranges, *every* + /// replacement (not just the first) must be primed before the next + /// selection. P0 emits `[1, 2]` then splits into P0a (next `3`) and P0b + /// (next `10, 20`); if P0b were left unprimed, `50` would be emitted + /// before `10, 20`. A large page cap keeps popping within one page. The + /// replacements carry a forwarded continuation (positioned past the `2` + /// boundary), so the merge wraps and orders them without re-resolving and + /// with no client-side discard. + #[tokio::test] + async fn split_during_pop_loop_primes_all_split_replacements() { + let p0 = ChildStream::fresh( + range("", "80"), + Box::new(MockLeaf::with_pages(vec![ + envelope_page(&[("d1", 1), ("d2", 2)], Some("p0-ct")), + // The split child yields two sub-range leaves that carry the + // forwarded continuation, so their pages already start past the + // `2` boundary. + split_page(vec![ + positioned_replacement_leaf("", "40", vec![envelope_page(&[("d3", 3)], None)]), + positioned_replacement_leaf( + "40", + "80", + vec![envelope_page(&[("d10", 10), ("d20", 20)], None)], + ), + ]), + ])), + ); + let p1 = ChildStream::fresh( + range("80", "FF"), + Box::new(MockLeaf::with_pages(vec![envelope_page( + &[("d50", 50)], + None, + )])), + ); + let mut node = merge(vec![p0, p1], vec![SortOrder::Ascending]); + + let emitted = drain_all_ids(&mut node).await; + assert_eq!( + emitted, + vec!["d1", "d2", "d3", "d10", "d20", "d50"] + .into_iter() + .map(str::to_owned) + .collect::>(), + "the second split replacement's smaller rows (10, 20) must precede 50" + ); + } + + /// Companion to the large-cap regression: with a page cap that fills + /// exactly as the split completes, the merge must checkpoint immediately + /// after the split — snapshotting *both* replacements — then resume in + /// order on the next page. P0 emits `[1, 2]`, splits, emits `3` + /// (cap = 3 reached), and P0b (`10, 20`) plus P1 (`50`) follow next page. + #[tokio::test] + async fn split_during_pop_loop_small_cap_checkpoints_after_split() { + let p0 = ChildStream::fresh( + range("", "80"), + Box::new(MockLeaf::with_pages(vec![ + envelope_page(&[("d1", 1), ("d2", 2)], Some("p0-ct")), + split_page(vec![ + positioned_replacement_leaf("", "40", vec![envelope_page(&[("d3", 3)], None)]), + positioned_replacement_leaf( + "40", + "80", + vec![envelope_page(&[("d10", 10), ("d20", 20)], None)], + ), + ]), + ])), + ); + let p1 = ChildStream::fresh( + range("80", "FF"), + Box::new(MockLeaf::with_pages(vec![envelope_page( + &[("d50", 50)], + None, + )])), + ); + let mut node = merge(vec![p0, p1], vec![SortOrder::Ascending]); + node.plain_operation = Arc::new((*node.plain_operation).clone().with_max_item_count( + crate::models::MaxItemCountHint::Limit(std::num::NonZeroU32::new(3).unwrap()), + )); + + // Page 1 stops right after the split, at the cap. + let page1 = next_page(&mut node).await; + let PageResult::Page { + response: r1, + is_terminal: t1, + } = page1 + else { + panic!("expected a page"); + }; + assert_eq!( + ids(&r1), + vec!["d1".to_owned(), "d2".to_owned(), "d3".to_owned()] + ); + assert!(!t1, "P0b and P1 rows remain buffered for the next page"); + + // The checkpoint must retain both post-split children (P0b + P1), + // with P0b carrying the forwarded resume boundary for its snapshot. + match node.snapshot_state().unwrap() { + PipelineNodeState::StreamingOrderedMerge { ranges, .. } => { + assert_eq!( + ranges.len(), + 2, + "both post-split children must be snapshotted" + ); + assert!( + ranges[0].boundary.is_some(), + "the surviving split replacement must resume from its value boundary" + ); + } + other => panic!("expected StreamingOrderedMerge snapshot, got {other:?}"), + } + + // Page 2 drains the rest in order. + let page2 = next_page(&mut node).await; + let PageResult::Page { + response: r2, + is_terminal: t2, + } = page2 + else { + panic!("expected a page"); + }; + assert_eq!( + ids(&r2), + vec!["d10".to_owned(), "d20".to_owned(), "d50".to_owned()] + ); + assert!(t2, "all children drained after the second page"); + } + + /// A split replacement may itself split before it yields a row (cascading + /// split). Each level supplies its own replacement leaves, each carrying a + /// forwarded continuation, so ordering and resume correctness hold no + /// matter how deep the cascade goes — with no client-side discard. + #[tokio::test] + async fn cascading_split_replacements_preserve_order_and_boundary() { + // P0 emits [1, 2] then splits into P0a + P0b; P0a splits *again* into + // P0a1 (next 3) and P0a2 (next 5) before yielding a row. + let p0a_cascade = split_page(vec![ + positioned_replacement_leaf("", "20", vec![envelope_page(&[("d3", 3)], None)]), + positioned_replacement_leaf("20", "40", vec![envelope_page(&[("d5", 5)], None)]), + ]); + let p0 = ChildStream::fresh( + range("", "80"), + Box::new(MockLeaf::with_pages(vec![ + envelope_page(&[("d1", 1), ("d2", 2)], Some("p0-ct")), + split_page(vec![ + positioned_replacement_leaf("", "40", vec![p0a_cascade]), + positioned_replacement_leaf( + "40", + "80", + vec![envelope_page(&[("d10", 10)], None)], + ), + ]), + ])), + ); + let mut node = merge(vec![p0], vec![SortOrder::Ascending]); + let emitted = drain_all_ids(&mut node).await; + assert_eq!( + emitted, + vec!["d1", "d2", "d3", "d5", "d10"] + .into_iter() + .map(str::to_owned) + .collect::>(), + "rows from a cascaded (twice-split) sub-range stay globally ordered" + ); + } + + #[tokio::test] + async fn ties_are_broken_deterministically_by_rid() { + // Both children have a row with rank=1; RID "a" < "b" must win. + let left = ChildStream::fresh( + range("", "80"), + Box::new(MockLeaf::with_pages(vec![envelope_page(&[("b", 1)], None)])), + ); + let right = ChildStream::fresh( + range("80", "FF"), + Box::new(MockLeaf::with_pages(vec![envelope_page(&[("a", 1)], None)])), + ); + let mut node = merge(vec![left, right], vec![SortOrder::Ascending]); + let PageResult::Page { response, .. } = next_page(&mut node).await else { + panic!("expected a page"); + }; + assert_eq!(ids(&response), vec!["a".to_owned(), "b".to_owned()]); + } + + #[tokio::test] + async fn multi_column_mixed_direction_drives_pop_order() { + // Column 0 ASC ties at 1; column 1 DESC means "b" sorts before "a". + fn two_col_page(rid: &str, c0: i64, c1: &str) -> crate::error::Result { + let body = serde_json::json!({ + "_rid": "", + "Documents": [{ + "_rid": rid, + "orderByItems": [{"item": c0}, {"item": c1}], + "payload": {"id": rid}, + }], + "_count": 1, + }); + Ok(PageResult::Page { + response: mocks::response(&serde_json::to_vec(&body).unwrap()), + is_terminal: true, + }) + } + let left = ChildStream::fresh( + range("", "80"), + Box::new(MockLeaf::with_pages(vec![two_col_page("left", 1, "a")])), + ); + let right = ChildStream::fresh( + range("80", "FF"), + Box::new(MockLeaf::with_pages(vec![two_col_page("right", 1, "b")])), + ); + let mut node = merge( + vec![left, right], + vec![SortOrder::Ascending, SortOrder::Descending], + ); + let PageResult::Page { response, .. } = next_page(&mut node).await else { + panic!("expected a page"); + }; + assert_eq!(ids(&response), vec!["right".to_owned(), "left".to_owned()]); + } + + #[tokio::test] + async fn undefined_sorts_before_defined_values_across_partitions() { + fn page(rid: &str, item: Option) -> crate::error::Result { + let order_by_items = match item { + Some(v) => serde_json::json!([{"item": v}]), + None => serde_json::json!([{}]), + }; + let body = serde_json::json!({ + "_rid": "", + "Documents": [{"_rid": rid, "orderByItems": order_by_items, "payload": {"id": rid}}], + "_count": 1, + }); + Ok(PageResult::Page { + response: mocks::response(&serde_json::to_vec(&body).unwrap()), + is_terminal: true, + }) + } + let left = ChildStream::fresh( + range("", "80"), + Box::new(MockLeaf::with_pages(vec![page("has-value", Some(1))])), + ); + let right = ChildStream::fresh( + range("80", "FF"), + Box::new(MockLeaf::with_pages(vec![page("undefined", None)])), + ); + let mut node = merge(vec![left, right], vec![SortOrder::Ascending]); + let PageResult::Page { response, .. } = next_page(&mut node).await else { + panic!("expected a page"); + }; + assert_eq!( + ids(&response), + vec!["undefined".to_owned(), "has-value".to_owned()] + ); + } + + /// Helper: an all-scalar single-column resume-boundary discard. `skip_count` + /// is the number of already-emitted rows sharing the exact `(value, rid)` + /// (usually `1` for a non-JOIN boundary). + fn number_boundary_discard(value: f64, last_rid: &str, skip_count: u32) -> PendingDiscard { + PendingDiscard::ResumeBoundary { + resume_values: vec![OrderByResumeValue::Number { + value: value.into(), + }], + last_rid: last_rid.to_owned(), + skip_count, + directions: vec![SortOrder::Ascending], + } + } + + #[tokio::test] + async fn resume_with_boundary_discards_already_emitted_ties_by_rid() { + // Rows tied on rank=5 with `_rid <= "tied-2"` were already emitted. + let mut child = ChildStream::fresh( + range("", "FF"), + Box::new(MockLeaf::with_pages(vec![envelope_page( + &[("tied-1", 5), ("tied-2", 5), ("tied-3", 5), ("new", 6)], + None, + )])), + ); + child.pending_discard = number_boundary_discard(5.0, "tied-2", 1); + let mut node = merge(vec![child], vec![SortOrder::Ascending]); + let PageResult::Page { response, .. } = next_page(&mut node).await else { + panic!("expected a page"); + }; + assert_eq!(ids(&response), vec!["tied-3".to_owned(), "new".to_owned()]); + } + + /// Regression: an empty leading page (with continuation) must keep the + /// boundary discard active, not clear it, or later tied rows leak through. + #[tokio::test] + async fn resume_boundary_discard_survives_empty_leading_page() { + let mut child = ChildStream::fresh( + range("", "FF"), + Box::new(MockLeaf::with_pages(vec![ + envelope_page(&[], Some("ct-empty")), + envelope_page(&[("tied-1", 5), ("new", 6)], None), + ])), + ); + child.pending_discard = number_boundary_discard(5.0, "tied-1", 1); + let mut node = merge(vec![child], vec![SortOrder::Ascending]); + let PageResult::Page { response, .. } = next_page(&mut node).await else { + panic!("expected a page"); + }; + assert_eq!( + ids(&response), + vec!["new".to_owned()], + "the tied row on the second page must still be discarded" + ); + } + + /// Regression: a tie run spanning a page boundary must stay fully + /// discarded, not just the portion on the first page. + #[tokio::test] + async fn resume_boundary_discard_survives_tie_run_spanning_pages() { + let mut child = ChildStream::fresh( + range("", "FF"), + Box::new(MockLeaf::with_pages(vec![ + envelope_page(&[("tied-1", 5), ("tied-2", 5)], Some("ct-mid")), + envelope_page(&[("tied-3", 5), ("new", 6)], None), + ])), + ); + child.pending_discard = number_boundary_discard(5.0, "tied-3", 1); + let mut node = merge(vec![child], vec![SortOrder::Ascending]); + let PageResult::Page { response, .. } = next_page(&mut node).await else { + panic!("expected a page"); + }; + assert_eq!( + ids(&response), + vec!["new".to_owned()], + "every tied row up to and including tied-3 was already emitted" + ); + } + + /// Regression: after a split, both sub-ranges resume from the same + /// boundary; the `_rid`-aware discard avoids dropping/duplicating rows. + #[tokio::test] + async fn split_resume_is_rid_aware_with_no_omissions_or_duplicates() { + // Pre-split emitted a,b,c tied on rank=5; left keeps a,c + e (unemitted tie) + m; + // right keeps b + z. + let left = { + let mut c = ChildStream::fresh( + range("", "80"), + Box::new(MockLeaf::with_pages(vec![envelope_page( + &[("a", 5), ("c", 5), ("e", 5), ("m", 7)], + None, + )])), + ); + c.pending_discard = number_boundary_discard(5.0, "c", 1); + c + }; + let right = { + let mut c = ChildStream::fresh( + range("80", "FF"), + Box::new(MockLeaf::with_pages(vec![envelope_page( + &[("b", 5), ("z", 6)], + None, + )])), + ); + c.pending_discard = number_boundary_discard(5.0, "c", 1); + c + }; + let mut node = merge(vec![left, right], vec![SortOrder::Ascending]); + let PageResult::Page { response, .. } = next_page(&mut node).await else { + panic!("expected a page"); + }; + assert_eq!( + ids(&response), + vec!["e".to_owned(), "z".to_owned(), "m".to_owned()], + "the unemitted tied row `e` must survive (no omission) and no \ + already-emitted row may reappear (no duplicate)" + ); + } + + /// A complex (array) boundary discards already-emitted ties via the + /// hash-based comparison, keeping same-hash rows past the rid cut-off. + /// This replaces the old positional rescan. (A differently-valued complex + /// row orders by the backend's hash order, which a `MockLeaf` can't + /// reproduce, so this exercises the tie run.) + #[tokio::test] + async fn resume_with_complex_boundary_discards_matching_hash_ties_by_rid() { + let boundary = + OrderByItem::Array(vec![OrderByItem::Number(5_i64.into())]).to_resume_value(); + let mut child = ChildStream::fresh( + range("", "FF"), + Box::new(MockLeaf::with_pages(vec![array_envelope_page( + &[("tied-1", 5), ("tied-2", 5), ("tied-3", 5)], + None, + )])), + ); + child.pending_discard = PendingDiscard::ResumeBoundary { + resume_values: vec![boundary], + last_rid: "tied-1".to_owned(), + skip_count: 1, + directions: vec![SortOrder::Ascending], + }; + let mut node = merge(vec![child], vec![SortOrder::Ascending]); + let PageResult::Page { response, .. } = next_page(&mut node).await else { + panic!("expected a page"); + }; + // `tied-1` (== boundary hash, rid <= last_rid) is dropped; `tied-2` + // and `tied-3` (same hash, rid > last_rid) survive. + assert_eq!( + ids(&response), + vec!["tied-2".to_owned(), "tied-3".to_owned()] + ); + } + + /// Regression for the complex-boundary discard: a *distinct* un-emitted + /// complex value must never be dropped, no matter how its MurmurHash + /// happens to order against the boundary's — hash order is not Cosmos + /// sort order. (The earlier hash-order comparison could classify `[1]` + /// as "before" the `[5]` boundary and silently drop it, ~50% of the + /// time, causing data loss.) + #[tokio::test] + async fn resume_complex_boundary_never_drops_distinct_unemitted_value() { + let boundary = + OrderByItem::Array(vec![OrderByItem::Number(5_i64.into())]).to_resume_value(); + let mut child = ChildStream::fresh( + range("", "FF"), + Box::new(MockLeaf::with_pages(vec![array_envelope_page( + &[("emitted", 5), ("distinct", 1)], + None, + )])), + ); + child.pending_discard = PendingDiscard::ResumeBoundary { + resume_values: vec![boundary], + last_rid: "emitted".to_owned(), + skip_count: 1, + directions: vec![SortOrder::Ascending], + }; + let mut node = merge(vec![child], vec![SortOrder::Ascending]); + let PageResult::Page { response, .. } = next_page(&mut node).await else { + panic!("expected a page"); + }; + // `emitted` (exact-hash tie, rid <= last_rid) is dropped; the + // distinct array `[1]` is kept — never inferred "before" from hash. + assert_eq!(ids(&response), vec!["distinct".to_owned()]); + } + + // ── skip_count: JOIN duplicate-RID resume ──────────────────────────── + + #[test] + fn record_emission_tracks_skip_count_for_duplicate_key_rid() { + let mut child = ChildStream::fresh(range("", "FF"), Box::new(MockLeaf::with_pages(vec![]))); + let key5 = [OrderByItem::Number(5_i64.into())]; + let key6 = [OrderByItem::Number(6_i64.into())]; + + // First emission of (5, docA): skip_count starts at 1. + child.record_emission(&key5, "docA").unwrap(); + assert_eq!(child.boundary().unwrap().skip_count, 1); + // Same (5, docA) again (a JOIN duplicate of one document): increments. + child.record_emission(&key5, "docA").unwrap(); + assert_eq!(child.boundary().unwrap().skip_count, 2); + // A new key with the same rid resets the count. + child.record_emission(&key6, "docA").unwrap(); + assert_eq!(child.boundary().unwrap().skip_count, 1); + // A new rid with the same key also resets. + child.record_emission(&key6, "docB").unwrap(); + let boundary = child.boundary().unwrap(); + assert_eq!(boundary.skip_count, 1); + assert_eq!(boundary.last_rid, "docB"); + } + + #[tokio::test] + async fn resume_skips_exactly_skip_count_duplicate_rid_rows() { + // Resume at (rank=5, rid=docA) after emitting 2 of docA's JOIN rows. + // The page re-returns all 3 docA rows plus a docB row; the discard + // drops exactly 2 (the emitted duplicates) and keeps the third, then + // the later document. + let mut child = ChildStream::fresh( + range("", "FF"), + Box::new(MockLeaf::with_pages(vec![join_envelope_page( + &[ + ("docA", 5, "a1"), + ("docA", 5, "a2"), + ("docA", 5, "a3"), + ("docB", 6, "b1"), + ], + None, + )])), + ); + child.pending_discard = number_boundary_discard(5.0, "docA", 2); + let mut node = merge(vec![child], vec![SortOrder::Ascending]); + let PageResult::Page { response, .. } = next_page(&mut node).await else { + panic!("expected a page"); + }; + assert_eq!(ids(&response), vec!["a3".to_owned(), "b1".to_owned()]); + } + + #[tokio::test] + async fn resume_skip_count_persists_across_pages() { + // skip_count = 3, but docA's JOIN rows straddle a page boundary: 2 on + // the first page, 2 on the second. The discard must carry the residual + // skip across the page break, dropping exactly 3 docA rows total. + let mut child = ChildStream::fresh( + range("", "FF"), + Box::new(MockLeaf::with_pages(vec![ + join_envelope_page(&[("docA", 5, "a1"), ("docA", 5, "a2")], Some("ct-mid")), + join_envelope_page( + &[("docA", 5, "a3"), ("docA", 5, "a4"), ("docB", 6, "b1")], + None, + ), + ])), + ); + child.pending_discard = number_boundary_discard(5.0, "docA", 3); + let mut node = merge(vec![child], vec![SortOrder::Ascending]); + let PageResult::Page { response, .. } = next_page(&mut node).await else { + panic!("expected a page"); + }; + assert_eq!( + ids(&response), + vec!["a4".to_owned(), "b1".to_owned()], + "3 already-emitted docA duplicates dropped across the page break; a4 and b1 survive" + ); + } + + #[tokio::test] + async fn legacy_boundary_missing_skip_count_discards_boundary_row() { + // A continuation token minted before `skip_count` existed omits the + // field; it must deserialize as skip_count == 1 so the single boundary + // row is still dropped on resume, never re-emitted. + let boundary: ValueBoundary = serde_json::from_str( + r#"{"resume_values":[{"type":"number","value":5.0}],"last_rid":"boundary"}"#, + ) + .unwrap(); + assert_eq!(boundary.skip_count, 1); + + let mut child = ChildStream::fresh( + range("", "FF"), + Box::new(MockLeaf::with_pages(vec![envelope_page( + &[("boundary", 5), ("new", 6)], + None, + )])), + ); + child.pending_discard = PendingDiscard::ResumeBoundary { + resume_values: boundary.resume_values, + last_rid: boundary.last_rid, + skip_count: boundary.skip_count, + directions: vec![SortOrder::Ascending], + }; + let mut node = merge(vec![child], vec![SortOrder::Ascending]); + let PageResult::Page { response, .. } = next_page(&mut node).await else { + panic!("expected a page"); + }; + assert_eq!(ids(&response), vec!["new".to_owned()]); + } + + // ── Live-split: forwarded-continuation resume ─────────────────────── + + /// Blocker regression: one JOIN document (rid `docA`, key 5) expands into + /// 150 result rows. Page 1 emits rows #1–100; the next fetch live-splits. + /// The replacement is a real `Request` in `Continuing` state carrying the + /// forwarded backend continuation (Cosmos contract: a parent partition's + /// continuation stays valid on a post-split child), so it resumes *after* + /// row #100 and its first page starts at #101. No client discard is + /// applied, so rows #101–150 all emit. The old code reinstalled the + /// boundary discard with `skip_count = 100` and would have silently dropped + /// all 50. + #[tokio::test] + async fn live_split_forwarded_continuation_emits_post_boundary_join_rows() { + let id_1_100: Vec = (1..=100).map(|i| format!("a{i}")).collect(); + let id_101_150: Vec = (101..=150).map(|i| format!("a{i}")).collect(); + let page1_rows: Vec<(&str, i64, &str)> = + id_1_100.iter().map(|s| ("docA", 5, s.as_str())).collect(); + let page2_rows: Vec<(&str, i64, &str)> = + id_101_150.iter().map(|s| ("docA", 5, s.as_str())).collect(); + + // The split child yields a real `Request` replacement carrying the + // forwarded continuation, mirroring `split_for_topology_change`. + let target = RequestTarget::effective_partition_key_range( + range("", "FF"), + "pk-0".to_owned(), + range("", "FF"), + ); + let replacement: Box = Box::new(Request::new( + Arc::new(mocks::operation()), + target, + Some("p0-ct".to_owned()), + )); + let split_child = ChildStream::fresh( + range("", "FF"), + Box::new(MockLeaf::with_pages(vec![ + join_envelope_page(&page1_rows, Some("p0-ct")), + split_page(vec![replacement]), + ])), + ); + let mut node = merge(vec![split_child], vec![SortOrder::Ascending]); + + // The replacement fetches once with the forwarded continuation; the + // backend returns only rows #101–150, already past the emitted prefix. + let mut executor = + mocks::MockRequestExecutor::new(vec![Ok(join_envelope_response(&page2_rows, None))]); + let mut topology = mocks::NoopTopologyProvider; + let mut context = PipelineContext::new(&mut executor, Some(&mut topology)); + + let PageResult::Page { response: r1, .. } = node.next_page(&mut context).await.unwrap() + else { + panic!("expected a page"); + }; + assert_eq!(ids(&r1), id_1_100, "page 1 emits the first 100 JOIN rows"); + + let PageResult::Page { response: r2, .. } = node.next_page(&mut context).await.unwrap() + else { + panic!("expected a page"); + }; + assert_eq!( + ids(&r2), + id_101_150, + "rows #101–150 must all emit — none dropped by a stale skip_count discard" + ); + + assert_eq!( + executor.continuation_calls, + vec![Some("p0-ct".to_owned())], + "the replacement resumed once from the forwarded continuation" + ); + } + + /// A live split whose replacement reports no forwarded continuation, yet + /// the split child had already emitted a row (boundary set), cannot be + /// safely repositioned — a generic node's resume position is unknown. + /// Reattaching the `skip_count` discard could drop or duplicate rows, so + /// the merge rejects with a typed `SPLIT_REPLACEMENT_INVALID` error. + #[tokio::test] + async fn live_split_replacement_without_continuation_with_boundary_errors() { + let split_child = ChildStream::fresh( + range("", "FF"), + Box::new(MockLeaf::with_pages(vec![split_page(vec![ + replacement_leaf("", "80", vec![envelope_page(&[("d3", 3)], None)]), + replacement_leaf("80", "FF", vec![envelope_page(&[("d9", 9)], None)]), + ])])), + ); + let mut node = merge(vec![split_child], vec![SortOrder::Ascending]); + // The split child had emitted a row before splitting (a plain boundary). + node.children[0].last_emitted = Some(LastEmitted { + resume_values: vec![OrderByResumeValue::Number { value: 2.0.into() }], + rid: "d2".to_owned(), + skip_count: 1, + }); + + let mut executor = mocks::NoopRequestExecutor; + let mut topology = mocks::NoopTopologyProvider; + let mut context = PipelineContext::new(&mut executor, Some(&mut topology)); + let err = node + .next_page(&mut context) + .await + .expect_err("a replacement with an unknown position must be rejected"); + assert_eq!( + err.status().sub_status(), + Some(crate::error::SubStatusCode::CLIENT_STREAMING_MERGE_SPLIT_REPLACEMENT_INVALID), + "must surface the typed split-replacement-invalid error (20215), got: {err}" + ); + } + + /// The companion accept case: a child that live-splits *before* emitting + /// any row has no boundary to protect, so generic replacements (no + /// forwarded continuation) are accepted and their rows stream through + /// fresh, in order. + #[tokio::test] + async fn live_split_initial_no_boundary_accepts_generic_replacements() { + let split_child = ChildStream::fresh( + range("", "FF"), + Box::new(MockLeaf::with_pages(vec![split_page(vec![ + replacement_leaf("", "80", vec![envelope_page(&[("d3", 3)], None)]), + replacement_leaf("80", "FF", vec![envelope_page(&[("d9", 9)], None)]), + ])])), + ); + // `last_emitted` stays `None` — no row emitted before the split. + let mut node = merge(vec![split_child], vec![SortOrder::Ascending]); + let emitted = drain_all_ids(&mut node).await; + assert_eq!( + emitted, + vec!["d3".to_owned(), "d9".to_owned()], + "an initial split with no boundary accepts fresh replacements" + ); + } + + /// A `ResumeFilterInjected` child that live-splits *before its first page* + /// (so its replacements carry no forwarded continuation) is rebuilt via the + /// boundary discard: each replacement re-runs the structured `resumeFilter` + /// (the backend `DistinctHash` seek), so the merge reinstalls the boundary + /// discard and orders the results. This is safe for a complex boundary + /// because the backend seek already excluded the emitted rows. The + /// replacement leaves are real `Request`s whose operation body carries that + /// `resumeFilter`, mirroring what `Request::split_for_topology_change` + /// clones on a live split. + #[tokio::test] + async fn resume_filter_injected_complex_boundary_live_split_is_accepted_and_ordered() { + let boundary = complex_boundary("rid-1"); + let filtered_body = query_response::with_resume_filter( + query_operation().body(), + &boundary.resume_values, + Some(&boundary.last_rid), + false, + ) + .expect("resume-filter body"); + let filtered_op = Arc::new((*query_operation()).clone().with_body(filtered_body)); + let replacement = |min: &'static str, max: &'static str| -> Box { + let target = RequestTarget::effective_partition_key_range( + range(min, max), + format!("pk-{min}-{max}"), + range(min, max), + ); + Box::new(Request::new(Arc::clone(&filtered_op), target, None)) + }; + // The split child (MockLeaf) yields two resume-filtered Request leaves. + let split_child = { + let mut c = ChildStream::fresh( + range("", "80"), + Box::new(MockLeaf::with_pages(vec![split_page(vec![ + replacement("", "40"), + replacement("40", "80"), + ])])), + ); + c.query_shape = ChildQueryShape::ResumeFilterInjected; + c.last_emitted = Some(LastEmitted { + resume_values: boundary.resume_values.clone(), + rid: boundary.last_rid.clone(), + skip_count: boundary.skip_count, + }); + c + }; + let mut node = merge(vec![split_child], vec![SortOrder::Ascending]); + + // The backend seek already excluded emitted rows, so each replacement + // returns only distinct, later complex values (d3 < d5). + let mut executor = mocks::MockRequestExecutor::new(vec![ + Ok(array_envelope_response(&[("d3", 3)], None)), + Ok(array_envelope_response(&[("d5", 5)], None)), + ]); + let mut topology = mocks::NoopTopologyProvider; + let mut context = PipelineContext::new(&mut executor, Some(&mut topology)); + let page = node + .next_page(&mut context) + .await + .expect("a resume-filtered complex-boundary split must be accepted, not rejected"); + let PageResult::Page { response, .. } = page else { + panic!("expected a page"); + }; + assert_eq!( + ids(&response), + vec!["d3".to_owned(), "d5".to_owned()], + "replacements must be accepted and ordered by key across the split" + ); + // Both replacement requests carried the structured resumeFilter seek. + assert!( + executor.body_text(0).contains("resumeFilter"), + "first replacement request must carry the resumeFilter body, got: {}", + executor.body_text(0) + ); + assert!( + executor.body_text(1).contains("resumeFilter"), + "second replacement request must carry the resumeFilter body, got: {}", + executor.body_text(1) + ); + } + + // ── build_children resume/topology paths ───────────────────────────── + + fn query_operation() -> Arc { + Arc::new( + mocks::operation().with_body( + br#"{"query":"SELECT * FROM c ORDER BY c.rank","parameters":[]}"#.to_vec(), + ), + ) + } + + /// One ascending sort direction, matching what the planner derives from + /// `QueryInfo::order_by`. + const ASC: &[SortOrder] = &[SortOrder::Ascending]; + + fn resolved_range(min: &str, max: &str, id: &str) -> ResolvedRange { + ResolvedRange { + partition_key_range_id: id.to_owned(), + range: range(min, max), + } + } + + fn scalar_boundary(value: f64, last_rid: &str) -> ValueBoundary { + ValueBoundary { + resume_values: vec![OrderByResumeValue::Number { + value: value.into(), + }], + last_rid: last_rid.to_owned(), + skip_count: 1, + } + } + + fn complex_boundary(last_rid: &str) -> ValueBoundary { + ValueBoundary { + resume_values: vec![ + OrderByItem::Array(vec![OrderByItem::Number(1.0.into())]).to_resume_value() + ], + last_rid: last_rid.to_owned(), + skip_count: 1, + } + } + + /// A scalar boundary crossing a split fans out into one resume-filtered + /// child per sub-range, each with the boundary discard installed. + #[test] + fn build_children_splits_scalar_boundary_into_resume_filtered_children() { + let op = query_operation(); + let scope = range("", "FF"); + let resolved = vec![ + resolved_range("", "80", "pk-left"), + resolved_range("80", "FF", "pk-right"), + ]; + let boundary = scalar_boundary(5.0, "c"); + let children = build_children(&resolved, &scope, &op, ASC, None, Some(&boundary)) + .expect("scalar boundary resumes across a split"); + assert_eq!(children.len(), 2); + for child in &children { + assert!( + matches!(child.pending_discard, PendingDiscard::ResumeBoundary { .. }), + "each split sub-range must resume via the boundary discard" + ); + assert_eq!( + child.query_shape, + ChildQueryShape::ResumeFilterInjected, + "a resume-filtered child's continuation must never be snapshotted" + ); + } + } + + /// A complex (array/object) boundary now also fans out across a split via + /// the structured `resumeFilter` — the old topology-change rejection is + /// gone because the backend seek is a per-row predicate. + #[test] + fn build_children_splits_complex_boundary_into_resume_filtered_children() { + let op = query_operation(); + let scope = range("", "FF"); + let resolved = vec![ + resolved_range("", "80", "pk-left"), + resolved_range("80", "FF", "pk-right"), + ]; + let boundary = complex_boundary("rid-1"); + let children = build_children(&resolved, &scope, &op, ASC, None, Some(&boundary)) + .expect("a complex boundary now resumes across a split"); + assert_eq!(children.len(), 2); + for child in &children { + assert!(matches!( + child.pending_discard, + PendingDiscard::ResumeBoundary { .. } + )); + } + } + + /// A merge resolves the saved sub-range to a wider physical range; it + /// must be clipped to scope before coverage validation, not rejected. + #[test] + fn build_children_clips_merged_physical_range_to_scope() { + let op = query_operation(); + let scope = range("", "80"); + // Post-merge: the saved [00,80) sub-range is now served by [00,FF). + let resolved = vec![resolved_range("", "FF", "pk-merged")]; + let boundary = scalar_boundary(5.0, "c"); + let children = build_children(&resolved, &scope, &op, ASC, None, Some(&boundary)) + .expect("a merged (widened) physical range clips to the saved scope"); + assert_eq!(children.len(), 1); + assert_eq!(children[0].range, scope); + } + + /// A complex boundary resumes across a merge (single clipped sub-range) + /// via the structured `resumeFilter`. + #[test] + fn build_children_allows_complex_boundary_across_merge() { + let op = query_operation(); + let scope = range("", "80"); + let resolved = vec![resolved_range("", "FF", "pk-merged")]; + let boundary = complex_boundary("rid-1"); + let children = build_children(&resolved, &scope, &op, ASC, None, Some(&boundary)) + .expect("a complex boundary resumes across a merge (single clipped range)"); + assert_eq!(children.len(), 1); + assert!(matches!( + children[0].pending_discard, + PendingDiscard::ResumeBoundary { .. } + )); + } + + #[tokio::test] + async fn page_size_cap_retains_unread_rows_for_next_call() { + let child = ChildStream::fresh( + range("", "FF"), + Box::new(MockLeaf::with_pages(vec![envelope_page( + &[("a", 1), ("b", 2), ("c", 3)], + None, + )])), + ); + let mut node = merge(vec![child], vec![SortOrder::Ascending]); + // Force a page size of 1 directly via the operation. + node.plain_operation = Arc::new((*node.plain_operation).clone().with_max_item_count( + crate::models::MaxItemCountHint::Limit(std::num::NonZeroU32::new(1).unwrap()), + )); + + let PageResult::Page { + response: r1, + is_terminal: t1, + } = next_page(&mut node).await + else { + panic!("expected a page"); + }; + assert_eq!(ids(&r1), vec!["a".to_owned()]); + assert!(!t1, "more rows remain buffered/available"); + + let PageResult::Page { + response: r2, + is_terminal: t2, + } = next_page(&mut node).await + else { + panic!("expected a page"); + }; + assert_eq!(ids(&r2), vec!["b".to_owned()]); + assert!(!t2); + + let PageResult::Page { + response: r3, + is_terminal: t3, + } = next_page(&mut node).await + else { + panic!("expected a page"); + }; + assert_eq!(ids(&r3), vec!["c".to_owned()]); + assert!(t3, "the child is drained and its buffer is now empty"); + } + + #[tokio::test] + async fn terminal_page_reports_drained_children() { + let child = ChildStream::fresh( + range("", "FF"), + Box::new(MockLeaf::with_pages(vec![envelope_page(&[("a", 1)], None)])), + ); + let mut node = merge(vec![child], vec![SortOrder::Ascending]); + let PageResult::Page { is_terminal, .. } = next_page(&mut node).await else { + panic!("expected a page"); + }; + assert!(is_terminal); + assert!(node.children.is_empty()); + + // Calling again on an already-fully-drained merge reports Drained. + assert!(matches!(next_page(&mut node).await, PageResult::Drained)); + } + + #[tokio::test] + async fn malformed_envelope_surfaces_typed_error() { + let body = serde_json::json!({ + "_rid": "", + "Documents": [{"_rid": "a", "orderByItems": [{"item": 1}]}], // missing payload + "_count": 1, + }); + let child = ChildStream::fresh( + range("", "FF"), + Box::new(MockLeaf::with_pages(vec![Ok(PageResult::Page { + response: mocks::response(&serde_json::to_vec(&body).unwrap()), + is_terminal: true, + })])), + ); + let mut node = merge(vec![child], vec![SortOrder::Ascending]); + let mut executor = mocks::NoopRequestExecutor; + let mut topology = mocks::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(), + crate::error::CosmosStatus::SERVICE_ORDER_BY_ENVELOPE_INVALID + ); + } + + #[tokio::test] + async fn snapshot_state_reports_drained_for_empty_merge() { + let node = merge(vec![], vec![SortOrder::Ascending]); + assert_eq!(node.snapshot_state().unwrap(), PipelineNodeState::Drained); + } + + /// Only a `ChildQueryShape::Plain` child may surface its live + /// continuation; a `ResumeFilterInjected` child's is bound to the + /// filtered text and must be suppressed, using its scalar boundary instead. + #[test] + fn snapshot_suppresses_resume_filtered_child_backend_continuation() { + fn child_with_live_continuation(shape: ChildQueryShape, token: &str) -> ChildStream { + // Empty buffer, not drained, with a live backend continuation. + let mut child = ChildStream::fresh( + range("", "FF"), + Box::new( + MockLeaf::with_pages(vec![]).with_snapshot(PipelineNodeState::Request { + server_continuation: Some(token.to_owned()), + }), + ), + ); + child.last_emitted = Some(LastEmitted { + resume_values: vec![OrderByResumeValue::Number { value: 1.0.into() }], + rid: "a".to_owned(), + skip_count: 1, + }); + child.query_shape = shape; + child + } + + let resume_filtered = merge( + vec![child_with_live_continuation( + ChildQueryShape::ResumeFilterInjected, + "resume-filtered-tok", + )], + vec![SortOrder::Ascending], + ); + match resume_filtered.snapshot_state().unwrap() { + PipelineNodeState::StreamingOrderedMerge { ranges, .. } => { + assert_eq!(ranges.len(), 1); + assert_eq!( + ranges[0].server_continuation, None, + "a resume-filtered child must never leak its backend continuation" + ); + assert!( + ranges[0].boundary.is_some(), + "it resumes from its scalar boundary instead" + ); + } + other => panic!("expected StreamingOrderedMerge, got {other:?}"), + } + + let plain = merge( + vec![child_with_live_continuation( + ChildQueryShape::Plain, + "plain-tok", + )], + vec![SortOrder::Ascending], + ); + match plain.snapshot_state().unwrap() { + PipelineNodeState::StreamingOrderedMerge { ranges, .. } => { + assert_eq!( + ranges[0].server_continuation, + Some("plain-tok".to_owned()), + "a plain-query child's backend continuation is portable and must be saved" + ); + } + other => panic!("expected StreamingOrderedMerge, got {other:?}"), + } + } + + #[tokio::test] + async fn cancellation_error_propagates_from_child_fetch() { + let child = ChildStream::fresh( + range("", "FF"), + Box::new(MockLeaf::with_pages(vec![Err( + mocks::non_topology_gone_error(), + )])), + ); + let mut node = merge(vec![child], vec![SortOrder::Ascending]); + let mut executor = mocks::NoopRequestExecutor; + let mut topology = mocks::NoopTopologyProvider; + let mut context = PipelineContext::new(&mut executor, Some(&mut topology)); + assert!(node.next_page(&mut context).await.is_err()); + } + + // ── Catalog-driven scenarios ───────────────────────────────────────── + // + // Reuses `tests/fixtures/streaming_order_by_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/streaming_order_by_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 = "expectedIds", default)] + expected_ids: Vec, + checkpoint: Option, + #[serde(rename = "expectedError")] + expected_error: Option, + } + + #[derive(serde::Deserialize)] + struct QueryFixture { + columns: Vec, + } + + #[derive(serde::Deserialize)] + struct ColumnFixture { + direction: String, + } + + #[derive(serde::Deserialize)] + struct MockFixture { + partitions: Vec, + } + + #[derive(serde::Deserialize)] + struct PartitionFixture { + pages: Vec, + } + + fn fixture_direction(s: &str) -> SortOrder { + match s { + "Ascending" => SortOrder::Ascending, + "Descending" => SortOrder::Descending, + other => panic!("unknown direction in fixture: {other}"), + } + } + + /// Runs catalog scenarios tagged `mockPipeline` through the real node, + /// asserting drained order matches `expectedIds`. Split checkpoints are + /// skipped (need the real planner — see `integration_tests::order_by_resume`); + /// malformed scenarios are covered by `malformed_envelope_surfaces_typed_error`. + #[tokio::test] + async fn catalog_mock_pipeline_scenarios_drain_in_expected_order() { + const CATALOG_JSON: &str = + include_str!("../../../tests/fixtures/streaming_order_by_scenarios.json"); + let catalog: CatalogFixture = + serde_json::from_str(CATALOG_JSON).expect("catalog must parse"); + + let mut ran_at_least_one = false; + let mut ran_a_resume_checkpoint = false; + for scenario in &catalog.scenarios { + if !scenario.layers.iter().any(|l| l == "mockPipeline") { + continue; + } + let Some(mock) = &scenario.mock else { + continue; + }; + if scenario.expected_error.is_some() { + // Covered by `malformed_envelope_surfaces_typed_error` instead. + continue; + } + + let directions: Vec = scenario + .query + .columns + .iter() + .map(|c| fixture_direction(&c.direction)) + .collect(); + + // Split checkpoints (`splitBeforeRow`/`replacementRanges`) need the + // real planner + topology, so they can't run on `MockLeaf` here and + // are skipped; mid-page-split execution coverage lives in the + // `split_during_*` node tests and `integration_tests::order_by_resume` + // (so no `splitBeforeRow > 0` fixture is added just to be skipped). A + // value-boundary resume becomes a `_rid`-aware discard on the single child. + let resume_discard: Option = match scenario.checkpoint.as_ref() { + None => None, + Some(cp) + if cp.get("splitBeforeRow").is_some() + || cp.get("replacementRanges").is_some() + || cp.get("splitReplacementRanges").is_some() => + { + continue; + } + Some(cp) => match ( + cp.get("resumeValues"), + cp.get("lastRid").and_then(|v| v.as_str()), + ) { + (Some(values), Some(last_rid)) => { + let resume_values: Vec = + serde_json::from_value(values.clone()) + .expect("checkpoint.resumeValues must parse as OrderByResumeValue"); + // A real boundary always emitted at least the boundary + // row, so a fixture that omits `skipCount` resumes as 1. + let skip_count = cp + .get("skipCount") + .and_then(|v| v.as_u64()) + .map_or(1, |n| n as u32); + ran_a_resume_checkpoint = true; + Some(PendingDiscard::ResumeBoundary { + resume_values, + last_rid: last_rid.to_owned(), + skip_count, + directions: directions.clone(), + }) + } + // A checkpoint that only affects token minting leaves + // the drain order unchanged: drive the full fresh drain. + _ => None, + }, + }; + ran_at_least_one = true; + + let mut children = Vec::new(); + for (idx, partition) in mock.partitions.iter().enumerate() { + let pages: Vec> = partition + .pages + .iter() + .map(|page_value| { + let continuation = page_value + .get("continuation") + .and_then(|c| c.as_str()) + .map(str::to_owned); + // Translate the fixture's `rid` field to the wire + // envelope's `_rid`. + let documents: Vec = page_value["rows"] + .as_array() + .cloned() + .unwrap_or_default() + .into_iter() + .map(|mut row| { + if let Some(obj) = row.as_object_mut() { + if let Some(rid) = obj.remove("rid") { + obj.insert("_rid".to_owned(), rid); + } + } + row + }) + .collect(); + let body = serde_json::json!({ + "_rid": "", + "Documents": documents, + "_count": documents.len(), + }); + Ok(PageResult::Page { + response: mocks::response_with_continuation( + &serde_json::to_vec(&body).unwrap(), + continuation.as_deref(), + ), + is_terminal: continuation.is_none(), + }) + }) + .collect(); + let lo = format!("{:02x}", idx * 0x10); + let hi = format!("{:02x}", (idx + 1) * 0x10); + children.push(ChildStream::fresh( + range(&lo, &hi), + Box::new(MockLeaf::with_pages(pages)), + )); + } + + if let Some(discard) = resume_discard { + // Value-boundary resume scenarios in the catalog are + // single-range (the boundary belongs to one range). + assert_eq!( + children.len(), + 1, + "scenario {} declares a value-boundary resume but has {} partitions", + scenario.id, + children.len(), + ); + children[0].pending_discard = discard; + } + + let mut node = merge(children, directions); + let mut collected_ids = Vec::new(); + loop { + let mut executor = mocks::NoopRequestExecutor; + let mut topology = mocks::NoopTopologyProvider; + let mut context = PipelineContext::new(&mut executor, Some(&mut topology)); + match node + .next_page(&mut context) + .await + .unwrap_or_else(|e| panic!("scenario {} failed: {e:?}", scenario.id)) + { + PageResult::Page { + response, + is_terminal, + } => { + collected_ids.extend(ids(&response)); + if is_terminal { + break; + } + } + PageResult::Drained => break, + PageResult::SplitRequired { .. } => { + panic!("scenario {} unexpectedly required a split", scenario.id) + } + } + } + + assert_eq!( + collected_ids, scenario.expected_ids, + "scenario {} drained ids do not match expectedIds", + scenario.id, + ); + } + assert!( + ran_at_least_one, + "no catalog scenario matched the mockPipeline filter; \ + the catalog-driven test wiring is broken" + ); + assert!( + ran_a_resume_checkpoint, + "expected at least one value-boundary resume checkpoint scenario \ + (e.g. equal_key_resume_requiring_skip_count) to run in the mock harness" + ); + } +} 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 f5cf6d3aebb..76f6ebf2eb8 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 @@ -516,12 +516,19 @@ impl SubStatusCode { 20211 => Some("ClientComputeRangeInvokedWithEmptyPartitionKey"), 20212 => Some("ClientChangeFeedPipelineUnexpectedlyDrained"), 20213 => Some("ClientContinuationTokenSavedRangeUnhonored"), + 20214 => Some("ClientContinuationTokenOrderByStateInvalid"), + 20215 => Some("ClientStreamingMergeSplitReplacementInvalid"), 20300 => Some("ClientNoOverlappingFeedRangesForSessionToken"), 20301 => Some("ClientNoThroughputOfferForResource"), 20302 => Some("ClientQueryPlanProducedEmptyRanges"), 20303 => Some("ServiceReturnedOfferWithoutId"), 20304 => Some("ClientThroughputPollerIncomplete"), 20305 => Some("ClientTopologyResolutionFailed"), + 20306 => Some("ServiceReturnedObjectWithoutRid"), + 20307 => Some("ClientQueryPlanRangeNotCoveredByTopology"), + 20308 => Some("ServiceOrderByEnvelopeInvalid"), + 20309 => Some("ServiceQueryPlanOrderByMissingRewrittenQuery"), + 20311 => Some("ServiceQueryPlanOrderByExpressionsMismatch"), // SDK Server-side codes (21xxx) - consistent across .NET and Java 21001 => Some("NameCacheIsStaleExceededRetryLimit"), @@ -1436,6 +1443,18 @@ impl SubStatusCode { /// see also 20200, 20203, 20204, 20205. pub const CLIENT_CONTINUATION_TOKEN_SAVED_RANGE_UNHONORED: SubStatusCode = SubStatusCode(20213); + /// A `StreamingOrderedMerge` continuation token is semantically invalid + /// (bad column/direction, hash, RID, or skip count) (20214). + pub const CLIENT_CONTINUATION_TOKEN_ORDER_BY_STATE_INVALID: SubStatusCode = + SubStatusCode(20214); + + /// `StreamingOrderedMerge` split replacement nodes are unusable (20215): + /// either their ranges do not exactly cover the prior range (a coverage + /// gap/overlap), or a replacement carries no continuation to resume a + /// mid-group boundary and cannot be safely repositioned. + pub const CLIENT_STREAMING_MERGE_SPLIT_REPLACEMENT_INVALID: SubStatusCode = + SubStatusCode(20215); + // ----- 20300-20349: SDK-detected service contract violations ----- /// The supplied session-token feed ranges contain no overlap with @@ -1481,6 +1500,22 @@ impl SubStatusCode { /// 503 — an internal client-side condition, not a transport failure. pub const CLIENT_TOPOLOGY_RESOLUTION_FAILED: SubStatusCode = SubStatusCode(20305); + /// A cross-partition streaming `ORDER BY` rewritten-query result item + /// did not match the expected envelope shape: missing `payload`, + /// missing/empty RID, or a mismatched `orderByItems` length (20308). + pub const SERVICE_ORDER_BY_ENVELOPE_INVALID: SubStatusCode = SubStatusCode(20308); + + /// The backend query plan reported `ORDER BY` columns but did not + /// supply a non-empty `rewrittenQuery` (20309). + pub const SERVICE_QUERY_PLAN_ORDER_BY_MISSING_REWRITTEN_QUERY: SubStatusCode = + SubStatusCode(20309); + + /// The backend query plan's `orderBy` and `orderByExpressions` arrays + /// had mismatched or zero length, so sort keys can't pair with + /// directions (20311). + pub const SERVICE_QUERY_PLAN_ORDER_BY_EXPRESSIONS_MISMATCH: SubStatusCode = + SubStatusCode(20311); + /// A topology range resolved for a query-plan EPK range did not /// overlap that range (20307). The query planner intersects each /// resolved partition with the query-plan range it was resolved @@ -2286,6 +2321,21 @@ impl CosmosStatus { sub_status: Some(SubStatusCode::CLIENT_CONTINUATION_TOKEN_SAVED_RANGE_UNHONORED), }; + /// 500 / 20214 — a `StreamingOrderedMerge` continuation token is + /// semantically invalid (bad column/direction, hash, RID, or count). + pub const CLIENT_CONTINUATION_TOKEN_ORDER_BY_STATE_INVALID: CosmosStatus = CosmosStatus { + status_code: StatusCode::InternalServerError, + sub_status: Some(SubStatusCode::CLIENT_CONTINUATION_TOKEN_ORDER_BY_STATE_INVALID), + }; + + /// 500 / 20215 — `StreamingOrderedMerge` split replacement nodes are + /// unusable: their ranges do not exactly cover the prior range, or a + /// replacement carries no continuation to resume a mid-group boundary. + pub const CLIENT_STREAMING_MERGE_SPLIT_REPLACEMENT_INVALID: CosmosStatus = CosmosStatus { + status_code: StatusCode::InternalServerError, + sub_status: Some(SubStatusCode::CLIENT_STREAMING_MERGE_SPLIT_REPLACEMENT_INVALID), + }; + // SDK-detected service contract violations (HTTP varies, sub-status 20300-20349) /// 410 / 20300 — the supplied session-token feed ranges contain no @@ -2334,6 +2384,27 @@ impl CosmosStatus { sub_status: Some(SubStatusCode::CLIENT_TOPOLOGY_RESOLUTION_FAILED), }; + /// 500 / 20308 — a rewritten-query result item didn't match the + /// expected envelope shape. + pub const SERVICE_ORDER_BY_ENVELOPE_INVALID: CosmosStatus = CosmosStatus { + status_code: StatusCode::InternalServerError, + sub_status: Some(SubStatusCode::SERVICE_ORDER_BY_ENVELOPE_INVALID), + }; + + /// 500 / 20309 — the query plan reported `ORDER BY` columns but no + /// `rewrittenQuery`. + pub const SERVICE_QUERY_PLAN_ORDER_BY_MISSING_REWRITTEN_QUERY: CosmosStatus = CosmosStatus { + status_code: StatusCode::InternalServerError, + sub_status: Some(SubStatusCode::SERVICE_QUERY_PLAN_ORDER_BY_MISSING_REWRITTEN_QUERY), + }; + + /// 500 / 20311 — the query plan's `orderBy`/`orderByExpressions` + /// arrays had mismatched or zero length. + pub const SERVICE_QUERY_PLAN_ORDER_BY_EXPRESSIONS_MISMATCH: CosmosStatus = CosmosStatus { + status_code: StatusCode::InternalServerError, + sub_status: Some(SubStatusCode::SERVICE_QUERY_PLAN_ORDER_BY_EXPRESSIONS_MISMATCH), + }; + /// 500 / 20307 — a topology range resolved for a query-plan EPK /// range did not overlap that range, a `resolve_ranges` contract /// violation. Returned instead of panicking the query worker (see 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 63987b1f76a..d8bf7c95bee 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 @@ -3185,7 +3185,13 @@ fn local_sort_order_to_dataflow( fn local_query_info_to_dataflow( info: crate::query::plan::LocalQueryInfo, + original_query: &str, ) -> crate::driver::dataflow::query_plan::QueryInfo { + let rewritten_query = if info.order_by.is_empty() { + Some(String::new()) + } else { + synthesize_order_by_rewritten_query(original_query, &info.order_by_expressions) + }; crate::driver::dataflow::query_plan::QueryInfo { distinct_type: local_distinct_type_to_dataflow(info.distinct_type), top: info.top.map(|v| v as u64), @@ -3205,12 +3211,86 @@ fn local_query_info_to_dataflow( .map(|a| format!("{a:?}")) .collect(), group_by_alias_to_aggregate_type: HashMap::new(), - rewritten_query: Some(String::new()), + rewritten_query, has_select_value: info.has_select_value, has_non_streaming_order_by: false, } } +/// Synthesizes a single-level rewritten `ORDER BY` envelope query, mirroring +/// what the real Gateway/native query-plan engine returns in +/// `rewrittenQuery`: +/// +/// ```text +/// SELECT VALUE {"_rid": ._rid, "orderByItems": [{"item": }, ...], "payload": } +/// +/// WHERE [() AND] {documentdb-formattableorderbyquery-filter} +/// +/// ``` +/// +/// The placeholder is substituted in place by the client with `true` +/// (fresh start) or a scalar `_rid`-aware resume filter (see +/// `driver::dataflow::query_response`) — no outer subquery wrapper, so the +/// result stays flat and directly evaluable. +/// +/// Scoped to `SELECT *`-shaped queries (`payload` is always the whole +/// document); returns `None` if no top-level `FROM` is found, which +/// shouldn't happen for a query with a non-empty ORDER BY. +fn synthesize_order_by_rewritten_query( + original_query: &str, + order_by_expressions: &[String], +) -> Option { + use crate::query::lexer::{Lexer, TokenKind}; + + // Kept in sync with `driver::dataflow::query_response`'s + // `ORDER_BY_FILTER_PLACEHOLDER` (this authors it; that substitutes it). + const ORDER_BY_FILTER_PLACEHOLDER: &str = "{documentdb-formattableorderbyquery-filter}"; + + let tokens = Lexer::tokenize(original_query); + let from_idx = tokens.iter().position(|t| t.kind == TokenKind::From)?; + let collection_token = tokens.get(from_idx + 1)?; + let alias = match tokens.get(from_idx + 2) { + Some(t) if t.kind == TokenKind::As => tokens.get(from_idx + 3)?.text, + Some(t) if t.kind == TokenKind::Identifier => t.text, + _ => collection_token.text, + }; + + let order_idx = tokens.iter().position(|t| t.kind == TokenKind::Order); + let clause_end = order_idx + .map(|i| tokens[i].span.start) + .unwrap_or(original_query.len()); + let order_by_text = order_idx.map(|i| original_query[tokens[i].span.start..].trim())?; + + // FROM is emitted verbatim; the placeholder is ANDed into WHERE + // (creating one if absent) — the slot every rewritten query carries. + let where_bound = order_idx.unwrap_or(tokens.len()); + let where_idx = (from_idx + 1..where_bound).find(|&i| tokens[i].kind == TokenKind::Where); + let from_end = where_idx + .map(|i| tokens[i].span.start) + .unwrap_or(clause_end); + let from_text = original_query[tokens[from_idx].span.start..from_end].trim(); + let where_clause = match where_idx { + Some(i) => { + let predicate = original_query[tokens[i].span.end..clause_end].trim(); + format!("WHERE ({predicate}) AND {ORDER_BY_FILTER_PLACEHOLDER}") + } + None => format!("WHERE {ORDER_BY_FILTER_PLACEHOLDER}"), + }; + + let order_by_items: Vec = order_by_expressions + .iter() + .map(|expr| format!(r#"{{"item": {expr}}}"#)) + .collect(); + + Some(format!( + r#"SELECT VALUE {{"_rid": {alias}._rid, "orderByItems": [{items}], "payload": {alias}}} {from_text} {where_clause} {order_by}"#, + items = order_by_items.join(", "), + from_text = from_text, + where_clause = where_clause, + order_by = order_by_text, + )) +} + fn full_query_range() -> crate::driver::dataflow::query_plan::QueryRange { crate::driver::dataflow::query_plan::QueryRange { min: Epk::MIN.to_hex(), @@ -3387,7 +3467,7 @@ fn handle_query_plan( let plan = crate::driver::dataflow::query_plan::QueryPlan { partitioned_query_execution_info_version: 2, - query_info: Some(local_query_info_to_dataflow(local_plan.query_info)), + query_info: Some(local_query_info_to_dataflow(local_plan.query_info, &query)), query_ranges, hybrid_search_query_info: None, }; diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/in_memory_emulator/rid.rs b/sdk/cosmos/azure_data_cosmos_driver/src/in_memory_emulator/rid.rs index 8d393082331..428a63ee390 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/in_memory_emulator/rid.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/in_memory_emulator/rid.rs @@ -61,14 +61,19 @@ impl RidGenerator { } /// Generates a new document RID (16 bytes: db_id + coll_id + doc_id with type nibble 0x0). + /// + /// The document segment (bytes `[8..16)`) is little-endian with the type + /// nibble in its top 4 bits, matching real Cosmos DB (and + /// `models::resource_id::document_ordinal`, which reads it the same + /// way for the `ORDER BY` tie-break) — a big-endian encoding would make + /// that tie-break non-monotonic once `doc_id >= 16`. pub fn next_document_rid(&self, db_id: u32, coll_id: u32) -> (u64, String) { let doc_id = self.doc_counter.fetch_add(1, Ordering::SeqCst); - let doc_with_type = doc_id << 4; // type nibble 0x0 in the low nibble let coll_with_high_bit = coll_id | 0x80000000; let mut bytes = [0u8; 16]; bytes[..4].copy_from_slice(&db_id.to_be_bytes()); bytes[4..8].copy_from_slice(&coll_with_high_bit.to_be_bytes()); - bytes[8..16].copy_from_slice(&doc_with_type.to_be_bytes()); + bytes[8..16].copy_from_slice(&doc_id.to_le_bytes()); (doc_id, encode_rid(&bytes)) } @@ -156,9 +161,17 @@ mod tests { let parent_db = u32::from_be_bytes(bytes[..4].try_into().unwrap()); assert_eq!(parent_db, db_id); - // Type nibble should be 0x0 for documents - let doc_raw = u64::from_be_bytes(bytes[8..16].try_into().unwrap()); - assert_eq!(doc_raw & 0x0F, 0x00); + // The document segment is little-endian (matching real Cosmos DB; + // see `next_document_rid`'s doc comment), so the ordinal itself + // round-trips via `from_le_bytes` and the type nibble is the top 4 + // bits, not the bottom 4. + let doc_raw = u64::from_le_bytes(bytes[8..16].try_into().unwrap()); + assert_eq!(doc_raw, doc_id); + assert_eq!( + doc_raw >> 60, + 0x0, + "type nibble should be 0x0 for documents" + ); } #[test] @@ -186,4 +199,29 @@ mod tests { assert_eq!(id1, 1); assert_eq!(id2, 2); } + + /// Regression: generated document RIDs must stay in creation order under + /// `document_ordinal`/`compare_document_rids` (the streaming `ORDER BY` + /// tie-break) well past `doc_id = 16` — the point at which a big-endian + /// document segment (the pre-fix encoding) stopped agreeing with the + /// little-endian reader, since the meaningful byte spills into a second + /// byte there. + #[test] + fn document_rids_stay_ordinal_ordered_past_16_documents() { + let gen = RidGenerator::new(); + let (db_id, _) = gen.next_database_rid(); + let (coll_id, _) = gen.next_collection_rid(db_id); + let rids: Vec = (0..40) + .map(|_| gen.next_document_rid(db_id, coll_id).1) + .collect(); + for pair in rids.windows(2) { + assert_eq!( + crate::models::resource_id::compare_document_rids(&pair[0], &pair[1]), + std::cmp::Ordering::Less, + "creation order must match document-ordinal order: {:?} then {:?}", + pair[0], + pair[1] + ); + } + } } diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/models/mod.rs b/sdk/cosmos/azure_data_cosmos_driver/src/models/mod.rs index 9e0addc336e..1ca5a4f558e 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/models/mod.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/models/mod.rs @@ -36,8 +36,9 @@ pub(crate) use finite_f64::FiniteF64; #[allow(dead_code)] pub mod effective_partition_key; mod feed_range; -#[allow(dead_code)] -mod murmur_hash; +// Used by `driver::dataflow::order_by` to hash complex resume values into +// a bounded 128-bit representation (mirrors .NET's complex-key encoding). +pub(crate) mod murmur_hash; #[allow(dead_code)] pub mod partition_key_range; #[allow(dead_code)] diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/models/resource_id.rs b/sdk/cosmos/azure_data_cosmos_driver/src/models/resource_id.rs index de3b34d58ed..3e2bf0285ae 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/models/resource_id.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/models/resource_id.rs @@ -69,6 +69,37 @@ pub(crate) fn encode_rid(bytes: &[u8]) -> String { STANDARD.encode(bytes).replace('/', "-") } +/// Extracts a document `_rid`'s document ordinal for `ORDER BY` tie-breaks: +/// the little-endian `u64` in the document segment (bytes `[8..16)` of the +/// decoded RID). Mirrors .NET `ResourceId.Document` (`BitConverter.ToUInt64`) +/// and Java `ResourceId.getDocument()` (`Long.reverseBytes`), both of which +/// read this segment as little-endian. +/// +/// Returns `None` when `rid` isn't a decodable RID with a document segment +/// (e.g. a synthetic test fixture like `"a"`); callers then fall back to +/// raw-string ordering. +pub(crate) fn document_ordinal(rid: &str) -> Option { + let bytes = decode_rid(rid).ok()?; + let segment: [u8; 8] = bytes.get(8..16)?.try_into().ok()?; + Some(u64::from_le_bytes(segment)) +} + +/// Compares two document `_rid`s in ascending Cosmos document order — the +/// order the backend uses to break `ORDER BY` key ties within a partition. +/// +/// When both decode as Cosmos RIDs, compares their document ordinals +/// numerically (see [`document_ordinal`]); a base64 `_rid` *string* +/// comparison would disagree with this (wrong byte order and a +/// non-monotonic alphabet), so it is never used. Values that aren't +/// decodable RIDs (test fixtures) fall back to raw-string comparison, +/// which is self-consistent within a single query. +pub(crate) fn compare_document_rids(a: &str, b: &str) -> std::cmp::Ordering { + match (document_ordinal(a), document_ordinal(b)) { + (Some(a_ord), Some(b_ord)) => a_ord.cmp(&b_ord), + _ => a.cmp(b), + } +} + /// A resource name (user-provided identifier). /// /// Used for human-readable identifiers like database names, container names, etc. @@ -592,4 +623,71 @@ mod tests { assert!(parsed.container_rid().is_some()); assert!(parsed.document_rid().is_some()); } + + // ===== Document-ordinal tie-break tests ===== + + /// Builds a document RID (16 bytes) with a fixed db/collection prefix and + /// `doc_id` as the little-endian document segment. + fn document_rid(doc_id: u64) -> String { + let mut bytes = [0u8; 16]; + bytes[0..4].copy_from_slice(&[0x0A, 0x0B, 0x0C, 0x0D]); + bytes[4..8].copy_from_slice(&[0x80, 0x01, 0x02, 0x03]); + bytes[8..16].copy_from_slice(&doc_id.to_le_bytes()); + encode_rid(&bytes) + } + + #[test] + fn document_ordinal_reads_little_endian_document_segment() { + assert_eq!(document_ordinal(&document_rid(0)), Some(0)); + assert_eq!(document_ordinal(&document_rid(1)), Some(1)); + assert_eq!(document_ordinal(&document_rid(u64::MAX)), Some(u64::MAX)); + assert_eq!( + document_ordinal(&document_rid(0x0102_0304_0506_0708)), + Some(0x0102_0304_0506_0708) + ); + // Not a decodable >=16-byte RID: falls back to `None`. + assert_eq!(document_ordinal("a"), None); + assert_eq!( + document_ordinal(&encode_rid(&[0x01, 0x02, 0x03, 0x04])), + None + ); + } + + #[test] + fn compare_document_rids_is_numeric_not_base64_string_order() { + use std::cmp::Ordering; + // Two document ids whose base64 RID strings sort opposite to their + // numeric order — proves we compare numerically, not by string. + let mut disagreement = None; + for hi in 0..64u64 { + let low = document_rid(hi); + let high = document_rid(hi + 1); + if low.cmp(&high) != Ordering::Less { + disagreement = Some((hi, low, high)); + break; + } + } + let (doc_id, low_rid, high_rid) = + disagreement.expect("some adjacent doc ids must disagree between string and numeric"); + // String order disagrees, but the numeric comparator is always correct. + assert_eq!( + compare_document_rids(&low_rid, &high_rid), + Ordering::Less, + "doc id {doc_id} < {} must compare Less numerically", + doc_id + 1 + ); + assert_eq!( + compare_document_rids(&high_rid, &low_rid), + Ordering::Greater + ); + assert_eq!(compare_document_rids(&low_rid, &low_rid), Ordering::Equal); + } + + #[test] + fn compare_document_rids_falls_back_to_string_for_non_rids() { + use std::cmp::Ordering; + // Synthetic test fixtures (not decodable RIDs) compare as raw strings. + assert_eq!(compare_document_rids("a", "b"), Ordering::Less); + assert_eq!(compare_document_rids("tied-1", "tied-2"), Ordering::Less); + } } 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 799e198b755..89be8eba875 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 @@ -580,15 +580,21 @@ fn project_group( // ─── ORDER BY helpers ──────────────────────────────────────────────────────── /// Type ordering for cross-type ORDER BY comparisons. +/// +/// Canonical Cosmos ascending type order: +/// `Undefined < Null < Boolean < Number < String < Array < Object`. +/// `Undefined` sorts *before* `Null`, not last — must match +/// `driver::dataflow::order_by`'s comparator to remain a valid ORDER BY +/// oracle. fn sort_type_order(v: &CosmosValue) -> u8 { match v { - CosmosValue::Null => 0, - CosmosValue::Boolean(_) => 1, - CosmosValue::Number(_) | CosmosValue::Integer(_) => 2, - CosmosValue::String(_) => 3, - CosmosValue::Array(_) => 4, - CosmosValue::Object(_) => 5, - CosmosValue::Undefined => 6, + CosmosValue::Undefined => 0, + CosmosValue::Null => 1, + CosmosValue::Boolean(_) => 2, + CosmosValue::Number(_) | CosmosValue::Integer(_) => 3, + CosmosValue::String(_) => 4, + CosmosValue::Array(_) => 5, + CosmosValue::Object(_) => 6, } } @@ -598,6 +604,31 @@ fn total_cmp_for_sort(a: &CosmosValue, b: &CosmosValue) -> Ordering { .unwrap_or_else(|| sort_type_order(a).cmp(&sort_type_order(b))) } +/// Deterministic full-key tie-break for the ORDER BY oracle: orders tied +/// rows by document `_rid` in the first sort column's direction (numeric +/// document-ordinal order via +/// [`crate::models::resource_id::compare_document_rids`]), matching the +/// backend and the production streaming merge. Returns `Ordering::Equal` +/// (leaving the stable sort untouched) when either row lacks a string +/// `_rid`, so projections without `_rid` are undisturbed. +fn order_by_rid_tiebreak( + a: &serde_json::Value, + b: &serde_json::Value, + order_by: &SqlOrderByClause, +) -> Ordering { + let (Some(a_rid), Some(b_rid)) = ( + a.get("_rid").and_then(serde_json::Value::as_str), + b.get("_rid").and_then(serde_json::Value::as_str), + ) else { + return Ordering::Equal; + }; + let ascending = crate::models::resource_id::compare_document_rids(a_rid, b_rid); + match order_by.items.first().map(|item| item.order) { + Some(SqlSortOrder::Descending) => ascending.reverse(), + _ => ascending, + } +} + /// Compare two documents according to an ORDER BY clause. /// /// F20 note: superseded by inline pre-computed-keys sort in `query_documents` @@ -944,7 +975,14 @@ pub fn query_documents( return cmp; } } - Ordering::Equal + // Full-key tie: order by document `_rid` to match the backend's + // deterministic tie order (non-grouped rows only; a group has no + // single `_rid`). + if groups.is_none() { + order_by_rid_tiebreak(&originals[a], &originals[b], order_by) + } else { + Ordering::Equal + } }); results = indices.iter().map(|&i| results[i].clone()).collect(); } @@ -2402,11 +2440,38 @@ mod tests { ]; let results = query_documents("SELECT * FROM c ORDER BY c.age ASC", &[], &docs).unwrap(); assert_eq!(results.len(), 3); - // Documents with defined age sort first in ASC - assert_eq!(results[0]["age"], 25); - assert_eq!(results[1]["age"], 30); - // Document missing age sorts last - assert_eq!(results[2]["name"], "Bob"); + // `Undefined` (missing property) sorts before every defined value. + assert_eq!(results[0]["name"], "Bob"); + assert_eq!(results[1]["age"], 25); + assert_eq!(results[2]["age"], 30); + } + + // Regression: pins the canonical Cosmos ascending type order so the + // emulator remains a valid ORDER BY oracle for `driver::dataflow::order_by`. + #[test] + fn order_by_undefined_sorts_before_null_and_every_other_type() { + let docs = vec![ + serde_json::json!({"name": "has-null", "val": null}), + serde_json::json!({"name": "has-bool", "val": false}), + serde_json::json!({"name": "has-undefined"}), + serde_json::json!({"name": "has-number", "val": 1}), + serde_json::json!({"name": "has-string", "val": "s"}), + ]; + let results = query_documents("SELECT * FROM c ORDER BY c.val ASC", &[], &docs).unwrap(); + let names: Vec<&str> = results + .iter() + .map(|d| d["name"].as_str().unwrap()) + .collect(); + assert_eq!( + names, + vec![ + "has-undefined", + "has-null", + "has-bool", + "has-number", + "has-string", + ] + ); } #[test] @@ -2441,6 +2506,78 @@ mod tests { // ── GROUP BY + Aggregates tests ───────────────────────────────────── + /// Pins that the evaluator correctly executes the emulator's + /// synthesized rewritten `ORDER BY` envelope shape (see + /// `in_memory_emulator::operations::synthesize_order_by_rewritten_query`). + #[test] + fn order_by_envelope_rewrite_single_level_query_evaluates_correctly() { + let docs = vec![ + serde_json::json!({"id": "a", "rank": 2}), + serde_json::json!({"id": "b", "rank": 1}), + ]; + let query = r#"SELECT VALUE {"_rid": c._rid, "orderByItems": [{"item": c.rank}], "payload": c} FROM c WHERE true ORDER BY c.rank ASC"#; + let results = query_documents(query, &[], &docs).unwrap(); + let ids: Vec<&str> = results + .iter() + .map(|r| r["payload"]["id"].as_str().unwrap()) + .collect(); + assert_eq!(ids, vec!["b", "a"]); + } + + /// Regression: a full-key tie must be broken by document `_rid` + /// (creation order), never by whatever order the documents were handed + /// to the evaluator in. In the real emulator, a logical partition's + /// documents are stored in a `BTreeMap` keyed by id (see + /// `PhysicalPartition::documents`), so its iteration order is + /// alphabetical-by-id — deliberately different here from the rids' + /// creation order, proving the tie-break really re-sorts rather than + /// just preserving input/storage order. + #[test] + fn order_by_tie_break_uses_rid_creation_order_not_input_order() { + fn real_rid(doc_id: u64) -> String { + let mut bytes = [0u8; 16]; + bytes[0..4].copy_from_slice(&[0x0A, 0x0B, 0x0C, 0x0D]); + bytes[4..8].copy_from_slice(&[0x80, 0x01, 0x02, 0x03]); + bytes[8..16].copy_from_slice(&doc_id.to_le_bytes()); + crate::models::resource_id::encode_rid(&bytes) + } + // Fed in alphabetical-by-id order (matching the store's `BTreeMap` + // iteration), but created in a different order: "bbb" first (rid + // ordinal 1), "aaa" second (ordinal 2), "ccc" third (ordinal 3). + let docs = vec![ + serde_json::json!({"id": "aaa", "rank": 5, "_rid": real_rid(2)}), + serde_json::json!({"id": "bbb", "rank": 5, "_rid": real_rid(1)}), + serde_json::json!({"id": "ccc", "rank": 5, "_rid": real_rid(3)}), + ]; + + let asc = query_documents("SELECT * FROM c ORDER BY c.rank ASC", &[], &docs).unwrap(); + let asc_ids: Vec<&str> = asc.iter().map(|d| d["id"].as_str().unwrap()).collect(); + assert_eq!( + asc_ids, + vec!["bbb", "aaa", "ccc"], + "ASC ties must follow creation (rid-ordinal) order, not alphabetical input order" + ); + + let desc = query_documents("SELECT * FROM c ORDER BY c.rank DESC", &[], &docs).unwrap(); + let desc_ids: Vec<&str> = desc.iter().map(|d| d["id"].as_str().unwrap()).collect(); + assert_eq!( + desc_ids, + vec!["ccc", "aaa", "bbb"], + "DESC ties must follow reverse creation (rid-ordinal) order, not input order" + ); + } + + /// Scope guard: `FROM (subquery)` remains unsupported. Unrelated to + /// ORDER BY resume, which never wraps a subquery. Remove if subquery + /// execution is later added. + #[test] + fn subquery_from_clause_is_not_yet_executable() { + let docs = vec![serde_json::json!({"id": "a", "rank": 1})]; + let query = r#"SELECT VALUE r FROM (SELECT c.id AS id FROM c) AS r"#; + let err = query_documents(query, &[], &docs).unwrap_err(); + assert!(format!("{err}").contains("FROM subqueries")); + } + #[test] fn group_by_count() { let docs = vec![ 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 4db63323cdd..9e9445db3ca 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/query/mod.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/query/mod.rs @@ -27,14 +27,10 @@ pub(crate) use parser::parse; /// advertises to the Cosmos DB Gateway via /// `x-ms-cosmos-supported-query-features`. /// -/// Advertised as `"None"`: the cross-partition -/// query pipeline does not yet support any of the advanced rewrite shapes the -/// Gateway can plan (Aggregate, CompositeAggregate, CountIf, DCount, Distinct, -/// GroupBy, HybridSearch, MultipleAggregates, MultipleOrderBy, -/// NonStreamingOrderBy, NonValueAggregate, OffsetAndLimit, OrderBy, Top, -/// WeightedRankFusion); advertising any of them in production would cause the -/// Gateway to return a plan we cannot execute. Add a feature here only after -/// the local pipeline gains support for the corresponding rewrite shape. +/// The production pipeline supports streaming single- and multi-column +/// `ORDER BY` rewrites. Other advanced rewrite shapes remain unadvertised +/// until their corresponding pipeline stages are implemented; advertising +/// one prematurely would cause the Gateway to return a plan we cannot execute. /// /// The value must be non-empty: the Gateway V2 thin-client proxy rejects /// QueryPlan requests where the `x-ms-cosmos-supported-query-features` header @@ -43,7 +39,7 @@ 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 = "None"; +pub(crate) const SUPPORTED_QUERY_FEATURES: &str = "OrderBy,MultipleOrderBy"; /// 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/tests/fixtures/streaming_order_by_scenarios.json b/sdk/cosmos/azure_data_cosmos_driver/tests/fixtures/streaming_order_by_scenarios.json new file mode 100644 index 00000000000..c46be360926 --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos_driver/tests/fixtures/streaming_order_by_scenarios.json @@ -0,0 +1,2054 @@ +{ + "schemaVersion": 1, + "scenarios": [ + { + "id": "single_column_asc_cross_partition", + "description": "Ascending single-column ORDER BY merges two partitions into one globally sorted stream.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.EmulatorTests/Query/OrderByQueryTests.cs", + "test": "TestOrderByCrossPartitionQueryAsync" + }, + { + "sdk": "java", + "path": "sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/OrderbyDocumentQueryTest.java", + "test": "queryDocumentsWithOrderBy" + } + ], + "layers": [ + "comparator", + "mockPipeline", + "inMemoryEmulator" + ], + "query": { + "text": "SELECT * FROM c ORDER BY c.rank ASC", + "parameters": [], + "columns": [ + { + "expression": "c.rank", + "direction": "Ascending" + } + ] + }, + "documents": [ + { + "id": "d1", + "pk": "a", + "rank": 1 + }, + { + "id": "d2", + "pk": "b", + "rank": 2 + }, + { + "id": "d3", + "pk": "a", + "rank": 3 + }, + { + "id": "d4", + "pk": "b", + "rank": 4 + } + ], + "mock": { + "partitions": [ + { + "range": { + "minEpk": "", + "maxEpk": "80" + }, + "pages": [ + { + "rows": [ + { + "rid": "p0-r1", + "orderByItems": [ + { + "item": 1 + } + ], + "payload": { + "id": "d1" + } + }, + { + "rid": "p0-r2", + "orderByItems": [ + { + "item": 3 + } + ], + "payload": { + "id": "d3" + } + } + ], + "continuation": null + } + ] + }, + { + "range": { + "minEpk": "80", + "maxEpk": "FF" + }, + "pages": [ + { + "rows": [ + { + "rid": "p1-r1", + "orderByItems": [ + { + "item": 2 + } + ], + "payload": { + "id": "d2" + } + }, + { + "rid": "p1-r2", + "orderByItems": [ + { + "item": 4 + } + ], + "payload": { + "id": "d4" + } + } + ], + "continuation": null + } + ] + } + ] + }, + "pageSizes": [ + 1, + 2, + 10 + ], + "expectedIds": [ + "d1", + "d2", + "d3", + "d4" + ], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "single_column_desc_cross_partition", + "description": "Descending single-column ORDER BY reverses the same-column comparator direction across partitions.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.EmulatorTests/Query/OrderByQueryTests.cs", + "test": "TestOrderByCrossPartitionQueryAsync" + }, + { + "sdk": "java", + "path": "sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/OrderbyDocumentQueryTest.java", + "test": "queryDocumentsWithOrderBy" + } + ], + "layers": [ + "comparator", + "mockPipeline", + "inMemoryEmulator" + ], + "query": { + "text": "SELECT * FROM c ORDER BY c.rank DESC", + "parameters": [], + "columns": [ + { + "expression": "c.rank", + "direction": "Descending" + } + ] + }, + "documents": [ + { + "id": "d1", + "pk": "a", + "rank": 1 + }, + { + "id": "d2", + "pk": "b", + "rank": 2 + }, + { + "id": "d3", + "pk": "a", + "rank": 3 + }, + { + "id": "d4", + "pk": "b", + "rank": 4 + } + ], + "mock": { + "partitions": [ + { + "range": { + "minEpk": "", + "maxEpk": "80" + }, + "pages": [ + { + "rows": [ + { + "rid": "p0-r1", + "orderByItems": [ + { + "item": 3 + } + ], + "payload": { + "id": "d3" + } + }, + { + "rid": "p0-r2", + "orderByItems": [ + { + "item": 1 + } + ], + "payload": { + "id": "d1" + } + } + ], + "continuation": null + } + ] + }, + { + "range": { + "minEpk": "80", + "maxEpk": "FF" + }, + "pages": [ + { + "rows": [ + { + "rid": "p1-r1", + "orderByItems": [ + { + "item": 4 + } + ], + "payload": { + "id": "d4" + } + }, + { + "rid": "p1-r2", + "orderByItems": [ + { + "item": 2 + } + ], + "payload": { + "id": "d2" + } + } + ], + "continuation": null + } + ] + } + ] + }, + "pageSizes": [ + 1, + 2, + 10 + ], + "expectedIds": [ + "d4", + "d3", + "d2", + "d1" + ], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "multi_column_mixed_direction_with_ties", + "description": "Validates all keys and deterministic tie breaking across two ORDER BY columns with opposite directions.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.EmulatorTests/Query/OrderByQueryTests.cs", + "test": "TestMultiOrderByQueriesAsync" + }, + { + "sdk": "java", + "path": "sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/MultiOrderByQueryTests.java", + "test": "queryDocumentsWithMultiOrder" + } + ], + "layers": [ + "comparator", + "mockPipeline", + "inMemoryEmulator" + ], + "query": { + "text": "SELECT * FROM c ORDER BY c.rank ASC, c.name DESC", + "parameters": [], + "columns": [ + { + "expression": "c.rank", + "direction": "Ascending" + }, + { + "expression": "c.name", + "direction": "Descending" + } + ] + }, + "documents": [ + { + "id": "d1", + "rank": 1, + "name": "b" + }, + { + "id": "d2", + "rank": 1, + "name": "a" + }, + { + "id": "d3", + "rank": 2, + "name": "z" + } + ], + "mock": { + "partitions": [ + { + "range": { + "minEpk": "", + "maxEpk": "FF" + }, + "pages": [ + { + "rows": [ + { + "rid": "p0-r1", + "orderByItems": [ + { + "item": 1 + }, + { + "item": "b" + } + ], + "payload": { + "id": "d1" + } + }, + { + "rid": "p0-r2", + "orderByItems": [ + { + "item": 1 + }, + { + "item": "a" + } + ], + "payload": { + "id": "d2" + } + }, + { + "rid": "p0-r3", + "orderByItems": [ + { + "item": 2 + }, + { + "item": "z" + } + ], + "payload": { + "id": "d3" + } + } + ], + "continuation": null + } + ] + } + ] + }, + "pageSizes": [ + 1, + 2, + 10 + ], + "expectedIds": [ + "d1", + "d2", + "d3" + ], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "exact_duplicate_keys_across_partitions", + "description": "Two partitions each contribute a row with the identical sort key; RID breaks the tie deterministically.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs", + "test": "TestDrainFully" + } + ], + "layers": [ + "comparator", + "mockPipeline" + ], + "query": { + "text": "SELECT * FROM c ORDER BY c.rank ASC", + "parameters": [], + "columns": [ + { + "expression": "c.rank", + "direction": "Ascending" + } + ] + }, + "documents": [], + "mock": { + "partitions": [ + { + "range": { + "minEpk": "", + "maxEpk": "80" + }, + "pages": [ + { + "rows": [ + { + "rid": "b-rid", + "orderByItems": [ + { + "item": 1 + } + ], + "payload": { + "id": "left" + } + } + ], + "continuation": null + } + ] + }, + { + "range": { + "minEpk": "80", + "maxEpk": "FF" + }, + "pages": [ + { + "rows": [ + { + "rid": "a-rid", + "orderByItems": [ + { + "item": 1 + } + ], + "payload": { + "id": "right" + } + } + ], + "continuation": null + } + ] + } + ] + }, + "pageSizes": [ + 1, + 10 + ], + "expectedIds": [ + "right", + "left" + ], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "undefined_field_single_column", + "description": "A document missing the sort property (Undefined) sorts before every defined value in single-column ASC order.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.Tests/Query/ContinuationTokens/OrderByQueryResultTests.cs", + "test": "TestOrderByUndefined" + } + ], + "layers": [ + "comparator", + "mockPipeline", + "inMemoryEmulator" + ], + "query": { + "text": "SELECT * FROM c ORDER BY c.rank ASC", + "parameters": [], + "columns": [ + { + "expression": "c.rank", + "direction": "Ascending" + } + ] + }, + "documents": [ + { + "id": "has-rank", + "rank": 1 + }, + { + "id": "no-rank" + } + ], + "mock": { + "partitions": [ + { + "range": { + "minEpk": "", + "maxEpk": "FF" + }, + "pages": [ + { + "rows": [ + { + "rid": "p0-r1", + "orderByItems": [ + {} + ], + "payload": { + "id": "no-rank" + } + }, + { + "rid": "p0-r2", + "orderByItems": [ + { + "item": 1 + } + ], + "payload": { + "id": "has-rank" + } + } + ], + "continuation": null + } + ] + } + ] + }, + "pageSizes": [ + 1, + 10 + ], + "expectedIds": [ + "no-rank", + "has-rank" + ], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "undefined_field_multi_column", + "description": "Undefined in the leading column sorts first regardless of the trailing column's value.", + "sources": [ + { + "sdk": "java", + "path": "sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/MultiOrderByQueryTests.java", + "test": "queryDocumentsWithMultiOrder" + } + ], + "layers": [ + "comparator", + "mockPipeline" + ], + "query": { + "text": "SELECT * FROM c ORDER BY c.rank ASC, c.name ASC", + "parameters": [], + "columns": [ + { + "expression": "c.rank", + "direction": "Ascending" + }, + { + "expression": "c.name", + "direction": "Ascending" + } + ] + }, + "documents": [], + "mock": { + "partitions": [ + { + "range": { + "minEpk": "", + "maxEpk": "FF" + }, + "pages": [ + { + "rows": [ + { + "rid": "p0-r1", + "orderByItems": [ + {}, + { + "item": "z" + } + ], + "payload": { + "id": "undefined-rank" + } + }, + { + "rid": "p0-r2", + "orderByItems": [ + { + "item": 1 + }, + { + "item": "a" + } + ], + "payload": { + "id": "defined-rank" + } + } + ], + "continuation": null + } + ] + } + ] + }, + "pageSizes": [ + 1, + 10 + ], + "expectedIds": [ + "undefined-rank", + "defined-rank" + ], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "mixed_type_ordering", + "description": "One key per Cosmos type validates the full ascending type-order chain: undefined, null, false, true, number, string, array, object.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.EmulatorTests/Query/OrderByQueryTests.cs", + "test": "TestOrderByCrossPartitionQueryAsync" + }, + { + "sdk": "java", + "path": "sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/OrderbyDocumentQueryTest.java", + "test": "queryDocumentsWithOrderByAndMixedTypes" + } + ], + "layers": [ + "comparator", + "mockPipeline" + ], + "query": { + "text": "SELECT * FROM c ORDER BY c.val ASC", + "parameters": [], + "columns": [ + { + "expression": "c.val", + "direction": "Ascending" + } + ] + }, + "documents": [], + "mock": { + "partitions": [ + { + "range": { + "minEpk": "", + "maxEpk": "FF" + }, + "pages": [ + { + "rows": [ + { + "rid": "r1", + "orderByItems": [ + {} + ], + "payload": { + "id": "v-undefined" + } + }, + { + "rid": "r2", + "orderByItems": [ + { + "item": null + } + ], + "payload": { + "id": "v-null" + } + }, + { + "rid": "r3", + "orderByItems": [ + { + "item": false + } + ], + "payload": { + "id": "v-false" + } + }, + { + "rid": "r4", + "orderByItems": [ + { + "item": true + } + ], + "payload": { + "id": "v-true" + } + }, + { + "rid": "r5", + "orderByItems": [ + { + "item": 5 + } + ], + "payload": { + "id": "v-number" + } + }, + { + "rid": "r6", + "orderByItems": [ + { + "item": "s" + } + ], + "payload": { + "id": "v-string" + } + }, + { + "rid": "r7", + "orderByItems": [ + { + "item": [ + 1, + 2 + ] + } + ], + "payload": { + "id": "v-array" + } + }, + { + "rid": "r8", + "orderByItems": [ + { + "item": { + "a": 1 + } + } + ], + "payload": { + "id": "v-object" + } + } + ], + "continuation": null + } + ] + } + ] + }, + "pageSizes": [ + 1, + 3, + 10 + ], + "expectedIds": [ + "v-undefined", + "v-null", + "v-false", + "v-true", + "v-number", + "v-string", + "v-array", + "v-object" + ], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "array_keys_ordering_and_continuation", + "description": "Array-typed sort keys compare lexicographically and round-trip through a bounded complex-value continuation.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.Tests/Query/ContinuationTokens/OrderByContinuationTokenTests.cs", + "test": "TestRoundTripComplexKeys" + }, + { + "sdk": "java", + "path": "sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/MultiOrderByQueryTests.java", + "test": "queryDocumentsWithMultiOrder" + } + ], + "layers": [ + "comparator", + "mockPipeline" + ], + "query": { + "text": "SELECT * FROM c ORDER BY c.tags ASC", + "parameters": [], + "columns": [ + { + "expression": "c.tags", + "direction": "Ascending" + } + ] + }, + "documents": [], + "mock": { + "partitions": [ + { + "range": { + "minEpk": "", + "maxEpk": "FF" + }, + "pages": [ + { + "rows": [ + { + "rid": "r1", + "orderByItems": [ + { + "item": [ + 1 + ] + } + ], + "payload": { + "id": "short" + } + }, + { + "rid": "r2", + "orderByItems": [ + { + "item": [ + 1, + 2 + ] + } + ], + "payload": { + "id": "prefix" + } + }, + { + "rid": "r3", + "orderByItems": [ + { + "item": [ + 2 + ] + } + ], + "payload": { + "id": "greater" + } + } + ], + "continuation": null + } + ] + } + ] + }, + "pageSizes": [ + 1, + 10 + ], + "expectedIds": [ + "short", + "prefix", + "greater" + ], + "checkpoint": { + "afterItems": 1 + }, + "expectedContinuation": { + "kind": "streaming_ordered_merge_range_boundary", + "resumeValueType": "complex", + "complexType": "array" + }, + "expectedError": null + }, + { + "id": "object_keys_ordering_and_continuation", + "description": "Object-typed sort keys compare by ascending key then value and round-trip through a bounded complex-value continuation.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.Tests/Query/ContinuationTokens/OrderByContinuationTokenTests.cs", + "test": "TestRoundTripComplexKeys" + }, + { + "sdk": "java", + "path": "sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/MultiOrderByQueryTests.java", + "test": "queryDocumentsWithMultiOrder" + } + ], + "layers": [ + "comparator", + "mockPipeline" + ], + "query": { + "text": "SELECT * FROM c ORDER BY c.meta ASC", + "parameters": [], + "columns": [ + { + "expression": "c.meta", + "direction": "Ascending" + } + ] + }, + "documents": [], + "mock": { + "partitions": [ + { + "range": { + "minEpk": "", + "maxEpk": "FF" + }, + "pages": [ + { + "rows": [ + { + "rid": "r1", + "orderByItems": [ + { + "item": { + "a": 1 + } + } + ], + "payload": { + "id": "a-1" + } + }, + { + "rid": "r2", + "orderByItems": [ + { + "item": { + "a": 2 + } + } + ], + "payload": { + "id": "a-2" + } + }, + { + "rid": "r3", + "orderByItems": [ + { + "item": { + "b": 0 + } + } + ], + "payload": { + "id": "b-0" + } + } + ], + "continuation": null + } + ] + } + ] + }, + "pageSizes": [ + 1, + 10 + ], + "expectedIds": [ + "a-1", + "a-2", + "b-0" + ], + "checkpoint": { + "afterItems": 1 + }, + "expectedContinuation": { + "kind": "streaming_ordered_merge_range_boundary", + "resumeValueType": "complex", + "complexType": "object" + }, + "expectedError": null + }, + { + "id": "equal_key_resume_requiring_skip_count", + "description": "A JOIN expands one document (RID `tied`) into three result rows sharing that RID and the same sort key `5`; two were already emitted (skipCount 2), so on resume the `(sort key, _rid, skipCount)` discard drops exactly those two duplicates and keeps the third plus the next document.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.Tests/Query/ContinuationTokens/OrderByContinuationTokenTests.cs", + "test": "TestOrderByContinuationTokenWithSkipCount" + } + ], + "layers": [ + "mockPipeline" + ], + "query": { + "text": "SELECT VALUE t FROM c JOIN t IN c.tags ORDER BY c.rank ASC", + "parameters": [], + "columns": [ + { + "expression": "c.rank", + "direction": "Ascending" + } + ] + }, + "documents": [], + "mock": { + "partitions": [ + { + "range": { + "minEpk": "", + "maxEpk": "FF" + }, + "pages": [ + { + "rows": [ + { + "rid": "tied", + "orderByItems": [ + { + "item": 5 + } + ], + "payload": { + "id": "dup-1" + } + }, + { + "rid": "tied", + "orderByItems": [ + { + "item": 5 + } + ], + "payload": { + "id": "dup-2" + } + }, + { + "rid": "tied", + "orderByItems": [ + { + "item": 5 + } + ], + "payload": { + "id": "dup-3" + } + }, + { + "rid": "new", + "orderByItems": [ + { + "item": 6 + } + ], + "payload": { + "id": "new" + } + } + ], + "continuation": null + } + ] + } + ] + }, + "pageSizes": [ + 1, + 10 + ], + "expectedIds": [ + "dup-3", + "new" + ], + "checkpoint": { + "resumeValues": [ + { + "type": "number", + "value": 5.0 + } + ], + "lastRid": "tied", + "skipCount": 2 + }, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "formatted_resume_filter_asc_desc_multi_column", + "description": "A multi-column ASC/DESC resume sends the boundary as the .NET-compatible structured resumeFilter (query text unchanged; direction lives in the ORDER BY clause), not an SDK-rewritten seek predicate.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs", + "test": "TestFormatFilter" + } + ], + "layers": [ + "comparator" + ], + "query": { + "text": "SELECT * FROM c ORDER BY c.rank ASC, c.name DESC", + "parameters": [], + "columns": [ + { + "expression": "c.rank", + "direction": "Ascending" + }, + { + "expression": "c.name", + "direction": "Descending" + } + ] + }, + "documents": [], + "mock": null, + "pageSizes": [], + "expectedIds": [], + "checkpoint": { + "resumeValues": [ + { + "type": "number", + "value": 5.0 + }, + { + "type": "string", + "value": "mid" + } + ] + }, + "expectedContinuation": { + "kind": "resume_filter", + "resumeFilter": { + "value": [ + 5.0, + "mid" + ], + "exclude": false + } + }, + "expectedError": null + }, + { + "id": "empty_total_result", + "description": "A query that matches nothing anywhere surfaces one empty, terminal page rather than an error.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.EmulatorTests/Query/OrderByQueryTests.cs", + "test": "TestOrderByCrossPartitionQueryAsync" + } + ], + "layers": [ + "mockPipeline", + "inMemoryEmulator" + ], + "query": { + "text": "SELECT * FROM c WHERE c.rank > 1000 ORDER BY c.rank ASC", + "parameters": [], + "columns": [ + { + "expression": "c.rank", + "direction": "Ascending" + } + ] + }, + "documents": [ + { + "id": "d1", + "rank": 1 + } + ], + "mock": { + "partitions": [ + { + "range": { + "minEpk": "", + "maxEpk": "FF" + }, + "pages": [ + { + "rows": [], + "continuation": null + } + ] + } + ] + }, + "pageSizes": [ + 10 + ], + "expectedIds": [], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "empty_backend_page_with_continuation_before_rows", + "description": "An empty page carrying a backend continuation must be transparently re-polled, not mistaken for drained.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.Tests/Query/Pipeline/OrderByCrossPartitionQueryPipelineStageTests.cs", + "test": "TestEmptyPages" + } + ], + "layers": [ + "mockPipeline", + "inMemoryEmulator" + ], + "query": { + "text": "SELECT * FROM c ORDER BY c.rank ASC", + "parameters": [], + "columns": [ + { + "expression": "c.rank", + "direction": "Ascending" + } + ] + }, + "documents": [ + { + "id": "d1", + "rank": 1 + } + ], + "mock": { + "partitions": [ + { + "range": { + "minEpk": "", + "maxEpk": "FF" + }, + "pages": [ + { + "rows": [], + "continuation": "ct-empty" + }, + { + "rows": [ + { + "rid": "r1", + "orderByItems": [ + { + "item": 1 + } + ], + "payload": { + "id": "d1" + } + } + ], + "continuation": null + } + ] + } + ] + }, + "pageSizes": [ + 10 + ], + "expectedIds": [ + "d1" + ], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "headers_request_charge_and_diagnostics_aggregation", + "description": "Request charge and diagnostics are summed across every backend page consumed to assemble one emitted page.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.Tests/Query/Pipeline/FullPipelineTests.cs", + "test": "TestRequestChargeAggregation" + } + ], + "layers": [ + "mockPipeline" + ], + "query": { + "text": "SELECT * FROM c ORDER BY c.rank ASC", + "parameters": [], + "columns": [ + { + "expression": "c.rank", + "direction": "Ascending" + } + ] + }, + "documents": [], + "mock": { + "partitions": [ + { + "range": { + "minEpk": "", + "maxEpk": "80" + }, + "pages": [ + { + "rows": [ + { + "rid": "l1", + "orderByItems": [ + { + "item": 1 + } + ], + "payload": { + "id": "l1" + } + } + ], + "continuation": null + } + ] + }, + { + "range": { + "minEpk": "80", + "maxEpk": "FF" + }, + "pages": [ + { + "rows": [ + { + "rid": "r1", + "orderByItems": [ + { + "item": 2 + } + ], + "payload": { + "id": "r1" + } + } + ], + "continuation": null + } + ] + } + ] + }, + "pageSizes": [ + 10 + ], + "expectedIds": [ + "l1", + "r1" + ], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "split_before_first_row", + "description": "A partition splits before it has produced any row; replacements start fresh with no resume boundary to protect.", + "sources": [ + { + "sdk": "java", + "path": "sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/QueryValidationTests.java", + "test": "splitQueryContinuationToken" + } + ], + "layers": [ + "mockPipeline" + ], + "query": { + "text": "SELECT * FROM c ORDER BY c.rank ASC", + "parameters": [], + "columns": [ + { + "expression": "c.rank", + "direction": "Ascending" + } + ] + }, + "documents": [], + "mock": { + "partitions": [ + { + "range": { + "minEpk": "", + "maxEpk": "FF" + }, + "pages": [ + { + "rows": [], + "continuation": null + } + ] + } + ] + }, + "pageSizes": [ + 10 + ], + "expectedIds": [], + "checkpoint": { + "splitBeforeRow": 0, + "replacementRanges": [ + { + "minEpk": "", + "maxEpk": "80" + }, + { + "minEpk": "80", + "maxEpk": "FF" + } + ] + }, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "split_after_resume", + "description": "A range resumed from a value boundary splits again before producing a row; the boundary must be forwarded to every new replacement.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.Tests/Query/OrderByQueryPartitionRangePageAsyncEnumeratorTests.cs", + "test": "TestSplitDuringDrain" + }, + { + "sdk": "java", + "path": "sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/QueryValidationTests.java", + "test": "splitQueryContinuationToken" + } + ], + "layers": [ + "mockPipeline" + ], + "query": { + "text": "SELECT * FROM c ORDER BY c.rank ASC", + "parameters": [], + "columns": [ + { + "expression": "c.rank", + "direction": "Ascending" + } + ] + }, + "documents": [], + "mock": { + "partitions": [ + { + "range": { + "minEpk": "", + "maxEpk": "FF" + }, + "pages": [ + { + "rows": [ + { + "rid": "after", + "orderByItems": [ + { + "item": 10 + } + ], + "payload": { + "id": "after" + } + } + ], + "continuation": null + } + ] + } + ] + }, + "pageSizes": [ + 10 + ], + "expectedIds": [ + "after" + ], + "checkpoint": { + "resumeValues": [ + { + "type": "number", + "value": 5.0 + } + ], + "lastRid": "before-split", + "skipCount": 0, + "splitReplacementRanges": [ + { + "minEpk": "", + "maxEpk": "80" + }, + { + "minEpk": "80", + "maxEpk": "FF" + } + ] + }, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "invalid_rewritten_envelope_missing_payload", + "description": "A rewritten-envelope item missing `payload` is a service contract violation, not silently skipped.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.Tests/Query/ContinuationTokens/OrderByQueryResultTests.cs", + "test": "TestOrderByUndefined" + } + ], + "layers": [ + "mockPipeline" + ], + "query": { + "text": "SELECT * FROM c ORDER BY c.rank ASC", + "parameters": [], + "columns": [ + { + "expression": "c.rank", + "direction": "Ascending" + } + ] + }, + "documents": [], + "mock": { + "partitions": [ + { + "range": { + "minEpk": "", + "maxEpk": "FF" + }, + "pages": [ + { + "rows": [ + { + "rid": "bad", + "orderByItems": [ + { + "item": 1 + } + ] + } + ], + "continuation": null + } + ] + } + ] + }, + "pageSizes": [ + 10 + ], + "expectedIds": [], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": { + "category": "service_order_by_envelope_invalid", + "messageFragment": "payload" + } + }, + { + "id": "invalid_rewritten_envelope_key_count_mismatch", + "description": "A rewritten-envelope item whose orderByItems length disagrees with the query's ORDER BY column count is rejected.", + "sources": [ + { + "sdk": "java", + "path": "sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/MultiOrderByQueryTests.java", + "test": "queryDocumentsWithMultiOrder" + } + ], + "layers": [ + "mockPipeline" + ], + "query": { + "text": "SELECT * FROM c ORDER BY c.rank ASC, c.name ASC", + "parameters": [], + "columns": [ + { + "expression": "c.rank", + "direction": "Ascending" + }, + { + "expression": "c.name", + "direction": "Ascending" + } + ] + }, + "documents": [], + "mock": { + "partitions": [ + { + "range": { + "minEpk": "", + "maxEpk": "FF" + }, + "pages": [ + { + "rows": [ + { + "rid": "bad", + "orderByItems": [ + { + "item": 1 + } + ], + "payload": { + "id": "bad" + } + } + ], + "continuation": null + } + ] + } + ] + }, + "pageSizes": [ + 10 + ], + "expectedIds": [], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": { + "category": "service_order_by_envelope_invalid", + "messageFragment": "orderByItems" + } + }, + { + "id": "malformed_token_wrong_node_kind", + "description": "A continuation token shaped for a different pipeline node kind must be rejected as a shape mismatch.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.Tests/Query/ContinuationTokens/OrderByContinuationTokenTests.cs", + "test": "TestNegativeCases" + } + ], + "layers": [ + "comparator" + ], + "query": { + "text": "SELECT * FROM c ORDER BY c.rank ASC", + "parameters": [], + "columns": [ + { + "expression": "c.rank", + "direction": "Ascending" + } + ] + }, + "documents": [], + "mock": null, + "pageSizes": [], + "expectedIds": [], + "checkpoint": { + "tokenShape": "sequential_drain" + }, + "expectedContinuation": null, + "expectedError": { + "category": "client_continuation_token_shape_mismatch", + "messageFragment": "does not match a streaming ORDER BY operation" + } + }, + { + "id": "malformed_token_direction_count_mismatch", + "description": "A continuation token whose column count/direction list disagrees with the current query must be rejected.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.Tests/Query/ContinuationTokens/OrderByContinuationTokenTests.cs", + "test": "TestNegativeCases" + } + ], + "layers": [ + "comparator" + ], + "query": { + "text": "SELECT * FROM c ORDER BY c.rank ASC", + "parameters": [], + "columns": [ + { + "expression": "c.rank", + "direction": "Ascending" + } + ] + }, + "documents": [], + "mock": null, + "pageSizes": [], + "expectedIds": [], + "checkpoint": { + "tokenShape": "streaming_ordered_merge", + "directions": [ + "Ascending", + "Descending" + ] + }, + "expectedContinuation": null, + "expectedError": { + "category": "client_continuation_token_order_by_state_invalid", + "messageFragment": "ORDER BY column count" + } + }, + { + "id": "malformed_token_invalid_range", + "description": "A continuation token with an invalid EPK range (min >= max) must be rejected.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.Tests/Query/ContinuationTokens/OrderByContinuationTokenTests.cs", + "test": "TestNegativeCases" + } + ], + "layers": [ + "comparator" + ], + "query": { + "text": "SELECT * FROM c ORDER BY c.rank ASC", + "parameters": [], + "columns": [ + { + "expression": "c.rank", + "direction": "Ascending" + } + ] + }, + "documents": [], + "mock": null, + "pageSizes": [], + "expectedIds": [], + "checkpoint": { + "tokenShape": "streaming_ordered_merge", + "invalidRange": { + "minEpk": "80", + "maxEpk": "40" + } + }, + "expectedContinuation": null, + "expectedError": { + "category": "client_continuation_token_order_by_state_invalid", + "messageFragment": "invalid range" + } + }, + { + "id": "single_logical_partition_remains_trivial_path", + "description": "An ORDER BY query scoped to a single, complete logical partition key stays on the trivial (server-ordered) path and never builds a StreamingOrderedMerge.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.Tests/Query/Pipeline/FullPipelineTests.cs", + "test": "TestSinglePartitionOrderBy" + } + ], + "layers": [ + "mockPipeline" + ], + "query": { + "text": "SELECT * FROM c WHERE c.pk = 'a' ORDER BY c.rank ASC", + "parameters": [], + "columns": [ + { + "expression": "c.rank", + "direction": "Ascending" + } + ] + }, + "documents": [ + { + "id": "d1", + "pk": "a", + "rank": 2 + }, + { + "id": "d2", + "pk": "a", + "rank": 1 + } + ], + "mock": null, + "pageSizes": [ + 10 + ], + "expectedIds": [ + "d2", + "d1" + ], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "order_by_with_filters_and_parameters", + "description": "A WHERE clause with a bound parameter is preserved verbatim in the rewritten per-partition request body.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.EmulatorTests/Query/OrderByQueryTests.cs", + "test": "TestOrderByCrossPartitionQueryAsync" + }, + { + "sdk": "java", + "path": "sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/OrderbyDocumentQueryTest.java", + "test": "queryDocumentsWithOrderBy" + } + ], + "layers": [ + "mockPipeline", + "inMemoryEmulator" + ], + "query": { + "text": "SELECT * FROM c WHERE c.category = @category ORDER BY c.rank ASC", + "parameters": [ + { + "name": "@category", + "value": "books" + } + ], + "columns": [ + { + "expression": "c.rank", + "direction": "Ascending" + } + ] + }, + "documents": [ + { + "id": "book-1", + "category": "books", + "rank": 2 + }, + { + "id": "book-2", + "category": "books", + "rank": 1 + }, + { + "id": "toy-1", + "category": "toys", + "rank": 0 + } + ], + "mock": { + "partitions": [ + { + "range": { + "minEpk": "", + "maxEpk": "FF" + }, + "pages": [ + { + "rows": [ + { + "rid": "r1", + "orderByItems": [ + { + "item": 1 + } + ], + "payload": { + "id": "book-2" + } + }, + { + "rid": "r2", + "orderByItems": [ + { + "item": 2 + } + ], + "payload": { + "id": "book-1" + } + } + ], + "continuation": null + } + ] + } + ] + }, + "pageSizes": [ + 10 + ], + "expectedIds": [ + "book-2", + "book-1" + ], + "checkpoint": null, + "expectedContinuation": null, + "expectedError": null + }, + { + "id": "unsupported_combination_top_with_streaming_order_by", + "description": "TOP combined with streaming ORDER BY remains precisely unsupported until the TOP composition stage (#4750) lands.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.Tests/Query/Pipeline/FullPipelineTests.cs", + "test": "TestUnsupportedCombinations" + } + ], + "layers": [ + "comparator" + ], + "query": { + "text": "SELECT TOP 5 * FROM c ORDER BY c.rank ASC", + "parameters": [], + "columns": [ + { + "expression": "c.rank", + "direction": "Ascending" + } + ] + }, + "documents": [], + "mock": null, + "pageSizes": [], + "expectedIds": [], + "checkpoint": { + "top": 5 + }, + "expectedContinuation": null, + "expectedError": { + "category": "client_unsupported_query_feature", + "messageFragment": "TOP combined with streaming ORDER BY" + } + }, + { + "id": "unsupported_combination_non_streaming_order_by", + "description": "A query the Gateway marks hasNonStreamingOrderBy remains rejected (issue #4755, not this issue).", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.Tests/Query/Pipeline/FullPipelineTests.cs", + "test": "TestUnsupportedCombinations" + } + ], + "layers": [ + "comparator" + ], + "query": { + "text": "SELECT * FROM c ORDER BY c.rank ASC", + "parameters": [], + "columns": [ + { + "expression": "c.rank", + "direction": "Ascending" + } + ] + }, + "documents": [], + "mock": null, + "pageSizes": [], + "expectedIds": [], + "checkpoint": { + "hasNonStreamingOrderBy": true + }, + "expectedContinuation": null, + "expectedError": { + "category": "client_unsupported_query_feature", + "messageFragment": "non-streaming ORDER BY" + } + }, + { + "id": "unsupported_combination_distinct_group_by_aggregates", + "description": "DISTINCT, GROUP BY, and aggregates combined with ORDER BY remain unsupported until their sibling issues land.", + "sources": [ + { + "sdk": "dotnet", + "path": "Microsoft.Azure.Cosmos.Tests/Query/Pipeline/FullPipelineTests.cs", + "test": "TestUnsupportedCombinations" + } + ], + "layers": [ + "comparator" + ], + "query": { + "text": "SELECT DISTINCT c.category FROM c ORDER BY c.category ASC", + "parameters": [], + "columns": [ + { + "expression": "c.category", + "direction": "Ascending" + } + ] + }, + "documents": [], + "mock": null, + "pageSizes": [], + "expectedIds": [], + "checkpoint": { + "distinctType": "Ordered" + }, + "expectedContinuation": null, + "expectedError": { + "category": "client_unsupported_query_feature", + "messageFragment": "DISTINCT combined with ORDER BY" + } + } + ] +} 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 2951e653697..92ffae9bdb1 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 @@ -18,6 +18,7 @@ pub mod hedging; pub mod host_recorder; pub mod multi_region; pub mod offers; +pub mod order_by; pub mod point_operations; pub mod ppaf_dynamic_enablement; pub mod query; diff --git a/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/order_by.rs b/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/order_by.rs new file mode 100644 index 00000000000..871590d356f --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/order_by.rs @@ -0,0 +1,333 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +//! In-memory-emulator integration tests for the cross-partition streaming +//! `ORDER BY` pipeline (`StreamingOrderedMerge`). +//! +//! Unlike the mock-pipeline tests, these exercise the real planner → +//! query-plan fetch → per-partition rewritten-query execution → envelope +//! parsing → merge path against the in-memory emulator's actual query +//! evaluator. +//! +//! Scoped to *fresh* (non-resumed) scenarios; resume/split behavior is +//! covered by the mock-pipeline and driver-level integration tests. + +use std::sync::Arc; + +use azure_core::http::Url; + +use azure_data_cosmos_driver::in_memory_emulator::{ + ConsistencyLevel, ContainerConfig, InMemoryEmulatorHttpClient, VirtualAccountConfig, + VirtualRegion, +}; +use azure_data_cosmos_driver::models::{ + CosmosOperation, FeedRange, ItemReference, MaxItemCountHint, PartitionKey, + PartitionKeyDefinition, +}; +use azure_data_cosmos_driver::options::{DriverOptions, OperationOptions}; + +const GATEWAY_URL: &str = "https://eastus.emulator.local"; + +/// 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) +} + +#[tokio::test] +async fn cross_partition_order_by_returns_globally_sorted_results() { + let (_emulator, driver) = setup().await; + let container = driver + .resolve_container("testdb", "testcoll") + .await + .expect("container resolves"); + + // Shuffled ranks across partitions so an unsorted merge would fail. + let seeds = [ + ("d1", "pk-a", 5), + ("d2", "pk-b", 1), + ("d3", "pk-c", 4), + ("d4", "pk-d", 2), + ("d5", "pk-e", 3), + ("d6", "pk-f", 0), + ]; + for (id, pk, rank) in seeds { + let item_ref = ItemReference::from_name(&container, PartitionKey::from(pk), id); + let body = serde_json::json!({"id": id, "pk": pk, "rank": rank}); + driver + .execute_singleton_operation( + CosmosOperation::create_item(item_ref) + .with_body(serde_json::to_vec(&body).unwrap()), + OperationOptions::default(), + ) + .await + .expect("seed item created"); + } + + let operation = CosmosOperation::query_items(container.clone(), Some(FeedRange::full())) + .with_body(br#"{"query":"SELECT * FROM c ORDER BY c.rank ASC","parameters":[]}"#.to_vec()); + + let mut plan = driver + .plan_operation(operation, &OperationOptions::default(), None) + .await + .expect("plan builds a StreamingOrderedMerge pipeline"); + + let mut ranks: Vec = Vec::new(); + while let Some(response) = driver + .execute_plan( + &mut plan, + Some(container.clone()), + OperationOptions::default(), + ) + .await + .expect("page executes") + { + let body: serde_json::Value = + serde_json::from_slice(&response.into_body().single().unwrap()).unwrap(); + for item in body["Documents"].as_array().unwrap() { + ranks.push(item["rank"].as_i64().unwrap()); + } + } + + assert_eq!(ranks, vec![0, 1, 2, 3, 4, 5]); +} + +#[tokio::test] +async fn cross_partition_order_by_paginates_with_small_page_size() { + let (_emulator, driver) = setup().await; + let container = driver + .resolve_container("testdb", "testcoll") + .await + .expect("container resolves"); + + for i in 0..8 { + let id = format!("item-{i}"); + let pk = format!("pk-{i}"); + // Descending seed order so a passthrough merge would fail. + let rank = 8 - i; + let item_ref = + ItemReference::from_name(&container, PartitionKey::from(pk.clone()), id.clone()); + let body = serde_json::json!({"id": id, "pk": pk, "rank": rank}); + driver + .execute_singleton_operation( + CosmosOperation::create_item(item_ref) + .with_body(serde_json::to_vec(&body).unwrap()), + OperationOptions::default(), + ) + .await + .expect("seed item created"); + } + + let operation = CosmosOperation::query_items(container.clone(), Some(FeedRange::full())) + .with_body(br#"{"query":"SELECT * FROM c ORDER BY c.rank ASC","parameters":[]}"#.to_vec()) + .with_max_item_count(MaxItemCountHint::Limit( + std::num::NonZeroU32::new(3).unwrap(), + )); + + let mut plan = driver + .plan_operation(operation, &OperationOptions::default(), None) + .await + .unwrap(); + + let mut ranks: Vec = Vec::new(); + let mut page_sizes: Vec = Vec::new(); + while let Some(response) = driver + .execute_plan( + &mut plan, + Some(container.clone()), + OperationOptions::default(), + ) + .await + .unwrap() + { + let body: serde_json::Value = + serde_json::from_slice(&response.into_body().single().unwrap()).unwrap(); + let documents = body["Documents"].as_array().unwrap(); + page_sizes.push(documents.len()); + for item in documents { + ranks.push(item["rank"].as_i64().unwrap()); + } + } + + assert_eq!(ranks, (1..=8).collect::>()); + assert!( + page_sizes.iter().all(|&n| n <= 3), + "no page should exceed the requested max_item_count: {page_sizes:?}" + ); + assert!( + page_sizes.len() >= 3, + "expected multiple pages: {page_sizes:?}" + ); +} + +/// Runs `query` (already-valid Cosmos SQL) to completion and returns every +/// result's `id`, in emitted order. +async fn run_query_collecting_ids( + driver: &azure_data_cosmos_driver::driver::CosmosDriver, + container: &azure_data_cosmos_driver::models::ContainerReference, + query: &str, +) -> Vec { + let body = serde_json::to_vec(&serde_json::json!({"query": query, "parameters": []})).unwrap(); + let operation = + CosmosOperation::query_items(container.clone(), Some(FeedRange::full())).with_body(body); + let mut plan = driver + .plan_operation(operation, &OperationOptions::default(), None) + .await + .expect("plan builds a StreamingOrderedMerge pipeline"); + + let mut ids = Vec::new(); + while let Some(response) = driver + .execute_plan( + &mut plan, + Some(container.clone()), + OperationOptions::default(), + ) + .await + .expect("page executes") + { + let body: serde_json::Value = + serde_json::from_slice(&response.into_body().single().unwrap()).unwrap(); + for item in body["Documents"].as_array().unwrap() { + ids.push(item["id"].as_str().unwrap().to_owned()); + } + } + ids +} + +/// Regression: within a single logical partition, a full-key tie must be +/// broken by document `_rid` (creation order) — never by the store's +/// internal iteration order. The in-memory emulator stores a logical +/// partition's documents in a `BTreeMap` keyed by id (alphabetical +/// iteration), so seeding out of alphabetical order ("bbb" first, then +/// "aaa", then "ccc", all under the same partition key) makes the two +/// orders disagree: a correct rid-based tie-break must reorder them. +#[tokio::test] +async fn tied_order_by_key_within_one_partition_returns_rid_order_not_storage_order() { + let (_emulator, driver) = setup().await; + let container = driver + .resolve_container("testdb", "testcoll") + .await + .expect("container resolves"); + + let pk = "tied-pk"; + for id in ["bbb", "aaa", "ccc"] { + let item_ref = ItemReference::from_name(&container, PartitionKey::from(pk), id); + let body = serde_json::json!({"id": id, "pk": pk, "rank": 5}); + driver + .execute_singleton_operation( + CosmosOperation::create_item(item_ref) + .with_body(serde_json::to_vec(&body).unwrap()), + OperationOptions::default(), + ) + .await + .expect("seed item created"); + } + + let asc_ids = + run_query_collecting_ids(&driver, &container, "SELECT * FROM c ORDER BY c.rank ASC").await; + assert_eq!( + asc_ids, + vec!["bbb", "aaa", "ccc"], + "ASC ties must follow creation (rid) order, not the store's alphabetical-by-id \ + iteration order: {asc_ids:?}" + ); + + let desc_ids = + run_query_collecting_ids(&driver, &container, "SELECT * FROM c ORDER BY c.rank DESC").await; + assert_eq!( + desc_ids, + vec!["ccc", "aaa", "bbb"], + "DESC ties must follow reverse creation (rid) order: {desc_ids:?}" + ); +} + +/// Regression: a full-key tie spanning *different* physical partitions +/// must also be broken by document `_rid` in the driver's cross-partition +/// merge — not by whichever partition happens to be polled/arrive first. +/// Seeded out of both alphabetical-id and partition-key order so neither +/// coincidentally matches rid (creation) order. +#[tokio::test] +async fn tied_order_by_key_across_partitions_returns_rid_order() { + let (_emulator, driver) = setup().await; + let container = driver + .resolve_container("testdb", "testcoll") + .await + .expect("container resolves"); + + // Same partition keys as `cross_partition_order_by_returns_globally_sorted_results`, + // known to spread across both of this container's physical partitions. + let seeds = [ + ("s6", "pk-f"), + ("s1", "pk-a"), + ("s5", "pk-e"), + ("s2", "pk-b"), + ("s4", "pk-d"), + ("s3", "pk-c"), + ]; + for (id, pk) in seeds { + let item_ref = ItemReference::from_name(&container, PartitionKey::from(pk), id); + let body = serde_json::json!({"id": id, "pk": pk, "rank": 7}); + driver + .execute_singleton_operation( + CosmosOperation::create_item(item_ref) + .with_body(serde_json::to_vec(&body).unwrap()), + OperationOptions::default(), + ) + .await + .expect("seed item created"); + } + let creation_order: Vec = seeds.iter().map(|(id, _)| id.to_string()).collect(); + + let asc_ids = + run_query_collecting_ids(&driver, &container, "SELECT * FROM c ORDER BY c.rank ASC").await; + assert_eq!( + asc_ids, creation_order, + "cross-partition ASC ties must follow creation (rid) order: {asc_ids:?}" + ); + + let desc_ids = + run_query_collecting_ids(&driver, &container, "SELECT * FROM c ORDER BY c.rank DESC").await; + let mut expected_desc = creation_order; + expected_desc.reverse(); + assert_eq!( + desc_ids, expected_desc, + "cross-partition DESC ties must follow reverse creation (rid) order: {desc_ids:?}" + ); +} diff --git a/sdk/cosmos/azure_data_cosmos_driver/tests/streaming_order_by_scenario_catalog.rs b/sdk/cosmos/azure_data_cosmos_driver/tests/streaming_order_by_scenario_catalog.rs new file mode 100644 index 00000000000..90c41730d76 --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos_driver/tests/streaming_order_by_scenario_catalog.rs @@ -0,0 +1,314 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +//! Fixture/schema tests for the streaming cross-partition `ORDER BY` +//! scenario catalog (`tests/fixtures/streaming_order_by_scenarios.json`). +//! +//! This is the single source-attributed catalog reused across every test +//! layer (comparator/continuation 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 and some +//! expected result; mock partitions are sorted, gapless, and tile +//! correctly; row `orderByItems` length matches the query's column count; +//! and required scenario-inventory categories are represented. + +use std::collections::BTreeSet; + +use serde::Deserialize; + +const CATALOG_JSON: &str = include_str!("fixtures/streaming_order_by_scenarios.json"); + +const KNOWN_LAYERS: &[&str] = &["comparator", "mockPipeline", "inMemoryEmulator", "recorded"]; + +#[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)] + #[allow(dead_code)] + page_sizes: Vec, + #[serde(rename = "expectedIds", default)] + expected_ids: Vec, + #[allow(dead_code)] + checkpoint: Option, + #[serde(rename = "expectedContinuation")] + expected_continuation: Option, + #[serde(rename = "expectedError")] + expected_error: Option, +} + +#[derive(Deserialize)] +struct Source { + #[allow(dead_code)] + sdk: String, + #[allow(dead_code)] + path: String, + #[allow(dead_code)] + test: String, +} + +#[derive(Deserialize)] +struct QuerySpec { + #[allow(dead_code)] + text: String, + #[serde(default)] + #[allow(dead_code)] + parameters: Vec, + columns: Vec, +} + +#[derive(Deserialize)] +struct ColumnSpec { + #[allow(dead_code)] + expression: String, + #[allow(dead_code)] + 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 { + #[allow(dead_code)] + rid: String, + #[serde(rename = "orderByItems")] + order_by_items: Vec, + #[allow(dead_code)] + payload: Option, +} + +#[derive(Deserialize)] +struct ExpectedError { + #[allow(dead_code)] + category: String, + #[allow(dead_code)] + #[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() { + let catalog = load_catalog(); + assert_eq!(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 + ); + } +} + +#[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_error.is_some() + || scenario.expected_continuation.is_some() + || all_mock_rows_empty; + assert!( + has_expected_result, + "scenario {} has no explicit expected result (expectedIds, expectedError, \ + expectedContinuation, or an all-empty mock)", + 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(); + } + } +} + +#[test] +fn mock_row_key_counts_match_the_query_column_count() { + let catalog = load_catalog(); + for scenario in &catalog.scenarios { + // Malformed-envelope scenarios deliberately mismatch on purpose. + if scenario.expected_error.is_some() { + continue; + } + let Some(mock) = &scenario.mock else { + continue; + }; + let expected_len = scenario.query.columns.len(); + for (p_idx, partition) in mock.partitions.iter().enumerate() { + for (pg_idx, page) in partition.pages.iter().enumerate() { + for (row_idx, row) in page.rows.iter().enumerate() { + assert_eq!( + row.order_by_items.len(), + expected_len, + "scenario {} partition {p_idx} page {pg_idx} row {row_idx}: \ + orderByItems length does not match the query's {expected_len} \ + ORDER BY column(s)", + scenario.id, + ); + } + } + } + } +} + +#[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; not every + // category needs a dedicated scenario, but each marker must appear. + let required_markers = [ + "single_column_asc", + "single_column_desc", + "multi_column", + "duplicate_keys", + "undefined_field", + "mixed_type_ordering", + "array_keys", + "object_keys", + "skip_count", + "formatted_resume_filter", + "empty_total_result", + "empty_backend_page", + "headers_request_charge", + "split_", + "invalid_rewritten_envelope", + "malformed_token", + "single_logical_partition", + "filters_and_parameters", + "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:?}" + ); + } +}