[Cosmos] cross partition continuation token #35511

topshot99 · 2025-08-08T05:04:24Z

Packages impacted by this PR

@azure/cosmos

Issues associated with this PR

Describe the problem that is addressed by this PR

Adds support for continuation token for cross partition queries

What are the possible designs available to address the problem? If there are more than one possible design, why was the one in this PR chosen?

DOC Link

Are there test cases added in this PR? (If not, why?)

Yes

Provide a list of related PRs (if any)

Command used to generate this PR:***(Applicable only to SDK release request PRs)*

Checklists

Added impacted package name to the issue description
Does this PR needs any fixes in the SDK Generator?** (If so, create an Issue in the Autorest/typescript repository and link it here)
Added a changelog (if necessary)

…tly across components

… enhance query iterator logging

…ing and refactor related components

…into feature/QueryContinuationToken

…y iterator recreation with continuation token

- Implement comprehensive tests for the _enableQueryControlFetchMoreImplementation method, covering various scenarios including: - Processing existing buffer with items and unprocessed ranges. - Fetching from endpoint when buffer is empty. - Handling cases where endpoint returns no data or an empty buffer. - Managing partial buffer processing and endpoint responses with orderByItemsArray. Enhance functional query tests to log continuation token details and verify query execution phases.

…ling, and improve partition key range mapping in query execution context.

…ate ContinuationTokenManager to handle these values

… limit handling.

…nt for partition splits and merges, and validate continuation token usage for unsupported query types.

…rameter and the rangesOverlap method for improved clarity and maintainability.

…en handling - Implemented a test class for ParallelQueryExecutionContextBase to validate continuation token filtering logic. - Added tests for detecting query types (parallel and OrderBy) based on query information. - Included tests for EPK value extraction and document producer creation with EPK values. - Developed integration scenarios to validate the complete continuation token filtering workflow for both parallel and OrderBy queries. - Implemented tests for handling partition splits and merges, ensuring correct behavior of continuation token management.

topshot99 · 2025-08-20T09:19:07Z

sdk/cosmosdb/cosmos/src/queryExecutionContext/pipelinedQueryExecutionContext.ts

+      // Pass shared continuation token manager via options
+      const optionsWithSharedManager = {
+        ...this.options,
+        continuationTokenManager: sharedContinuationTokenManager


create an extended feed option for passing continuation token manager

Or pass continuationTokenManager directly to the EndpointComponents.

topshot99 · 2025-08-20T09:21:49Z

sdk/cosmosdb/cosmos/src/queryExecutionContext/pipelinedQueryExecutionContext.ts

+    // 1. There are items in the fetch buffer, OR
+    // 2. There are unprocessed ranges in the partition key range map, OR
+    // 3. The endpoint has more results
+    if (this.options.enableQueryControl) {


throw error in older flow

sdk/cosmosdb/cosmos/src/queryExecutionContext/pipelinedQueryExecutionContext.ts

sdk/cosmosdb/cosmos/src/queryExecutionContext/parallelQueryExecutionContextBase.ts

…rove continuation token management - Updated GroupByEndpointComponent, GroupByValueEndpointComponent, NonStreamingOrderByDistinctEndpointComponent, NonStreamingOrderByEndpointComponent, and OffsetLimitEndpointComponent to eliminate buffer references and streamline response handling. - Introduced continuation token management in OrderByEndpointComponent and OrderedDistinctEndpointComponent to enhance query execution efficiency. - Removed SimplifiedTargetPartitionRangeManager and integrated its functionality into TargetPartitionRangeManager for better code organization. - Adjusted test cases in parallelQueryExecutionContextBase.spec.ts and targetPartitionRangeManager.spec.ts to reflect changes in response structure and ensure accurate validation of filtered ranges. - Enhanced partition range filtering strategies to support direct return types instead of promises, simplifying the interface and improving performance.

niteshvijay1995 · 2025-08-21T07:53:08Z

sdk/cosmosdb/cosmos/src/client/Item/Items.ts

@@ -147,6 +147,10 @@ export class Items {
   */
  public query<T>(query: string | SqlQuerySpec, options?: FeedOptions): QueryIterator<T>;
  public query<T>(query: string | SqlQuerySpec, options: FeedOptions = {}): QueryIterator<T> {
+    console.log("==========================================");


add logger wherever needed.

…ngeStrategy - Improve validation logic to check for null, undefined, or empty string continuation tokens. - Ensure all range mappings in continuation tokens have valid partitionKeyRange. - Update filterPartitionRanges method to handle empty and null target ranges gracefully. - Refactor tests to remove async/await where unnecessary and ensure proper handling of continuation tokens.

amanrao23 · 2025-08-21T05:30:55Z

sdk/cosmosdb/cosmos/src/queryExecutionContext/QueryRangeMapping.ts

+ * @returns True if EPK boundaries are defined
+ */
+export function hasEpkBoundaries(partitionKeyRange: ExtendedPartitionKeyRange): boolean {
+  return !!(partitionKeyRange.epkMin && partitionKeyRange.epkMax);


why are there two !!?

sdk/cosmosdb/cosmos/src/queryExecutionContext/QueryRangeMapping.ts

sdk/cosmosdb/cosmos/src/queryExecutionContext/parallelQueryExecutionContextBase.ts

amanrao23 · 2025-08-21T07:59:45Z

sdk/cosmosdb/cosmos/src/queryExecutionContext/parallelQueryExecutionContextBase.ts

+   * @param originalDocumentProducer - The document producer for the original partition that was split/merged
+   * @param replacementPartitionKeyRanges - The new partition ranges after the split/merge
+   */
+  private _updateContinuationTokenForPartitionSplit(


Update name to include merge too

sdk/cosmosdb/cosmos/src/queryExecutionContext/parallelQueryExecutionContextBase.ts

sdk/cosmosdb/cosmos/src/documents/ContinuationToken/OrderByQueryContinuationToken.ts

sdk/cosmosdb/cosmos/src/queryExecutionContext/QueryRangeMapping.ts

sdk/cosmosdb/cosmos/src/queryExecutionContext/ContinuationTokenManager.ts

…erage - Updated test cases to reject tokens with empty orderByItems array. - Changed async tests to synchronous where applicable for better performance. - Introduced new test scenarios for handling continuation tokens, including edge cases and exhausted tokens. - Enhanced validation checks for malformed and empty composite tokens. - Improved assertions to ensure accurate validation of filtering conditions and continuation tokens. - Organized tests into logical groups for better readability and maintainability.

…thod and streamline offset/limit handling in query execution

…tion queries

…tinuationTokenManager to use new token creation and serialization functions

…functions related to partition key ranges

…andling and enhance error reporting; simplify condition in HybridQueryExecutionContext for result validation

niteshvijay1995 · 2025-08-26T10:33:59Z

sdk/cosmosdb/cosmos/src/queryExecutionContext/EndpointComponent/OffsetLimitEndpointComponent.ts

+    options?: FeedOptions,
+  ) {
+    // Get the continuation token manager from options if available
+    this.continuationTokenManager = options.continuationTokenManager;


nit: options is optional

niteshvijay1995 · 2025-08-26T11:44:09Z

sdk/cosmosdb/cosmos/src/request/FeedOptions.ts

+   * Shared continuation token manager for handling query pagination state.
+   * This is used internally to coordinate continuation tokens across query execution contexts.
+   */
+  continuationTokenManager?: ContinuationTokenManager;


Is FeedOptions best place to keep this?

niteshvijay1995 · 2025-08-26T11:44:47Z

sdk/cosmosdb/cosmos/src/queryIterator.ts

@@ -62,6 +62,9 @@ export class QueryIterator<T> {
    private resourceLink?: string,
    private resourceType?: ResourceType,
  ) {
+    console.log("==========================================");


this will be removed?

sdk/cosmosdb/cosmos/src/queryExecutionContext/CompositeQueryContinuationToken.ts

niteshvijay1995 · 2025-08-26T12:04:25Z

sdk/cosmosdb/cosmos/src/queryExecutionContext/pipelinedQueryExecutionContext.ts

@@ -98,36 +148,47 @@ export class PipelinedQueryExecutionContext implements ExecutionContext {
        );
      }
    } else {
+      // Create shared continuation token manager for streaming execution contexts
+      const sharedContinuationTokenManager = new ContinuationTokenManager(


this is confusing. Why do we have 2 continuationTokenManager?

niteshvijay1995 · 2025-08-26T12:04:48Z

sdk/cosmosdb/cosmos/src/queryExecutionContext/pipelinedQueryExecutionContext.ts

      }
-
+      


nit: revert

niteshvijay1995 · 2025-08-26T12:20:43Z

sdk/cosmosdb/cosmos/src/queryExecutionContext/ContinuationTokenManager.ts

+          `Failed to parse continuation token: ${error.message}, initializing empty token`,
+        );
+        // Fallback to empty continuation token if parsing fails
+        this.compositeContinuationToken = new CompositeQueryContinuationTokenClass(


Don't fallback. Better to fail the request.
Implementation should be transparent and deterministic for users.

sdk/cosmosdb/cosmos/src/queryExecutionContext/ContinuationTokenManager.ts

niteshvijay1995 · 2025-08-26T12:28:20Z

sdk/cosmosdb/cosmos/src/queryExecutionContext/pipelinedQueryExecutionContext.ts

      // If offset+limit then add that to the pipeline
      const limit = partitionedQueryExecutionInfo.queryInfo.limit;
      const offset = partitionedQueryExecutionInfo.queryInfo.offset;
      if (typeof limit === "number" && typeof offset === "number") {
-        this.endpoint = new OffsetLimitEndpointComponent(this.endpoint, offset, limit);
+        this.endpoint = new OffsetLimitEndpointComponent(this.endpoint, offset, limit, optionsWithSharedManager);


can a user pass offset and limit with continuation token?

niteshvijay1995 · 2025-08-26T12:32:10Z

sdk/cosmosdb/cosmos/src/queryExecutionContext/pipelinedQueryExecutionContext.ts


+      console.log("hasBufferedItems:", hasBufferedItems);


We are not planning to keep these console.log?

- Changed CompositeQueryContinuationToken from a class to an interface and introduced factory functions for creation and serialization. - Updated ContinuationTokenManager to utilize the new interface and functions for managing composite continuation tokens. - Introduced PartitionRangeManager to encapsulate partition key range mapping logic, improving separation of concerns. - Refactored OrderByQueryRangeStrategy and ParallelQueryRangeStrategy to use the new composite token handling methods. - Updated unit tests to reflect changes in the CompositeQueryContinuationToken structure and ensure proper functionality.

… clarity

sdk/cosmosdb/cosmos/src/queryExecutionContext/parallelQueryExecutionContextBase.ts

topshot99 · 2025-08-27T07:59:32Z

sdk/cosmosdb/cosmos/src/queryExecutionContext/EndpointComponent/OffsetLimitEndpointComponent.ts

+    // Check continuation token for offset/limit values during initialization
+    if (options?.continuationToken) {
+      try {
+        const parsedToken = JSON.parse(options.continuationToken);    


pass it from QEC

…ntext

sdk/cosmosdb/cosmos/src/queryExecutionContext/CompositeQueryContinuationToken.ts

amanrao23 · 2025-08-27T08:33:15Z

sdk/cosmosdb/cosmos/src/queryExecutionContext/pipelinedQueryExecutionContext.ts

+    // Initialize continuation token manager early so it's available for OffsetLimitEndpointComponent
+    const sortOrders = partitionedQueryExecutionInfo.queryInfo.orderBy;
+    const isOrderByQuery = Array.isArray(sortOrders) && sortOrders.length > 0;
+    this.continuationTokenManager = new ContinuationTokenManager(


add a check for validating that continuation token belongs to this container only using rid.

continuation token manager is getting created in this class? why do we need rid check?

sdk/cosmosdb/cosmos/src/documents/ContinuationToken/OrderByQueryContinuationToken.ts

sdk/cosmosdb/cosmos/src/queryExecutionContext/CompositeQueryContinuationToken.ts

amanrao23 · 2025-08-27T09:37:12Z

sdk/cosmosdb/cosmos/src/queryExecutionContext/ContinuationTokenManager.ts

+        console.warn(
+          `Failed to parse continuation token: ${error.message}, initializing empty token`,
+        );
+        // Fallback to empty continuation token if parsing fails


Parsing could also fail if someone tampers with the continuation token string. We should fail the request here with an error.

amanrao23 · 2025-08-27T10:18:01Z

sdk/cosmosdb/cosmos/src/queryExecutionContext/PartitionRangeManager.ts

+   * @param rangeId - Unique identifier for the partition range
+   * @param mapping - The QueryRangeMapping to add
+   */
+  public updatePartitionRangeMapping(rangeId: string, mapping: QueryRangeMapping): void {


Should this be renamed to addPartitionRangeMapping? update is misleading here

amanrao23 · 2025-08-27T10:22:33Z

sdk/cosmosdb/cosmos/src/queryExecutionContext/PartitionRangeManager.ts

+   * @param rangeMappings - Array of range mappings to filter
+   * @returns Filtered array without exhausted ranges
+   */
+  public removeExhaustedRanges(rangeMappings: QueryRangeMapping[]): QueryRangeMapping[] {


function name is misleading, IMO it should not return anything and just remove the ranges. Or we should update the name accordingly.

amanrao23 · 2025-08-27T10:23:48Z

sdk/cosmosdb/cosmos/src/queryExecutionContext/PartitionRangeManager.ts

+    // ORDER BY queries require orderByItemsArray to be present and non-empty
+    if (!orderByItemsArray || orderByItemsArray.length === 0) {
+      throw new Error(
+        "ORDER BY query processing failed: orderByItemsArray is required but was not provided or is empty"


Add what should Cx do if such a case arises.

amanrao23 · 2025-08-27T10:26:23Z

sdk/cosmosdb/cosmos/src/queryExecutionContext/PartitionRangeManager.ts

+    console.log("=== Processing ORDER BY Query (Sequential Mode) ===");
+
+    // ORDER BY queries require orderByItemsArray to be present and non-empty
+    if (!orderByItemsArray || orderByItemsArray.length === 0) {


is orderByItemsArray passed here just for checking? I don't see any other use of it in this function.

amanrao23 · 2025-08-27T10:37:27Z

sdk/cosmosdb/cosmos/src/queryExecutionContext/PartitionRangeManager.ts

+   * @param initialLimit - Initial limit value  
+   * @returns Updated partition key range map with offset/limit values for each range
+   */
+  public calculateOffsetLimitForEachPartitionRange(


is this getting used outside this class? If not, make it private

…prove clarity and support for enableQueryControl option

- Updated OrderByQueryRangeStrategy and ParallelQueryRangeStrategy to utilize PartitionRangeWithContinuationToken for improved handling of continuation tokens. - Removed legacy continuation token validation and parsing methods, simplifying the filterPartitionRanges method. - Enhanced TargetPartitionRangeManager to accept range-token pairs and additional query information. - Implemented partition split/merge detection in ParallelQueryExecutionContextBase, allowing for dynamic updates to continuation tokens based on partition topology changes. - Adjusted related interfaces and types to accommodate new structure for handling partition ranges and continuation tokens. - Updated unit tests to reflect changes in continuation token handling and ensure correct functionality.

- Introduced a new `ParallelQueryResult` interface to standardize the structure of results returned from parallel queries. - Updated `GroupByValueEndpointComponent`, `NonStreamingOrderByDistinctEndpointComponent`, and `NonStreamingOrderByEndpointComponent` to utilize the new `ParallelQueryResult` structure. - Enhanced error handling for undefined or empty results in query components. - Created `TargetPartitionRangeManager` and associated strategies (`ParallelQueryRangeStrategy`, `OrderByQueryRangeStrategy`) to manage partition range filtering based on query type and continuation tokens. - Added unit tests for the new filtering strategies and refactored existing tests to accommodate changes in import paths. - Ensured backward compatibility by maintaining existing functionality while introducing new structures and strategies.

…prove continuation token handling

…ents

- Introduced CompositeQueryContinuationToken to handle parallel query execution across multiple partition ranges. - Updated OrderByQueryContinuationToken to utilize rangeMappings instead of a composite token. - Enhanced ContinuationTokenManager to manage both composite and order by continuation tokens, streamlining offset and limit handling. - Removed deprecated continuation token parsing logic and integrated new serialization/deserialization methods. - Updated various endpoint components to align with the new continuation token structure, ensuring compatibility with offset and limit queries. - Cleaned up FeedOptions by removing the continuationTokenManager property, as it is now managed internally. - Added unit tests to validate the new continuation token management logic.

…g ORDER BY processing

…t structure and improve offset/limit handling

mkhandelwal-123 added 7 commits July 7, 2025 13:20

add continuation token

a4b69f7

Refactor query execution context to handle response buffers consisten…

4722fcf

…tly across components

Implement query control feature with composite continuation token and…

ee19dc5

… enhance query iterator logging

Add ContinuationTokenManager for improved multi-partition query handl…

8d90660

…ing and refactor related components

Merge branch 'main' of https://github.com/topshot99/azure-sdk-for-js …

4402ac0

…into feature/QueryContinuationToken

Enhance ORDER BY query handling and continuation token management

b826254

add new tests for fetchBufferEndIndexForCurrentPage and ORDER BY quer…

865e0a2

…y iterator recreation with continuation token

topshot99 added the Cosmos label Aug 8, 2025

mkhandelwal-123 added 10 commits August 8, 2025 14:03

Implement code changes to enhance functionality and improve performance

70f7482

Enhance continuation token management by adding offset and limit hand…

adb598e

…ling, and improve partition key range mapping in query execution context.

Enhance continuation token management by adding offset and limit hand…

5586a14

…ling, and improve partition key range mapping in query execution context.

Add offset and limit support to OrderByQueryContinuationToken and upd…

b1998f9

…ate ContinuationTokenManager to handle these values

enhance OffsetLimitEndpointComponent to accept options for offset and…

c6d38ce

… limit handling.

add hashed last result support for distinct order queries

0109bb6

Enhance query execution context by adding continuation token manageme…

92eebc5

…nt for partition splits and merges, and validate continuation token usage for unsupported query types.

Refactor filterPartitionRanges method by removing unused queryInfo pa…

e1da4f8

…rameter and the rangesOverlap method for improved clarity and maintainability.

topshot99 commented Aug 20, 2025

View reviewed changes

amanrao23 reviewed Aug 20, 2025

View reviewed changes

sdk/cosmosdb/cosmos/src/queryExecutionContext/pipelinedQueryExecutionContext.ts Show resolved Hide resolved

sdk/cosmosdb/cosmos/src/queryExecutionContext/parallelQueryExecutionContextBase.ts Outdated Show resolved Hide resolved

topshot99 commented Aug 20, 2025

View reviewed changes

sdk/cosmosdb/cosmos/src/queryExecutionContext/parallelQueryExecutionContextBase.ts Outdated Show resolved Hide resolved

topshot99 commented Aug 20, 2025

View reviewed changes

sdk/cosmosdb/cosmos/src/queryExecutionContext/parallelQueryExecutionContextBase.ts Outdated Show resolved Hide resolved

topshot99 commented Aug 20, 2025

View reviewed changes

sdk/cosmosdb/cosmos/src/queryExecutionContext/parallelQueryExecutionContextBase.ts Show resolved Hide resolved

niteshvijay1995 reviewed Aug 21, 2025

View reviewed changes

amanrao23 reviewed Aug 21, 2025

View reviewed changes

mkhandelwal-123 added 6 commits August 26, 2025 11:54

Refactor ContinuationTokenManager to encapsulate updateOffsetLimit me…

a43970a

…thod and streamline offset/limit handling in query execution

Add error handling for unsupported continuation tokens in cross-parti…

a295fbd

…tion queries

Add CompositeQueryContinuationToken class and update related imports

8111cd0

Refactor OrderByQueryContinuationToken to an interface and update Con…

d5e80cd

…tinuationTokenManager to use new token creation and serialization functions

Refactor QueryRangeMapping interface documentation and remove unused …

527d975

…functions related to partition key ranges

Refactor OffsetLimitEndpointComponent to improve continuation token h…

d5a24b8

…andling and enhance error reporting; simplify condition in HybridQueryExecutionContext for result validation

niteshvijay1995 reviewed Aug 26, 2025

View reviewed changes

sdk/cosmosdb/cosmos/src/queryExecutionContext/CompositeQueryContinuationToken.ts Outdated Show resolved Hide resolved

niteshvijay1995 reviewed Aug 26, 2025

View reviewed changes

sdk/cosmosdb/cosmos/src/queryExecutionContext/pipelinedQueryExecutionContext.ts

}

Copy link

Member

niteshvijay1995 Aug 26, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: revert

niteshvijay1995 reviewed Aug 26, 2025

View reviewed changes

sdk/cosmosdb/cosmos/src/queryExecutionContext/ContinuationTokenManager.ts Show resolved Hide resolved

niteshvijay1995 reviewed Aug 26, 2025

View reviewed changes

mkhandelwal-123 added 2 commits August 26, 2025 19:36

Refactor offset/limit calculation in PartitionRangeManager to improve…

db7cae5

… clarity

aditishree1 reviewed Aug 27, 2025

View reviewed changes

sdk/cosmosdb/cosmos/src/queryExecutionContext/parallelQueryExecutionContextBase.ts Outdated Show resolved Hide resolved

topshot99 commented Aug 27, 2025

View reviewed changes

Enhance continuation token management in pipelined query execution co…

faed188

…ntext

amanrao23 reviewed Aug 27, 2025

View reviewed changes

mkhandelwal-123 added 8 commits August 28, 2025 12:55

Refactor continuation token handling in query execution context to im…

5d21731

…prove clarity and support for enableQueryControl option

Refactor endpoint components to support parallel query results and im…

8ea6701

…prove continuation token handling

Add orderByItems support to parallel query results and related compon…

9c5feea

…ents

Remove error throwing for insufficient orderByItemsArray length durin…

5820d6e

…g ORDER BY processing

Refactor query execution components to support new ParallelQueryResul…

945bfe9

…t structure and improve offset/limit handling

[Cosmos] cross partition continuation token #35511

Are you sure you want to change the base?

[Cosmos] cross partition continuation token #35511

Uh oh!

Conversation

topshot99 commented Aug 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Packages impacted by this PR

Issues associated with this PR

Describe the problem that is addressed by this PR

What are the possible designs available to address the problem? If there are more than one possible design, why was the one in this PR chosen?

Are there test cases added in this PR? (If not, why?)

Provide a list of related PRs (if any)

Command used to generate this PR:**(Applicable only to SDK release request PRs)

Checklists

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

topshot99 Aug 20, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

niteshvijay1995 Aug 26, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

topshot99 commented Aug 8, 2025 •

edited

Loading

Command used to generate this PR:***(Applicable only to SDK release request PRs)*

topshot99 Aug 20, 2025 •

edited

Loading

niteshvijay1995 Aug 26, 2025 •

edited

Loading