Cosmos: Add cross-partition vector queries - #5130
Open
Simon Moreno (simorenoh) wants to merge 7 commits into
Open
Cosmos: Add cross-partition vector queries#5130Simon Moreno (simorenoh) wants to merge 7 commits into
Simon Moreno (simorenoh) wants to merge 7 commits into
Conversation
Add live-only coverage for parameterized vector distance queries scoped to one logical partition. Document raw SQL usage and verify float-vector serialization. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Add a bounded non-streaming merge pipeline for finite pure-vector ORDER BY queries across Cosmos partitions. Reject unsupported continuation, hybrid, unbounded, and oversized query shapes before fan-out. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 26e6962d-c1cb-4663-955f-f3427bcfe2b7
|
Azure Pipelines: Successfully started running 1 pipeline(s). 3 pipeline(s) were filtered out due to trigger conditions. There may be pipelines that require an authorized user to comment /azp run to run. |
Contributor
There was a problem hiding this comment.
Pull request overview
Adds bounded cross-partition vector query execution to the Cosmos SDK and driver.
Changes:
- Adds SDK and driver buffer-limit configuration APIs.
- Implements non-streaming vector ordering, pagination, metrics aggregation, and typed errors.
- Adds unit, documentation, and live-account coverage.
Reviewed changes
Copilot reviewed 17 out of 17 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
azure_data_cosmos/tests/emulator_tests/mod.rs |
Registers vector tests. |
azure_data_cosmos/tests/emulator_tests/cosmos_vector_query.rs |
Tests vector query scenarios. |
azure_data_cosmos/src/options/feed.rs |
Exposes buffer-limit options. |
azure_data_cosmos/src/feed/query.rs |
Documents vector parameters. |
azure_data_cosmos/src/clients/container_client.rs |
Documents cross-partition vectors. |
azure_data_cosmos/CHANGELOG.md |
Records SDK support. |
azure_data_cosmos_driver/src/query/mod.rs |
Classifies vector query shapes. |
azure_data_cosmos_driver/src/query_plan_native/mod.rs |
Requests formattable rewrites. |
azure_data_cosmos_driver/src/options/plan_options.rs |
Defines driver buffer limits. |
azure_data_cosmos_driver/src/options/mod.rs |
Exports new options. |
azure_data_cosmos_driver/src/error/cosmos_status.rs |
Adds typed statuses. |
azure_data_cosmos_driver/src/driver/dataflow/query_response.rs |
Aggregates response metadata. |
azure_data_cosmos_driver/src/driver/dataflow/planner.rs |
Builds buffered merge plans. |
azure_data_cosmos_driver/src/driver/dataflow/non_streaming_ordered_merge.rs |
Implements bounded ordered merge. |
azure_data_cosmos_driver/src/driver/dataflow/mod.rs |
Registers the merge stage. |
azure_data_cosmos_driver/src/driver/cosmos_driver.rs |
Routes non-streaming plans. |
azure_data_cosmos_driver/CHANGELOG.md |
Records driver support. |
💡 Add a code-review agent skill for context-aware, tailored reviews. Learn more in the docs.
Port the valid 300-dimensional .NET vector corpus into live cross-partition coverage for all supported distance functions and both indexed and brute-force execution. Remove the unreleased public buffer-ceiling option while retaining mandatory finite TOP or OFFSET/LIMIT windows. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 26e6962d-c1cb-4663-955f-f3427bcfe2b7
Cap retained query and index metrics by entry count and byte size, and mark truncated output explicitly so buffered fan-out diagnostics cannot grow with backend page count. Make the vector parameter doctest compare against the serialized f32 source values, avoiding architecture-independent precision mismatches. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 26e6962d-c1cb-4663-955f-f3427bcfe2b7
Keep the vector-search PR focused by restoring PageAggregator query and index metrics handling exactly to main. Cross-page metrics aggregation and synthetic truncation markers belong in a separate design change. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 26e6962d-c1cb-4663-955f-f3427bcfe2b7
Member
Author
|
/azp run rust - cosmos - weekly |
|
Azure Pipelines: Successfully started running 1 pipeline(s). |
Simon Moreno (simorenoh)
changed the title
Simon Moreno (simorenoh)
marked this pull request as ready for review
August 27, 2026 14:13
|
Azure Pipelines: Successfully started running 1 pipeline(s). 3 pipeline(s) were filtered out due to trigger conditions. There may be pipelines that require an authorized user to comment /azp run to run. |
Extract comparator-driven binary heap operations for the streaming and buffered ordered merges so both pipelines share the same mechanics while retaining their runtime ordering rules. Update query_items documentation to describe the current cross-partition query surface and distinguish finite buffered vector ordering from streaming features. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 26e6962d-c1cb-4663-955f-f3427bcfe2b7
Route buffered ORDER BY queries from authoritative plan metadata instead of reparsing SQL locally. Match peer SDK native plan options while retaining the existing finite-window and unsupported-feature guards. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 26e6962d-c1cb-4663-955f-f3427bcfe2b7
Member
Author
|
/azp run rust - cosmos - weekly |
|
Azure Pipelines: Successfully started running 1 pipeline(s). |
| //! Comparator-driven binary-heap operations used by ordered merge stages. | ||
|
|
||
| /// Pushes `item` while preserving the heap defined by `precedes`. | ||
| pub(crate) fn push_by<T>(heap: &mut Vec<T>, item: T, precedes: impl Fn(&T, &T) -> bool) { |
There was a problem hiding this comment.
Could we use std::collections::BinaryHeap instead of a full custom implementation
|
|
||
| TestClient::run_with_unique_db( | ||
| async |run_context, db_client| { | ||
| let container = seed_vector_container(run_context, db_client).await?; |
There was a problem hiding this comment.
from what I can tell this is creating a container with one partition, so it is not exercising the cross partition portion. Same with this test precomputed_pure_vector_search
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds finite-window, cross-partition
ORDER BY VectorDistance(...)execution to the Cosmos DB Rust driver. Plans markedhasNonStreamingOrderBypreviously failed because the SDK could not combine per-partition results into a global ranking.Design
NonStreamingOrderedMergedrains the Gateway/native provider's rewritten query across every target range and retains the bestKrows in a max-heap, whereKisTOPorOFFSET + LIMIT. Routing trusts authoritative query-plan metadata rather than reparsing SQL, matching .NET, Python, Java, and JavaScript. The merge supports plan-provided sort-key tuples and preserves stable service arrival order for ties.TOPorOFFSET/LIMITorderingDISTINCT, aggregate,GROUP BY, or hybrid plansThe requested window is fully buffered before the first result page. The SDK imposes no additional item ceiling; callers control memory through the finite window and projected row size.
max_item_countcontrols emitted pages after the merge, not the buffered window.Native query planning always requests the formattable, non-continuation shape used by .NET. The driver replaces the Gateway's
{documentdb-formattableorderbyquery-filter}placeholder before execution. Request charges, session tokens, and diagnostics retain the existing page-aggregation behavior; complete cross-page query/index metrics are tracked separately in #5170.Validation
The live suite covers single- and cross-partition vector queries with filters, bound vectors, indexed and brute-force execution,
TOP,OFFSET/LIMIT, small output pages, and duplicate/missing-row detection. It also validates Euclidean, Cosine, and DotProduct ranking against nine valid 300-dimensional embeddings adapted from .NET'sNonStreamingOrderByQueryTests; the malformed 279-dimensionalyear_NOUNfixture is documented and omitted.