Skip to content

Commit 01ad248

Browse files
mbhaskarmbhaskar
authored
** DONOT MERGE ** Enabling queryplan cache by default (Azure#24673)
* Enabling queryplan cache by default - Queryplan cache, an LRU cache of size 1000 would try to cache query plan for simple queries
1 parent e0d5e3e commit 01ad248

File tree

9 files changed

+85
-25
lines changed

9 files changed

+85
-25
lines changed

‎sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@
2727
import java.net.URI;
2828
import java.net.URISyntaxException;
2929
import java.util.List;
30-
import java.util.concurrent.ConcurrentMap;
30+
import java.util.Map;
3131

3232
/**
3333
* Provides a client-side logical representation of the Azure Cosmos DB
@@ -1570,7 +1570,7 @@ Flux<FeedResponse<Document>> readAllDocuments(
15701570
CosmosQueryRequestOptions options
15711571
);
15721572

1573-
ConcurrentMap<String, PartitionedQueryExecutionInfo> getQueryPlanCache();
1573+
Map<String, PartitionedQueryExecutionInfo> getQueryPlanCache();
15741574

15751575
/**
15761576
* Gets the collection cache.

‎sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -235,7 +235,7 @@ public static String getEnvironmentName() {
235235

236236
public static boolean isQueryPlanCachingEnabled() {
237237
// Queryplan caching will be disabled by default
238-
return getJVMConfigAsBoolean(QUERYPLAN_CACHING_ENABLED, false);
238+
return getJVMConfigAsBoolean(QUERYPLAN_CACHING_ENABLED, true);
239239
}
240240

241241
public static int getAddressRefreshResponseTimeoutInSeconds() {

‎sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Constants.java

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -245,4 +245,6 @@ public static final class PartitionedQueryExecutionInfo {
245245
public static final class QueryExecutionContext {
246246
public static final String INCREMENTAL_FEED_HEADER_VALUE = "Incremental feed";
247247
}
248+
249+
public static final int QUERYPLAN_CACHE_SIZE = 1000;
248250
}

‎sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
import com.azure.cosmos.implementation.batch.PartitionKeyRangeServerBatchRequest;
1717
import com.azure.cosmos.implementation.batch.ServerBatchRequest;
1818
import com.azure.cosmos.implementation.batch.SinglePartitionKeyServerBatchRequest;
19+
import com.azure.cosmos.implementation.caches.SizeLimitingLRUCache;
1920
import com.azure.cosmos.implementation.caches.RxClientCollectionCache;
2021
import com.azure.cosmos.implementation.caches.RxCollectionCache;
2122
import com.azure.cosmos.implementation.caches.RxPartitionKeyRangeCache;
@@ -142,7 +143,7 @@ public class RxDocumentClientImpl implements AsyncDocumentClient, IAuthorization
142143
private RxPartitionKeyRangeCache partitionKeyRangeCache;
143144
private Map<String, List<PartitionKeyAndResourceTokenPair>> resourceTokensMap;
144145
private final boolean contentResponseOnWriteEnabled;
145-
private ConcurrentMap<String, PartitionedQueryExecutionInfo> queryPlanCache;
146+
private Map<String, PartitionedQueryExecutionInfo> queryPlanCache;
146147

147148
private final AtomicBoolean closed = new AtomicBoolean(false);
148149
private final int clientId;
@@ -362,7 +363,7 @@ private RxDocumentClientImpl(URI serviceEndpoint,
362363
this.retryPolicy = new RetryPolicy(this, this.globalEndpointManager, this.connectionPolicy);
363364
this.resetSessionTokenRetryPolicy = retryPolicy;
364365
CpuMemoryMonitor.register(this);
365-
this.queryPlanCache = new ConcurrentHashMap<>();
366+
this.queryPlanCache = Collections.synchronizedMap(new SizeLimitingLRUCache(Constants.QUERYPLAN_CACHE_SIZE));
366367
} catch (RuntimeException e) {
367368
logger.error("unexpected failure in initializing client.", e);
368369
close();
@@ -2462,7 +2463,7 @@ public Flux<FeedResponse<Document>> readAllDocuments(
24622463
}
24632464

24642465
@Override
2465-
public ConcurrentMap<String, PartitionedQueryExecutionInfo> getQueryPlanCache() {
2466+
public Map<String, PartitionedQueryExecutionInfo> getQueryPlanCache() {
24662467
return queryPlanCache;
24672468
}
24682469

‎sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/caches/SizeLimitingLRUCache.java

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
// Copyright (c) Microsoft Corporation. All rights reserved.
2+
// Licensed under the MIT License.
3+
package com.azure.cosmos.implementation.caches;
4+
5+
import com.azure.cosmos.implementation.query.PartitionedQueryExecutionInfo;
6+
7+
import java.util.LinkedHashMap;
8+
import java.util.Map;
9+
10+
/**
11+
* LRU Cache using LinkedHashMap that limits the number of entries
12+
*/
13+
public class SizeLimitingLRUCache extends LinkedHashMap<String, PartitionedQueryExecutionInfo> {
14+
15+
private static final long serialVersionUID = 1L;
16+
private final int maxEntries;
17+
18+
public SizeLimitingLRUCache(int maxEntries) {
19+
this.maxEntries = maxEntries;
20+
}
21+
22+
public SizeLimitingLRUCache(int initialCapacity, float loadFactor, int maxEntries) {
23+
super(initialCapacity, loadFactor);
24+
this.maxEntries = maxEntries;
25+
}
26+
27+
public SizeLimitingLRUCache(
28+
Map<? extends String, ? extends PartitionedQueryExecutionInfo> m, int maxEntries) {
29+
super(m);
30+
this.maxEntries = maxEntries;
31+
}
32+
33+
public SizeLimitingLRUCache(int initialCapacity, float loadFactor, boolean accessOrder, int maxEntries) {
34+
super(initialCapacity, loadFactor, accessOrder);
35+
this.maxEntries = maxEntries;
36+
}
37+
38+
public SizeLimitingLRUCache(int initialCapacity, int maxEntries) {
39+
super(initialCapacity);
40+
this.maxEntries = maxEntries;
41+
}
42+
43+
@Override
44+
protected boolean removeEldestEntry(
45+
Map.Entry<String, PartitionedQueryExecutionInfo> eldest) {
46+
return size() > maxEntries;
47+
}
48+
}

‎sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java

Lines changed: 8 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,6 @@
3434
import java.util.List;
3535
import java.util.Map;
3636
import java.util.UUID;
37-
import java.util.concurrent.ConcurrentMap;
3837
import java.util.stream.Collectors;
3938

4039
/**
@@ -71,7 +70,7 @@ private static <T extends Resource> Mono<Pair<List<Range<String>>,QueryInfo>> ge
7170
String resourceLink,
7271
DocumentCollection collection,
7372
DefaultDocumentQueryExecutionContext<T> queryExecutionContext, boolean queryPlanCachingEnabled,
74-
ConcurrentMap<String, PartitionedQueryExecutionInfo> queryPlanCache) {
73+
Map<String, PartitionedQueryExecutionInfo> queryPlanCache) {
7574

7675
// The partitionKeyRangeIdInternal is no more a public API on
7776
// FeedOptions, but have the below condition
@@ -99,6 +98,7 @@ private static <T extends Resource> Mono<Pair<List<Range<String>>,QueryInfo>> ge
9998
Instant endTime = Instant.now(); // endTime for query plan diagnostics
10099
PartitionedQueryExecutionInfo partitionedQueryExecutionInfo = queryPlanCache.get(query.getQueryText());
101100
if (partitionedQueryExecutionInfo != null) {
101+
logger.info("Skipping query plan round trip by using the cached plan");
102102
return getTargetRangesFromQueryPlan(cosmosQueryRequestOptions, collection, queryExecutionContext,
103103
partitionedQueryExecutionInfo, startTime, endTime);
104104
}
@@ -117,7 +117,7 @@ private static <T extends Resource> Mono<Pair<List<Range<String>>,QueryInfo>> ge
117117

118118
Instant endTime = Instant.now();
119119

120-
if (queryPlanCachingEnabled) {
120+
if (queryPlanCachingEnabled && isScopedToSinglePartition(cosmosQueryRequestOptions)) {
121121
tryCacheQueryPlan(query, partitionedQueryExecutionInfo, queryPlanCache);
122122
}
123123

@@ -169,14 +169,9 @@ private static <T extends Resource> Mono<Pair<List<Range<String>>, QueryInfo>> g
169169
private static void tryCacheQueryPlan(
170170
SqlQuerySpec query,
171171
PartitionedQueryExecutionInfo partitionedQueryExecutionInfo,
172-
ConcurrentMap<String, PartitionedQueryExecutionInfo> queryPlanCache) {
172+
Map<String, PartitionedQueryExecutionInfo> queryPlanCache) {
173173
QueryInfo queryInfo = partitionedQueryExecutionInfo.getQueryInfo();
174174
if (canCacheQuery(queryInfo) && !queryPlanCache.containsKey(query.getQueryText())) {
175-
if (queryPlanCache.size() > MAX_CACHE_SIZE) {
176-
// Clearing query plan cache if size is above max size. This can be optimized in future by using
177-
// a threadsafe LRU cache
178-
queryPlanCache.clear();
179-
}
180175
queryPlanCache.put(query.getQueryText(), partitionedQueryExecutionInfo);
181176
}
182177
}
@@ -188,7 +183,9 @@ private static boolean canCacheQuery(QueryInfo queryInfo) {
188183
&& !queryInfo.hasGroupBy()
189184
&& !queryInfo.hasLimit()
190185
&& !queryInfo.hasTop()
191-
&& !queryInfo.hasOffset();
186+
&& !queryInfo.hasOffset()
187+
&& !queryInfo.hasDCount()
188+
&& !queryInfo.hasOrderBy();
192189
}
193190

194191
private static boolean isScopedToSinglePartition(CosmosQueryRequestOptions cosmosQueryRequestOptions) {
@@ -208,7 +205,7 @@ public static <T extends Resource> Flux<? extends IDocumentQueryExecutionContext
208205
boolean isContinuationExpected,
209206
UUID correlatedActivityId,
210207
boolean queryPlanCachingEnabled,
211-
ConcurrentMap<String, PartitionedQueryExecutionInfo> queryPlanCache) {
208+
Map<String, PartitionedQueryExecutionInfo> queryPlanCache) {
212209

213210
// return proxy
214211
Flux<Utils.ValueHolder<DocumentCollection>> collectionObs = Flux.just(new Utils.ValueHolder<>(null));

‎sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/QueryValidationTests.java

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -239,7 +239,6 @@ public void queryOptionNullValidation() {
239239
private Object[][] query() {
240240
return new Object[][]{
241241
new Object[] { "Select * from c "},
242-
new Object[] { "select * from c order by c.prop ASC"},
243242
};
244243
}
245244

@@ -299,7 +298,6 @@ public void queryPlanCacheSinglePartitionParameterizedQueriesCorrectness() {
299298
createdContainer);
300299
documentsInserted.addAll(pk2Docs);
301300

302-
303301
CosmosQueryRequestOptions options = new CosmosQueryRequestOptions();
304302
options.setPartitionKey(new PartitionKey(pk2));
305303

@@ -325,6 +323,20 @@ public void queryPlanCacheSinglePartitionParameterizedQueriesCorrectness() {
325323
// Top query should not be cached
326324
assertThat(contextClient.getQueryPlanCache().containsKey(sqlQuerySpec.getQueryText())).isFalse();
327325

326+
// group by should not be cached
327+
sqlQuerySpec.setQueryText("select max(c.id) from c order by c.name group by c.name");
328+
values1 = queryAndGetResults(sqlQuerySpec, options, TestObject.class);
329+
assertThat(contextClient.getQueryPlanCache().containsKey(sqlQuerySpec.getQueryText())).isFalse();
330+
331+
// distinct queries should not be cached
332+
sqlQuerySpec.setQueryText("SELECT distinct c.name from c");
333+
values1 = queryAndGetResults(sqlQuerySpec, options, TestObject.class);
334+
assertThat(contextClient.getQueryPlanCache().containsKey(sqlQuerySpec.getQueryText())).isFalse();
335+
336+
//order by query should not be cached
337+
sqlQuerySpec.setQueryText("select * from c order by c.name");
338+
values1 = queryAndGetResults(sqlQuerySpec, options, TestObject.class);
339+
assertThat(contextClient.getQueryPlanCache().containsKey(sqlQuerySpec.getQueryText())).isFalse();
328340
}
329341

330342
@Test(groups = {"simple"}, timeOut = TIMEOUT * 40)

‎sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/CosmosTemplatePartitionIT.java

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -42,8 +42,8 @@
4242

4343
import java.util.Collections;
4444
import java.util.List;
45+
import java.util.Map;
4546
import java.util.UUID;
46-
import java.util.concurrent.ConcurrentMap;
4747

4848
import static com.azure.spring.data.cosmos.common.TestConstants.ADDRESSES;
4949
import static com.azure.spring.data.cosmos.common.TestConstants.FIRST_NAME;
@@ -144,7 +144,7 @@ public void testFindWithPartitionWithQueryPlanCachingEnabled() {
144144

145145
CosmosAsyncClient cosmosAsyncClient = cosmosFactory.getCosmosAsyncClient();
146146
AsyncDocumentClient asyncDocumentClient = CosmosBridgeInternal.getAsyncDocumentClient(cosmosAsyncClient);
147-
ConcurrentMap<String, PartitionedQueryExecutionInfo> initialCache = asyncDocumentClient.getQueryPlanCache();
147+
Map<String, PartitionedQueryExecutionInfo> initialCache = asyncDocumentClient.getQueryPlanCache();
148148
assertThat(initialCache.containsKey(sqlQuerySpec.getQueryText())).isTrue();
149149
int initialSize = initialCache.size();
150150

@@ -158,7 +158,7 @@ public void testFindWithPartitionWithQueryPlanCachingEnabled() {
158158
result = TestUtils.toList(cosmosTemplate.find(query, PartitionPerson.class,
159159
PartitionPerson.class.getSimpleName()));
160160

161-
ConcurrentMap<String, PartitionedQueryExecutionInfo> postQueryCallCache = asyncDocumentClient.getQueryPlanCache();
161+
Map<String, PartitionedQueryExecutionInfo> postQueryCallCache = asyncDocumentClient.getQueryPlanCache();
162162
assertThat(postQueryCallCache.containsKey(sqlQuerySpec.getQueryText())).isTrue();
163163
assertThat(postQueryCallCache.size()).isEqualTo(initialSize);
164164
assertThat(result.size()).isEqualTo(1);

‎sdk/cosmos/azure-spring-data-cosmos-test/src/test/java/com/azure/spring/data/cosmos/core/ReactiveCosmosTemplatePartitionIT.java

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -39,8 +39,8 @@
3939
import reactor.test.StepVerifier;
4040

4141
import java.util.Collections;
42+
import java.util.Map;
4243
import java.util.UUID;
43-
import java.util.concurrent.ConcurrentMap;
4444

4545
import static org.assertj.core.api.Assertions.assertThat;
4646
import static org.hamcrest.Matchers.equalTo;
@@ -126,7 +126,7 @@ public void testFindWithPartitionWithQueryPlanCachingEnabled() {
126126

127127
CosmosAsyncClient cosmosAsyncClient = cosmosFactory.getCosmosAsyncClient();
128128
AsyncDocumentClient asyncDocumentClient = CosmosBridgeInternal.getAsyncDocumentClient(cosmosAsyncClient);
129-
ConcurrentMap<String, PartitionedQueryExecutionInfo> initialCache = asyncDocumentClient.getQueryPlanCache();
129+
Map<String, PartitionedQueryExecutionInfo> initialCache = asyncDocumentClient.getQueryPlanCache();
130130
assertThat(initialCache.containsKey(sqlQuerySpec.getQueryText())).isTrue();
131131
int initialSize = initialCache.size();
132132

@@ -145,7 +145,7 @@ public void testFindWithPartitionWithQueryPlanCachingEnabled() {
145145
Assert.assertThat(actual.getZipCode(), is(equalTo(TEST_PERSON_2.getZipCode())));
146146
}).verifyComplete();
147147

148-
ConcurrentMap<String, PartitionedQueryExecutionInfo> postQueryCallCache = asyncDocumentClient.getQueryPlanCache();
148+
Map<String, PartitionedQueryExecutionInfo> postQueryCallCache = asyncDocumentClient.getQueryPlanCache();
149149
assertThat(postQueryCallCache.containsKey(sqlQuerySpec.getQueryText())).isTrue();
150150
assertThat(postQueryCallCache.size()).isEqualTo(initialSize);
151151
}

0 commit comments

Comments
 (0)