diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/HierarchicalPartitionUtils.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/HierarchicalPartitionUtils.cs index 2b91c7cb40..34771387fe 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/HierarchicalPartitionUtils.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/HierarchicalPartitionUtils.cs @@ -5,10 +5,13 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition { using System; + using System.Collections.Generic; using Microsoft.Azure.Cosmos.Query.Core.QueryClient; + using Microsoft.Azure.Documents.Routing; internal static class HierarchicalPartitionUtils { + private static readonly bool IsLengthAwareComparisonEnabled = ConfigurationManager.IsLengthAwareRangeComparatorEnabled(); /// /// Updates the FeedRange to limit the scope of incoming feedRange to logical partition within a single physical partition. /// Generally speaking, a subpartitioned container can experience split partition at any level of hierarchical partition key. @@ -50,7 +53,16 @@ public static FeedRangeInternal LimitFeedRangeToSinglePartition(PartitionKey? pa String overlappingMax; bool maxInclusive; - if (Documents.Routing.Range.MinComparer.Instance.Compare( + //LengthAwareComparer is the default Range comparer and flag is used to ovverride the default comparer to legacy Min/Max comparer. + IComparer> minComparer = IsLengthAwareComparisonEnabled + ? Documents.Routing.Range.LengthAwareMinComparer.Instance + : Documents.Routing.Range.MinComparer.Instance; + + IComparer> maxComparer = IsLengthAwareComparisonEnabled + ? Documents.Routing.Range.LengthAwareMaxComparer.Instance + : Documents.Routing.Range.MaxComparer.Instance; + + if (minComparer.Compare( epkForPartitionKey, feedRangeEpk.Range) < 0) { @@ -63,7 +75,7 @@ public static FeedRangeInternal LimitFeedRangeToSinglePartition(PartitionKey? pa minInclusive = epkForPartitionKey.IsMinInclusive; } - if (Documents.Routing.Range.MaxComparer.Instance.Compare( + if (maxComparer.Compare( epkForPartitionKey, feedRangeEpk.Range) > 0) { diff --git a/Microsoft.Azure.Cosmos/src/Routing/CollectionRoutingMap.cs b/Microsoft.Azure.Cosmos/src/Routing/CollectionRoutingMap.cs index ef387dd61d..e88c090248 100644 --- a/Microsoft.Azure.Cosmos/src/Routing/CollectionRoutingMap.cs +++ b/Microsoft.Azure.Cosmos/src/Routing/CollectionRoutingMap.cs @@ -32,18 +32,7 @@ internal sealed class CollectionRoutingMap internal int HighestNonOfflinePkRangeId { get; private set; } - public CollectionRoutingMap( - CollectionRoutingMap collectionRoutingMap, - string changeFeedNextIfNoneMatch) - { - this.rangeById = new Dictionary>(collectionRoutingMap.rangeById); - this.orderedPartitionKeyRanges = new List(collectionRoutingMap.orderedPartitionKeyRanges); - this.orderedRanges = new List>(collectionRoutingMap.orderedRanges); - this.goneRanges = new HashSet(collectionRoutingMap.goneRanges); - this.HighestNonOfflinePkRangeId = collectionRoutingMap.HighestNonOfflinePkRangeId; - this.CollectionUniqueId = collectionRoutingMap.CollectionUniqueId; - this.ChangeFeedNextIfNoneMatch = changeFeedNextIfNoneMatch; - } + private readonly (IComparer> MinComparer, IComparer> MaxComparer) comparers; private CollectionRoutingMap( Dictionary> rangeById, @@ -82,6 +71,13 @@ private CollectionRoutingMap( } return range.Status == PartitionKeyRangeStatus.Offline ? CollectionRoutingMap.InvalidPkRangeId : pkId; }); + + //LengthAwareComparer is the default Range comparer and flag is used to ovverride the default comparer to legacy Min/Max comparer. + bool useLengthAwareComparer = ConfigurationManager.IsLengthAwareRangeComparatorEnabled(); + + this.comparers = useLengthAwareComparer + ? (Range.LengthAwareMinComparer.Instance, Range.LengthAwareMaxComparer.Instance) + : (Range.MinComparer.Instance, Range.MaxComparer.Instance); } public static CollectionRoutingMap TryCreateCompleteRoutingMap( @@ -142,13 +138,13 @@ public IReadOnlyList GetOverlappingRanges(IReadOnlyList providedRange in providedPartitionKeyRanges) { - int minIndex = this.orderedRanges.BinarySearch(providedRange, Range.MinComparer.Instance); + int minIndex = this.orderedRanges.BinarySearch(providedRange, this.comparers.MinComparer); if (minIndex < 0) { minIndex = Math.Max(0, (~minIndex) - 1); } - int maxIndex = this.orderedRanges.BinarySearch(providedRange, Range.MaxComparer.Instance); + int maxIndex = this.orderedRanges.BinarySearch(providedRange, this.comparers.MaxComparer); if (maxIndex < 0) { maxIndex = Math.Min(this.OrderedPartitionKeyRanges.Count - 1, ~maxIndex); diff --git a/Microsoft.Azure.Cosmos/src/Routing/PartitionRoutingHelper.cs b/Microsoft.Azure.Cosmos/src/Routing/PartitionRoutingHelper.cs index 1c0f859eab..f6a346abb6 100644 --- a/Microsoft.Azure.Cosmos/src/Routing/PartitionRoutingHelper.cs +++ b/Microsoft.Azure.Cosmos/src/Routing/PartitionRoutingHelper.cs @@ -25,6 +25,7 @@ namespace Microsoft.Azure.Cosmos.Routing internal class PartitionRoutingHelper { + private static readonly bool IsLengthAwareComparisonEnabled = ConfigurationManager.IsLengthAwareRangeComparatorEnabled(); public static IReadOnlyList> GetProvidedPartitionKeyRanges( string querySpecJsonString, bool enableCrossPartitionQuery, @@ -231,9 +232,11 @@ await routingMapProvider.TryGetRangeByEffectivePartitionKeyAsync( return new ResolvedRangeInfo(lastPartitionKeyRange, suppliedTokens); } + (IComparer> minComparer, _) = this.GetComparers(IsLengthAwareComparisonEnabled); + Range minimumRange = PartitionRoutingHelper.Min( providedPartitionKeyRanges, - Range.MinComparer.Instance); + minComparer); return new ResolvedRangeInfo( await routingMapProvider.TryGetRangeByEffectivePartitionKeyAsync(collectionRid, minimumRange.Min, trace), @@ -347,6 +350,8 @@ public virtual async Task TryAddPartitionKeyRangeToContinuationTokenAsync( // We only need to get the next range if we have to if (string.IsNullOrEmpty(backendResponseHeaders[HttpConstants.HttpHeaders.Continuation])) { + (IComparer> minComparer, IComparer> maxComparer) = this.GetComparers(IsLengthAwareComparisonEnabled); + if (direction == RntdbEnumerationDirection.Reverse) { rangeToUse = PartitionRoutingHelper.MinBefore( @@ -354,14 +359,15 @@ public virtual async Task TryAddPartitionKeyRangeToContinuationTokenAsync( collectionRid, providedPartitionKeyRanges.Single(), trace)).ToList(), - currentRange); + currentRange, + minComparer); } else { Range nextProvidedRange = PartitionRoutingHelper.MinAfter( providedPartitionKeyRanges, currentRange.ToRange(), - Range.MaxComparer.Instance); + maxComparer); if (nextProvidedRange == null) { @@ -547,14 +553,14 @@ private static T MinAfter(IReadOnlyList values, T minValue, IComparer c return min; } - private static PartitionKeyRange MinBefore(IReadOnlyList values, PartitionKeyRange minValue) + private static PartitionKeyRange MinBefore(IReadOnlyList values, PartitionKeyRange minValue, + IComparer> comparer) { if (values.Count == 0) { throw new ArgumentException(nameof(values)); } - IComparer> comparer = Range.MinComparer.Instance; PartitionKeyRange min = null; foreach (PartitionKeyRange value in values) { @@ -566,6 +572,15 @@ private static PartitionKeyRange MinBefore(IReadOnlyList valu return min; } + + //LengthAwareComparer is the default Range comparer and flag is used to ovverride the default comparer to legacy Min/Max comparer. + private (IComparer> minComparer, IComparer> maxComparer) GetComparers(bool useLengthAwareComparison) + { + return ( + useLengthAwareComparison ? Range.LengthAwareMinComparer.Instance : Range.MinComparer.Instance, + useLengthAwareComparison ? Range.LengthAwareMaxComparer.Instance : Range.MaxComparer.Instance + ); + } public readonly struct ResolvedRangeInfo { diff --git a/Microsoft.Azure.Cosmos/src/Util/ConfigurationManager.cs b/Microsoft.Azure.Cosmos/src/Util/ConfigurationManager.cs index 0923257e20..b5b4d8eec2 100644 --- a/Microsoft.Azure.Cosmos/src/Util/ConfigurationManager.cs +++ b/Microsoft.Azure.Cosmos/src/Util/ConfigurationManager.cs @@ -117,6 +117,14 @@ internal static class ConfigurationManager /// internal static readonly string BypassQueryParsing = "AZURE_COSMOS_BYPASS_QUERY_PARSING"; + /// + /// A read-only string containing the environment variable name for disabling length aware range comparator. + /// Length aware range comparators were intorduced 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. + /// + internal static readonly string UseLengthAwareRangeComparator = "AZURE_COSMOS_USE_LENGTH_AWARE_RANGE_COMPARATOR"; + public static T GetEnvironmentVariable(string variable, T defaultValue) { string value = Environment.GetEnvironmentVariable(variable); @@ -376,5 +384,22 @@ public static bool ForceBypassQueryParsing() variable: ConfigurationManager.BypassQueryParsing, defaultValue: false); } + + /// + /// Gets the boolean value indicating if length-aware range comparator is enabled. + /// Default: true for preview , false for GA. + /// + /// A boolean flag indicating if length-aware range comparator is enabled. + public static bool IsLengthAwareRangeComparatorEnabled() + { + bool defaultValue = false; +#if PREVIEW && !INTERNAL + defaultValue = true; +#endif + return ConfigurationManager + .GetEnvironmentVariable( + variable: ConfigurationManager.UseLengthAwareRangeComparator, + defaultValue: defaultValue); + } } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CollectionRoutingMapTest.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CollectionRoutingMapTest.cs index 7e92be0f9f..cb7baf2de5 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CollectionRoutingMapTest.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CollectionRoutingMapTest.cs @@ -116,6 +116,299 @@ public void TestCollectionRoutingMap() Assert.AreEqual("2", partitionKeyRanges1.ElementAt(1).Id); } + /// + /// Validates that CollectionRoutingMap correctly identifies overlapping partition key ranges + /// when using length-aware range comparators. + /// This test ensures that EPK advanced comparison logic are applied as expected, + /// and that the routing map's behavior is consistent regardless if the input EPK is fully or partially specified. + /// The test covers scenarios where input EPKs are partial or fall on range boundaries, + /// verifying that the correct partition key ranges are returned when using the new LengthAware comparators. + /// + [TestMethod] + [DataRow(false)] + [DataRow(true)] + public void TestCollectionRoutingMapWithLengthAwareRangeComparators(bool isRoutingMapFullySpecified) + { + try + { + // Arrange: Set environment variable to "true" since the default is only true for Preview. + Environment.SetEnvironmentVariable(ConfigurationManager.UseLengthAwareRangeComparator, "true"); + + CollectionRoutingMap routingMap = this.GenerateRoutingMap(isRoutingMapFullySpecified); + + // Test scenario 1.1: Input EPK is partial and falls on the boundary between two overlapping ranges. + // The LengthAware comparators are able to correctly compare partial and full EPK ranges.Routing map is hybrid of fully specified and partially specified EPK ranges. + // Input Min EPK 06AB34CFE4E482236BCACBBF50E234AB matches (significant bytes) with maxEPK of pkrangeid 1 and minEPK of pkrangeid 2. + Range inputPkRange = new Range( + "06AB34CFE4E482236BCACBBF50E234AB", + "06AB34CFE4E482236BCACBBF50E234ABFF", + true, + false); + + // Expected outcome: Only partition key range with id 2 overlaps, as the LengthAware comparator correctly handles the partial EPK. + IReadOnlyList partitionKeyRanges1 = routingMap.GetOverlappingRanges(inputPkRange); + Assert.AreEqual(1, partitionKeyRanges1.Count); + Assert.AreEqual("2", partitionKeyRanges1[0].Id); + + // Test scenario 1.2: Input EPK falls on a boundary and maxEPK also matches the next range's max. + // The LengthAware comparator should return only the correct overlapping range. + inputPkRange = new Range( + "0BD3FBE846AF75790CE63F78B1A81631", + "0BD3FBE846AF75790CE63F78B1A81631FF", + true, + false); + + partitionKeyRanges1 = routingMap.GetOverlappingRanges(inputPkRange); + Assert.AreEqual(1, partitionKeyRanges1.Count); + CollectionAssert.AreEquivalent(new[] { "11" }, partitionKeyRanges1.Select(r => r.Id).ToArray()); + + inputPkRange = new Range( + "0D4DC2CD8F49C65A8E0C5306B61B43440D4DC2CD8F49C65A8E0C5306B61B4343", + "0D4DC2CD8F49C65A8E0C5306B61B43440D4DC2CD8F49C65A8E0C5306B61B4344", + true, + false); + + partitionKeyRanges1 = routingMap.GetOverlappingRanges(inputPkRange); + Assert.AreEqual(1, partitionKeyRanges1.Count); + CollectionAssert.AreEquivalent(new[] { "4" }, partitionKeyRanges1.Select(r => r.Id).ToArray()); + + // Test scenario 1.2 (continued): Input EPK falls in boundary and maxEPK also matches the next range's max. + inputPkRange = new Range( + "0BD3FBE846AF75790CE63F78B1A81620", + "0BD3FBE846AF75790CE63F78B1A81631", + true, + false); + + partitionKeyRanges1 = routingMap.GetOverlappingRanges(inputPkRange); + Assert.AreEqual(1, partitionKeyRanges1.Count); + CollectionAssert.AreEquivalent(new[] { "3" }, partitionKeyRanges1.Select(r => r.Id).ToArray()); + + // Test scenario 1.3: Input EPK is partial and spans two overlapping ranges. + /// Input Min EPK 0DCEB8CE51C6BFE84F4BD9409F69B9BB falls in both pkrangeid 4 and pkrangeid 5. + inputPkRange = new Range( + "0DCEB8CE51C6BFE84F4BD9409F69B9BB", + "0DCEB8CE51C6BFE84F4BD9409F69B9BBFF", + true, + false); + + partitionKeyRanges1 = routingMap.GetOverlappingRanges(inputPkRange); + Assert.AreEqual(2, partitionKeyRanges1.Count); + CollectionAssert.AreEquivalent(new[] { "24", "5" }, partitionKeyRanges1.Select(r => r.Id).ToArray()); + + + ///Test scenario 1.4: Input EPK is partial and falls in a single range in the middle. Routing map is hybrid of fully specified and partially specified ranges. + inputPkRange = new Range( + "02559A67F2724111B5E565DFA8711A00", + "02559A67F2724111B5E565DFA8711A00", + true, + true); + + partitionKeyRanges1 = routingMap.GetOverlappingRanges(inputPkRange); + Assert.AreEqual(1, partitionKeyRanges1.Count); + Assert.AreEqual("0", partitionKeyRanges1[0].Id); + + + ///Test scenario 1.5: Input EPK is partial and falls in a single range in the middle. Routing map targeted range has partial EPK values only. + inputPkRange = new Range( + "0D4DC2CD8F49C65A8E0C5306B61B4345", + "0D4DC2CD8F49C65A8E0C5306B61B4345", + true, + true); + + partitionKeyRanges1 = routingMap.GetOverlappingRanges(inputPkRange); + Assert.AreEqual(1, partitionKeyRanges1.Count); + Assert.AreEqual("4", partitionKeyRanges1[0].Id); + + + // The following part of the test case verifies the routing map values i.e.backend ranges when they are not fully specified. + if (!isRoutingMapFullySpecified) + { + // Test scenario 1.6: Input EPK is fully specified and backend range is partially specified. + // The LengthAware comparator correctly matches the fully specified input to the partially specified backend range. + inputPkRange = new Range( + "0D4DC2CD8F49C65A8E0C5306B61B434300000000000000000000000000000000", + "0D4EC2CD8F49C65A8E0C5306B61B434300000000000000000000000000000000", + true, + false); + + // LengthAware comparator yields only the correct range. + partitionKeyRanges1 = routingMap.GetOverlappingRanges(inputPkRange); + Assert.AreEqual(1, partitionKeyRanges1.Count); + CollectionAssert.AreEquivalent(new[] { "4" }, partitionKeyRanges1.Select(r => r.Id).ToArray()); + } + } + finally + { + // Clean up: Remove the environment variable after the test. + Environment.SetEnvironmentVariable(ConfigurationManager.UseLengthAwareRangeComparator, null); + } + } + + // Test GetOverlappingRanges behavior when the UseLengthAwareRangeComparator environment flag is set to false, + // which forces the use of legacy Min/Max comparators. + [TestMethod] + public void TestLegacyComparatorsUsedWhenLengthAwareComparatorFlagIsFalse() + { + try + { + // Arrange: Set environment variable to force legacy comparator usage. + Environment.SetEnvironmentVariable(ConfigurationManager.UseLengthAwareRangeComparator, "false"); + CollectionRoutingMap routingMap = this.GenerateRoutingMap(false); + + + // Test scenario: Input EPK is partial and falls on the boundary between two overlapping ranges. + // With the environment flag set, the routing map uses legacy Min/Max comparators, which do not distinguish + // between partial and full EPKs. As a result, both partition key ranges with ids 1 and 2 are considered overlapping. + // Input Min EPK 06AB34CFE4E482236BCACBBF50E234AB matches (significant bytes) with maxEPK of pkrangeid 1 and minEPK of pkrangeid 2. + Range inputPkRange = new Range( + "06AB34CFE4E482236BCACBBF50E234AB", + "06AB34CFE4E482236BCACBBF50E234ABFF", + true, + false); + IReadOnlyList partitionKeyRanges1 = routingMap.GetOverlappingRanges(inputPkRange); + Assert.AreEqual(2, partitionKeyRanges1.Count); + CollectionAssert.AreEquivalent(new[] { "1", "2" }, partitionKeyRanges1.Select(r => r.Id).ToArray()); + } + finally + { + Environment.SetEnvironmentVariable(ConfigurationManager.UseLengthAwareRangeComparator, null); + } + } + + private CollectionRoutingMap GenerateRoutingMap(bool isFullySpecified) + { + IEnumerable> partitionKeyRangeTuples = new[] + { + Tuple.Create( + new PartitionKeyRange + { + Id = "0", + MinInclusive = "", + MaxExclusive = "03559A67F2724111B5E565DFA8711A00" + }, + (ServiceIdentity)null), + + Tuple.Create( + new PartitionKeyRange + { + Id = "1", + MinInclusive = "03559A67F2724111B5E565DFA8711A00", + MaxExclusive = "06AB34CFE4E482236BCACBBF50E234AB00000000000000000000000000000000" + }, + (ServiceIdentity)null), + + Tuple.Create( + new PartitionKeyRange + { + Id = "2", + MinInclusive = "06AB34CFE4E482236BCACBBF50E234AB00000000000000000000000000000000", + MaxExclusive = "0BD3FBE846AF75790CE63F78B1A81620" + }, + (ServiceIdentity)null), + + Tuple.Create( + new PartitionKeyRange + { + Id = "3", + MinInclusive = "0BD3FBE846AF75790CE63F78B1A81620", + MaxExclusive = "0BD3FBE846AF75790CE63F78B1A8163100000000000000000000000000000000" + }, + (ServiceIdentity)null), + Tuple.Create( + new PartitionKeyRange + { + Id = "11", + MinInclusive = "0BD3FBE846AF75790CE63F78B1A8163100000000000000000000000000000000", + MaxExclusive = "0BD3FBE846AF75790CE63F78B1A81631FF" + }, + (ServiceIdentity)null), + Tuple.Create( + new PartitionKeyRange + { + Id = "12", + MinInclusive = "0BD3FBE846AF75790CE63F78B1A81631FF", + MaxExclusive = "0D4DC2CD8F49C65A8E0C5306B61B4343" + }, + (ServiceIdentity)null), + + Tuple.Create( + new PartitionKeyRange + { + Id = "4", + MinInclusive = "0D4DC2CD8F49C65A8E0C5306B61B4343", + MaxExclusive = "0D4EC2CD8F49C65A8E0C5306B61B4343" + }, + (ServiceIdentity)null), + + Tuple.Create( + new PartitionKeyRange + { + Id = "44", + MinInclusive = "0D4EC2CD8F49C65A8E0C5306B61B4343", + MaxExclusive = "0D5DC2CD8F49C65A8E0C5306B61B4343" + }, + (ServiceIdentity)null), + + Tuple.Create( + new PartitionKeyRange + { + Id = "24", + MinInclusive = "0D5DC2CD8F49C65A8E0C5306B61B4343", + MaxExclusive = "0DCEB8CE51C6BFE84F4BD9409F69B9BB2164DEBD78C50C850E0C1E3E3F0579ED" + }, + (ServiceIdentity)null), + + Tuple.Create( + new PartitionKeyRange + { + Id = "5", + MinInclusive = "0DCEB8CE51C6BFE84F4BD9409F69B9BB2164DEBD78C50C850E0C1E3E3F0579ED", + MaxExclusive = "1080F600C27CF98DC13F8639E94E7676" + }, + (ServiceIdentity)null), + Tuple.Create( + new PartitionKeyRange + { + Id = "9", + MinInclusive = "1080F600C27CF98DC13F8639E94E7676", + MaxExclusive = "FF" + }, + (ServiceIdentity)null), + }; + + if (isFullySpecified) + { + partitionKeyRangeTuples = partitionKeyRangeTuples + .Select(tuple => + { + PartitionKeyRange range = tuple.Item1; + // Pad right to 64 bytes (128 hex chars) for MinInclusive and MaxExclusive if not empty + string PadTo64(string value) + { + if (string.IsNullOrEmpty(value) || value == "FF") + return value; + return value.PadRight(64, '0'); + } + return Tuple.Create( + new PartitionKeyRange + { + Id = range.Id, + MinInclusive = PadTo64(range.MinInclusive), + MaxExclusive = PadTo64(range.MaxExclusive) + }, + tuple.Item2 + ); + }) + .ToList(); + } + + CollectionRoutingMap routingMap = CollectionRoutingMap.TryCreateCompleteRoutingMap( + partitionKeyRangeTuples, + string.Empty); + + return routingMap; + } + [TestMethod] [ExpectedException(typeof(InvalidOperationException))] public void TestInvalidRoutingMap()