Add result level caching to Brokers

docs/content/configuration/broker.md

@@ -108,6 +108,9 @@ You can optionally only configure caching to be enabled on the broker by setting
 |--------|---------------|-----------|-------|
 |`druid.broker.cache.useCache`|true, false|Enable the cache on the broker.|false|
 |`druid.broker.cache.populateCache`|true, false|Populate the cache on the broker.|false|
+|`druid.broker.cache.useResultLevelCache`|true, false|Enable result level caching on the broker.|false|
+|`druid.broker.cache.populateResultLevelCache`|true, false|Populate the result level cache on the broker.|false|
+|`druid.broker.cache.resultLevelCacheLimit`|positive integer|Maximum size of query response that can be cached.|`Integer.MAX_VALUE`|
 |`druid.broker.cache.unCacheable`|All druid query types|All query types to not cache.|`["groupBy", "select"]`|
 |`druid.broker.cache.cacheBulkMergeLimit`|positive integer or 0|Queries with more segments than this number will not attempt to fetch from cache at the broker level, leaving potential caching fetches (and cache result merging) to the historicals|`Integer.MAX_VALUE`|
 

docs/content/querying/caching.md

@@ -3,9 +3,10 @@ layout: doc_page
 ---
 # Query Caching
 
-Druid supports query result caching through an LRU cache. Results are stored on a per segment basis, along with the 
-parameters of a given query. This allows Druid to return final results based partially on segment results in the cache and partially 
-on segment results from scanning historical/real-time segments.
+Druid supports query result caching through an LRU cache. Results are stored as a whole or either on a per segment basis along with the 
+parameters of a given query. Segment level caching allows Druid to return final results based partially on segment results in the cache 
+and partially on segment results from scanning historical/real-time segments. Result level caching enables Druid to cache the entire 
+result set, so that query results can be completely retrieved from the cache for identical queries.
 
 Segment results can be stored in a local heap cache or in an external distributed key/value store. Segment query caches 
 can be enabled at either the Historical and Broker level (it is not recommended to enable caching on both).
@@ -15,6 +16,7 @@ can be enabled at either the Historical and Broker level (it is not recommended
 Enabling caching on the broker can yield faster results than if query caches were enabled on Historicals for small clusters. This is 
 the recommended setup for smaller production clusters (< 20 servers). Take note that when caching is enabled on the Broker, 
 results from Historicals are returned on a per segment basis, and Historicals will not be able to do any local result merging.
+Result level caching is enabled only on the Broker side.
 
 ## Query caching on Historicals
 

docs/content/querying/query-context.md

@@ -15,6 +15,8 @@ The query context is used for various query configuration parameters. The follow
 |queryId          | auto-generated                         | Unique identifier given to this query. If a query ID is set or known, this can be used to cancel the query |
 |useCache         | `true`                                 | Flag indicating whether to leverage the query cache for this query. When set to false, it disables reading from the query cache for this query. When set to true, Druid uses druid.broker.cache.useCache or druid.historical.cache.useCache to determine whether or not to read from the query cache |
 |populateCache    | `true`                                 | Flag indicating whether to save the results of the query to the query cache. Primarily used for debugging. When set to false, it disables saving the results of this query to the query cache. When set to true, Druid uses druid.broker.cache.populateCache or druid.historical.cache.populateCache to determine whether or not to save the results of this query to the query cache |
+|useResultLevelCache         | `true`                                 | Flag indicating whether to leverage the result level cache for this query. When set to false, it disables reading from the query cache for this query. When set to true, Druid uses druid.broker.cache.useResultLevelCache to determine whether or not to read from the query cache |
+|populateResultLevelCache    | `true`                                 | Flag indicating whether to save the results of the query to the result level cache. Primarily used for debugging. When set to false, it disables saving the results of this query to the query cache. When set to true, Druid uses druid.broker.cache.populateCache to determine whether or not to save the results of this query to the query cache |
 |bySegment        | `false`                                | Return "by segment" results. Primarily used for debugging, setting it to `true` returns results associated with the data segment they came from |
 |finalize         | `true`                                 | Flag indicating whether to "finalize" aggregation results. Primarily used for debugging. For instance, the `hyperUnique` aggregator will return the full HyperLogLog sketch instead of the estimated cardinality when this flag is set to `false` |
 |chunkPeriod      | `P0D` (off)                            | At the broker node level, long interval queries (of any type) may be broken into shorter interval queries to parallelize merging more than normal. Broken up queries will use a larger share of cluster resources, but may be able to complete faster as a result. Use ISO 8601 periods. For example, if this property is set to `P1M` (one month), then a query covering a year would be broken into 12 smaller queries. The broker uses its query processing executor service to initiate processing for query chunks, so make sure "druid.processing.numThreads" is configured appropriately on the broker. [groupBy queries](groupbyquery.html) do not support chunkPeriod by default, although they do if using the legacy "v1" engine. |

processing/src/main/java/io/druid/query/CacheStrategy.java

@@ -61,14 +61,16 @@ public interface CacheStrategy<T, CacheType, QueryType extends Query<T>>
    *
    * The resulting function must be thread-safe.
    *
+   * @param isResultLevelCache indicates whether the function is invoked for result-level caching or segment-level caching
    * @return a thread-safe function that converts the QueryType's result type into something cacheable
    */
-  Function<T, CacheType> prepareForCache();
+  Function<T, CacheType> prepareForCache(boolean isResultLevelCache);
 
   /**
    * A function that does the inverse of the operation that the function prepareForCache returns
    *
+   * @param isResultLevelCache indicates whether the function is invoked for result-level caching or segment-level caching
    * @return A function that does the inverse of the operation that the function prepareForCache returns
    */
-  Function<CacheType, T> pullFromCache();
+  Function<CacheType, T> pullFromCache(boolean isResultLevelCache);
 }

processing/src/main/java/io/druid/query/QueryContexts.java

@@ -37,6 +37,8 @@ public class QueryContexts
   public static final boolean DEFAULT_BY_SEGMENT = false;
   public static final boolean DEFAULT_POPULATE_CACHE = true;
   public static final boolean DEFAULT_USE_CACHE = true;
+  public static final boolean DEFAULT_POPULATE_RESULTLEVEL_CACHE = true;
+  public static final boolean DEFAULT_USE_RESULTLEVEL_CACHE = true;
   public static final int DEFAULT_PRIORITY = 0;
   public static final int DEFAULT_UNCOVERED_INTERVALS_LIMIT = 0;
   public static final long DEFAULT_TIMEOUT_MILLIS = 300_000; // 5 minutes
@@ -72,6 +74,26 @@ public static <T> boolean isUseCache(Query<T> query, boolean defaultValue)
     return parseBoolean(query, "useCache", defaultValue);
   }
 
+  public static <T> boolean isPopulateResultLevelCache(Query<T> query)
+  {
+    return isPopulateResultLevelCache(query, DEFAULT_POPULATE_RESULTLEVEL_CACHE);
+  }
+
+  public static <T> boolean isPopulateResultLevelCache(Query<T> query, boolean defaultValue)
+  {
+    return parseBoolean(query, "populateResultLevelCache", defaultValue);
+  }
+
+  public static <T> boolean isUseResultLevelCache(Query<T> query)
+  {
+    return isUseResultLevelCache(query, DEFAULT_USE_RESULTLEVEL_CACHE);
+  }
+
+  public static <T> boolean isUseResultLevelCache(Query<T> query, boolean defaultValue)
+  {
+    return parseBoolean(query, "useResultLevelCache", defaultValue);
+  }
+
   public static <T> boolean isFinalize(Query<T> query, boolean defaultValue)
   {
     return parseBoolean(query, "finalize", defaultValue);

processing/src/main/java/io/druid/query/groupby/GroupByQueryQueryToolChest.java

@@ -50,6 +50,7 @@
 import io.druid.query.aggregation.AggregatorFactory;
 import io.druid.query.aggregation.MetricManipulationFn;
 import io.druid.query.aggregation.MetricManipulatorFns;
+import io.druid.query.aggregation.PostAggregator;
 import io.druid.query.cache.CacheKeyBuilder;
 import io.druid.query.dimension.DefaultDimensionSpec;
 import io.druid.query.dimension.DimensionSpec;
@@ -408,7 +409,7 @@ public TypeReference<Object> getCacheObjectClazz()
       }
 
       @Override
-      public Function<Row, Object> prepareForCache()
+      public Function<Row, Object> prepareForCache(boolean isResultLevelCache)
       {
         return new Function<Row, Object>()
         {
@@ -426,6 +427,11 @@ public Object apply(Row input)
               for (AggregatorFactory agg : aggs) {
                 retVal.add(event.get(agg.getName()));
               }
+              if (isResultLevelCache) {
+                for (PostAggregator postAgg : query.getPostAggregatorSpecs()) {
+                  retVal.add(event.get(postAgg.getName()));
+                }
+              }
               return retVal;
             }
 
@@ -435,7 +441,7 @@ public Object apply(Row input)
       }
 
       @Override
-      public Function<Object, Row> pullFromCache()
+      public Function<Object, Row> pullFromCache(boolean isResultLevelCache)
       {
         return new Function<Object, Row>()
         {
@@ -460,7 +466,12 @@ public Row apply(Object input)
               final AggregatorFactory factory = aggsIter.next();
               event.put(factory.getName(), factory.deserialize(results.next()));
             }
-
+            if (isResultLevelCache) {
+              Iterator<PostAggregator> postItr = query.getPostAggregatorSpecs().iterator();
+              while (postItr.hasNext() && results.hasNext()) {
+                event.put(postItr.next().getName(), results.next());
+              }
+            }
             if (dimsIter.hasNext() || aggsIter.hasNext() || results.hasNext()) {
               throw new ISE(
                   "Found left over objects while reading from cache!! dimsIter[%s] aggsIter[%s] results[%s]",

processing/src/main/java/io/druid/query/metadata/SegmentMetadataQueryQueryToolChest.java

@@ -198,7 +198,7 @@ public TypeReference<SegmentAnalysis> getCacheObjectClazz()
       }
 
       @Override
-      public Function<SegmentAnalysis, SegmentAnalysis> prepareForCache()
+      public Function<SegmentAnalysis, SegmentAnalysis> prepareForCache(boolean isResultLevelCache)
       {
         return new Function<SegmentAnalysis, SegmentAnalysis>()
         {
@@ -211,7 +211,7 @@ public SegmentAnalysis apply(@Nullable SegmentAnalysis input)
       }
 
       @Override
-      public Function<SegmentAnalysis, SegmentAnalysis> pullFromCache()
+      public Function<SegmentAnalysis, SegmentAnalysis> pullFromCache(boolean isResultLevelCache)
       {
         return new Function<SegmentAnalysis, SegmentAnalysis>()
         {

processing/src/main/java/io/druid/query/search/SearchQueryQueryToolChest.java

@@ -206,7 +206,7 @@ public TypeReference<Object> getCacheObjectClazz()
       }
 
       @Override
-      public Function<Result<SearchResultValue>, Object> prepareForCache()
+      public Function<Result<SearchResultValue>, Object> prepareForCache(boolean isResultLevelCache)
       {
         return new Function<Result<SearchResultValue>, Object>()
         {
@@ -221,7 +221,7 @@ public Object apply(Result<SearchResultValue> input)
       }
 
       @Override
-      public Function<Object, Result<SearchResultValue>> pullFromCache()
+      public Function<Object, Result<SearchResultValue>> pullFromCache(boolean isResultLevelCache)
       {
         return new Function<Object, Result<SearchResultValue>>()
         {

processing/src/main/java/io/druid/query/select/SelectQueryQueryToolChest.java

@@ -243,7 +243,7 @@ public TypeReference<Object> getCacheObjectClazz()
       }
 
       @Override
-      public Function<Result<SelectResultValue>, Object> prepareForCache()
+      public Function<Result<SelectResultValue>, Object> prepareForCache(boolean isResultLevelCache)
       {
         return new Function<Result<SelectResultValue>, Object>()
         {
@@ -272,7 +272,7 @@ public Object apply(final Result<SelectResultValue> input)
       }
 
       @Override
-      public Function<Object, Result<SelectResultValue>> pullFromCache()
+      public Function<Object, Result<SelectResultValue>> pullFromCache(boolean isResultLevelCache)
       {
         return new Function<Object, Result<SelectResultValue>>()
         {

processing/src/main/java/io/druid/query/timeboundary/TimeBoundaryQueryQueryToolChest.java

@@ -171,7 +171,7 @@ public TypeReference<Object> getCacheObjectClazz()
       }
 
       @Override
-      public Function<Result<TimeBoundaryResultValue>, Object> prepareForCache()
+      public Function<Result<TimeBoundaryResultValue>, Object> prepareForCache(boolean isResultLevelCache)
       {
         return new Function<Result<TimeBoundaryResultValue>, Object>()
         {
@@ -184,7 +184,7 @@ public Object apply(Result<TimeBoundaryResultValue> input)
       }
 
       @Override
-      public Function<Object, Result<TimeBoundaryResultValue>> pullFromCache()
+      public Function<Object, Result<TimeBoundaryResultValue>> pullFromCache(boolean isResultLevelCache)
       {
         return new Function<Object, Result<TimeBoundaryResultValue>>()
         {

processing/src/main/java/io/druid/query/timeseries/TimeseriesQueryQueryToolChest.java

@@ -174,7 +174,7 @@ public TypeReference<Object> getCacheObjectClazz()
       }
 
       @Override
-      public Function<Result<TimeseriesResultValue>, Object> prepareForCache()
+      public Function<Result<TimeseriesResultValue>, Object> prepareForCache(boolean isResultLevelCache)
       {
         return new Function<Result<TimeseriesResultValue>, Object>()
         {
@@ -188,14 +188,18 @@ public Object apply(final Result<TimeseriesResultValue> input)
             for (AggregatorFactory agg : aggs) {
               retVal.add(results.getMetric(agg.getName()));
             }
-
+            if (isResultLevelCache) {
+              for (PostAggregator postAgg : query.getPostAggregatorSpecs()) {
+                retVal.add(results.getMetric(postAgg.getName()));
+              }
+            }
             return retVal;
           }
         };
       }
 
       @Override
-      public Function<Object, Result<TimeseriesResultValue>> pullFromCache()
+      public Function<Object, Result<TimeseriesResultValue>> pullFromCache(boolean isResultLevelCache)
       {
         return new Function<Object, Result<TimeseriesResultValue>>()
         {
@@ -216,6 +220,12 @@ public Result<TimeseriesResultValue> apply(@Nullable Object input)
               final AggregatorFactory factory = aggsIter.next();
               retVal.put(factory.getName(), factory.deserialize(resultIter.next()));
             }
+            if (isResultLevelCache) {
+              Iterator<PostAggregator> postItr = query.getPostAggregatorSpecs().iterator();
+              while (postItr.hasNext() && resultIter.hasNext()) {
+                retVal.put(postItr.next().getName(), resultIter.next());
+              }
+            }
 
             return new Result<TimeseriesResultValue>(
                 timestamp,

processing/src/main/java/io/druid/query/topn/TopNQueryQueryToolChest.java

@@ -293,7 +293,6 @@ public TypeReference<Result<TopNResultValue>> getResultTypeReference()
   }
 
 
-
   @Override
   public CacheStrategy<Result<TopNResultValue>, Object, TopNQuery> getCacheStrategy(final TopNQuery query)
   {
@@ -341,7 +340,7 @@ public TypeReference<Object> getCacheObjectClazz()
       }
 
       @Override
-      public Function<Result<TopNResultValue>, Object> prepareForCache()
+      public Function<Result<TopNResultValue>, Object> prepareForCache(boolean isResultLevelCache)
       {
         return new Function<Result<TopNResultValue>, Object>()
         {
@@ -361,6 +360,11 @@ public Object apply(final Result<TopNResultValue> input)
               for (String aggName : aggFactoryNames) {
                 vals.add(result.getMetric(aggName));
               }
+              if (isResultLevelCache) {
+                for (PostAggregator postAgg : query.getPostAggregatorSpecs()) {
+                  vals.add(result.getMetric(postAgg.getName()));
+                }
+              }
               retVal.add(vals);
             }
             return retVal;
@@ -369,7 +373,7 @@ public Object apply(final Result<TopNResultValue> input)
       }
 
       @Override
-      public Function<Object, Result<TopNResultValue>> pullFromCache()
+      public Function<Object, Result<TopNResultValue>> pullFromCache(boolean isResultLevelCache)
       {
         return new Function<Object, Result<TopNResultValue>>()
         {
@@ -401,7 +405,12 @@ public Result<TopNResultValue> apply(Object input)
               for (PostAggregator postAgg : postAggs) {
                 vals.put(postAgg.getName(), postAgg.compute(vals));
               }
-
+              if (isResultLevelCache) {
+                Iterator<PostAggregator> postItr = query.getPostAggregatorSpecs().iterator();
+                while (postItr.hasNext() && resultIter.hasNext()) {
+                  vals.put(postItr.next().getName(), resultIter.next());
+                }
+              }
               retVal.add(vals);
             }
 

processing/src/test/java/io/druid/query/metadata/SegmentMetadataQueryQueryToolChestTest.java

@@ -89,15 +89,15 @@ public void testCacheStrategy() throws Exception
         null
     );
 
-    Object preparedValue = strategy.prepareForCache().apply(result);
+    Object preparedValue = strategy.prepareForCache(false).apply(result);
 
     ObjectMapper objectMapper = new DefaultObjectMapper();
     SegmentAnalysis fromCacheValue = objectMapper.readValue(
         objectMapper.writeValueAsBytes(preparedValue),
         strategy.getCacheObjectClazz()
     );
 
-    SegmentAnalysis fromCacheResult = strategy.pullFromCache().apply(fromCacheValue);
+    SegmentAnalysis fromCacheResult = strategy.pullFromCache(false).apply(fromCacheValue);
 
     Assert.assertEquals(result, fromCacheResult);
   }

processing/src/test/java/io/druid/query/search/SearchQueryQueryToolChestTest.java

@@ -59,7 +59,7 @@ public void testCacheStrategy() throws Exception
         new SearchResultValue(ImmutableList.of(new SearchHit("dim1", "a")))
     );
 
-    Object preparedValue = strategy.prepareForCache().apply(
+    Object preparedValue = strategy.prepareForCache(false).apply(
         result
     );
 
@@ -69,7 +69,7 @@ public void testCacheStrategy() throws Exception
         strategy.getCacheObjectClazz()
     );
 
-    Result<SearchResultValue> fromCacheResult = strategy.pullFromCache().apply(fromCacheValue);
+    Result<SearchResultValue> fromCacheResult = strategy.pullFromCache(false).apply(fromCacheValue);
 
     Assert.assertEquals(result, fromCacheResult);
   }

processing/src/test/java/io/druid/query/timeboundary/TimeBoundaryQueryQueryToolChestTest.java

@@ -215,7 +215,7 @@ public void testCacheStrategy() throws Exception
     )
     );
 
-    Object preparedValue = strategy.prepareForCache().apply(
+    Object preparedValue = strategy.prepareForCache(false).apply(
         result
     );
 
@@ -225,7 +225,7 @@ public void testCacheStrategy() throws Exception
         strategy.getCacheObjectClazz()
     );
 
-    Result<TimeBoundaryResultValue> fromCacheResult = strategy.pullFromCache().apply(fromCacheValue);
+    Result<TimeBoundaryResultValue> fromCacheResult = strategy.pullFromCache(false).apply(fromCacheValue);
 
     Assert.assertEquals(result, fromCacheResult);
   }