Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
59 commits
Select commit Hold shift + click to select a range
53b1f9f
Local QueryPlan generation and emulator execution
FabianMeiswinkel Apr 29, 2026
6618876
Fixed build issues
FabianMeiswinkel Apr 29, 2026
ebe227a
Merge branch 'release/azure_data_cosmos-previews' of https://github.c…
FabianMeiswinkel Apr 29, 2026
bbce203
Build updates
FabianMeiswinkel Apr 29, 2026
150973b
Apply suggestion from @Copilot
FabianMeiswinkel Apr 29, 2026
067cc09
Apply suggestion from @FabianMeiswinkel
FabianMeiswinkel Apr 29, 2026
7960e5e
Build fixes
FabianMeiswinkel Apr 29, 2026
c5abeb5
Merge branch 'users/fabianm/LocalQueryPlanPOC' of https://github.com/…
FabianMeiswinkel Apr 29, 2026
3425f2a
Update sdk/cosmos/query-engine-porting-plan.md
FabianMeiswinkel Apr 29, 2026
e1413e7
Update sdk/cosmos/azure_data_cosmos/src/query/plan.rs
FabianMeiswinkel Apr 29, 2026
3eec8a5
Addressing code review feedback
FabianMeiswinkel Apr 30, 2026
80b6244
Merge branch 'users/fabianm/LocalQueryPlanPOC' of https://github.com/…
FabianMeiswinkel Apr 30, 2026
3a8476a
Update mod.rs
FabianMeiswinkel Apr 30, 2026
3c62f3c
Iterating
FabianMeiswinkel Apr 30, 2026
d9752e0
More Tests
FabianMeiswinkel Apr 30, 2026
5e0b15a
Fixed sanity check errors
FabianMeiswinkel Apr 30, 2026
beca9ff
Update mod.rs
FabianMeiswinkel Apr 30, 2026
18a0229
Merge branch 'release/azure_data_cosmos-previews' into users/fabianm/…
FabianMeiswinkel Apr 30, 2026
9c26a24
Update Cargo.toml
FabianMeiswinkel Apr 30, 2026
5372b82
Update pack.yml
FabianMeiswinkel Apr 30, 2026
af6fea3
Update Cargo.toml
FabianMeiswinkel Apr 30, 2026
f92bfc8
Merge branch 'release/azure_data_cosmos-previews' of https://github.c…
FabianMeiswinkel May 4, 2026
a29ca01
Refactored query submodule to live within the driver crate
FabianMeiswinkel May 4, 2026
f4e9360
Iterating
FabianMeiswinkel May 4, 2026
4e7333a
Fix left-overs after refactoring
FabianMeiswinkel May 4, 2026
eb2b331
Update query-engine-porting-plan.md
FabianMeiswinkel May 4, 2026
4d4348a
Fix review comments
FabianMeiswinkel May 4, 2026
a989193
Fixing build issues
FabianMeiswinkel May 4, 2026
db69e3e
Adding back comparison tests against Gateway
FabianMeiswinkel May 4, 2026
f503fa3
Fixing parameterized TOP, OFFSET and LIMIT
FabianMeiswinkel May 4, 2026
17fc107
Adresses CR feedback
FabianMeiswinkel May 4, 2026
d291f96
Fix cspell errors
FabianMeiswinkel May 4, 2026
57d7ce8
Update location_state_store.rs
FabianMeiswinkel May 5, 2026
22a0d99
Update location_state_store.rs
FabianMeiswinkel May 5, 2026
9075c2b
Adding regression test
FabianMeiswinkel May 5, 2026
11410ea
Update cosmos_operation.rs
FabianMeiswinkel May 5, 2026
bbd9b49
Make regression test more robust
FabianMeiswinkel May 5, 2026
eb99322
CR feedback
FabianMeiswinkel May 5, 2026
f911413
Merge branch 'release/azure_data_cosmos-previews' of https://github.c…
FabianMeiswinkel May 5, 2026
ed27863
CR fixes
FabianMeiswinkel May 5, 2026
0b01895
Addressing code review feedback
FabianMeiswinkel May 6, 2026
2adf69c
Fixed code review feedback
FabianMeiswinkel May 6, 2026
78cc3d2
Small refactoring to gate eval submodule behind feature flag
FabianMeiswinkel May 6, 2026
d265c25
Addressed review feedback
FabianMeiswinkel May 6, 2026
c6d1cec
Addressing review comments
FabianMeiswinkel May 6, 2026
03daf63
Merge branch 'release/azure_data_cosmos-previews' into users/fabianm/…
FabianMeiswinkel May 6, 2026
9d8430b
Fix build issues
FabianMeiswinkel May 6, 2026
58497aa
Merge branch 'users/fabianm/LocalQueryPlanPOC' of https://github.com/…
FabianMeiswinkel May 6, 2026
5832760
Fixed build issues
FabianMeiswinkel May 7, 2026
595e4ed
Update operation_options.rs
FabianMeiswinkel May 7, 2026
413b3c1
Update gateway_query_plan_comparison.rs
FabianMeiswinkel May 7, 2026
b278306
Fixes test issues
FabianMeiswinkel May 7, 2026
bfe67d0
Update mod.rs
FabianMeiswinkel May 7, 2026
4f6c5cd
Update operation_pipeline.rs
FabianMeiswinkel May 7, 2026
93ec596
Merge branch 'release/azure_data_cosmos-previews' of https://github.c…
FabianMeiswinkel May 7, 2026
a58aa74
Merge branch 'release/azure_data_cosmos-previews' into users/fabianm/…
FabianMeiswinkel May 7, 2026
a278c97
Merge branch 'release/azure_data_cosmos-previews' into users/fabianm/…
FabianMeiswinkel May 7, 2026
06c47ea
Merge branch 'users/fabianm/LocalQueryPlanPOC' of https://github.com/…
FabianMeiswinkel May 8, 2026
ffbf615
Update Cargo.toml
FabianMeiswinkel May 8, 2026
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@

### Features Added

- Added local query-plan generator scaffolding under `crate::query` (lexer, parser, AST, planner, and in-memory evaluator). The scaffolding is **not wired into the production query path** yet — production callers still issue Gateway query-plan requests via `CosmosOperation::query_plan`. The `__internal_testing` cargo feature exposes `query::__test_only_generate_query_plan_for_pk_paths`, `query::__TEST_ONLY_SUPPORTED_QUERY_FEATURES`, and `CosmosOperation::query_plan` for cross-crate gateway-comparison tests; this feature is intentionally unstable and **not covered by SemVer**.
- Added per-partition automatic failover (PPAF) for writes on single-master accounts. On 403/3 WriteForbidden, 503 ServiceUnavailable, 429/3092 SystemResourceUnavailable, 410/1022 Gone, or 408 RequestTimeout from a region, the affected partition is failed over to the next preferred region; subsequent writes for that partition skip the failed region. ([#4156](https://github.com/Azure/azure-sdk-for-rust/pull/4156))
- Added per-partition circuit breaker (PPCB) for reads (any account) and writes (multi-master accounts). Tracks failure counts per `(partition_key_range_id, region)` and routes to an alternate region once the threshold (default 10 reads, 5 writes) is exceeded. A background failback loop probes the original region for recovery. ([#4156](https://github.com/Azure/azure-sdk-for-rust/pull/4156))
- Added `OperationOptions` fields for tuning PPCB: `circuit_breaker_failure_count_for_reads`, `circuit_breaker_failure_count_for_writes`, `circuit_breaker_timeout_counter_reset_window_in_minutes`, `allowed_partition_unavailability_duration_in_seconds`, `ppcb_stale_partition_unavailability_refresh_interval_in_seconds`, and `per_partition_circuit_breaker_enabled` (each also configurable via the corresponding `AZURE_COSMOS_*` environment variable). ([#4156](https://github.com/Azure/azure-sdk-for-rust/pull/4156))
Expand Down
13 changes: 13 additions & 0 deletions sdk/cosmos/azure_data_cosmos_driver/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -76,8 +76,21 @@ reqwest = [
rustls = ["reqwest", "reqwest/rustls", "__tls"]
native_tls = ["reqwest", "reqwest/native-tls", "__tls"]
fault_injection = ["dep:rand"]
# `__internal_in_memory_emulator` exposes the in-memory Cosmos DB emulator
# (`crate::in_memory_emulator`) and its query evaluator
# (`crate::query::eval`, `crate::query::value`). The evaluator intentionally
# trades full Cosmos parity for emulator usability (see
# `docs/IN_MEMORY_EMULATOR_SPEC.md` and the doc comments on the `eval` module).
# Production code MUST NOT enable this feature; it is not covered by SemVer
# and may change or disappear at any time.
__internal_in_memory_emulator = ["dep:tokio", "dep:time", "dep:percent-encoding"]
__internal_mocking = []
# `__internal_testing` exposes a small, intentionally-unstable surface
# (`CosmosOperation::query_plan` and `query::__TEST_ONLY_SUPPORTED_QUERY_FEATURES`,
# plus `query::__test_only_generate_query_plan_for_pk_paths`) for cross-crate
# gateway-comparison tests. Production code MUST NOT enable this feature; it is
# not covered by SemVer and may change or disappear at any time.
__internal_testing = []
__tls = []

[package.metadata.docs.rs]
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,186 @@
<!-- cspell:ignore queryengines LALR WCHAR bitflags STARTSWITH ENDSWITH LTRIM RTRIM sqlparser -->
# Cosmos DB Query Engine — Rust Implementation

## Summary

A subset of the C++ query engine has been ported to Rust, enabling:

1. **Client-side query plan generation** — Parse SQL text, extract partition key filters, and produce structural query info (aggregates, ORDER BY, GROUP BY, DISTINCT, etc.) without a Gateway roundtrip.
2. **In-memory query evaluation** — Match JSON documents against SQL WHERE clauses and apply SELECT projections, for use in test emulators.

The implementation lives entirely inside the `azure_data_cosmos_driver` crate. In normal builds the query subsystem remains crate-private; test builds and the `__internal_testing` feature expose temporary validation entry points (`query` and `__test_only_generate_query_plan_for_pk_paths`) so parity tests can exercise the local planner without making it part of the supported surface.

The supported SDK query path still uses Gateway query plans today. The local planner and evaluator are scaffolding that is validated in isolation, but they are not yet wired into production query execution.

---

## Architecture

```
SQL Text
→ Lexer (hand-crafted tokenizer)
→ Parser (recursive descent with Pratt precedence)
→ QueryPlan { pk_filters, query_info }
├── pk_filters: PartitionKeyFilter (Equality / InList / Unconstrained / Contradictory / NotEvaluated)
└── query_info: LocalQueryInfo (structural analysis from the AST)

Gateway response (when issued)
→ GatewayQueryPlan { partition_key_ranges, query_info: GatewayQueryInfo }
```

The `LocalQueryInfo` and `GatewayQueryInfo` types are intentionally **not**
unified (see commit marker `F21`). `LocalQueryInfo` carries fields the AST
can populate (`has_join`, `has_subquery`, `has_where`, `has_udf`,
`has_select_value`, …). `GatewayQueryInfo` carries fields only the Gateway
can populate (`rewritten_query`, `group_by_aliases`, `d_count_info`,
`has_non_streaming_order_by`, …). The fields they share are compared by
`gateway_plan::shared_fields_match`, which is the parity surface the
`tests/gateway_query_plan_comparison.rs` suite asserts against. Splitting
the types avoids silently fabricating `false` for local-only booleans on
Gateway responses (and vice versa).

The pipeline goes directly from SQL AST to partition key extraction and structural analysis. No IL layer, no VM — direct AST interpretation.

---

## Module Structure

All modules live under `azure_data_cosmos_driver::query`. The module is `pub(crate)` in normal builds and exposed only for tests / `__internal_testing` validation:

```
sdk/cosmos/azure_data_cosmos_driver/src/query/
├── mod.rs # Module root, re-exports parse()
├── ast/mod.rs # SQL AST types (SqlProgram, SqlQuery, SqlScalarExpression, etc.)
├── lexer/mod.rs # Hand-crafted tokenizer (TokenKind, Lexer, keyword lookup)
├── parser/mod.rs # Recursive descent parser, Pratt precedence for expressions
├── plan/
│ ├── mod.rs # Query plan generation + LocalQueryInfo type
│ └── tests/
│ └── query_plan_comparison.rs # Exhaustive structural comparison tests
├── eval/mod.rs # In-memory evaluator (gated on `__internal_in_memory_emulator`)
├── gateway_plan.rs # Gateway response envelope (GatewayQueryPlan / GatewayQueryInfo + shared_fields_match)
├── common.rs # Shared utilities (root alias extraction)
└── value.rs # CosmosValue: type-aware comparison semantics (gated on `__internal_in_memory_emulator`)
```

### Why Inside the Driver Crate?

- Query plan generation is an internal implementation detail — no external consumer needs the types.
- The driver already has all required dependencies (`serde`, `serde_json`, `azure_core`).
- Keeps the supported public API surface at zero in normal builds; only test/internal feature gates expose validation hooks.
- The split `LocalQueryInfo` / `GatewayQueryInfo` types live next to the
pieces that produce them (plan generator vs. response deserialization)
while `gateway_plan::shared_fields_match` keeps the parity contract in
one place.

---

## Implemented Features

### SQL Parser

Full recursive descent parser for the Cosmos DB SQL dialect:
- SELECT (star, list, VALUE), DISTINCT, TOP
- FROM with aliases, JOINs, array iterators, subqueries
- WHERE with all scalar expression types
- GROUP BY, ORDER BY, OFFSET/LIMIT
- Operators: arithmetic, comparison, logical, bitwise, string concat, coalesce, ternary
- IN, BETWEEN, LIKE (with ESCAPE), IS NULL / IS NOT NULL
- EXISTS, ARRAY, scalar subqueries
- UDF calls (`udf.name(args)`)
- Parameters (`@name`)
- Max nesting depth: 128

### Query Plan Generation

- Partition key filter extraction from WHERE clauses
- Single PK equality, IN lists, hierarchical PK (2 and 3 components)
- AND intersection logic (contradictory, redundant, narrowing)
- OR union logic (equality + equality, equality + IN, IN + IN) with duplicate-value deduplication
- Nested PK paths (e.g., `/address/city`)
- FROM alias resolution
- Full structural analysis populated by the AST: `LocalQueryInfo` with
`distinct`, `top`, `offset`, `limit`, `order_by`, `group_by`,
`aggregates`, `has_join`, `has_subquery`, `has_where`, `has_udf`,
`has_select_value`.

### LocalQueryInfo / GatewayQueryInfo split

- **`LocalQueryInfo`** is produced by the local plan generator from the
AST. Only fields the AST can populate are present (`has_join`,
`has_subquery`, `has_where`, `has_udf`, `has_select_value`, plus the
structural fields above).
- **`GatewayQueryInfo`** is what the Gateway returns over the wire. It
carries Gateway-only fields (`rewritten_query`, `group_by_aliases`,
`d_count_info`, `has_non_streaming_order_by`, …) in addition to the
shared fields.
- `gateway_plan::shared_fields_match(&LocalQueryInfo)` is the comparison
surface: it explicitly ignores the disjoint Gateway-only / local-only
fields and only compares the fields both sides can populate. The
parity test suite (`tests/gateway_query_plan_comparison.rs`) asserts
through this contract so future divergences are caught early.

### In-Memory Evaluator

Gated behind the `__internal_in_memory_emulator` feature flag. Used by the
in-memory Cosmos DB emulator and inline unit tests. The evaluator
intentionally trades full Cosmos parity for emulator usability — see
`docs/IN_MEMORY_EMULATOR_SPEC.md` for the documented trade-offs.

- `matches_query()`: WHERE clause evaluation against JSON documents
- `project()`: SELECT clause projection
- `query_documents()`: Full query execution (WHERE + SELECT + JOIN + GROUP BY + ORDER BY + TOP + OFFSET/LIMIT)
- 30+ built-in functions (CONTAINS, UPPER, ABS, ARRAY_CONTAINS, etc.)
- SQL LIKE with DP-based pattern matching
- Three-valued logic (undefined AND/OR semantics)
- Cosmos DB comparison semantics (type ordering, cross-type = undefined)
- JOIN expansion with multiple iterator bindings
- GROUP BY with aggregate evaluation (COUNT, SUM, AVG, MIN, MAX)

---

## Testing

- **Exhaustive structural plan comparison tests** covering every `QueryInfo` field, PK extraction pattern, hierarchical PK, AND/OR intersection, nested paths, aliases, and edge cases
- **Inline unit tests** in each module (lexer, parser, plan, eval, value), including typed `GatewayQueryPlan` deserialization coverage in `gateway_plan.rs`
- **Live Gateway validation tests** in `tests/gateway_query_plan_comparison.rs`, behind `__internal_testing`, comparing local plans against Gateway responses using `CosmosOperation::query_plan`

---

## Known Limitations / Parity Gaps

These are deliberate (and small) divergences from the Gateway, tracked here so
a future PR can close the gap without re-discovering it from scratch.

| Area | Gap | Notes |
| --------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------------------------- |
| PK extraction | `c["pk"]` and `c.address["city"]` style indexer references are not extracted as PK references; the local plan falls back to cross-partition routing for them. The Gateway recognizes these forms. | F5 in the post-review notes. |
| `LENGTH` builtin | The local evaluator counts Unicode scalar values; the Gateway returns UTF-16 code-unit count (matching JS / .NET `string.Length`). Surrogate-pair characters diverge. | F35. |
| Bitwise ops on `f64` | `&` `\|` `^` `<<` `>>` use `f64 as i64` saturating cast; the Gateway uses C++/JS int32 truncation. Documented inline in `eval::int_op`. | F23. |
| Parameterized `TOP @n` | Locally accepted when the parameter is bound; the Gateway rejects parameterized `TOP` with HTTP 400 even when bound. The integration layer must avoid sending such queries to the Gateway. | F14. |
| `LIKE … ESCAPE 'xy'` (multi-char) | Local evaluator returns `Undefined` (row does not match); the Gateway rejects the query. Plan-level shape is unaffected. | F15. |
| `~` on fractional `Number` | Local evaluator returns `Undefined`; the Gateway rejects non-integral bitwise input. | F22. |

## What Is Explicitly Not Implemented

| Component | Reason |
| -------------------------------------------- | ------------------------------------------------------------------------ |
| IL compilation pipeline | Direct AST interpretation suffices |
| VM runtime / bytecode execution | Backend-only concern |
| Index plans / physical plans | Backend-only concern |
| Distributed query coordination | Gateway's responsibility |
| KQL / JavaScript query support | Not needed |
| Full ORDER BY / GROUP BY in plan routing | Plan generation detects these features; execution is server-side |
| Production query execution using local plans | Still pending; the supported SDK path continues to request Gateway plans |

## Alternatives considered

This implementation is a port of the Cosmos SQL native engine; an off-the-shelf parser like
[`sqlparser-rs`](https://crates.io/crates/sqlparser) was not adopted because (a) Cosmos SQL
has dialect-specific JSON-path syntax (`c.address.city`, array subscripts, `IN` over arrays)
and operators (`??`, ternary, `EXISTS`/`ARRAY` subqueries) that don't map cleanly onto a
generic SQL parser's AST, (b) the porting strategy validates correctness against the
Gateway via `tests/gateway_query_plan_comparison.rs`
for end-to-end parity, and (c) hand-written parsing keeps the AST under tight control for
the partition-key extraction and plan-generation passes that are the main reason the local
plan generator exists.
24 changes: 24 additions & 0 deletions sdk/cosmos/azure_data_cosmos_driver/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,30 @@ pub mod fault_injection;
pub mod in_memory_emulator;
pub mod models;
pub mod options;
// The `query` module is local-plan scaffolding. Many helpers (gateway response
// envelope, value comparison helpers, etc.) are temporarily unused in the driver
// proper because no production caller wires the local plan generator in yet. The
// `#[allow(dead_code)]` annotation is intentional and should be removed once the
// driver pipeline starts consuming the local plan output. Until then, individual
// per-item `#[allow(dead_code)]` would mean ~50 annotations across lexer/parser/
// eval/plan scaffolding without changing what the compiler actually checks.
//
// The two `mod query;` declarations differ only in visibility, which is gated on
// the `__internal_testing` feature: when that feature is on we expose a small,
// `#[doc(hidden)]` test-only surface (`__test_only_generate_query_plan_for_pk_paths`,
// `__TEST_ONLY_SUPPORTED_QUERY_FEATURES`) so cross-crate gateway-comparison
// tests can drive the local plan generator without depending on internal types;
// otherwise the module is `pub(crate)` and nothing leaks out of the crate.
// Keep both arms in sync if you add another item under `mod query`.
//
// TODO(local-plan-wire-up): drop `allow(dead_code)` once the driver wires the
// local plan generator into the query execution path.
#[cfg(any(test, feature = "__internal_testing"))]
#[allow(dead_code)]
pub mod query;
#[cfg(not(any(test, feature = "__internal_testing")))]
#[allow(dead_code)]
pub(crate) mod query;
pub(crate) mod system;
#[cfg(feature = "__internal_mocking")]
pub mod testing;
Expand Down
Loading
Loading