UNPKG

@azure/cosmos

Version:
156 lines 5.87 kB
// Copyright (c) Microsoft Corporation. // Licensed under the MIT License. /** * Manages partition key range mappings for query execution. * Handles range operations, offset/limit processing, and distinct query logic. * @hidden */ export class PartitionRangeManager { partitionKeyRangeMap = new Map(); constructor(initialPartitionKeyRangeMap) { if (initialPartitionKeyRangeMap) { this.partitionKeyRangeMap = new Map(initialPartitionKeyRangeMap); } } /** * Gets a copy of the current partition key range map for constructor pattern */ getPartitionKeyRangeMap() { return new Map(this.partitionKeyRangeMap); } /** * Checks if a continuation token indicates an exhausted partition * @param continuationToken - The continuation token to check * @returns true if the partition is exhausted (null, empty, or "null" string) */ isPartitionExhausted(continuationToken) { return (!continuationToken || continuationToken === "" || continuationToken === "null" || continuationToken.toLowerCase() === "null"); } /** * Adds a range mapping to the partition key range map * Does not allow updates to existing keys - only new additions * @param rangeId - Unique identifier for the partition range * @param mapping - The QueryRangeMapping to add */ addPartitionRangeMapping(rangeId, mapping) { if (!this.partitionKeyRangeMap.has(rangeId)) { this.partitionKeyRangeMap.set(rangeId, mapping); } } /** * Removes a range mapping from the partition key range map */ removePartitionRangeMapping(rangeId) { this.partitionKeyRangeMap.delete(rangeId); } /** * Updates the partition key range map with new mappings from the endpoint response * @param partitionKeyRangeMap - Map of range IDs to QueryRangeMapping objects */ addPartitionKeyRangeMap(partitionKeyRangeMap) { if (partitionKeyRangeMap) { for (const [rangeId, mapping] of partitionKeyRangeMap) { this.addPartitionRangeMapping(rangeId, mapping); } } } /** * Checks if there are any unprocessed ranges in the sliding window */ hasUnprocessedRanges() { return this.partitionKeyRangeMap.size > 0; } /** * Removes exhausted(fully drained) ranges from the given range mappings * @param rangeMappings - Array of range mappings to filter * @returns Filtered array without exhausted ranges */ removeExhaustedRanges(rangeMappings) { if (!rangeMappings || !Array.isArray(rangeMappings)) { return []; } return rangeMappings.filter((mapping) => { if (!mapping) { return false; } const isExhausted = this.isPartitionExhausted(mapping.continuationToken); if (isExhausted) { return false; } return true; }); } /** * Processes ranges for ORDER BY queries */ processOrderByRanges(pageSize) { let endIndex = 0; const processedRanges = []; let lastRangeBeforePageLimit = null; let rangeIndex = 0; for (const [rangeId, value] of this.partitionKeyRangeMap) { rangeIndex++; const { itemCount } = value; // Check if this complete range fits within remaining page size capacity if (endIndex + itemCount <= pageSize) { lastRangeBeforePageLimit = value; endIndex += itemCount; processedRanges.push(rangeId); } else { break; } } return { endIndex, processedRanges, lastRangeBeforePageLimit }; } processEmptyOrderByRanges(ranges) { const endIndex = 0; const processedRanges = []; let lastRangeBeforePageLimit; // since there is no data returned add all the ids to processed ranges for (const [rangeId, _] of this.partitionKeyRangeMap) { processedRanges.push(rangeId); } // search for matching range in the map(min max value exact match) and return that as lastRangeBeforePageLimit for (const [_, mapping] of this.partitionKeyRangeMap) { if (mapping.partitionKeyRange.minInclusive === ranges[0].queryRange.min && mapping.partitionKeyRange.maxExclusive === ranges[0].queryRange.max) { lastRangeBeforePageLimit = mapping; break; } } return { endIndex, processedRanges, lastRangeBeforePageLimit }; } /** * Processes ranges for parallel queries - multi-range aggregation */ processParallelRanges(pageSize) { let endIndex = 0; const processedRanges = []; const processedRangeMappings = []; let rangesAggregatedInCurrentToken = 0; let lastPartitionBeforeCutoff; for (const [rangeId, value] of this.partitionKeyRangeMap) { rangesAggregatedInCurrentToken++; // Validate range data if (!value || value.itemCount === undefined) { continue; } const { itemCount } = value; if (endIndex + itemCount <= pageSize) { lastPartitionBeforeCutoff = { rangeId, mapping: value }; endIndex += itemCount; processedRanges.push(rangeId); processedRangeMappings.push(value); } else { break; // No more ranges can fit, exit loop } } return { endIndex, processedRanges, processedRangeMappings, lastPartitionBeforeCutoff }; } } //# sourceMappingURL=PartitionRangeManager.js.map