From fa5dc6b372d54ecb8312077832a859bf5a35e6af Mon Sep 17 00:00:00 2001 From: Annie Liang Date: Mon, 29 Jun 2020 19:27:33 -0700 Subject: [PATCH 1/6] FixForClassCastExceptionForTopDocumentQueryContext --- ...ggregateDocumentQueryExecutionContext.java | 9 +- ...DistinctDocumentQueryExecutionContext.java | 9 +- .../DocumentQueryExecutionContextFactory.java | 29 ++-- .../GroupByDocumentQueryExecutionContext.java | 9 +- .../OrderByDocumentQueryExecutionContext.java | 53 +++--- ...ParallelDocumentQueryExecutionContext.java | 50 +++--- ...ipelinedDocumentQueryExecutionContext.java | 96 ++++++----- .../query/PipelinedDocumentQueryParams.java | 157 ++++++++++++++++++ .../SkipDocumentQueryExecutionContext.java | 9 +- .../TopDocumentQueryExecutionContext.java | 59 ++----- .../cosmos/rx/OffsetLimitQueryTests.java | 27 +++ .../com/azure/cosmos/rx/TopQueryTests.java | 48 ++++++ 12 files changed, 373 insertions(+), 182 deletions(-) create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/PipelinedDocumentQueryParams.java 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..03cb6c313387 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{ @@ -104,15 +104,16 @@ 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) + .apply(continuationToken, documentQueryParams) .map(component -> { return new AggregateDocumentQueryExecutionContext(component, new ArrayList<>(aggregates), 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..b535065fdd3e 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,7 +71,7 @@ public static Flux> cre final UInt128 continuationTokenLastHash = distinctContinuationToken.getLastHash(); - return createSourceComponentFunction.apply(distinctContinuationToken.getSourceToken()).map(component -> { + return createSourceComponentFunction.apply(distinctContinuationToken.getSourceToken(), documentQueryParams).map(component -> { return new DistinctDocumentQueryExecutionContext(component, distinctQueryType, continuationTokenLastHash); 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..9d216b897e31 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..503b7cd1e6cd 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.getGetLazyResponseFeed(), + 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..32b2840964f2 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.getGetLazyResponseFeed(), + 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..44f1eb6266c4 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,114 @@ 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) { + IDocumentQueryClient client, + PipelinedDocumentQueryParams initParams) { + // Use nested callback pattern to unwrap the continuation token at each level. - Function>> createBaseComponentFunction; + BiFunction, Flux>> createBaseComponentFunction; - QueryInfo queryInfo = partitionedQueryExecutionInfo.getQueryInfo(); + 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) -> { + createAggregateComponentFunction = (continuationToken, documentQueryParams) -> { return AggregateDocumentQueryExecutionContext.createAsync(createBaseComponentFunction, queryInfo.getAggregates(), queryInfo.getGroupByAliasToAggregateType(), queryInfo.getGroupByAliases(), queryInfo.hasSelectValue(), - continuationToken); + 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) -> { + return DistinctDocumentQueryExecutionContext.createAsync( + createAggregateComponentFunction, + queryInfo.getDistinctQueryType(), + continuationToken, + documentQueryParams); }; } else { createDistinctComponentFunction = createAggregateComponentFunction; } - Function>> createGroupByComponentFunction; + BiFunction, Flux>> createGroupByComponentFunction; if (queryInfo.hasGroupBy()) { - createGroupByComponentFunction = (continuationToken) -> { + createGroupByComponentFunction = (continuationToken, documentQueryParams) -> { return GroupByDocumentQueryExecutionContext.createAsync(createDistinctComponentFunction, continuationToken, queryInfo.getGroupByAliasToAggregateType(), queryInfo.getGroupByAliases(), - queryInfo.hasSelectValue()); + queryInfo.hasSelectValue(), + documentQueryParams); }; } else{ createGroupByComponentFunction = createDistinctComponentFunction; } - Function>> createSkipComponentFunction; + BiFunction, Flux>> createSkipComponentFunction; if (queryInfo.hasOffset()) { - createSkipComponentFunction = (continuationToken) -> { + createSkipComponentFunction = (continuationToken, documentQueryParams) -> { return SkipDocumentQueryExecutionContext.createAsync(createGroupByComponentFunction, queryInfo.getOffset(), - continuationToken); + continuationToken, + documentQueryParams); }; } else { createSkipComponentFunction = createGroupByComponentFunction; } - Function>> createTopComponentFunction; + BiFunction, Flux>> createTopComponentFunction; if (queryInfo.hasTop()) { - createTopComponentFunction = (continuationToken) -> { + createTopComponentFunction = (continuationToken, documentQueryParams) -> { return TopDocumentQueryExecutionContext.createAsync(createSkipComponentFunction, - queryInfo.getTop(), queryInfo.getTop(), continuationToken); + 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 +166,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..787a8c9a5ae8 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/PipelinedDocumentQueryParams.java @@ -0,0 +1,157 @@ +// 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 int initialPageSize; + private boolean isContinuationExpected; + private String collectionRid; + private ResourceType resourceTypeEnum; + private Class resourceType; + private SqlQuerySpec query; + private String resourceLink; + private Boolean getLazyResponseFeed; + private UUID correlatedActivityId; + private CosmosQueryRequestOptions cosmosQueryRequestOptions; + private final List partitionKeyRanges; + private QueryInfo queryInfo; + + public PipelinedDocumentQueryParams( + ResourceType resourceTypeEnum, + Class resourceType, + SqlQuerySpec query, + String resourceLink, + String collectionRid, + boolean getLazyFeedResponse, + 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 = getLazyFeedResponse; + 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 void setInitialPageSize(int initialPageSize) { + this.initialPageSize = initialPageSize; + } + + public boolean isContinuationExpected() { + return isContinuationExpected; + } + + public void setContinuationExpected(boolean continuationExpected) { + isContinuationExpected = continuationExpected; + } + + public String getCollectionRid() { + return collectionRid; + } + + public void setCollectionRid(String collectionRid) { + this.collectionRid = collectionRid; + } + + public ResourceType getResourceTypeEnum() { + return resourceTypeEnum; + } + + public void setResourceTypeEnum(ResourceType resourceTypeEnum) { + this.resourceTypeEnum = resourceTypeEnum; + } + + public Class getResourceType() { + return resourceType; + } + + public void setResourceType(Class resourceType) { + this.resourceType = resourceType; + } + + public SqlQuerySpec getQuery() { + return query; + } + + public void setQuery(SqlQuerySpec query) { + this.query = query; + } + + public String getResourceLink() { + return resourceLink; + } + + public void setResourceLink(String resourceLink) { + this.resourceLink = resourceLink; + } + + public Boolean getGetLazyResponseFeed() { + return getLazyResponseFeed; + } + + public void setGetLazyResponseFeed(Boolean getLazyResponseFeed) { + this.getLazyResponseFeed = getLazyResponseFeed; + } + + public UUID getCorrelatedActivityId() { + return correlatedActivityId; + } + + public void setCorrelatedActivityId(UUID correlatedActivityId) { + this.correlatedActivityId = correlatedActivityId; + } + + public CosmosQueryRequestOptions getCosmosQueryRequestOptions() { + return cosmosQueryRequestOptions; + } + + public void setCosmosQueryRequestOptions(CosmosQueryRequestOptions cosmosQueryRequestOptions) { + this.cosmosQueryRequestOptions = cosmosQueryRequestOptions; + } + + public List getPartitionKeyRanges() { + return partitionKeyRanges; + } + + public QueryInfo getQueryInfo() { + return queryInfo; + } + + public void setQueryInfo(QueryInfo queryInfo) { + this.queryInfo = 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..309ea188ef8a 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,7 @@ import com.azure.cosmos.CosmosAsyncContainer; import com.azure.cosmos.CosmosAsyncDatabase; import com.azure.cosmos.CosmosClientBuilder; +import com.azure.cosmos.models.ModelBridgeInternal; import com.azure.cosmos.util.CosmosPagedFlux; import com.azure.cosmos.implementation.InternalObjectNode; import com.azure.cosmos.models.CosmosQueryRequestOptions; @@ -23,11 +24,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 +133,30 @@ 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); + } + private void queryWithContinuationTokensAndPageSizes(String query, int[] pageSizes, int takeCount) { for (int pageSize : pageSizes) { List receivedDocuments = this.queryWithContinuationTokens(query, pageSize); diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/TopQueryTests.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/TopQueryTests.java index ba9bbf9912d0..c99c96296d9c 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/TopQueryTests.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/TopQueryTests.java @@ -124,6 +124,54 @@ public void queryDocumentsWithTopContinuationTokens() throws Exception { this.queryWithContinuationTokensAndPageSizes(query, new int[] { 1, 5, 10 }, 8); } + @Test(groups = { "simple" }, timeOut = TIMEOUT, dataProvider = "queryMetricsArgProvider", retryAnalyzer = RetryAnalyzer.class) + public void queryDocumentsWithTopAndAggregate(boolean qmEnabled) throws Exception { + + CosmosQueryRequestOptions options = new CosmosQueryRequestOptions(); + options.setMaxDegreeOfParallelism(2); + options.setQueryMetricsEnabled(qmEnabled); + + int expectedTotalSize = 20; + int expectedNumberOfPages = 3; + int[] expectedPageLengths = new int[] { 9, 9, 2 }; + + for (int i = 0; i < 2; i++) { + CosmosPagedFlux queryObservable1 = createdCollection.queryItems("SELECT TOP 0 value AVG(c.field) from c", + options, + JsonNode.class); + + FeedResponseListValidator validator1 = new FeedResponseListValidator.Builder() + .totalSize(0).build(); + + validateQuerySuccess(queryObservable1.byPage(9), validator1, TIMEOUT); + + CosmosPagedFlux queryObservable2 = createdCollection.queryItems("SELECT TOP 1 value AVG(c.field) from c", + options, + JsonNode.class); + + FeedResponseListValidator validator2 = new FeedResponseListValidator.Builder() + .totalSize(1).build(); + + validateQuerySuccess(queryObservable2.byPage(), validator2, TIMEOUT); + + CosmosPagedFlux queryObservable3 = createdCollection.queryItems("SELECT TOP 20 * from c", options, InternalObjectNode.class); + + FeedResponseListValidator validator3 = new FeedResponseListValidator.Builder() + .totalSize(expectedTotalSize).numberOfPages(expectedNumberOfPages).pageLengths(expectedPageLengths) + .hasValidQueryMetrics(qmEnabled).build(); + + validateQuerySuccess(queryObservable3.byPage(), validator3, TIMEOUT); + + if (i == 0) { + options.setPartitionKey(new PartitionKey(firstPk)); + expectedTotalSize = 10; + expectedNumberOfPages = 2; + expectedPageLengths = new int[] { 9, 1 }; + + } + } + } + private void queryWithContinuationTokensAndPageSizes(String query, int[] pageSizes, int topCount) { for (int pageSize : pageSizes) { List receivedDocuments = this.queryWithContinuationTokens(query, pageSize); From a41f6a4a8f0853037e079a0b3e5bbe355d0b8392 Mon Sep 17 00:00:00 2001 From: Annie Liang Date: Mon, 29 Jun 2020 19:33:08 -0700 Subject: [PATCH 2/6] FixForClassCastException --- .../query/PipelinedDocumentQueryParams.java | 66 +++++-------------- 1 file changed, 17 insertions(+), 49 deletions(-) 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 index 787a8c9a5ae8..5f5d74bf4ac5 100644 --- 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 @@ -16,16 +16,16 @@ public class PipelinedDocumentQueryParams { private int top = -1; private int initialPageSize; private boolean isContinuationExpected; - private String collectionRid; - private ResourceType resourceTypeEnum; - private Class resourceType; - private SqlQuerySpec query; - private String resourceLink; - private Boolean getLazyResponseFeed; - private UUID correlatedActivityId; - private CosmosQueryRequestOptions cosmosQueryRequestOptions; + private 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 final CosmosQueryRequestOptions cosmosQueryRequestOptions; private final List partitionKeyRanges; - private QueryInfo queryInfo; + private final QueryInfo queryInfo; public PipelinedDocumentQueryParams( ResourceType resourceTypeEnum, @@ -79,70 +79,42 @@ public void setContinuationExpected(boolean continuationExpected) { isContinuationExpected = continuationExpected; } - public String getCollectionRid() { - return collectionRid; + public boolean isGetLazyResponseFeed() { + return getLazyResponseFeed; } - public void setCollectionRid(String collectionRid) { - this.collectionRid = collectionRid; + public void setGetLazyResponseFeed(boolean getLazyResponseFeed) { + this.getLazyResponseFeed = getLazyResponseFeed; } - public ResourceType getResourceTypeEnum() { - return resourceTypeEnum; + public String getCollectionRid() { + return collectionRid; } - public void setResourceTypeEnum(ResourceType resourceTypeEnum) { - this.resourceTypeEnum = resourceTypeEnum; + public ResourceType getResourceTypeEnum() { + return resourceTypeEnum; } public Class getResourceType() { return resourceType; } - public void setResourceType(Class resourceType) { - this.resourceType = resourceType; - } - public SqlQuerySpec getQuery() { return query; } - public void setQuery(SqlQuerySpec query) { - this.query = query; - } - public String getResourceLink() { return resourceLink; } - public void setResourceLink(String resourceLink) { - this.resourceLink = resourceLink; - } - - public Boolean getGetLazyResponseFeed() { - return getLazyResponseFeed; - } - - public void setGetLazyResponseFeed(Boolean getLazyResponseFeed) { - this.getLazyResponseFeed = getLazyResponseFeed; - } - public UUID getCorrelatedActivityId() { return correlatedActivityId; } - public void setCorrelatedActivityId(UUID correlatedActivityId) { - this.correlatedActivityId = correlatedActivityId; - } - public CosmosQueryRequestOptions getCosmosQueryRequestOptions() { return cosmosQueryRequestOptions; } - public void setCosmosQueryRequestOptions(CosmosQueryRequestOptions cosmosQueryRequestOptions) { - this.cosmosQueryRequestOptions = cosmosQueryRequestOptions; - } - public List getPartitionKeyRanges() { return partitionKeyRanges; } @@ -150,8 +122,4 @@ public List getPartitionKeyRanges() { public QueryInfo getQueryInfo() { return queryInfo; } - - public void setQueryInfo(QueryInfo queryInfo) { - this.queryInfo = queryInfo; - } } From 2834b705d33752d749e6b0eefbe0211e92acdc70 Mon Sep 17 00:00:00 2001 From: Annie Liang Date: Mon, 29 Jun 2020 19:35:59 -0700 Subject: [PATCH 3/6] revert --- .../com/azure/cosmos/rx/TopQueryTests.java | 48 ------------------- 1 file changed, 48 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/TopQueryTests.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/TopQueryTests.java index c99c96296d9c..ba9bbf9912d0 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/TopQueryTests.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/TopQueryTests.java @@ -124,54 +124,6 @@ public void queryDocumentsWithTopContinuationTokens() throws Exception { this.queryWithContinuationTokensAndPageSizes(query, new int[] { 1, 5, 10 }, 8); } - @Test(groups = { "simple" }, timeOut = TIMEOUT, dataProvider = "queryMetricsArgProvider", retryAnalyzer = RetryAnalyzer.class) - public void queryDocumentsWithTopAndAggregate(boolean qmEnabled) throws Exception { - - CosmosQueryRequestOptions options = new CosmosQueryRequestOptions(); - options.setMaxDegreeOfParallelism(2); - options.setQueryMetricsEnabled(qmEnabled); - - int expectedTotalSize = 20; - int expectedNumberOfPages = 3; - int[] expectedPageLengths = new int[] { 9, 9, 2 }; - - for (int i = 0; i < 2; i++) { - CosmosPagedFlux queryObservable1 = createdCollection.queryItems("SELECT TOP 0 value AVG(c.field) from c", - options, - JsonNode.class); - - FeedResponseListValidator validator1 = new FeedResponseListValidator.Builder() - .totalSize(0).build(); - - validateQuerySuccess(queryObservable1.byPage(9), validator1, TIMEOUT); - - CosmosPagedFlux queryObservable2 = createdCollection.queryItems("SELECT TOP 1 value AVG(c.field) from c", - options, - JsonNode.class); - - FeedResponseListValidator validator2 = new FeedResponseListValidator.Builder() - .totalSize(1).build(); - - validateQuerySuccess(queryObservable2.byPage(), validator2, TIMEOUT); - - CosmosPagedFlux queryObservable3 = createdCollection.queryItems("SELECT TOP 20 * from c", options, InternalObjectNode.class); - - FeedResponseListValidator validator3 = new FeedResponseListValidator.Builder() - .totalSize(expectedTotalSize).numberOfPages(expectedNumberOfPages).pageLengths(expectedPageLengths) - .hasValidQueryMetrics(qmEnabled).build(); - - validateQuerySuccess(queryObservable3.byPage(), validator3, TIMEOUT); - - if (i == 0) { - options.setPartitionKey(new PartitionKey(firstPk)); - expectedTotalSize = 10; - expectedNumberOfPages = 2; - expectedPageLengths = new int[] { 9, 1 }; - - } - } - } - private void queryWithContinuationTokensAndPageSizes(String query, int[] pageSizes, int topCount) { for (int pageSize : pageSizes) { List receivedDocuments = this.queryWithContinuationTokens(query, pageSize); From b043fec00fdaf3c20c460e6e6ac98aa22c7588ab Mon Sep 17 00:00:00 2001 From: Annie Liang Date: Tue, 30 Jun 2020 09:46:56 -0700 Subject: [PATCH 4/6] clean --- .../OrderByDocumentQueryExecutionContext.java | 2 +- ...ParallelDocumentQueryExecutionContext.java | 2 +- ...ipelinedDocumentQueryExecutionContext.java | 75 +++++++++---------- .../query/PipelinedDocumentQueryParams.java | 28 +++---- .../cosmos/rx/OffsetLimitQueryTests.java | 28 +++++++ 5 files changed, 77 insertions(+), 58 deletions(-) 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 503b7cd1e6cd..ca1ee0b2b234 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 @@ -89,7 +89,7 @@ public static Flux> cre initParams.getResourceLink(), initParams.getQueryInfo().getRewrittenQuery(), initParams.isContinuationExpected(), - initParams.getGetLazyResponseFeed(), + initParams.isGetLazyResponseFeed(), new OrderbyRowComparer(initParams.getQueryInfo().getOrderBy()), initParams.getCollectionRid(), initParams.getCorrelatedActivityId()); 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 32b2840964f2..9450d46b17ae 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 @@ -73,7 +73,7 @@ public static Flux> cre initParams.getQueryInfo().getRewrittenQuery(), initParams.getCollectionRid(), initParams.isContinuationExpected(), - initParams.getGetLazyResponseFeed(), + initParams.isGetLazyResponseFeed(), initParams.getCorrelatedActivityId()); context.setTop(initParams.getTop()); 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 44f1eb6266c4..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 @@ -48,7 +48,7 @@ public static Flux initParams) { - // Use nested callback pattern to unwrap the continuation token at each level. + // Use nested callback pattern to unwrap the continuation token and query params at each level. BiFunction, Flux>> createBaseComponentFunction; QueryInfo queryInfo = initParams.getQueryInfo(); @@ -75,67 +75,66 @@ public static Flux, Flux>> createAggregateComponentFunction; if (queryInfo.hasAggregates() && !queryInfo.hasGroupBy()) { - createAggregateComponentFunction = (continuationToken, documentQueryParams) -> { - return AggregateDocumentQueryExecutionContext.createAsync(createBaseComponentFunction, - queryInfo.getAggregates(), - queryInfo.getGroupByAliasToAggregateType(), - queryInfo.getGroupByAliases(), - queryInfo.hasSelectValue(), - continuationToken, - documentQueryParams); - }; + createAggregateComponentFunction = + (continuationToken, documentQueryParams) -> + AggregateDocumentQueryExecutionContext.createAsync(createBaseComponentFunction, + queryInfo.getAggregates(), + queryInfo.getGroupByAliasToAggregateType(), + queryInfo.getGroupByAliases(), + queryInfo.hasSelectValue(), + continuationToken, + documentQueryParams); } else { createAggregateComponentFunction = createBaseComponentFunction; } BiFunction, Flux>> createDistinctComponentFunction; if (queryInfo.hasDistinct()) { - createDistinctComponentFunction = (continuationToken, documentQueryParams) -> { - return DistinctDocumentQueryExecutionContext.createAsync( - createAggregateComponentFunction, - queryInfo.getDistinctQueryType(), - continuationToken, - documentQueryParams); - }; + createDistinctComponentFunction = + (continuationToken, documentQueryParams) -> + DistinctDocumentQueryExecutionContext.createAsync(createAggregateComponentFunction, + queryInfo.getDistinctQueryType(), + continuationToken, + documentQueryParams); } else { createDistinctComponentFunction = createAggregateComponentFunction; } BiFunction, Flux>> createGroupByComponentFunction; if (queryInfo.hasGroupBy()) { - createGroupByComponentFunction = (continuationToken, documentQueryParams) -> { - return GroupByDocumentQueryExecutionContext.createAsync(createDistinctComponentFunction, - continuationToken, - queryInfo.getGroupByAliasToAggregateType(), - queryInfo.getGroupByAliases(), - queryInfo.hasSelectValue(), - documentQueryParams); - }; + createGroupByComponentFunction = + (continuationToken, documentQueryParams) -> + GroupByDocumentQueryExecutionContext.createAsync(createDistinctComponentFunction, + continuationToken, + queryInfo.getGroupByAliasToAggregateType(), + queryInfo.getGroupByAliases(), + queryInfo.hasSelectValue(), + documentQueryParams); } else{ createGroupByComponentFunction = createDistinctComponentFunction; } BiFunction, Flux>> createSkipComponentFunction; if (queryInfo.hasOffset()) { - createSkipComponentFunction = (continuationToken, documentQueryParams) -> { - return SkipDocumentQueryExecutionContext.createAsync(createGroupByComponentFunction, - queryInfo.getOffset(), - continuationToken, - documentQueryParams); - }; + createSkipComponentFunction = + (continuationToken, documentQueryParams) -> + SkipDocumentQueryExecutionContext.createAsync(createGroupByComponentFunction, + queryInfo.getOffset(), + continuationToken, + documentQueryParams); } else { createSkipComponentFunction = createGroupByComponentFunction; } BiFunction, Flux>> createTopComponentFunction; if (queryInfo.hasTop()) { - createTopComponentFunction = (continuationToken, documentQueryParams) -> { - return TopDocumentQueryExecutionContext.createAsync(createSkipComponentFunction, - queryInfo.getTop(), - queryInfo.getTop(), - continuationToken, - documentQueryParams); - }; + createTopComponentFunction = + (continuationToken, documentQueryParams) -> + TopDocumentQueryExecutionContext.createAsync(createSkipComponentFunction, + queryInfo.getTop(), + queryInfo.getTop(), + continuationToken, + documentQueryParams); } else { createTopComponentFunction = createSkipComponentFunction; } 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 index 5f5d74bf4ac5..ac11edc397d8 100644 --- 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 @@ -14,16 +14,16 @@ public class PipelinedDocumentQueryParams { private int top = -1; - private int initialPageSize; - private boolean isContinuationExpected; - private boolean getLazyResponseFeed; + 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 final CosmosQueryRequestOptions cosmosQueryRequestOptions; + private CosmosQueryRequestOptions cosmosQueryRequestOptions; private final List partitionKeyRanges; private final QueryInfo queryInfo; @@ -33,7 +33,7 @@ public PipelinedDocumentQueryParams( SqlQuerySpec query, String resourceLink, String collectionRid, - boolean getLazyFeedResponse, + boolean getLazyResponseFeed, boolean isContinuationExpected, int initialPageSize, List partitionKeyRanges, @@ -46,7 +46,7 @@ public PipelinedDocumentQueryParams( this.query = query; this.resourceLink = resourceLink; this.collectionRid = collectionRid; - this.getLazyResponseFeed = getLazyFeedResponse; + this.getLazyResponseFeed = getLazyResponseFeed; this.isContinuationExpected = isContinuationExpected; this.initialPageSize = initialPageSize; this.correlatedActivityId = correlatedActivityId; @@ -67,26 +67,14 @@ public int getInitialPageSize() { return initialPageSize; } - public void setInitialPageSize(int initialPageSize) { - this.initialPageSize = initialPageSize; - } - public boolean isContinuationExpected() { return isContinuationExpected; } - public void setContinuationExpected(boolean continuationExpected) { - isContinuationExpected = continuationExpected; - } - public boolean isGetLazyResponseFeed() { return getLazyResponseFeed; } - public void setGetLazyResponseFeed(boolean getLazyResponseFeed) { - this.getLazyResponseFeed = getLazyResponseFeed; - } - public String getCollectionRid() { return collectionRid; } @@ -115,6 +103,10 @@ public CosmosQueryRequestOptions getCosmosQueryRequestOptions() { return cosmosQueryRequestOptions; } + public void setCosmosQueryRequestOptions(CosmosQueryRequestOptions cosmosQueryRequestOptions) { + this.cosmosQueryRequestOptions = cosmosQueryRequestOptions; + } + public List getPartitionKeyRanges() { return partitionKeyRanges; } 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 309ea188ef8a..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,7 @@ 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; @@ -17,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; @@ -157,6 +159,32 @@ public void queryDocumentsWithDistinct(boolean qmEnabled) { 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); From f65853f3dd679d43deee586d9cb62220263d440e Mon Sep 17 00:00:00 2001 From: Annie Liang Date: Tue, 30 Jun 2020 09:54:29 -0700 Subject: [PATCH 5/6] clean --- ...AggregateDocumentQueryExecutionContext.java | 18 ++++++++---------- .../DistinctDocumentQueryExecutionContext.java | 8 +++----- 2 files changed, 11 insertions(+), 15 deletions(-) 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 03cb6c313387..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 @@ -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) { @@ -114,14 +114,12 @@ public static Flux> cre return createSourceComponentFunction .apply(continuationToken, documentQueryParams) - .map(component -> { - return new AggregateDocumentQueryExecutionContext(component, - new ArrayList<>(aggregates), - groupByAliasToAggregateType, - groupByAliases, - hasSelectValue, - continuationToken); - }); + .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 b535065fdd3e..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 @@ -71,11 +71,9 @@ public static Flux> cre final UInt128 continuationTokenLastHash = distinctContinuationToken.getLastHash(); - return createSourceComponentFunction.apply(distinctContinuationToken.getSourceToken(), documentQueryParams).map(component -> { - return new DistinctDocumentQueryExecutionContext(component, - distinctQueryType, - continuationTokenLastHash); - }); + return createSourceComponentFunction + .apply(distinctContinuationToken.getSourceToken(), documentQueryParams) + .map(component -> new DistinctDocumentQueryExecutionContext(component, distinctQueryType, continuationTokenLastHash)); } IDocumentQueryExecutionComponent getComponent() { From 67500145529a27805ed3b572c574f083a44fb201 Mon Sep 17 00:00:00 2001 From: Annie Liang Date: Tue, 30 Jun 2020 10:22:41 -0700 Subject: [PATCH 6/6] fix compilation --- .../query/DocumentQueryExecutionContextFactory.java | 2 +- .../query/OrderByDocumentQueryExecutionContext.java | 2 +- .../query/ParallelDocumentQueryExecutionContext.java | 2 +- .../implementation/query/PipelinedDocumentQueryParams.java | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) 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 9d216b897e31..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,7 +195,7 @@ public static Flux documentQueryParams = new PipelinedDocumentQueryParams( + PipelinedDocumentQueryParams documentQueryParams = new PipelinedDocumentQueryParams( resourceTypeEnum, resourceType, query, 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 ca1ee0b2b234..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,7 +77,7 @@ private OrderByDocumentQueryExecutionContext( public static Flux> createAsync( IDocumentQueryClient client, - PipelinedDocumentQueryParams initParams) { + PipelinedDocumentQueryParams initParams) { OrderByDocumentQueryExecutionContext context = new OrderByDocumentQueryExecutionContext( client, 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 9450d46b17ae..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,7 +60,7 @@ private ParallelDocumentQueryExecutionContext( public static Flux> createAsync( IDocumentQueryClient client, - PipelinedDocumentQueryParams initParams) { + PipelinedDocumentQueryParams initParams) { ParallelDocumentQueryExecutionContext context = new ParallelDocumentQueryExecutionContext( client, 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 index ac11edc397d8..f9b256d5482a 100644 --- 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 @@ -29,7 +29,7 @@ public class PipelinedDocumentQueryParams { public PipelinedDocumentQueryParams( ResourceType resourceTypeEnum, - Class resourceType, + Class resourceType, SqlQuerySpec query, String resourceLink, String collectionRid,