Introduction
AI-assisted database querying is becoming useful for developers who work with large and unfamiliar data models. Instead of manually writing every query, a developer can describe the required information in natural language and let an AI system help construct the database query.
For Azure Cosmos DB applications, this can be especially useful when developers need to explore documents, filter records, or build queries against containers they did not design themselves.
However, a generated query should not be considered successful simply because it executes without an error.
A useful query must satisfy at least two requirements:
It should return the correct data.
It should use database resources efficiently.
For Azure Cosmos DB, Request Units (RUs) provide an important measure of database consumption. Two queries can return the same records while consuming different amounts of RUs.
This makes AI-assisted query generation an interesting engineering problem:
Can AI-generated Cosmos DB queries achieve the same correctness as developer-written queries while maintaining reasonable RU efficiency?
The answer should come from measurement rather than assumption.
What Is AI Query Assistance?
AI query assistance allows a developer or application to describe a data requirement in natural language.
For example:
Find the five most recent completed orders
for customer 1001.An AI system may generate a query similar to:
SELECT TOP 5
c.id,
c.customerId,
c.status,
c.total,
c.createdAt
FROM c
WHERE c.customerId = @customerId
AND c.status = @status
ORDER BY c.createdAt DESCThe application can then validate and execute the query.
The important distinction is:
Natural Language
|
v
AI Query Generation
|
v
Validation
|
v
Cosmos DB
|
+--> Result
+--> RU ChargeThe AI generates the candidate query, but the application should remain responsible for validation and execution policy.
Why Accuracy and RU Efficiency Must Be Measured Together
Consider two queries that return exactly the same five orders.
Query A
Correct result
Lower RU consumption
Query B
Correct result
Higher RU consumptionBoth are functionally correct.
However, Query A may be preferable for a frequently executed production workload.
Now consider:
Query C
Lower RU consumption
Incorrect resultQuery C is not a successful optimization.
Therefore, the benchmark should use a two-stage evaluation:
Stage 1
Correctness
Stage 2
RU EfficiencyNever optimize a query before confirming that it answers the business requirement correctly.
Define a Representative Cosmos DB Dataset
Start with a realistic document model.
For example:
{
"id": "order-1001",
"customerId": "customer-42",
"status": "Completed",
"total": 749.50,
"createdAt": "2026-08-20T10:15:00Z",
"items": [
{
"productId": "p-100",
"quantity": 2
}
]
}The benchmark should contain enough data to exercise realistic query behavior.
Consider multiple dataset sizes:
| Dataset | Approximate Size |
|---|---|
| Small | 1,000 documents |
| Medium | 100,000 documents |
| Large | 1,000,000+ documents |
These are test configurations, not universal requirements.
The important point is to evaluate behavior beyond a tiny development dataset.
Keep the Database Configuration Consistent
The AI-generated query and expert-written query should run against the same:
Container
Partition key
Documents
Indexing configuration
Consistency configuration
Parameters
Database environment
The benchmark should look like:
Same Requirement
|
+---------+---------+
| |
v v
Expert Query AI Query
| |
+---------+---------+
|
v
Same Cosmos DB
|
+---------+---------+
| |
v v
Results RU ChargeChanging database configuration between tests makes the comparison unreliable.
Define the Ground Truth
Before generating an AI query, define what the correct result should contain.
For example:
Requirement:
Return the five most recent completed orders
for customer-42.
Expected:
- Only customer-42
- Only Completed orders
- Sorted newest first
- Maximum five resultsThis becomes the ground truth.
The AI-generated query can then be evaluated against it.
Test Query Correctness
Correctness should include more than syntax.
Check:
Filtering
Does the query return the correct customer?
Status
Does it return only the requested status?
Ordering
Are the results sorted correctly?
Result Count
Does it return no more than the requested number?
Projection
Does it return the required fields?
Null Handling
Does it behave correctly when optional properties are missing?
Boundary Conditions
Does date filtering behave correctly around the beginning and end of a range?
A query that executes successfully can still fail any of these tests.
Compare Results Programmatically
For repeatable testing, compare the results automatically.
A simple C# representation might be:
public sealed record OrderResult(
string Id,
string CustomerId,
string Status,
decimal Total,
DateTime CreatedAt);After executing both queries:
bool SameResults(
IReadOnlyList<OrderResult> expected,
IReadOnlyList<OrderResult> actual)
{
return expected.SequenceEqual(actual);
}In a real benchmark, the comparison should account for the required ordering and fields rather than relying on object equality alone.
Capture the RU Charge
Cosmos DB query responses expose the request charge.
For example:
using FeedIterator<OrderResult> iterator =
container.GetItemQueryIterator<OrderResult>(
queryDefinition);
while (iterator.HasMoreResults)
{
FeedResponse<OrderResult> response =
await iterator.ReadNextAsync();
var requestCharge = response.RequestCharge;
Console.WriteLine(
$"RU Charge: {requestCharge}");
}The benchmark should record the charge for every execution.
For a multi-page query, do not record only the first page.
Calculate the total:
Total RU =
Page 1 RU
+ Page 2 RU
+ Page 3 RU
+ ...This gives a more realistic measurement of the complete query operation.
Measure RU Per Successful Task
A useful metric is:
RU per Successful QueryFor example:
| Query | Correct? | Total RU |
|---|---|---|
| Expert Query | Yes | Measure |
| AI Query | Yes | Measure |
Do not compare RU consumption when one query returns incorrect data unless the result is explicitly categorized as a correctness failure.
Test Different Query Types
A benchmark should contain more than simple point-like filters.
Include categories such as:
| Query Type | Example |
|---|---|
| Filtering | Find active orders |
| Sorting | Latest orders |
| Aggregation | Total sales |
| Projection | Return selected fields |
| Array filtering | Orders containing a product |
| Date filtering | Orders within a range |
| Pagination | Retrieve the next page |
| Multi-condition | Customer + status + date |
Different query shapes can produce different behavior.
Test Partition-Key Awareness
Partitioning is one of the most important aspects of Cosmos DB query design.
Suppose the container uses:
/customerIdas its partition key.
A query that includes the partition key can be fundamentally different from one that does not.
For example:
SELECT *
FROM c
WHERE c.customerId = @customerId
AND c.status = @statusThe AI should understand the data model well enough to use the relevant partition information when the application's requirements permit it.
The benchmark should therefore include:
Partition-Aware Query
vs.
Query Without Partition Contextand measure both correctness and RU behavior.
Do not assume that every query must contain the partition key. Some legitimate queries are intentionally broader.
The benchmark should evaluate whether the query matches the actual workload.
Test Queries With Different Selectivity
Data distribution matters.
Consider:
Customer A
10 orders
Customer B
10,000 ordersThe same query structure may behave differently depending on the requested customer.
Test multiple parameter distributions:
No Matches
Few Matches
Typical Matches
Many MatchesThis is more informative than testing a single customer.
Compare Projection Strategies
Suppose the requirement asks for:
Order ID
Status
TotalAn AI system may generate:
SELECT *
FROM c
WHERE c.customerId = @customerIdwhile a more focused query might use:
SELECT
c.id,
c.status,
c.total
FROM c
WHERE c.customerId = @customerIdThe benchmark should compare the resulting behavior.
The point is not to assume that selecting fewer fields always produces a particular RU reduction. Measure the actual workload.
Test Pagination
Pagination is another useful benchmark category.
For example:
SELECT TOP @pageSize
c.id,
c.createdAt,
c.total
FROM c
WHERE c.customerId = @customerId
ORDER BY c.createdAt DESCThe benchmark should test multiple pages and record:
Page Number
Rows Returned
RU Charge
LatencyA query that performs well for the first page may behave differently for later pages depending on the pagination strategy.
Test Aggregations
AI-generated aggregation queries deserve additional validation.
For example:
Join the conversation! Your thoughts help the community grow.