Add binary encoding for Cosmos query pages

sdk/cosmos/azure_data_cosmos/CHANGELOG.md

@@ -11,6 +11,7 @@
 - Reading and querying items by RID now works end-to-end, including a parent-database cross-check that rejects a container RID belonging to a different database. ([#4687](https://github.com/Azure/azure-sdk-for-rust/pull/4687))
 - Added cross-partition `DISTINCT` query support. `SELECT DISTINCT` now deduplicates structurally equal values across every physical partition and page, rather than failing as an unsupported query feature. A `DISTINCT` query of the exact form `SELECT DISTINCT VALUE <path> … ORDER BY <same path>` (for example `SELECT DISTINCT VALUE c.city FROM c ORDER BY c.city`) is resumable from a continuation token; every other shape — including a list projection such as `SELECT DISTINCT c.city …` and any multi-column `ORDER BY` — is not, and requesting a token for it returns an error explaining how to rewrite the query. ([#5026](https://github.com/Azure/azure-sdk-for-rust/pull/5026))
 - Added an SDK-generated `x-ms-client-id` header that remains stable for each `CosmosClient`. ([#4844](https://github.com/Azure/azure-sdk-for-rust/pull/4844))
+- Added opt-in Cosmos binary JSON encoding for SQL query pages (including cross-partition `DISTINCT` and streaming `ORDER BY`) and document read feeds. Query plans and change feeds remain text.
 - Added opt-in Cosmos binary JSON encoding for item operations (`create`/`read`/`replace`/`upsert`). Enable it via `CosmosClientBuilder::with_binary_encoding_options` (or the `AZURE_COSMOS_BINARY_ENCODING_ENABLED` environment-variable fallback). Off by default; when disabled, requests and responses are byte-for-byte unchanged. ([#4671](https://github.com/Azure/azure-sdk-for-rust/pull/4671))
 - Added `FeedOptions::max_fan_out` (and `FeedOptions::with_max_fan_out`) to cap how many physical partitions a cross-partition query or change feed may fan out to. Applies to `ContainerClient::query_items` and `ContainerClient::query_change_feed`. The cap is enforced only at initial query setup; a partition that splits mid-execution and pushes the fan-out higher does not abort the operation. ([#4855](https://github.com/Azure/azure-sdk-for-rust/pull/4855))
 - Added resumable cross-partition streaming `ORDER BY` query support. ([#4800](https://github.com/Azure/azure-sdk-for-rust/pull/4800))

sdk/cosmos/azure_data_cosmos/tests/binary_roundtrip_fuzzer.rs

@@ -45,9 +45,12 @@ use azure_data_cosmos::options::{
     OperationOptions, Region, ServerCertificateValidation,
 };
 use azure_data_cosmos::{
-    AccountEndpoint, AccountReference, CosmosClient, CosmosRuntime, RoutingStrategy, SubStatusCode,
+    AccountEndpoint, AccountReference, CosmosClient, CosmosRuntime, FeedScope, Query,
+    RoutingStrategy, SubStatusCode,
 };
 use azure_data_cosmos_driver::models::ConnectionString;
+use futures::TryStreamExt;
+use serde::de::DeserializeOwned;
 use serde::{Deserialize, Serialize};
 use serde_json::{Map, Number, Value};
 use sha2::{Digest, Sha256};
@@ -1880,6 +1883,11 @@ struct IntProbe {
     wide: u64,
 }
 
+#[derive(Deserialize, Debug, PartialEq)]
+struct QueryIntProbe {
+    int: i64,
+}
+
 /// Round-trips an [`IntProbe`] so `deserialize_integer` (the production change)
 /// is exercised live, then asserts the typed values survived. Only meaningful on
 /// the pure-binary config (see the call site).
@@ -1964,11 +1972,108 @@ where
                 );
                 tokio::time::sleep(backoff).await;
             }
+
             Err(e) => return Err(format!("{context}: {op_name} failed: {e}").into()),
         }
     }
 }
 
+async fn query_values<T: DeserializeOwned + Send + 'static>(
+    container: &ContainerClient,
+    sql: &str,
+    run_id: &str,
+    context: &str,
+) -> Result<Vec<T>, Box<dyn Error>> {
+    let mut attempt = 0;
+    loop {
+        attempt += 1;
+        let query = Query::from(sql).with_parameter("@run", run_id)?;
+        let result = match container
+            .query_items(query, FeedScope::full_container(), None)
+            .await
+        {
+            Ok(iterator) => Box::pin(iterator.try_collect()).await,
+            Err(err) => Err(err),
+        };
+        match result {
+            Ok(values) => return Ok(values),
+            Err(err) if is_transient(&err) && attempt < MAX_OP_ATTEMPTS => {
+                let backoff =
+                    std::time::Duration::from_millis(200u64 * (1u64 << (attempt - 1)).min(16));
+                eprintln!(
+                    "{context}: query transient failure (attempt {attempt}/{MAX_OP_ATTEMPTS}), \
+                     retrying in {backoff:?}: {err}"
+                );
+                tokio::time::sleep(backoff).await;
+            }
+            Err(err) => return Err(format!("{context}: query failed: {err}").into()),
+        }
+    }
+}
+
+fn canonical_query_results(values: Vec<Value>, ordered: bool) -> Vec<String> {
+    let mut canonical: Vec<String> = values
+        .into_iter()
+        .map(|value| {
+            let value = match value {
+                Value::Object(mut map) => {
+                    strip_reserved_fields(&mut map);
+                    Value::Object(map)
+                }
+                value => value,
+            };
+            structural_query_key(&value)
+        })
+        .collect();
+    if !ordered {
+        canonical.sort();
+    }
+    canonical
+}
+
+fn structural_query_key(value: &Value) -> String {
+    fn write(value: &Value, output: &mut String) {
+        match value {
+            Value::Null => output.push('n'),
+            Value::Bool(value) => output.push(if *value { 't' } else { 'f' }),
+            Value::Number(number) => {
+                if let Some(value) = number.as_i64() {
+                    output.push_str(&format!("i:{value};"));
+                } else if let Some(value) = number.as_u64() {
+                    output.push_str(&format!("u:{value};"));
+                } else {
+                    let value = number.as_f64().expect("finite JSON number");
+                    output.push_str(&format!("d:{:016x};", value.to_bits()));
+                }
+            }
+            Value::String(value) => {
+                output.push_str("s:");
+                output.push_str(&serde_json::to_string(value).expect("string serializes"));
+                output.push(';');
+            }
+            Value::Array(values) => {
+                output.push('[');
+                values.iter().for_each(|value| write(value, output));
+                output.push(']');
+            }
+            Value::Object(values) => {
+                output.push('{');
+                let mut entries: Vec<_> = values.iter().collect();
+                entries.sort_unstable_by_key(|(key, _)| *key);
+                for (key, value) in entries {
+                    write(&Value::String(key.clone()), output);
+                    write(value, output);
+                }
+                output.push('}');
+            }
+        }
+    }
+
+    let mut output = String::new();
+    write(value, &mut output);
+    output
+}
+
 // ─────────────────────────────────────────────────────────────────────────────
 // The fuzzer test
 // ─────────────────────────────────────────────────────────────────────────────
@@ -2027,7 +2132,8 @@ async fn binary_encoding_roundtrip_fuzz() -> Result<(), Box<dyn Error>> {
         // below — the same document stored under multiple configs would
         // otherwise collide on the `(pk, id)` key and fail with 409 Conflict.
         let base_doc = gen_object(&mut rng, &cfg);
-        let pk = format!("pk-{}", rng.below(16));
+        let partition_bucket = rng.below(16);
+        let run_id = format!("fuzz-{:016x}-{iter}", cfg.seed);
 
         // Optionally print the generated document (pretty JSON) so a run can be
         // eyeballed. Enable with `AZURE_COSMOS_FUZZ_PRINT=true`.
@@ -2050,10 +2156,14 @@ async fn binary_encoding_roundtrip_fuzz() -> Result<(), Box<dyn Error>> {
             // touching the document RNG stream so a rerun with the same
             // AZURE_COSMOS_FUZZ_SEED reproduces the exact document *and* its
             // canonical form (a random `Uuid` here would defeat that promise).
-            let id = format!("fuzz-{:016x}-{iter}-{config_idx}", cfg.seed);
+            let id = format!("{run_id}-{config_idx}");
+            // Reuse 16 logical partitions per encoding config so point
+            // operations cover multiple items sharing the same partition key.
+            let pk = format!("fuzz-pk-{partition_bucket}-{config_idx}");
             let mut doc = base_doc.clone();
             doc.insert("id".to_string(), Value::String(id.clone()));
             doc.insert("pk".to_string(), Value::String(pk.clone()));
+            doc.insert("fuzzRun".to_string(), Value::String(run_id.clone()));
             // Never send Cosmos-reserved system properties (`_rid`, `_self`,
             // `_etag`, `_ts`, `_attachments`): the service **owns** these and
             // overwrites/assigns them, so a random value we send would come
@@ -2179,22 +2289,87 @@ async fn binary_encoding_roundtrip_fuzz() -> Result<(), Box<dyn Error>> {
             // The four ops above decode into `serde_json::Value` (→
             // `deserialize_any`), so they do NOT cover the native typed-integer
             // path (`deserialize_integer`) this PR ships. A typed probe covers it
-            // live on the pure-binary config — the only mode that returns binary
-            // for an integer field (text modes return text, which `serde_json`
-            // rejects into an integer).
+            // live on the pure-binary config — the only mode that exercises the
+            // binary deserializer's integral-Double coercion directly.
             if *label == "binary" {
                 assert_typed_integer_probe(&container, &pk, iter, cfg.seed, &context).await?;
                 checked += 1;
             }
         }
 
+        let query_cases = [
+            (
+                "SELECT * FROM c WHERE c.fuzzRun = @run",
+                false,
+                "select-all",
+            ),
+            (
+                "SELECT DISTINCT VALUE c._sampler.int FROM c WHERE c.fuzzRun = @run",
+                false,
+                "distinct",
+            ),
+            (
+                "SELECT DISTINCT VALUE c._sampler.int FROM c WHERE c.fuzzRun = @run \
+                 ORDER BY c._sampler.int",
+                true,
+                "distinct-order-by",
+            ),
+        ];
+        for (sql, ordered, phase) in query_cases {
+            let mut expected: Option<Vec<String>> = None;
+            for (label, client) in &clients {
+                let container = client
+                    .database_client(&database_name)
+                    .container_client(&container_name)
+                    .await?;
+                let context = format!("iter={iter} config={label} query={phase} seed={}", cfg.seed);
+                let actual = canonical_query_results(
+                    query_values(&container, sql, &run_id, &context).await?,
+                    ordered,
+                );
+                if let Some(expected) = &expected {
+                    assert_eq!(
+                        &actual, expected,
+                        "{context}: query result diverged across encoding configurations"
+                    );
+                } else {
+                    expected = Some(actual);
+                }
+
+                checked += 1;
+            }
+        }
+
+        for (label, client) in &clients {
+            let container = client
+                .database_client(&database_name)
+                .container_client(&container_name)
+                .await?;
+            let context = format!(
+                "iter={iter} config={label} query=typed-integer seed={}",
+                cfg.seed
+            );
+            let values: Vec<QueryIntProbe> = query_values(
+                &container,
+                "SELECT VALUE {\"int\": 7} FROM c WHERE c.fuzzRun = @run",
+                &run_id,
+                &context,
+            )
+            .await?;
+            assert!(
+                !values.is_empty() && values.iter().all(|value| value.int == 7),
+                "{context}: typed integer query returned unexpected values: {values:?}"
+            );
+            checked += 1;
+        }
+
         if (iter + 1) % 100 == 0 {
             println!("... {} iterations, {checked} round-trips OK", iter + 1);
         }
     }
 
     println!(
-        "binary_roundtrip_fuzzer: DONE — {} documents × {} configs × 4 point ops = {checked} round-trips, all canonical-equal (seed={})",
+        "binary_roundtrip_fuzzer: DONE — {} documents × {} configs × 4 point ops + 4 queries/config = {checked} canonical comparisons, all equal (seed={})",
         cfg.iterations,
         configs.len(),
         cfg.seed
@@ -2444,6 +2619,14 @@ mod tests {
         assert_eq!(canon(&serde_json::json!(-0.0)), "0");
     }
 
+    #[test]
+    fn query_comparison_preserves_integral_float_representation() {
+        assert_ne!(
+            structural_query_key(&serde_json::json!(1)),
+            structural_query_key(&serde_json::json!(1.0))
+        );
+    }
+
     #[test]
     fn canonicalize_keeps_non_integral_floats() {
         assert_eq!(canon(&serde_json::json!(3.5)), "3.5");

sdk/cosmos/azure_data_cosmos/tests/split_tests/cosmos_query_distinct_split.rs

@@ -1,7 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-//! Live-only split coverage for cross-partition `DISTINCT`.
+//! Live split coverage for text and binary cross-partition `DISTINCT`.
 //!
 //! Neither .NET nor Java tests `DISTINCT` against a real partition split
 //! (.NET's `FullPipelineTests.TestMerge` covers `ORDER BY` only), so this is
@@ -39,7 +39,7 @@ use azure_data_cosmos::{
     clients::ContainerClient,
     feed::FeedScope,
     models::{ContainerProperties, CosmosStatus, ThroughputProperties},
-    options::{MaxItemCountHint, QueryOptions},
+    options::{BinaryEncodingOptions, MaxItemCountHint, QueryOptions},
 };
 use framework::{TestClient, TestOptions};
 use futures::StreamExt;
@@ -124,12 +124,9 @@ where
 /// Part 2 captures an ordered `DISTINCT` continuation token before the split
 /// (already taken, since the split happened in part 1) and resumes it against
 /// the post-split topology.
-#[tokio::test]
-#[cfg_attr(
-    not(test_category = "split"),
-    ignore = "requires test_category 'split'"
-)]
-pub async fn distinct_query_across_split_returns_each_value_once() -> Result<(), Box<dyn Error>> {
+async fn run_distinct_query_across_split_returns_each_value_once(
+    binary: bool,
+) -> Result<(), Box<dyn Error>> {
     TestClient::run_with_unique_db(
         async |run_context, db_client| {
             let properties =
@@ -316,7 +313,33 @@ pub async fn distinct_query_across_split_returns_each_value_once() -> Result<(),
         },
         // A real split takes minutes; the 80s default would abort mid-poll.
         // Matches the other split tests in this directory.
-        Some(TestOptions::new().with_timeout(Duration::from_secs(40 * 60))),
+        Some({
+            let options = TestOptions::new().with_timeout(Duration::from_secs(40 * 60));
+            if binary {
+                options.with_binary_encoding(BinaryEncodingOptions::new().with_enabled(true))
+            } else {
+                options
+            }
+        }),
     )
     .await
 }
+
+#[tokio::test]
+#[cfg_attr(
+    not(test_category = "split"),
+    ignore = "requires test_category 'split'"
+)]
+pub async fn distinct_query_across_split_returns_each_value_once() -> Result<(), Box<dyn Error>> {
+    run_distinct_query_across_split_returns_each_value_once(false).await
+}
+
+#[tokio::test]
+#[cfg_attr(
+    not(test_category = "split"),
+    ignore = "requires test_category 'split'"
+)]
+pub async fn binary_distinct_query_across_split_returns_each_value_once(
+) -> Result<(), Box<dyn Error>> {
+    run_distinct_query_across_split_returns_each_value_once(true).await
+}

sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md

@@ -4,6 +4,8 @@
 
 ### Features Added
 
+- Added Cosmos binary JSON negotiation for SQL query pages and document read feeds. Query plans and change feeds remain text.
+- Extended the binary round trip fuzzer with plain, `DISTINCT`, and `DISTINCT` + `ORDER BY` query parity across text, binary, and binary-with-text-response modes.
 - Added binary round trip fuzzer. As a part of the implementation, binary JSON responses now deserialize a service-echoed integral `Double` into a signed or unsigned integer field (previously a type error); this is intentionally lossy for integers the service cannot represent exactly, while a fractional `Double` remains a type error. Does not yet cover integer elements inside a uniform `Float64` array or an enum variant. ([#4976](https://github.com/Azure/azure-sdk-for-rust/pull/4976))
 - Added driver-internal resolution of containers by resource id (RID). `CosmosDriver::resolve_container_by_rid` reads a container's metadata addressing it purely by RID (deriving the parent database RID from the container RID, so no `read_database` round-trip is needed) and caches the result in a by-RID index. References are validated for consistent name/RID addressing in `plan_operation` — the single choke point every executable operation passes through, including multi-page queries — returning a deterministic `CLIENT_MIXED_NAME_RID_ADDRESSING` error before signing instead of letting the gateway reject a mixed name/RID request with an opaque `401`. The new `CLIENT_INVALID_RESOURCE_ID` and `CLIENT_MIXED_NAME_RID_ADDRESSING` client statuses carry searchable names for diagnostics. ([#4663](https://github.com/Azure/azure-sdk-for-rust/pull/4663))
 - Added `models::is_database_rid`, which reports whether a RID string decodes to a database-level RID (4 bytes). Lets callers that reuse a supplied RID as a database identity reject a wrong-hierarchy RID before it addresses the wrong resource. ([#4640](https://github.com/Azure/azure-sdk-for-rust/pull/4640))