[Cosmos] cross partition continuation token (#35511)

### 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](https://loop.cloud.microsoft/p/eyJ3Ijp7InUiOiJodHRwczovL21pY3Jvc29mdC5zaGFyZXBvaW50LmNvbS8%2FbmF2PWN6MGxNa1ltWkQxaUlWTlROemhDVTNaS2VVVlhhMjlZVEZrME1FazNhek5PVjFkek0xbGpVemxPYmtNeFoyOVpkRVkzTWpOYVMyTlBOVzVIT1VaVGNWWm5lbmxCUkhGeFEzY21aajB3TVZGRlIweFFWVlkwUlVkUVFVNUxNMGN6VWtnelMwbzNVak0yTmxrMFVVNUZKbU05Sm1ac2RXbGtQVEUlM0QiLCJyIjpmYWxzZX0sInAiOnsidSI6Imh0dHBzOi8vbWljcm9zb2Z0LnNoYXJlcG9pbnQuY29tLzpmbDovci9zaXRlcy9iZGJkZjc1NS03ZjViLTRjOTEtOTk5OS0wMjg3NzNhNTA5OGMvU2hhcmVkJTIwRG9jdW1lbnRzL0xvb3BBcHBEYXRhL1VudGl0bGVkJTIwMi5sb29wP2Q9dzZhNTA0Y2I1MTg1MzQ5NWFhZjg0NGFjYjJhNjFlMTM3JmNzZj0xJndlYj0xJm5hdj1jejBsTWtaemFYUmxjeVV5Um1Ka1ltUm1OelUxTFRkbU5XSXROR001TVMwNU9UazVMVEF5T0RjM00yRTFNRGs0WXlaa1BXSWhVMU0zT0VKVGRrcDVSVmRyYjFoTVdUUXdTVGRyTTA1WFYzTXpXV05UT1U1dVF6Rm5iMWwwUmpjeU0xcExZMDgxYmtjNVJsTnhWbWQ2ZVVGRWNYRkRkeVptUFRBeFVVVkhURkJWVmxaS1VrbEhWVlZaV1V4S1JUSTNRa05MV2sxV1IwUlpTbGdtWXowbE1rWW1abXgxYVdROU1TWmhQVXh2YjNCQmNIQW1jRDBsTkRCbWJIVnBaSGdsTWtac2IyOXdMWEJoWjJVdFkyOXVkR0ZwYm1WeUpuZzlKVGRDSlRJeWR5VXlNaVV6UVNVeU1sUXdVbFJWU0hoMFlWZE9lV0l6VG5aYWJsRjFZekpvYUdOdFZuZGlNbXgxWkVNMWFtSXlNVGhaYVVaVVZYcGpORkZzVGpKVGJteEdWakowZGxkRmVGcE9SRUpLVGpKemVsUnNaRmhqZWs1YVdURk5OVlJ0TlVSTlYyUjJWMWhTUjA1NlNYcFhhM1JxVkhwV2RWSjZiRWRWTTBaWFdqTndOVkZWVW5oalZVNHpaa1JCZUZWVlZraFVSa0pXVm1wU1JsSXhRa0pVYTNONlVucE9VMU5FVGt4VGFtUlRUWHBaTWxkVVVsSlVhMVVsTTBRbE1qSWxNa01sTWpKcEpUSXlKVE5CSlRJeVpqazFNakF6T1RndE5Ua3dOQzAwT0RCakxXRXhNRGd0TTJGbU56UmtNVEV6WmpGbEpUSXlKVGRFIiwiciI6ZmFsc2V9LCJpIjp7ImkiOiJmOTUyMDM5OC01OTA0LTQ4MGMtYTEwOC0zYWY3NGQxMTNmMWUifX0%3D)

### 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](https://github.com/Azure/autorest.typescript)
repository and link it here)_
- [ ] Added a changelog (if necessary)

---------

Co-authored-by: Manik Khandelwal <[email protected]>

Author: topshot99

sdk/cosmosdb/cosmos/src/common/constants.ts

@@ -310,6 +310,23 @@ export const AAD_DEFAULT_SCOPE = "https://cosmos.azure.com/.default";
 export const AAD_AUTH_PREFIX = "type=aad&ver=1.0&sig=";
 export const AAD_RESOURCE_NOT_FOUND_ERROR = "AADSTS500011";
 
+/**
+ * @internal
+ * Internal query execution constants - not part of public API
+ */
+const QueryExecution = {
+  /** Default page size for query execution when maxItemCount is not specified */
+  DEFAULT_PAGE_SIZE: 10,
+  /** Default maximum buffer size for vector search queries */
+  DEFAULT_MAX_VECTOR_SEARCH_BUFFER_SIZE: 50000,
+} as const;
+
+/**
+ * @internal
+ * Export for internal SDK use only
+ */
+export { QueryExecution };
+
 /**
  * @hidden
  */

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

@@ -0,0 +1,177 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+import { QueryRange } from "../../routing/QueryRange.js";
+import type { QueryRangeMapping } from "../../queryExecutionContext/queryRangeMapping.js";
+
+/**
+ * @hidden
+ */
+export interface RangeBoundary {
+  /**
+   * Minimum boundary (inclusive)
+   */
+  min: string;
+  /**
+   * Maximum boundary (exclusive)
+   */
+  max: string;
+}
+
+/**
+ * @hidden
+ * Base interface for all continuation tokens containing common fields
+ */
+export interface BaseContinuationToken {
+  /**
+   * Resource ID of the container for which the continuation token is issued
+   */
+  rid: string;
+
+  /**
+   * List of query ranges with their continuation tokens
+   */
+  rangeMappings: QueryRangeWithContinuationToken[];
+
+  /**
+   * Current offset value for OFFSET/LIMIT queries
+   */
+  offset?: number;
+
+  /**
+   * Current limit value for OFFSET/LIMIT queries
+   */
+  limit?: number;
+}
+
+/**
+ * @hidden
+ * Composite continuation token for parallel query execution across multiple partition ranges
+ */
+export interface CompositeQueryContinuationToken extends BaseContinuationToken {}
+
+/**
+ * Creates a new CompositeQueryContinuationToken
+ * @hidden
+ */
+export function createCompositeQueryContinuationToken(
+  rid: string,
+  rangeMappings: QueryRangeWithContinuationToken[],
+  offset?: number,
+  limit?: number,
+): CompositeQueryContinuationToken {
+  if (!rangeMappings || rangeMappings.length === 0) {
+    throw new Error(
+      "Failed to create composite continuation token: No partition range mappings provided. " +
+        "This typically indicates an issue with query execution context initialization or partition key range resolution. " +
+        "Ensure the query is properly configured and the container has valid partition ranges.",
+    );
+  }
+
+  return {
+    rid,
+    rangeMappings: rangeMappings,
+    offset,
+    limit,
+  };
+}
+
+/**
+ * Serializes the composite continuation token to a JSON string
+ * @hidden
+ */
+export function serializeCompositeToken(token: CompositeQueryContinuationToken): string {
+  return JSON.stringify(token);
+}
+
+/**
+ * Deserializes a JSON string to a CompositeQueryContinuationToken
+ * @hidden
+ */
+export function parseCompositeQueryContinuationToken(
+  tokenString: string,
+): CompositeQueryContinuationToken {
+  return JSON.parse(tokenString);
+}
+
+/**
+ * Deserializes a JSON string to a CompositeQueryContinuationToken
+ * @hidden
+ */
+export function parseBaseContinuationToken(tokenString: string): BaseContinuationToken {
+  return JSON.parse(tokenString);
+}
+
+/**
+ * @hidden
+ * Represents a query range with its associated continuation token
+ */
+export interface QueryRangeWithContinuationToken {
+  /**
+   * The simplified query range containing min/max boundaries
+   */
+  queryRange: RangeBoundary;
+
+  /**
+   * The continuation token for this specific range
+   */
+  continuationToken: string | undefined;
+}
+
+/**
+ * Converts QueryRangeMapping to QueryRangeWithContinuationToken using simplified range format
+ * @param rangeMapping - The QueryRangeMapping to convert
+ * @returns QueryRangeWithContinuationToken with simplified boundaries and continuation token
+ * @hidden
+ */
+export function convertRangeMappingToQueryRange(
+  rangeMapping: QueryRangeMapping,
+): QueryRangeWithContinuationToken {
+  if (!rangeMapping.partitionKeyRange) {
+    throw new Error(
+      "Failed to convert range mapping: Missing partition key range information. " +
+        "The QueryRangeMapping object must contain a valid partitionKeyRange with min and max boundaries. " +
+        "This may indicate an incomplete partition key range resolution during query setup.",
+    );
+  }
+
+  const pkRange = rangeMapping.partitionKeyRange;
+
+  // Create simplified range assuming min is inclusive and max is exclusive
+  const simplifiedRange: RangeBoundary = {
+    min: pkRange.minInclusive,
+    max: pkRange.maxExclusive,
+  };
+
+  return {
+    queryRange: simplifiedRange,
+    continuationToken: rangeMapping.continuationToken,
+  };
+}
+
+/**
+ * Converts an array of QueryRangeMapping to an array of QueryRangeWithContinuationToken
+ * @param rangeMappings - Array of QueryRangeMapping to convert
+ * @returns Array of QueryRangeWithContinuationToken with simplified boundaries and continuation tokens
+ * @hidden
+ */
+export function convertRangeMappingsToQueryRangesWithTokens(
+  rangeMappings: QueryRangeMapping[],
+): QueryRangeWithContinuationToken[] {
+  return rangeMappings.map((mapping) => convertRangeMappingToQueryRange(mapping));
+}
+
+/**
+ * Converts a SimplifiedQueryRange back to a QueryRange for internal use
+ * @param simplifiedRange - The simplified range to convert
+ * @returns QueryRange with standard assumptions (min inclusive, max exclusive)
+ * @hidden
+ */
+export function convertSimplifiedRangeToQueryRange(simplifiedRange: RangeBoundary): QueryRange {
+  return new QueryRange(
+    simplifiedRange.min,
+    simplifiedRange.max,
+    true, // minInclusive = true (assumption)
+    false, // maxInclusive = false (max is exclusive, assumption)
+  );
+}

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

@@ -0,0 +1,88 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+import type {
+  QueryRangeWithContinuationToken,
+  BaseContinuationToken,
+} from "./CompositeQueryContinuationToken.js";
+
+/**
+ * Continuation token for order by queries.
+ * @hidden
+ */
+export interface OrderByQueryContinuationToken extends BaseContinuationToken {
+  /**
+   * Order by items for the query
+   */
+  orderByItems: any[];
+
+  /**
+   * Number of items to skip in the query
+   */
+  skipCount: number;
+
+  /**
+   * Document ID of the last document result
+   */
+  documentRid: string;
+
+  /**
+   * Hash of the last document result for distinct order queries
+   * Used to ensure duplicates are not returned across continuation boundaries
+   */
+  hashedLastResult?: string;
+}
+
+/**
+ * Creates an OrderByQueryContinuationToken
+ * @hidden
+ */
+export function createOrderByQueryContinuationToken(
+  rangeMappings: QueryRangeWithContinuationToken[],
+  orderByItems: any[],
+  rid: string,
+  skipCount: number,
+  documentRid?: string,
+  offset?: number,
+  limit?: number,
+  hashedLastResult?: string,
+): OrderByQueryContinuationToken {
+  if (!rangeMappings || rangeMappings.length === 0) {
+    throw new Error("rangeMappings must contain at least one element");
+  }
+
+  if (!orderByItems || orderByItems.length === 0) {
+    throw new Error("orderByItems must contain at least one element");
+  }
+
+  return {
+    rangeMappings,
+    orderByItems,
+    rid,
+    skipCount,
+    documentRid,
+    offset,
+    limit,
+    hashedLastResult,
+  } as OrderByQueryContinuationToken;
+}
+
+/**
+ * Serializes an OrderByQueryContinuationToken to a JSON string
+ * @hidden
+ */
+export function serializeOrderByQueryContinuationToken(
+  token: OrderByQueryContinuationToken,
+): string {
+  return JSON.stringify(token);
+}
+
+/**
+ * Deserializes a JSON string to an OrderByQueryContinuationToken
+ * @hidden
+ */
+export function parseOrderByQueryContinuationToken(
+  tokenString: string,
+): OrderByQueryContinuationToken {
+  return JSON.parse(tokenString);
+}

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

@@ -0,0 +1,25 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+import type { QueryRange } from "../../index.js";
+
+/**
+ * Represents information about a partition range update that occurred during query execution.
+ * This includes the original range, new ranges after split/merge, and the continuation token.
+ * @hidden
+ */
+export interface PartitionRangeUpdate {
+  /** The original partition key range before the split/merge operation */
+  oldRange: QueryRange;
+  /** The new partition key ranges after the split/merge operation */
+  newRanges: QueryRange[];
+  /** The continuation token associated with this range update */
+  continuationToken: string | undefined;
+}
+
+/**
+ * A collection of partition range updates indexed by range keys.
+ * The key is typically in the format "minInclusive-maxExclusive".
+ * @hidden
+ */
+export type PartitionRangeUpdates = Record<string, PartitionRangeUpdate>;