From c6634d7c2c6bbb570b697ddcffdc69ca06709d62 Mon Sep 17 00:00:00 2001
From: Himanshu Gupta <g.himanshu@gmail.com>
Date: Tue, 5 Jan 2016 00:01:19 -0600
Subject: [PATCH] adding json for thetaSketch Memory object representation

---
 .../theta/MemoryJsonSerializer.java           | 40 +++++++++++++++++++
 .../datasketches/theta/SketchModule.java      |  7 +++-
 .../datasketches/theta/SketchOperations.java  |  7 +++-
 3 files changed, 52 insertions(+), 2 deletions(-)
 create mode 100644 extensions/datasketches/src/main/java/io/druid/query/aggregation/datasketches/theta/MemoryJsonSerializer.java

diff --git a/extensions/datasketches/src/main/java/io/druid/query/aggregation/datasketches/theta/MemoryJsonSerializer.java b/extensions/datasketches/src/main/java/io/druid/query/aggregation/datasketches/theta/MemoryJsonSerializer.java
new file mode 100644
index 00000000000..2590c1fc3fb
--- /dev/null
+++ b/extensions/datasketches/src/main/java/io/druid/query/aggregation/datasketches/theta/MemoryJsonSerializer.java
@@ -0,0 +1,40 @@
+/*
+ * Licensed to Metamarkets Group Inc. (Metamarkets) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. Metamarkets licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package io.druid.query.aggregation.datasketches.theta;
+
+import com.fasterxml.jackson.core.JsonGenerator;
+import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.JsonSerializer;
+import com.fasterxml.jackson.databind.SerializerProvider;
+import com.yahoo.sketches.memory.Memory;
+
+import java.io.IOException;
+
+/**
+ */
+public class MemoryJsonSerializer extends JsonSerializer<Memory>
+{
+  @Override
+  public void serialize(Memory mem, JsonGenerator jgen, SerializerProvider provider)
+      throws IOException, JsonProcessingException
+  {
+    jgen.writeBinary(SketchOperations.deserializeFromMemory(mem).toByteArray());
+  }
+}
diff --git a/extensions/datasketches/src/main/java/io/druid/query/aggregation/datasketches/theta/SketchModule.java b/extensions/datasketches/src/main/java/io/druid/query/aggregation/datasketches/theta/SketchModule.java
index 1969494963d..84b0853cf22 100644
--- a/extensions/datasketches/src/main/java/io/druid/query/aggregation/datasketches/theta/SketchModule.java
+++ b/extensions/datasketches/src/main/java/io/druid/query/aggregation/datasketches/theta/SketchModule.java
@@ -23,6 +23,7 @@ import com.fasterxml.jackson.databind.Module;
 import com.fasterxml.jackson.databind.jsontype.NamedType;
 import com.fasterxml.jackson.databind.module.SimpleModule;
 import com.google.inject.Binder;
+import com.yahoo.sketches.memory.Memory;
 import com.yahoo.sketches.theta.Sketch;
 import io.druid.initialization.DruidModule;
 import io.druid.segment.serde.ComplexMetrics;
@@ -66,7 +67,11 @@ public class SketchModule implements DruidModule
                 new NamedType(SketchEstimatePostAggregator.class, THETA_SKETCH_ESTIMATE_POST_AGG),
                 new NamedType(SketchSetPostAggregator.class, THETA_SKETCH_SET_OP_POST_AGG)
             )
-            .addSerializer(Sketch.class, new SketchJsonSerializer())
+            .addSerializer(
+                Sketch.class, new SketchJsonSerializer()
+            )
+            .addSerializer(
+                Memory.class, new MemoryJsonSerializer())
     );
   }
 }
diff --git a/extensions/datasketches/src/main/java/io/druid/query/aggregation/datasketches/theta/SketchOperations.java b/extensions/datasketches/src/main/java/io/druid/query/aggregation/datasketches/theta/SketchOperations.java
index a40931bb030..b86ef333bc0 100644
--- a/extensions/datasketches/src/main/java/io/druid/query/aggregation/datasketches/theta/SketchOperations.java
+++ b/extensions/datasketches/src/main/java/io/druid/query/aggregation/datasketches/theta/SketchOperations.java
@@ -22,6 +22,7 @@ package io.druid.query.aggregation.datasketches.theta;
 import com.google.common.base.Charsets;
 import com.metamx.common.logger.Logger;
 import com.yahoo.sketches.Family;
+import com.yahoo.sketches.memory.Memory;
 import com.yahoo.sketches.memory.NativeMemory;
 import com.yahoo.sketches.theta.AnotB;
 import com.yahoo.sketches.theta.Intersection;
@@ -72,7 +73,11 @@ public class SketchOperations
 
   public static Sketch deserializeFromByteArray(byte[] data)
   {
-    NativeMemory mem = new NativeMemory(data);
+    return deserializeFromMemory(new NativeMemory(data));
+  }
+
+  public static Sketch deserializeFromMemory(Memory mem)
+  {
     if (Sketch.getSerializationVersion(mem) < 3) {
       return Sketches.heapifySketch(mem);
     } else {