Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[Internal] Query: Adds single physical partition check for OptimisticDirectExecution queries #3699

Merged
merged 9 commits into from
Feb 17, 2023
Merged

[Internal] Query: Adds single physical partition check for OptimisticDirectExecution queries #3699

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
132ab7b
Added single physical partition check for Ode queries. Updated test i…
akotalwar Feb 7, 2023
bc23bfa
Refactored emulator tests to have all the test cases at the top of th…
akotalwar Feb 10, 2023
02bda3c
Updated TryGetTargetRangeOptimisticDirectExecutionAsync()
akotalwar Feb 13, 2023
94902df
Uodated logic on how many times CreateIngestQueryDeleteAsync() gets c…
akotalwar Feb 15, 2023
f6986e1
Added debug asserts for partitionKeyDefinition
akotalwar Feb 15, 2023
27ea927
Added pageSizeOptions parameter in CreateInput() in EmulatorTests
akotalwar Feb 16, 2023
7762bae
Fixed comments
akotalwar Feb 16, 2023
e4eaf8c
Merge branch 'master' into users/akotalwar/SinglePhysicalPartitionCheck
akotalwar Feb 17, 2023
54a257d
Merge branch 'master' into users/akotalwar/SinglePhysicalPartitionCheck
akotalwar Feb 17, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
52 changes: 24 additions & 28 deletions Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -753,54 +753,50 @@ private static Documents.PartitionKeyDefinition GetPartitionKeyDefinition(InputP
ContainerQueryProperties containerQueryProperties,
ITrace trace)
{
if (!inputParameters.EnableOptimisticDirectExecution) return null;

// case 1: Is query going to a single partition
bool hasPartitionKey = inputParameters.PartitionKey.HasValue
&& inputParameters.PartitionKey != PartitionKey.Null
&& inputParameters.PartitionKey != PartitionKey.None;

// case 2: does query execution plan have a single query range
bool hasQueryRanges = partitionedQueryExecutionInfo != null
&& partitionedQueryExecutionInfo.QueryRanges.Count == 1
&& partitionedQueryExecutionInfo.QueryRanges[0].IsSingleValue;

if (!hasPartitionKey && !hasQueryRanges) return null;

//TODO: does collection have only one physical partition

List<Documents.PartitionKeyRange> targetRanges = new List<Documents.PartitionKeyRange>();

if (!inputParameters.EnableOptimisticDirectExecution) return null;

List<Documents.PartitionKeyRange> targetRanges = new List<Documents.PartitionKeyRange>();
akotalwar marked this conversation as resolved.
Show resolved Hide resolved
if (partitionedQueryExecutionInfo != null)
{
targetRanges = await CosmosQueryExecutionContextFactory.GetTargetPartitionKeyRangesAsync(
cosmosQueryContext.QueryClient,
cosmosQueryContext.ResourceLink,
partitionedQueryExecutionInfo,
containerQueryProperties,
inputParameters.Properties,
inputParameters.InitialFeedRange,
trace);
cosmosQueryContext.QueryClient,
cosmosQueryContext.ResourceLink,
partitionedQueryExecutionInfo,
containerQueryProperties,
inputParameters.Properties,
inputParameters.InitialFeedRange,
trace);
}
else
{
Documents.PartitionKeyDefinition partitionKeyDefinition = GetPartitionKeyDefinition(inputParameters, containerQueryProperties);
if (partitionKeyDefinition != null && containerQueryProperties.ResourceId != null && inputParameters.PartitionKey != null)
if (containerQueryProperties.ResourceId != null)
{
akotalwar marked this conversation as resolved.
Show resolved Hide resolved
targetRanges = await cosmosQueryContext.QueryClient.GetTargetPartitionKeyRangesByEpkStringAsync(
if (inputParameters.PartitionKey != null)
{
targetRanges = await cosmosQueryContext.QueryClient.GetTargetPartitionKeyRangesByEpkStringAsync(
cosmosQueryContext.ResourceLink,
containerQueryProperties.ResourceId,
inputParameters.PartitionKey.Value.InternalKey.GetEffectivePartitionKeyString(partitionKeyDefinition),
forceRefresh: false,
trace);
akotalwar marked this conversation as resolved.
Show resolved Hide resolved
}
else
{
targetRanges = await cosmosQueryContext.QueryClient.GetTargetPartitionKeyRangesAsync(
cosmosQueryContext.ResourceLink,
containerQueryProperties.ResourceId,
new List<Documents.Routing.Range<string>> { FeedRangeEpk.FullRange.Range },
akotalwar marked this conversation as resolved.
Show resolved Hide resolved
forceRefresh: false,
trace);
akotalwar marked this conversation as resolved.
Show resolved Hide resolved
}
}
}
akotalwar marked this conversation as resolved.
Show resolved Hide resolved

if (targetRanges.Count == 1)
{
return targetRanges.Single();
}

return null;
}

Expand Down
159 changes: 61 additions & 98 deletions ...s/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/OptimisticDirectExecutionQueryTests.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,30 +40,35 @@ public async Task TestOptimisticDirectExecutionQueries()
NullField = nullField,
};

await this.CreateIngestQueryDeleteAsync<SinglePartitionWithContinuationsArgs>(
ConnectionModes.Direct | ConnectionModes.Gateway,
CollectionTypes.SinglePartition | CollectionTypes.MultiPartition,
documents,
RunTests,
args,
"/" + partitionKey);
CollectionTypes[] pipelineTypes = new CollectionTypes[] { CollectionTypes.SinglePartition, CollectionTypes.MultiPartition };
foreach (CollectionTypes type in pipelineTypes)
{
await this.CreateIngestQueryDeleteAsync<SinglePartitionWithContinuationsArgs>(
ConnectionModes.Direct | ConnectionModes.Gateway,
type,
documents,
(container, documents, args) => RunTests(container, documents, args, type),
args,
"/" + partitionKey);
}
}

private static async Task RunTests(Container container, IReadOnlyList<CosmosObject> documents, SinglePartitionWithContinuationsArgs args)
private static async Task RunTests(Container container, IReadOnlyList<CosmosObject> documents, SinglePartitionWithContinuationsArgs args, CollectionTypes pipelineType)
{
await TestPositiveOptimisticDirectExecutionOutput(container, args);
await TestPositiveOptimisticDirectExecutionOutput(container, args, pipelineType);
await TestNegativeOptimisticDirectExecutionOutput(container);
}

private static async Task TestPositiveOptimisticDirectExecutionOutput(
Container container,
SinglePartitionWithContinuationsArgs args)
SinglePartitionWithContinuationsArgs args,
CollectionTypes pipelineType)
{
int documentCount = args.NumberOfDocuments;
string partitionKey = args.PartitionKey;
string numberField = args.NumberField;
string nullField = args.NullField;

QueryRequestOptions feedOptions = new QueryRequestOptions
{
MaxItemCount = -1,
Expand All @@ -78,98 +83,54 @@ private static async Task TestPositiveOptimisticDirectExecutionOutput(

Assert.AreEqual(null, responseWithEmptyContinuationExpected.ContinuationToken);

List<string> queries = new List<string>
{
$"SELECT TOP 5 VALUE r.numberField FROM r ORDER BY r.{partitionKey}",
$"SELECT VALUE r.numberField FROM r",
$"SELECT VALUE r.numberField FROM r",
$"SELECT TOP 4 VALUE r.numberField FROM r ORDER BY r.{numberField}",
$"SELECT TOP 3 VALUE r.numberField FROM r WHERE r.{numberField} BETWEEN 0 AND {documentCount} ORDER BY r.{numberField} DESC",
$"SELECT VALUE r.numberField FROM r WHERE r.{numberField} BETWEEN 0 AND {documentCount} OFFSET 1 LIMIT 1",
$"SELECT DISTINCT VALUE r.{numberField} FROM r ORDER BY r.{numberField}",
$"SELECT TOP 3 VALUE r.numberField FROM r WHERE r.{numberField} BETWEEN 0 AND {documentCount} ORDER BY r.{numberField} DESC",
$"SELECT TOP 4 VALUE r.numberField FROM r ORDER BY r.{numberField}",
$"SELECT VALUE r.numberField FROM r",
};

List<List<long>> results = new List<List<long>>
{
new List<long> { 0, 1, 2, 3, 4 },
new List<long> { 0, 1, 2, 3, 4, 5, 6, 7 },
new List<long> { 0, 1, 2, 3, 4, 5, 6, 7 },
new List<long> { 0, 1, 2, 3},
new List<long> { 7, 6, 5},
new List<long> { 1},
new List<long> { 0, 1, 2, 3, 4, 5, 6, 7},
new List<long> { 7, 6, 5},
new List<long> { 0, 1, 2, 3},
new List<long> { 0, 1, 2, 3, 4, 5, 6, 7 },
};

List<string> partitionKeys = new List<string>
{
"/value",
null,
"/value",
"/value",
"/value",
null,
null,
"/value",
"/value",
"/value",
};

List<TestInjections.PipelineType> expectedPipelineType = new List<TestInjections.PipelineType>
List<QueryResultsAndPipelineType> queryAndResults = new List<QueryResultsAndPipelineType>()
{
TestInjections.PipelineType.OptimisticDirectExecution,
TestInjections.PipelineType.Passthrough,
TestInjections.PipelineType.OptimisticDirectExecution,
TestInjections.PipelineType.OptimisticDirectExecution,
TestInjections.PipelineType.OptimisticDirectExecution,
TestInjections.PipelineType.Specialized,
TestInjections.PipelineType.Specialized,
TestInjections.PipelineType.Specialized,
TestInjections.PipelineType.Specialized,
TestInjections.PipelineType.Passthrough,
// Tests for bool enableOptimisticDirectExecution
new QueryResultsAndPipelineType { Query = $"SELECT TOP 5 VALUE r.numberField FROM r ORDER BY r.{partitionKey}", Result = new List<long> { 0, 1, 2, 3, 4 }, PartitionKey = "/value", Partition = CollectionTypes.SinglePartition, EnableOptimisticDirectExecution = true, ExpectedPipelineType = TestInjections.PipelineType.OptimisticDirectExecution},
akotalwar marked this conversation as resolved.
Show resolved Hide resolved
new QueryResultsAndPipelineType { Query = $"SELECT TOP 5 VALUE r.numberField FROM r ORDER BY r.{partitionKey}", Result = new List<long> { 0, 1, 2, 3, 4 }, PartitionKey = "/value", Partition = CollectionTypes.SinglePartition, EnableOptimisticDirectExecution = false, ExpectedPipelineType = TestInjections.PipelineType.Specialized},

// Simple Query
new QueryResultsAndPipelineType { Query = $"SELECT VALUE r.numberField FROM r", Result = new List<long> { 0, 1, 2, 3, 4, 5, 6, 7 }, PartitionKey = "/value", Partition = CollectionTypes.SinglePartition, EnableOptimisticDirectExecution = true, ExpectedPipelineType = TestInjections.PipelineType.OptimisticDirectExecution},
akotalwar marked this conversation as resolved.
Show resolved Hide resolved
new QueryResultsAndPipelineType { Query = $"SELECT VALUE r.numberField FROM r", Result = new List<long> { 0, 1, 2, 3, 4, 5, 6, 7 }, PartitionKey = null, Partition = CollectionTypes.SinglePartition, EnableOptimisticDirectExecution = true, ExpectedPipelineType = TestInjections.PipelineType.OptimisticDirectExecution},
new QueryResultsAndPipelineType { Query = $"SELECT VALUE r.numberField FROM r", Result = new List<long> { 0, 1, 2, 3, 4, 5, 6, 7 }, PartitionKey = "/value", Partition = CollectionTypes.MultiPartition, EnableOptimisticDirectExecution = true, ExpectedPipelineType = TestInjections.PipelineType.OptimisticDirectExecution},
new QueryResultsAndPipelineType { Query = $"SELECT VALUE r.numberField FROM r", Result = new List<long> { 0, 1, 2, 3, 4, 5, 6, 7 }, PartitionKey = null, Partition = CollectionTypes.MultiPartition, EnableOptimisticDirectExecution = true, ExpectedPipelineType = TestInjections.PipelineType.Passthrough},

// DISTINCT with ORDER BY
new QueryResultsAndPipelineType { Query = $"SELECT DISTINCT VALUE r.{numberField} FROM r ORDER BY r.{numberField} DESC", Result = new List<long> { 7, 6, 5, 4, 3, 2, 1, 0 }, PartitionKey = "/value", Partition = CollectionTypes.SinglePartition, EnableOptimisticDirectExecution = true, ExpectedPipelineType = TestInjections.PipelineType.OptimisticDirectExecution},
new QueryResultsAndPipelineType { Query = $"SELECT DISTINCT VALUE r.{numberField} FROM r ORDER BY r.{numberField} DESC", Result = new List<long> { 7, 6, 5, 4, 3, 2, 1, 0 }, PartitionKey = null, Partition = CollectionTypes.SinglePartition, EnableOptimisticDirectExecution = true, ExpectedPipelineType = TestInjections.PipelineType.OptimisticDirectExecution},
new QueryResultsAndPipelineType { Query = $"SELECT DISTINCT VALUE r.{numberField} FROM r ORDER BY r.{numberField} DESC", Result = new List<long> { 7, 6, 5, 4, 3, 2, 1, 0 }, PartitionKey = "/value", Partition = CollectionTypes.MultiPartition, EnableOptimisticDirectExecution = true, ExpectedPipelineType = TestInjections.PipelineType.OptimisticDirectExecution},
new QueryResultsAndPipelineType { Query = $"SELECT DISTINCT VALUE r.{numberField} FROM r ORDER BY r.{numberField} DESC", Result = new List<long> { 7, 6, 5, 4, 3, 2, 1, 0 }, PartitionKey = null, Partition = CollectionTypes.MultiPartition, EnableOptimisticDirectExecution = true, ExpectedPipelineType = TestInjections.PipelineType.Specialized},

// TOP with GROUP BY
new QueryResultsAndPipelineType { Query = $"SELECT TOP 5 VALUE r.{numberField} FROM r GROUP BY r.{numberField}", Result = new List<long> { 0, 1, 2, 3, 4 }, PartitionKey = "/value", Partition = CollectionTypes.SinglePartition, EnableOptimisticDirectExecution = true, ExpectedPipelineType = TestInjections.PipelineType.OptimisticDirectExecution},
new QueryResultsAndPipelineType { Query = $"SELECT TOP 5 VALUE r.{numberField} FROM r GROUP BY r.{numberField}", Result = new List<long> { 0, 1, 2, 3, 4 }, PartitionKey = null, Partition = CollectionTypes.SinglePartition, EnableOptimisticDirectExecution = true, ExpectedPipelineType = TestInjections.PipelineType.OptimisticDirectExecution},
new QueryResultsAndPipelineType { Query = $"SELECT TOP 5 VALUE r.{numberField} FROM r GROUP BY r.{numberField}", Result = new List<long> { 0, 1, 2, 3, 4 }, PartitionKey = "/value", Partition = CollectionTypes.MultiPartition, EnableOptimisticDirectExecution = true, ExpectedPipelineType = TestInjections.PipelineType.OptimisticDirectExecution},
new QueryResultsAndPipelineType { Query = $"SELECT TOP 5 VALUE r.{numberField} FROM r GROUP BY r.{numberField}", Result = new List<long> { 0, 1, 2, 3, 4 }, PartitionKey = null, Partition = CollectionTypes.MultiPartition, EnableOptimisticDirectExecution = true, ExpectedPipelineType = TestInjections.PipelineType.Specialized},

// OFFSET LIMIT with WHERE and BETWEEN
new QueryResultsAndPipelineType { Query = $"SELECT VALUE r.numberField FROM r WHERE r.{numberField} BETWEEN 0 AND {documentCount} OFFSET 1 LIMIT 1", Result = new List<long> { 1 }, PartitionKey = "/value", Partition = CollectionTypes.SinglePartition, EnableOptimisticDirectExecution = true, ExpectedPipelineType = TestInjections.PipelineType.OptimisticDirectExecution},
new QueryResultsAndPipelineType { Query = $"SELECT VALUE r.numberField FROM r WHERE r.{numberField} BETWEEN 0 AND {documentCount} OFFSET 1 LIMIT 1", Result = new List<long> { 1 }, PartitionKey = null, Partition = CollectionTypes.SinglePartition, EnableOptimisticDirectExecution = true, ExpectedPipelineType = TestInjections.PipelineType.OptimisticDirectExecution},
new QueryResultsAndPipelineType { Query = $"SELECT VALUE r.numberField FROM r WHERE r.{numberField} BETWEEN 0 AND {documentCount} OFFSET 1 LIMIT 1", Result = new List<long> { 1 }, PartitionKey = "/value", Partition = CollectionTypes.MultiPartition, EnableOptimisticDirectExecution = true, ExpectedPipelineType = TestInjections.PipelineType.OptimisticDirectExecution},
new QueryResultsAndPipelineType { Query = $"SELECT VALUE r.numberField FROM r WHERE r.{numberField} BETWEEN 0 AND {documentCount} OFFSET 1 LIMIT 1", Result = new List<long> { 1 }, PartitionKey = null, Partition = CollectionTypes.MultiPartition, EnableOptimisticDirectExecution = true, ExpectedPipelineType = TestInjections.PipelineType.Specialized},
};
akotalwar marked this conversation as resolved.
Show resolved Hide resolved

List<bool> enabledOptimisticDirectExecution = new List<bool>
{
true,
true,
true,
true,
true,
true,
true,
false,
false,
false,
};

List<QueryResultsAndPipelineType> queryAndResults = new List<QueryResultsAndPipelineType>();

for (int i = 0; i < queries.Count(); i++)
{
QueryResultsAndPipelineType queryAndResult = new QueryResultsAndPipelineType()
{
Query = queries[i],
Result = results[i],
PartitionKey = partitionKeys[i],
ExpectedPipelineType = expectedPipelineType[i],
EnableOptimisticDirectExecution = enabledOptimisticDirectExecution[i],
};

queryAndResults.Add(queryAndResult);
}

int[] pageSizeOptions = new[] { -1, 1, 2, 10, 100 };

for (int i = 0; i < pageSizeOptions.Length; i++)
akotalwar marked this conversation as resolved.
Show resolved Hide resolved
{
for(int j = 0; j < queryAndResults.Count(); j++)
{
if (pipelineType != queryAndResults[j].Partition)
{
continue;
}

// Added check because "Continuation token is not supported for queries with GROUP BY."
if (queryAndResults[j].Query.Contains("GROUP BY"))
akotalwar marked this conversation as resolved.
Show resolved Hide resolved
{
if (pipelineType == CollectionTypes.MultiPartition) continue;
if (pageSizeOptions[i] != -1) continue;
}

feedOptions = new QueryRequestOptions
{
MaxItemCount = pageSizeOptions[i],
Expand All @@ -186,15 +147,16 @@ private static async Task TestPositiveOptimisticDirectExecutionOutput(
feedOptions);

long[] actual = items.Cast<CosmosNumber>().Select(x => Number64.ToLong(x.Value)).ToArray();

Assert.IsTrue(queryAndResults[j].Result.SequenceEqual(actual));

if (queryAndResults[j].EnableOptimisticDirectExecution)
{
Assert.AreEqual(queryAndResults[j].ExpectedPipelineType, feedOptions.TestSettings.Stats.PipelineType.Value);
akotalwar marked this conversation as resolved.
Show resolved Hide resolved
}
else
{ // test if pipeline is called if TestInjection.EnableOptimisticDirectExecution is false
{
// test if Ode is called if TestInjection.EnableOptimisticDirectExecution is false
Assert.AreNotEqual(TestInjections.PipelineType.OptimisticDirectExecution, feedOptions.TestSettings.Stats.PipelineType.Value);
}
}
Expand Down Expand Up @@ -254,8 +216,9 @@ private struct QueryResultsAndPipelineType
public string Query { get; set; }
public List<long> Result { get; set; }
akotalwar marked this conversation as resolved.
Show resolved Hide resolved
public string PartitionKey { get; set; }
public TestInjections.PipelineType ExpectedPipelineType { get; set; }
public CollectionTypes Partition { get; set; }
public bool EnableOptimisticDirectExecution { get; set; }
public TestInjections.PipelineType ExpectedPipelineType { get; set; }
}
}
}
Loading