Monomorphic processing of TopN queries with 1 and 2 aggregators (key part of #3798)

extensions-contrib/distinctcount/src/main/java/io/druid/query/aggregation/distinctcount/DistinctCountAggregatorFactory.java

@@ -83,7 +83,7 @@ public BufferAggregator factorizeBuffered(ColumnSelectorFactory columnFactory)
   {
     DimensionSelector selector = makeDimensionSelector(columnFactory);
     if (selector == null) {
-      return new EmptyDistinctCountBufferAggregator();
+      return EmptyDistinctCountBufferAggregator.instance();
     } else {
       return new DistinctCountBufferAggregator(makeDimensionSelector(columnFactory));
     }

extensions-contrib/distinctcount/src/main/java/io/druid/query/aggregation/distinctcount/DistinctCountBufferAggregator.java

@@ -22,16 +22,17 @@
 import io.druid.collections.bitmap.MutableBitmap;
 import io.druid.collections.bitmap.WrappedRoaringBitmap;
 import io.druid.query.aggregation.BufferAggregator;
+import io.druid.query.monomorphicprocessing.RuntimeShapeInspector;
 import io.druid.segment.DimensionSelector;
+import it.unimi.dsi.fastutil.ints.Int2ObjectMap;
+import it.unimi.dsi.fastutil.ints.Int2ObjectOpenHashMap;
 
 import java.nio.ByteBuffer;
-import java.util.HashMap;
-import java.util.Map;
 
 public class DistinctCountBufferAggregator implements BufferAggregator
 {
   private final DimensionSelector selector;
-  private final Map<Integer, MutableBitmap> mutableBitmapCollection = new HashMap<>();
+  private final Int2ObjectMap<MutableBitmap> mutableBitmapCollection = new Int2ObjectOpenHashMap<>();
 
   public DistinctCountBufferAggregator(
       DimensionSelector selector
@@ -89,4 +90,10 @@ public void close()
   {
     mutableBitmapCollection.clear();
   }
+
+  @Override
+  public void inspectRuntimeShape(RuntimeShapeInspector inspector)
+  {
+    inspector.visit("selector", selector);
+  }
 }

extensions-contrib/distinctcount/src/main/java/io/druid/query/aggregation/distinctcount/EmptyDistinctCountBufferAggregator.java

@@ -20,13 +20,24 @@
 package io.druid.query.aggregation.distinctcount;
 
 import io.druid.query.aggregation.BufferAggregator;
+import io.druid.query.monomorphicprocessing.RuntimeShapeInspector;
 
 import java.nio.ByteBuffer;
 
-public class EmptyDistinctCountBufferAggregator implements BufferAggregator
+/**
+ * The difference from {@link io.druid.query.aggregation.NoopBufferAggregator} is that
+ * EmptyDistinctCountBufferAggregator returns 0 instead of null from {@link #get(ByteBuffer, int)}.
+ */
+public final class EmptyDistinctCountBufferAggregator implements BufferAggregator
 {
+  private static final EmptyDistinctCountBufferAggregator INSTANCE = new EmptyDistinctCountBufferAggregator();
+
+  static EmptyDistinctCountBufferAggregator instance()
+  {
+    return INSTANCE;
+  }
 
-  public EmptyDistinctCountBufferAggregator()
+  private EmptyDistinctCountBufferAggregator()
   {
   }
 
@@ -62,4 +73,9 @@ public long getLong(ByteBuffer buf, int position)
   public void close()
   {
   }
+
+  @Override
+  public void inspectRuntimeShape(RuntimeShapeInspector inspector)
+  {
+  }
 }

extensions-contrib/time-min-max/src/main/java/io/druid/query/aggregation/TimestampBufferAggregator.java

@@ -20,6 +20,7 @@
 package io.druid.query.aggregation;
 
 import io.druid.data.input.impl.TimestampSpec;
+import io.druid.query.monomorphicprocessing.RuntimeShapeInspector;
 import io.druid.segment.ObjectColumnSelector;
 
 import java.nio.ByteBuffer;
@@ -79,7 +80,14 @@ public long getLong(ByteBuffer buf, int position)
   }
 
   @Override
-  public void close() {
+  public void close()
+  {
+  }
 
+  @Override
+  public void inspectRuntimeShape(RuntimeShapeInspector inspector)
+  {
+    inspector.visit("selector", selector);
+    inspector.visit("comparator", comparator);
   }
 }

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

@@ -20,12 +20,20 @@
 package io.druid.query.aggregation.datasketches.theta;
 
 import io.druid.query.aggregation.BufferAggregator;
+import io.druid.query.monomorphicprocessing.RuntimeShapeInspector;
 
 import java.nio.ByteBuffer;
 
-public class EmptySketchBufferAggregator implements BufferAggregator
+public final class EmptySketchBufferAggregator implements BufferAggregator
 {
-  public EmptySketchBufferAggregator()
+  private static final EmptySketchBufferAggregator INSTANCE = new EmptySketchBufferAggregator();
+
+  public static EmptySketchBufferAggregator instance()
+  {
+    return INSTANCE;
+  }
+
+  private EmptySketchBufferAggregator()
   {
   }
 
@@ -61,4 +69,9 @@ public long getLong(ByteBuffer buf, int position)
   public void close()
   {
   }
+
+  @Override
+  public void inspectRuntimeShape(RuntimeShapeInspector inspector)
+  {
+  }
 }

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

@@ -73,7 +73,7 @@ public BufferAggregator factorizeBuffered(ColumnSelectorFactory metricFactory)
   {
     ObjectColumnSelector selector = metricFactory.makeObjectColumnSelector(fieldName);
     if (selector == null) {
-      return new EmptySketchBufferAggregator();
+      return EmptySketchBufferAggregator.instance();
     } else {
       return new SketchBufferAggregator(selector, size, getMaxIntermediateSize());
     }

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

@@ -26,6 +26,7 @@
 import com.yahoo.sketches.theta.SetOperation;
 import com.yahoo.sketches.theta.Union;
 import io.druid.query.aggregation.BufferAggregator;
+import io.druid.query.monomorphicprocessing.RuntimeShapeInspector;
 import io.druid.segment.ObjectColumnSelector;
 
 import java.nio.ByteBuffer;
@@ -113,4 +114,9 @@ public void close()
     unions.clear();
   }
 
+  @Override
+  public void inspectRuntimeShape(RuntimeShapeInspector inspector)
+  {
+    inspector.visit("selector", selector);
+  }
 }

extensions-core/histogram/src/main/java/io/druid/query/aggregation/histogram/ApproximateHistogramBufferAggregator.java

@@ -20,6 +20,7 @@
 package io.druid.query.aggregation.histogram;
 
 import io.druid.query.aggregation.BufferAggregator;
+import io.druid.query.monomorphicprocessing.RuntimeShapeInspector;
 import io.druid.segment.FloatColumnSelector;
 
 import java.nio.ByteBuffer;
@@ -99,4 +100,10 @@ public void close()
   {
     // no resources to cleanup
   }
+
+  @Override
+  public void inspectRuntimeShape(RuntimeShapeInspector inspector)
+  {
+    inspector.visit("selector", selector);
+  }
 }

extensions-core/histogram/src/main/java/io/druid/query/aggregation/histogram/ApproximateHistogramFoldingBufferAggregator.java

@@ -20,6 +20,7 @@
 package io.druid.query.aggregation.histogram;
 
 import io.druid.query.aggregation.BufferAggregator;
+import io.druid.query.monomorphicprocessing.RuntimeShapeInspector;
 import io.druid.segment.ObjectColumnSelector;
 
 import java.nio.ByteBuffer;
@@ -102,4 +103,10 @@ public void close()
   {
     // no resources to cleanup
   }
+
+  @Override
+  public void inspectRuntimeShape(RuntimeShapeInspector inspector)
+  {
+    inspector.visit("selector", selector);
+  }
 }

extensions-core/stats/src/main/java/io/druid/query/aggregation/variance/VarianceAggregatorFactory.java

@@ -23,14 +23,14 @@
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.fasterxml.jackson.annotation.JsonTypeName;
 import com.google.common.base.Preconditions;
-
 import io.druid.java.util.common.IAE;
 import io.druid.java.util.common.StringUtils;
 import io.druid.query.aggregation.Aggregator;
 import io.druid.query.aggregation.AggregatorFactory;
 import io.druid.query.aggregation.AggregatorFactoryNotMergeableException;
-import io.druid.query.aggregation.Aggregators;
 import io.druid.query.aggregation.BufferAggregator;
+import io.druid.query.aggregation.NoopAggregator;
+import io.druid.query.aggregation.NoopBufferAggregator;
 import io.druid.segment.ColumnSelectorFactory;
 import io.druid.segment.ObjectColumnSelector;
 import org.apache.commons.codec.binary.Base64;
@@ -95,7 +95,7 @@ public Aggregator factorize(ColumnSelectorFactory metricFactory)
   {
     ObjectColumnSelector selector = metricFactory.makeObjectColumnSelector(fieldName);
     if (selector == null) {
-      return Aggregators.noopAggregator();
+      return NoopAggregator.instance();
     }
 
     if ("float".equalsIgnoreCase(inputType)) {
@@ -115,7 +115,7 @@ public BufferAggregator factorizeBuffered(ColumnSelectorFactory metricFactory)
   {
     ObjectColumnSelector selector = metricFactory.makeObjectColumnSelector(fieldName);
     if (selector == null) {
-      return Aggregators.noopBufferAggregator();
+      return NoopBufferAggregator.instance();
     }
     if ("float".equalsIgnoreCase(inputType)) {
       return new VarianceBufferAggregator.FloatVarianceAggregator(

extensions-core/stats/src/main/java/io/druid/query/aggregation/variance/VarianceBufferAggregator.java

@@ -22,6 +22,7 @@
 import com.google.common.primitives.Doubles;
 import com.google.common.primitives.Longs;
 import io.druid.query.aggregation.BufferAggregator;
+import io.druid.query.monomorphicprocessing.RuntimeShapeInspector;
 import io.druid.segment.FloatColumnSelector;
 import io.druid.segment.LongColumnSelector;
 import io.druid.segment.ObjectColumnSelector;
@@ -102,6 +103,12 @@ public void aggregate(ByteBuffer buf, int position)
         buf.putDouble(position + NVARIANCE_OFFSET, variance);
       }
     }
+
+    @Override
+    public void inspectRuntimeShape(RuntimeShapeInspector inspector)
+    {
+      inspector.visit("selector", selector);
+    }
   }
 
   public static final class LongVarianceAggregator extends VarianceBufferAggregator
@@ -128,6 +135,12 @@ public void aggregate(ByteBuffer buf, int position)
         buf.putDouble(position + NVARIANCE_OFFSET, variance);
       }
     }
+
+    @Override
+    public void inspectRuntimeShape(RuntimeShapeInspector inspector)
+    {
+      inspector.visit("selector", selector);
+    }
   }
 
   public static final class ObjectVarianceAggregator extends VarianceBufferAggregator
@@ -167,5 +180,11 @@ public void aggregate(ByteBuffer buf, int position)
       buf.putDouble(position + SUM_OFFSET, sum);
       buf.putDouble(position + NVARIANCE_OFFSET, nvariance);
     }
+
+    @Override
+    public void inspectRuntimeShape(RuntimeShapeInspector inspector)
+    {
+      inspector.visit("selector", selector);
+    }
   }
 }

extensions-core/stats/src/test/java/io/druid/query/aggregation/variance/VarianceAggregatorCollectorTest.java

@@ -20,10 +20,9 @@
 package io.druid.query.aggregation.variance;
 
 import com.google.common.collect.Lists;
-
 import io.druid.java.util.common.Pair;
-import io.druid.segment.FloatColumnSelector;
 import io.druid.segment.ObjectColumnSelector;
+import io.druid.segment.TestFloatColumnSelector;
 import org.junit.Assert;
 import org.junit.Test;
 
@@ -142,7 +141,7 @@ public void testVariance()
     }
   }
 
-  private static class FloatHandOver implements FloatColumnSelector
+  private static class FloatHandOver extends TestFloatColumnSelector
   {
     float v;
 

extensions-core/stats/src/test/java/io/druid/query/aggregation/variance/VarianceTopNQueryTest.java

@@ -54,7 +54,7 @@ public class VarianceTopNQueryTest
   @Parameterized.Parameters(name="{0}")
   public static Iterable<Object[]> constructorFeeder() throws IOException
   {
-    return TopNQueryRunnerTest.constructorFeeder();
+    return QueryRunnerTestHelper.transformToConstructionFeeder(TopNQueryRunnerTest.queryRunners());
   }
 
   private final QueryRunner runner;

pom.xml

@@ -638,6 +638,16 @@
                 <artifactId>RoaringBitmap</artifactId>
                 <version>0.5.18</version>
             </dependency>
+            <dependency>
+                <groupId>org.ow2.asm</groupId>
+                <artifactId>asm</artifactId>
+                <version>5.2</version>
+            </dependency>
+            <dependency>
+                <groupId>org.ow2.asm</groupId>
+                <artifactId>asm-commons</artifactId>
+                <version>5.2</version>
+            </dependency>
             <dependency>
                 <groupId>org.apache.calcite</groupId>
                 <artifactId>calcite-core</artifactId>
@@ -815,10 +825,11 @@
                                 <version>1.0</version>
                             </signature>
                             <ignores>
-                                <!-- Some of our code uses DirectBuffer & Cleaner directly, which are not part of
+                                <!-- Some of our code uses sun.* classes directly, which are not part of
                                      the JDK signature (although they are there anyway). -->
                                 <ignore>sun.nio.ch.DirectBuffer</ignore>
                                 <ignore>sun.misc.Cleaner</ignore>
+                                <ignore>sun.misc.Unsafe</ignore>
                             </ignores>
                         </configuration>
                     </execution>

processing/pom.xml

@@ -93,6 +93,14 @@
             <groupId>commons-lang</groupId>
             <artifactId>commons-lang</artifactId>
         </dependency>
+        <dependency>
+            <groupId>org.ow2.asm</groupId>
+            <artifactId>asm</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>org.ow2.asm</groupId>
+            <artifactId>asm-commons</artifactId>
+        </dependency>
 
         <!-- Tests -->
         <dependency>

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

@@ -97,6 +97,13 @@ private static <T> boolean parseBoolean(Query<T> query, String key, boolean defa
     }
   }
 
+  public static void checkInterrupted()
+  {
+    if (Thread.interrupted()) {
+      throw new QueryInterruptedException(new InterruptedException());
+    }
+  }
+
   public static final String QUERYID = "queryId";
   private final DataSource dataSource;
   private final boolean descending;