@azure/cosmos 4.8.0 → 4.9.0

package/dist/react-native/queryExecutionContext/parallelQueryExecutionContextBase.js

@@ -2,15 +2,15 @@
 // Licensed under the MIT License.
 import PriorityQueue from "priorityqueuejs";
 import semaphore from "semaphore";
-import { createClientLogger } from "@azure/logger";
 import { StatusCodes, SubStatusCodes } from "../common/statusCodes.js";
+import { ErrorResponse } from "../request/ErrorResponse.js";
 import { QueryRange } from "../routing/QueryRange.js";
 import { SmartRoutingMapProvider } from "../routing/smartRoutingMapProvider.js";
 import { DocumentProducer } from "./documentProducer.js";
 import { getInitialHeader, mergeHeaders } from "./headerUtils.js";
+import { RidSkipCountFilter } from "./queryFilteringStrategy/RidSkipCountFilter.js";
 import { DiagnosticNodeInternal, DiagnosticNodeType, } from "../diagnostics/DiagnosticNodeInternal.js";
-/** @hidden */
-const logger = createClientLogger("parallelQueryExecutionContextBase");
+import { createParallelQueryResult } from "./parallelQueryResult.js";
 /** @hidden */
 export var ParallelQueryExecutionContextBaseStates;
 (function (ParallelQueryExecutionContextBaseStates) {
@@ -26,17 +26,22 @@ export class ParallelQueryExecutionContextBase {
     options;
     partitionedQueryExecutionInfo;
     correlatedActivityId;
+    rangeManager;
+    queryProcessingStrategy;
+    documentProducerComparator;
     err;
     state;
     static STATES = ParallelQueryExecutionContextBaseStates;
     routingProvider;
-    sortOrders;
     requestContinuation;
     respHeaders;
     unfilledDocumentProducersQueue;
     bufferedDocumentProducersQueue;
     // TODO: update type of buffer from any --> generic can be used here
     buffer;
+    partitionDataPatchMap = new Map();
+    patchCounter = 0;
+    updatedContinuationRanges = new Map();
     sem;
     diagnosticNodeWrapper;
     /**
@@ -52,13 +57,16 @@ export class ParallelQueryExecutionContextBase {
      * @param partitionedQueryExecutionInfo - PartitionedQueryExecutionInfo
      * @hidden
      */
-    constructor(clientContext, collectionLink, query, options, partitionedQueryExecutionInfo, correlatedActivityId) {
+    constructor(clientContext, collectionLink, query, options, partitionedQueryExecutionInfo, correlatedActivityId, rangeManager, queryProcessingStrategy, documentProducerComparator) {
         this.clientContext = clientContext;
         this.collectionLink = collectionLink;
         this.query = query;
         this.options = options;
         this.partitionedQueryExecutionInfo = partitionedQueryExecutionInfo;
         this.correlatedActivityId = correlatedActivityId;
+        this.rangeManager = rangeManager;
+        this.queryProcessingStrategy = queryProcessingStrategy;
+        this.documentProducerComparator = documentProducerComparator;
         this.clientContext = clientContext;
         this.collectionLink = collectionLink;
         this.query = query;
@@ -73,64 +81,377 @@ export class ParallelQueryExecutionContextBase {
         this.err = undefined;
         this.state = ParallelQueryExecutionContextBase.STATES.started;
         this.routingProvider = new SmartRoutingMapProvider(this.clientContext);
-        this.sortOrders = this.partitionedQueryExecutionInfo.queryInfo.orderBy;
         this.buffer = [];
-        this.requestContinuation = options ? options.continuationToken || options.continuation : null;
+        this.requestContinuation = options
+            ? options.continuationToken || options.continuation
+            : undefined;
+        // Validate continuation token usage immediately
+        if (this.requestContinuation && !this.options.enableQueryControl) {
+            throw new Error("Continuation tokens are supported when enableQueryControl is set true in FeedOptions");
+        }
         // response headers of undergoing operation
         this.respHeaders = getInitialHeader();
         // Make priority queue for documentProducers
-        this.unfilledDocumentProducersQueue = new PriorityQueue((a, b) => a.generation - b.generation);
-        // The comparator is supplied by the derived class
+        this.unfilledDocumentProducersQueue = new PriorityQueue((a, b) => this.compareDocumentProducersByRange(a, b));
         this.bufferedDocumentProducersQueue = new PriorityQueue((a, b) => this.documentProducerComparator(b, a));
         // Creating the documentProducers
         this.sem = semaphore(1);
-        const createDocumentProducersAndFillUpPriorityQueueFunc = async () => {
-            // ensure the lock is released after finishing up
+        this.sem.take(() => this._initializeDocumentProducers());
+    }
+    /**
+     * Determine if there are still remaining resources to processs based on the value of the continuation
+     * token or the elements remaining on the current batch in the QueryIterator.
+     * @returns true if there is other elements to process in the ParallelQueryExecutionContextBase.
+     */
+    hasMoreResults() {
+        return (!this.err &&
+            (this.buffer.length > 0 || this.state !== ParallelQueryExecutionContextBase.STATES.ended));
+    }
+    /**
+     * Fetches more results from the query execution context.
+     * @param diagnosticNode - Optional diagnostic node for tracing.
+     * @returns A promise that resolves to the fetched results.
+     * @hidden
+     */
+    async fetchMore(diagnosticNode) {
+        await this.bufferDocumentProducers(diagnosticNode);
+        await this.fillBufferFromBufferQueue();
+        return this.drainBufferedItems();
+    }
+    /**
+     * Processes buffered document producers
+     * @returns A promise that resolves when processing is complete.
+     */
+    async processBufferedDocumentProducers() {
+        while (this.hasBufferedProducers() &&
+            this.shouldProcessBufferedProducers(this.isUnfilledQueueEmpty())) {
+            const producer = this.getNextBufferedProducer();
+            if (!producer)
+                break;
+            await this.processDocumentProducer(producer);
+        }
+    }
+    /**
+     * Processes a single document producer using template method pattern.
+     * Common structure with query-specific processing delegated to subclasses.
+     */
+    async processDocumentProducer(producer) {
+        const response = await this.fetchFromProducer(producer);
+        this._mergeWithActiveResponseHeaders(response.headers);
+        if (response.result) {
+            this.addToBuffer(response.result);
+            this.handlePartitionMapping(producer, response.result);
+        }
+        // Handle producer lifecycle
+        if (producer.peakNextItem() !== undefined) {
+            this.requeueProducer(producer);
+        }
+        else if (producer.hasMoreResults()) {
+            this.moveToUnfilledQueue(producer);
+        }
+    }
+    /**
+     * Handles partition mapping updates - implemented in base class using template method pattern.
+     * Child classes provide query-specific parameters through abstract methods.
+     */
+    handlePartitionMapping(producer, result) {
+        const itemCount = result?.length || 0;
+        const continuationToken = this.getContinuationToken(producer);
+        const mapping = {
+            itemCount,
+            partitionKeyRange: producer.targetPartitionKeyRange,
+            continuationToken,
+        };
+        this.updatePartitionMapping(mapping);
+    }
+    /**
+     * Gets the continuation token to use - implemented by subclasses.
+     */
+    getContinuationToken(producer) {
+        const hasMoreBufferedItems = producer.peakNextItem() !== undefined;
+        return hasMoreBufferedItems ? producer.previousContinuationToken : producer.continuationToken;
+    }
+    /**
+     * Updates partition mapping - creates new entry or merges with existing for ORDER BY queries.
+     */
+    updatePartitionMapping(mapping) {
+        const currentPatch = this.partitionDataPatchMap.get(this.patchCounter.toString());
+        const isSamePartition = currentPatch?.partitionKeyRange?.id === mapping.partitionKeyRange.id;
+        if (isSamePartition && currentPatch) {
+            currentPatch.itemCount += mapping.itemCount;
+            currentPatch.continuationToken = mapping.continuationToken;
+            return;
+        }
+        // Create new partition mapping entry
+        this.partitionDataPatchMap.set((++this.patchCounter).toString(), mapping);
+    }
+    /**
+     * Checks if the unfilled queue is empty (used by ORDER BY for processing control).
+     */
+    isUnfilledQueueEmpty() {
+        return this.unfilledDocumentProducersQueue.size() === 0;
+    }
+    /**
+     * Initializes document producers and fills the priority queue.
+     * Handles both continuation token and fresh query scenarios.
+     */
+    async _initializeDocumentProducers() {
+        try {
+            const targetPartitionRanges = await this._onTargetPartitionRanges();
+            const documentProducers = this.requestContinuation
+                ? await this._createDocumentProducersFromContinuation(targetPartitionRanges)
+                : this._createDocumentProducersFromFresh(targetPartitionRanges);
+            // Fill up our priority queue with documentProducers
+            this._enqueueDocumentProducers(documentProducers);
+            this.sem.leave();
+        }
+        catch (err) {
+            this.err = err;
+            this.sem.leave();
+        }
+    }
+    /**
+     * Creates document producers from continuation token scenario.
+     */
+    async _createDocumentProducersFromContinuation(targetPartitionRanges) {
+        // Parse continuation token to get range mappings and check for split/merge scenarios
+        const parsedToken = this._parseContinuationToken(this.requestContinuation);
+        const continuationRanges = await this._handlePartitionRangeChanges(parsedToken);
+        // Use strategy to create additional query info from parsed token
+        const additionalQueryInfo = this.queryProcessingStrategy.createAdditionalQueryInfo(parsedToken);
+        const filterResult = this.rangeManager.filterPartitionRanges(targetPartitionRanges, continuationRanges, additionalQueryInfo);
+        // Extract ranges and tokens from the combined result
+        const rangeTokenPairs = filterResult.rangeTokenPairs;
+        // Use strategy to create filter context for continuation token processing
+        const filterContext = this.queryProcessingStrategy.createFilterContext(parsedToken);
+        return rangeTokenPairs.map((rangeTokenPair) => this._createDocumentProducerFromRangeTokenPair(rangeTokenPair, continuationRanges, filterContext));
+    }
+    /**
+     * Creates document producers from fresh query scenario (no continuation token).
+     */
+    _createDocumentProducersFromFresh(targetPartitionRanges) {
+        return targetPartitionRanges.map((partitionTargetRange) => this._createTargetPartitionQueryExecutionContext(partitionTargetRange, undefined));
+    }
+    /**
+     * Creates a document producer from a range token pair (continuation token scenario).
+     */
+    _createDocumentProducerFromRangeTokenPair(rangeTokenPair, continuationRanges, filterContext) {
+        const partitionTargetRange = rangeTokenPair.range;
+        const continuationToken = rangeTokenPair.continuationToken;
+        const filterCondition = rangeTokenPair.filteringCondition || undefined;
+        // Find EPK ranges for this partition range from processed continuation response
+        const matchingContinuationRange = continuationRanges.find((cr) => cr.range.id === partitionTargetRange.id);
+        const startEpk = matchingContinuationRange?.epkMin;
+        const endEpk = matchingContinuationRange?.epkMax;
+        // Use strategy to determine partition-specific filter context
+        const targetPartitionId = continuationRanges.length > 0 && continuationRanges[continuationRanges.length - 1].range
+            ? continuationRanges[continuationRanges.length - 1].range.id
+            : undefined;
+        const partitionFilterContext = this.queryProcessingStrategy.getPartitionFilterContext(filterContext, targetPartitionId, partitionTargetRange.id);
+        return this._createTargetPartitionQueryExecutionContext(partitionTargetRange, continuationToken, startEpk, endEpk, !!(startEpk && endEpk), // populateEpkRangeHeaders - true if both EPK values are present
+        filterCondition, partitionFilterContext);
+    }
+    /**
+     * Enqueues document producers into the unfilled queue.
+     */
+    _enqueueDocumentProducers(documentProducers) {
+        documentProducers.forEach((documentProducer) => {
             try {
-                const targetPartitionRanges = await this._onTargetPartitionRanges();
-                const maxDegreeOfParallelism = options.maxDegreeOfParallelism === undefined || options.maxDegreeOfParallelism < 1
-                    ? targetPartitionRanges.length
-                    : Math.min(options.maxDegreeOfParallelism, targetPartitionRanges.length);
-                logger.info("Query starting against " +
-                    targetPartitionRanges.length +
-                    " ranges with parallelism of " +
-                    maxDegreeOfParallelism);
-                let filteredPartitionKeyRanges = [];
-                // The document producers generated from filteredPartitionKeyRanges
-                const targetPartitionQueryExecutionContextList = [];
-                if (this.requestContinuation) {
-                    throw new Error("Continuation tokens are not yet supported for cross partition queries");
+                this.unfilledDocumentProducersQueue.enq(documentProducer);
+            }
+            catch (e) {
+                this.err = e;
+            }
+        });
+    }
+    /**
+     * Checks if there are buffered document producers ready for processing.
+     * Encapsulates queue size checking.
+     */
+    hasBufferedProducers() {
+        return this.bufferedDocumentProducersQueue.size() > 0;
+    }
+    /**
+     * Gets the next buffered document producer for processing.
+     * Encapsulates queue dequeuing logic.
+     */
+    getNextBufferedProducer() {
+        if (this.bufferedDocumentProducersQueue.size() > 0) {
+            return this.bufferedDocumentProducersQueue.deq();
+        }
+        return undefined;
+    }
+    /**
+     * Adds items to the result buffer. Handles both single items and arrays.
+     */
+    addToBuffer(items) {
+        if (Array.isArray(items)) {
+            if (items.length > 0) {
+                this.buffer.push(...items);
+            }
+        }
+        else if (items) {
+            this.buffer.push(items);
+        }
+    }
+    /**
+     * Moves a producer to the unfilled queue for later processing.
+     */
+    moveToUnfilledQueue(producer) {
+        this.unfilledDocumentProducersQueue.enq(producer);
+    }
+    /**
+     * Re-queues a producer to the buffered queue for further processing.
+     */
+    requeueProducer(producer) {
+        this.bufferedDocumentProducersQueue.enq(producer);
+    }
+    /**
+     * Compares two document producers based on their partition key ranges and EPK values.
+     * Primary comparison: minInclusive values for left-to-right range traversal
+     * Secondary comparison: EPK ranges when minInclusive values are identical
+     * @param a - First document producer
+     * @param b - Second document producer
+     * @returns Comparison result for priority queue ordering
+     * @hidden
+     */
+    compareDocumentProducersByRange(a, b) {
+        const aMinInclusive = a.targetPartitionKeyRange.minInclusive;
+        const bMinInclusive = b.targetPartitionKeyRange.minInclusive;
+        const minInclusiveComparison = bMinInclusive.localeCompare(aMinInclusive);
+        // If minInclusive values are the same, check minEPK ranges if they exist
+        if (minInclusiveComparison === 0) {
+            const aMinEpk = a.startEpk;
+            const bMinEpk = b.startEpk;
+            if (aMinEpk && bMinEpk) {
+                return bMinEpk.localeCompare(aMinEpk);
+            }
+        }
+        return minInclusiveComparison;
+    }
+    /**
+     * Detects partition splits/merges by analyzing parsed continuation token ranges and comparing with current topology
+     * @param parsed - The continuation token containing range mappings to analyze
+     * @returns Array of processed ranges with EPK info
+     */
+    async _handlePartitionRangeChanges(parsed) {
+        const processedRanges = [];
+        // Extract range mappings from the already parsed token
+        const rangeMappings = parsed.rangeMappings;
+        if (!rangeMappings || rangeMappings.length === 0) {
+            return [];
+        }
+        // Check each range mapping for potential splits/merges
+        for (const rangeWithToken of rangeMappings) {
+            // Create a new QueryRange instance from the simplified range data
+            const range = rangeWithToken.queryRange;
+            const queryRange = new QueryRange(range.min, range.max, true, // isMinInclusive - assumption: always true
+            false);
+            const rangeMin = queryRange.min;
+            const rangeMax = queryRange.max;
+            // Get current overlapping ranges for this continuation token range
+            const overlappingRanges = await this.routingProvider.getOverlappingRanges(this.collectionLink, [queryRange], this.getDiagnosticNode());
+            // Detect split/merge scenario based on the number of overlapping ranges
+            if (overlappingRanges.length === 0) {
+                continue;
+            }
+            else if (overlappingRanges.length === 1) {
+                // Check if it's the same range (no change) or a merge scenario
+                const currentRange = overlappingRanges[0];
+                if (currentRange.minInclusive !== rangeMin || currentRange.maxExclusive !== rangeMax) {
+                    // Merge scenario - include EPK ranges from original continuation token range
+                    await this._handleContinuationTokenMerge(rangeWithToken, currentRange);
+                    processedRanges.push({
+                        range: currentRange,
+                        continuationToken: rangeWithToken.continuationToken,
+                        epkMin: rangeMin, // Original range min becomes EPK min
+                        epkMax: rangeMax, // Original range max becomes EPK max
+                    });
                 }
                 else {
-                    filteredPartitionKeyRanges = targetPartitionRanges;
+                    // Same range - no merge, no EPK ranges needed
+                    processedRanges.push({
+                        range: currentRange,
+                        continuationToken: rangeWithToken.continuationToken,
+                    });
                 }
-                // Create one documentProducer for each partitionTargetRange
-                filteredPartitionKeyRanges.forEach((partitionTargetRange) => {
-                    // TODO: any partitionTargetRange
-                    // no async callback
-                    targetPartitionQueryExecutionContextList.push(this._createTargetPartitionQueryExecutionContext(partitionTargetRange, undefined));
-                });
-                // Fill up our priority queue with documentProducers
-                targetPartitionQueryExecutionContextList.forEach((documentProducer) => {
-                    // has async callback
-                    try {
-                        this.unfilledDocumentProducersQueue.enq(documentProducer);
-                    }
-                    catch (e) {
-                        this.err = e;
-                    }
-                });
-                this.sem.leave();
             }
-            catch (err) {
-                this.err = err;
-                // release the lock
-                this.sem.leave();
-                return;
+            else {
+                // Split scenario - one range from continuation token now maps to multiple ranges
+                await this._handleContinuationTokenSplit(rangeWithToken, overlappingRanges);
+                // Add all overlapping ranges with the same continuation token to processed ranges
+                overlappingRanges.forEach((rangeValue) => {
+                    processedRanges.push({
+                        range: rangeValue,
+                        continuationToken: rangeWithToken.continuationToken,
+                    });
+                });
             }
-        };
-        this.sem.take(createDocumentProducersAndFillUpPriorityQueueFunc);
+        }
+        return processedRanges;
     }
+    /**
+     * Parses the continuation token based on query type
+     * @param continuationToken - The continuation token string to parse
+     * @returns Parsed continuation token object (ORDER BY or Parallel query token)
+     * @throws ErrorResponse when continuation token is malformed or cannot be parsed
+     */
+    _parseContinuationToken(continuationToken) {
+        try {
+            return this.queryProcessingStrategy.parseContinuationToken(continuationToken);
+        }
+        catch (e) {
+            throw new ErrorResponse(`Invalid continuation token format. Expected token with rangeMappings property. ` +
+                `Ensure the continuation token was generated by a compatible query and has not been modified.`);
+        }
+    }
+    /**
+     * Handles partition merge scenario for continuation token ranges
+     */
+    async _handleContinuationTokenMerge(rangeWithToken, _newMergedRange) {
+        const rangeKey = `${rangeWithToken.queryRange.min}-${rangeWithToken.queryRange.max}`;
+        this.updatedContinuationRanges.set(rangeKey, {
+            oldRange: {
+                min: rangeWithToken.queryRange.min,
+                max: rangeWithToken.queryRange.max,
+                isMinInclusive: true, // Assumption: min is always inclusive
+                isMaxInclusive: false, // Assumption: max is always exclusive
+            },
+            newRanges: [
+                {
+                    min: rangeWithToken.queryRange.min,
+                    max: rangeWithToken.queryRange.max,
+                    isMinInclusive: true, // Assumption: min is always inclusive
+                    isMaxInclusive: false, // Assumption: max is always exclusive
+                },
+            ],
+            continuationToken: rangeWithToken.continuationToken,
+        });
+    }
+    /**
+     * Handles partition split scenario for continuation token ranges
+     */
+    async _handleContinuationTokenSplit(rangeWithToken, overlappingRanges) {
+        const rangeKey = `${rangeWithToken.queryRange.min}-${rangeWithToken.queryRange.max}`;
+        this.updatedContinuationRanges.set(rangeKey, {
+            oldRange: {
+                min: rangeWithToken.queryRange.min,
+                max: rangeWithToken.queryRange.max,
+                isMinInclusive: true, // Assumption: min is always inclusive
+                isMaxInclusive: false, // Assumption: max is always exclusive
+            },
+            newRanges: overlappingRanges.map((range) => ({
+                min: range.minInclusive,
+                max: range.maxExclusive,
+                isMinInclusive: true,
+                isMaxInclusive: false,
+            })),
+            continuationToken: rangeWithToken.continuationToken,
+        });
+    }
+    /**
+     * Handles partition merge scenario for continuation token ranges
+     */
     _mergeWithActiveResponseHeaders(headers) {
         mergeHeaders(this.respHeaders, headers);
     }
@@ -165,7 +486,11 @@ export class ParallelQueryExecutionContextBase {
         if (replacementPartitionKeyRanges.length === 0) {
             throw error;
         }
-        else if (replacementPartitionKeyRanges.length === 1) {
+        if (this.requestContinuation) {
+            // Update composite continuation token to handle partition split
+            this._updateContinuationTokenOnPartitionChange(documentProducer, replacementPartitionKeyRanges);
+        }
+        if (replacementPartitionKeyRanges.length === 1) {
             // Partition is gone due to Merge
             // Create the replacement documentProducer with populateEpkRangeHeaders Flag set to true to set startEpk and endEpk headers
             const replacementDocumentProducer = this._createTargetPartitionQueryExecutionContext(replacementPartitionKeyRanges[0], documentProducer.continuationToken, documentProducer.startEpk, documentProducer.endEpk, true);
@@ -188,26 +513,44 @@ export class ParallelQueryExecutionContextBase {
             });
         }
     }
+    _updateContinuationTokenOnPartitionChange(originalDocumentProducer, replacementPartitionKeyRanges) {
+        const rangeWithToken = this._createQueryRangeWithContinuationToken(originalDocumentProducer);
+        if (replacementPartitionKeyRanges.length === 1) {
+            this._handleContinuationTokenMerge(rangeWithToken, replacementPartitionKeyRanges[0]);
+        }
+        else {
+            this._handleContinuationTokenSplit(rangeWithToken, replacementPartitionKeyRanges);
+        }
+    }
+    /**
+     * Creates a QueryRangeWithContinuationToken object from a DocumentProducer.
+     * Uses the DocumentProducer's target partition key range and continuation token.
+     * @param documentProducer - The DocumentProducer to convert
+     * @returns QueryRangeWithContinuationToken object for token operations
+     */
+    _createQueryRangeWithContinuationToken(documentProducer) {
+        const partitionRange = documentProducer.targetPartitionKeyRange;
+        // Create a simplified QueryRange using the partition key range boundaries
+        const simplifiedQueryRange = {
+            min: documentProducer.startEpk || partitionRange.minInclusive,
+            max: documentProducer.endEpk || partitionRange.maxExclusive,
+        };
+        return {
+            queryRange: simplifiedQueryRange,
+            continuationToken: documentProducer.continuationToken,
+        };
+    }
     static _needPartitionKeyRangeCacheRefresh(error) {
         // TODO: any error
         return (error.code === StatusCodes.Gone &&
             "substatus" in error &&
             error["substatus"] === SubStatusCodes.PartitionKeyRangeGone);
     }
-    /**
-     * Determine if there are still remaining resources to processs based on the value of the continuation
-     * token or the elements remaining on the current batch in the QueryIterator.
-     * @returns true if there is other elements to process in the ParallelQueryExecutionContextBase.
-     */
-    hasMoreResults() {
-        return (!this.err &&
-            (this.buffer.length > 0 || this.state !== ParallelQueryExecutionContextBase.STATES.ended));
-    }
     /**
      * Creates target partition range Query Execution Context
      */
-    _createTargetPartitionQueryExecutionContext(partitionKeyTargetRange, continuationToken, startEpk, endEpk, populateEpkRangeHeaders) {
-        let rewrittenQuery = this.partitionedQueryExecutionInfo.queryInfo.rewrittenQuery;
+    _createTargetPartitionQueryExecutionContext(partitionKeyTargetRange, continuationToken, startEpk, endEpk, populateEpkRangeHeaders, filterCondition, filterContext) {
+        let rewrittenQuery = this.partitionedQueryExecutionInfo.queryInfo?.rewrittenQuery;
         let sqlQuerySpec;
         const query = this.query;
         if (typeof query === "string") {
@@ -219,13 +562,18 @@ export class ParallelQueryExecutionContextBase {
         const formatPlaceHolder = "{documentdb-formattableorderbyquery-filter}";
         if (rewrittenQuery) {
             sqlQuerySpec = JSON.parse(JSON.stringify(sqlQuerySpec));
-            // We hardcode the formattable filter to true for now
-            rewrittenQuery = rewrittenQuery.replace(formatPlaceHolder, "true");
+            rewrittenQuery = filterCondition
+                ? rewrittenQuery.replace(formatPlaceHolder, filterCondition)
+                : rewrittenQuery.replace(formatPlaceHolder, "true");
             sqlQuerySpec["query"] = rewrittenQuery;
         }
         const options = { ...this.options };
         options.continuationToken = continuationToken;
-        return new DocumentProducer(this.clientContext, this.collectionLink, sqlQuerySpec, partitionKeyTargetRange, options, this.correlatedActivityId, startEpk, endEpk, populateEpkRangeHeaders);
+        let filter;
+        if (filterContext) {
+            filter = new RidSkipCountFilter(filterContext);
+        }
+        return new DocumentProducer(this.clientContext, this.collectionLink, sqlQuerySpec, partitionKeyTargetRange, options, this.correlatedActivityId, startEpk, endEpk, populateEpkRangeHeaders, filter);
     }
     async drainBufferedItems() {
         return new Promise((resolve, reject) => {
@@ -240,19 +588,33 @@ export class ParallelQueryExecutionContextBase {
                 // return undefined if there is no more results
                 if (this.buffer.length === 0) {
                     this.sem.leave();
+                    const partitionDataPatchMap = this.partitionDataPatchMap;
+                    this.partitionDataPatchMap = new Map();
+                    this.patchCounter = 0;
+                    // Get and reset updated continuation ranges
+                    const updatedContinuationRanges = Object.fromEntries(this.updatedContinuationRanges);
+                    this.updatedContinuationRanges.clear();
+                    const result = createParallelQueryResult([], partitionDataPatchMap, updatedContinuationRanges, undefined);
                     return resolve({
-                        result: this.state === ParallelQueryExecutionContextBase.STATES.ended ? undefined : [],
+                        result: this.state === ParallelQueryExecutionContextBase.STATES.ended ? undefined : result,
                         headers: this._getAndResetActiveResponseHeaders(),
                     });
                 }
                 // draing the entire buffer object and return that in result of return object
                 const bufferedResults = this.buffer;
                 this.buffer = [];
+                // reset the patchToRangeMapping
+                const partitionDataPatchMap = this.partitionDataPatchMap;
+                this.partitionDataPatchMap = new Map();
+                this.patchCounter = 0;
+                // Get and reset updated continuation ranges
+                const updatedContinuationRanges = Object.fromEntries(this.updatedContinuationRanges);
+                this.updatedContinuationRanges.clear();
                 // release the lock before returning
                 this.sem.leave();
-                // invoke the callback on the item
+                const result = createParallelQueryResult(bufferedResults, partitionDataPatchMap, updatedContinuationRanges, undefined);
                 return resolve({
-                    result: bufferedResults,
+                    result,
                     headers: this._getAndResetActiveResponseHeaders(),
                 });
             });
@@ -286,7 +648,7 @@ export class ParallelQueryExecutionContextBase {
                 try {
                     const maxDegreeOfParallelism = this.options.maxDegreeOfParallelism === undefined ||
                         this.options.maxDegreeOfParallelism < 1
-                        ? this.unfilledDocumentProducersQueue.size()
+                        ? this.unfilledDocumentProducersQueue.size() // number of partitions
                         : Math.min(this.options.maxDegreeOfParallelism, this.unfilledDocumentProducersQueue.size());
                     const documentProducers = [];
                     while (documentProducers.length < maxDegreeOfParallelism &&
@@ -307,13 +669,30 @@ export class ParallelQueryExecutionContextBase {
                         try {
                             const headers = await documentProducer.bufferMore(diagnosticNode);
                             this._mergeWithActiveResponseHeaders(headers);
-                            // if buffer of document producer is filled, add it to the buffered document producers queue
+                            // Always track this document producer in patchToRangeMapping, even if it has no results
+                            // This ensures we maintain a record of all partition ranges that were scanned
                             const nextItem = documentProducer.peakNextItem();
                             if (nextItem !== undefined) {
                                 this.bufferedDocumentProducersQueue.enq(documentProducer);
                             }
-                            else if (documentProducer.hasMoreResults()) {
-                                this.unfilledDocumentProducersQueue.enq(documentProducer);
+                            else {
+                                // Track document producer with no results in patchToRangeMapping
+                                // This represents a scanned partition that yielded no results
+                                // IMPORTANT: Only include if continuation token is NOT null/exhausted
+                                // Document producers with no data in buffer and no continuation token are exhausted and should not be added to partitionDataPatchMap to prevent infinite loops in order by queries
+                                if (documentProducer.continuationToken &&
+                                    documentProducer.continuationToken !== "" &&
+                                    documentProducer.continuationToken.toLowerCase() !== "null") {
+                                    const patchKey = `empty-${documentProducer.targetPartitionKeyRange.id}-${documentProducer.targetPartitionKeyRange.minInclusive}`;
+                                    this.partitionDataPatchMap.set(patchKey, {
+                                        itemCount: 0, // 0 items for empty result set
+                                        partitionKeyRange: documentProducer.targetPartitionKeyRange,
+                                        continuationToken: documentProducer.continuationToken,
+                                    });
+                                }
+                                if (documentProducer.hasMoreResults()) {
+                                    this.unfilledDocumentProducersQueue.enq(documentProducer);
+                                }
                             }
                         }
                         catch (err) {
@@ -355,10 +734,10 @@ export class ParallelQueryExecutionContextBase {
     }
     /**
      * Drains the buffer of filled document producers and appends their items to the main buffer.
-     * @param isOrderBy - Indicates if the query is an order by query.
+     * Uses template method pattern - delegates actual processing to subclasses.
      * @returns A promise that resolves when the buffer is filled.
      */
-    async fillBufferFromBufferQueue(isOrderBy = false) {
+    async fillBufferFromBufferQueue() {
         return new Promise((resolve, reject) => {
             this.sem.take(async () => {
                 if (this.err) {
@@ -375,39 +754,7 @@ export class ParallelQueryExecutionContextBase {
                     return;
                 }
                 try {
-                    if (isOrderBy) {
-                        while (this.unfilledDocumentProducersQueue.isEmpty() &&
-                            this.bufferedDocumentProducersQueue.size() > 0) {
-                            const documentProducer = this.bufferedDocumentProducersQueue.deq();
-                            const { result, headers } = await documentProducer.fetchNextItem();
-                            this._mergeWithActiveResponseHeaders(headers);
-                            if (result) {
-                                this.buffer.push(result);
-                            }
-                            if (documentProducer.peakNextItem() !== undefined) {
-                                this.bufferedDocumentProducersQueue.enq(documentProducer);
-                            }
-                            else if (documentProducer.hasMoreResults()) {
-                                this.unfilledDocumentProducersQueue.enq(documentProducer);
-                            }
-                            else {
-                                // no more results in document producer
-                            }
-                        }
-                    }
-                    else {
-                        while (this.bufferedDocumentProducersQueue.size() > 0) {
-                            const documentProducer = this.bufferedDocumentProducersQueue.deq();
-                            const { result, headers } = await documentProducer.fetchBufferedItems();
-                            this._mergeWithActiveResponseHeaders(headers);
-                            if (result) {
-                                this.buffer.push(...result);
-                            }
-                            if (documentProducer.hasMoreResults()) {
-                                this.unfilledDocumentProducersQueue.enq(documentProducer);
-                            }
-                        }
-                    }
+                    await this.processBufferedDocumentProducers();
                     this.updateStates(this.err);
                 }
                 catch (err) {