diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/AggregateDocumentQueryExecutionContext.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/AggregateDocumentQueryExecutionContext.java index 6c253ba3857d..bcb43e6a7ba6 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/AggregateDocumentQueryExecutionContext.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/AggregateDocumentQueryExecutionContext.java @@ -20,7 +20,7 @@ import java.util.Map; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; -import java.util.function.Function; +import java.util.function.BiFunction; public class AggregateDocumentQueryExecutionContext implements IDocumentQueryExecutionComponent{ @@ -57,8 +57,8 @@ public Flux> drainAsync(int maxPageSize) { .map( superList -> { double requestCharge = 0; - List aggregateResults = new ArrayList(); - HashMap headers = new HashMap(); + List aggregateResults = new ArrayList<>(); + HashMap headers = new HashMap<>(); for(FeedResponse page : superList) { @@ -104,23 +104,22 @@ public Flux> drainAsync(int maxPageSize) { } public static Flux> createAsync( - Function>> createSourceComponentFunction, + BiFunction, Flux>> createSourceComponentFunction, Collection aggregates, Map groupByAliasToAggregateType, List groupByAliases, boolean hasSelectValue, - String continuationToken) { + String continuationToken, + PipelinedDocumentQueryParams documentQueryParams) { return createSourceComponentFunction - .apply(continuationToken) - .map(component -> { - return new AggregateDocumentQueryExecutionContext(component, - new ArrayList<>(aggregates), - groupByAliasToAggregateType, - groupByAliases, - hasSelectValue, - continuationToken); - }); + .apply(continuationToken, documentQueryParams) + .map(component -> new AggregateDocumentQueryExecutionContext(component, + new ArrayList<>(aggregates), + groupByAliasToAggregateType, + groupByAliases, + hasSelectValue, + continuationToken)); } public IDocumentQueryExecutionComponent getComponent() { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DistinctDocumentQueryExecutionContext.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DistinctDocumentQueryExecutionContext.java index fa5ef23280be..39e61580b298 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DistinctDocumentQueryExecutionContext.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DistinctDocumentQueryExecutionContext.java @@ -18,7 +18,7 @@ import java.util.List; import java.util.Map; import java.util.concurrent.atomic.AtomicReference; -import java.util.function.Function; +import java.util.function.BiFunction; public class DistinctDocumentQueryExecutionContext implements IDocumentQueryExecutionComponent { private final IDocumentQueryExecutionComponent component; @@ -43,9 +43,10 @@ private DistinctDocumentQueryExecutionContext( } public static Flux> createAsync( - Function>> createSourceComponentFunction, + BiFunction, Flux>> createSourceComponentFunction, DistinctQueryType distinctQueryType, - String continuationToken) { + String continuationToken, + PipelinedDocumentQueryParams documentQueryParams) { Utils.ValueHolder outDistinctcontinuationtoken = new Utils.ValueHolder<>(); DistinctContinuationToken distinctContinuationToken = new DistinctContinuationToken(null /*lasthash*/, @@ -70,11 +71,9 @@ public static Flux> cre final UInt128 continuationTokenLastHash = distinctContinuationToken.getLastHash(); - return createSourceComponentFunction.apply(distinctContinuationToken.getSourceToken()).map(component -> { - return new DistinctDocumentQueryExecutionContext(component, - distinctQueryType, - continuationTokenLastHash); - }); + return createSourceComponentFunction + .apply(distinctContinuationToken.getSourceToken(), documentQueryParams) + .map(component -> new DistinctDocumentQueryExecutionContext(component, distinctQueryType, continuationTokenLastHash)); } IDocumentQueryExecutionComponent getComponent() { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java index 41b48c65fb50..19b225bde665 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java @@ -195,20 +195,21 @@ public static Flux documentQueryParams = new PipelinedDocumentQueryParams( + resourceTypeEnum, + resourceType, + query, + resourceLink, + collectionRid, + getLazyFeedResponse, + isContinuationExpected, + initialPageSize, + targetRanges, + partitionedQueryExecutionInfo.getQueryInfo(), + cosmosQueryRequestOptions, + correlatedActivityId); + + return PipelinedDocumentQueryExecutionContext.createAsync(client, documentQueryParams); } public static Flux> createReadManyQueryAsync( diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/GroupByDocumentQueryExecutionContext.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/GroupByDocumentQueryExecutionContext.java index 83bf12d464fd..d9f63ba24e92 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/GroupByDocumentQueryExecutionContext.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/GroupByDocumentQueryExecutionContext.java @@ -19,7 +19,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; -import java.util.function.Function; +import java.util.function.BiFunction; public final class GroupByDocumentQueryExecutionContext implements IDocumentQueryExecutionComponent { @@ -40,11 +40,12 @@ public final class GroupByDocumentQueryExecutionContext impl } public static Flux> createAsync( - Function>> createSourceComponentFunction, + BiFunction, Flux>> createSourceComponentFunction, String continuationToken, Map groupByAliasToAggregateType, List orderedAliases, - boolean hasSelectValue) { + boolean hasSelectValue, + PipelinedDocumentQueryParams documentQueryParams) { if (continuationToken != null) { CosmosException dce = new BadRequestException(CONTINUATION_TOKEN_NOT_SUPPORTED_WITH_GROUP_BY); return Flux.error(dce); @@ -57,7 +58,7 @@ public static Flux> cre } GroupingTable table = new GroupingTable(groupByAliasToAggregateType, orderedAliases, hasSelectValue); // Have to pass non-null continuation token once supported - return createSourceComponentFunction.apply(null) + return createSourceComponentFunction.apply(null, documentQueryParams) .map(component -> new GroupByDocumentQueryExecutionContext<>(component, table)); } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/OrderByDocumentQueryExecutionContext.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/OrderByDocumentQueryExecutionContext.java index 4786c7648ec6..623b5ea7c5ec 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/OrderByDocumentQueryExecutionContext.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/OrderByDocumentQueryExecutionContext.java @@ -77,39 +77,32 @@ private OrderByDocumentQueryExecutionContext( public static Flux> createAsync( IDocumentQueryClient client, - ResourceType resourceTypeEnum, - Class resourceType, - SqlQuerySpec expression, - CosmosQueryRequestOptions cosmosQueryRequestOptions, - String resourceLink, - String collectionRid, - PartitionedQueryExecutionInfo partitionedQueryExecutionInfo, - List partitionKeyRanges, - int initialPageSize, - boolean isContinuationExpected, - boolean getLazyFeedResponse, - UUID correlatedActivityId) { + PipelinedDocumentQueryParams initParams) { - OrderByDocumentQueryExecutionContext context = new OrderByDocumentQueryExecutionContext(client, - partitionKeyRanges, - resourceTypeEnum, - resourceType, - expression, - cosmosQueryRequestOptions, - resourceLink, - partitionedQueryExecutionInfo.getQueryInfo().getRewrittenQuery(), - isContinuationExpected, - getLazyFeedResponse, - new OrderbyRowComparer(partitionedQueryExecutionInfo.getQueryInfo().getOrderBy()), - collectionRid, - correlatedActivityId); + OrderByDocumentQueryExecutionContext context = new OrderByDocumentQueryExecutionContext( + client, + initParams.getPartitionKeyRanges(), + initParams.getResourceTypeEnum(), + initParams.getResourceType(), + initParams.getQuery(), + initParams.getCosmosQueryRequestOptions(), + initParams.getResourceLink(), + initParams.getQueryInfo().getRewrittenQuery(), + initParams.isContinuationExpected(), + initParams.isGetLazyResponseFeed(), + new OrderbyRowComparer(initParams.getQueryInfo().getOrderBy()), + initParams.getCollectionRid(), + initParams.getCorrelatedActivityId()); + + context.setTop(initParams.getTop()); try { - context.initialize(partitionKeyRanges, - partitionedQueryExecutionInfo.getQueryInfo().getOrderBy(), - partitionedQueryExecutionInfo.getQueryInfo().getOrderByExpressions(), - initialPageSize, - ModelBridgeInternal.getRequestContinuationFromQueryRequestOptions(cosmosQueryRequestOptions)); + context.initialize( + initParams.getPartitionKeyRanges(), + initParams.getQueryInfo().getOrderBy(), + initParams.getQueryInfo().getOrderByExpressions(), + initParams.getInitialPageSize(), + ModelBridgeInternal.getRequestContinuationFromQueryRequestOptions(initParams.getCosmosQueryRequestOptions())); return Flux.just(context); } catch (CosmosException dce) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/ParallelDocumentQueryExecutionContext.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/ParallelDocumentQueryExecutionContext.java index 30e1d3636e40..e80b277aa4fc 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/ParallelDocumentQueryExecutionContext.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/ParallelDocumentQueryExecutionContext.java @@ -60,37 +60,29 @@ private ParallelDocumentQueryExecutionContext( public static Flux> createAsync( IDocumentQueryClient client, - ResourceType resourceTypeEnum, - Class resourceType, - SqlQuerySpec query, - CosmosQueryRequestOptions cosmosQueryRequestOptions, - String resourceLink, - String collectionRid, - PartitionedQueryExecutionInfo partitionedQueryExecutionInfo, - List targetRanges, - int initialPageSize, - boolean isContinuationExpected, - boolean getLazyFeedResponse, - UUID correlatedActivityId) { - - ParallelDocumentQueryExecutionContext context = new ParallelDocumentQueryExecutionContext(client, - targetRanges, - resourceTypeEnum, - resourceType, - query, - cosmosQueryRequestOptions, - resourceLink, - partitionedQueryExecutionInfo.getQueryInfo().getRewrittenQuery(), - collectionRid, - isContinuationExpected, - getLazyFeedResponse, - correlatedActivityId); + PipelinedDocumentQueryParams initParams) { + + ParallelDocumentQueryExecutionContext context = new ParallelDocumentQueryExecutionContext( + client, + initParams.getPartitionKeyRanges(), + initParams.getResourceTypeEnum(), + initParams.getResourceType(), + initParams.getQuery(), + initParams.getCosmosQueryRequestOptions(), + initParams.getResourceLink(), + initParams.getQueryInfo().getRewrittenQuery(), + initParams.getCollectionRid(), + initParams.isContinuationExpected(), + initParams.isGetLazyResponseFeed(), + initParams.getCorrelatedActivityId()); + context.setTop(initParams.getTop()); try { - context.initialize(collectionRid, - targetRanges, - initialPageSize, - ModelBridgeInternal.getRequestContinuationFromQueryRequestOptions(cosmosQueryRequestOptions)); + context.initialize( + initParams.getCollectionRid(), + initParams.getPartitionKeyRanges(), + initParams.getInitialPageSize(), + ModelBridgeInternal.getRequestContinuationFromQueryRequestOptions(initParams.getCosmosQueryRequestOptions())); return Flux.just(context); } catch (CosmosException dce) { return Flux.error(dce); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/PipelinedDocumentQueryExecutionContext.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/PipelinedDocumentQueryExecutionContext.java index 7be73c08e7bf..1b8269440cfd 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/PipelinedDocumentQueryExecutionContext.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/PipelinedDocumentQueryExecutionContext.java @@ -2,20 +2,19 @@ // Licensed under the MIT License. package com.azure.cosmos.implementation.query; +import com.azure.cosmos.implementation.PartitionKeyRange; +import com.azure.cosmos.implementation.Resource; +import com.azure.cosmos.implementation.ResourceType; +import com.azure.cosmos.implementation.Utils; import com.azure.cosmos.models.CosmosQueryRequestOptions; import com.azure.cosmos.models.FeedResponse; import com.azure.cosmos.models.ModelBridgeInternal; -import com.azure.cosmos.implementation.Resource; import com.azure.cosmos.models.SqlQuerySpec; -import com.azure.cosmos.implementation.PartitionKeyRange; -import com.azure.cosmos.implementation.ResourceType; -import com.azure.cosmos.implementation.Utils; import reactor.core.publisher.Flux; -import java.util.List; import java.util.Map; import java.util.UUID; -import java.util.function.Function; +import java.util.function.BiFunction; /** * While this class is public, but it is not part of our published public APIs. @@ -46,105 +45,113 @@ private PipelinedDocumentQueryExecutionContext(IDocumentQueryExecutionComponent< } public static Flux> createAsync( - IDocumentQueryClient client, ResourceType resourceTypeEnum, Class resourceType, SqlQuerySpec expression, - CosmosQueryRequestOptions cosmosQueryRequestOptions, String resourceLink, String collectionRid, - PartitionedQueryExecutionInfo partitionedQueryExecutionInfo, List targetRanges, - int initialPageSize, boolean isContinuationExpected, boolean getLazyFeedResponse, - UUID correlatedActivityId) { - // Use nested callback pattern to unwrap the continuation token at each level. - Function>> createBaseComponentFunction; + IDocumentQueryClient client, + PipelinedDocumentQueryParams initParams) { - QueryInfo queryInfo = partitionedQueryExecutionInfo.getQueryInfo(); + // Use nested callback pattern to unwrap the continuation token and query params at each level. + BiFunction, Flux>> createBaseComponentFunction; + + QueryInfo queryInfo = initParams.getQueryInfo(); + UUID correlatedActivityId = initParams.getCorrelatedActivityId(); + CosmosQueryRequestOptions cosmosQueryRequestOptions = initParams.getCosmosQueryRequestOptions(); if (queryInfo.hasOrderBy()) { - createBaseComponentFunction = (continuationToken) -> { + createBaseComponentFunction = (continuationToken, documentQueryParams) -> { CosmosQueryRequestOptions orderByCosmosQueryRequestOptions = ModelBridgeInternal.createQueryRequestOptions(cosmosQueryRequestOptions); ModelBridgeInternal.setQueryRequestOptionsContinuationToken(orderByCosmosQueryRequestOptions, continuationToken); - return OrderByDocumentQueryExecutionContext.createAsync(client, resourceTypeEnum, resourceType, - expression, orderByCosmosQueryRequestOptions, resourceLink, collectionRid, partitionedQueryExecutionInfo, - targetRanges, initialPageSize, isContinuationExpected, getLazyFeedResponse, - correlatedActivityId); + initParams.setCosmosQueryRequestOptions(orderByCosmosQueryRequestOptions); + + return OrderByDocumentQueryExecutionContext.createAsync(client, documentQueryParams); }; } else { - createBaseComponentFunction = (continuationToken) -> { + createBaseComponentFunction = (continuationToken, documentQueryParams) -> { CosmosQueryRequestOptions parallelCosmosQueryRequestOptions = ModelBridgeInternal.createQueryRequestOptions(cosmosQueryRequestOptions); ModelBridgeInternal.setQueryRequestOptionsContinuationToken(parallelCosmosQueryRequestOptions, continuationToken); - return ParallelDocumentQueryExecutionContext.createAsync(client, resourceTypeEnum, resourceType, - expression, parallelCosmosQueryRequestOptions, resourceLink, collectionRid, partitionedQueryExecutionInfo, - targetRanges, initialPageSize, isContinuationExpected, getLazyFeedResponse, - correlatedActivityId); + initParams.setCosmosQueryRequestOptions(parallelCosmosQueryRequestOptions); + + return ParallelDocumentQueryExecutionContext.createAsync(client, documentQueryParams); }; } - Function>> createAggregateComponentFunction; + BiFunction, Flux>> createAggregateComponentFunction; if (queryInfo.hasAggregates() && !queryInfo.hasGroupBy()) { - createAggregateComponentFunction = (continuationToken) -> { - return AggregateDocumentQueryExecutionContext.createAsync(createBaseComponentFunction, - queryInfo.getAggregates(), - queryInfo.getGroupByAliasToAggregateType(), - queryInfo.getGroupByAliases(), - queryInfo.hasSelectValue(), - continuationToken); - }; + createAggregateComponentFunction = + (continuationToken, documentQueryParams) -> + AggregateDocumentQueryExecutionContext.createAsync(createBaseComponentFunction, + queryInfo.getAggregates(), + queryInfo.getGroupByAliasToAggregateType(), + queryInfo.getGroupByAliases(), + queryInfo.hasSelectValue(), + continuationToken, + documentQueryParams); } else { createAggregateComponentFunction = createBaseComponentFunction; } - Function>> createDistinctComponentFunction; + BiFunction, Flux>> createDistinctComponentFunction; if (queryInfo.hasDistinct()) { - createDistinctComponentFunction = (continuationToken) -> { - return DistinctDocumentQueryExecutionContext.createAsync(createAggregateComponentFunction, - queryInfo.getDistinctQueryType(), continuationToken); - }; + createDistinctComponentFunction = + (continuationToken, documentQueryParams) -> + DistinctDocumentQueryExecutionContext.createAsync(createAggregateComponentFunction, + queryInfo.getDistinctQueryType(), + continuationToken, + documentQueryParams); } else { createDistinctComponentFunction = createAggregateComponentFunction; } - Function>> createGroupByComponentFunction; + BiFunction, Flux>> createGroupByComponentFunction; if (queryInfo.hasGroupBy()) { - createGroupByComponentFunction = (continuationToken) -> { - return GroupByDocumentQueryExecutionContext.createAsync(createDistinctComponentFunction, - continuationToken, - queryInfo.getGroupByAliasToAggregateType(), - queryInfo.getGroupByAliases(), - queryInfo.hasSelectValue()); - }; + createGroupByComponentFunction = + (continuationToken, documentQueryParams) -> + GroupByDocumentQueryExecutionContext.createAsync(createDistinctComponentFunction, + continuationToken, + queryInfo.getGroupByAliasToAggregateType(), + queryInfo.getGroupByAliases(), + queryInfo.hasSelectValue(), + documentQueryParams); } else{ createGroupByComponentFunction = createDistinctComponentFunction; } - Function>> createSkipComponentFunction; + BiFunction, Flux>> createSkipComponentFunction; if (queryInfo.hasOffset()) { - createSkipComponentFunction = (continuationToken) -> { - return SkipDocumentQueryExecutionContext.createAsync(createGroupByComponentFunction, - queryInfo.getOffset(), - continuationToken); - }; + createSkipComponentFunction = + (continuationToken, documentQueryParams) -> + SkipDocumentQueryExecutionContext.createAsync(createGroupByComponentFunction, + queryInfo.getOffset(), + continuationToken, + documentQueryParams); } else { createSkipComponentFunction = createGroupByComponentFunction; } - Function>> createTopComponentFunction; + BiFunction, Flux>> createTopComponentFunction; if (queryInfo.hasTop()) { - createTopComponentFunction = (continuationToken) -> { - return TopDocumentQueryExecutionContext.createAsync(createSkipComponentFunction, - queryInfo.getTop(), queryInfo.getTop(), continuationToken); - }; + createTopComponentFunction = + (continuationToken, documentQueryParams) -> + TopDocumentQueryExecutionContext.createAsync(createSkipComponentFunction, + queryInfo.getTop(), + queryInfo.getTop(), + continuationToken, + documentQueryParams); } else { createTopComponentFunction = createSkipComponentFunction; } - Function>> createTakeComponentFunction; + BiFunction, Flux>> createTakeComponentFunction; if (queryInfo.hasLimit()) { - createTakeComponentFunction = (continuationToken) -> { + createTakeComponentFunction = (continuationToken, documentQueryParams) -> { int totalLimit = queryInfo.getLimit(); if (queryInfo.hasOffset()) { // This is being done to match the limit from rewritten query totalLimit = queryInfo.getOffset() + queryInfo.getLimit(); } return TopDocumentQueryExecutionContext.createAsync(createTopComponentFunction, - queryInfo.getLimit(), totalLimit, - continuationToken); + queryInfo.getLimit(), + totalLimit, + continuationToken, + documentQueryParams); }; } else { createTakeComponentFunction = createTopComponentFunction; @@ -158,7 +165,7 @@ public static Flux new PipelinedDocumentQueryExecutionContext<>(c, pageSize, correlatedActivityId, queryInfo)); } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/PipelinedDocumentQueryParams.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/PipelinedDocumentQueryParams.java new file mode 100644 index 000000000000..f9b256d5482a --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/PipelinedDocumentQueryParams.java @@ -0,0 +1,117 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.query; + +import com.azure.cosmos.implementation.PartitionKeyRange; +import com.azure.cosmos.implementation.Resource; +import com.azure.cosmos.implementation.ResourceType; +import com.azure.cosmos.models.CosmosQueryRequestOptions; +import com.azure.cosmos.models.SqlQuerySpec; + +import java.util.List; +import java.util.UUID; + +public class PipelinedDocumentQueryParams { + private int top = -1; + private final int initialPageSize; + private final boolean isContinuationExpected; + private final boolean getLazyResponseFeed; + private final String collectionRid; + private final ResourceType resourceTypeEnum; + private final Class resourceType; + private final SqlQuerySpec query; + private final String resourceLink; + private final UUID correlatedActivityId; + private CosmosQueryRequestOptions cosmosQueryRequestOptions; + private final List partitionKeyRanges; + private final QueryInfo queryInfo; + + public PipelinedDocumentQueryParams( + ResourceType resourceTypeEnum, + Class resourceType, + SqlQuerySpec query, + String resourceLink, + String collectionRid, + boolean getLazyResponseFeed, + boolean isContinuationExpected, + int initialPageSize, + List partitionKeyRanges, + QueryInfo queryInfo, + CosmosQueryRequestOptions cosmosQueryRequestOptions, + UUID correlatedActivityId) { + + this.resourceTypeEnum = resourceTypeEnum; + this.resourceType = resourceType; + this.query = query; + this.resourceLink = resourceLink; + this.collectionRid = collectionRid; + this.getLazyResponseFeed = getLazyResponseFeed; + this.isContinuationExpected = isContinuationExpected; + this.initialPageSize = initialPageSize; + this.correlatedActivityId = correlatedActivityId; + this.partitionKeyRanges = partitionKeyRanges; + this.queryInfo = queryInfo; + this.cosmosQueryRequestOptions = cosmosQueryRequestOptions; + } + + public int getTop() { + return top; + } + + public void setTop(int top) { + this.top = top; + } + + public int getInitialPageSize() { + return initialPageSize; + } + + public boolean isContinuationExpected() { + return isContinuationExpected; + } + + public boolean isGetLazyResponseFeed() { + return getLazyResponseFeed; + } + + public String getCollectionRid() { + return collectionRid; + } + + public ResourceType getResourceTypeEnum() { + return resourceTypeEnum; + } + + public Class getResourceType() { + return resourceType; + } + + public SqlQuerySpec getQuery() { + return query; + } + + public String getResourceLink() { + return resourceLink; + } + + public UUID getCorrelatedActivityId() { + return correlatedActivityId; + } + + public CosmosQueryRequestOptions getCosmosQueryRequestOptions() { + return cosmosQueryRequestOptions; + } + + public void setCosmosQueryRequestOptions(CosmosQueryRequestOptions cosmosQueryRequestOptions) { + this.cosmosQueryRequestOptions = cosmosQueryRequestOptions; + } + + public List getPartitionKeyRanges() { + return partitionKeyRanges; + } + + public QueryInfo getQueryInfo() { + return queryInfo; + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/SkipDocumentQueryExecutionContext.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/SkipDocumentQueryExecutionContext.java index 0c4e559d17de..741a20cd64e3 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/SkipDocumentQueryExecutionContext.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/SkipDocumentQueryExecutionContext.java @@ -14,7 +14,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; -import java.util.function.Function; +import java.util.function.BiFunction; import java.util.stream.Collectors; public final class SkipDocumentQueryExecutionContext implements IDocumentQueryExecutionComponent { @@ -31,9 +31,10 @@ public final class SkipDocumentQueryExecutionContext impleme } public static Flux> createAsync( - Function>> createSourceComponentFunction, + BiFunction, Flux>> createSourceComponentFunction, int skipCount, - String continuationToken) { + String continuationToken, + PipelinedDocumentQueryParams documentQueryParams) { OffsetContinuationToken offsetContinuationToken; Utils.ValueHolder outOffsetContinuationToken = new Utils.ValueHolder<>(); if (continuationToken != null) { @@ -51,7 +52,7 @@ public static Flux> cre offsetContinuationToken = new OffsetContinuationToken(skipCount, null); } - return createSourceComponentFunction.apply(offsetContinuationToken.getSourceToken()) + return createSourceComponentFunction.apply(offsetContinuationToken.getSourceToken(), documentQueryParams) .map(component -> new SkipDocumentQueryExecutionContext<>(component, offsetContinuationToken.getOffset())); } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/TopDocumentQueryExecutionContext.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/TopDocumentQueryExecutionContext.java index 43ca0baafa44..e3a61abe7474 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/TopDocumentQueryExecutionContext.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/TopDocumentQueryExecutionContext.java @@ -13,6 +13,7 @@ import java.util.HashMap; import java.util.Map; +import java.util.function.BiFunction; import java.util.function.Function; import java.util.function.Predicate; @@ -30,8 +31,11 @@ public TopDocumentQueryExecutionContext(IDocumentQueryExecutionComponent comp } public static Flux> createAsync( - Function>> createSourceComponentFunction, - int topCount, int limit, String topContinuationToken) { + BiFunction, Flux>> createSourceComponentFunction, + int topCount, + int limit, + String topContinuationToken, + PipelinedDocumentQueryParams documentQueryParams) { TakeContinuationToken takeContinuationToken; if (topContinuationToken == null) { @@ -57,19 +61,17 @@ public static Flux> cre return Flux.error(dce); } + // The top value setting here will be propagated down to document producer. + documentQueryParams.setTop(limit); + return createSourceComponentFunction - .apply(takeContinuationToken.getSourceToken()) + .apply(takeContinuationToken.getSourceToken(), documentQueryParams) .map(component -> new TopDocumentQueryExecutionContext<>(component, takeContinuationToken.getTakeCount(), limit)); } @Override public Flux> drainAsync(int maxPageSize) { - ParallelDocumentQueryExecutionContextBase context; - // These checks below are getting messier. Needs redesign - context = getContextFromComponent(); - - context.setTop(this.limit); return this.component.drainAsync(maxPageSize).takeUntil(new Predicate>() { @@ -124,45 +126,4 @@ public FeedResponse apply(FeedResponse t) { } }); } - - // TODO: This could be restructured eventually with better design to avoid these many checks - private ParallelDocumentQueryExecutionContextBase getContextFromComponent() { - ParallelDocumentQueryExecutionContextBase context; - if (this.component instanceof GroupByDocumentQueryExecutionContext) { - context = - (ParallelDocumentQueryExecutionContextBase) ((GroupByDocumentQueryExecutionContext) this.component) - .getComponent(); - } else if (this.component instanceof DistinctDocumentQueryExecutionContext) { - context = - (ParallelDocumentQueryExecutionContextBase) ((DistinctDocumentQueryExecutionContext) this.component) - .getComponent(); - } else if (this.component instanceof AggregateDocumentQueryExecutionContext) { - context = - (ParallelDocumentQueryExecutionContextBase) ((AggregateDocumentQueryExecutionContext) this.component) - .getComponent(); - } else if (this.component instanceof SkipDocumentQueryExecutionContext) { - - IDocumentQueryExecutionComponent component = ((SkipDocumentQueryExecutionContext) this.component) - .getComponent(); - if (component instanceof DistinctDocumentQueryExecutionContext) { - context = - (ParallelDocumentQueryExecutionContextBase) ((DistinctDocumentQueryExecutionContext) component) - .getComponent(); - } else if (component instanceof AggregateDocumentQueryExecutionContext) { - context = - (ParallelDocumentQueryExecutionContextBase) ((AggregateDocumentQueryExecutionContext) component) - .getComponent(); - } else if (component instanceof GroupByDocumentQueryExecutionContext) { - context = - (ParallelDocumentQueryExecutionContextBase) ((GroupByDocumentQueryExecutionContext) component) - .getComponent(); - } else { - context = - (ParallelDocumentQueryExecutionContextBase) component; - } - } else { - context = (ParallelDocumentQueryExecutionContextBase) this.component; - } - return context; - } } diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/OffsetLimitQueryTests.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/OffsetLimitQueryTests.java index 2b96d2ff8b8b..84d4612ea6d3 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/OffsetLimitQueryTests.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/OffsetLimitQueryTests.java @@ -8,6 +8,8 @@ import com.azure.cosmos.CosmosAsyncContainer; import com.azure.cosmos.CosmosAsyncDatabase; import com.azure.cosmos.CosmosClientBuilder; +import com.azure.cosmos.implementation.Resource; +import com.azure.cosmos.models.ModelBridgeInternal; import com.azure.cosmos.util.CosmosPagedFlux; import com.azure.cosmos.implementation.InternalObjectNode; import com.azure.cosmos.models.CosmosQueryRequestOptions; @@ -16,6 +18,7 @@ import com.azure.cosmos.implementation.FeedResponseValidator; import com.azure.cosmos.implementation.Utils; import com.azure.cosmos.implementation.query.OffsetContinuationToken; +import com.fasterxml.jackson.databind.JsonNode; import io.reactivex.subscribers.TestSubscriber; import org.testng.annotations.AfterClass; import org.testng.annotations.BeforeClass; @@ -23,11 +26,13 @@ import org.testng.annotations.Test; import java.util.ArrayList; +import java.util.Comparator; import java.util.HashSet; import java.util.Iterator; import java.util.List; import java.util.Set; import java.util.concurrent.TimeUnit; +import java.util.stream.Collectors; import static org.assertj.core.api.Assertions.assertThat; @@ -130,6 +135,56 @@ public void queryDocumentsWithOffsetContinuationTokens() { this.queryWithContinuationTokensAndPageSizes(query, new int[] {1, 5, 15}, takeCount); } + @Test(groups = {"simple"}, timeOut = TIMEOUT, dataProvider = "queryMetricsArgProvider") + public void queryDocumentsWithDistinct(boolean qmEnabled) { + int skipCount = 4; + int takeCount = 10; + String query = + String.format("SELECT DISTINCT c.id from c OFFSET %s LIMIT %s", skipCount, takeCount); + CosmosQueryRequestOptions options = new CosmosQueryRequestOptions(); + options.setQueryMetricsEnabled(qmEnabled); + options.setMaxDegreeOfParallelism(2); + CosmosPagedFlux queryObservable = createdCollection.queryItems(query, options, InternalObjectNode.class); + + List expectedIds = + docs.stream().skip(4).limit(10).map(doc -> doc.getResourceId()).collect(Collectors.toList()); + + FeedResponseListValidator validator = + new FeedResponseListValidator.Builder() + .containsExactly(expectedIds) + .numberOfPages(3) + .hasValidQueryMetrics(qmEnabled) + .build(); + + validateQuerySuccess(queryObservable.byPage(5), validator, TIMEOUT); + } + + @Test(groups = {"simple"}, timeOut = TIMEOUT, dataProvider = "queryMetricsArgProvider") + public void queryDocumentsWithAggregate(boolean qmEnabled) { + int skipCount = 0; + int takeCount = 10; + String query = + String.format("SELECT VALUE MAX(c.%s) from c OFFSET %s LIMIT %s", field, skipCount, takeCount); + CosmosQueryRequestOptions options = new CosmosQueryRequestOptions(); + options.setQueryMetricsEnabled(qmEnabled); + CosmosPagedFlux queryObservable = createdCollection.queryItems(query, options, JsonNode.class); + + // The pipeline execution sequence is Aggregrate, skip, and top/limit, hence finding the max from among the docs + InternalObjectNode expectedDoc = docs + .stream() + .max(Comparator.comparing(r -> ModelBridgeInternal.getIntFromJsonSerializable(r, field))) + .get(); + + FeedResponseListValidator validator = + new FeedResponseListValidator.Builder() + .withAggregateValue(ModelBridgeInternal.getIntFromJsonSerializable(expectedDoc, field)) + .numberOfPages(1) + .hasValidQueryMetrics(qmEnabled) + .build(); + + validateQuerySuccess(queryObservable.byPage(5), validator, TIMEOUT); + } + private void queryWithContinuationTokensAndPageSizes(String query, int[] pageSizes, int takeCount) { for (int pageSize : pageSizes) { List receivedDocuments = this.queryWithContinuationTokens(query, pageSize);