SketchAggregatorFactory.combine(..) returns Union object now so that it can be reused across multiple combine(..) calls (#3471)

extensions-core/datasketches/src/main/java/io/druid/query/aggregation/datasketches/theta/SketchAggregator.java

@@ -107,6 +107,8 @@ public class SketchAggregator implements Aggregator
       union.update((Memory) update);
     } else if (update instanceof Sketch) {
       union.update((Sketch) update);
+    } else if (update instanceof Union) {
+      union.update(((Union) update).getResult(false, null));
     } else if (update instanceof String) {
       union.update((String) update);
     } else if (update instanceof byte[]) {

extensions-core/datasketches/src/main/java/io/druid/query/aggregation/datasketches/theta/SketchAggregatorFactory.java

@@ -110,10 +110,21 @@ public abstract class SketchAggregatorFactory extends AggregatorFactory
   @Override
   public Object combine(Object lhs, Object rhs)
   {
-    Union union = (Union) SetOperation.builder().build(size, Family.UNION);
+    final Union union;
-    updateUnion(union, lhs);
+    if (lhs instanceof Union) {
-    updateUnion(union, rhs);
+      union = (Union) lhs;
-    return union.getResult(false, null);
+      updateUnion(union, rhs);
+    } else if (rhs instanceof Union) {
+      union = (Union) rhs;
+      updateUnion(union, lhs);
+    } else {
+      union = (Union) SetOperation.builder().build(size, Family.UNION);
+      updateUnion(union, lhs);
+      updateUnion(union, rhs);
+    }
+
+
+    return union;
   }
 
   private void updateUnion(Union union, Object obj)
@@ -124,6 +135,8 @@ public abstract class SketchAggregatorFactory extends AggregatorFactory
       union.update((Memory) obj);
     } else if (obj instanceof Sketch) {
       union.update((Sketch) obj);
+    } else if (obj instanceof Union) {
+      union.update(((Union) obj).getResult(false, null));
     } else {
       throw new IAE("Object of type [%s] can not be unioned", obj.getClass().getName());
     }

extensions-core/datasketches/src/main/java/io/druid/query/aggregation/datasketches/theta/SketchEstimatePostAggregator.java

@@ -82,7 +82,7 @@ public class SketchEstimatePostAggregator implements PostAggregator
   @Override
   public Object compute(Map<String, Object> combinedAggregators)
   {
-    Sketch sketch = (Sketch) field.compute(combinedAggregators);
+    Sketch sketch = SketchSetPostAggregator.toSketch(field.compute(combinedAggregators));
     if (errorBoundsStdDev != null) {
       SketchEstimateWithErrorBounds result = new SketchEstimateWithErrorBounds(
           sketch.getEstimate(),

extensions-core/datasketches/src/main/java/io/druid/query/aggregation/datasketches/theta/SketchModule.java

@@ -25,6 +25,7 @@ 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 com.yahoo.sketches.theta.Union;
 import io.druid.initialization.DruidModule;
 import io.druid.segment.serde.ComplexMetrics;
 
@@ -71,7 +72,11 @@ public class SketchModule implements DruidModule
                 Sketch.class, new SketchJsonSerializer()
             )
             .addSerializer(
-                Memory.class, new MemoryJsonSerializer())
+                Memory.class, new MemoryJsonSerializer()
+            )
+            .addSerializer(
+                Union.class, new UnionJsonSerializer()
+            )
     );
   }
 }

extensions-core/datasketches/src/main/java/io/druid/query/aggregation/datasketches/theta/SketchObjectStrategy.java

@@ -26,6 +26,7 @@ import com.yahoo.sketches.memory.MemoryRegion;
 import com.yahoo.sketches.memory.NativeMemory;
 import com.yahoo.sketches.theta.Sketch;
 import com.yahoo.sketches.theta.Sketches;
+import com.yahoo.sketches.theta.Union;
 import io.druid.segment.data.ObjectStrategy;
 
 import java.nio.ByteBuffer;
@@ -98,6 +99,8 @@ public class SketchObjectStrategy implements ObjectStrategy
       byte[] retVal = new byte[(int) mem.getCapacity()];
       mem.getByteArray(0, retVal, 0, (int) mem.getCapacity());
       return retVal;
+    }  else if (obj instanceof Union) {
+      return toBytes(((Union) obj).getResult(true, null));
     } else if (obj == null) {
       return EMPTY_BYTES;
     } else {

extensions-core/datasketches/src/main/java/io/druid/query/aggregation/datasketches/theta/SketchSetPostAggregator.java

@@ -26,6 +26,7 @@ import com.metamx.common.IAE;
 import com.metamx.common.logger.Logger;
 import com.yahoo.sketches.Util;
 import com.yahoo.sketches.theta.Sketch;
+import com.yahoo.sketches.theta.Union;
 import io.druid.query.aggregation.PostAggregator;
 
 import java.util.Comparator;
@@ -83,12 +84,23 @@ public class SketchSetPostAggregator implements PostAggregator
   {
     Sketch[] sketches = new Sketch[fields.size()];
     for (int i = 0; i < sketches.length; i++) {
-      sketches[i] = (Sketch) fields.get(i).compute(combinedAggregators);
+      sketches[i] = toSketch(fields.get(i).compute(combinedAggregators));
     }
 
     return SketchOperations.sketchSetOperation(func, maxSketchSize, sketches);
   }
 
+  public final static Sketch toSketch(Object obj)
+  {
+    if (obj instanceof Sketch) {
+      return (Sketch) obj;
+    } else if (obj instanceof Union) {
+      return ((Union) obj).getResult(true, null);
+    } else {
+      throw new IAE("Can't convert to Sketch object [%s]", obj.getClass());
+    }
+  }
+
   @Override
   @JsonProperty
   public String getName()

extensions-core/datasketches/src/main/java/io/druid/query/aggregation/datasketches/theta/UnionJsonSerializer.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.theta.Union;
+
+import java.io.IOException;
+
+/**
+ */
+public class UnionJsonSerializer extends JsonSerializer<Union>
+{
+  @Override
+  public void serialize(Union union, JsonGenerator jgen, SerializerProvider provider)
+      throws IOException, JsonProcessingException
+  {
+    jgen.writeBinary(union.getResult(true, null).toByteArray());
+  }
+}