Sketch cache key should include size, isInputThetaSketch. (apache#2893)

SnowCicala · Apr 28, 2016 · 909abd1 · 909abd1
1 parent 90ce03c
commit 909abd1
Show file tree

Hide file tree

Showing 3 changed files with 55 additions and 5 deletions.
diff --git a/.../src/main/java/io/druid/query/aggregation/datasketches/theta/SketchAggregatorFactory.java b/.../src/main/java/io/druid/query/aggregation/datasketches/theta/SketchAggregatorFactory.java
@@ -22,6 +22,7 @@
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.google.common.base.Preconditions;
 import com.google.common.primitives.Doubles;
+import com.google.common.primitives.Ints;
 import com.metamx.common.IAE;
 import com.yahoo.sketches.Family;
 import com.yahoo.sketches.Util;
@@ -169,7 +170,11 @@ public List<String> requiredFields()
   public byte[] getCacheKey()
   {
     byte[] fieldNameBytes = fieldName.getBytes();
-    return ByteBuffer.allocate(1 + fieldNameBytes.length).put(cacheId).put(fieldNameBytes).array();
+    return ByteBuffer.allocate(1 + Ints.BYTES + fieldNameBytes.length)
+                     .put(cacheId)
+                     .putInt(size)
+                     .put(fieldNameBytes)
+                     .array();
   }
 
   @Override

diff --git a/...main/java/io/druid/query/aggregation/datasketches/theta/SketchMergeAggregatorFactory.java b/...main/java/io/druid/query/aggregation/datasketches/theta/SketchMergeAggregatorFactory.java
@@ -25,6 +25,7 @@
 import io.druid.query.aggregation.AggregatorFactory;
 import io.druid.query.aggregation.AggregatorFactoryNotMergeableException;
 
+import java.nio.ByteBuffer;
 import java.util.Collections;
 import java.util.List;
 
@@ -104,7 +105,7 @@ public boolean getIsInputThetaSketch()
   {
     return isInputThetaSketch;
   }
-  
+
   @JsonProperty
   public Integer getErrorBoundsStdDev()
   {
@@ -126,9 +127,9 @@ public Object finalizeComputation(Object object)
       Sketch sketch = (Sketch) object;
       if (errorBoundsStdDev != null) {
         SketchEstimateWithErrorBounds result = new SketchEstimateWithErrorBounds(
-            sketch.getEstimate(), 
-            sketch.getUpperBound(errorBoundsStdDev), 
-            sketch.getLowerBound(errorBoundsStdDev), 
+            sketch.getEstimate(),
+            sketch.getUpperBound(errorBoundsStdDev),
+            sketch.getLowerBound(errorBoundsStdDev),
             errorBoundsStdDev);
         return result;
       } else {
@@ -149,6 +150,17 @@ public String getTypeName()
     }
   }
 
+  @Override
+  public byte[] getCacheKey()
+  {
+    final byte[] superCacheKey = super.getCacheKey();
+
+    return ByteBuffer.allocate(superCacheKey.length + 1)
+                     .put(superCacheKey)
+                     .put(isInputThetaSketch ? (byte) 1 : (byte) 0)
+                     .array();
+  }
+
   @Override
   public boolean equals(Object o)
   {

diff --git a/...es/src/test/java/io/druid/query/aggregation/datasketches/theta/SketchAggregationTest.java b/...es/src/test/java/io/druid/query/aggregation/datasketches/theta/SketchAggregationTest.java
@@ -46,6 +46,7 @@
 import java.io.File;
 import java.io.IOException;
 import java.nio.charset.Charset;
+import java.util.Arrays;
 import java.util.List;
 
 /**
@@ -384,6 +385,38 @@ public void testSketchSetPostAggregatorSerde() throws Exception
     );
   }
 
+  @Test
+  public void testCacheKey()
+  {
+    final SketchMergeAggregatorFactory factory1 = new SketchMergeAggregatorFactory(
+        "name",
+        "fieldName",
+        16,
+        null,
+        null,
+        null
+    );
+    final SketchMergeAggregatorFactory factory2 = new SketchMergeAggregatorFactory(
+        "name",
+        "fieldName",
+        16,
+        null,
+        null,
+        null
+    );
+    final SketchMergeAggregatorFactory factory3 = new SketchMergeAggregatorFactory(
+        "name",
+        "fieldName",
+        32,
+        null,
+        null,
+        null
+    );
+
+    Assert.assertTrue(Arrays.equals(factory1.getCacheKey(), factory2.getCacheKey()));
+    Assert.assertFalse(Arrays.equals(factory1.getCacheKey(), factory3.getCacheKey()));
+  }
+
   private void assertPostAggregatorSerde(PostAggregator agg) throws Exception
   {
     Assert.assertEquals(