Local QueryPlan generation and emulator execution

sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md

@@ -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))

sdk/cosmos/azure_data_cosmos_driver/Cargo.toml

@@ -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]

sdk/cosmos/azure_data_cosmos_driver/docs/query-engine-porting-plan.md

@@ -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.

sdk/cosmos/azure_data_cosmos_driver/src/lib.rs

@@ -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;