HPK: Add internal CosmosClientOptions flag UseLengthAwareRangeComparer for length aware range comparer rollout.

Microsoft.Azure.Cosmos/src/CosmosClientOptions.cs

@@ -466,7 +466,26 @@ public System.Text.Json.JsonSerializerOptions UseSystemTextJsonSerializerWithOpt
         /// The default value for this parameter is 'false'.
         /// </summary>
         internal bool EnableStreamPassThrough { get; set; } = false;
-
+
+        /// <summary>
+        /// Gets or sets a value indicating whether to use length-aware range comparators for EPK range comparisons.
+        /// Length-aware range comparators were introduced in Range class to handle EPK range comparisons correctly 
+        /// in the case of a container's physical partition set consisting of fully and partially specified EPK values.
+        /// By default, length-aware range comparator is enabled. Refer to Range.cs in Msdata project for more details. 
+        /// Range.LengthAwareMinComparer/LengthAwareMaxComparer.
+        /// Setting the value to false will disable length-aware range comparator and switch to using the regular 
+        /// Range.MinComparer/MaxComparer.
+        /// </summary>
+        /// <value>
+        /// The default value is true.
+        /// </value>
+        internal bool UseLengthAwareRangeComparer { get; set; } =
+#if !INTERNAL
+            true;
+#else
+            false;
+#endif
+
         /// <summary>
         /// (Direct/TCP) Controls the amount of idle time after which unused connections are closed.
         /// </summary>

Microsoft.Azure.Cosmos/src/DocumentClient.cs

@@ -457,6 +457,7 @@ internal DocumentClient(Uri serviceEndpoint,
         /// <param name="cosmosClientTelemetryOptions">This is distributed tracing flag</param>
         /// <param name="chaosInterceptorFactory">This is the chaos interceptor used for fault injection</param>
         /// <param name="enableAsyncCacheExceptionNoSharing">A boolean flag indicating if stack trace optimization is enabled.</param>
+        /// <param name="useLengthAwareRangeComparer">A boolean flag indicating if length-aware range comparators should be used for EPK range comparisons.</param>
         /// <remarks>
         /// The service endpoint can be obtained from the Azure Management Portal.
         /// If you are connecting using one of the Master Keys, these can be obtained along with the endpoint from the Azure Management Portal
@@ -486,7 +487,8 @@ internal DocumentClient(Uri serviceEndpoint,
                               RemoteCertificateValidationCallback remoteCertificateValidationCallback = null,
                               CosmosClientTelemetryOptions cosmosClientTelemetryOptions = null,
                               IChaosInterceptorFactory chaosInterceptorFactory = null,
-                              bool enableAsyncCacheExceptionNoSharing = true)
+                              bool enableAsyncCacheExceptionNoSharing = true,
+                              bool useLengthAwareRangeComparer = false)
         {
             if (sendingRequestEventArgs != null)
             {
@@ -514,6 +516,7 @@ internal DocumentClient(Uri serviceEndpoint,
                 enableAsyncCacheExceptionNoSharing: this.enableAsyncCacheExceptionNoSharing);
             this.chaosInterceptorFactory = chaosInterceptorFactory;
             this.chaosInterceptor = chaosInterceptorFactory?.CreateInterceptor(this);
+            this.UseLengthAwareRangeComparer = useLengthAwareRangeComparer;
 
             this.Initialize(
                 serviceEndpoint: serviceEndpoint,
@@ -1101,7 +1104,7 @@ private async Task<bool> GetInitializationTaskAsync(IStoreClientFactory storeCli
                     retryPolicy: this.retryPolicy,
                     telemetryToServiceHelper: this.telemetryToServiceHelper,
                     enableAsyncCacheExceptionNoSharing: this.enableAsyncCacheExceptionNoSharing);
-            this.partitionKeyRangeCache = new PartitionKeyRangeCache(this, this.GatewayStoreModel, this.collectionCache, this.GlobalEndpointManager, this.enableAsyncCacheExceptionNoSharing);
+            this.partitionKeyRangeCache = new PartitionKeyRangeCache(this, this.GatewayStoreModel, this.collectionCache, this.GlobalEndpointManager, this.UseLengthAwareRangeComparer, this.enableAsyncCacheExceptionNoSharing);
             this.ResetSessionTokenRetryPolicy = new ResetSessionTokenRetryPolicyFactory(this.sessionContainer, this.collectionCache, this.retryPolicy);
 
             gatewayStoreModel.SetCaches(this.partitionKeyRangeCache, this.collectionCache);
@@ -1229,7 +1232,9 @@ internal ApiType ApiType
             get; private set;
         }
 
-        internal bool UseMultipleWriteLocations { get; private set; }
+        internal bool UseMultipleWriteLocations { get; private set; }
+
+        internal bool UseLengthAwareRangeComparer { get; private set; }
 
         /// <summary>
         /// Gets the endpoint Uri for the service endpoint from the Azure Cosmos DB service.

Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs

@@ -517,9 +517,9 @@ private static TryCatch<IQueryPipelineStage> TryCreatePassthroughQueryExecutionC
                             isMinInclusive: true,
                             isMaxInclusive: false)))
                     .ToList(),
+                partitionKey: inputParameters.PartitionKey,
                 queryPaginationOptions: new QueryExecutionOptions(
                     pageSizeHint: inputParameters.MaxItemCount),
-                partitionKey: inputParameters.PartitionKey,
                 containerQueryProperties: containerQueryProperties,
                 maxConcurrency: inputParameters.MaxConcurrency,
                 prefetchPolicy: PrefetchPolicy.PrefetchSinglePage,
@@ -926,6 +926,7 @@ public static InputParameters Create(
             public bool EnableOptimisticDirectExecution { get; }
             public bool IsHybridSearchQueryPlanOptimizationDisabled { get; }
             public bool EnableDistributedQueryGatewayMode { get; }
+            public bool UseLengthAwareRangeComparer { get; }
 
             public InputParameters WithContinuationToken(CosmosElement token)
             {

Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/ContainerQueryProperties.cs

@@ -15,13 +15,15 @@ public ContainerQueryProperties(
             IReadOnlyList<Range<string>> effectivePartitionKeyRanges,
             PartitionKeyDefinition partitionKeyDefinition,
             Cosmos.VectorEmbeddingPolicy vectorEmbeddingPolicy,
-            Cosmos.GeospatialType geospatialType)
+            Cosmos.GeospatialType geospatialType,
+            bool useLengthAwareRangeComparer)
         {
             this.ResourceId = resourceId;
             this.EffectiveRangesForPartitionKey = effectivePartitionKeyRanges;
             this.PartitionKeyDefinition = partitionKeyDefinition;
             this.VectorEmbeddingPolicy = vectorEmbeddingPolicy;
-            this.GeospatialType = geospatialType;
+            this.GeospatialType = geospatialType;
+            this.UseLengthAwareRangeComparer = useLengthAwareRangeComparer;
         }
 
         public string ResourceId { get; }
@@ -34,6 +36,7 @@ public ContainerQueryProperties(
 
         public Cosmos.VectorEmbeddingPolicy VectorEmbeddingPolicy { get; }
 
-        public Cosmos.GeospatialType GeospatialType { get; }
+        public Cosmos.GeospatialType GeospatialType { get; }
+        public bool UseLengthAwareRangeComparer { get; }
     }
 }

Microsoft.Azure.Cosmos/src/Query/Core/QueryRangeUtils.cs

@@ -7,7 +7,8 @@ namespace Microsoft.Azure.Cosmos.Query.Core
     using System;
     using System.Collections.Generic;
     using System.Diagnostics;
-    using Microsoft.Azure.Cosmos.Query.Core.QueryClient;
+    using Microsoft.Azure.Cosmos.Query.Core.QueryClient;
+    using Microsoft.Azure.Cosmos.Routing;
     using Microsoft.Azure.Documents.Routing;
 
     internal static class QueryRangeUtils
@@ -54,15 +55,8 @@ public static FeedRangeInternal LimitHpkFeedRangeToPartition(PartitionKey? parti
                                 bool minInclusive;
                                 String overlappingMax;
                                 bool maxInclusive;
-
-                                //LengthAwareComparer is the default Range comparer and flag <see cref="ConfigurationManager.UseLengthAwareRangeComparator"/>  is used to ovverride the default comparer to legacy Min/Max comparer.
-                                IComparer<Range<string>> minComparer = IsLengthAwareComparisonEnabled
-                                    ? Documents.Routing.Range<string>.LengthAwareMinComparer.Instance
-                                    : Documents.Routing.Range<string>.MinComparer.Instance;
-
-                                IComparer<Range<string>> maxComparer = IsLengthAwareComparisonEnabled
-                                    ? Documents.Routing.Range<string>.LengthAwareMaxComparer.Instance
-                                    : Documents.Routing.Range<string>.MaxComparer.Instance;
+
+                                (IComparer<Range<string>> minComparer, IComparer<Range<string>> maxComparer) = RangeComparerProvider.GetComparers(containerQueryProperties.UseLengthAwareRangeComparer);
 
                                 if (minComparer.Compare(
                                         epkForPartitionKey,
@@ -117,18 +111,14 @@ public static FeedRangeInternal LimitHpkFeedRangeToPartition(PartitionKey? parti
         /// </summary>
         /// <param name="partitionKeyRanges">The list of partition key ranges to trim</param>
         /// <param name="providedRanges">The EPK ranges to use as boundaries</param>
+        /// <param name="useLengthAwareComparer">Whether to use length-aware range comparers</param>
         /// <returns>A list of trimmed partition key ranges that fit within the provided ranges</returns>
         public static List<Documents.PartitionKeyRange> LimitPartitionKeyRangesToProvidedRanges(
             List<Documents.PartitionKeyRange> partitionKeyRanges,
-            IReadOnlyList<Documents.Routing.Range<string>> providedRanges)
+            IReadOnlyList<Documents.Routing.Range<string>> providedRanges,
+            bool useLengthAwareComparer = true)
         {
-            IComparer<Range<string>> minComparer = IsLengthAwareComparisonEnabled
-                ? Documents.Routing.Range<string>.LengthAwareMinComparer.Instance
-                : Documents.Routing.Range<string>.MinComparer.Instance;
-
-            IComparer<Range<string>> maxComparer = IsLengthAwareComparisonEnabled
-                ? Documents.Routing.Range<string>.LengthAwareMaxComparer.Instance
-                : Documents.Routing.Range<string>.MaxComparer.Instance;
+            (IComparer<Range<string>> minComparer, IComparer<Range<string>> maxComparer) = RangeComparerProvider.GetComparers(useLengthAwareComparer);
 
             // Compute the overall min and max from providedRanges
             string overallMin = providedRanges[0].Min;

Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs

@@ -85,7 +85,8 @@ public override async Task<ContainerQueryProperties> GetCachedContainerQueryProp
                 effectivePartitionKeyRange,
                 containerProperties.PartitionKey,
                 containerProperties.VectorEmbeddingPolicy,
-                containerProperties.GeospatialConfig.GeospatialType);
+                containerProperties.GeospatialConfig.GeospatialType,
+                this.documentClient.UseLengthAwareRangeComparer);
         }
 
         public override async Task<TryCatch<PartitionedQueryExecutionInfo>> TryGetPartitionedQueryExecutionInfoAsync(
@@ -240,7 +241,7 @@ public override async Task<List<PartitionKeyRange>> GetTargetPartitionKeyRangeBy
                     forceRefresh,
                     childTrace);
 
-                return QueryRangeUtils.LimitPartitionKeyRangesToProvidedRanges(ranges, providedRanges);
+                return QueryRangeUtils.LimitPartitionKeyRangesToProvidedRanges(ranges, providedRanges, this.clientContext.ClientOptions.UseLengthAwareRangeComparer);
             }
         }
 

Microsoft.Azure.Cosmos/src/Resource/ClientContextCore.cs

@@ -88,7 +88,8 @@ internal static CosmosClientContext Create(
                remoteCertificateValidationCallback: ClientContextCore.SslCustomValidationCallBack(clientOptions.GetServerCertificateCustomValidationCallback()),
                cosmosClientTelemetryOptions: clientOptions.CosmosClientTelemetryOptions,
                chaosInterceptorFactory: clientOptions.ChaosInterceptorFactory,
-               enableAsyncCacheExceptionNoSharing: clientOptions.EnableAsyncCacheExceptionNoSharing);
+               enableAsyncCacheExceptionNoSharing: clientOptions.EnableAsyncCacheExceptionNoSharing,
+               useLengthAwareRangeComparer: clientOptions.UseLengthAwareRangeComparer);
 
             return ClientContextCore.Create(
                 cosmosClient,

Microsoft.Azure.Cosmos/src/Routing/CollectionRoutingMap.cs

@@ -38,7 +38,8 @@ private CollectionRoutingMap(
             Dictionary<string, Tuple<PartitionKeyRange, ServiceIdentity>> rangeById,
             List<PartitionKeyRange> orderedPartitionKeyRanges,
             string collectionUniqueId,
-            string changeFeedNextIfNoneMatch)
+            string changeFeedNextIfNoneMatch,
+            bool useLengthAwareRangeComparer)
         {
             this.rangeById = rangeById;
             this.orderedPartitionKeyRanges = orderedPartitionKeyRanges;
@@ -71,18 +72,13 @@ private CollectionRoutingMap(
                     }
                     return range.Status == PartitionKeyRangeStatus.Offline ? CollectionRoutingMap.InvalidPkRangeId : pkId;
                 });
-
-            //LengthAwareComparer is the default Range comparer and flag <see cref="ConfigurationManager.UseLengthAwareRangeComparator"/>  is used to ovverride the default comparer to legacy Min/Max comparer.
-            bool useLengthAwareComparer = ConfigurationManager.IsLengthAwareRangeComparatorEnabled();
-
-            this.comparers = useLengthAwareComparer
-                ? (Range<string>.LengthAwareMinComparer.Instance, Range<string>.LengthAwareMaxComparer.Instance)
-                : (Range<string>.MinComparer.Instance, Range<string>.MaxComparer.Instance);
+            this.comparers = RangeComparerProvider.GetComparers(useLengthAwareRangeComparer);
         }
 
         public static CollectionRoutingMap TryCreateCompleteRoutingMap(
             IEnumerable<Tuple<PartitionKeyRange, ServiceIdentity>> ranges,
             string collectionUniqueId,
+            bool useLengthAwareRangeComparer,
             string changeFeedNextIfNoneMatch = null)
         {
             Dictionary<string, Tuple<PartitionKeyRange, ServiceIdentity>> rangeById =
@@ -102,7 +98,7 @@ public static CollectionRoutingMap TryCreateCompleteRoutingMap(
                 return null;
             }
 
-            return new CollectionRoutingMap(rangeById, orderedRanges, collectionUniqueId, changeFeedNextIfNoneMatch);
+            return new CollectionRoutingMap(rangeById, orderedRanges, collectionUniqueId, changeFeedNextIfNoneMatch, useLengthAwareRangeComparer);
         }
 
         public string CollectionUniqueId { get; private set; }
@@ -212,7 +208,8 @@ public ServiceIdentity TryGetInfoByPartitionKeyRangeId(string partitionKeyRangeI
 
         public CollectionRoutingMap TryCombine(
             IEnumerable<Tuple<PartitionKeyRange, ServiceIdentity>> ranges,
-            string changeFeedNextIfNoneMatch)
+            string changeFeedNextIfNoneMatch,
+            bool useLengthAwareComparer)
         {
             HashSet<string> newGoneRanges = new HashSet<string>(ranges.SelectMany(tuple => tuple.Item1.Parents ?? Enumerable.Empty<string>()));
             newGoneRanges.UnionWith(this.goneRanges);
@@ -239,7 +236,7 @@ public CollectionRoutingMap TryCombine(
                 return null;
             }
 
-            return new CollectionRoutingMap(newRangeById, newOrderedRanges, this.CollectionUniqueId, changeFeedNextIfNoneMatch);
+            return new CollectionRoutingMap(newRangeById, newOrderedRanges, this.CollectionUniqueId, changeFeedNextIfNoneMatch, useLengthAwareComparer);
         }
 
         private class MinPartitionKeyTupleComparer : IComparer<Tuple<PartitionKeyRange, ServiceIdentity>>

Microsoft.Azure.Cosmos/src/Routing/PartitionKeyRangeCache.cs

@@ -30,13 +30,15 @@ internal class PartitionKeyRangeCache : IRoutingMapProvider, ICollectionRoutingM
         private readonly ICosmosAuthorizationTokenProvider authorizationTokenProvider;
         private readonly IStoreModel storeModel;
         private readonly CollectionCache collectionCache;
-        private readonly IGlobalEndpointManager endpointManager;
+        private readonly IGlobalEndpointManager endpointManager;
+        private readonly bool useLengthAwareRangeComparer;
 
         public PartitionKeyRangeCache(
             ICosmosAuthorizationTokenProvider authorizationTokenProvider,
             IStoreModel storeModel,
             CollectionCache collectionCache,
-            IGlobalEndpointManager endpointManager,
+            IGlobalEndpointManager endpointManager,
+            bool useLengthAwareRangeComparer,
             bool enableAsyncCacheExceptionNoSharing = true)
         {
             this.routingMapCache = new AsyncCacheNonBlocking<string, CollectionRoutingMap>(
@@ -45,7 +47,8 @@ public PartitionKeyRangeCache(
             this.authorizationTokenProvider = authorizationTokenProvider;
             this.storeModel = storeModel;
             this.collectionCache = collectionCache;
-            this.endpointManager = endpointManager;
+            this.endpointManager = endpointManager;
+            this.useLengthAwareRangeComparer = useLengthAwareRangeComparer;
         }
 
         public virtual async Task<IReadOnlyList<PartitionKeyRange>> TryGetOverlappingRangesAsync(
@@ -241,12 +244,13 @@ private async Task<CollectionRoutingMap> GetRoutingMapForCollectionAsync(
                 HashSet<string> goneRanges = new HashSet<string>(ranges.SelectMany(range => range.Parents ?? Enumerable.Empty<string>()));
                 routingMap = CollectionRoutingMap.TryCreateCompleteRoutingMap(
                     tuples.Where(tuple => !goneRanges.Contains(tuple.Item1.Id)),
-                    string.Empty,
+                    string.Empty,
+                    false,
                     changeFeedNextIfNoneMatch);
             }
             else
             {
-                routingMap = previousRoutingMap.TryCombine(tuples, changeFeedNextIfNoneMatch);
+                routingMap = previousRoutingMap.TryCombine(tuples, changeFeedNextIfNoneMatch, this.useLengthAwareRangeComparer);
             }
 
             if (routingMap == null)

Microsoft.Azure.Cosmos/src/Routing/RangeComparerProvider.cs

@@ -0,0 +1,24 @@
+//------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+//------------------------------------------------------------
+
+namespace Microsoft.Azure.Cosmos.Routing
+{
+    using System.Collections.Generic;
+    using Microsoft.Azure.Documents.Routing;
+
+    internal class RangeComparerProvider
+    {
+        /// <summary>
+        /// Gets the minimum and maximum range comparers based on the current configuration.
+        /// </summary>
+        /// <returns>Tuple of (minComparer, maxComparer)</returns>
+        public static (IComparer<Range<string>> minComparer, IComparer<Range<string>> maxComparer) GetComparers(bool useLengthAwareComparison)
+        {
+            return (
+                useLengthAwareComparison ? Range<string>.LengthAwareMinComparer.Instance : Range<string>.MinComparer.Instance,
+                useLengthAwareComparison ? Range<string>.LengthAwareMaxComparer.Instance : Range<string>.MaxComparer.Instance
+            );
+        }
+    }
+}

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemTests.cs

@@ -2175,7 +2175,7 @@ public async Task ItemEpkQuerySingleKeyRangeValidation()
                     //new List<Documents.Routing.Range<string>> { new Documents.Routing.Range<string>("AA", "AA", true, true) },
                     containerResponse.Resource.PartitionKey,
                     vectorEmbeddingPolicy: null,
-                    containerResponse.Resource.GeospatialConfig.GeospatialType);
+                    containerResponse.Resource.GeospatialConfig.GeospatialType, false);
 
                 // There should only be one range since the EPK option is set.
                 List<PartitionKeyRange> partitionKeyRanges = await CosmosQueryExecutionContextFactory.GetTargetPartitionKeyRangesAsync(