diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs index fdca4ee452..d778db2438 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs @@ -1,470 +1,472 @@ -//------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -//------------------------------------------------------------ -namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext -{ - using System; - using System.Collections.Generic; - using System.Diagnostics; - using System.Linq; - using System.Text.RegularExpressions; - using System.Threading; - using System.Threading.Tasks; - using Microsoft.Azure.Cosmos; - using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Pagination; - using Microsoft.Azure.Cosmos.Query.Core; - using Microsoft.Azure.Cosmos.Query.Core.Exceptions; - using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Cosmos.Query.Core.Parser; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.Parallel; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.OptimisticDirectExecutionQuery; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Pagination; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Tokens; - using Microsoft.Azure.Cosmos.Query.Core.QueryClient; - using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; - 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 - { - internal const string ClientDisableOptimisticDirectExecution = "clientDisableOptimisticDirectExecution"; - private const string InternalPartitionKeyDefinitionProperty = "x-ms-query-partitionkey-definition"; - private const string QueryInspectionPattern = @"\s*(GROUP\s+BY\s+|COUNT\s*\(|MIN\s*\(|MAX\s*\(|AVG\s*\(|SUM\s*\(|DISTINCT\s+)"; - private const string OptimisticDirectExecution = "OptimisticDirectExecution"; - private const string Passthrough = "Passthrough"; - private const string Specialized = "Specialized"; - private const int PageSizeFactorForTop = 5; - private static readonly Regex QueryInspectionRegex = new Regex(QueryInspectionPattern, RegexOptions.IgnoreCase | RegexOptions.Compiled); - - public static IQueryPipelineStage Create( - DocumentContainer documentContainer, - CosmosQueryContext cosmosQueryContext, - InputParameters inputParameters, - ITrace trace) - { - if (cosmosQueryContext == null) - { - throw new ArgumentNullException(nameof(cosmosQueryContext)); - } - - if (inputParameters == null) - { - throw new ArgumentNullException(nameof(inputParameters)); - } - - if (trace == null) - { - throw new ArgumentNullException(nameof(trace)); - } - - NameCacheStaleRetryQueryPipelineStage nameCacheStaleRetryQueryPipelineStage = new NameCacheStaleRetryQueryPipelineStage( - cosmosQueryContext: cosmosQueryContext, - queryPipelineStageFactory: () => - { - // Query Iterator requires that the creation of the query context is deferred until the user calls ReadNextAsync - AsyncLazy> lazyTryCreateStage = new AsyncLazy>( +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ +namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext +{ + using System; + using System.Collections.Generic; + using System.Diagnostics; + using System.Linq; + using System.Text.RegularExpressions; + using System.Threading; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos; + using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Pagination; + using Microsoft.Azure.Cosmos.Query.Core; + using Microsoft.Azure.Cosmos.Query.Core.Exceptions; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Parser; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.Parallel; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.OptimisticDirectExecutionQuery; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Pagination; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Tokens; + using Microsoft.Azure.Cosmos.Query.Core.QueryClient; + using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; + 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 + { + internal const string ClientDisableOptimisticDirectExecution = "clientDisableOptimisticDirectExecution"; + private const string InternalPartitionKeyDefinitionProperty = "x-ms-query-partitionkey-definition"; + private const string QueryInspectionPattern = @"\s*(GROUP\s+BY\s+|COUNT\s*\(|MIN\s*\(|MAX\s*\(|AVG\s*\(|SUM\s*\(|DISTINCT\s+)"; + private const string OptimisticDirectExecution = "OptimisticDirectExecution"; + private const string Passthrough = "Passthrough"; + private const string Specialized = "Specialized"; + private const int PageSizeFactorForTop = 5; + private static readonly Regex QueryInspectionRegex = new Regex(QueryInspectionPattern, RegexOptions.IgnoreCase | RegexOptions.Compiled); + + public static IQueryPipelineStage Create( + DocumentContainer documentContainer, + CosmosQueryContext cosmosQueryContext, + InputParameters inputParameters, + ITrace trace) + { + if (cosmosQueryContext == null) + { + throw new ArgumentNullException(nameof(cosmosQueryContext)); + } + + if (inputParameters == null) + { + throw new ArgumentNullException(nameof(inputParameters)); + } + + if (trace == null) + { + throw new ArgumentNullException(nameof(trace)); + } + + NameCacheStaleRetryQueryPipelineStage nameCacheStaleRetryQueryPipelineStage = new NameCacheStaleRetryQueryPipelineStage( + cosmosQueryContext: cosmosQueryContext, + queryPipelineStageFactory: () => + { + // Query Iterator requires that the creation of the query context is deferred until the user calls ReadNextAsync + AsyncLazy> lazyTryCreateStage = new AsyncLazy>( valueFactory: (trace, innerCancellationToken) => TryCreateCoreContextAsync( - documentContainer, - cosmosQueryContext, - inputParameters, - trace, - innerCancellationToken)); - + documentContainer, + cosmosQueryContext, + inputParameters, + trace, + innerCancellationToken)); + LazyQueryPipelineStage lazyQueryPipelineStage = new LazyQueryPipelineStage(lazyTryCreateStage: lazyTryCreateStage); - return lazyQueryPipelineStage; - }); - + return lazyQueryPipelineStage; + }); + CatchAllQueryPipelineStage catchAllQueryPipelineStage = new CatchAllQueryPipelineStage(nameCacheStaleRetryQueryPipelineStage); - return catchAllQueryPipelineStage; - } - - private static async Task> TryCreateCoreContextAsync( - DocumentContainer documentContainer, - CosmosQueryContext cosmosQueryContext, - InputParameters inputParameters, - ITrace trace, - CancellationToken cancellationToken) - { - // The default - using (ITrace createQueryPipelineTrace = trace.StartChild("Create Query Pipeline", TraceComponent.Query, Tracing.TraceLevel.Info)) - { - // Try to parse the continuation token. - CosmosElement continuationToken = inputParameters.InitialUserContinuationToken; - PartitionedQueryExecutionInfo queryPlanFromContinuationToken = inputParameters.PartitionedQueryExecutionInfo; - if (continuationToken != null) - { - if (!PipelineContinuationToken.TryCreateFromCosmosElement( - continuationToken, - out PipelineContinuationToken pipelineContinuationToken)) - { - return TryCatch.FromException( - new MalformedContinuationTokenException( - $"Malformed {nameof(PipelineContinuationToken)}: {continuationToken}.")); - } - - if (PipelineContinuationToken.IsTokenFromTheFuture(pipelineContinuationToken)) - { - return TryCatch.FromException( - new MalformedContinuationTokenException( - $"{nameof(PipelineContinuationToken)} Continuation token is from a newer version of the SDK. " + - $"Upgrade the SDK to avoid this issue." + - $"{continuationToken}.")); - } - - if (!PipelineContinuationToken.TryConvertToLatest( - pipelineContinuationToken, - out PipelineContinuationTokenV1_1 latestVersionPipelineContinuationToken)) - { - return TryCatch.FromException( - new MalformedContinuationTokenException( - $"{nameof(PipelineContinuationToken)}: '{continuationToken}' is no longer supported.")); - } - - continuationToken = latestVersionPipelineContinuationToken.SourceContinuationToken; - if (latestVersionPipelineContinuationToken.QueryPlan != null) - { - queryPlanFromContinuationToken = latestVersionPipelineContinuationToken.QueryPlan; - } - } - - CosmosQueryClient cosmosQueryClient = cosmosQueryContext.QueryClient; - - ContainerQueryProperties containerQueryProperties = await cosmosQueryClient.GetCachedContainerQueryPropertiesAsync( - cosmosQueryContext.ResourceLink, - inputParameters.PartitionKey, - createQueryPipelineTrace, - cancellationToken); - cosmosQueryContext.ContainerResourceId = containerQueryProperties.ResourceId; - - Documents.PartitionKeyRange targetRange = await TryGetTargetRangeOptimisticDirectExecutionAsync( - inputParameters, - queryPlanFromContinuationToken, - cosmosQueryContext, - containerQueryProperties, - trace); - - if (targetRange != null) - { - return await TryCreateSinglePartitionExecutionContextAsync( - documentContainer, - partitionedQueryExecutionInfo: null, - cosmosQueryContext, - containerQueryProperties, - inputParameters, - targetRange, - createQueryPipelineTrace, - cancellationToken); - } - - PartitionedQueryExecutionInfo partitionedQueryExecutionInfo; - if (queryPlanFromContinuationToken != null) - { - partitionedQueryExecutionInfo = queryPlanFromContinuationToken; - } - else - { - // If the query would go to gateway, but we have a partition key, - // then try seeing if we can execute as a passthrough using client side only logic. - // This is to short circuit the need to go to the gateway to get the query plan. - if (cosmosQueryContext.QueryClient.BypassQueryParsing() - && inputParameters.PartitionKey.HasValue) - { - bool parsed; - SqlQuery sqlQuery; - using (ITrace queryParseTrace = createQueryPipelineTrace.StartChild("Parse Query", TraceComponent.Query, Tracing.TraceLevel.Info)) - { - parsed = SqlQueryParser.TryParse(inputParameters.SqlQuerySpec.QueryText, out sqlQuery); - } - - if (parsed) - { - bool hasDistinct = sqlQuery.SelectClause.HasDistinct; - bool hasGroupBy = sqlQuery.GroupByClause != default; - bool hasAggregates = AggregateProjectionDetector.HasAggregate(sqlQuery.SelectClause.SelectSpec); - bool createPassthroughQuery = !hasAggregates && !hasDistinct && !hasGroupBy; - - if (createPassthroughQuery) - { - SetTestInjectionPipelineType(inputParameters, Passthrough); - - // Only thing that matters is that we target the correct range. - Documents.PartitionKeyDefinition partitionKeyDefinition = GetPartitionKeyDefinition(inputParameters, containerQueryProperties); - List targetRanges = await cosmosQueryContext.QueryClient.GetTargetPartitionKeyRangesAsync( - cosmosQueryContext.ResourceLink, - containerQueryProperties.ResourceId, - containerQueryProperties.EffectiveRangesForPartitionKey, - forceRefresh: false, - createQueryPipelineTrace); - + return catchAllQueryPipelineStage; + } + + private static async Task> TryCreateCoreContextAsync( + DocumentContainer documentContainer, + CosmosQueryContext cosmosQueryContext, + InputParameters inputParameters, + ITrace trace, + CancellationToken cancellationToken) + { + // The default + using (ITrace createQueryPipelineTrace = trace.StartChild("Create Query Pipeline", TraceComponent.Query, Tracing.TraceLevel.Info)) + { + // Try to parse the continuation token. + CosmosElement continuationToken = inputParameters.InitialUserContinuationToken; + PartitionedQueryExecutionInfo queryPlanFromContinuationToken = inputParameters.PartitionedQueryExecutionInfo; + if (continuationToken != null) + { + if (!PipelineContinuationToken.TryCreateFromCosmosElement( + continuationToken, + out PipelineContinuationToken pipelineContinuationToken)) + { + return TryCatch.FromException( + new MalformedContinuationTokenException( + $"Malformed {nameof(PipelineContinuationToken)}: {continuationToken}.")); + } + + if (PipelineContinuationToken.IsTokenFromTheFuture(pipelineContinuationToken)) + { + return TryCatch.FromException( + new MalformedContinuationTokenException( + $"{nameof(PipelineContinuationToken)} Continuation token is from a newer version of the SDK. " + + $"Upgrade the SDK to avoid this issue." + + $"{continuationToken}.")); + } + + if (!PipelineContinuationToken.TryConvertToLatest( + pipelineContinuationToken, + out PipelineContinuationTokenV1_1 latestVersionPipelineContinuationToken)) + { + return TryCatch.FromException( + new MalformedContinuationTokenException( + $"{nameof(PipelineContinuationToken)}: '{continuationToken}' is no longer supported.")); + } + + continuationToken = latestVersionPipelineContinuationToken.SourceContinuationToken; + if (latestVersionPipelineContinuationToken.QueryPlan != null) + { + queryPlanFromContinuationToken = latestVersionPipelineContinuationToken.QueryPlan; + } + } + + CosmosQueryClient cosmosQueryClient = cosmosQueryContext.QueryClient; + + ContainerQueryProperties containerQueryProperties = await cosmosQueryClient.GetCachedContainerQueryPropertiesAsync( + cosmosQueryContext.ResourceLink, + inputParameters.PartitionKey, + createQueryPipelineTrace, + cancellationToken); + cosmosQueryContext.ContainerResourceId = containerQueryProperties.ResourceId; + + Documents.PartitionKeyRange targetRange = await TryGetTargetRangeOptimisticDirectExecutionAsync( + inputParameters, + queryPlanFromContinuationToken, + cosmosQueryContext, + containerQueryProperties, + trace); + + if (targetRange != null) + { + return await TryCreateSinglePartitionExecutionContextAsync( + documentContainer, + partitionedQueryExecutionInfo: null, + cosmosQueryContext, + containerQueryProperties, + inputParameters, + targetRange, + createQueryPipelineTrace, + cancellationToken); + } + + PartitionedQueryExecutionInfo partitionedQueryExecutionInfo; + if (queryPlanFromContinuationToken != null) + { + partitionedQueryExecutionInfo = queryPlanFromContinuationToken; + } + else + { + // If the query would go to gateway, but we have a partition key, + // then try seeing if we can execute as a passthrough using client side only logic. + // This is to short circuit the need to go to the gateway to get the query plan. + if (cosmosQueryContext.QueryClient.BypassQueryParsing() + && inputParameters.PartitionKey.HasValue) + { + bool parsed; + SqlQuery sqlQuery; + using (ITrace queryParseTrace = createQueryPipelineTrace.StartChild("Parse Query", TraceComponent.Query, Tracing.TraceLevel.Info)) + { + parsed = SqlQueryParser.TryParse(inputParameters.SqlQuerySpec.QueryText, out sqlQuery); + } + + if (parsed) + { + bool hasDistinct = sqlQuery.SelectClause.HasDistinct; + bool hasGroupBy = sqlQuery.GroupByClause != default; + bool hasAggregates = AggregateProjectionDetector.HasAggregate(sqlQuery.SelectClause.SelectSpec); + bool createPassthroughQuery = !hasAggregates && !hasDistinct && !hasGroupBy; + + if (createPassthroughQuery) + { + SetTestInjectionPipelineType(inputParameters, Passthrough); + + // Only thing that matters is that we target the correct range. + Documents.PartitionKeyDefinition partitionKeyDefinition = GetPartitionKeyDefinition(inputParameters, containerQueryProperties); + List targetRanges = await cosmosQueryContext.QueryClient.GetTargetPartitionKeyRangesAsync( + cosmosQueryContext.ResourceLink, + containerQueryProperties.ResourceId, + containerQueryProperties.EffectiveRangesForPartitionKey, + forceRefresh: false, + createQueryPipelineTrace); + return TryCreatePassthroughQueryExecutionContext( - documentContainer, - inputParameters, - targetRanges); - } - } - } - - partitionedQueryExecutionInfo = await GetPartitionedQueryExecutionInfoAsync( - cosmosQueryContext, - inputParameters, - containerQueryProperties, - createQueryPipelineTrace, - cancellationToken); - } - - return await TryCreateFromPartitionedQueryExecutionInfoAsync( - documentContainer, - partitionedQueryExecutionInfo, - containerQueryProperties, - cosmosQueryContext, - inputParameters, - createQueryPipelineTrace, - cancellationToken); - } - } - - private static async Task> TryCreateFromPartitionedQueryExecutionInfoAsync( - DocumentContainer documentContainer, - PartitionedQueryExecutionInfo partitionedQueryExecutionInfo, - ContainerQueryProperties containerQueryProperties, - CosmosQueryContext cosmosQueryContext, - InputParameters inputParameters, - ITrace trace, - CancellationToken cancellationToken) - { - cancellationToken.ThrowIfCancellationRequested(); - + documentContainer, + inputParameters, + targetRanges, + containerQueryProperties); + } + } + } + + partitionedQueryExecutionInfo = await GetPartitionedQueryExecutionInfoAsync( + cosmosQueryContext, + inputParameters, + containerQueryProperties, + createQueryPipelineTrace, + cancellationToken); + } + + return await TryCreateFromPartitionedQueryExecutionInfoAsync( + documentContainer, + partitionedQueryExecutionInfo, + containerQueryProperties, + cosmosQueryContext, + inputParameters, + createQueryPipelineTrace, + cancellationToken); + } + } + + private static async Task> TryCreateFromPartitionedQueryExecutionInfoAsync( + DocumentContainer documentContainer, + PartitionedQueryExecutionInfo partitionedQueryExecutionInfo, + ContainerQueryProperties containerQueryProperties, + CosmosQueryContext cosmosQueryContext, + InputParameters inputParameters, + ITrace trace, + CancellationToken cancellationToken) + { + cancellationToken.ThrowIfCancellationRequested(); + List targetRanges = await GetTargetPartitionKeyRangesAsync( - cosmosQueryContext.QueryClient, - cosmosQueryContext.ResourceLink, - partitionedQueryExecutionInfo, - containerQueryProperties, - inputParameters.Properties, - inputParameters.InitialFeedRange, - trace); - - TryCatch tryCreatePipelineStage; - Documents.PartitionKeyRange targetRange = await TryGetTargetRangeOptimisticDirectExecutionAsync( - inputParameters, - partitionedQueryExecutionInfo, - cosmosQueryContext, - containerQueryProperties, - trace); - - if (targetRange != null) - { - tryCreatePipelineStage = await TryCreateSinglePartitionExecutionContextAsync( - documentContainer, - partitionedQueryExecutionInfo, - cosmosQueryContext, - containerQueryProperties, - inputParameters, - targetRange, - trace, - cancellationToken); - } - else - { - bool singleLogicalPartitionKeyQuery = inputParameters.PartitionKey.HasValue - || ((partitionedQueryExecutionInfo.QueryRanges.Count == 1) - && partitionedQueryExecutionInfo.QueryRanges[0].IsSingleValue); - bool serverStreamingQuery = !partitionedQueryExecutionInfo.QueryInfo.HasAggregates - && !partitionedQueryExecutionInfo.QueryInfo.HasDistinct - && !partitionedQueryExecutionInfo.QueryInfo.HasGroupBy; - bool streamingSinglePartitionQuery = singleLogicalPartitionKeyQuery && serverStreamingQuery; - - bool clientStreamingQuery = serverStreamingQuery - && !partitionedQueryExecutionInfo.QueryInfo.HasOrderBy - && !partitionedQueryExecutionInfo.QueryInfo.HasTop - && !partitionedQueryExecutionInfo.QueryInfo.HasLimit - && !partitionedQueryExecutionInfo.QueryInfo.HasOffset; - bool streamingCrossContinuationQuery = !singleLogicalPartitionKeyQuery && clientStreamingQuery; - bool createPassthroughQuery = streamingSinglePartitionQuery || streamingCrossContinuationQuery; - - if (createPassthroughQuery) - { - SetTestInjectionPipelineType(inputParameters, Passthrough); - + cosmosQueryContext.QueryClient, + cosmosQueryContext.ResourceLink, + partitionedQueryExecutionInfo, + containerQueryProperties, + inputParameters.Properties, + inputParameters.InitialFeedRange, + trace); + + TryCatch tryCreatePipelineStage; + Documents.PartitionKeyRange targetRange = await TryGetTargetRangeOptimisticDirectExecutionAsync( + inputParameters, + partitionedQueryExecutionInfo, + cosmosQueryContext, + containerQueryProperties, + trace); + + if (targetRange != null) + { + tryCreatePipelineStage = await TryCreateSinglePartitionExecutionContextAsync( + documentContainer, + partitionedQueryExecutionInfo, + cosmosQueryContext, + containerQueryProperties, + inputParameters, + targetRange, + trace, + cancellationToken); + } + else + { + bool singleLogicalPartitionKeyQuery = inputParameters.PartitionKey.HasValue + || ((partitionedQueryExecutionInfo.QueryRanges.Count == 1) + && partitionedQueryExecutionInfo.QueryRanges[0].IsSingleValue); + bool serverStreamingQuery = !partitionedQueryExecutionInfo.QueryInfo.HasAggregates + && !partitionedQueryExecutionInfo.QueryInfo.HasDistinct + && !partitionedQueryExecutionInfo.QueryInfo.HasGroupBy; + bool streamingSinglePartitionQuery = singleLogicalPartitionKeyQuery && serverStreamingQuery; + + bool clientStreamingQuery = serverStreamingQuery + && !partitionedQueryExecutionInfo.QueryInfo.HasOrderBy + && !partitionedQueryExecutionInfo.QueryInfo.HasTop + && !partitionedQueryExecutionInfo.QueryInfo.HasLimit + && !partitionedQueryExecutionInfo.QueryInfo.HasOffset; + bool streamingCrossContinuationQuery = !singleLogicalPartitionKeyQuery && clientStreamingQuery; + bool createPassthroughQuery = streamingSinglePartitionQuery || streamingCrossContinuationQuery; + + if (createPassthroughQuery) + { + SetTestInjectionPipelineType(inputParameters, Passthrough); + tryCreatePipelineStage = TryCreatePassthroughQueryExecutionContext( - documentContainer, - inputParameters, - targetRanges); - } - else - { + documentContainer, + inputParameters, + targetRanges, + containerQueryProperties); + } + else + { + tryCreatePipelineStage = TryCreateSpecializedDocumentQueryExecutionContext(documentContainer, cosmosQueryContext, inputParameters, targetRanges, containerQueryProperties, partitionedQueryExecutionInfo); + } + } + + return tryCreatePipelineStage; + } + + private static async Task> TryCreateSinglePartitionExecutionContextAsync( + DocumentContainer documentContainer, + PartitionedQueryExecutionInfo partitionedQueryExecutionInfo, + CosmosQueryContext cosmosQueryContext, + ContainerQueryProperties containerQueryProperties, + InputParameters inputParameters, + Documents.PartitionKeyRange targetRange, + ITrace trace, + CancellationToken cancellationToken) + { + // Retrieve the query plan in a subset of cases to ensure the query is valid before creating the Ode pipeline + if (partitionedQueryExecutionInfo == null && QueryInspectionRegex.IsMatch(inputParameters.SqlQuerySpec.QueryText)) + { + partitionedQueryExecutionInfo = await GetPartitionedQueryExecutionInfoAsync( + cosmosQueryContext, + inputParameters, + containerQueryProperties, + trace, + cancellationToken); + } + + // Test code added to confirm the correct pipeline is being utilized + SetTestInjectionPipelineType(inputParameters, OptimisticDirectExecution); + + TryCatch tryCreatePipelineStage = TryCreateOptimisticDirectExecutionContext( + documentContainer, + cosmosQueryContext, + containerQueryProperties, + inputParameters, + targetRange, + cancellationToken); + + // A malformed continuation token exception would happen for 2 reasons here + // 1. the token is actually malformed + // 2. Its a non Ode continuation token + // In both cases, Ode pipeline delegates the work to the Specialized pipeline + // as Ode pipeline should not take over execution while some other pipeline is already handling it + if (tryCreatePipelineStage.Failed && tryCreatePipelineStage.InnerMostException is MalformedContinuationTokenException) + { + SetTestInjectionPipelineType(inputParameters, Specialized); + + if (partitionedQueryExecutionInfo != null) + { + List targetRanges = new List + { + targetRange + }; + tryCreatePipelineStage = TryCreateSpecializedDocumentQueryExecutionContext( documentContainer, cosmosQueryContext, inputParameters, targetRanges, - partitionedQueryExecutionInfo); - } - } - - return tryCreatePipelineStage; - } - - private static async Task> TryCreateSinglePartitionExecutionContextAsync( - DocumentContainer documentContainer, - PartitionedQueryExecutionInfo partitionedQueryExecutionInfo, - CosmosQueryContext cosmosQueryContext, - ContainerQueryProperties containerQueryProperties, - InputParameters inputParameters, - Documents.PartitionKeyRange targetRange, - ITrace trace, - CancellationToken cancellationToken) - { - // Retrieve the query plan in a subset of cases to ensure the query is valid before creating the Ode pipeline - if (partitionedQueryExecutionInfo == null && QueryInspectionRegex.IsMatch(inputParameters.SqlQuerySpec.QueryText)) - { - partitionedQueryExecutionInfo = await GetPartitionedQueryExecutionInfoAsync( - cosmosQueryContext, - inputParameters, - containerQueryProperties, - trace, - cancellationToken); - } - - // Test code added to confirm the correct pipeline is being utilized - SetTestInjectionPipelineType(inputParameters, OptimisticDirectExecution); - - TryCatch tryCreatePipelineStage = TryCreateOptimisticDirectExecutionContext( - documentContainer, - cosmosQueryContext, - containerQueryProperties, - inputParameters, - targetRange, - cancellationToken); - - // A malformed continuation token exception would happen for 2 reasons here - // 1. the token is actually malformed - // 2. Its a non Ode continuation token - // In both cases, Ode pipeline delegates the work to the Specialized pipeline - // as Ode pipeline should not take over execution while some other pipeline is already handling it - if (tryCreatePipelineStage.Failed && tryCreatePipelineStage.InnerMostException is MalformedContinuationTokenException) - { - SetTestInjectionPipelineType(inputParameters, Specialized); - - if (partitionedQueryExecutionInfo != null) - { - List targetRanges = new List - { - targetRange - }; - - tryCreatePipelineStage = TryCreateSpecializedDocumentQueryExecutionContext( - documentContainer, - cosmosQueryContext, - inputParameters, - targetRanges, - partitionedQueryExecutionInfo); - } - else - { - tryCreatePipelineStage = await TryCreateSpecializedDocumentQueryExecutionContextAsync( - documentContainer, - cosmosQueryContext, - containerQueryProperties, - inputParameters, - trace, - cancellationToken); - } - } - - return tryCreatePipelineStage; - } - - private static TryCatch TryCreateSpecializedDocumentQueryExecutionContext( - DocumentContainer documentContainer, - CosmosQueryContext cosmosQueryContext, - InputParameters inputParameters, - List targetRanges, - PartitionedQueryExecutionInfo partitionedQueryExecutionInfo) - { - SetTestInjectionPipelineType(inputParameters, Specialized); - - if (!string.IsNullOrEmpty(partitionedQueryExecutionInfo.QueryInfo.RewrittenQuery)) - { - // We need pass down the rewritten query. - SqlQuerySpec rewrittenQuerySpec = new SqlQuerySpec() - { - QueryText = partitionedQueryExecutionInfo.QueryInfo.RewrittenQuery, - Parameters = inputParameters.SqlQuerySpec.Parameters - }; - - inputParameters = new InputParameters( - rewrittenQuerySpec, - inputParameters.InitialUserContinuationToken, - inputParameters.InitialFeedRange, - inputParameters.MaxConcurrency, - inputParameters.MaxItemCount, - inputParameters.MaxBufferedItemCount, - inputParameters.PartitionKey, - inputParameters.Properties, - inputParameters.PartitionedQueryExecutionInfo, - inputParameters.ExecutionEnvironment, - inputParameters.ReturnResultsInDeterministicOrder, - inputParameters.EnableOptimisticDirectExecution, - inputParameters.TestInjections); - } - + containerQueryProperties, + partitionedQueryExecutionInfo); + } + else + { + tryCreatePipelineStage = await TryCreateSpecializedDocumentQueryExecutionContextAsync( + documentContainer, + cosmosQueryContext, + containerQueryProperties, + inputParameters, + trace, + cancellationToken); + } + } + + return tryCreatePipelineStage; + } + + private static TryCatch TryCreateSpecializedDocumentQueryExecutionContext( + DocumentContainer documentContainer, + CosmosQueryContext cosmosQueryContext, + InputParameters inputParameters, + List targetRanges, + ContainerQueryProperties containerQueryProperties, + PartitionedQueryExecutionInfo partitionedQueryExecutionInfo) + { + SetTestInjectionPipelineType(inputParameters, Specialized); + + if (!string.IsNullOrEmpty(partitionedQueryExecutionInfo.QueryInfo.RewrittenQuery)) + { + // We need pass down the rewritten query. + SqlQuerySpec rewrittenQuerySpec = new SqlQuerySpec() + { + QueryText = partitionedQueryExecutionInfo.QueryInfo.RewrittenQuery, + Parameters = inputParameters.SqlQuerySpec.Parameters + }; + + inputParameters = new InputParameters( + rewrittenQuerySpec, + inputParameters.InitialUserContinuationToken, + inputParameters.InitialFeedRange, + inputParameters.MaxConcurrency, + inputParameters.MaxItemCount, + inputParameters.MaxBufferedItemCount, + inputParameters.PartitionKey, + inputParameters.Properties, + inputParameters.PartitionedQueryExecutionInfo, + inputParameters.ExecutionEnvironment, + inputParameters.ReturnResultsInDeterministicOrder, + inputParameters.EnableOptimisticDirectExecution, + inputParameters.TestInjections); + } + return TryCreateSpecializedDocumentQueryExecutionContext( - documentContainer, - cosmosQueryContext, - inputParameters, - partitionedQueryExecutionInfo, - targetRanges); - } - - private static async Task> TryCreateSpecializedDocumentQueryExecutionContextAsync( - DocumentContainer documentContainer, - CosmosQueryContext cosmosQueryContext, - ContainerQueryProperties containerQueryProperties, - InputParameters inputParameters, - ITrace trace, - CancellationToken cancellationToken) - { - PartitionedQueryExecutionInfo partitionedQueryExecutionInfo = await GetPartitionedQueryExecutionInfoAsync( - cosmosQueryContext, - inputParameters, - containerQueryProperties, - trace, - cancellationToken); - + documentContainer, + cosmosQueryContext, + inputParameters, + partitionedQueryExecutionInfo, + targetRanges, + containerQueryProperties); + } + + private static async Task> TryCreateSpecializedDocumentQueryExecutionContextAsync( + DocumentContainer documentContainer, + CosmosQueryContext cosmosQueryContext, + ContainerQueryProperties containerQueryProperties, + InputParameters inputParameters, + ITrace trace, + CancellationToken cancellationToken) + { + PartitionedQueryExecutionInfo partitionedQueryExecutionInfo = await GetPartitionedQueryExecutionInfoAsync( + cosmosQueryContext, + inputParameters, + containerQueryProperties, + trace, + cancellationToken); + List targetRanges = await GetTargetPartitionKeyRangesAsync( - cosmosQueryContext.QueryClient, - cosmosQueryContext.ResourceLink, - partitionedQueryExecutionInfo, - containerQueryProperties, - inputParameters.Properties, - inputParameters.InitialFeedRange, - trace); - - return TryCreateSpecializedDocumentQueryExecutionContext( - documentContainer, - cosmosQueryContext, - inputParameters, - targetRanges, - partitionedQueryExecutionInfo); - } - - private static TryCatch TryCreateOptimisticDirectExecutionContext( - DocumentContainer documentContainer, - CosmosQueryContext cosmosQueryContext, - ContainerQueryProperties containerQueryProperties, - InputParameters inputParameters, - Documents.PartitionKeyRange targetRange, - CancellationToken cancellationToken) - { - // Return a OptimisticDirectExecution context - return OptimisticDirectExecutionQueryPipelineStage.MonadicCreate( - documentContainer: documentContainer, - inputParameters: inputParameters, - targetRange: new FeedRangeEpk(targetRange.ToRange()), + cosmosQueryContext.QueryClient, + cosmosQueryContext.ResourceLink, + partitionedQueryExecutionInfo, + containerQueryProperties, + inputParameters.Properties, + inputParameters.InitialFeedRange, + trace); + + return TryCreateSpecializedDocumentQueryExecutionContext( + documentContainer, + cosmosQueryContext, + inputParameters, + targetRanges, + containerQueryProperties, + partitionedQueryExecutionInfo); + } + + private static TryCatch TryCreateOptimisticDirectExecutionContext( + DocumentContainer documentContainer, + CosmosQueryContext cosmosQueryContext, + ContainerQueryProperties containerQueryProperties, + InputParameters inputParameters, + Documents.PartitionKeyRange targetRange, + CancellationToken cancellationToken) + { + // Return a OptimisticDirectExecution context + return OptimisticDirectExecutionQueryPipelineStage.MonadicCreate( + documentContainer: documentContainer, + inputParameters: inputParameters, + targetRange: new FeedRangeEpk(targetRange.ToRange()), + containerQueryProperties: containerQueryProperties, fallbackQueryPipelineStageFactory: (continuationToken) => // In fallback scenario, the Specialized pipeline is always invoked TryCreateSpecializedDocumentQueryExecutionContextAsync( @@ -474,628 +476,632 @@ private static TryCatch TryCreateOptimisticDirectExecutionC inputParameters.WithContinuationToken(continuationToken), NoOpTrace.Singleton, cancellationToken), - cancellationToken: cancellationToken); - } - - private static TryCatch TryCreatePassthroughQueryExecutionContext( - DocumentContainer documentContainer, - InputParameters inputParameters, - List targetRanges) - { - // Return a parallel context, since we still want to be able to handle splits and concurrency / buffering. - return ParallelCrossPartitionQueryPipelineStage.MonadicCreate( - documentContainer: documentContainer, - sqlQuerySpec: inputParameters.SqlQuerySpec, - targetRanges: targetRanges - .Select(range => new FeedRangeEpk( - new Documents.Routing.Range( - min: range.MinInclusive, - max: range.MaxExclusive, - isMinInclusive: true, - isMaxInclusive: false))) - .ToList(), - queryPaginationOptions: new QueryPaginationOptions( - pageSizeHint: inputParameters.MaxItemCount), - partitionKey: inputParameters.PartitionKey, - prefetchPolicy: PrefetchPolicy.PrefetchSinglePage, - maxConcurrency: inputParameters.MaxConcurrency, + cancellationToken: cancellationToken); + } + + private static TryCatch TryCreatePassthroughQueryExecutionContext( + DocumentContainer documentContainer, + InputParameters inputParameters, + List targetRanges, + ContainerQueryProperties containerQueryProperties) + { + // Return a parallel context, since we still want to be able to handle splits and concurrency / buffering. + return ParallelCrossPartitionQueryPipelineStage.MonadicCreate( + documentContainer: documentContainer, + sqlQuerySpec: inputParameters.SqlQuerySpec, + targetRanges: targetRanges + .Select(range => new FeedRangeEpk( + new Documents.Routing.Range( + min: range.MinInclusive, + max: range.MaxExclusive, + isMinInclusive: true, + isMaxInclusive: false))) + .ToList(), + queryPaginationOptions: new QueryPaginationOptions( + pageSizeHint: inputParameters.MaxItemCount), + partitionKey: inputParameters.PartitionKey, + containerQueryProperties: containerQueryProperties, + maxConcurrency: inputParameters.MaxConcurrency, + prefetchPolicy: PrefetchPolicy.PrefetchSinglePage, continuationToken: inputParameters.InitialUserContinuationToken); - } - - private static TryCatch TryCreateSpecializedDocumentQueryExecutionContext( - DocumentContainer documentContainer, - CosmosQueryContext cosmosQueryContext, - InputParameters inputParameters, - PartitionedQueryExecutionInfo partitionedQueryExecutionInfo, - List targetRanges) - { - QueryInfo queryInfo = partitionedQueryExecutionInfo.QueryInfo; - - // We need to compute the optimal initial page size for order-by queries - long optimalPageSize = inputParameters.MaxItemCount; - if (queryInfo.HasOrderBy) - { - int top; - if (queryInfo.HasTop && (partitionedQueryExecutionInfo.QueryInfo.Top.Value > 0)) - { - top = partitionedQueryExecutionInfo.QueryInfo.Top.Value; - } - else if (queryInfo.HasLimit && (partitionedQueryExecutionInfo.QueryInfo.Limit.Value > 0)) - { - top = (partitionedQueryExecutionInfo.QueryInfo.Offset ?? 0) + partitionedQueryExecutionInfo.QueryInfo.Limit.Value; - } - else - { - top = 0; - } - - if (top > 0) - { - // All partitions should initially fetch about 1/nth of the top value. - long pageSizeWithTop = (long)Math.Min( - Math.Ceiling(top / (double)targetRanges.Count) * CosmosQueryExecutionContextFactory.PageSizeFactorForTop, - top); - - optimalPageSize = Math.Min(pageSizeWithTop, optimalPageSize); - } - else if (cosmosQueryContext.IsContinuationExpected) - { - optimalPageSize = (long)Math.Min( - Math.Ceiling(optimalPageSize / (double)targetRanges.Count) * CosmosQueryExecutionContextFactory.PageSizeFactorForTop, - optimalPageSize); - } - } - - Debug.Assert( - (optimalPageSize > 0) && (optimalPageSize <= int.MaxValue), - $"Invalid MaxItemCount {optimalPageSize}"); - - return PipelineFactory.MonadicCreate( - executionEnvironment: inputParameters.ExecutionEnvironment, - documentContainer: documentContainer, - sqlQuerySpec: inputParameters.SqlQuerySpec, - targetRanges: targetRanges - .Select(range => new FeedRangeEpk( - new Documents.Routing.Range( - min: range.MinInclusive, - max: range.MaxExclusive, - isMinInclusive: true, - isMaxInclusive: false))) - .ToList(), - partitionKey: inputParameters.PartitionKey, - queryInfo: partitionedQueryExecutionInfo.QueryInfo, - queryPaginationOptions: new QueryPaginationOptions( - pageSizeHint: (int)optimalPageSize), - maxConcurrency: inputParameters.MaxConcurrency, + } + + private static TryCatch TryCreateSpecializedDocumentQueryExecutionContext( + DocumentContainer documentContainer, + CosmosQueryContext cosmosQueryContext, + InputParameters inputParameters, + PartitionedQueryExecutionInfo partitionedQueryExecutionInfo, + List targetRanges, + ContainerQueryProperties containerQueryProperties) + { + QueryInfo queryInfo = partitionedQueryExecutionInfo.QueryInfo; + + // We need to compute the optimal initial page size for order-by queries + long optimalPageSize = inputParameters.MaxItemCount; + if (queryInfo.HasOrderBy) + { + int top; + if (queryInfo.HasTop && (partitionedQueryExecutionInfo.QueryInfo.Top.Value > 0)) + { + top = partitionedQueryExecutionInfo.QueryInfo.Top.Value; + } + else if (queryInfo.HasLimit && (partitionedQueryExecutionInfo.QueryInfo.Limit.Value > 0)) + { + top = (partitionedQueryExecutionInfo.QueryInfo.Offset ?? 0) + partitionedQueryExecutionInfo.QueryInfo.Limit.Value; + } + else + { + top = 0; + } + + if (top > 0) + { + // All partitions should initially fetch about 1/nth of the top value. + long pageSizeWithTop = (long)Math.Min( + Math.Ceiling(top / (double)targetRanges.Count) * CosmosQueryExecutionContextFactory.PageSizeFactorForTop, + top); + + optimalPageSize = Math.Min(pageSizeWithTop, optimalPageSize); + } + else if (cosmosQueryContext.IsContinuationExpected) + { + optimalPageSize = (long)Math.Min( + Math.Ceiling(optimalPageSize / (double)targetRanges.Count) * CosmosQueryExecutionContextFactory.PageSizeFactorForTop, + optimalPageSize); + } + } + + Debug.Assert( + (optimalPageSize > 0) && (optimalPageSize <= int.MaxValue), + $"Invalid MaxItemCount {optimalPageSize}"); + + return PipelineFactory.MonadicCreate( + executionEnvironment: inputParameters.ExecutionEnvironment, + documentContainer: documentContainer, + sqlQuerySpec: inputParameters.SqlQuerySpec, + targetRanges: targetRanges + .Select(range => new FeedRangeEpk( + new Documents.Routing.Range( + min: range.MinInclusive, + max: range.MaxExclusive, + isMinInclusive: true, + isMaxInclusive: false))) + .ToList(), + partitionKey: inputParameters.PartitionKey, + queryInfo: partitionedQueryExecutionInfo.QueryInfo, + queryPaginationOptions: new QueryPaginationOptions( + pageSizeHint: (int)optimalPageSize), + containerQueryProperties: containerQueryProperties, + maxConcurrency: inputParameters.MaxConcurrency, requestContinuationToken: inputParameters.InitialUserContinuationToken); - } - - private static async Task GetPartitionedQueryExecutionInfoAsync( - CosmosQueryContext cosmosQueryContext, - InputParameters inputParameters, - ContainerQueryProperties containerQueryProperties, - ITrace trace, - CancellationToken cancellationToken) - { - PartitionedQueryExecutionInfo partitionedQueryExecutionInfo; - if (cosmosQueryContext.QueryClient.BypassQueryParsing()) - { - // For non-Windows platforms(like Linux and OSX) in .NET Core SDK, we cannot use ServiceInterop, so need to bypass in that case. - // We are also now bypassing this for 32 bit host process running even on Windows as there are many 32 bit apps that will not work without this - partitionedQueryExecutionInfo = await QueryPlanRetriever.GetQueryPlanThroughGatewayAsync( - cosmosQueryContext, - inputParameters.SqlQuerySpec, - cosmosQueryContext.ResourceLink, - inputParameters.PartitionKey, - trace, - cancellationToken); - } - else - { - Documents.PartitionKeyDefinition partitionKeyDefinition = GetPartitionKeyDefinition(inputParameters, containerQueryProperties); - - partitionedQueryExecutionInfo = await QueryPlanRetriever.GetQueryPlanWithServiceInteropAsync( - cosmosQueryContext.QueryClient, - inputParameters.SqlQuerySpec, - cosmosQueryContext.ResourceTypeEnum, - partitionKeyDefinition, - inputParameters.PartitionKey != null, - containerQueryProperties.GeospatialType, - cosmosQueryContext.UseSystemPrefix, - trace, - cancellationToken); - } - - return partitionedQueryExecutionInfo; - } - - /// - /// Gets the list of partition key ranges. - /// 1. Check partition key range id - /// 2. Check Partition key - /// 3. Check the effective partition key - /// 4. Get the range from the FeedToken - /// 5. Get the range from the PartitionedQueryExecutionInfo - /// - internal static async Task> GetTargetPartitionKeyRangesAsync( - CosmosQueryClient queryClient, - string resourceLink, - PartitionedQueryExecutionInfo partitionedQueryExecutionInfo, - ContainerQueryProperties containerQueryProperties, - IReadOnlyDictionary properties, - FeedRangeInternal feedRangeInternal, - ITrace trace) - { - List targetRanges; - if (containerQueryProperties.EffectiveRangesForPartitionKey != null) - { - targetRanges = await queryClient.GetTargetPartitionKeyRangesAsync( - resourceLink, - containerQueryProperties.ResourceId, - containerQueryProperties.EffectiveRangesForPartitionKey, - forceRefresh: false, - trace); - } - else if (TryGetEpkProperty(properties, out string effectivePartitionKeyString)) - { - //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> effectiveRanges = new List> - { Range.GetPointRange(effectivePartitionKeyString) }; - - targetRanges = await queryClient.GetTargetPartitionKeyRangesAsync( - resourceLink, - containerQueryProperties.ResourceId, - effectiveRanges, - forceRefresh: false, - trace); - } - else if (feedRangeInternal != null) - { - targetRanges = await queryClient.GetTargetPartitionKeyRangeByFeedRangeAsync( - resourceLink, - containerQueryProperties.ResourceId, - containerQueryProperties.PartitionKeyDefinition, - feedRangeInternal, - forceRefresh: false, - trace); - } - else - { - targetRanges = await queryClient.GetTargetPartitionKeyRangesAsync( - resourceLink, - containerQueryProperties.ResourceId, - partitionedQueryExecutionInfo.QueryRanges, - forceRefresh: false, - trace); - } - - return targetRanges; - } - - private static bool TryGetEpkProperty( - IReadOnlyDictionary properties, - out string effectivePartitionKeyString) - { - if (properties != null - && properties.TryGetValue( - Documents.WFConstants.BackendHeaders.EffectivePartitionKeyString, - out object effectivePartitionKeyStringObject)) - { - effectivePartitionKeyString = effectivePartitionKeyStringObject as string; - if (string.IsNullOrEmpty(effectivePartitionKeyString)) - { - throw new ArgumentOutOfRangeException(nameof(effectivePartitionKeyString)); - } - - return true; - } - - effectivePartitionKeyString = null; - 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 - //if collection is deleted/created with same name. - //need to make it not rely on information from collection cache. - - Documents.PartitionKeyDefinition partitionKeyDefinition; - if ((inputParameters.Properties != null) - && inputParameters.Properties.TryGetValue(InternalPartitionKeyDefinitionProperty, out object partitionKeyDefinitionObject)) - { - if (!(partitionKeyDefinitionObject is Documents.PartitionKeyDefinition definition)) - { - throw new ArgumentException( - "partitionkeydefinition has invalid type", - nameof(partitionKeyDefinitionObject)); - } - - partitionKeyDefinition = definition; - } - else - { - partitionKeyDefinition = containerQueryProperties.PartitionKeyDefinition; - } - - return partitionKeyDefinition; - } - - private static async Task TryGetTargetRangeOptimisticDirectExecutionAsync( - InputParameters inputParameters, - PartitionedQueryExecutionInfo partitionedQueryExecutionInfo, - CosmosQueryContext cosmosQueryContext, - ContainerQueryProperties containerQueryProperties, - ITrace trace) - { - bool clientDisableOptimisticDirectExecution = await cosmosQueryContext.QueryClient.GetClientDisableOptimisticDirectExecutionAsync(); - - // Use the Ode code path only if ClientDisableOptimisticDirectExecution is false and EnableOptimisticDirectExecution is true - if (clientDisableOptimisticDirectExecution || !inputParameters.EnableOptimisticDirectExecution) - { - if (inputParameters.InitialUserContinuationToken != null - && OptimisticDirectExecutionContinuationToken.IsOptimisticDirectExecutionContinuationToken(inputParameters.InitialUserContinuationToken)) - { - string errorMessage = "Execution of this query using the supplied continuation token requires EnableOptimisticDirectExecution to be set in QueryRequestOptions. " + - "If the error persists after that, contact system administrator."; - - throw new MalformedContinuationTokenException($"{errorMessage} Continuation Token: {inputParameters.InitialUserContinuationToken}"); - } - - return null; - } - - Debug.Assert(containerQueryProperties.ResourceId != null, "CosmosQueryExecutionContextFactory Assert!", "Container ResourceId cannot be null!"); - - List targetRanges; - if (partitionedQueryExecutionInfo != null || inputParameters.InitialFeedRange != null) - { + } + + private static async Task GetPartitionedQueryExecutionInfoAsync( + CosmosQueryContext cosmosQueryContext, + InputParameters inputParameters, + ContainerQueryProperties containerQueryProperties, + ITrace trace, + CancellationToken cancellationToken) + { + PartitionedQueryExecutionInfo partitionedQueryExecutionInfo; + if (cosmosQueryContext.QueryClient.BypassQueryParsing()) + { + // For non-Windows platforms(like Linux and OSX) in .NET Core SDK, we cannot use ServiceInterop, so need to bypass in that case. + // We are also now bypassing this for 32 bit host process running even on Windows as there are many 32 bit apps that will not work without this + partitionedQueryExecutionInfo = await QueryPlanRetriever.GetQueryPlanThroughGatewayAsync( + cosmosQueryContext, + inputParameters.SqlQuerySpec, + cosmosQueryContext.ResourceLink, + inputParameters.PartitionKey, + trace, + cancellationToken); + } + else + { + Documents.PartitionKeyDefinition partitionKeyDefinition = GetPartitionKeyDefinition(inputParameters, containerQueryProperties); + + partitionedQueryExecutionInfo = await QueryPlanRetriever.GetQueryPlanWithServiceInteropAsync( + cosmosQueryContext.QueryClient, + inputParameters.SqlQuerySpec, + cosmosQueryContext.ResourceTypeEnum, + partitionKeyDefinition, + inputParameters.PartitionKey != null, + containerQueryProperties.GeospatialType, + cosmosQueryContext.UseSystemPrefix, + trace, + cancellationToken); + } + + return partitionedQueryExecutionInfo; + } + + /// + /// Gets the list of partition key ranges. + /// 1. Check partition key range id + /// 2. Check Partition key + /// 3. Check the effective partition key + /// 4. Get the range from the FeedToken + /// 5. Get the range from the PartitionedQueryExecutionInfo + /// + internal static async Task> GetTargetPartitionKeyRangesAsync( + CosmosQueryClient queryClient, + string resourceLink, + PartitionedQueryExecutionInfo partitionedQueryExecutionInfo, + ContainerQueryProperties containerQueryProperties, + IReadOnlyDictionary properties, + FeedRangeInternal feedRangeInternal, + ITrace trace) + { + List targetRanges; + if (containerQueryProperties.EffectiveRangesForPartitionKey != null) + { + targetRanges = await queryClient.GetTargetPartitionKeyRangesAsync( + resourceLink, + containerQueryProperties.ResourceId, + containerQueryProperties.EffectiveRangesForPartitionKey, + forceRefresh: false, + trace); + } + else if (TryGetEpkProperty(properties, out string effectivePartitionKeyString)) + { + //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> effectiveRanges = new List> + { Range.GetPointRange(effectivePartitionKeyString) }; + + targetRanges = await queryClient.GetTargetPartitionKeyRangesAsync( + resourceLink, + containerQueryProperties.ResourceId, + effectiveRanges, + forceRefresh: false, + trace); + } + else if (feedRangeInternal != null) + { + targetRanges = await queryClient.GetTargetPartitionKeyRangeByFeedRangeAsync( + resourceLink, + containerQueryProperties.ResourceId, + containerQueryProperties.PartitionKeyDefinition, + feedRangeInternal, + forceRefresh: false, + trace); + } + else + { + targetRanges = await queryClient.GetTargetPartitionKeyRangesAsync( + resourceLink, + containerQueryProperties.ResourceId, + partitionedQueryExecutionInfo.QueryRanges, + forceRefresh: false, + trace); + } + + return targetRanges; + } + + private static bool TryGetEpkProperty( + IReadOnlyDictionary properties, + out string effectivePartitionKeyString) + { + if (properties != null + && properties.TryGetValue( + Documents.WFConstants.BackendHeaders.EffectivePartitionKeyString, + out object effectivePartitionKeyStringObject)) + { + effectivePartitionKeyString = effectivePartitionKeyStringObject as string; + if (string.IsNullOrEmpty(effectivePartitionKeyString)) + { + throw new ArgumentOutOfRangeException(nameof(effectivePartitionKeyString)); + } + + return true; + } + + effectivePartitionKeyString = null; + 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 + //if collection is deleted/created with same name. + //need to make it not rely on information from collection cache. + + Documents.PartitionKeyDefinition partitionKeyDefinition; + if ((inputParameters.Properties != null) + && inputParameters.Properties.TryGetValue(InternalPartitionKeyDefinitionProperty, out object partitionKeyDefinitionObject)) + { + if (!(partitionKeyDefinitionObject is Documents.PartitionKeyDefinition definition)) + { + throw new ArgumentException( + "partitionkeydefinition has invalid type", + nameof(partitionKeyDefinitionObject)); + } + + partitionKeyDefinition = definition; + } + else + { + partitionKeyDefinition = containerQueryProperties.PartitionKeyDefinition; + } + + return partitionKeyDefinition; + } + + private static async Task TryGetTargetRangeOptimisticDirectExecutionAsync( + InputParameters inputParameters, + PartitionedQueryExecutionInfo partitionedQueryExecutionInfo, + CosmosQueryContext cosmosQueryContext, + ContainerQueryProperties containerQueryProperties, + ITrace trace) + { + bool clientDisableOptimisticDirectExecution = await cosmosQueryContext.QueryClient.GetClientDisableOptimisticDirectExecutionAsync(); + + // Use the Ode code path only if ClientDisableOptimisticDirectExecution is false and EnableOptimisticDirectExecution is true + if (clientDisableOptimisticDirectExecution || !inputParameters.EnableOptimisticDirectExecution) + { + if (inputParameters.InitialUserContinuationToken != null + && OptimisticDirectExecutionContinuationToken.IsOptimisticDirectExecutionContinuationToken(inputParameters.InitialUserContinuationToken)) + { + string errorMessage = "Execution of this query using the supplied continuation token requires EnableOptimisticDirectExecution to be set in QueryRequestOptions. " + + "If the error persists after that, contact system administrator."; + + throw new MalformedContinuationTokenException($"{errorMessage} Continuation Token: {inputParameters.InitialUserContinuationToken}"); + } + + return null; + } + + Debug.Assert(containerQueryProperties.ResourceId != null, "CosmosQueryExecutionContextFactory Assert!", "Container ResourceId cannot be null!"); + + List targetRanges; + if (partitionedQueryExecutionInfo != null || inputParameters.InitialFeedRange != null) + { targetRanges = await GetTargetPartitionKeyRangesAsync( - cosmosQueryContext.QueryClient, - cosmosQueryContext.ResourceLink, - partitionedQueryExecutionInfo, - containerQueryProperties, - inputParameters.Properties, - inputParameters.InitialFeedRange, - trace); - } - else - { - Documents.PartitionKeyDefinition partitionKeyDefinition = GetPartitionKeyDefinition(inputParameters, containerQueryProperties); - if (inputParameters.PartitionKey.HasValue) - { - Debug.Assert(partitionKeyDefinition != null, "CosmosQueryExecutionContextFactory Assert!", "PartitionKeyDefinition cannot be null if partitionKey is defined"); - targetRanges = await cosmosQueryContext.QueryClient.GetTargetPartitionKeyRangesAsync( - cosmosQueryContext.ResourceLink, - containerQueryProperties.ResourceId, - containerQueryProperties.EffectiveRangesForPartitionKey, - forceRefresh: false, - trace); - } - else - { - targetRanges = await cosmosQueryContext.QueryClient.GetTargetPartitionKeyRangesAsync( - cosmosQueryContext.ResourceLink, - containerQueryProperties.ResourceId, - new List> { FeedRangeEpk.FullRange.Range }, - forceRefresh: false, - trace); - } - } - - if (targetRanges.Count == 1) - { - return targetRanges.Single(); - } - - return null; - } - - public sealed class InputParameters - { - private const int DefaultMaxConcurrency = 0; - private const int DefaultMaxItemCount = 1000; - private const int DefaultMaxBufferedItemCount = 1000; - private const bool DefaultReturnResultsInDeterministicOrder = true; - private const ExecutionEnvironment DefaultExecutionEnvironment = ExecutionEnvironment.Client; - - public InputParameters( - SqlQuerySpec sqlQuerySpec, - CosmosElement initialUserContinuationToken, - FeedRangeInternal initialFeedRange, - int? maxConcurrency, - int? maxItemCount, - int? maxBufferedItemCount, - PartitionKey? partitionKey, - IReadOnlyDictionary properties, - PartitionedQueryExecutionInfo partitionedQueryExecutionInfo, - ExecutionEnvironment? executionEnvironment, - bool? returnResultsInDeterministicOrder, - bool enableOptimisticDirectExecution, - TestInjections testInjections) - { - this.SqlQuerySpec = sqlQuerySpec ?? throw new ArgumentNullException(nameof(sqlQuerySpec)); - this.InitialUserContinuationToken = initialUserContinuationToken; - this.InitialFeedRange = initialFeedRange; - - int resolvedMaxConcurrency = maxConcurrency.GetValueOrDefault(InputParameters.DefaultMaxConcurrency); - if (resolvedMaxConcurrency < 0) - { - resolvedMaxConcurrency = int.MaxValue; - } - this.MaxConcurrency = resolvedMaxConcurrency; - - int resolvedMaxItemCount = maxItemCount.GetValueOrDefault(InputParameters.DefaultMaxItemCount); - if (resolvedMaxItemCount < 0) - { - resolvedMaxItemCount = int.MaxValue; - } - this.MaxItemCount = resolvedMaxItemCount; - - int resolvedMaxBufferedItemCount = maxBufferedItemCount.GetValueOrDefault(InputParameters.DefaultMaxBufferedItemCount); - if (resolvedMaxBufferedItemCount < 0) - { - resolvedMaxBufferedItemCount = int.MaxValue; - } - this.MaxBufferedItemCount = resolvedMaxBufferedItemCount; - - this.PartitionKey = partitionKey; - this.Properties = properties; - this.PartitionedQueryExecutionInfo = partitionedQueryExecutionInfo; - this.ExecutionEnvironment = executionEnvironment.GetValueOrDefault(InputParameters.DefaultExecutionEnvironment); - this.ReturnResultsInDeterministicOrder = returnResultsInDeterministicOrder.GetValueOrDefault(InputParameters.DefaultReturnResultsInDeterministicOrder); - this.EnableOptimisticDirectExecution = enableOptimisticDirectExecution; - this.TestInjections = testInjections; - } - - public SqlQuerySpec SqlQuerySpec { get; } - public CosmosElement InitialUserContinuationToken { get; } - public FeedRangeInternal InitialFeedRange { get; } - public int MaxConcurrency { get; } - public int MaxItemCount { get; } - public int MaxBufferedItemCount { get; } - public PartitionKey? PartitionKey { get; } - public IReadOnlyDictionary Properties { get; } - public PartitionedQueryExecutionInfo PartitionedQueryExecutionInfo { get; } - public ExecutionEnvironment ExecutionEnvironment { get; } - public bool ReturnResultsInDeterministicOrder { get; } - public TestInjections TestInjections { get; } - public bool EnableOptimisticDirectExecution { get; } - - public InputParameters WithContinuationToken(CosmosElement token) - { - return new InputParameters( - this.SqlQuerySpec, - token, - this.InitialFeedRange, - this.MaxConcurrency, - this.MaxItemCount, - this.MaxBufferedItemCount, - this.PartitionKey, - this.Properties, - this.PartitionedQueryExecutionInfo, - this.ExecutionEnvironment, - this.ReturnResultsInDeterministicOrder, - this.EnableOptimisticDirectExecution, - this.TestInjections); - } - } - - internal sealed class AggregateProjectionDetector - { - /// - /// Determines whether or not the SqlSelectSpec has an aggregate in the outer most query. - /// - /// The select spec to traverse. - /// Whether or not the SqlSelectSpec has an aggregate in the outer most query. - public static bool HasAggregate(SqlSelectSpec selectSpec) - { - return selectSpec.Accept(AggregateProjectionDectorVisitor.Singleton); - } - - private sealed class AggregateProjectionDectorVisitor : SqlSelectSpecVisitor - { - public static readonly AggregateProjectionDectorVisitor Singleton = new AggregateProjectionDectorVisitor(); - - public override bool Visit(SqlSelectListSpec selectSpec) - { - bool hasAggregates = false; - foreach (SqlSelectItem selectItem in selectSpec.Items) - { - hasAggregates |= selectItem.Expression.Accept(AggregateScalarExpressionDetector.Singleton); - } - - return hasAggregates; - } - - public override bool Visit(SqlSelectValueSpec selectSpec) - { - return selectSpec.Expression.Accept(AggregateScalarExpressionDetector.Singleton); - } - - public override bool Visit(SqlSelectStarSpec selectSpec) - { - return false; - } - - /// - /// Determines if there is an aggregate in a scalar expression. - /// - private sealed class AggregateScalarExpressionDetector : SqlScalarExpressionVisitor - { - private enum Aggregate - { - Min, - Max, - Sum, - Count, - Avg, - } - - public static readonly AggregateScalarExpressionDetector Singleton = new AggregateScalarExpressionDetector(); - - public override bool Visit(SqlAllScalarExpression sqlAllScalarExpression) - { - // No need to worry about aggregates within the subquery (they will recursively get rewritten). - return false; - } - - public override bool Visit(SqlArrayCreateScalarExpression sqlArrayCreateScalarExpression) - { - bool hasAggregates = false; - foreach (SqlScalarExpression item in sqlArrayCreateScalarExpression.Items) - { - hasAggregates |= item.Accept(this); - } - - return hasAggregates; - } - - public override bool Visit(SqlArrayScalarExpression sqlArrayScalarExpression) - { - // No need to worry about aggregates in the subquery (they will recursively get rewritten). - return false; - } - - public override bool Visit(SqlBetweenScalarExpression sqlBetweenScalarExpression) - { - return sqlBetweenScalarExpression.Expression.Accept(this) || - sqlBetweenScalarExpression.StartInclusive.Accept(this) || - sqlBetweenScalarExpression.EndInclusive.Accept(this); - } - - public override bool Visit(SqlBinaryScalarExpression sqlBinaryScalarExpression) - { - return sqlBinaryScalarExpression.LeftExpression.Accept(this) || - sqlBinaryScalarExpression.RightExpression.Accept(this); - } - - public override bool Visit(SqlCoalesceScalarExpression sqlCoalesceScalarExpression) - { - return sqlCoalesceScalarExpression.Left.Accept(this) || - sqlCoalesceScalarExpression.Right.Accept(this); - } - - public override bool Visit(SqlConditionalScalarExpression sqlConditionalScalarExpression) - { - return sqlConditionalScalarExpression.Condition.Accept(this) || - sqlConditionalScalarExpression.Consequent.Accept(this) || - sqlConditionalScalarExpression.Alternative.Accept(this); - } - - public override bool Visit(SqlExistsScalarExpression sqlExistsScalarExpression) - { - // No need to worry about aggregates within the subquery (they will recursively get rewritten). - return false; - } - - public override bool Visit(SqlFirstScalarExpression sqlFirstScalarExpression) - { - // No need to worry about aggregates within the subquery (they will recursively get rewritten). - return false; - } - - public override bool Visit(SqlFunctionCallScalarExpression sqlFunctionCallScalarExpression) - { - return !sqlFunctionCallScalarExpression.IsUdf && - Enum.TryParse(value: sqlFunctionCallScalarExpression.Name.Value, ignoreCase: true, result: out _); - } - - public override bool Visit(SqlInScalarExpression sqlInScalarExpression) - { - bool hasAggregates = false; - for (int i = 0; i < sqlInScalarExpression.Haystack.Length; i++) - { - hasAggregates |= sqlInScalarExpression.Haystack[i].Accept(this); - } - - return hasAggregates; - } - - public override bool Visit(SqlLastScalarExpression sqlLastScalarExpression) - { - // No need to worry about aggregates within the subquery (they will recursively get rewritten). - return false; - } - - public override bool Visit(SqlLiteralScalarExpression sqlLiteralScalarExpression) - { - return false; - } - - public override bool Visit(SqlMemberIndexerScalarExpression sqlMemberIndexerScalarExpression) - { - return sqlMemberIndexerScalarExpression.Member.Accept(this) || - sqlMemberIndexerScalarExpression.Indexer.Accept(this); - } - - public override bool Visit(SqlObjectCreateScalarExpression sqlObjectCreateScalarExpression) - { - bool hasAggregates = false; - foreach (SqlObjectProperty property in sqlObjectCreateScalarExpression.Properties) - { - hasAggregates |= property.Value.Accept(this); - } - - return hasAggregates; - } - - public override bool Visit(SqlPropertyRefScalarExpression sqlPropertyRefScalarExpression) - { - bool hasAggregates = false; - if (sqlPropertyRefScalarExpression.Member != null) - { - hasAggregates = sqlPropertyRefScalarExpression.Member.Accept(this); - } - - return hasAggregates; - } - - public override bool Visit(SqlSubqueryScalarExpression sqlSubqueryScalarExpression) - { - // No need to worry about the aggregates within the subquery since they get recursively evaluated. - return false; - } - - public override bool Visit(SqlUnaryScalarExpression sqlUnaryScalarExpression) - { - return sqlUnaryScalarExpression.Expression.Accept(this); - } - - public override bool Visit(SqlParameterRefScalarExpression scalarExpression) - { - return false; - } - - public override bool Visit(SqlLikeScalarExpression scalarExpression) - { - return false; - } - } - } - } - } + cosmosQueryContext.QueryClient, + cosmosQueryContext.ResourceLink, + partitionedQueryExecutionInfo, + containerQueryProperties, + inputParameters.Properties, + inputParameters.InitialFeedRange, + trace); + } + else + { + Documents.PartitionKeyDefinition partitionKeyDefinition = GetPartitionKeyDefinition(inputParameters, containerQueryProperties); + if (inputParameters.PartitionKey.HasValue) + { + Debug.Assert(partitionKeyDefinition != null, "CosmosQueryExecutionContextFactory Assert!", "PartitionKeyDefinition cannot be null if partitionKey is defined"); + targetRanges = await cosmosQueryContext.QueryClient.GetTargetPartitionKeyRangesAsync( + cosmosQueryContext.ResourceLink, + containerQueryProperties.ResourceId, + containerQueryProperties.EffectiveRangesForPartitionKey, + forceRefresh: false, + trace); + } + else + { + targetRanges = await cosmosQueryContext.QueryClient.GetTargetPartitionKeyRangesAsync( + cosmosQueryContext.ResourceLink, + containerQueryProperties.ResourceId, + new List> { FeedRangeEpk.FullRange.Range }, + forceRefresh: false, + trace); + } + } + + if (targetRanges.Count == 1) + { + return targetRanges.Single(); + } + + return null; + } + + public sealed class InputParameters + { + private const int DefaultMaxConcurrency = 0; + private const int DefaultMaxItemCount = 1000; + private const int DefaultMaxBufferedItemCount = 1000; + private const bool DefaultReturnResultsInDeterministicOrder = true; + private const ExecutionEnvironment DefaultExecutionEnvironment = ExecutionEnvironment.Client; + + public InputParameters( + SqlQuerySpec sqlQuerySpec, + CosmosElement initialUserContinuationToken, + FeedRangeInternal initialFeedRange, + int? maxConcurrency, + int? maxItemCount, + int? maxBufferedItemCount, + PartitionKey? partitionKey, + IReadOnlyDictionary properties, + PartitionedQueryExecutionInfo partitionedQueryExecutionInfo, + ExecutionEnvironment? executionEnvironment, + bool? returnResultsInDeterministicOrder, + bool enableOptimisticDirectExecution, + TestInjections testInjections) + { + this.SqlQuerySpec = sqlQuerySpec ?? throw new ArgumentNullException(nameof(sqlQuerySpec)); + this.InitialUserContinuationToken = initialUserContinuationToken; + this.InitialFeedRange = initialFeedRange; + + int resolvedMaxConcurrency = maxConcurrency.GetValueOrDefault(InputParameters.DefaultMaxConcurrency); + if (resolvedMaxConcurrency < 0) + { + resolvedMaxConcurrency = int.MaxValue; + } + this.MaxConcurrency = resolvedMaxConcurrency; + + int resolvedMaxItemCount = maxItemCount.GetValueOrDefault(InputParameters.DefaultMaxItemCount); + if (resolvedMaxItemCount < 0) + { + resolvedMaxItemCount = int.MaxValue; + } + this.MaxItemCount = resolvedMaxItemCount; + + int resolvedMaxBufferedItemCount = maxBufferedItemCount.GetValueOrDefault(InputParameters.DefaultMaxBufferedItemCount); + if (resolvedMaxBufferedItemCount < 0) + { + resolvedMaxBufferedItemCount = int.MaxValue; + } + this.MaxBufferedItemCount = resolvedMaxBufferedItemCount; + + this.PartitionKey = partitionKey; + this.Properties = properties; + this.PartitionedQueryExecutionInfo = partitionedQueryExecutionInfo; + this.ExecutionEnvironment = executionEnvironment.GetValueOrDefault(InputParameters.DefaultExecutionEnvironment); + this.ReturnResultsInDeterministicOrder = returnResultsInDeterministicOrder.GetValueOrDefault(InputParameters.DefaultReturnResultsInDeterministicOrder); + this.EnableOptimisticDirectExecution = enableOptimisticDirectExecution; + this.TestInjections = testInjections; + } + + public SqlQuerySpec SqlQuerySpec { get; } + public CosmosElement InitialUserContinuationToken { get; } + public FeedRangeInternal InitialFeedRange { get; } + public int MaxConcurrency { get; } + public int MaxItemCount { get; } + public int MaxBufferedItemCount { get; } + public PartitionKey? PartitionKey { get; } + public IReadOnlyDictionary Properties { get; } + public PartitionedQueryExecutionInfo PartitionedQueryExecutionInfo { get; } + public ExecutionEnvironment ExecutionEnvironment { get; } + public bool ReturnResultsInDeterministicOrder { get; } + public TestInjections TestInjections { get; } + public bool EnableOptimisticDirectExecution { get; } + + public InputParameters WithContinuationToken(CosmosElement token) + { + return new InputParameters( + this.SqlQuerySpec, + token, + this.InitialFeedRange, + this.MaxConcurrency, + this.MaxItemCount, + this.MaxBufferedItemCount, + this.PartitionKey, + this.Properties, + this.PartitionedQueryExecutionInfo, + this.ExecutionEnvironment, + this.ReturnResultsInDeterministicOrder, + this.EnableOptimisticDirectExecution, + this.TestInjections); + } + } + + internal sealed class AggregateProjectionDetector + { + /// + /// Determines whether or not the SqlSelectSpec has an aggregate in the outer most query. + /// + /// The select spec to traverse. + /// Whether or not the SqlSelectSpec has an aggregate in the outer most query. + public static bool HasAggregate(SqlSelectSpec selectSpec) + { + return selectSpec.Accept(AggregateProjectionDectorVisitor.Singleton); + } + + private sealed class AggregateProjectionDectorVisitor : SqlSelectSpecVisitor + { + public static readonly AggregateProjectionDectorVisitor Singleton = new AggregateProjectionDectorVisitor(); + + public override bool Visit(SqlSelectListSpec selectSpec) + { + bool hasAggregates = false; + foreach (SqlSelectItem selectItem in selectSpec.Items) + { + hasAggregates |= selectItem.Expression.Accept(AggregateScalarExpressionDetector.Singleton); + } + + return hasAggregates; + } + + public override bool Visit(SqlSelectValueSpec selectSpec) + { + return selectSpec.Expression.Accept(AggregateScalarExpressionDetector.Singleton); + } + + public override bool Visit(SqlSelectStarSpec selectSpec) + { + return false; + } + + /// + /// Determines if there is an aggregate in a scalar expression. + /// + private sealed class AggregateScalarExpressionDetector : SqlScalarExpressionVisitor + { + private enum Aggregate + { + Min, + Max, + Sum, + Count, + Avg, + } + + public static readonly AggregateScalarExpressionDetector Singleton = new AggregateScalarExpressionDetector(); + + public override bool Visit(SqlAllScalarExpression sqlAllScalarExpression) + { + // No need to worry about aggregates within the subquery (they will recursively get rewritten). + return false; + } + + public override bool Visit(SqlArrayCreateScalarExpression sqlArrayCreateScalarExpression) + { + bool hasAggregates = false; + foreach (SqlScalarExpression item in sqlArrayCreateScalarExpression.Items) + { + hasAggregates |= item.Accept(this); + } + + return hasAggregates; + } + + public override bool Visit(SqlArrayScalarExpression sqlArrayScalarExpression) + { + // No need to worry about aggregates in the subquery (they will recursively get rewritten). + return false; + } + + public override bool Visit(SqlBetweenScalarExpression sqlBetweenScalarExpression) + { + return sqlBetweenScalarExpression.Expression.Accept(this) || + sqlBetweenScalarExpression.StartInclusive.Accept(this) || + sqlBetweenScalarExpression.EndInclusive.Accept(this); + } + + public override bool Visit(SqlBinaryScalarExpression sqlBinaryScalarExpression) + { + return sqlBinaryScalarExpression.LeftExpression.Accept(this) || + sqlBinaryScalarExpression.RightExpression.Accept(this); + } + + public override bool Visit(SqlCoalesceScalarExpression sqlCoalesceScalarExpression) + { + return sqlCoalesceScalarExpression.Left.Accept(this) || + sqlCoalesceScalarExpression.Right.Accept(this); + } + + public override bool Visit(SqlConditionalScalarExpression sqlConditionalScalarExpression) + { + return sqlConditionalScalarExpression.Condition.Accept(this) || + sqlConditionalScalarExpression.Consequent.Accept(this) || + sqlConditionalScalarExpression.Alternative.Accept(this); + } + + public override bool Visit(SqlExistsScalarExpression sqlExistsScalarExpression) + { + // No need to worry about aggregates within the subquery (they will recursively get rewritten). + return false; + } + + public override bool Visit(SqlFirstScalarExpression sqlFirstScalarExpression) + { + // No need to worry about aggregates within the subquery (they will recursively get rewritten). + return false; + } + + public override bool Visit(SqlFunctionCallScalarExpression sqlFunctionCallScalarExpression) + { + return !sqlFunctionCallScalarExpression.IsUdf && + Enum.TryParse(value: sqlFunctionCallScalarExpression.Name.Value, ignoreCase: true, result: out _); + } + + public override bool Visit(SqlInScalarExpression sqlInScalarExpression) + { + bool hasAggregates = false; + for (int i = 0; i < sqlInScalarExpression.Haystack.Length; i++) + { + hasAggregates |= sqlInScalarExpression.Haystack[i].Accept(this); + } + + return hasAggregates; + } + + public override bool Visit(SqlLastScalarExpression sqlLastScalarExpression) + { + // No need to worry about aggregates within the subquery (they will recursively get rewritten). + return false; + } + + public override bool Visit(SqlLiteralScalarExpression sqlLiteralScalarExpression) + { + return false; + } + + public override bool Visit(SqlMemberIndexerScalarExpression sqlMemberIndexerScalarExpression) + { + return sqlMemberIndexerScalarExpression.Member.Accept(this) || + sqlMemberIndexerScalarExpression.Indexer.Accept(this); + } + + public override bool Visit(SqlObjectCreateScalarExpression sqlObjectCreateScalarExpression) + { + bool hasAggregates = false; + foreach (SqlObjectProperty property in sqlObjectCreateScalarExpression.Properties) + { + hasAggregates |= property.Value.Accept(this); + } + + return hasAggregates; + } + + public override bool Visit(SqlPropertyRefScalarExpression sqlPropertyRefScalarExpression) + { + bool hasAggregates = false; + if (sqlPropertyRefScalarExpression.Member != null) + { + hasAggregates = sqlPropertyRefScalarExpression.Member.Accept(this); + } + + return hasAggregates; + } + + public override bool Visit(SqlSubqueryScalarExpression sqlSubqueryScalarExpression) + { + // No need to worry about the aggregates within the subquery since they get recursively evaluated. + return false; + } + + public override bool Visit(SqlUnaryScalarExpression sqlUnaryScalarExpression) + { + return sqlUnaryScalarExpression.Expression.Accept(this); + } + + public override bool Visit(SqlParameterRefScalarExpression scalarExpression) + { + return false; + } + + public override bool Visit(SqlLikeScalarExpression scalarExpression) + { + return false; + } + } + } + } + } } \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/Parallel/ParallelCrossPartitionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/Parallel/ParallelCrossPartitionQueryPipelineStage.cs index d9b564b5f3..7f1d1079c1 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/Parallel/ParallelCrossPartitionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/Parallel/ParallelCrossPartitionQueryPipelineStage.cs @@ -1,273 +1,277 @@ -// ------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -// ------------------------------------------------------------ - -namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.Parallel -{ - using System; - using System.Collections.Generic; - using System.Linq; - using System.Threading; - using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Pagination; - using Microsoft.Azure.Cosmos.Query.Core.Exceptions; - using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Pagination; - using Microsoft.Azure.Cosmos.Tracing; - using static Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.PartitionMapper; - - /// - /// is an implementation of that drain results from multiple remote nodes. - /// This class is responsible for draining cross partition queries that do not have order by conditions. - /// The way parallel queries work is that it drains from the left most partition first. - /// This class handles draining in the correct order and can also stop and resume the query - /// by generating a continuation token and resuming from said continuation token. - /// - internal sealed class ParallelCrossPartitionQueryPipelineStage : IQueryPipelineStage - { - private readonly CrossPartitionRangePageAsyncEnumerator crossPartitionRangePageAsyncEnumerator; - - private ParallelCrossPartitionQueryPipelineStage( +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.Parallel +{ + using System; + using System.Collections.Generic; + using System.Linq; + using System.Threading; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Pagination; + using Microsoft.Azure.Cosmos.Query.Core.Exceptions; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Pagination; + using Microsoft.Azure.Cosmos.Query.Core.QueryClient; + using Microsoft.Azure.Cosmos.Tracing; + using static Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.PartitionMapper; + + /// + /// is an implementation of that drain results from multiple remote nodes. + /// This class is responsible for draining cross partition queries that do not have order by conditions. + /// The way parallel queries work is that it drains from the left most partition first. + /// This class handles draining in the correct order and can also stop and resume the query + /// by generating a continuation token and resuming from said continuation token. + /// + internal sealed class ParallelCrossPartitionQueryPipelineStage : IQueryPipelineStage + { + private readonly CrossPartitionRangePageAsyncEnumerator crossPartitionRangePageAsyncEnumerator; + + private ParallelCrossPartitionQueryPipelineStage( CrossPartitionRangePageAsyncEnumerator crossPartitionRangePageAsyncEnumerator) - { - this.crossPartitionRangePageAsyncEnumerator = crossPartitionRangePageAsyncEnumerator ?? throw new ArgumentNullException(nameof(crossPartitionRangePageAsyncEnumerator)); - } - - public TryCatch Current { get; private set; } - - public ValueTask DisposeAsync() - { - return this.crossPartitionRangePageAsyncEnumerator.DisposeAsync(); - } - - // In order to maintain the continuation token for the user we must drain with a few constraints - // 1) We fully drain from the left most partition before moving on to the next partition - // 2) We drain only full pages from the document producer so we aren't left with a partial page - // otherwise we would need to add to the continuation token how many items to skip over on that page. + { + this.crossPartitionRangePageAsyncEnumerator = crossPartitionRangePageAsyncEnumerator ?? throw new ArgumentNullException(nameof(crossPartitionRangePageAsyncEnumerator)); + } + + public TryCatch Current { get; private set; } + + public ValueTask DisposeAsync() + { + return this.crossPartitionRangePageAsyncEnumerator.DisposeAsync(); + } + + // In order to maintain the continuation token for the user we must drain with a few constraints + // 1) We fully drain from the left most partition before moving on to the next partition + // 2) We drain only full pages from the document producer so we aren't left with a partial page + // otherwise we would need to add to the continuation token how many items to skip over on that page. public async ValueTask MoveNextAsync(ITrace trace, CancellationToken cancellationToken) - { - if (trace == null) - { - throw new ArgumentNullException(nameof(trace)); - } - + { + if (trace == null) + { + throw new ArgumentNullException(nameof(trace)); + } + if (!await this.crossPartitionRangePageAsyncEnumerator.MoveNextAsync(trace, cancellationToken)) - { - this.Current = default; - return false; - } - - TryCatch> currentCrossPartitionPage = this.crossPartitionRangePageAsyncEnumerator.Current; - if (currentCrossPartitionPage.Failed) - { - this.Current = TryCatch.FromException(currentCrossPartitionPage.Exception); - return true; - } - - CrossFeedRangePage crossPartitionPageResult = currentCrossPartitionPage.Result; - QueryPage backendQueryPage = crossPartitionPageResult.Page; - CrossFeedRangeState crossPartitionState = crossPartitionPageResult.State; - - QueryState queryState; - if (crossPartitionState == null) - { - queryState = null; - } - else - { - // left most and any non null continuations - IOrderedEnumerable> feedRangeStates = crossPartitionState - .Value - .ToArray() - .OrderBy(tuple => ((FeedRangeEpk)tuple.FeedRange).Range.Min); - - List activeParallelContinuationTokens = new List(); - { - FeedRangeState firstState = feedRangeStates.First(); - ParallelContinuationToken firstParallelContinuationToken = new ParallelContinuationToken( - token: firstState.State != null ? ((CosmosString)firstState.State.Value).Value : null, - range: ((FeedRangeEpk)firstState.FeedRange).Range); - - activeParallelContinuationTokens.Add(firstParallelContinuationToken); - } - - foreach (FeedRangeState feedRangeState in feedRangeStates.Skip(1)) - { + { + this.Current = default; + return false; + } + + TryCatch> currentCrossPartitionPage = this.crossPartitionRangePageAsyncEnumerator.Current; + if (currentCrossPartitionPage.Failed) + { + this.Current = TryCatch.FromException(currentCrossPartitionPage.Exception); + return true; + } + + CrossFeedRangePage crossPartitionPageResult = currentCrossPartitionPage.Result; + QueryPage backendQueryPage = crossPartitionPageResult.Page; + CrossFeedRangeState crossPartitionState = crossPartitionPageResult.State; + + QueryState queryState; + if (crossPartitionState == null) + { + queryState = null; + } + else + { + // left most and any non null continuations + IOrderedEnumerable> feedRangeStates = crossPartitionState + .Value + .ToArray() + .OrderBy(tuple => ((FeedRangeEpk)tuple.FeedRange).Range.Min); + + List activeParallelContinuationTokens = new List(); + { + FeedRangeState firstState = feedRangeStates.First(); + ParallelContinuationToken firstParallelContinuationToken = new ParallelContinuationToken( + token: firstState.State != null ? ((CosmosString)firstState.State.Value).Value : null, + range: ((FeedRangeEpk)firstState.FeedRange).Range); + + activeParallelContinuationTokens.Add(firstParallelContinuationToken); + } + + foreach (FeedRangeState feedRangeState in feedRangeStates.Skip(1)) + { cancellationToken.ThrowIfCancellationRequested(); - - if (feedRangeState.State != null) - { - ParallelContinuationToken parallelContinuationToken = new ParallelContinuationToken( - token: feedRangeState.State != null ? ((CosmosString)feedRangeState.State.Value).Value : null, - range: ((FeedRangeEpk)feedRangeState.FeedRange).Range); - - activeParallelContinuationTokens.Add(parallelContinuationToken); - } - } - - IEnumerable cosmosElementContinuationTokens = activeParallelContinuationTokens - .Select(token => ParallelContinuationToken.ToCosmosElement(token)); - CosmosArray cosmosElementParallelContinuationTokens = CosmosArray.Create(cosmosElementContinuationTokens); - - queryState = new QueryState(cosmosElementParallelContinuationTokens); - } - - QueryPage crossPartitionQueryPage = new QueryPage( - backendQueryPage.Documents, - backendQueryPage.RequestCharge, - backendQueryPage.ActivityId, - backendQueryPage.CosmosQueryExecutionInfo, - distributionPlanSpec: default, - backendQueryPage.DisallowContinuationTokenMessage, - backendQueryPage.AdditionalHeaders, + + if (feedRangeState.State != null) + { + ParallelContinuationToken parallelContinuationToken = new ParallelContinuationToken( + token: feedRangeState.State != null ? ((CosmosString)feedRangeState.State.Value).Value : null, + range: ((FeedRangeEpk)feedRangeState.FeedRange).Range); + + activeParallelContinuationTokens.Add(parallelContinuationToken); + } + } + + IEnumerable cosmosElementContinuationTokens = activeParallelContinuationTokens + .Select(token => ParallelContinuationToken.ToCosmosElement(token)); + CosmosArray cosmosElementParallelContinuationTokens = CosmosArray.Create(cosmosElementContinuationTokens); + + queryState = new QueryState(cosmosElementParallelContinuationTokens); + } + + QueryPage crossPartitionQueryPage = new QueryPage( + backendQueryPage.Documents, + backendQueryPage.RequestCharge, + backendQueryPage.ActivityId, + backendQueryPage.CosmosQueryExecutionInfo, + distributionPlanSpec: default, + backendQueryPage.DisallowContinuationTokenMessage, + backendQueryPage.AdditionalHeaders, queryState, backendQueryPage.Streaming); - - this.Current = TryCatch.FromResult(crossPartitionQueryPage); - return true; - } - - public static TryCatch MonadicCreate( - IDocumentContainer documentContainer, - SqlQuerySpec sqlQuerySpec, - IReadOnlyList targetRanges, - Cosmos.PartitionKey? partitionKey, - QueryPaginationOptions queryPaginationOptions, - int maxConcurrency, - PrefetchPolicy prefetchPolicy, + + this.Current = TryCatch.FromResult(crossPartitionQueryPage); + return true; + } + + public static TryCatch MonadicCreate( + IDocumentContainer documentContainer, + SqlQuerySpec sqlQuerySpec, + IReadOnlyList targetRanges, + Cosmos.PartitionKey? partitionKey, + QueryPaginationOptions queryPaginationOptions, + ContainerQueryProperties containerQueryProperties, + int maxConcurrency, + PrefetchPolicy prefetchPolicy, CosmosElement continuationToken) - { - if (targetRanges == null) - { - throw new ArgumentNullException(nameof(targetRanges)); - } - - if (targetRanges.Count == 0) - { - throw new ArgumentException($"{nameof(targetRanges)} must have some elements"); - } - - TryCatch> monadicExtractState = MonadicExtractState(continuationToken, targetRanges); - if (monadicExtractState.Failed) - { - return TryCatch.FromException(monadicExtractState.Exception); - } - - CrossFeedRangeState state = monadicExtractState.Result; - - CrossPartitionRangePageAsyncEnumerator crossPartitionPageEnumerator = new CrossPartitionRangePageAsyncEnumerator( - feedRangeProvider: documentContainer, - createPartitionRangeEnumerator: ParallelCrossPartitionQueryPipelineStage.MakeCreateFunction(documentContainer, sqlQuerySpec, queryPaginationOptions, partitionKey), - comparer: Comparer.Singleton, - maxConcurrency: maxConcurrency, - prefetchPolicy: prefetchPolicy, + { + if (targetRanges == null) + { + throw new ArgumentNullException(nameof(targetRanges)); + } + + if (targetRanges.Count == 0) + { + throw new ArgumentException($"{nameof(targetRanges)} must have some elements"); + } + + TryCatch> monadicExtractState = MonadicExtractState(continuationToken, targetRanges); + if (monadicExtractState.Failed) + { + return TryCatch.FromException(monadicExtractState.Exception); + } + + CrossFeedRangeState state = monadicExtractState.Result; + + CrossPartitionRangePageAsyncEnumerator crossPartitionPageEnumerator = new CrossPartitionRangePageAsyncEnumerator( + feedRangeProvider: documentContainer, + createPartitionRangeEnumerator: ParallelCrossPartitionQueryPipelineStage.MakeCreateFunction(documentContainer, sqlQuerySpec, queryPaginationOptions, partitionKey, containerQueryProperties), + comparer: Comparer.Singleton, + maxConcurrency: maxConcurrency, + prefetchPolicy: prefetchPolicy, state: state); - + ParallelCrossPartitionQueryPipelineStage stage = new ParallelCrossPartitionQueryPipelineStage(crossPartitionPageEnumerator); - return TryCatch.FromResult(stage); - } - - private static TryCatch> MonadicExtractState( - CosmosElement continuationToken, - IReadOnlyList ranges) - { - if (continuationToken == null) - { - // Full fan out to the ranges with null continuations - CrossFeedRangeState fullFanOutState = new CrossFeedRangeState(ranges.Select(range => new FeedRangeState(range, (QueryState)null)).ToArray()); - return TryCatch>.FromResult(fullFanOutState); - } - - if (!(continuationToken is CosmosArray parallelContinuationTokenListRaw)) - { - return TryCatch>.FromException( - new MalformedContinuationTokenException( - $"Invalid format for continuation token {continuationToken} for {nameof(ParallelCrossPartitionQueryPipelineStage)}")); - } - - if (parallelContinuationTokenListRaw.Count == 0) - { - return TryCatch>.FromException( - new MalformedContinuationTokenException( - $"Invalid format for continuation token {continuationToken} for {nameof(ParallelCrossPartitionQueryPipelineStage)}")); - } - - List parallelContinuationTokens = new List(); - foreach (CosmosElement parallelContinuationTokenRaw in parallelContinuationTokenListRaw) - { - TryCatch tryCreateParallelContinuationToken = ParallelContinuationToken.TryCreateFromCosmosElement(parallelContinuationTokenRaw); - if (tryCreateParallelContinuationToken.Failed) - { - return TryCatch>.FromException( - tryCreateParallelContinuationToken.Exception); - } - - parallelContinuationTokens.Add(tryCreateParallelContinuationToken.Result); - } - - TryCatch> partitionMappingMonad = PartitionMapper.MonadicGetPartitionMapping( - ranges, - parallelContinuationTokens); - if (partitionMappingMonad.Failed) - { - return TryCatch>.FromException( - partitionMappingMonad.Exception); - } - - PartitionMapping partitionMapping = partitionMappingMonad.Result; - List> feedRangeStates = new List>(); - - List> rangesToInitialize = new List>() - { - // Skip all the partitions left of the target range, since they have already been drained fully. - partitionMapping.TargetMapping, - partitionMapping.MappingRightOfTarget, - }; - - foreach (IReadOnlyDictionary rangeToInitalize in rangesToInitialize) - { - foreach (KeyValuePair kvp in rangeToInitalize) - { - FeedRangeState feedRangeState = new FeedRangeState(kvp.Key, kvp.Value?.Token != null ? new QueryState(CosmosString.Create(kvp.Value.Token)) : null); - feedRangeStates.Add(feedRangeState); - } - } - - CrossFeedRangeState crossPartitionState = new CrossFeedRangeState(feedRangeStates.ToArray()); - - return TryCatch>.FromResult(crossPartitionState); - } - - private static CreatePartitionRangePageAsyncEnumerator MakeCreateFunction( - IQueryDataSource queryDataSource, - SqlQuerySpec sqlQuerySpec, - QueryPaginationOptions queryPaginationOptions, - Cosmos.PartitionKey? partitionKey) => (FeedRangeState feedRangeState) => new QueryPartitionRangePageAsyncEnumerator( - queryDataSource, - sqlQuerySpec, - feedRangeState, - partitionKey, - queryPaginationOptions); - - private sealed class Comparer : IComparer> - { - public static readonly Comparer Singleton = new Comparer(); - - public int Compare( - PartitionRangePageAsyncEnumerator partitionRangePageEnumerator1, - PartitionRangePageAsyncEnumerator partitionRangePageEnumerator2) - { - if (object.ReferenceEquals(partitionRangePageEnumerator1, partitionRangePageEnumerator2)) - { - return 0; - } - - // Either both don't have results or both do. - return string.CompareOrdinal( - ((FeedRangeEpk)partitionRangePageEnumerator1.FeedRangeState.FeedRange).Range.Min, - ((FeedRangeEpk)partitionRangePageEnumerator2.FeedRangeState.FeedRange).Range.Min); - } - } - } + return TryCatch.FromResult(stage); + } + + private static TryCatch> MonadicExtractState( + CosmosElement continuationToken, + IReadOnlyList ranges) + { + if (continuationToken == null) + { + // Full fan out to the ranges with null continuations + CrossFeedRangeState fullFanOutState = new CrossFeedRangeState(ranges.Select(range => new FeedRangeState(range, (QueryState)null)).ToArray()); + return TryCatch>.FromResult(fullFanOutState); + } + + if (!(continuationToken is CosmosArray parallelContinuationTokenListRaw)) + { + return TryCatch>.FromException( + new MalformedContinuationTokenException( + $"Invalid format for continuation token {continuationToken} for {nameof(ParallelCrossPartitionQueryPipelineStage)}")); + } + + if (parallelContinuationTokenListRaw.Count == 0) + { + return TryCatch>.FromException( + new MalformedContinuationTokenException( + $"Invalid format for continuation token {continuationToken} for {nameof(ParallelCrossPartitionQueryPipelineStage)}")); + } + + List parallelContinuationTokens = new List(); + foreach (CosmosElement parallelContinuationTokenRaw in parallelContinuationTokenListRaw) + { + TryCatch tryCreateParallelContinuationToken = ParallelContinuationToken.TryCreateFromCosmosElement(parallelContinuationTokenRaw); + if (tryCreateParallelContinuationToken.Failed) + { + return TryCatch>.FromException( + tryCreateParallelContinuationToken.Exception); + } + + parallelContinuationTokens.Add(tryCreateParallelContinuationToken.Result); + } + + TryCatch> partitionMappingMonad = PartitionMapper.MonadicGetPartitionMapping( + ranges, + parallelContinuationTokens); + if (partitionMappingMonad.Failed) + { + return TryCatch>.FromException( + partitionMappingMonad.Exception); + } + + PartitionMapping partitionMapping = partitionMappingMonad.Result; + List> feedRangeStates = new List>(); + + List> rangesToInitialize = new List>() + { + // Skip all the partitions left of the target range, since they have already been drained fully. + partitionMapping.TargetMapping, + partitionMapping.MappingRightOfTarget, + }; + + foreach (IReadOnlyDictionary rangeToInitalize in rangesToInitialize) + { + foreach (KeyValuePair kvp in rangeToInitalize) + { + FeedRangeState feedRangeState = new FeedRangeState(kvp.Key, kvp.Value?.Token != null ? new QueryState(CosmosString.Create(kvp.Value.Token)) : null); + feedRangeStates.Add(feedRangeState); + } + } + + CrossFeedRangeState crossPartitionState = new CrossFeedRangeState(feedRangeStates.ToArray()); + + return TryCatch>.FromResult(crossPartitionState); + } + + private static CreatePartitionRangePageAsyncEnumerator MakeCreateFunction( + IQueryDataSource queryDataSource, + SqlQuerySpec sqlQuerySpec, + QueryPaginationOptions queryPaginationOptions, + Cosmos.PartitionKey? partitionKey, + ContainerQueryProperties containerQueryProperties) => (FeedRangeState feedRangeState) => new QueryPartitionRangePageAsyncEnumerator( + queryDataSource, + sqlQuerySpec, + feedRangeState, + partitionKey, + queryPaginationOptions, + containerQueryProperties); + + private sealed class Comparer : IComparer> + { + public static readonly Comparer Singleton = new Comparer(); + + public int Compare( + PartitionRangePageAsyncEnumerator partitionRangePageEnumerator1, + PartitionRangePageAsyncEnumerator partitionRangePageEnumerator2) + { + if (object.ReferenceEquals(partitionRangePageEnumerator1, partitionRangePageEnumerator2)) + { + return 0; + } + + // Either both don't have results or both do. + return string.CompareOrdinal( + ((FeedRangeEpk)partitionRangePageEnumerator1.FeedRangeState.FeedRange).Range.Min, + ((FeedRangeEpk)partitionRangePageEnumerator2.FeedRangeState.FeedRange).Range.Min); + } + } + } } \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/Parallel/QueryPartitionRangePageAsyncEnumerator.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/Parallel/QueryPartitionRangePageAsyncEnumerator.cs index 5744ed734b..7ff976d7a2 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/Parallel/QueryPartitionRangePageAsyncEnumerator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/Parallel/QueryPartitionRangePageAsyncEnumerator.cs @@ -10,6 +10,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.Parallel using Microsoft.Azure.Cosmos.Pagination; using Microsoft.Azure.Cosmos.Query.Core.Monads; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Pagination; + using Microsoft.Azure.Cosmos.Query.Core.QueryClient; using Microsoft.Azure.Cosmos.Tracing; internal sealed class QueryPartitionRangePageAsyncEnumerator : PartitionRangePageAsyncEnumerator @@ -17,6 +18,7 @@ internal sealed class QueryPartitionRangePageAsyncEnumerator : PartitionRangePag private readonly IQueryDataSource queryDataSource; private readonly SqlQuerySpec sqlQuerySpec; private readonly QueryPaginationOptions queryPaginationOptions; + private readonly ContainerQueryProperties containerQueryProperties; private readonly Cosmos.PartitionKey? partitionKey; public QueryPartitionRangePageAsyncEnumerator( @@ -24,13 +26,15 @@ public QueryPartitionRangePageAsyncEnumerator( SqlQuerySpec sqlQuerySpec, FeedRangeState feedRangeState, Cosmos.PartitionKey? partitionKey, - QueryPaginationOptions queryPaginationOptions) + QueryPaginationOptions queryPaginationOptions, + ContainerQueryProperties containerQueryProperties) : base(feedRangeState) { this.queryDataSource = queryDataSource ?? throw new ArgumentNullException(nameof(queryDataSource)); this.sqlQuerySpec = sqlQuerySpec ?? throw new ArgumentNullException(nameof(sqlQuerySpec)); this.queryPaginationOptions = queryPaginationOptions; this.partitionKey = partitionKey; + this.containerQueryProperties = containerQueryProperties; } public override ValueTask DisposeAsync() => default; @@ -42,9 +46,7 @@ protected override Task> GetNextPageAsync(ITrace trace, Canc throw new ArgumentNullException(nameof(trace)); } - // We sadly need to check the partition key, since a user can set a partition key in the request options with a different continuation token. - // In the future the partition filtering and continuation information needs to be a tightly bounded contract (like cross feed range state). - FeedRangeInternal feedRange = this.partitionKey.HasValue ? new FeedRangePartitionKey(this.partitionKey.Value) : this.FeedRangeState.FeedRange; + FeedRangeInternal feedRange = this.LimitFeedRangeToSinglePartition(); return this.queryDataSource.MonadicQueryAsync( sqlQuerySpec: this.sqlQuerySpec, feedRangeState: new FeedRangeState(feedRange, this.FeedRangeState.State), @@ -52,5 +54,95 @@ protected override Task> GetNextPageAsync(ITrace trace, Canc trace: trace, cancellationToken); } + + /// + /// Updates the FeedRange to limit the scope of this enumerator to single physical partition. + /// Generally speaking, a subpartitioned container can experience split partition at any level of hierarchical partition key. + /// This could cause a situation where more than one physical partition contains the data for a partial partition key. + /// Currently, enumerator instantiation does not honor physical partition boundary and allocates entire epk range which could spans across multiple physical partitions to the enumerator. + /// Since such an epk range does not exist at the container level, Service generates a GoneException. + /// This method restrics the range of each container by shrinking the ends of the range so that they do not span across physical partition. + /// + private FeedRangeInternal LimitFeedRangeToSinglePartition() + { + // We sadly need to check the partition key, since a user can set a partition key in the request options with a different continuation token. + // In the future the partition filtering and continuation information needs to be a tightly bounded contract (like cross feed range state). + FeedRangeInternal feedRange = this.FeedRangeState.FeedRange; + if (this.partitionKey.HasValue) + { + // ISSUE-HACK-adityasa-3/25/2024 - We should not update the original feed range inside this class. + // Instead we should guarantee that when enumerator is instantiated it is limited to a single physical partition. + // Ultimately we should remove enumerator's dependency on PartitionKey. + if ((this.containerQueryProperties.PartitionKeyDefinition.Paths.Count > 1) && + (this.partitionKey.Value.InternalKey.Components.Count != this.containerQueryProperties.PartitionKeyDefinition.Paths.Count) && + (feedRange is FeedRangeEpk feedRangeEpk)) + { + if (this.containerQueryProperties.EffectiveRangesForPartitionKey == null || + this.containerQueryProperties.EffectiveRangesForPartitionKey.Count == 0) + { + throw new InvalidOperationException( + "EffectiveRangesForPartitionKey should be populated when PK is specified in request options."); + } + + foreach (Documents.Routing.Range epkForPartitionKey in + this.containerQueryProperties.EffectiveRangesForPartitionKey) + { + if (Documents.Routing.Range.CheckOverlapping( + feedRangeEpk.Range, + epkForPartitionKey)) + { + if (!feedRangeEpk.Range.Equals(epkForPartitionKey)) + { + String overlappingMin; + bool minInclusive; + String overlappingMax; + bool maxInclusive; + + if (Documents.Routing.Range.MinComparer.Instance.Compare( + epkForPartitionKey, + feedRangeEpk.Range) < 0) + { + overlappingMin = feedRangeEpk.Range.Min; + minInclusive = feedRangeEpk.Range.IsMinInclusive; + } + else + { + overlappingMin = epkForPartitionKey.Min; + minInclusive = epkForPartitionKey.IsMinInclusive; + } + + if (Documents.Routing.Range.MaxComparer.Instance.Compare( + epkForPartitionKey, + feedRangeEpk.Range) > 0) + { + overlappingMax = feedRangeEpk.Range.Max; + maxInclusive = feedRangeEpk.Range.IsMaxInclusive; + } + else + { + overlappingMax = epkForPartitionKey.Max; + maxInclusive = epkForPartitionKey.IsMaxInclusive; + } + + feedRange = new FeedRangeEpk( + new Documents.Routing.Range( + overlappingMin, + overlappingMax, + minInclusive, + maxInclusive)); + } + + break; + } + } + } + else + { + feedRange = new FeedRangePartitionKey(this.partitionKey.Value); + } + } + + return feedRange; + } } } \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs index 8bad485a32..ff5f6bb438 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs @@ -1,318 +1,323 @@ -// ------------------------------------------------------------ -// Copyright (c) Microsoft Corporation. All rights reserved. -// ------------------------------------------------------------ - -namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.OptimisticDirectExecutionQuery -{ - using System; - using System.Collections.Generic; - using System.Diagnostics; - using System.Linq; - using System.Threading; - using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Pagination; - using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; - using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.Parallel; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Pagination; - using Microsoft.Azure.Cosmos.Tracing; - using Microsoft.Azure.Documents; - - internal sealed class OptimisticDirectExecutionQueryPipelineStage : IQueryPipelineStage - { - private enum ExecutionState - { - OptimisticDirectExecution, - SpecializedDocumentQueryExecution, - } - - private const string OptimisticDirectExecutionToken = "OptimisticDirectExecutionToken"; - private readonly FallbackQueryPipelineStageFactory queryPipelineStageFactory; - private TryCatch inner; - private CosmosElement continuationToken; - private ExecutionState executionState; - private bool? previousRequiresDistribution; - - private OptimisticDirectExecutionQueryPipelineStage(TryCatch inner, FallbackQueryPipelineStageFactory queryPipelineStageFactory, CosmosElement continuationToken) - { - this.inner = inner; - this.queryPipelineStageFactory = queryPipelineStageFactory; - this.continuationToken = continuationToken; - this.executionState = ExecutionState.OptimisticDirectExecution; - - if (this.continuationToken != null) - { - this.previousRequiresDistribution = false; - } - } - - public delegate Task> FallbackQueryPipelineStageFactory(CosmosElement continuationToken); - - public TryCatch Current => this.inner.Try(pipelineStage => pipelineStage.Current); - - public ValueTask DisposeAsync() - { - return this.inner.Failed ? default : this.inner.Result.DisposeAsync(); - } - +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.OptimisticDirectExecutionQuery +{ + using System; + using System.Collections.Generic; + using System.Diagnostics; + using System.Linq; + using System.Threading; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Pagination; + using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.Parallel; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Pagination; + using Microsoft.Azure.Cosmos.Query.Core.QueryClient; + using Microsoft.Azure.Cosmos.Tracing; + using Microsoft.Azure.Documents; + + internal sealed class OptimisticDirectExecutionQueryPipelineStage : IQueryPipelineStage + { + private enum ExecutionState + { + OptimisticDirectExecution, + SpecializedDocumentQueryExecution, + } + + private const string OptimisticDirectExecutionToken = "OptimisticDirectExecutionToken"; + private readonly FallbackQueryPipelineStageFactory queryPipelineStageFactory; + private TryCatch inner; + private CosmosElement continuationToken; + private ExecutionState executionState; + private bool? previousRequiresDistribution; + + private OptimisticDirectExecutionQueryPipelineStage(TryCatch inner, FallbackQueryPipelineStageFactory queryPipelineStageFactory, CosmosElement continuationToken) + { + this.inner = inner; + this.queryPipelineStageFactory = queryPipelineStageFactory; + this.continuationToken = continuationToken; + this.executionState = ExecutionState.OptimisticDirectExecution; + + if (this.continuationToken != null) + { + this.previousRequiresDistribution = false; + } + } + + public delegate Task> FallbackQueryPipelineStageFactory(CosmosElement continuationToken); + + public TryCatch Current => this.inner.Try(pipelineStage => pipelineStage.Current); + + public ValueTask DisposeAsync() + { + return this.inner.Failed ? default : this.inner.Result.DisposeAsync(); + } + public async ValueTask MoveNextAsync(ITrace trace, CancellationToken cancellationToken) - { + { TryCatch hasNext = await this.inner.TryAsync(pipelineStage => pipelineStage.MoveNextAsync(trace, cancellationToken)); - bool success = hasNext.Succeeded && hasNext.Result; - if (this.executionState == ExecutionState.OptimisticDirectExecution) - { - bool isPartitionSplitException = hasNext.Succeeded && this.Current.Failed && this.Current.InnerMostException.IsPartitionSplitException(); - if (success && !isPartitionSplitException) - { - this.continuationToken = this.Current.Succeeded ? this.Current.Result.State?.Value : null; - if (this.continuationToken != null) - { - bool requiresDistribution; - if (this.Current.Result.AdditionalHeaders.TryGetValue(HttpConstants.HttpHeaders.RequiresDistribution, out string requiresDistributionHeaderValue)) - { - requiresDistribution = bool.Parse(requiresDistributionHeaderValue); - } - else - { - requiresDistribution = true; - } - - if (this.previousRequiresDistribution.HasValue && this.previousRequiresDistribution != requiresDistribution) - { - // We should never come here as requiresDistribution flag can never switch mid execution. - // Hence, this exception should never be thrown. - throw new InvalidOperationException($"Unexpected switch in {HttpConstants.HttpHeaders.RequiresDistribution} value. Previous value : {this.previousRequiresDistribution} Current value : {requiresDistribution}."); - } - - if (requiresDistribution) - { - // This is where we will unwrap tne continuation token and extract the client distribution plan - // Pipelines to handle client distribution would be generated here + bool success = hasNext.Succeeded && hasNext.Result; + if (this.executionState == ExecutionState.OptimisticDirectExecution) + { + bool isPartitionSplitException = hasNext.Succeeded && this.Current.Failed && this.Current.InnerMostException.IsPartitionSplitException(); + if (success && !isPartitionSplitException) + { + this.continuationToken = this.Current.Succeeded ? this.Current.Result.State?.Value : null; + if (this.continuationToken != null) + { + bool requiresDistribution; + if (this.Current.Result.AdditionalHeaders.TryGetValue(HttpConstants.HttpHeaders.RequiresDistribution, out string requiresDistributionHeaderValue)) + { + requiresDistribution = bool.Parse(requiresDistributionHeaderValue); + } + else + { + requiresDistribution = true; + } + + if (this.previousRequiresDistribution.HasValue && this.previousRequiresDistribution != requiresDistribution) + { + // We should never come here as requiresDistribution flag can never switch mid execution. + // Hence, this exception should never be thrown. + throw new InvalidOperationException($"Unexpected switch in {HttpConstants.HttpHeaders.RequiresDistribution} value. Previous value : {this.previousRequiresDistribution} Current value : {requiresDistribution}."); + } + + if (requiresDistribution) + { + // This is where we will unwrap tne continuation token and extract the client distribution plan + // Pipelines to handle client distribution would be generated here success = await this.SwitchToFallbackPipelineAsync(continuationToken: null, trace, cancellationToken); - } - - this.previousRequiresDistribution = requiresDistribution; - } - } - else if (isPartitionSplitException) - { + } + + this.previousRequiresDistribution = requiresDistribution; + } + } + else if (isPartitionSplitException) + { success = await this.SwitchToFallbackPipelineAsync(continuationToken: UnwrapContinuationToken(this.continuationToken), trace, cancellationToken); - } - } - - return success; - } - - private static CosmosElement UnwrapContinuationToken(CosmosElement continuationToken) - { - if (continuationToken == null) return null; - - CosmosObject cosmosObject = continuationToken as CosmosObject; - CosmosElement backendContinuationToken = cosmosObject[OptimisticDirectExecutionToken]; - Debug.Assert(backendContinuationToken != null); - - return CosmosArray.Create(backendContinuationToken); - } - + } + } + + return success; + } + + private static CosmosElement UnwrapContinuationToken(CosmosElement continuationToken) + { + if (continuationToken == null) return null; + + CosmosObject cosmosObject = continuationToken as CosmosObject; + CosmosElement backendContinuationToken = cosmosObject[OptimisticDirectExecutionToken]; + Debug.Assert(backendContinuationToken != null); + + return CosmosArray.Create(backendContinuationToken); + } + private async Task SwitchToFallbackPipelineAsync(CosmosElement continuationToken, ITrace trace, CancellationToken cancellationToken) - { - Debug.Assert(this.executionState == ExecutionState.OptimisticDirectExecution, "OptimisticDirectExecuteQueryPipelineStage Assert!", "Only OptimisticDirectExecute pipeline can create this fallback pipeline"); - this.executionState = ExecutionState.SpecializedDocumentQueryExecution; - this.inner = continuationToken != null - ? await this.queryPipelineStageFactory(continuationToken) - : await this.queryPipelineStageFactory(null); - - if (this.inner.Failed) - { - return false; - } - + { + Debug.Assert(this.executionState == ExecutionState.OptimisticDirectExecution, "OptimisticDirectExecuteQueryPipelineStage Assert!", "Only OptimisticDirectExecute pipeline can create this fallback pipeline"); + this.executionState = ExecutionState.SpecializedDocumentQueryExecution; + this.inner = continuationToken != null + ? await this.queryPipelineStageFactory(continuationToken) + : await this.queryPipelineStageFactory(null); + + if (this.inner.Failed) + { + return false; + } + return await this.inner.Result.MoveNextAsync(trace, cancellationToken); - } - - public static TryCatch MonadicCreate( - DocumentContainer documentContainer, - CosmosQueryExecutionContextFactory.InputParameters inputParameters, - FeedRangeEpk targetRange, - FallbackQueryPipelineStageFactory fallbackQueryPipelineStageFactory, - CancellationToken cancellationToken) - { - QueryPaginationOptions paginationOptions = new QueryPaginationOptions(pageSizeHint: inputParameters.MaxItemCount, optimisticDirectExecute: true); - TryCatch pipelineStage = OptimisticDirectExecutionQueryPipelineImpl.MonadicCreate( - documentContainer: documentContainer, - sqlQuerySpec: inputParameters.SqlQuerySpec, - targetRange: targetRange, - queryPaginationOptions: paginationOptions, - partitionKey: inputParameters.PartitionKey, - continuationToken: inputParameters.InitialUserContinuationToken, - cancellationToken: cancellationToken); - - if (pipelineStage.Failed) - { - return pipelineStage; - } - - OptimisticDirectExecutionQueryPipelineStage odePipelineStageMonadicCreate = new OptimisticDirectExecutionQueryPipelineStage(pipelineStage, fallbackQueryPipelineStageFactory, inputParameters.InitialUserContinuationToken); - return TryCatch.FromResult(odePipelineStageMonadicCreate); - } - - private sealed class OptimisticDirectExecutionQueryPipelineImpl : IQueryPipelineStage - { - private const int ClientQLCompatibilityLevel = 1; - private readonly QueryPartitionRangePageAsyncEnumerator queryPartitionRangePageAsyncEnumerator; - - private OptimisticDirectExecutionQueryPipelineImpl( - QueryPartitionRangePageAsyncEnumerator queryPartitionRangePageAsyncEnumerator) - { - this.queryPartitionRangePageAsyncEnumerator = queryPartitionRangePageAsyncEnumerator ?? throw new ArgumentNullException(nameof(queryPartitionRangePageAsyncEnumerator)); - } - - public TryCatch Current { get; private set; } - - public ValueTask DisposeAsync() - { - return this.queryPartitionRangePageAsyncEnumerator.DisposeAsync(); - } - + } + + public static TryCatch MonadicCreate( + DocumentContainer documentContainer, + CosmosQueryExecutionContextFactory.InputParameters inputParameters, + FeedRangeEpk targetRange, + ContainerQueryProperties containerQueryProperties, + FallbackQueryPipelineStageFactory fallbackQueryPipelineStageFactory, + CancellationToken cancellationToken) + { + QueryPaginationOptions paginationOptions = new QueryPaginationOptions(pageSizeHint: inputParameters.MaxItemCount, optimisticDirectExecute: true); + TryCatch pipelineStage = OptimisticDirectExecutionQueryPipelineImpl.MonadicCreate( + documentContainer: documentContainer, + sqlQuerySpec: inputParameters.SqlQuerySpec, + targetRange: targetRange, + queryPaginationOptions: paginationOptions, + partitionKey: inputParameters.PartitionKey, + containerQueryProperties: containerQueryProperties, + continuationToken: inputParameters.InitialUserContinuationToken, + cancellationToken: cancellationToken); + + if (pipelineStage.Failed) + { + return pipelineStage; + } + + OptimisticDirectExecutionQueryPipelineStage odePipelineStageMonadicCreate = new OptimisticDirectExecutionQueryPipelineStage(pipelineStage, fallbackQueryPipelineStageFactory, inputParameters.InitialUserContinuationToken); + return TryCatch.FromResult(odePipelineStageMonadicCreate); + } + + private sealed class OptimisticDirectExecutionQueryPipelineImpl : IQueryPipelineStage + { + private const int ClientQLCompatibilityLevel = 1; + private readonly QueryPartitionRangePageAsyncEnumerator queryPartitionRangePageAsyncEnumerator; + + private OptimisticDirectExecutionQueryPipelineImpl( + QueryPartitionRangePageAsyncEnumerator queryPartitionRangePageAsyncEnumerator) + { + this.queryPartitionRangePageAsyncEnumerator = queryPartitionRangePageAsyncEnumerator ?? throw new ArgumentNullException(nameof(queryPartitionRangePageAsyncEnumerator)); + } + + public TryCatch Current { get; private set; } + + public ValueTask DisposeAsync() + { + return this.queryPartitionRangePageAsyncEnumerator.DisposeAsync(); + } + public async ValueTask MoveNextAsync(ITrace trace, CancellationToken cancellationToken) - { - if (trace == null) - { - throw new ArgumentNullException(nameof(trace)); - } - + { + if (trace == null) + { + throw new ArgumentNullException(nameof(trace)); + } + if (!await this.queryPartitionRangePageAsyncEnumerator.MoveNextAsync(trace, cancellationToken)) - { - this.Current = default; - return false; - } - - TryCatch partitionPage = this.queryPartitionRangePageAsyncEnumerator.Current; - if (partitionPage.Failed) - { - this.Current = TryCatch.FromException(partitionPage.Exception); - return true; - } - - QueryPage backendQueryPage = partitionPage.Result; - - QueryState queryState; - if (backendQueryPage.State == null) - { - queryState = null; - } - else - { - QueryState backendQueryState = backendQueryPage.State; - ParallelContinuationToken parallelContinuationToken = new ParallelContinuationToken( - token: (backendQueryState?.Value as CosmosString)?.Value, - range: ((FeedRangeEpk)this.queryPartitionRangePageAsyncEnumerator.FeedRangeState.FeedRange).Range); - - OptimisticDirectExecutionContinuationToken optimisticDirectExecutionContinuationToken = new OptimisticDirectExecutionContinuationToken(parallelContinuationToken); - CosmosElement cosmosElementContinuationToken = OptimisticDirectExecutionContinuationToken.ToCosmosElement(optimisticDirectExecutionContinuationToken); - queryState = new QueryState(cosmosElementContinuationToken); - } - - QueryPage queryPage = new QueryPage( - backendQueryPage.Documents, - backendQueryPage.RequestCharge, - backendQueryPage.ActivityId, - backendQueryPage.CosmosQueryExecutionInfo, - backendQueryPage.DistributionPlanSpec, - disallowContinuationTokenMessage: null, - backendQueryPage.AdditionalHeaders, + { + this.Current = default; + return false; + } + + TryCatch partitionPage = this.queryPartitionRangePageAsyncEnumerator.Current; + if (partitionPage.Failed) + { + this.Current = TryCatch.FromException(partitionPage.Exception); + return true; + } + + QueryPage backendQueryPage = partitionPage.Result; + + QueryState queryState; + if (backendQueryPage.State == null) + { + queryState = null; + } + else + { + QueryState backendQueryState = backendQueryPage.State; + ParallelContinuationToken parallelContinuationToken = new ParallelContinuationToken( + token: (backendQueryState?.Value as CosmosString)?.Value, + range: ((FeedRangeEpk)this.queryPartitionRangePageAsyncEnumerator.FeedRangeState.FeedRange).Range); + + OptimisticDirectExecutionContinuationToken optimisticDirectExecutionContinuationToken = new OptimisticDirectExecutionContinuationToken(parallelContinuationToken); + CosmosElement cosmosElementContinuationToken = OptimisticDirectExecutionContinuationToken.ToCosmosElement(optimisticDirectExecutionContinuationToken); + queryState = new QueryState(cosmosElementContinuationToken); + } + + QueryPage queryPage = new QueryPage( + backendQueryPage.Documents, + backendQueryPage.RequestCharge, + backendQueryPage.ActivityId, + backendQueryPage.CosmosQueryExecutionInfo, + backendQueryPage.DistributionPlanSpec, + disallowContinuationTokenMessage: null, + backendQueryPage.AdditionalHeaders, queryState, backendQueryPage.Streaming); - - this.Current = TryCatch.FromResult(queryPage); - return true; - } - - public static TryCatch MonadicCreate( - IDocumentContainer documentContainer, - SqlQuerySpec sqlQuerySpec, - FeedRangeEpk targetRange, - Cosmos.PartitionKey? partitionKey, - QueryPaginationOptions queryPaginationOptions, - CosmosElement continuationToken, - CancellationToken cancellationToken) - { - if (targetRange == null) - { - throw new ArgumentNullException(nameof(targetRange)); - } - - TryCatch> monadicExtractState; - if (continuationToken == null) - { - FeedRangeState getState = new (targetRange, (QueryState)null); - monadicExtractState = TryCatch>.FromResult(getState); - } - else - { - monadicExtractState = MonadicExtractState(continuationToken, targetRange); - } - - if (monadicExtractState.Failed) - { - return TryCatch.FromException(monadicExtractState.Exception); - } - - SqlQuerySpec updatedSqlQuerySpec = new SqlQuerySpec(sqlQuerySpec.QueryText, sqlQuerySpec.Parameters) - { - ClientQLCompatibilityLevel = ClientQLCompatibilityLevel - }; - - FeedRangeState feedRangeState = monadicExtractState.Result; - QueryPartitionRangePageAsyncEnumerator partitionPageEnumerator = new QueryPartitionRangePageAsyncEnumerator( - documentContainer, - updatedSqlQuerySpec, - feedRangeState, - partitionKey, - queryPaginationOptions); - - OptimisticDirectExecutionQueryPipelineImpl stage = new OptimisticDirectExecutionQueryPipelineImpl(partitionPageEnumerator); - return TryCatch.FromResult(stage); - } - - private static TryCatch> MonadicExtractState( - CosmosElement continuationToken, - FeedRangeEpk range) - { - if (continuationToken == null) - { - throw new ArgumentNullException(nameof(continuationToken)); - } - - TryCatch tryCreateContinuationToken = OptimisticDirectExecutionContinuationToken.TryCreateFromCosmosElement(continuationToken); - if (tryCreateContinuationToken.Failed) - { - return TryCatch>.FromException(tryCreateContinuationToken.Exception); - } - - TryCatch> partitionMappingMonad = PartitionMapper.MonadicGetPartitionMapping( - range, - tryCreateContinuationToken.Result); - - if (partitionMappingMonad.Failed) - { - return TryCatch>.FromException( - partitionMappingMonad.Exception); - } - - PartitionMapper.PartitionMapping partitionMapping = partitionMappingMonad.Result; - - KeyValuePair kvpRange = new KeyValuePair( - partitionMapping.TargetMapping.Keys.First(), - partitionMapping.TargetMapping.Values.First()); - - FeedRangeState feedRangeState = new FeedRangeState(kvpRange.Key, kvpRange.Value?.Token != null ? new QueryState(CosmosString.Create(kvpRange.Value.Token.Token)) : null); - - return TryCatch>.FromResult(feedRangeState); - } - } - } + + this.Current = TryCatch.FromResult(queryPage); + return true; + } + + public static TryCatch MonadicCreate( + IDocumentContainer documentContainer, + SqlQuerySpec sqlQuerySpec, + FeedRangeEpk targetRange, + Cosmos.PartitionKey? partitionKey, + QueryPaginationOptions queryPaginationOptions, + ContainerQueryProperties containerQueryProperties, + CosmosElement continuationToken, + CancellationToken cancellationToken) + { + if (targetRange == null) + { + throw new ArgumentNullException(nameof(targetRange)); + } + + TryCatch> monadicExtractState; + if (continuationToken == null) + { + FeedRangeState getState = new (targetRange, (QueryState)null); + monadicExtractState = TryCatch>.FromResult(getState); + } + else + { + monadicExtractState = MonadicExtractState(continuationToken, targetRange); + } + + if (monadicExtractState.Failed) + { + return TryCatch.FromException(monadicExtractState.Exception); + } + + SqlQuerySpec updatedSqlQuerySpec = new SqlQuerySpec(sqlQuerySpec.QueryText, sqlQuerySpec.Parameters) + { + ClientQLCompatibilityLevel = ClientQLCompatibilityLevel + }; + + FeedRangeState feedRangeState = monadicExtractState.Result; + QueryPartitionRangePageAsyncEnumerator partitionPageEnumerator = new QueryPartitionRangePageAsyncEnumerator( + documentContainer, + updatedSqlQuerySpec, + feedRangeState, + partitionKey, + queryPaginationOptions, + containerQueryProperties); + + OptimisticDirectExecutionQueryPipelineImpl stage = new OptimisticDirectExecutionQueryPipelineImpl(partitionPageEnumerator); + return TryCatch.FromResult(stage); + } + + private static TryCatch> MonadicExtractState( + CosmosElement continuationToken, + FeedRangeEpk range) + { + if (continuationToken == null) + { + throw new ArgumentNullException(nameof(continuationToken)); + } + + TryCatch tryCreateContinuationToken = OptimisticDirectExecutionContinuationToken.TryCreateFromCosmosElement(continuationToken); + if (tryCreateContinuationToken.Failed) + { + return TryCatch>.FromException(tryCreateContinuationToken.Exception); + } + + TryCatch> partitionMappingMonad = PartitionMapper.MonadicGetPartitionMapping( + range, + tryCreateContinuationToken.Result); + + if (partitionMappingMonad.Failed) + { + return TryCatch>.FromException( + partitionMappingMonad.Exception); + } + + PartitionMapper.PartitionMapping partitionMapping = partitionMappingMonad.Result; + + KeyValuePair kvpRange = new KeyValuePair( + partitionMapping.TargetMapping.Keys.First(), + partitionMapping.TargetMapping.Values.First()); + + FeedRangeState feedRangeState = new FeedRangeState(kvpRange.Key, kvpRange.Value?.Token != null ? new QueryState(CosmosString.Create(kvpRange.Value.Token.Token)) : null); + + return TryCatch>.FromResult(feedRangeState); + } + } + } } \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs index 1ae5312c2e..814682ba72 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/PipelineFactory.cs @@ -20,6 +20,7 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Pagination; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Skip; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Take; + using Microsoft.Azure.Cosmos.Query.Core.QueryClient; using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; internal static class PipelineFactory @@ -32,6 +33,7 @@ public static TryCatch MonadicCreate( PartitionKey? partitionKey, QueryInfo queryInfo, QueryPaginationOptions queryPaginationOptions, + ContainerQueryProperties containerQueryProperties, int maxConcurrency, CosmosElement requestContinuationToken) { @@ -87,6 +89,7 @@ public static TryCatch MonadicCreate( targetRanges: targetRanges, queryPaginationOptions: queryPaginationOptions, partitionKey: partitionKey, + containerQueryProperties: containerQueryProperties, prefetchPolicy: prefetchPolicy, maxConcurrency: maxConcurrency, continuationToken: continuationToken); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/SubpartitionTests.TestQueriesOnSplitContainer.xml b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/SubpartitionTests.TestQueriesOnSplitContainer.xml new file mode 100644 index 0000000000..192bfd9a20 --- /dev/null +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/BaselineTest/TestBaseline/SubpartitionTests.TestQueriesOnSplitContainer.xml @@ -0,0 +1,60 @@ + + + + SELECT + + True + + + + + + + + SELECT without ODE + + False + + + + + + \ No newline at end of file diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Microsoft.Azure.Cosmos.Tests.csproj b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Microsoft.Azure.Cosmos.Tests.csproj index 81d0b6797a..c9ad45d555 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Microsoft.Azure.Cosmos.Tests.csproj +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Microsoft.Azure.Cosmos.Tests.csproj @@ -31,6 +31,7 @@ + @@ -349,6 +350,9 @@ + + PreserveNewest + PreserveNewest diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryContainer.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryContainer.cs index bfcfb0c725..4c0a62b72d 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryContainer.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Pagination/InMemoryContainer.cs @@ -8,15 +8,18 @@ namespace Microsoft.Azure.Cosmos.Tests.Pagination using System.Collections; using System.Collections.Generic; using System.Collections.Immutable; + using Debug = System.Diagnostics.Debug; using System.IO; using System.Linq; using System.Reflection; + using System.Text; using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos; using Microsoft.Azure.Cosmos.ChangeFeed.Pagination; using Microsoft.Azure.Cosmos.CosmosElements; using Microsoft.Azure.Cosmos.CosmosElements.Numbers; + using Microsoft.Azure.Cosmos.Handlers; using Microsoft.Azure.Cosmos.Json; using Microsoft.Azure.Cosmos.Pagination; using Microsoft.Azure.Cosmos.Query.Core; @@ -35,6 +38,7 @@ namespace Microsoft.Azure.Cosmos.Tests.Pagination using static Microsoft.Azure.Cosmos.Query.Core.SqlQueryResumeFilter; using ResourceIdentifier = Cosmos.Pagination.ResourceIdentifier; using UInt128 = UInt128; + using Microsoft.Azure.Documents.Routing; // Collection useful for mocking requests and repartitioning (splits / merge). internal class InMemoryContainer : IMonadicDocumentContainer @@ -46,9 +50,13 @@ internal class InMemoryContainer : IMonadicDocumentContainer private PartitionKeyHashRangeDictionary> partitionedChanges; private Dictionary partitionKeyRangeIdToHashRange; private Dictionary cachedPartitionKeyRangeIdToHashRange; + private readonly bool createSplitForMultiHashAtSecondlevel; + private readonly bool resolvePartitionsBasedOnPrefix; public InMemoryContainer( - PartitionKeyDefinition partitionKeyDefinition) + PartitionKeyDefinition partitionKeyDefinition, + bool createSplitForMultiHashAtSecondlevel = false, + bool resolvePartitionsBasedOnPrefix = false) { this.partitionKeyDefinition = partitionKeyDefinition ?? throw new ArgumentNullException(nameof(partitionKeyDefinition)); PartitionKeyHashRange fullRange = new PartitionKeyHashRange(startInclusive: null, endExclusive: new PartitionKeyHash(Cosmos.UInt128.MaxValue)); @@ -66,6 +74,8 @@ public InMemoryContainer( { 0, fullRange } }; this.parentToChildMapping = new Dictionary(); + this.createSplitForMultiHashAtSecondlevel = createSplitForMultiHashAtSecondlevel; + this.resolvePartitionsBasedOnPrefix = resolvePartitionsBasedOnPrefix; } public Task>> MonadicGetFeedRangesAsync( @@ -472,7 +482,10 @@ public virtual Task> MonadicQueryAsync( using (ITrace childTrace = trace.StartChild("Query Transport", TraceComponent.Transport, TraceLevel.Info)) { - TryCatch monadicPartitionKeyRangeId = this.MonadicGetPartitionKeyRangeIdFromFeedRange(feedRangeState.FeedRange); + FeedRange feedRange = this.resolvePartitionsBasedOnPrefix ? + ResolveFeedRangeBasedOnPrefixContainer(feedRangeState.FeedRange, this.partitionKeyDefinition) : + feedRangeState.FeedRange; + TryCatch monadicPartitionKeyRangeId = this.MonadicGetPartitionKeyRangeIdFromFeedRange(feedRange); if (monadicPartitionKeyRangeId.Failed) { return Task.FromResult(TryCatch.FromException(monadicPartitionKeyRangeId.Exception)); @@ -943,6 +956,29 @@ public Task MonadicSplitAsync( return Task.FromResult(TryCatch.FromResult()); } + internal static FeedRange ResolveFeedRangeBasedOnPrefixContainer( + FeedRange feedRange, + PartitionKeyDefinition partitionKeyDefinition) + { + if (feedRange is FeedRangePartitionKey feedRangePartitionKey) + { + if (partitionKeyDefinition != null && partitionKeyDefinition.Kind == PartitionKind.MultiHash + && feedRangePartitionKey.PartitionKey.InternalKey?.Components?.Count < partitionKeyDefinition.Paths?.Count) + { + PartitionKeyHash partitionKeyHash = feedRangePartitionKey.PartitionKey.InternalKey.Components[0] switch + { + null => PartitionKeyHash.V2.HashUndefined(), + StringPartitionKeyComponent stringPartitionKey => PartitionKeyHash.V2.Hash((string)stringPartitionKey.ToObject()), + NumberPartitionKeyComponent numberPartitionKey => PartitionKeyHash.V2.Hash(Number64.ToDouble(numberPartitionKey.Value)), + _ => throw new ArgumentOutOfRangeException(), + }; + feedRange = new FeedRangeEpk(new Documents.Routing.Range(min: partitionKeyHash.Value, max: partitionKeyHash.Value + "-FF", isMinInclusive:true, isMaxInclusive: false)); + } + } + + return feedRange; + } + public Task MonadicMergeAsync( FeedRangeInternal feedRange1, FeedRangeInternal feedRange2, @@ -1337,7 +1373,7 @@ private static bool IsRecordWithinFeedRange( } } - private TryCatch MonadicGetPartitionKeyRangeIdFromFeedRange(FeedRange feedRange) + internal TryCatch MonadicGetPartitionKeyRangeIdFromFeedRange(FeedRange feedRange) { int partitionKeyRangeId; if (feedRange is FeedRangeEpk feedRangeEpk) @@ -1406,12 +1442,118 @@ private TryCatch MonadicGetPartitionKeyRangeIdFromFeedRange(FeedRange feedR private static PartitionKeyHashRange FeedRangeEpkToHashRange(FeedRangeEpk feedRangeEpk) { - PartitionKeyHash? start = feedRangeEpk.Range.Min == string.Empty ? (PartitionKeyHash?)null : PartitionKeyHash.Parse(feedRangeEpk.Range.Min); - PartitionKeyHash? end = feedRangeEpk.Range.Max == string.Empty || feedRangeEpk.Range.Max == "FF" ? (PartitionKeyHash?)null : PartitionKeyHash.Parse(feedRangeEpk.Range.Max); + PartitionKeyHash? start = + feedRangeEpk.Range.Min == string.Empty ? + (PartitionKeyHash?)null : + FromHashString(feedRangeEpk.Range.Min); + PartitionKeyHash? end = + feedRangeEpk.Range.Max == string.Empty || feedRangeEpk.Range.Max == "FF" ? + (PartitionKeyHash?)null : + FromHashString(feedRangeEpk.Range.Max); PartitionKeyHashRange hashRange = new PartitionKeyHashRange(start, end); return hashRange; } + /// + /// Creates a partition key hash from a rangeHash value. Supports if the rangeHash is over a hierarchical partition key. + /// + private static PartitionKeyHash FromHashString(string rangeHash) + { + List hashes = new(); + foreach(string hashComponent in GetHashComponents(rangeHash)) + { + // Hash FF has a special meaning in CosmosDB stack. It represents the max range which needs to be correctly represented for UInt128 parsing. + string value = hashComponent.Equals("FF", StringComparison.OrdinalIgnoreCase) ? + "FF-FF-FF-FF-FF-FF-FF-FF-FF-FF-FF-FF-FF-FF-FF-FF" : + hashComponent; + + bool success = UInt128.TryParse(value, out UInt128 uInt128); + Debug.Assert(success, "InMemoryContainer Assert!", "UInt128 parsing must succeed"); + hashes.Add(uInt128); + } + + return new PartitionKeyHash(hashes.ToArray()); + } + + /// + /// PartitionKeyHash.Parse requires a UInt128 parse-able string which itself requires hyphens to be present between subsequent byte values. + /// The hash values generated by rest of the (test) code may or may not honor this. + /// Furthermore, in case of hierarchical partitions, the hash values are concatenated together and therefore need to be broken into separate segments for parsing each one individually. + /// + /// + /// + private static IEnumerable GetHashComponents(string rangeValue) + { + int start = 0; + + while (start < rangeValue.Length) + { + string uInt128Segment = FixupUInt128(rangeValue, ref start); + yield return uInt128Segment; + } + } + + private static string FixupUInt128(string buffer, ref int start) + { + string result; + if (buffer.Length <= start + 2) + { + result = buffer.Substring(start); + start = buffer.Length; + } + else + { + StringBuilder stringBuilder = new StringBuilder(); + int index = start; + bool done = false; + int count = 0; + while (!done) + { + Debug.Assert(buffer[index] != '-', "InMemoryContainer Assert!", "First character of a chunk cannot be a hyphen"); + stringBuilder.Append(buffer[index]); + index++; + + Debug.Assert(index < buffer.Length, "InMemoryContainer Assert!", "At least 2 characters must be found in a chunk"); + Debug.Assert(buffer[index] != '-', "InMemoryContainer Assert!", "Second character of a chunk cannot be a hyphen"); + stringBuilder.Append(buffer[index]); + index++; + + if ((index < buffer.Length) && (buffer[index] == '-')) + { + index++; + } + + count++; + done = count == 16 || (index >= buffer.Length); + + if (!done) + { + stringBuilder.Append('-'); + } + } + + start = index; + + result = stringBuilder.ToString(); + Debug.Assert( + result.Length >= 2, + "InMemoryContainer Assert!", + "At least 1 byte must be present in hash value"); + Debug.Assert( + result[0] != '-' && result[result.Length - 1] != '-', + "InMemoryContainer Assert!", + "Hyphens should NOT be present at the start of end of the string"); + Debug.Assert( + Enumerable + .Range(1, result.Length - 1) + .All(i => (i % 3 == 2) == (result[i] == '-')), + "InMemoryContainer Assert!", + "Hyphens should be (only) present after every subsequent byte value"); + } + + return result; + } + private static FeedRangeEpk HashRangeToFeedRangeEpk(PartitionKeyHashRange hashRange) { return new FeedRangeEpk( @@ -1441,7 +1583,7 @@ private PartitionKeyHash ComputeMedianSplitPointAmongDocumentsInPKRange(Partitio // For MultiHash Collection, split at top level to ensure documents for top level key exist across partitions // after split - if (medianPkHash.HashValues.Count > 1) + if (medianPkHash.HashValues.Count > 1 && !this.createSplitForMultiHashAtSecondlevel) { return new PartitionKeyHash(medianPkHash.HashValues[0]); } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs index 7e697260f1..e0f94e3e64 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs @@ -1257,18 +1257,18 @@ public override Task ForceRefreshCollectionCacheAsync(string collectionLink, Can public override Task GetCachedContainerQueryPropertiesAsync(string containerLink, Cosmos.PartitionKey? partitionKey, ITrace trace, CancellationToken cancellationToken) { - return Task.FromResult(new ContainerQueryProperties( - "test", - new List> - { - new Range( - PartitionKeyInternal.MinimumInclusiveEffectivePartitionKey, - PartitionKeyInternal.MaximumExclusiveEffectivePartitionKey, - true, - true) - }, - new PartitionKeyDefinition(), - Cosmos.GeospatialType.Geometry)); + return Task.FromResult(new ContainerQueryProperties( + "test", + new List> + { + new Range( + PartitionKeyInternal.MinimumInclusiveEffectivePartitionKey, + PartitionKeyInternal.MaximumExclusiveEffectivePartitionKey, + true, + true) + }, + new PartitionKeyDefinition(), + Cosmos.GeospatialType.Geometry)); } public override Task GetClientDisableOptimisticDirectExecutionAsync() diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/FactoryTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/FactoryTests.cs index 54e5b99899..bc0c4c94a3 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/FactoryTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/FactoryTests.cs @@ -28,6 +28,7 @@ public void TestCreate() sqlQuerySpec: new SqlQuerySpec("SELECT * FROM c"), targetRanges: new List() { FeedRangeEpk.FullRange }, partitionKey: null, + containerQueryProperties: new Cosmos.Query.Core.QueryClient.ContainerQueryProperties(), queryInfo: new QueryInfo() { }, queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: 10), maxConcurrency: 10, diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/FullPipelineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/FullPipelineTests.cs index 68b5b7d9cd..2b68a66ce3 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/FullPipelineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/FullPipelineTests.cs @@ -590,6 +590,7 @@ private static async Task CreatePipelineAsync( partitionKey: null, GetQueryPlan(query), queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: pageSize), + containerQueryProperties: new ContainerQueryProperties(), maxConcurrency: 10, requestContinuationToken: state); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/ParallelCrossPartitionQueryPipelineStageTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/ParallelCrossPartitionQueryPipelineStageTests.cs index e4e06635b3..e3a7554d2f 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/ParallelCrossPartitionQueryPipelineStageTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/ParallelCrossPartitionQueryPipelineStageTests.cs @@ -36,6 +36,7 @@ public void MonadicCreate_NullContinuationToken() targetRanges: new List() { FeedRangeEpk.FullRange }, queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: 10), partitionKey: null, + containerQueryProperties: new Cosmos.Query.Core.QueryClient.ContainerQueryProperties(), maxConcurrency: 10, prefetchPolicy: PrefetchPolicy.PrefetchSinglePage, continuationToken: null); @@ -53,6 +54,7 @@ public void MonadicCreate_NonCosmosArrayContinuationToken() targetRanges: new List() { FeedRangeEpk.FullRange }, queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: 10), partitionKey: null, + containerQueryProperties: new Cosmos.Query.Core.QueryClient.ContainerQueryProperties(), maxConcurrency: 10, prefetchPolicy: PrefetchPolicy.PrefetchSinglePage, continuationToken: CosmosObject.Create(new Dictionary())); @@ -71,6 +73,7 @@ public void MonadicCreate_EmptyArrayContinuationToken() targetRanges: new List() { FeedRangeEpk.FullRange }, queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: 10), partitionKey: null, + containerQueryProperties: new Cosmos.Query.Core.QueryClient.ContainerQueryProperties(), maxConcurrency: 10, prefetchPolicy: PrefetchPolicy.PrefetchSinglePage, continuationToken: CosmosArray.Create(new List())); @@ -89,6 +92,7 @@ public void MonadicCreate_NonParallelContinuationToken() targetRanges: new List() { FeedRangeEpk.FullRange }, queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: 10), partitionKey: null, + containerQueryProperties: new Cosmos.Query.Core.QueryClient.ContainerQueryProperties(), maxConcurrency: 10, prefetchPolicy: PrefetchPolicy.PrefetchSinglePage, continuationToken: CosmosArray.Create(new List() { CosmosString.Create("asdf") })); @@ -111,6 +115,7 @@ public void MonadicCreate_SingleParallelContinuationToken() targetRanges: new List() { new FeedRangeEpk(new Documents.Routing.Range(min: "A", max: "B", isMinInclusive: true, isMaxInclusive: false)) }, queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: 10), partitionKey: null, + containerQueryProperties: new Cosmos.Query.Core.QueryClient.ContainerQueryProperties(), maxConcurrency: 10, prefetchPolicy: PrefetchPolicy.PrefetchSinglePage, continuationToken: CosmosArray.Create(new List() { ParallelContinuationToken.ToCosmosElement(token) })); @@ -140,6 +145,7 @@ public void MonadicCreate_MultipleParallelContinuationToken() }, queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: 10), partitionKey: null, + containerQueryProperties: new Cosmos.Query.Core.QueryClient.ContainerQueryProperties(), maxConcurrency: 10, prefetchPolicy: PrefetchPolicy.PrefetchSinglePage, continuationToken: CosmosArray.Create( @@ -180,6 +186,7 @@ async Task CreatePipelineStateAsync(IDocumentContainer docu cancellationToken: default), queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: 10), partitionKey: null, + containerQueryProperties: new Cosmos.Query.Core.QueryClient.ContainerQueryProperties(), maxConcurrency: 10, prefetchPolicy: aggressivePrefetch ? PrefetchPolicy.PrefetchAll : PrefetchPolicy.PrefetchSinglePage, continuationToken: continuationToken); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPartitionRangePageEnumeratorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPartitionRangePageEnumeratorTests.cs index 2bb6813ff6..7954de30b7 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPartitionRangePageEnumeratorTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPartitionRangePageEnumeratorTests.cs @@ -1,175 +1,178 @@ -namespace Microsoft.Azure.Cosmos.Tests.Query -{ - using System; - using System.Collections.Generic; - using System.Threading; - using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.CosmosElements; - using Microsoft.Azure.Cosmos.Pagination; - using Microsoft.Azure.Cosmos.Query.Core.Monads; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.Parallel; - using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Pagination; - using Microsoft.Azure.Cosmos.Tests.Pagination; - using Microsoft.Azure.Cosmos.Tests.Query.Pipeline; - using Microsoft.Azure.Cosmos.Tracing; - using Microsoft.Azure.Documents; - using Microsoft.VisualStudio.TestTools.UnitTesting; - - [TestClass] - public class QueryPartitionRangePageAsyncEnumeratorTests - { - [TestMethod] - public async Task Test429sAsync() - { - Implementation implementation = new Implementation(); - await implementation.Test429sAsync(false); - } - - [TestMethod] - public async Task Test429sWithContinuationsAsync() - { - Implementation implementation = new Implementation(); - await implementation.Test429sWithContinuationsAsync(false, false); - } - - [TestMethod] - public async Task TestDrainFullyAsync() - { - Implementation implementation = new Implementation(); - await implementation.TestDrainFullyAsync(false); - } - - [TestMethod] - public async Task TestEmptyPages() - { - Implementation implementation = new Implementation(); - await implementation.TestEmptyPages(false); - } - - [TestMethod] - public async Task TestResumingFromStateAsync() - { - Implementation implementation = new Implementation(); - await implementation.TestResumingFromStateAsync(false, false); - } - - [TestMethod] - public async Task TestSplitAsync() - { - Implementation implementation = new Implementation(); - await implementation.TestSplitAsync(); - } - - [TestClass] - private sealed class Implementation : PartitionRangeEnumeratorTests - { - public Implementation() - : base(singlePartition: true) - { - } - - [TestMethod] - public async Task TestSplitAsync() - { - int numItems = 100; - IDocumentContainer documentContainer = await this.CreateDocumentContainerAsync(numItems); - IAsyncEnumerator> enumerator = await this.CreateEnumeratorAsync(documentContainer); - - (HashSet parentIdentifiers, QueryState state) = await this.PartialDrainAsync(enumerator, numIterations: 3); - - // Split the partition - await documentContainer.SplitAsync(new FeedRangePartitionKeyRange("0"), cancellationToken: default); - - // Try To read from the partition that is gone. - await enumerator.MoveNextAsync(); - Assert.IsTrue(enumerator.Current.Failed); - - // Resume on the children using the parent continuaiton token - HashSet childIdentifiers = new HashSet(); - - await documentContainer.RefreshProviderAsync(NoOpTrace.Singleton, cancellationToken: default); - List ranges = await documentContainer.GetFeedRangesAsync( - trace: NoOpTrace.Singleton, - cancellationToken: default); - foreach (FeedRangeEpk range in ranges) - { - IAsyncEnumerable> enumerable = new PartitionRangePageAsyncEnumerable( - feedRangeState: new FeedRangeState(range, state), - (feedRangeState) => new QueryPartitionRangePageAsyncEnumerator( - queryDataSource: documentContainer, - sqlQuerySpec: new Cosmos.Query.Core.SqlQuerySpec("SELECT * FROM c"), - feedRangeState: feedRangeState, - partitionKey: null, - queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: 10)), - trace: NoOpTrace.Singleton); - HashSet resourceIdentifiers = await this.DrainFullyAsync(enumerable); - - childIdentifiers.UnionWith(resourceIdentifiers); - } - - Assert.AreEqual(numItems, parentIdentifiers.Count + childIdentifiers.Count); - } - - public override IReadOnlyList GetRecordsFromPage(QueryPage page) - { - List records = new List(page.Documents.Count); - foreach (CosmosElement element in page.Documents) - { - CosmosObject document = (CosmosObject)element; - ResourceId resourceIdentifier = ResourceId.Parse(((CosmosString)document["_rid"]).Value); - long ticks = Number64.ToLong(((CosmosNumber)document["_ts"]).Value); - string identifer = ((CosmosString)document["id"]).Value; - - records.Add(new Record(resourceIdentifier, new DateTime(ticks: ticks, DateTimeKind.Utc), identifer, document)); - } - - return records; - } - - protected override IAsyncEnumerable> CreateEnumerable( - IDocumentContainer documentContainer, - bool aggressivePrefetch = false, - QueryState state = null) - { - List ranges = documentContainer.GetFeedRangesAsync( - trace: NoOpTrace.Singleton, - cancellationToken: default).Result; - Assert.AreEqual(1, ranges.Count); - return new PartitionRangePageAsyncEnumerable( - feedRangeState: new FeedRangeState(ranges[0], state), - (feedRangeState) => new QueryPartitionRangePageAsyncEnumerator( - queryDataSource: documentContainer, - sqlQuerySpec: new Cosmos.Query.Core.SqlQuerySpec("SELECT * FROM c"), - feedRangeState: feedRangeState, - partitionKey: null, - queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: 10)), - trace: NoOpTrace.Singleton); - } - - protected override Task>> CreateEnumeratorAsync( - IDocumentContainer documentContainer, - bool aggressivePrefetch = false, - bool exercisePrefetch = false, - QueryState state = default, - CancellationToken cancellationToken = default) - { - List ranges = documentContainer.GetFeedRangesAsync( - trace: NoOpTrace.Singleton, - cancellationToken: default).Result; - Assert.AreEqual(1, ranges.Count); - - IAsyncEnumerator> enumerator = new TracingAsyncEnumerator>( - enumerator: new QueryPartitionRangePageAsyncEnumerator( - queryDataSource: documentContainer, - sqlQuerySpec: new Cosmos.Query.Core.SqlQuerySpec("SELECT * FROM c"), - feedRangeState: new FeedRangeState(ranges[0], state), - partitionKey: null, - queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: 10)), +namespace Microsoft.Azure.Cosmos.Tests.Query +{ + using System; + using System.Collections.Generic; + using System.Threading; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Pagination; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.Parallel; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Pagination; + using Microsoft.Azure.Cosmos.Tests.Pagination; + using Microsoft.Azure.Cosmos.Tests.Query.Pipeline; + using Microsoft.Azure.Cosmos.Tracing; + using Microsoft.Azure.Documents; + using Microsoft.VisualStudio.TestTools.UnitTesting; + + [TestClass] + public class QueryPartitionRangePageAsyncEnumeratorTests + { + [TestMethod] + public async Task Test429sAsync() + { + Implementation implementation = new Implementation(); + await implementation.Test429sAsync(false); + } + + [TestMethod] + public async Task Test429sWithContinuationsAsync() + { + Implementation implementation = new Implementation(); + await implementation.Test429sWithContinuationsAsync(false, false); + } + + [TestMethod] + public async Task TestDrainFullyAsync() + { + Implementation implementation = new Implementation(); + await implementation.TestDrainFullyAsync(false); + } + + [TestMethod] + public async Task TestEmptyPages() + { + Implementation implementation = new Implementation(); + await implementation.TestEmptyPages(false); + } + + [TestMethod] + public async Task TestResumingFromStateAsync() + { + Implementation implementation = new Implementation(); + await implementation.TestResumingFromStateAsync(false, false); + } + + [TestMethod] + public async Task TestSplitAsync() + { + Implementation implementation = new Implementation(); + await implementation.TestSplitAsync(); + } + + [TestClass] + private sealed class Implementation : PartitionRangeEnumeratorTests + { + public Implementation() + : base(singlePartition: true) + { + } + + [TestMethod] + public async Task TestSplitAsync() + { + int numItems = 100; + IDocumentContainer documentContainer = await this.CreateDocumentContainerAsync(numItems); + IAsyncEnumerator> enumerator = await this.CreateEnumeratorAsync(documentContainer); + + (HashSet parentIdentifiers, QueryState state) = await this.PartialDrainAsync(enumerator, numIterations: 3); + + // Split the partition + await documentContainer.SplitAsync(new FeedRangePartitionKeyRange("0"), cancellationToken: default); + + // Try To read from the partition that is gone. + await enumerator.MoveNextAsync(); + Assert.IsTrue(enumerator.Current.Failed); + + // Resume on the children using the parent continuaiton token + HashSet childIdentifiers = new HashSet(); + + await documentContainer.RefreshProviderAsync(NoOpTrace.Singleton, cancellationToken: default); + List ranges = await documentContainer.GetFeedRangesAsync( + trace: NoOpTrace.Singleton, + cancellationToken: default); + foreach (FeedRangeEpk range in ranges) + { + IAsyncEnumerable> enumerable = new PartitionRangePageAsyncEnumerable( + feedRangeState: new FeedRangeState(range, state), + (feedRangeState) => new QueryPartitionRangePageAsyncEnumerator( + queryDataSource: documentContainer, + sqlQuerySpec: new Cosmos.Query.Core.SqlQuerySpec("SELECT * FROM c"), + feedRangeState: feedRangeState, + partitionKey: null, + containerQueryProperties: new Cosmos.Query.Core.QueryClient.ContainerQueryProperties(), + queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: 10)), + trace: NoOpTrace.Singleton); + HashSet resourceIdentifiers = await this.DrainFullyAsync(enumerable); + + childIdentifiers.UnionWith(resourceIdentifiers); + } + + Assert.AreEqual(numItems, parentIdentifiers.Count + childIdentifiers.Count); + } + + public override IReadOnlyList GetRecordsFromPage(QueryPage page) + { + List records = new List(page.Documents.Count); + foreach (CosmosElement element in page.Documents) + { + CosmosObject document = (CosmosObject)element; + ResourceId resourceIdentifier = ResourceId.Parse(((CosmosString)document["_rid"]).Value); + long ticks = Number64.ToLong(((CosmosNumber)document["_ts"]).Value); + string identifer = ((CosmosString)document["id"]).Value; + + records.Add(new Record(resourceIdentifier, new DateTime(ticks: ticks, DateTimeKind.Utc), identifer, document)); + } + + return records; + } + + protected override IAsyncEnumerable> CreateEnumerable( + IDocumentContainer documentContainer, + bool aggressivePrefetch = false, + QueryState state = null) + { + List ranges = documentContainer.GetFeedRangesAsync( + trace: NoOpTrace.Singleton, + cancellationToken: default).Result; + Assert.AreEqual(1, ranges.Count); + return new PartitionRangePageAsyncEnumerable( + feedRangeState: new FeedRangeState(ranges[0], state), + (feedRangeState) => new QueryPartitionRangePageAsyncEnumerator( + queryDataSource: documentContainer, + sqlQuerySpec: new Cosmos.Query.Core.SqlQuerySpec("SELECT * FROM c"), + feedRangeState: feedRangeState, + partitionKey: null, + containerQueryProperties: new Cosmos.Query.Core.QueryClient.ContainerQueryProperties(), + queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: 10)), + trace: NoOpTrace.Singleton); + } + + protected override Task>> CreateEnumeratorAsync( + IDocumentContainer documentContainer, + bool aggressivePrefetch = false, + bool exercisePrefetch = false, + QueryState state = default, + CancellationToken cancellationToken = default) + { + List ranges = documentContainer.GetFeedRangesAsync( + trace: NoOpTrace.Singleton, + cancellationToken: default).Result; + Assert.AreEqual(1, ranges.Count); + + IAsyncEnumerator> enumerator = new TracingAsyncEnumerator>( + enumerator: new QueryPartitionRangePageAsyncEnumerator( + queryDataSource: documentContainer, + sqlQuerySpec: new Cosmos.Query.Core.SqlQuerySpec("SELECT * FROM c"), + feedRangeState: new FeedRangeState(ranges[0], state), + partitionKey: null, + containerQueryProperties: new Cosmos.Query.Core.QueryClient.ContainerQueryProperties(), + queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: 10)), trace: NoOpTrace.Singleton, - cancellationToken: cancellationToken); - - return Task.FromResult(enumerator); - } - } - } -} + cancellationToken: default); + + return Task.FromResult(enumerator); + } + } + } +} diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/SubpartitionTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/SubpartitionTests.cs new file mode 100644 index 0000000000..35869ebbc8 --- /dev/null +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/SubpartitionTests.cs @@ -0,0 +1,392 @@ +namespace Microsoft.Azure.Cosmos.Tests.Query +{ + using System; + using System.Collections.Generic; + using System.IO; + using System.Linq; + using System.Threading; + using System.Threading.Tasks; + using System.Xml; + using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Pagination; + using Microsoft.Azure.Cosmos.Query; + using Microsoft.Azure.Cosmos.Query.Core; + using Microsoft.Azure.Cosmos.Query.Core.ExecutionContext; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Pagination; + using Microsoft.Azure.Cosmos.Query.Core.QueryClient; + using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; + using Microsoft.Azure.Cosmos.Routing; + using Microsoft.Azure.Cosmos.Test.BaselineTest; + using Microsoft.Azure.Cosmos.Tests.Pagination; + using Microsoft.Azure.Cosmos.Tracing; + using Microsoft.Azure.Documents; + using Microsoft.VisualStudio.TestTools.UnitTesting; + + [TestClass] + public class SubpartitionTests : BaselineTests + { + private const int DocumentCount = 100; + private const int SplitPartitionKey = 2; + + [TestMethod] + public void TestQueriesOnSplitContainer() + { + List inputs = new List + { + new SubpartitionTestInput("SELECT", query: @"SELECT c.id, c.value2 FROM c", ode: true), + new SubpartitionTestInput("SELECT without ODE", query: @"SELECT c.id, c.value2 FROM c", ode: false), + }; + this.ExecuteTestSuite(inputs); + } + + /// + /// The test is a baseline for mock framework which splits the container at the top level of a hierarchical partition key. + /// After split, it is expected that more than one physical partitions contain data for some value of a top level path of partition key. + /// Please note that this does NOT occur in a single-partition key scenario where all data for a given value of a partition key + /// is contained within single physical partition. + /// This situation is known to create issues, especially while running queries due to inconsistent handling of FeedRangePartitionKey and FeedRangeEpk in the SDK stack. + /// Test framework's behavior in being able to replicate this situation is critical to for ensuring that tests provide sufficient protection against regressions. + /// + [TestMethod] + public async Task VerifyTestFrameworkSupportsPartitionSplit() + { + PartitionKeyDefinition partitionKeyDefinition = CreatePartitionKeyDefinition(); + InMemoryContainer inMemoryContainer = await CreateSplitInMemoryDocumentContainerAsync(DocumentCount, partitionKeyDefinition); + Cosmos.PartitionKey partitionKey = new Cosmos.PartitionKeyBuilder().Add(SplitPartitionKey.ToString()).Build(); + FeedRangePartitionKey feedRangePartitionKey = new FeedRangePartitionKey(partitionKey); + FeedRangeEpk feedRangeEpk = InMemoryContainer.ResolveFeedRangeBasedOnPrefixContainer(feedRangePartitionKey, partitionKeyDefinition) as FeedRangeEpk; + Assert.IsNotNull(feedRangeEpk); + TryCatch pkRangeId = inMemoryContainer.MonadicGetPartitionKeyRangeIdFromFeedRange(feedRangeEpk); + Assert.IsTrue(pkRangeId.Failed, $"Expected to fail for partition key {SplitPartitionKey}"); + Assert.IsTrue(pkRangeId.Exception.InnerException.Message.StartsWith("Epk Range: [B5-D7-B7-26-D6-EA-DB-11-F1-EF-AD-92-12-15-D6-60,B5-D7-B7-26-D6-EA-DB-11-F1-EF-AD-92-12-15-D6-60-FF) is gone."), "Gone exception is expected!"); + } + + public override SubpartitionTestOutput ExecuteTest(SubpartitionTestInput input) + { + IMonadicDocumentContainer monadicDocumentContainer = CreateSplitDocumentContainerAsync(DocumentCount).Result; + DocumentContainer documentContainer = new DocumentContainer(monadicDocumentContainer); + + List documents = new List(); + QueryRequestOptions queryRequestOptions = new QueryRequestOptions() + { + PartitionKey = new PartitionKeyBuilder().Add(SplitPartitionKey.ToString()).Build() + }; + (CosmosQueryExecutionContextFactory.InputParameters inputParameters, CosmosQueryContextCore cosmosQueryContextCore) = + CreateInputParamsAndQueryContext(input, queryRequestOptions); + IQueryPipelineStage queryPipelineStage = CosmosQueryExecutionContextFactory.Create( + documentContainer, + cosmosQueryContextCore, + inputParameters, + NoOpTrace.Singleton); + while (queryPipelineStage.MoveNextAsync(NoOpTrace.Singleton, cancellationToken: default).Result) + { + TryCatch tryGetPage = queryPipelineStage.Current; + + if (tryGetPage.Failed) + { + Assert.Fail("Unexpected error. Gone Exception should not reach till here"); + } + + documents.AddRange(tryGetPage.Result.Documents); + } + + return new SubpartitionTestOutput(documents); + } + + private static Tuple CreateInputParamsAndQueryContext(SubpartitionTestInput input, QueryRequestOptions queryRequestOptions) + { + string query = input.Query; + CosmosElement continuationToken = null; + PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition() + { + Paths = new System.Collections.ObjectModel.Collection() + { + "/id", + "/value1", + "/value2" + }, + Kind = PartitionKind.MultiHash, + Version = PartitionKeyDefinitionVersion.V2, + }; + + queryRequestOptions.EnableOptimisticDirectExecution = input.ODE; + + CosmosSerializerCore serializerCore = new(); + using StreamReader streamReader = new(serializerCore.ToStreamSqlQuerySpec(new SqlQuerySpec(query), Documents.ResourceType.Document)); + string sqlQuerySpecJsonString = streamReader.ReadToEnd(); + + (PartitionedQueryExecutionInfo partitionedQueryExecutionInfo, QueryPartitionProvider queryPartitionProvider) = GetPartitionedQueryExecutionInfoAndPartitionProvider(sqlQuerySpecJsonString, partitionKeyDefinition); + CosmosQueryExecutionContextFactory.InputParameters inputParameters = new CosmosQueryExecutionContextFactory.InputParameters( + sqlQuerySpec: new SqlQuerySpec(query), + initialUserContinuationToken: continuationToken, + initialFeedRange: null, + maxConcurrency: queryRequestOptions.MaxConcurrency, + maxItemCount: queryRequestOptions.MaxItemCount, + maxBufferedItemCount: queryRequestOptions.MaxBufferedItemCount, + partitionKey: queryRequestOptions.PartitionKey, + properties: new Dictionary() { { "x-ms-query-partitionkey-definition", partitionKeyDefinition } }, + partitionedQueryExecutionInfo: null, + executionEnvironment: null, + returnResultsInDeterministicOrder: null, + enableOptimisticDirectExecution: queryRequestOptions.EnableOptimisticDirectExecution, + testInjections: queryRequestOptions.TestSettings); + + string databaseId = "db1234"; + string resourceLink = $"dbs/{databaseId}/colls"; + CosmosQueryContextCore cosmosQueryContextCore = new CosmosQueryContextCore( + client: new TestCosmosQueryClient(queryPartitionProvider), + resourceTypeEnum: Documents.ResourceType.Document, + operationType: Documents.OperationType.Query, + resourceType: typeof(QueryResponseCore), + resourceLink: resourceLink, + isContinuationExpected: true, + allowNonValueAggregateQuery: true, + useSystemPrefix: false, + correlatedActivityId: Guid.NewGuid()); + + return Tuple.Create(inputParameters, cosmosQueryContextCore); + } + + internal static Tuple GetPartitionedQueryExecutionInfoAndPartitionProvider(string querySpecJsonString, PartitionKeyDefinition pkDefinition) + { + QueryPartitionProvider queryPartitionProvider = CreateCustomQueryPartitionProvider(); + TryCatch tryGetQueryPlan = queryPartitionProvider.TryGetPartitionedQueryExecutionInfo( + querySpecJsonString: querySpecJsonString, + partitionKeyDefinition: pkDefinition, + requireFormattableOrderByQuery: true, + isContinuationExpected: true, + allowNonValueAggregateQuery: true, + hasLogicalPartitionKey: false, + allowDCount: true, + useSystemPrefix: false, + geospatialType: Cosmos.GeospatialType.Geography); + + PartitionedQueryExecutionInfo partitionedQueryExecutionInfo = tryGetQueryPlan.Succeeded ? tryGetQueryPlan.Result : throw tryGetQueryPlan.Exception; + return Tuple.Create(partitionedQueryExecutionInfo, queryPartitionProvider); + } + + private static QueryPartitionProvider CreateCustomQueryPartitionProvider() + { + Dictionary queryEngineConfiguration = new Dictionary() + { + {"maxSqlQueryInputLength", 262144}, + {"maxJoinsPerSqlQuery", 5}, + {"maxLogicalAndPerSqlQuery", 2000}, + {"maxLogicalOrPerSqlQuery", 2000}, + {"maxUdfRefPerSqlQuery", 10}, + {"maxInExpressionItemsCount", 16000}, + {"queryMaxGroupByTableCellCount", 500000 }, + {"queryMaxInMemorySortDocumentCount", 500}, + {"maxQueryRequestTimeoutFraction", 0.90}, + {"sqlAllowNonFiniteNumbers", false}, + {"sqlAllowAggregateFunctions", true}, + {"sqlAllowSubQuery", true}, + {"sqlAllowScalarSubQuery", true}, + {"allowNewKeywords", true}, + {"sqlAllowLike", true}, + {"sqlAllowGroupByClause", true}, + {"maxSpatialQueryCells", 12}, + {"spatialMaxGeometryPointCount", 256}, + {"sqlDisableQueryILOptimization", false}, + {"sqlDisableFilterPlanOptimization", false}, + {"clientDisableOptimisticDirectExecution", false} + }; + + return new QueryPartitionProvider(queryEngineConfiguration); + } + + internal static PartitionKeyDefinition CreatePartitionKeyDefinition() + { + PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition() + { + Paths = new System.Collections.ObjectModel.Collection() + { + "/id", + "/value1", + "/value2" + }, + Kind = PartitionKind.MultiHash, + Version = PartitionKeyDefinitionVersion.V2, + }; + + return partitionKeyDefinition; + } + + private static async Task CreateSplitDocumentContainerAsync(int numItems) + { + PartitionKeyDefinition partitionKeyDefinition = CreatePartitionKeyDefinition(); + InMemoryContainer inMemoryContainer = await CreateSplitInMemoryDocumentContainerAsync(numItems, partitionKeyDefinition); + DocumentContainer documentContainer = new DocumentContainer(inMemoryContainer); + return documentContainer; + } + + private static async Task CreateSplitInMemoryDocumentContainerAsync(int numItems, PartitionKeyDefinition partitionKeyDefinition) + { + InMemoryContainer inMemoryContainer = new InMemoryContainer(partitionKeyDefinition, createSplitForMultiHashAtSecondlevel: true, resolvePartitionsBasedOnPrefix: true); + for (int i = 0; i < numItems; i++) + { + CosmosObject item = CosmosObject.Parse($"{{\"id\" : \"{i % 5}\", \"value1\" : \"{Guid.NewGuid()}\", \"value2\" : \"{i}\" }}"); + while (true) + { + TryCatch monadicCreateRecord = await inMemoryContainer.MonadicCreateItemAsync(item, cancellationToken: default); + if (monadicCreateRecord.Succeeded) + { + break; + } + } + } + + await inMemoryContainer.MonadicSplitAsync(FeedRangeEpk.FullRange, cancellationToken: default); + + return inMemoryContainer; + } + internal class TestCosmosQueryClient : CosmosQueryClient + { + private readonly QueryPartitionProvider queryPartitionProvider; + + public TestCosmosQueryClient(QueryPartitionProvider queryPartitionProvider) + { + this.queryPartitionProvider = queryPartitionProvider; + } + + public override Action OnExecuteScalarQueryCallback => throw new NotImplementedException(); + + public override bool BypassQueryParsing() + { + return false; + } + + public override void ClearSessionTokenCache(string collectionFullName) + { + throw new NotImplementedException(); + } + + public override Task> ExecuteItemQueryAsync(string resourceUri, ResourceType resourceType, OperationType operationType, Cosmos.FeedRange feedRange, QueryRequestOptions requestOptions, AdditionalRequestHeaders additionalRequestHeaders, SqlQuerySpec sqlQuerySpec, string continuationToken, int pageSize, ITrace trace, CancellationToken cancellationToken) + { + throw new NotImplementedException(); + } + + public override Task ExecuteQueryPlanRequestAsync(string resourceUri, ResourceType resourceType, OperationType operationType, SqlQuerySpec sqlQuerySpec, Cosmos.PartitionKey? partitionKey, string supportedQueryFeatures, Guid clientQueryCorrelationId, ITrace trace, CancellationToken cancellationToken) + { + return Task.FromResult(new PartitionedQueryExecutionInfo()); + } + + public override Task ForceRefreshCollectionCacheAsync(string collectionLink, CancellationToken cancellationToken) + { + throw new NotImplementedException(); + } + + public override Task GetCachedContainerQueryPropertiesAsync(string containerLink, Cosmos.PartitionKey? partitionKey, ITrace trace, CancellationToken cancellationToken) + { + List hashes = new(); + foreach (Documents.Routing.IPartitionKeyComponent component in partitionKey.Value.InternalKey.Components) + { + PartitionKeyHash partitionKeyHash = component switch + { + null => PartitionKeyHash.V2.HashUndefined(), + Documents.Routing.StringPartitionKeyComponent stringPartitionKey => PartitionKeyHash.V2.Hash((string)stringPartitionKey.ToObject()), + Documents.Routing.NumberPartitionKeyComponent numberPartitionKey => PartitionKeyHash.V2.Hash(Number64.ToDouble(numberPartitionKey.Value)), + _ => throw new ArgumentOutOfRangeException(), + }; + hashes.Add(partitionKeyHash.Value); + } + + string min = string.Join(string.Empty, hashes); + string max = min + "-FF"; + return Task.FromResult(new ContainerQueryProperties( + "test", + new List> + { + new Documents.Routing.Range( + min, + max, + true, + true) + }, + SubpartitionTests.CreatePartitionKeyDefinition(), + Cosmos.GeospatialType.Geometry)); + } + + public override async Task GetClientDisableOptimisticDirectExecutionAsync() + { + return this.queryPartitionProvider.ClientDisableOptimisticDirectExecution; + } + + public override Task> GetTargetPartitionKeyRangeByFeedRangeAsync(string resourceLink, string collectionResourceId, PartitionKeyDefinition partitionKeyDefinition, FeedRangeInternal feedRangeInternal, bool forceRefresh, ITrace trace) + { + throw new NotImplementedException(); + } + + public override Task> GetTargetPartitionKeyRangesAsync(string resourceLink, string collectionResourceId, IReadOnlyList> providedRanges, bool forceRefresh, ITrace trace) + { + return Task.FromResult(new List + { + new PartitionKeyRange() + { + MinInclusive = Documents.Routing.PartitionKeyInternal.MinimumInclusiveEffectivePartitionKey, + MaxExclusive = Documents.Routing.PartitionKeyInternal.MaximumExclusiveEffectivePartitionKey + } + }); + } + + public override Task> TryGetOverlappingRangesAsync(string collectionResourceId, Documents.Routing.Range range, bool forceRefresh = false) + { + throw new NotImplementedException(); + } + + public override async Task> TryGetPartitionedQueryExecutionInfoAsync(SqlQuerySpec sqlQuerySpec, ResourceType resourceType, PartitionKeyDefinition partitionKeyDefinition, bool requireFormattableOrderByQuery, bool isContinuationExpected, bool allowNonValueAggregateQuery, bool hasLogicalPartitionKey, bool allowDCount, bool useSystemPrefix, Cosmos.GeospatialType geospatialType, CancellationToken cancellationToken) + { + CosmosSerializerCore serializerCore = new(); + using StreamReader streamReader = new(serializerCore.ToStreamSqlQuerySpec(sqlQuerySpec, Documents.ResourceType.Document)); + string sqlQuerySpecJsonString = streamReader.ReadToEnd(); + + (PartitionedQueryExecutionInfo partitionedQueryExecutionInfo, QueryPartitionProvider queryPartitionProvider) = OptimisticDirectExecutionQueryBaselineTests.GetPartitionedQueryExecutionInfoAndPartitionProvider(sqlQuerySpecJsonString, partitionKeyDefinition); + return TryCatch.FromResult(partitionedQueryExecutionInfo); + } + } + } + + public class SubpartitionTestInput : BaselineTestInput + { + public SubpartitionTestInput(string description, string query, bool ode) + :base(description) + { + this.Query = query; + this.ODE = ode; + } + + internal string Query { get; } + + internal bool ODE { get; } + + public override void SerializeAsXml(XmlWriter xmlWriter) + { + xmlWriter.WriteElementString("Description", this.Description); + xmlWriter.WriteStartElement("Query"); + xmlWriter.WriteCData(this.Query); + xmlWriter.WriteEndElement(); + xmlWriter.WriteElementString("ODE", this.ODE.ToString()); + } + } + + public class SubpartitionTestOutput : BaselineTestOutput + { + private readonly List documents; + + internal SubpartitionTestOutput(IReadOnlyList documents) + { + this.documents = documents.ToList(); + } + + public override void SerializeAsXml(XmlWriter xmlWriter) + { + xmlWriter.WriteStartElement("Documents"); + string content = string.Join($",{Environment.NewLine}", + this.documents.Select(doc => doc.ToString()).OrderBy(serializedDoc => serializedDoc)); + xmlWriter.WriteCData(content); + xmlWriter.WriteEndElement(); + } + } +} diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Tracing/TraceWriterBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Tracing/TraceWriterBaselineTests.cs index aa2e145850..39082e940d 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Tracing/TraceWriterBaselineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Tracing/TraceWriterBaselineTests.cs @@ -756,6 +756,7 @@ private static IQueryPipelineStage CreatePipeline(IDocumentContainer documentCon partitionKey: null, GetQueryPlan(query), new QueryPaginationOptions(pageSizeHint: pageSize), + containerQueryProperties: new Cosmos.Query.Core.QueryClient.ContainerQueryProperties(), maxConcurrency: 10, requestContinuationToken: state);