Skip to content

Commit

Permalink
Subpartitioning: Fixes bug for queries on subpartitioned containers (#…
Browse files Browse the repository at this point in the history 
…3934) (#4085)

* initial fix, needs testing on prod

* test fix

* clean up pr

* query rework

* refactors previous changes

* requested changes and bug fixes

* nits

* requested changes

* bug fixes

* start of test

* added test

* nit: changed name of EffectivePartitionKeyRanges to EffectiveRangesForPartitionKey

* Address code comments

* Address code comments

* saving work

* addresses code comments

* nit, spacing

* PartitionKeyHash fixes

* Fixes bugs in tests

* Removed bad method, added additional test coverage

* Removed EffectivePartitionKeyString use

* test fix

* requested changes

* Requested changes

* fixed test

* Test fix

* Added comment

---------

Co-authored-by: SrinikhilReddy <[email protected]>
  • Loading branch information
@NaluTripician @SrinikhilReddy
NaluTripician and SrinikhilReddy authored Sep 15, 2023
1 parent aff544b commit 8c2c4c8
Show file tree
Hide file tree
Showing 16 changed files with 471 additions and 141 deletions.
94 changes: 51 additions & 43 deletions Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using global::Azure;
using Microsoft.Azure.Cosmos;
using Microsoft.Azure.Cosmos.CosmosElements;
using Microsoft.Azure.Cosmos.Pagination;
Expand All @@ -27,6 +28,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext
using Microsoft.Azure.Cosmos.SqlObjects;
using Microsoft.Azure.Cosmos.SqlObjects.Visitors;
using Microsoft.Azure.Cosmos.Tracing;
using Microsoft.Azure.Documents.Routing;

internal static class CosmosQueryExecutionContextFactory
{
Expand Down Expand Up @@ -211,10 +213,10 @@ private static async Task<TryCatch<IQueryPipelineStage>> TryCreateCoreContextAsy

// Only thing that matters is that we target the correct range.
Documents.PartitionKeyDefinition partitionKeyDefinition = GetPartitionKeyDefinition(inputParameters, containerQueryProperties);
List<Documents.PartitionKeyRange> targetRanges = await cosmosQueryContext.QueryClient.GetTargetPartitionKeyRangesByEpkStringAsync(
List<Documents.PartitionKeyRange> targetRanges = await cosmosQueryContext.QueryClient.GetTargetPartitionKeyRangesAsync(
cosmosQueryContext.ResourceLink,
containerQueryProperties.ResourceId,
containerQueryProperties.EffectivePartitionKeyString,
containerQueryProperties.EffectiveRangesForPartitionKey,
forceRefresh: false,
createQueryPipelineTrace);

Expand Down Expand Up @@ -635,67 +637,54 @@ private static async Task<PartitionedQueryExecutionInfo> GetPartitionedQueryExec
ITrace trace)
{
List<Documents.PartitionKeyRange> targetRanges;
if (containerQueryProperties.EffectivePartitionKeyString != null)
if (containerQueryProperties.EffectiveRangesForPartitionKey != null)
{
targetRanges = await queryClient.GetTargetPartitionKeyRangesByEpkStringAsync(
targetRanges = await queryClient.GetTargetPartitionKeyRangesAsync(
resourceLink,
containerQueryProperties.ResourceId,
containerQueryProperties.EffectivePartitionKeyString,
containerQueryProperties.EffectiveRangesForPartitionKey,
forceRefresh: false,
trace);
}
else if (TryGetEpkProperty(properties, out string effectivePartitionKeyString))
{
targetRanges = await queryClient.GetTargetPartitionKeyRangesByEpkStringAsync(
//Note that here we have no way to consume the EPK string as there is no way to convert
//the string to the partition key type to evaulate the number of components which needs to be done for the
//multihahs methods/classes. This is particually important for queries with prefix partition key.
//the EPK sting header is only for internal use but this needs to be fixed in the future.
List<Range<string>> effectiveRanges = new List<Range<string>>
{ Range<string>.GetPointRange(effectivePartitionKeyString) };

targetRanges = await queryClient.GetTargetPartitionKeyRangesAsync(
resourceLink,
containerQueryProperties.ResourceId,
effectivePartitionKeyString,
effectiveRanges,
forceRefresh: false,
trace);
}
else if (feedRangeInternal != null)
{
targetRanges = await queryClient.GetTargetPartitionKeyRangeByFeedRangeAsync(
resourceLink,
containerQueryProperties.ResourceId,
containerQueryProperties.PartitionKeyDefinition,
feedRangeInternal,
forceRefresh: false,
trace);
resourceLink,
containerQueryProperties.ResourceId,
containerQueryProperties.PartitionKeyDefinition,
feedRangeInternal,
forceRefresh: false,
trace);
}
else
{
targetRanges = await queryClient.GetTargetPartitionKeyRangesAsync(
resourceLink,
containerQueryProperties.ResourceId,
partitionedQueryExecutionInfo.QueryRanges,
forceRefresh: false,
trace);
targetRanges = await queryClient.GetTargetPartitionKeyRangesAsync(
resourceLink,
containerQueryProperties.ResourceId,
partitionedQueryExecutionInfo.QueryRanges,
forceRefresh: false,
trace);
}

return targetRanges;
}

private static void SetTestInjectionPipelineType(InputParameters inputParameters, string pipelineType)
{
TestInjections.ResponseStats responseStats = inputParameters?.TestInjections?.Stats;
if (responseStats != null)
{
if (pipelineType == OptimisticDirectExecution)
{
responseStats.PipelineType = TestInjections.PipelineType.OptimisticDirectExecution;
}
else if (pipelineType == Specialized)
{
responseStats.PipelineType = TestInjections.PipelineType.Specialized;
}
else
{
responseStats.PipelineType = TestInjections.PipelineType.Passthrough;
}
}
}

private static bool TryGetEpkProperty(
IReadOnlyDictionary<string, object> properties,
out string effectivePartitionKeyString)
Expand All @@ -718,6 +707,26 @@ private static bool TryGetEpkProperty(
return false;
}

private static void SetTestInjectionPipelineType(InputParameters inputParameters, string pipelineType)
{
TestInjections.ResponseStats responseStats = inputParameters?.TestInjections?.Stats;
if (responseStats != null)
{
if (pipelineType == OptimisticDirectExecution)
{
responseStats.PipelineType = TestInjections.PipelineType.OptimisticDirectExecution;
}
else if (pipelineType == Specialized)
{
responseStats.PipelineType = TestInjections.PipelineType.Specialized;
}
else
{
responseStats.PipelineType = TestInjections.PipelineType.Passthrough;
}
}
}

private static Documents.PartitionKeyDefinition GetPartitionKeyDefinition(InputParameters inputParameters, ContainerQueryProperties containerQueryProperties)
{
//todo:elasticcollections this may rely on information from collection cache which is outdated
Expand Down Expand Up @@ -771,14 +780,13 @@ private static Documents.PartitionKeyDefinition GetPartitionKeyDefinition(InputP
else
{
Documents.PartitionKeyDefinition partitionKeyDefinition = GetPartitionKeyDefinition(inputParameters, containerQueryProperties);
if (inputParameters.PartitionKey != null)
if (inputParameters.PartitionKey.HasValue)
{
Debug.Assert(partitionKeyDefinition != null, "CosmosQueryExecutionContextFactory Assert!", "PartitionKeyDefinition cannot be null if partitionKey is defined");

targetRanges = await cosmosQueryContext.QueryClient.GetTargetPartitionKeyRangesByEpkStringAsync(
targetRanges = await cosmosQueryContext.QueryClient.GetTargetPartitionKeyRangesAsync(
cosmosQueryContext.ResourceLink,
containerQueryProperties.ResourceId,
containerQueryProperties.EffectivePartitionKeyString,
containerQueryProperties.EffectiveRangesForPartitionKey,
forceRefresh: false,
trace);
}
Expand Down
11 changes: 8 additions & 3 deletions Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/ContainerQueryProperties.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,24 +4,29 @@

namespace Microsoft.Azure.Cosmos.Query.Core.QueryClient
{
using System.Collections.Generic;
using Microsoft.Azure.Documents;
using Microsoft.Azure.Documents.Routing;

internal readonly struct ContainerQueryProperties
{
public ContainerQueryProperties(
string resourceId,
string effectivePartitionKeyString,
IReadOnlyList<Range<string>> effectivePartitionKeyRanges,
PartitionKeyDefinition partitionKeyDefinition,
Cosmos.GeospatialType geospatialType)
{
this.ResourceId = resourceId;
this.EffectivePartitionKeyString = effectivePartitionKeyString;
this.EffectiveRangesForPartitionKey = effectivePartitionKeyRanges;
this.PartitionKeyDefinition = partitionKeyDefinition;
this.GeospatialType = geospatialType;
}

public string ResourceId { get; }
public string EffectivePartitionKeyString { get; }

//A PartitionKey has one range when it is a full PartitionKey value.
//It can span many it is a prefix PartitionKey for a sub-partitioned container.
public IReadOnlyList<Range<string>> EffectiveRangesForPartitionKey { get; }
public PartitionKeyDefinition PartitionKeyDefinition { get; }
public Cosmos.GeospatialType GeospatialType { get; }
}
Expand Down
9 changes: 1 addition & 8 deletions Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/CosmosQueryClient.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -76,13 +76,6 @@ public abstract Task<PartitionedQueryExecutionInfo> ExecuteQueryPlanRequestAsync

public abstract void ClearSessionTokenCache(string collectionFullName);

public abstract Task<List<Documents.PartitionKeyRange>> GetTargetPartitionKeyRangesByEpkStringAsync(
string resourceLink,
string collectionResourceId,
string effectivePartitionKeyString,
bool forceRefresh,
ITrace trace);

public abstract Task<List<Documents.PartitionKeyRange>> GetTargetPartitionKeyRangeByFeedRangeAsync(
string resourceLink,
string collectionResourceId,
Expand All @@ -94,7 +87,7 @@ public abstract Task<PartitionedQueryExecutionInfo> ExecuteQueryPlanRequestAsync
public abstract Task<List<Documents.PartitionKeyRange>> GetTargetPartitionKeyRangesAsync(
string resourceLink,
string collectionResourceId,
List<Documents.Routing.Range<string>> providedRanges,
IReadOnlyList<Documents.Routing.Range<string>> providedRanges,
bool forceRefresh,
ITrace trace);

Expand Down
35 changes: 13 additions & 22 deletions Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -63,17 +63,26 @@ public override async Task<ContainerQueryProperties> GetCachedContainerQueryProp
trace,
cancellationToken);

string effectivePartitionKeyString = null;
List<Range<string>> effectivePartitionKeyRange = null;
if (partitionKey != null)
{
// Dis-ambiguate the NonePK if used
PartitionKeyInternal partitionKeyInternal = partitionKey.Value.IsNone ? containerProperties.GetNoneValue() : partitionKey.Value.InternalKey;
effectivePartitionKeyString = partitionKeyInternal.GetEffectivePartitionKeyString(containerProperties.PartitionKey);
effectivePartitionKeyRange = new List<Range<string>>
{
PartitionKeyInternal.GetEffectivePartitionKeyRange(
containerProperties.PartitionKey,
new Range<PartitionKeyInternal>(
min: partitionKeyInternal,
max: partitionKeyInternal,
isMinInclusive: true,
isMaxInclusive: true))
};
}

return new ContainerQueryProperties(
containerProperties.ResourceId,
effectivePartitionKeyString,
effectivePartitionKeyRange,
containerProperties.PartitionKey,
containerProperties.GeospatialConfig.GeospatialType);
}
Expand Down Expand Up @@ -200,24 +209,6 @@ public override async Task<PartitionedQueryExecutionInfo> ExecuteQueryPlanReques
return partitionedQueryExecutionInfo;
}

public override Task<List<PartitionKeyRange>> GetTargetPartitionKeyRangesByEpkStringAsync(
string resourceLink,
string collectionResourceId,
string effectivePartitionKeyString,
bool forceRefresh,
ITrace trace)
{
return this.GetTargetPartitionKeyRangesAsync(
resourceLink,
collectionResourceId,
new List<Range<string>>
{
Range<string>.GetPointRange(effectivePartitionKeyString)
},
forceRefresh,
trace);
}

public override async Task<List<PartitionKeyRange>> GetTargetPartitionKeyRangeByFeedRangeAsync(
string resourceLink,
string collectionResourceId,
Expand All @@ -243,7 +234,7 @@ public override async Task<List<PartitionKeyRange>> GetTargetPartitionKeyRangeBy
public override async Task<List<PartitionKeyRange>> GetTargetPartitionKeyRangesAsync(
string resourceLink,
string collectionResourceId,
List<Range<string>> providedRanges,
IReadOnlyList<Range<string>> providedRanges,
bool forceRefresh,
ITrace trace)
{
Expand Down
30 changes: 26 additions & 4 deletions Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHash.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
namespace Microsoft.Azure.Cosmos.Routing
{
using System;
using System.Runtime.CompilerServices;
using System.Collections.Generic;
using System.Runtime.InteropServices;
using System.Text;
using Microsoft.Azure.Documents.Routing;
Expand Down Expand Up @@ -35,12 +35,34 @@ namespace Microsoft.Azure.Cosmos.Routing
/// </example>
internal readonly struct PartitionKeyHash : IComparable<PartitionKeyHash>, IEquatable<PartitionKeyHash>
{
private readonly IReadOnlyList<UInt128> values;

public PartitionKeyHash(UInt128 value)
: this(new UInt128[] { value })
{
this.Value = value;
}

public UInt128 Value { get; }
public PartitionKeyHash(UInt128[] values)
{
StringBuilder stringBuilder = new StringBuilder();
foreach (UInt128 value in values)
{
if (stringBuilder.Length > 0)
{
stringBuilder.Append('-');
}
stringBuilder.Append(value.ToString());
}

this.Value = stringBuilder.ToString();
this.values = values;
}

public readonly static PartitionKeyHash None = new PartitionKeyHash(0);

public string Value { get; }

internal readonly IReadOnlyList<UInt128> HashValues => this.values;

public int CompareTo(PartitionKeyHash other)
{
Expand All @@ -66,7 +88,7 @@ public override bool Equals(object obj)

public override int GetHashCode() => this.Value.GetHashCode();

public override string ToString() => this.Value.ToString();
public override string ToString() => this.Value;

public static bool TryParse(string value, out PartitionKeyHash parsedValue)
{
Expand Down
2 changes: 1 addition & 1 deletion Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHashRange.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -181,4 +181,4 @@ public override string ToString()
return stringBuilder.ToString();
}
}
}
}
14 changes: 7 additions & 7 deletions Microsoft.Azure.Cosmos/src/Routing/PartitionKeyHashRanges.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -132,9 +132,9 @@ public static CreateOutcome TryCreate(
{
if (partitionKeyHashRange.StartInclusive.HasValue)
{
if (partitionKeyHashRange.StartInclusive.Value.Value < minStart)
if (partitionKeyHashRange.StartInclusive.Value.HashValues[0] < minStart)
{
minStart = partitionKeyHashRange.StartInclusive.Value.Value;
minStart = partitionKeyHashRange.StartInclusive.Value.HashValues[0];
}
}
else
Expand All @@ -144,18 +144,18 @@ public static CreateOutcome TryCreate(

if (partitionKeyHashRange.EndExclusive.HasValue)
{
if (partitionKeyHashRange.EndExclusive.Value.Value > maxEnd)
if (partitionKeyHashRange.EndExclusive.Value.HashValues[0] > maxEnd)
{
maxEnd = partitionKeyHashRange.EndExclusive.Value.Value;
maxEnd = partitionKeyHashRange.EndExclusive.Value.HashValues[0];
}
}
else
{
maxEnd = UInt128.MaxValue;
}

UInt128 width = partitionKeyHashRange.EndExclusive.GetValueOrDefault(new PartitionKeyHash(UInt128.MaxValue)).Value
- partitionKeyHashRange.StartInclusive.GetValueOrDefault(new PartitionKeyHash(UInt128.MinValue)).Value;
UInt128 width = partitionKeyHashRange.EndExclusive.GetValueOrDefault(new PartitionKeyHash(UInt128.MaxValue)).HashValues[0]
- partitionKeyHashRange.StartInclusive.GetValueOrDefault(new PartitionKeyHash(UInt128.MinValue)).HashValues[0];
sumOfWidth += width;
if (sumOfWidth < width)
{
Expand Down Expand Up @@ -223,4 +223,4 @@ public enum CreateOutcome
Success,
}
}
}
}
Loading

0 comments on commit 8c2c4c8

Please sign in to comment.