Cosmos: Adds Binary Encoding Support for Queries and Streaming order_by

sdk/cosmos/AGENTS.md

@@ -534,6 +534,32 @@ Every public API should document:
 - **Performance**: RU/s implications, if relevant
 - **Partition Key**: Whether the operation is partition-scoped
 
+### Comment Brevity (IMPORTANT)
+
+Keep comments **short and dense**. Long, chatty comment blocks are noise: they age badly, bury the
+signal, and read as filler.
+
+- **Inline comments**: 1–2 lines. 3 is the hard ceiling. Explain *why*, never *what* the code already says.
+- **Doc comments**: lead with a one-sentence summary. Add detail only when a caller genuinely needs it
+  to use the API correctly.
+- **Test comments**: usually unnecessary — a good test name plus a clear assertion message says it.
+  Comment only a non-obvious setup or a subtle invariant.
+- **Never** restate an argument three ways, narrate the diff ("this used to be X, now it's Y"),
+  re-explain something already covered by a nearby doc comment, or write a mini design doc inline.
+  Link to the spec/HLD instead.
+- When a rationale really needs paragraphs, it belongs in `docs/`, not in the source.
+
+```rust
+// ❌ BAD: five lines to say one thing
+// This is a client-side re-encode of a payload the driver itself produced,
+// not a malformed service envelope, so it is a serialization fault (not the
+// 500 `SERVICE_ORDER_BY_ENVELOPE_INVALID`). Include the item ordinal to
+// correlate on a large page.
+
+// ✅ GOOD
+// Client-side re-encode failure, not a malformed service envelope.
+```
+
 ## Additional Resources
 
 - [Azure Cosmos DB REST API Reference](https://learn.microsoft.com/rest/api/cosmos-db/)

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 an SDK-generated `x-ms-client-id` header that remains stable for each `CosmosClient`. ([#4844](https://github.com/Azure/azure-sdk-for-rust/pull/4844))
 - Added opt-in Cosmos binary JSON encoding for item operations (`create`/`read`/`replace`/`upsert`). Enable it via `CosmosClientBuilder::with_binary_encoding_options` (or the `AZURE_COSMOS_BINARY_ENCODING_ENABLED` environment-variable fallback). Off by default; when disabled, requests and responses are byte-for-byte unchanged. ([#4671](https://github.com/Azure/azure-sdk-for-rust/pull/4671))
+- Extended binary JSON encoding to `query_items`: when binary encoding is enabled, queries negotiate a binary response (the request body stays text `application/query+json`) and decode binary feed pages, including the streaming cross-partition `ORDER BY` merge. Off by default and negotiated on the standard-gateway path. ([#5040](https://github.com/Azure/azure-sdk-for-rust/pull/5040))
 - 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))
 - Added a pluggable client-side diagnostics emission layer — the `DiagnosticsHandler` trait and ordered `DiagnosticsHandlerChain` (registered via `CosmosClientBuilder::with_diagnostics_handler`) — invoked once per operation (singleton and paginated, on success and failure) with the completed `DiagnosticsContext` plus an SDK-supplied `CosmosOperationContext`; the empty default chain is a zero-overhead no-op. ([#4789](https://github.com/Azure/azure-sdk-for-rust/pull/4789))

sdk/cosmos/azure_data_cosmos/src/clients/container_client.rs

@@ -922,6 +922,14 @@ impl ContainerClient {
         let options = options.unwrap_or_default();
         let query = query.into();
 
+        // Resolve binary encoding so the driver advertises a binary *response*
+        // via the negotiation header. Unlike point item writes, the query
+        // request body stays text (`application/query+json` is a query spec,
+        // not a document), so we do not touch body serialization here — the
+        // driver's request-body gate excludes query.
+        let (operation_options, _binary) =
+            resolve_binary_encoding(options.operation, &self.context.binary_encoding);
+
         let container_ref = self.container_ref.clone();
 
         // The first operation to execute in the query items flow.
@@ -948,7 +956,7 @@ impl ContainerClient {
             .driver
             .plan_operation(
                 initial_operation,
-                &options.operation,
+                &operation_options,
                 options.feed.continuation_token.as_ref(),
                 &options.feed.to_plan_options(),
             )
@@ -957,7 +965,7 @@ impl ContainerClient {
             self.context.driver.clone(),
             Some(self.container_ref.clone()),
             plan,
-            options.operation,
+            operation_options,
             self.context.diagnostics_handlers.clone(),
             self.operation_context("query_items"),
         ))

sdk/cosmos/azure_data_cosmos/tests/binary_roundtrip_fuzzer.rs

@@ -45,9 +45,11 @@ 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::{Deserialize, Serialize};
 use serde_json::{Map, Number, Value};
 use sha2::{Digest, Sha256};
@@ -1917,6 +1919,34 @@ async fn assert_typed_integer_probe(
         got, sent,
         "{context}: typed integer probe round-trip changed"
     );
+
+    // Also decode the same wide-integer probe through a full-container binary
+    // ORDER BY query, driving the streaming-merge envelope decode (`build_page`
+    // re-encodes to binary so `deserialize_integer` runs). A merge that emitted
+    // text would fail here with `invalid type: floating point, expected u64`.
+    //
+    // This container has default throughput (one physical partition), so the
+    // merge runs with a single child — multi-child interleave is covered by the
+    // 3-partition emulator tests.
+    let order_by = with_transient_retry("int-probe-order-by", context, || async {
+        let query = Query::from("SELECT * FROM c WHERE c.id = @id ORDER BY c.id")
+            .with_parameter("@id", id.as_str())?;
+        let iter = container
+            .query_items::<IntProbe>(query, FeedScope::full_container(), None)
+            .await?;
+        iter.try_collect::<Vec<_>>().await
+    })
+    .await?;
+    assert_eq!(
+        order_by.len(),
+        1,
+        "{context}: typed integer ORDER BY probe expected exactly one item, got {}",
+        order_by.len(),
+    );
+    assert_eq!(
+        order_by[0], sent,
+        "{context}: typed integer probe changed through the binary ORDER BY merge"
+    );
     Ok(())
 }
 
@@ -1969,6 +1999,94 @@ where
     }
 }
 
+/// Queries the just-written item back and asserts it round-trips, covering the
+/// query binary-response decode path (which point ops do not exercise). Always
+/// runs a single-partition query (the passthrough decode path). When
+/// `include_order_by` is set, also runs a full-container streaming `ORDER BY`
+/// query, whose per-page envelope decode is the binary path added for query
+/// support; it adds no binary coverage on text configs, so callers gate it.
+///
+/// The live container has default throughput (one physical partition), so the
+/// merge runs with a single child — multi-child interleave is covered by the
+/// 3-partition emulator tests.
+///
+/// Both filter on the unique `id`, so each returns exactly this item.
+// The sent-value trio (`sent_canon`, `sent_hash`, `doc`) plus routing context is
+// threaded verbatim into `assert_query_hit`/`assert_roundtrip`; bundling it into
+// a struct here would only move the argument count to those helpers, so the
+// eighth parameter is allowed on this test-only function.
+#[allow(clippy::too_many_arguments)]
+async fn assert_query_roundtrip(
+    container: &ContainerClient,
+    pk: &str,
+    id: &str,
+    sent_canon: &str,
+    sent_hash: &[u8; 32],
+    doc: &Map<String, Value>,
+    context: &str,
+    include_order_by: bool,
+) -> Result<usize, Box<dyn Error>> {
+    let single_partition = with_transient_retry("query", context, || async {
+        let query = Query::from("SELECT * FROM c WHERE c.id = @id").with_parameter("@id", id)?;
+        let iter = container
+            .query_items::<Value>(query, FeedScope::partition(pk.to_string()), None)
+            .await?;
+        iter.try_collect::<Vec<_>>().await
+    })
+    .await?;
+    assert_query_hit(
+        &single_partition,
+        doc,
+        sent_canon,
+        sent_hash,
+        context,
+        "query",
+    );
+
+    if !include_order_by {
+        return Ok(1);
+    }
+
+    let order_by = with_transient_retry("query-order-by", context, || async {
+        let query = Query::from("SELECT * FROM c WHERE c.id = @id ORDER BY c.id")
+            .with_parameter("@id", id)?;
+        let iter = container
+            .query_items::<Value>(query, FeedScope::full_container(), None)
+            .await?;
+        iter.try_collect::<Vec<_>>().await
+    })
+    .await?;
+    assert_query_hit(
+        &order_by,
+        doc,
+        sent_canon,
+        sent_hash,
+        context,
+        "query-order-by",
+    );
+
+    Ok(2)
+}
+
+/// Asserts a query returned exactly the one expected item and that it
+/// round-trips against the sent canonical form.
+fn assert_query_hit(
+    results: &[Value],
+    doc: &Map<String, Value>,
+    sent_canon: &str,
+    sent_hash: &[u8; 32],
+    context: &str,
+    phase: &str,
+) {
+    assert_eq!(
+        results.len(),
+        1,
+        "{context}: {phase} expected exactly one item, got {}",
+        results.len()
+    );
+    assert_roundtrip(doc, &results[0], sent_canon, sent_hash, context, phase);
+}
+
 // ─────────────────────────────────────────────────────────────────────────────
 // The fuzzer test
 // ─────────────────────────────────────────────────────────────────────────────
@@ -2176,6 +2294,25 @@ async fn binary_encoding_roundtrip_fuzz() -> Result<(), Box<dyn Error>> {
             // Four point-op round-trips this config: create, read, replace, upsert.
             checked += 4;
 
+            // QUERY the item back. The single-partition passthrough query runs on
+            // every config; the expensive ORDER BY merge runs only on pure
+            // `binary` — `text-control` adds no binary coverage, and
+            // `binary+text-response` suppresses query negotiation entirely, so
+            // both would merge over text.
+            let include_order_by = *label == "binary";
+            let queries_checked = assert_query_roundtrip(
+                &container,
+                &pk,
+                &id,
+                &sent_canon,
+                &sent_hash,
+                &doc,
+                &context,
+                include_order_by,
+            )
+            .await?;
+            checked += queries_checked as u64;
+
             // 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
@@ -2194,7 +2331,7 @@ async fn binary_encoding_roundtrip_fuzz() -> Result<(), Box<dyn Error>> {
     }
 
     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 each + 1–2 queries; full-container ORDER BY on binary configs only) = {checked} round-trips, all canonical-equal (seed={})",
         cfg.iterations,
         configs.len(),
         cfg.seed