Add OakIncrementalIndex to Druid

benchmarks/src/main/java/org/apache/druid/benchmark/indexing/IncrementalIndexReadBenchmark.java

@@ -57,8 +57,13 @@
 import org.openjdk.jmh.annotations.Scope;
 import org.openjdk.jmh.annotations.Setup;
 import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.TearDown;
 import org.openjdk.jmh.annotations.Warmup;
 import org.openjdk.jmh.infra.Blackhole;
+import org.openjdk.jmh.runner.Runner;
+import org.openjdk.jmh.runner.RunnerException;
+import org.openjdk.jmh.runner.options.Options;
+import org.openjdk.jmh.runner.options.OptionsBuilder;
 
 import java.io.IOException;
 import java.util.ArrayList;
@@ -82,6 +87,9 @@ public class IncrementalIndexReadBenchmark
   @Param({"true", "false"})
   private boolean rollup;
 
+  @Param({"oak", "onheap"})
+  private String indexType;
+
   private static final Logger log = new Logger(IncrementalIndexReadBenchmark.class);
   private static final int RNG_SEED = 9999;
   private IncrementalIndex incIndex;
@@ -118,18 +126,30 @@ public void setup() throws IOException
 
   }
 
+  @TearDown
+  public void tearDown()
+  {
+    incIndex.close();
+  }
+
   private IncrementalIndex makeIncIndex()
   {
-    return new IncrementalIndex.Builder()
+    IncrementalIndex.Builder builder = new IncrementalIndex.Builder()
         .setIndexSchema(
             new IncrementalIndexSchema.Builder()
                 .withMetrics(schemaInfo.getAggsArray())
                 .withRollup(rollup)
                 .build()
         )
         .setReportParseExceptions(false)
-        .setMaxRowCount(rowsPerSegment)
-        .buildOnheap();
+        .setMaxRowCount(rowsPerSegment);
+    switch (indexType) {
+      case "onheap":
+        return builder.buildOnheap();
+      case "oak":
+        return builder.buildOak();
+    }
+    return null;
   }
 
   @Benchmark
@@ -195,7 +215,7 @@ public void readWithFilters(Blackhole blackhole)
   private Sequence<Cursor> makeCursors(IncrementalIndexStorageAdapter sa, DimFilter filter)
   {
     return sa.makeCursors(
-        filter.toFilter(),
+        filter == null ? null : filter.toFilter(),
         schemaInfo.getDataInterval(),
         VirtualColumns.EMPTY,
         Granularities.ALL,
@@ -208,4 +228,18 @@ private static DimensionSelector makeDimensionSelector(Cursor cursor, String nam
   {
     return cursor.getColumnSelectorFactory().makeDimensionSelector(new DefaultDimensionSpec(name, null));
   }
+
+  public static void main(String[] args) throws RunnerException
+  {
+    Options opt = new OptionsBuilder()
+            .include(IncrementalIndexReadBenchmark.class.getSimpleName())
+            .forks(0)
+            .threads(1)
+            .param("indexType", "oak")
+            .build();
+
+    new Runner(opt).run();
+  }
+
+
 }

benchmarks/src/main/java/org/apache/druid/benchmark/indexing/IndexIngestionBenchmark.java

@@ -40,6 +40,7 @@
 import org.openjdk.jmh.annotations.Scope;
 import org.openjdk.jmh.annotations.Setup;
 import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.TearDown;
 import org.openjdk.jmh.annotations.Warmup;
 import org.openjdk.jmh.infra.Blackhole;
 
@@ -61,6 +62,9 @@ public class IndexIngestionBenchmark
   @Param({"true", "false"})
   private boolean rollup;
 
+  @Param({"onheap", "oak"})
+  private String indexType;
+
   private static final Logger log = new Logger(IndexIngestionBenchmark.class);
   private static final int RNG_SEED = 9999;
 
@@ -98,18 +102,30 @@ public void setup2()
     incIndex = makeIncIndex();
   }
 
+  @TearDown(Level.Invocation)
+  public void tearDown()
+  {
+    incIndex.close();
+  }
+
   private IncrementalIndex makeIncIndex()
   {
-    return new IncrementalIndex.Builder()
+    IncrementalIndex.Builder builder = new IncrementalIndex.Builder()
         .setIndexSchema(
             new IncrementalIndexSchema.Builder()
                 .withMetrics(schemaInfo.getAggsArray())
                 .withRollup(rollup)
                 .build()
         )
         .setReportParseExceptions(false)
-        .setMaxRowCount(rowsPerSegment * 2)
-        .buildOnheap();
+        .setMaxRowCount(rowsPerSegment * 2);
+    switch (indexType) {
+      case "onheap":
+        return builder.buildOnheap();
+      case "oak":
+        return builder.buildOak();
+    }
+    return null;
   }
 
   @Benchmark
@@ -123,4 +139,5 @@ public void addRows(Blackhole blackhole) throws Exception
       blackhole.consume(rv);
     }
   }
+
 }

benchmarks/src/main/java/org/apache/druid/benchmark/indexing/IndexIngestionMultithreadedBenchmark.java

@@ -0,0 +1,176 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF 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 org.apache.druid.benchmark.indexing;
+
+import org.apache.druid.benchmark.datagen.BenchmarkDataGenerator;
+import org.apache.druid.benchmark.datagen.BenchmarkSchemaInfo;
+import org.apache.druid.benchmark.datagen.BenchmarkSchemas;
+import org.apache.druid.data.input.InputRow;
+import org.apache.druid.hll.HyperLogLogHash;
+import org.apache.druid.java.util.common.logger.Logger;
+import org.apache.druid.query.aggregation.hyperloglog.HyperUniquesSerde;
+import org.apache.druid.segment.incremental.IncrementalIndex;
+import org.apache.druid.segment.incremental.IncrementalIndexSchema;
+import org.apache.druid.segment.serde.ComplexMetrics;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Level;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Param;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.TearDown;
+import org.openjdk.jmh.annotations.Warmup;
+import org.openjdk.jmh.infra.Blackhole;
+import org.openjdk.jmh.runner.Runner;
+import org.openjdk.jmh.runner.RunnerException;
+import org.openjdk.jmh.runner.options.Options;
+import org.openjdk.jmh.runner.options.OptionsBuilder;
+
+import java.util.ArrayList;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+
+@State(Scope.Benchmark)
+@Fork(value = 1)
+@Warmup(iterations = 10)
+@Measurement(iterations = 25)
+public class IndexIngestionMultithreadedBenchmark
+{
+  @Param({"75000"})
+  private int rowsPerSegment;
+
+  @Param({"basic"})
+  private String schema;
+
+  @Param({"true", "false"})
+  private boolean rollup;
+
+  @Param({"onheap", "oak"})
+  private String indexType;
+
+  private static final Logger log = new Logger(IndexIngestionBenchmark.class);
+  private static final int RNG_SEED = 9999;
+
+  private IncrementalIndex incIndex;
+  private ArrayList<InputRow> rows;
+  private BenchmarkSchemaInfo schemaInfo;
+  AtomicInteger threadIdAllocator = new AtomicInteger(0);
+
+  @Setup
+  public void setup()
+  {
+    ComplexMetrics.registerSerde("hyperUnique", new HyperUniquesSerde(HyperLogLogHash.getDefault()));
+
+    rows = new ArrayList<InputRow>();
+    schemaInfo = BenchmarkSchemas.SCHEMA_MAP.get(schema);
+
+    BenchmarkDataGenerator gen = new BenchmarkDataGenerator(
+            schemaInfo.getColumnSchemas(),
+            RNG_SEED,
+            schemaInfo.getDataInterval(),
+            rowsPerSegment
+    );
+
+    for (int i = 0; i < rowsPerSegment; i++) {
+      InputRow row = gen.nextRow();
+      if (i % 10000 == 0) {
+        log.info(i + " rows generated.");
+      }
+
+      rows.add(row);
+    }
+  }
+
+  @State(Scope.Thread)
+  public static class ThreadState
+  {
+    int id = 0;
+
+    @Setup
+    public void setup(IndexIngestionMultithreadedBenchmark globalState)
+    {
+      id = globalState.threadIdAllocator.getAndIncrement();
+    }
+  }
+
+  @Setup(Level.Invocation)
+  public void setup2()
+  {
+    incIndex = makeIncIndex();
+  }
+
+  @TearDown(Level.Iteration)
+  public void tearDown()
+  {
+    incIndex.close();
+  }
+
+  private IncrementalIndex makeIncIndex()
+  {
+    IncrementalIndex.Builder builder = new IncrementalIndex.Builder()
+            .setIndexSchema(
+                    new IncrementalIndexSchema.Builder()
+                           .withMetrics(schemaInfo.getAggsArray())
+                           .withRollup(rollup)
+                           .build()
+            )
+            .setReportParseExceptions(false)
+            .setMaxRowCount(rowsPerSegment * 2);
+    switch (indexType) {
+      case "onheap":
+        return builder.buildOnheap();
+      case "oak":
+        return builder.buildOak();
+    }
+    return null;
+  }
+
+  @Benchmark
+  @BenchmarkMode(Mode.SingleShotTime)
+  @OutputTimeUnit(TimeUnit.MICROSECONDS)
+  public void addRows(Blackhole blackhole, ThreadState threadState) throws Exception
+  {
+    int threads = threadIdAllocator.get();
+    for (int i = threadState.id; i < rowsPerSegment; i += threads) {
+      InputRow row = rows.get(i);
+      int rv = incIndex.add(row).getRowCount();
+      blackhole.consume(rv);
+    }
+  }
+
+  public static void main(String[] args) throws RunnerException
+  {
+    Options opt = new OptionsBuilder()
+            .include(IndexIngestionMultithreadedBenchmark.class.getSimpleName())
+            .forks(1)
+            .threads(4)
+            .param("indexType", "oak")
+            .param("rollup", "false")
+            .build();
+
+    new Runner(opt).run();
+  }
+
+}

benchmarks/src/main/java/org/apache/druid/benchmark/query/SearchBenchmark.java

@@ -121,6 +121,9 @@ public class SearchBenchmark
   @Param({"1000"})
   private int limit;
 
+  @Param({"oak", "onheap"})
+  private String indexType;
+
   private static final Logger log = new Logger(SearchBenchmark.class);
   private static final IndexMergerV9 INDEX_MERGER_V9;
   private static final IndexIO INDEX_IO;
@@ -383,15 +386,22 @@ public void setup() throws IOException
   public void tearDown() throws IOException
   {
     FileUtils.deleteDirectory(tmpDir);
+    incIndexes.forEach(index -> index.close());
   }
 
   private IncrementalIndex makeIncIndex()
   {
-    return new IncrementalIndex.Builder()
-        .setSimpleTestingIndexSchema(schemaInfo.getAggsArray())
-        .setReportParseExceptions(false)
-        .setMaxRowCount(rowsPerSegment)
-        .buildOnheap();
+    IncrementalIndex.Builder builder = new IncrementalIndex.Builder()
+            .setSimpleTestingIndexSchema(schemaInfo.getAggsArray())
+            .setReportParseExceptions(false)
+            .setMaxRowCount(rowsPerSegment);
+    switch (indexType) {
+      case "onheap":
+        return builder.buildOnheap();
+      case "oak":
+        return builder.buildOak();
+    }
+    return null;
   }
 
   private static <T> List<T> runQuery(QueryRunnerFactory factory, QueryRunner runner, Query<T> query)

benchmarks/src/main/java/org/apache/druid/benchmark/query/TopNBenchmark.java

@@ -112,6 +112,9 @@ public class TopNBenchmark
   @Param({"10"})
   private int threshold;
 
+  @Param({"oak", "onheap"})
+  private String indexType;
+
   private static final Logger log = new Logger(TopNBenchmark.class);
   private static final int RNG_SEED = 9999;
   private static final IndexMergerV9 INDEX_MERGER_V9;
@@ -290,15 +293,22 @@ public void setup() throws IOException
   public void tearDown() throws IOException
   {
     FileUtils.deleteDirectory(tmpDir);
+    incIndexes.forEach(index -> index.close());
   }
 
   private IncrementalIndex makeIncIndex()
   {
-    return new IncrementalIndex.Builder()
+    IncrementalIndex.Builder builder = new IncrementalIndex.Builder()
         .setSimpleTestingIndexSchema(schemaInfo.getAggsArray())
         .setReportParseExceptions(false)
-        .setMaxRowCount(rowsPerSegment)
-        .buildOnheap();
+        .setMaxRowCount(rowsPerSegment);
+    switch (indexType) {
+      case "onheap":
+        return builder.buildOnheap();
+      case "oak":
+        return builder.buildOak();
+    }
+    return null;
   }
 
   private static <T> List<T> runQuery(QueryRunnerFactory factory, QueryRunner runner, Query<T> query)