Merge pull request #1700 from himanshug/update_agg_test_helper

update indexing in the helper to use multiple persists and merge

extensions/histogram/src/test/java/io/druid/query/aggregation/histogram/ApproximateHistogramAggregationTest.java

@@ -26,7 +26,9 @@ import io.druid.data.input.MapBasedRow;
 import io.druid.granularity.QueryGranularity;
 import io.druid.query.aggregation.AggregationTestHelper;
 import org.junit.Assert;
+import org.junit.Rule;
 import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
 
 import java.io.File;
 
@@ -36,11 +38,14 @@ public class ApproximateHistogramAggregationTest
 {
   private AggregationTestHelper helper;
 
+  @Rule
+  public final TemporaryFolder tempFolder = new TemporaryFolder();
+
   public ApproximateHistogramAggregationTest()
   {
     ApproximateHistogramDruidModule module = new ApproximateHistogramDruidModule();
     module.configure(null);
-    helper = new AggregationTestHelper(Lists.newArrayList(module.getJacksonModules()));
+    helper = new AggregationTestHelper(Lists.newArrayList(module.getJacksonModules()), tempFolder);
   }
 
   @Test

processing/src/test/java/io/druid/query/aggregation/AggregationTestHelper.java

@@ -28,9 +28,7 @@ import com.google.common.base.Suppliers;
 import com.google.common.base.Throwables;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
-import com.google.common.io.CharSource;
 import com.google.common.io.Closeables;
-import com.google.common.io.Files;
 import com.google.common.util.concurrent.ListenableFuture;
 import com.metamx.common.guava.CloseQuietly;
 import com.metamx.common.guava.Sequence;
@@ -58,20 +56,22 @@ import io.druid.query.groupby.GroupByQueryRunnerFactory;
 import io.druid.segment.IndexIO;
 import io.druid.segment.IndexMerger;
 import io.druid.segment.IndexSpec;
+import io.druid.segment.QueryableIndex;
 import io.druid.segment.QueryableIndexSegment;
 import io.druid.segment.Segment;
 import io.druid.segment.incremental.IncrementalIndex;
+import io.druid.segment.incremental.IndexSizeExceededException;
 import io.druid.segment.incremental.OnheapIncrementalIndex;
-import org.apache.commons.io.FileUtils;
 import org.apache.commons.io.IOUtils;
 import org.apache.commons.io.LineIterator;
+import org.junit.rules.TemporaryFolder;
 
 import java.io.File;
 import java.io.FileInputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.nio.ByteBuffer;
-import java.nio.charset.Charset;
+import java.util.ArrayList;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
@@ -88,8 +88,11 @@ public class AggregationTestHelper
   private final GroupByQueryQueryToolChest toolChest;
   private final GroupByQueryRunnerFactory factory;
 
-  public AggregationTestHelper(List<? extends Module> jsonModulesToRegister)
+  private final TemporaryFolder tempFolder;
+
+  public AggregationTestHelper(List<? extends Module> jsonModulesToRegister, TemporaryFolder tempFoler)
   {
+    this.tempFolder = tempFoler;
     mapper = new DefaultObjectMapper();
 
     for(Module mod : jsonModulesToRegister) {
@@ -140,13 +143,9 @@ public class AggregationTestHelper
       String groupByQueryJson
   ) throws Exception
   {
-    File segmentDir = Files.createTempDir();
+    File segmentDir = tempFolder.newFolder();
-    try {
     createIndex(inputDataFile, parserJson, aggregators, segmentDir, minTimestamp, gran, maxRowCount);
     return runQueryOnSegments(Lists.newArrayList(segmentDir), groupByQueryJson);
-    } finally {
-      FileUtils.deleteDirectory(segmentDir);
-    }
   }
 
   public Sequence<Row> createIndexAndRunQueryOnSegment(
@@ -159,13 +158,9 @@ public class AggregationTestHelper
       String groupByQueryJson
   ) throws Exception
   {
-    File segmentDir = Files.createTempDir();
+    File segmentDir = tempFolder.newFolder();
-    try {
     createIndex(inputDataStream, parserJson, aggregators, segmentDir, minTimestamp, gran, maxRowCount);
     return runQueryOnSegments(Lists.newArrayList(segmentDir), groupByQueryJson);
-    } finally {
-      FileUtils.deleteDirectory(segmentDir);
-    }
   }
 
   public void createIndex(
@@ -237,10 +232,14 @@ public class AggregationTestHelper
       int maxRowCount
   ) throws Exception
   {
-    try(IncrementalIndex index = new OnheapIncrementalIndex(minTimestamp, gran, metrics, deserializeComplexMetrics, maxRowCount)) {
+    IncrementalIndex index = null;
-      while (rows.hasNext()) {
+    List<File> toMerge = new ArrayList<>();
 
+    try {
+      index = new OnheapIncrementalIndex(minTimestamp, gran, metrics, deserializeComplexMetrics, maxRowCount);
+      while (rows.hasNext()) {
         Object row = rows.next();
+        try {
           if (row instanceof String && parser instanceof StringInputRowParser) {
             //Note: this is required because StringInputRowParser is InputRowParser<ByteBuffer> as opposed to
             //InputRowsParser<String>
@@ -249,9 +248,39 @@ public class AggregationTestHelper
             index.add(parser.parse(row));
           }
         }
+        catch (IndexSizeExceededException ex) {
+          File tmp = tempFolder.newFolder();
+          toMerge.add(tmp);
+          IndexMerger.persist(index, tmp, null, new IndexSpec());
+          index.close();
+          index = new OnheapIncrementalIndex(minTimestamp, gran, metrics, deserializeComplexMetrics, maxRowCount);
+        }
+      }
+
+      if (toMerge.size() > 0) {
+        File tmp = tempFolder.newFolder();
+        toMerge.add(tmp);
+        IndexMerger.persist(index, tmp, null, new IndexSpec());
+
+        List<QueryableIndex> indexes = new ArrayList<>(toMerge.size());
+        for (File file : toMerge) {
+          indexes.add(IndexIO.loadIndex(file));
+        }
+        IndexMerger.mergeQueryableIndex(indexes, metrics, outDir, new IndexSpec());
+
+        for (QueryableIndex qi : indexes) {
+          qi.close();
+        }
+      } else {
         IndexMerger.persist(index, outDir, null, new IndexSpec());
       }
     }
+    finally {
+      if (index != null) {
+        index.close();
+      }
+    }
+  }
 
   //Simulates running group-by query on individual segments as historicals would do, json serialize the results
   //from each segment, later deserialize and merge and finally return the results

processing/src/test/java/io/druid/query/aggregation/hyperloglog/HyperUniquesAggregationTest.java

@@ -27,16 +27,21 @@ import io.druid.granularity.QueryGranularity;
 import io.druid.jackson.AggregatorsModule;
 import io.druid.query.aggregation.AggregationTestHelper;
 import org.junit.Assert;
+import org.junit.Rule;
 import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
 
 import java.io.File;
 
 public class HyperUniquesAggregationTest
 {
+  @Rule
+  public final TemporaryFolder tempFolder = new TemporaryFolder();
+
   @Test
   public void testIngestAndQuery() throws Exception
   {
-    AggregationTestHelper helper = new AggregationTestHelper(Lists.newArrayList(new AggregatorsModule()));
+    AggregationTestHelper helper = new AggregationTestHelper(Lists.newArrayList(new AggregatorsModule()), tempFolder);
 
     String metricSpec = "[{"
                         + "\"type\": \"hyperUnique\","