protected void AggregateGroupings(IReadOnlyList <CosmosElement> cosmosElements)
{
foreach (CosmosElement result in cosmosElements)
{
// Aggregate the values for all groupings across all continuations.
RewrittenGroupByProjection groupByItem = new RewrittenGroupByProjection(result);
this.groupingTable.AddPayload(groupByItem);
}
}
public void AddPayload(RewrittenGroupByProjection rewrittenGroupByProjection)
{
UInt128 groupByKeysHash = DistinctHash.GetHash(rewrittenGroupByProjection.GroupByItems);
if (!this.table.TryGetValue(groupByKeysHash, out SingleGroupAggregator singleGroupAggregator))
{
singleGroupAggregator = SingleGroupAggregator.TryCreate(
EmptyAggregateOperators,
this.groupByAliasToAggregateType,
this.orderedAliases,
this.hasSelectValue,
continuationToken: null).Result;
this.table[groupByKeysHash] = singleGroupAggregator;
}
CosmosElement payload = rewrittenGroupByProjection.Payload;
singleGroupAggregator.AddValues(payload);
}
public void AddPayload(RewrittenGroupByProjection rewrittenGroupByProjection)
{
// For VALUE queries the payload will be undefined if the field was undefined.
if (rewrittenGroupByProjection.TryGetPayload(out CosmosElement payload))
{
UInt128 groupByKeysHash = DistinctHash.GetHash(rewrittenGroupByProjection.GroupByItems);
if (!this.table.TryGetValue(groupByKeysHash, out SingleGroupAggregator singleGroupAggregator))
{
singleGroupAggregator = SingleGroupAggregator.TryCreate(
EmptyAggregateOperators,
this.groupByAliasToAggregateType,
this.orderedAliases,
this.hasSelectValue,
continuationToken: null).Result;
this.table[groupByKeysHash] = singleGroupAggregator;
}
singleGroupAggregator.AddValues(payload);
}
}
public override async Task <QueryResponseCore> DrainAsync(
int maxElements,
CancellationToken cancellationToken)
{
cancellationToken.ThrowIfCancellationRequested();
// Draining GROUP BY is broken down into two stages:
QueryResponseCore response;
if (!this.Source.IsDone)
{
// Stage 1:
// Drain the groupings fully from all continuation and all partitions
QueryResponseCore sourceResponse = await base.DrainAsync(int.MaxValue, cancellationToken);
if (!sourceResponse.IsSuccess)
{
return(sourceResponse);
}
foreach (CosmosElement result in sourceResponse.CosmosElements)
{
// Aggregate the values for all groupings across all continuations.
RewrittenGroupByProjection groupByItem = new RewrittenGroupByProjection(result);
UInt192 groupByKeysHash = DistinctHash.GetHash(groupByItem.GroupByItems);
if (!this.groupingTable.TryGetValue(groupByKeysHash, out SingleGroupAggregator singleGroupAggregator))
{
singleGroupAggregator = SingleGroupAggregator.Create(
this.cosmosQueryClient,
EmptyAggregateOperators,
this.groupByAliasToAggregateType,
this.orderedAliases,
this.hasSelectValue,
continuationToken: null);
this.groupingTable[groupByKeysHash] = singleGroupAggregator;
}
CosmosElement payload = groupByItem.Payload;
singleGroupAggregator.AddValues(payload);
}
// We need to give empty pages until the results are fully drained.
response = QueryResponseCore.CreateSuccess(
result: EmptyResults,
continuationToken: null,
disallowContinuationTokenMessage: GroupByDocumentQueryExecutionComponent.ContinuationTokenNotSupportedWithGroupBy,
activityId: sourceResponse.ActivityId,
requestCharge: sourceResponse.RequestCharge,
diagnostics: sourceResponse.Diagnostics,
responseLengthBytes: sourceResponse.ResponseLengthBytes);
this.isDone = false;
}
else
{
// Stage 2:
// Emit the results from the grouping table page by page
IEnumerable <SingleGroupAggregator> groupByValuesList = this.groupingTable
.OrderBy(kvp => kvp.Key)
.Skip(this.numPagesDrainedFromGroupingTable * maxElements)
.Take(maxElements)
.Select(kvp => kvp.Value);
List <CosmosElement> results = new List <CosmosElement>();
foreach (SingleGroupAggregator groupByValues in groupByValuesList)
{
results.Add(groupByValues.GetResult());
}
response = QueryResponseCore.CreateSuccess(
result: results,
continuationToken: null,
disallowContinuationTokenMessage: GroupByDocumentQueryExecutionComponent.ContinuationTokenNotSupportedWithGroupBy,
activityId: null,
requestCharge: 0,
diagnostics: QueryResponseCore.EmptyDiagnostics,
responseLengthBytes: 0);
this.numPagesDrainedFromGroupingTable++;
if (this.numPagesDrainedFromGroupingTable * maxElements >= this.groupingTable.Count)
{
this.isDone = true;
}
}
return(response);
}
public override async Task <QueryResponse> DrainAsync(
int maxElements,
CancellationToken cancellationToken)
{
cancellationToken.ThrowIfCancellationRequested();
// Draining GROUP BY is broken down into two stages:
QueryResponse response;
if (!this.Source.IsDone)
{
// Stage 1:
// Drain the groupings fully from all continuation and all partitions
QueryResponse sourceResponse = await base.DrainAsync(int.MaxValue, cancellationToken);
if (!sourceResponse.IsSuccessStatusCode)
{
return(sourceResponse);
}
this.containerRid = sourceResponse.QueryHeaders.ContainerRid;
foreach (CosmosElement result in sourceResponse.CosmosElements)
{
// Aggregate the values for all groupings across all continuations.
RewrittenGroupByProjection groupByItem = new RewrittenGroupByProjection(result);
this.distinctMap.Add(groupByItem.GroupByItems, out UInt192? groupByKeysHash);
if (!groupByKeysHash.HasValue)
{
throw new InvalidOperationException("hash invariant was broken");
}
if (!this.groupingTable.TryGetValue(groupByKeysHash.Value, out SingleGroupAggregator singleGroupAggregator))
{
singleGroupAggregator = SingleGroupAggregator.Create(
EmptyAggregateOperators,
this.groupByAliasToAggregateType,
this.hasSelectValue);
this.groupingTable[groupByKeysHash.Value] = singleGroupAggregator;
}
CosmosElement payload = groupByItem.Payload;
singleGroupAggregator.AddValues(payload);
}
string updatedContinuationToken = null;
// We need to give empty pages until the results are fully drained.
response = QueryResponse.CreateSuccess(
EmptyResults,
EmptyResults.Count,
sourceResponse.ResponseLengthBytes,
sourceResponse.QueryHeaders.CloneKnownProperties(updatedContinuationToken, GroupByDocumentQueryExecutionComponent.ContinuationTokenNotSupportedWithGroupBy));
this.isDone = false;
}
else
{
// Stage 2:
// Emit the results from the grouping table page by page
IEnumerable <SingleGroupAggregator> groupByValuesList = this.groupingTable
.OrderBy(kvp => kvp.Key)
.Skip(this.numPagesDrainedFromGroupingTable * maxElements)
.Take(maxElements)
.Select(kvp => kvp.Value);
List <CosmosElement> results = new List <CosmosElement>();
foreach (SingleGroupAggregator groupByValues in groupByValuesList)
{
results.Add(groupByValues.GetResult());
}
response = QueryResponse.CreateSuccess(
results,
results.Count,
responseLengthBytes: 0,
responseHeaders: new CosmosQueryResponseMessageHeaders(
continauationToken: null,
disallowContinuationTokenMessage: GroupByDocumentQueryExecutionComponent.ContinuationTokenNotSupportedWithGroupBy,
resourceType: Documents.ResourceType.Document,
containerRid: this.containerRid));
this.numPagesDrainedFromGroupingTable++;
if (this.numPagesDrainedFromGroupingTable * maxElements >= this.groupingTable.Count)
{
this.isDone = true;
}
}
return(response);
}
