Cosmos: Add cross-partition vector queries

sdk/cosmos/azure_data_cosmos/CHANGELOG.md

@@ -13,6 +13,7 @@
 - RID-addressed databases skip the extra database read when resolving throughput offers, reusing the addressed RID directly. ([#4687](https://github.com/Azure/azure-sdk-for-rust/pull/4687))
 - Reading and querying items by RID now works end-to-end, including a parent-database cross-check that rejects a container RID belonging to a different database. ([#4687](https://github.com/Azure/azure-sdk-for-rust/pull/4687))
 - Added cross-partition `DISTINCT` query support. `SELECT DISTINCT` now deduplicates structurally equal values across every physical partition and page, rather than failing as an unsupported query feature. A `DISTINCT` query of the exact form `SELECT DISTINCT VALUE <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 finite cross-partition `ORDER BY VectorDistance(...)` queries with `TOP` or `OFFSET`/`LIMIT`. Results are fully buffered before the first page and cannot be resumed from continuation tokens. Hybrid/full-text vector ranking remains unsupported. ([#5130](https://github.com/Azure/azure-sdk-for-rust/pull/5130))
 - 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))

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

@@ -858,15 +858,13 @@ impl ContainerClient {
         ))
     }
 
-    /// Executes a single-partition query against items in the container.
+    /// Executes a query against items in the container.
     ///
     /// The resulting document will be deserialized into the type provided as `T`.
     /// If you want to deserialize the document to a direct representation of the JSON returned, use [`serde_json::Value`] as the target type.
     ///
     /// We recommend using ["turbofish" syntax](https://doc.rust-lang.org/book/appendix-02-operators.html#:~:text=turbofish) (`query_items::<SomeTargetType>(...)`) to specify the target type, as it makes type inference easier.
     ///
-    /// **NOTE:** Currently, the Azure Cosmos DB SDK for Rust only supports single-partition querying. Cross-partition queries may be supported in the future.
-    ///
     /// # Arguments
     ///
     /// * `query` - The query to execute.
@@ -875,9 +873,20 @@ impl ContainerClient {
     ///
     /// # Cross Partition Queries
     ///
-    /// Cross-partition queries are significantly limited in the current version of the Cosmos DB SDK.
-    /// They are run on the gateway and limited to simple projections (`SELECT`) and filtering (`WHERE`).
-    /// For more details, see [the Cosmos DB documentation page on cross-partition queries](https://learn.microsoft.com/en-us/rest/api/cosmos-db/querying-cosmosdb-resources-using-the-rest-api#queries-that-cannot-be-served-by-gateway).
+    /// When `scope` spans multiple partitions, the SDK obtains a query plan and composes the
+    /// client-side pipeline needed to execute it. Supported features include ordinary projections
+    /// and filters, `TOP`, `OFFSET`/`LIMIT`, streaming single- and multiple-column `ORDER BY`, and
+    /// ordered or unordered `DISTINCT`.
+    ///
+    /// Cross-partition vector ordering supports pure `ORDER BY VectorDistance(...)` queries with a
+    /// finite `TOP N` or `OFFSET x LIMIT y` window. The SDK buffers that result window before
+    /// returning the first page, so use narrow projections and choose a result window appropriate
+    /// for the memory available to the application.
+    ///
+    /// The buffered result can be iterated in pages, but it does not support continuation tokens
+    /// for resuming in another process. Aggregates, `GROUP BY`, and hybrid/full-text ranking remain
+    /// unsupported when their query plans require client-side stages that have not been
+    /// implemented.
     ///
     /// # Examples
     ///
@@ -919,6 +928,38 @@ impl ContainerClient {
     /// # }
     /// ```
     ///
+    /// A raw SQL vector search can bind the query vector as a parameter. Vector ordering must not
+    /// specify `ASC` or `DESC`, and must use `TOP N` or `OFFSET`/`LIMIT` when querying across
+    /// partitions:
+    ///
+    /// ```rust,no_run
+    /// # async fn doc() -> Result<(), Box<dyn std::error::Error>> {
+    /// use azure_data_cosmos::{feed::FeedScope, Query};
+    /// use futures::TryStreamExt;
+    /// # let container_client: azure_data_cosmos::clients::ContainerClient = panic!("this is a non-running example");
+    /// #[derive(serde::Deserialize)]
+    /// struct VectorMatch {
+    ///     id: String,
+    ///     score: f64,
+    /// }
+    ///
+    /// let query_vector = vec![0.1_f32, 0.2, 0.3];
+    /// let query = Query::from(
+    ///     "SELECT TOP 5 c.id, VectorDistance(c.embedding, @vector, false) AS score \
+    ///      FROM c ORDER BY VectorDistance(c.embedding, @vector, false)",
+    /// )
+    /// .with_parameter("@vector", &query_vector)?;
+    /// let mut matches = container_client
+    ///     .query_items::<VectorMatch>(query, FeedScope::full_container(), None)
+    ///     .await?;
+    ///
+    /// while let Some(item) = matches.try_next().await? {
+    ///     println!("{}: {}", item.id, item.score);
+    /// }
+    /// # Ok(())
+    /// # }
+    /// ```
+    ///
     /// See [`PartitionKey`](crate::PartitionKey) for more information on how to specify a partition key, and [`Query`] for more information on how to specify a query.
     pub async fn query_items<T: DeserializeOwned + Send + 'static>(
         &self,
@@ -927,6 +968,7 @@ impl ContainerClient {
         options: Option<QueryOptions>,
     ) -> crate::Result<QueryItemIterator<T>> {
         let options = options.unwrap_or_default();
+        let plan_options = options.to_plan_options();
         let query = query.into();
 
         // Resolve binary encoding so the driver advertises a binary *response*
@@ -965,7 +1007,7 @@ impl ContainerClient {
                 initial_operation,
                 &operation_options,
                 options.feed.continuation_token.as_ref(),
-                &options.feed.to_plan_options(),
+                &plan_options,
             )
             .await?;
         Ok(QueryItemIterator::new(

sdk/cosmos/azure_data_cosmos/src/feed/query.rs

@@ -124,6 +124,21 @@ impl FeedScope {
 ///     .with_parameter("@customer_info", CustomerInfo { id: 42, name: "Contoso".into() }).unwrap();
 /// # assert_eq!(serde_json::to_string(&query).unwrap(), "{\"query\":\"\\n    SELECT * FROM c\\n    WHERE c.id = @customer_info.id\\n    AND c.name = @customer_info.name\",\"parameters\":[{\"name\":\"@customer_info\",\"value\":{\"id\":42,\"name\":\"Contoso\"}}]}");
 /// ```
+///
+/// Vectors can be bound as array parameters rather than interpolated into the
+/// query text:
+///
+/// ```rust
+/// # use azure_data_cosmos::Query;
+/// let query_vector = vec![0.1_f32, 0.2, 0.3];
+/// let query = Query::from(
+///     "SELECT TOP 5 c.id, VectorDistance(c.embedding, @vector, false) AS score \
+///      FROM c ORDER BY VectorDistance(c.embedding, @vector, false)",
+/// )
+/// .with_parameter("@vector", &query_vector).unwrap();
+/// # let serialized = serde_json::to_value(&query).unwrap();
+/// # assert_eq!(serialized["parameters"][0]["value"], serde_json::to_value(&query_vector).unwrap());
+/// ```
 #[derive(Clone, Debug, Serialize)]
 pub struct Query {
     /// The query text itself.
@@ -242,6 +257,29 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    pub fn serialize_query_with_float_vector_parameter() -> Result<(), Box<dyn Error>> {
+        let query_vector = vec![0.25_f32, -1.5, 2.0];
+        let query = Query::from(
+            "SELECT TOP 3 * FROM c ORDER BY VectorDistance(c.embedding, @vector, false)",
+        )
+        .with_parameter("@vector", query_vector)?;
+
+        assert_eq!(
+            serde_json::to_value(query)?,
+            serde_json::json!({
+                "query": "SELECT TOP 3 * FROM c ORDER BY VectorDistance(c.embedding, @vector, false)",
+                "parameters": [
+                    {
+                        "name": "@vector",
+                        "value": [0.25, -1.5, 2.0]
+                    }
+                ]
+            })
+        );
+        Ok(())
+    }
+
     #[test]
     pub fn with_text_replaces_query_text() {
         let query = Query::from("SELECT * FROM c").with_text("SELECT c.id FROM c".to_string());

sdk/cosmos/azure_data_cosmos/src/options/feed.rs

@@ -194,6 +194,10 @@ impl QueryOptions {
         self.feed = self.feed.with_continuation_token(continuation_token);
         self
     }
+
+    pub(crate) fn to_plan_options(&self) -> PlanOptions {
+        self.feed.to_plan_options()
+    }
 }
 
 #[cfg(test)]