Merge pull request #1700 from himanshug/update_agg_test_helper

update indexing in the helper to use multiple persists and merge
This commit is contained in:
Fangjin Yang 2015-09-14 06:56:29 -07:00
commit 34ef81572d
3 changed files with 69 additions and 30 deletions

View File

@ -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

View File

@ -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();
try {
createIndex(inputDataFile, parserJson, aggregators, segmentDir, minTimestamp, gran, maxRowCount);
return runQueryOnSegments(Lists.newArrayList(segmentDir), groupByQueryJson);
} finally {
FileUtils.deleteDirectory(segmentDir);
}
File segmentDir = tempFolder.newFolder();
createIndex(inputDataFile, parserJson, aggregators, segmentDir, minTimestamp, gran, maxRowCount);
return runQueryOnSegments(Lists.newArrayList(segmentDir), groupByQueryJson);
}
public Sequence<Row> createIndexAndRunQueryOnSegment(
@ -159,13 +158,9 @@ public class AggregationTestHelper
String groupByQueryJson
) throws Exception
{
File segmentDir = Files.createTempDir();
try {
createIndex(inputDataStream, parserJson, aggregators, segmentDir, minTimestamp, gran, maxRowCount);
return runQueryOnSegments(Lists.newArrayList(segmentDir), groupByQueryJson);
} finally {
FileUtils.deleteDirectory(segmentDir);
}
File segmentDir = tempFolder.newFolder();
createIndex(inputDataStream, parserJson, aggregators, segmentDir, minTimestamp, gran, maxRowCount);
return runQueryOnSegments(Lists.newArrayList(segmentDir), groupByQueryJson);
}
public void createIndex(
@ -237,19 +232,53 @@ public class AggregationTestHelper
int maxRowCount
) throws Exception
{
try(IncrementalIndex index = new OnheapIncrementalIndex(minTimestamp, gran, metrics, deserializeComplexMetrics, maxRowCount)) {
while (rows.hasNext()) {
IncrementalIndex index = null;
List<File> toMerge = new ArrayList<>();
try {
index = new OnheapIncrementalIndex(minTimestamp, gran, metrics, deserializeComplexMetrics, maxRowCount);
while (rows.hasNext()) {
Object row = rows.next();
if (row instanceof String && parser instanceof StringInputRowParser) {
//Note: this is required because StringInputRowParser is InputRowParser<ByteBuffer> as opposed to
//InputRowsParser<String>
index.add(((StringInputRowParser) parser).parse((String) row));
} else {
index.add(parser.parse(row));
try {
if (row instanceof String && parser instanceof StringInputRowParser) {
//Note: this is required because StringInputRowParser is InputRowParser<ByteBuffer> as opposed to
//InputRowsParser<String>
index.add(((StringInputRowParser) parser).parse((String) row));
} else {
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);
}
}
IndexMerger.persist(index, outDir, null, new IndexSpec());
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();
}
}
}

View File

@ -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\","