Query: Fixes handling of pipeline execution on partition merge

[ealsur]

Pull Request Template

Description

The CrossPartitionRangePageAsyncEnumerator when facing a merge, was adding the merge result as part of the list of Enumerators, instead of maintaining the current Range and draining it.

In this sample scenario, let's assume a 2 partition collection (partition A and B).

When the query pipeline starts, it generates a queue with 2 enumerators, one for A, one for B. Assume each partition has 10 documents.

Assume we fetch the first page which contains 10 documents from A, which returns a ContinuationToken like:

{token: null, range {B}}

Which normally would make the query continue with B and drain it.

If there is a merge and A and B merge into C (with C having all the 20 documents), the current logic adds a new enumerator for C in the queue copying the continuation from the parent (B), so the query would not continue on C (full range) with continuation null (all the documents).

The effect is that the query end ups returning 30 documents (10 from the first page, 20 from the second because it used C).

What is the fix

In the case of a merge where the result is a single parent range, instead of creating a new Queue Enumerator for the merge result, we re-enqueue the current Range Enumerator.

When this happens, the logic in the RequestInvokerHandler kicks in and applies the EPK filtering:

https://github.com/Azure/azure-cosmos-dotnet-v3/blob/master/Microsoft.Azure.Cosmos/src/Handler/RequestInvokerHandler.cs#L244-L263

There was a test currently trying to cover this scenario but it was failing to catch the bug due to a bug in the test itself and how the InMemoryContainer behaved on a merge.

Type of change

• Bug fix (non-breaking change which fixes an issue)

[FabianMeiswinkel]

LGTM - Thanks