From da914b835ef9407cb703744b883751ff4d1f0bba Mon Sep 17 00:00:00 2001
From: Eric Tschetter <cheddar@metamarketsgroup.com>
Date: Tue, 29 Jan 2013 13:04:57 -0600
Subject: [PATCH] 1) Add caching for SegmentMetadataQuery

---
 .../druid/client/CachingClusteredClient.java  |  7 ++-
 .../com/metamx/druid/query/CacheStrategy.java |  9 ++-
 .../metamx/druid/query/QueryToolChest.java    |  2 +-
 .../group/GroupByQueryQueryToolChest.java     |  2 +-
 .../query/metadata/AllColumnIncluderator.java |  6 ++
 .../query/metadata/ColumnIncluderator.java    |  5 ++
 .../metadata/ListColumnIncluderator.java      | 26 ++++++++
 .../metadata/NoneColumnIncluderator.java      |  6 ++
 .../SegmentMetadataQueryQueryToolChest.java   | 59 +++++++++++++++++--
 .../search/SearchQueryQueryToolChest.java     | 14 ++++-
 .../TimeBoundaryQueryQueryToolChest.java      | 13 +++-
 .../TimeseriesQueryQueryToolChest.java        | 15 +++--
 .../druid/coordination/ServerManagerTest.java |  2 +-
 13 files changed, 145 insertions(+), 21 deletions(-)

diff --git a/client/src/main/java/com/metamx/druid/client/CachingClusteredClient.java b/client/src/main/java/com/metamx/druid/client/CachingClusteredClient.java
index 5d640dc4e38..163f1986a53 100644
--- a/client/src/main/java/com/metamx/druid/client/CachingClusteredClient.java
+++ b/client/src/main/java/com/metamx/druid/client/CachingClusteredClient.java
@@ -54,6 +54,7 @@ import com.metamx.druid.query.segment.SegmentDescriptor;
 import com.metamx.druid.result.BySegmentResultValueClass;
 import com.metamx.druid.result.Result;
 import org.codehaus.jackson.map.ObjectMapper;
+import org.codehaus.jackson.type.TypeReference;
 import org.joda.time.DateTime;
 import org.joda.time.Interval;
 
@@ -110,7 +111,7 @@ public class CachingClusteredClient<T> implements QueryRunner<T>
   public Sequence<T> run(final Query<T> query)
   {
     final QueryToolChest<T, Query<T>> toolChest = warehouse.getToolChest(query);
-    final CacheStrategy<T, Query<T>> strategy = toolChest.getCacheStrategy(query);
+    final CacheStrategy<T, Object, Query<T>> strategy = toolChest.getCacheStrategy(query);
 
     final Map<DruidServer, List<SegmentDescriptor>> serverSegments = Maps.newTreeMap();
 
@@ -241,6 +242,7 @@ public class CachingClusteredClient<T> implements QueryRunner<T>
             }
 
             final Function<Object, T> pullFromCacheFunction = strategy.pullFromCache();
+            final TypeReference<Object> cacheObjectClazz = strategy.getCacheObjectClazz();
             for (Pair<DateTime, byte[]> cachedResultPair : cachedResults) {
               final byte[] cachedResult = cachedResultPair.rhs;
               Sequence<Object> cachedSequence = new BaseSequence<Object, Iterator<Object>>(
@@ -255,7 +257,8 @@ public class CachingClusteredClient<T> implements QueryRunner<T>
                         }
 
                         return objectMapper.readValues(
-                            objectMapper.getJsonFactory().createJsonParser(cachedResult), Object.class
+                            objectMapper.getJsonFactory().createJsonParser(cachedResult),
+                            cacheObjectClazz
                         );
                       }
                       catch (IOException e) {
diff --git a/client/src/main/java/com/metamx/druid/query/CacheStrategy.java b/client/src/main/java/com/metamx/druid/query/CacheStrategy.java
index abdbe4da259..f8f5098f6ca 100644
--- a/client/src/main/java/com/metamx/druid/query/CacheStrategy.java
+++ b/client/src/main/java/com/metamx/druid/query/CacheStrategy.java
@@ -22,16 +22,19 @@ package com.metamx.druid.query;
 import com.google.common.base.Function;
 import com.metamx.common.guava.Sequence;
 import com.metamx.druid.Query;
+import org.codehaus.jackson.type.TypeReference;
 
 /**
 */
-public interface CacheStrategy<T, QueryType extends Query<T>>
+public interface CacheStrategy<T, CacheType, QueryType extends Query<T>>
 {
   public byte[] computeCacheKey(QueryType query);
 
-  public Function<T, Object> prepareForCache();
+  public TypeReference<CacheType> getCacheObjectClazz();
 
-  public Function<Object, T> pullFromCache();
+  public Function<T, CacheType> prepareForCache();
+
+  public Function<CacheType, T> pullFromCache();
 
   public Sequence<T> mergeSequences(Sequence<Sequence<T>> seqOfSequences);
 }
diff --git a/client/src/main/java/com/metamx/druid/query/QueryToolChest.java b/client/src/main/java/com/metamx/druid/query/QueryToolChest.java
index ebf77f64af4..bec2170ec92 100644
--- a/client/src/main/java/com/metamx/druid/query/QueryToolChest.java
+++ b/client/src/main/java/com/metamx/druid/query/QueryToolChest.java
@@ -44,7 +44,7 @@ public interface QueryToolChest<ResultType, QueryType extends Query<ResultType>>
   public ServiceMetricEvent.Builder makeMetricBuilder(QueryType query);
   public Function<ResultType, ResultType> makeMetricManipulatorFn(QueryType query, MetricManipulationFn fn);
   public TypeReference<ResultType> getResultTypeReference();
-  public CacheStrategy<ResultType, QueryType> getCacheStrategy(QueryType query);
+  public <T> CacheStrategy<ResultType, T, QueryType> getCacheStrategy(QueryType query);
   public QueryRunner<ResultType> preMergeQueryDecoration(QueryRunner<ResultType> runner);
   public QueryRunner<ResultType> postMergeQueryDecoration(QueryRunner<ResultType> runner);
 }
diff --git a/client/src/main/java/com/metamx/druid/query/group/GroupByQueryQueryToolChest.java b/client/src/main/java/com/metamx/druid/query/group/GroupByQueryQueryToolChest.java
index 6ce071f8068..9dcf6110322 100644
--- a/client/src/main/java/com/metamx/druid/query/group/GroupByQueryQueryToolChest.java
+++ b/client/src/main/java/com/metamx/druid/query/group/GroupByQueryQueryToolChest.java
@@ -178,7 +178,7 @@ public class GroupByQueryQueryToolChest implements QueryToolChest<Row, GroupByQu
   }
 
   @Override
-  public CacheStrategy<Row, GroupByQuery> getCacheStrategy(GroupByQuery query)
+  public CacheStrategy<Row, Object, GroupByQuery> getCacheStrategy(GroupByQuery query)
   {
     return null;
   }
diff --git a/client/src/main/java/com/metamx/druid/query/metadata/AllColumnIncluderator.java b/client/src/main/java/com/metamx/druid/query/metadata/AllColumnIncluderator.java
index 8687f213503..cd96b5d718a 100644
--- a/client/src/main/java/com/metamx/druid/query/metadata/AllColumnIncluderator.java
+++ b/client/src/main/java/com/metamx/druid/query/metadata/AllColumnIncluderator.java
@@ -28,4 +28,10 @@ public class AllColumnIncluderator implements ColumnIncluderator
   {
     return true;
   }
+
+  @Override
+  public byte[] getCacheKey()
+  {
+    return ALL_CACHE_PREFIX;
+  }
 }
diff --git a/client/src/main/java/com/metamx/druid/query/metadata/ColumnIncluderator.java b/client/src/main/java/com/metamx/druid/query/metadata/ColumnIncluderator.java
index 466167d48fd..90533c4eaca 100644
--- a/client/src/main/java/com/metamx/druid/query/metadata/ColumnIncluderator.java
+++ b/client/src/main/java/com/metamx/druid/query/metadata/ColumnIncluderator.java
@@ -32,5 +32,10 @@ import org.codehaus.jackson.annotate.JsonTypeInfo;
 })
 public interface ColumnIncluderator
 {
+  public static final byte[] NONE_CACHE_PREFIX = new byte[]{0x0};
+  public static final byte[] ALL_CACHE_PREFIX = new byte[]{0x1};
+  public static final byte[] LIST_CACHE_PREFIX = new byte[]{0x2};
+
   public boolean include(String columnName);
+  public byte[] getCacheKey();
 }
diff --git a/client/src/main/java/com/metamx/druid/query/metadata/ListColumnIncluderator.java b/client/src/main/java/com/metamx/druid/query/metadata/ListColumnIncluderator.java
index 4048dcd420a..e74661d6822 100644
--- a/client/src/main/java/com/metamx/druid/query/metadata/ListColumnIncluderator.java
+++ b/client/src/main/java/com/metamx/druid/query/metadata/ListColumnIncluderator.java
@@ -19,10 +19,16 @@
 
 package com.metamx.druid.query.metadata;
 
+import com.google.common.base.Charsets;
+import com.google.common.base.Throwables;
+import com.google.common.collect.Lists;
 import com.google.common.collect.Sets;
 import org.codehaus.jackson.annotate.JsonCreator;
 import org.codehaus.jackson.annotate.JsonProperty;
 
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.nio.ByteBuffer;
 import java.util.Collections;
 import java.util.List;
 import java.util.Set;
@@ -53,4 +59,24 @@ public class ListColumnIncluderator implements ColumnIncluderator
   {
     return columns.contains(columnName);
   }
+
+  @Override
+  public byte[] getCacheKey()
+  {
+    int size = 1;
+    List<byte[]> columns = Lists.newArrayListWithExpectedSize(this.columns.size());
+
+    for (String column : this.columns) {
+      final byte[] bytes = column.getBytes(Charsets.UTF_8);
+      columns.add(bytes);
+      size += bytes.length;
+    }
+
+    final ByteBuffer bytes = ByteBuffer.allocate(size).put(LIST_CACHE_PREFIX);
+    for (byte[] column : columns) {
+      bytes.put(column);
+    }
+
+    return bytes.array();
+  }
 }
diff --git a/client/src/main/java/com/metamx/druid/query/metadata/NoneColumnIncluderator.java b/client/src/main/java/com/metamx/druid/query/metadata/NoneColumnIncluderator.java
index 56e6842b7ad..d1d66d26778 100644
--- a/client/src/main/java/com/metamx/druid/query/metadata/NoneColumnIncluderator.java
+++ b/client/src/main/java/com/metamx/druid/query/metadata/NoneColumnIncluderator.java
@@ -28,4 +28,10 @@ public class NoneColumnIncluderator implements ColumnIncluderator
   {
     return false;
   }
+
+  @Override
+  public byte[] getCacheKey()
+  {
+    return NONE_CACHE_PREFIX;
+  }
 }
diff --git a/client/src/main/java/com/metamx/druid/query/metadata/SegmentMetadataQueryQueryToolChest.java b/client/src/main/java/com/metamx/druid/query/metadata/SegmentMetadataQueryQueryToolChest.java
index 447bba2e5c6..160c23cd958 100644
--- a/client/src/main/java/com/metamx/druid/query/metadata/SegmentMetadataQueryQueryToolChest.java
+++ b/client/src/main/java/com/metamx/druid/query/metadata/SegmentMetadataQueryQueryToolChest.java
@@ -27,7 +27,7 @@ import com.google.common.collect.Maps;
 import com.google.common.collect.Ordering;
 import com.google.common.collect.Sets;
 import com.metamx.common.ISE;
-import com.metamx.common.guava.Comparators;
+import com.metamx.common.guava.MergeSequence;
 import com.metamx.common.guava.Sequence;
 import com.metamx.common.guava.nary.BinaryFn;
 import com.metamx.druid.Query;
@@ -44,8 +44,7 @@ import org.joda.time.Interval;
 import org.joda.time.Minutes;
 
 import javax.annotation.Nullable;
-import java.util.Comparator;
-import java.util.Iterator;
+import java.nio.ByteBuffer;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -54,6 +53,7 @@ import java.util.Set;
 public class SegmentMetadataQueryQueryToolChest implements QueryToolChest<SegmentAnalysis, SegmentMetadataQuery>
 {
   private static final TypeReference<SegmentAnalysis> TYPE_REFERENCE = new TypeReference<SegmentAnalysis>(){};
+  private static final byte[] SEGMENT_METADATA_CACHE_PREFIX = new byte[]{0x4};
 
   @Override
   public QueryRunner<SegmentAnalysis> mergeResults(final QueryRunner<SegmentAnalysis> runner)
@@ -165,9 +165,58 @@ public class SegmentMetadataQueryQueryToolChest implements QueryToolChest<Segmen
   }
 
   @Override
-  public CacheStrategy<SegmentAnalysis, SegmentMetadataQuery> getCacheStrategy(SegmentMetadataQuery query)
+  public CacheStrategy<SegmentAnalysis, SegmentAnalysis, SegmentMetadataQuery> getCacheStrategy(SegmentMetadataQuery query)
   {
-    return null;
+    return new CacheStrategy<SegmentAnalysis, SegmentAnalysis, SegmentMetadataQuery>()
+    {
+      @Override
+      public byte[] computeCacheKey(SegmentMetadataQuery query)
+      {
+        byte[] includerBytes = query.getToInclude().getCacheKey();
+        return ByteBuffer.allocate(1 + includerBytes.length)
+                         .put(SEGMENT_METADATA_CACHE_PREFIX)
+                         .put(includerBytes)
+                         .array();
+      }
+
+      @Override
+      public TypeReference<SegmentAnalysis> getCacheObjectClazz()
+      {
+        return getResultTypeReference();
+      }
+
+      @Override
+      public Function<SegmentAnalysis, SegmentAnalysis> prepareForCache()
+      {
+        return new Function<SegmentAnalysis, SegmentAnalysis>()
+        {
+          @Override
+          public SegmentAnalysis apply(@Nullable SegmentAnalysis input)
+          {
+            return input;
+          }
+        };
+      }
+
+      @Override
+      public Function<SegmentAnalysis, SegmentAnalysis> pullFromCache()
+      {
+        return new Function<SegmentAnalysis, SegmentAnalysis>()
+        {
+          @Override
+          public SegmentAnalysis apply(@Nullable SegmentAnalysis input)
+          {
+            return input;
+          }
+        };
+      }
+
+      @Override
+      public Sequence<SegmentAnalysis> mergeSequences(Sequence<Sequence<SegmentAnalysis>> seqOfSequences)
+      {
+        return new MergeSequence<SegmentAnalysis>(getOrdering(), seqOfSequences);
+      }
+    };
   }
 
   @Override
diff --git a/client/src/main/java/com/metamx/druid/query/search/SearchQueryQueryToolChest.java b/client/src/main/java/com/metamx/druid/query/search/SearchQueryQueryToolChest.java
index 95757fc60b1..ce3fcc86114 100644
--- a/client/src/main/java/com/metamx/druid/query/search/SearchQueryQueryToolChest.java
+++ b/client/src/main/java/com/metamx/druid/query/search/SearchQueryQueryToolChest.java
@@ -82,6 +82,10 @@ public class SearchQueryQueryToolChest implements QueryToolChest<Result<SearchRe
     maxSearchLimit = PropUtils.getPropertyAsInt(props, "com.metamx.query.search.maxSearchLimit", 1000);
   }
 
+  private static final TypeReference<Object> OBJECT_TYPE_REFERENCE = new TypeReference<Object>()
+  {
+  };
+
   @Override
   public QueryRunner<Result<SearchResultValue>> mergeResults(QueryRunner<Result<SearchResultValue>> runner)
   {
@@ -143,9 +147,9 @@ public class SearchQueryQueryToolChest implements QueryToolChest<Result<SearchRe
   }
 
   @Override
-  public CacheStrategy<Result<SearchResultValue>, SearchQuery> getCacheStrategy(SearchQuery query)
+  public CacheStrategy<Result<SearchResultValue>, Object, SearchQuery> getCacheStrategy(SearchQuery query)
   {
-    return new CacheStrategy<Result<SearchResultValue>, SearchQuery>()
+    return new CacheStrategy<Result<SearchResultValue>, Object, SearchQuery>()
     {
       @Override
       public byte[] computeCacheKey(SearchQuery query)
@@ -183,6 +187,12 @@ public class SearchQueryQueryToolChest implements QueryToolChest<Result<SearchRe
         return queryCacheKey.array();
       }
 
+      @Override
+      public TypeReference<Object> getCacheObjectClazz()
+      {
+        return OBJECT_TYPE_REFERENCE;
+      }
+
       @Override
       public Function<Result<SearchResultValue>, Object> prepareForCache()
       {
diff --git a/client/src/main/java/com/metamx/druid/query/timeboundary/TimeBoundaryQueryQueryToolChest.java b/client/src/main/java/com/metamx/druid/query/timeboundary/TimeBoundaryQueryQueryToolChest.java
index 9d65ab9b47c..5ee6c321bbb 100644
--- a/client/src/main/java/com/metamx/druid/query/timeboundary/TimeBoundaryQueryQueryToolChest.java
+++ b/client/src/main/java/com/metamx/druid/query/timeboundary/TimeBoundaryQueryQueryToolChest.java
@@ -53,6 +53,9 @@ public class TimeBoundaryQueryQueryToolChest
   private static final TypeReference<Result<TimeBoundaryResultValue>> TYPE_REFERENCE = new TypeReference<Result<TimeBoundaryResultValue>>()
   {
   };
+  private static final TypeReference<Object> OBJECT_TYPE_REFERENCE = new TypeReference<Object>()
+  {
+  };
 
   @Override
   public QueryRunner<Result<TimeBoundaryResultValue>> mergeResults(
@@ -106,9 +109,9 @@ public class TimeBoundaryQueryQueryToolChest
   }
 
   @Override
-  public CacheStrategy<Result<TimeBoundaryResultValue>, TimeBoundaryQuery> getCacheStrategy(TimeBoundaryQuery query)
+  public CacheStrategy<Result<TimeBoundaryResultValue>, Object, TimeBoundaryQuery> getCacheStrategy(TimeBoundaryQuery query)
   {
-    return new CacheStrategy<Result<TimeBoundaryResultValue>, TimeBoundaryQuery>()
+    return new CacheStrategy<Result<TimeBoundaryResultValue>, Object, TimeBoundaryQuery>()
     {
       @Override
       public byte[] computeCacheKey(TimeBoundaryQuery query)
@@ -119,6 +122,12 @@ public class TimeBoundaryQueryQueryToolChest
                          .array();
       }
 
+      @Override
+      public TypeReference<Object> getCacheObjectClazz()
+      {
+        return OBJECT_TYPE_REFERENCE;
+      }
+
       @Override
       public Function<Result<TimeBoundaryResultValue>, Object> prepareForCache()
       {
diff --git a/client/src/main/java/com/metamx/druid/query/timeseries/TimeseriesQueryQueryToolChest.java b/client/src/main/java/com/metamx/druid/query/timeseries/TimeseriesQueryQueryToolChest.java
index d8ef9802dad..9c633507ec5 100644
--- a/client/src/main/java/com/metamx/druid/query/timeseries/TimeseriesQueryQueryToolChest.java
+++ b/client/src/main/java/com/metamx/druid/query/timeseries/TimeseriesQueryQueryToolChest.java
@@ -68,6 +68,9 @@ public class TimeseriesQueryQueryToolChest implements QueryToolChest<Result<Time
   private static final TypeReference<Result<TimeseriesResultValue>> TYPE_REFERENCE = new TypeReference<Result<TimeseriesResultValue>>()
   {
   };
+  private static final TypeReference<Object> OBJECT_TYPE_REFERENCE = new TypeReference<Object>()
+  {
+  };
 
   @Override
   public QueryRunner<Result<TimeseriesResultValue>> mergeResults(QueryRunner<Result<TimeseriesResultValue>> queryRunner)
@@ -155,9 +158,9 @@ public class TimeseriesQueryQueryToolChest implements QueryToolChest<Result<Time
   }
 
   @Override
-  public CacheStrategy<Result<TimeseriesResultValue>, TimeseriesQuery> getCacheStrategy(final TimeseriesQuery query)
+  public CacheStrategy<Result<TimeseriesResultValue>, Object, TimeseriesQuery> getCacheStrategy(final TimeseriesQuery query)
   {
-    return new CacheStrategy<Result<TimeseriesResultValue>, TimeseriesQuery>()
+    return new CacheStrategy<Result<TimeseriesResultValue>, Object, TimeseriesQuery>()
     {
       private final List<AggregatorFactory> aggs = query.getAggregatorSpecs();
       private final List<PostAggregator> postAggs = query.getPostAggregatorSpecs();
@@ -179,6 +182,12 @@ public class TimeseriesQueryQueryToolChest implements QueryToolChest<Result<Time
             .array();
       }
 
+      @Override
+      public TypeReference<Object> getCacheObjectClazz()
+      {
+        return OBJECT_TYPE_REFERENCE;
+      }
+
       @Override
       public Function<Result<TimeseriesResultValue>, Object> prepareForCache()
       {
@@ -259,6 +268,4 @@ public class TimeseriesQueryQueryToolChest implements QueryToolChest<Result<Time
   {
     return Ordering.natural();
   }
-
-
 }
diff --git a/server/src/test/java/com/metamx/druid/coordination/ServerManagerTest.java b/server/src/test/java/com/metamx/druid/coordination/ServerManagerTest.java
index 43dc956c40e..ed8912d3371 100644
--- a/server/src/test/java/com/metamx/druid/coordination/ServerManagerTest.java
+++ b/server/src/test/java/com/metamx/druid/coordination/ServerManagerTest.java
@@ -421,7 +421,7 @@ public class ServerManagerTest
     }
 
     @Override
-    public CacheStrategy<T, QueryType> getCacheStrategy(QueryType query)
+    public <Typer> CacheStrategy<T, Typer, QueryType> getCacheStrategy(QueryType query)
     {
       return null;
     }