diff --git a/processing/src/test/java/io/druid/segment/IndexBuilder.java b/processing/src/test/java/io/druid/segment/IndexBuilder.java index 702e690b2be..2fdfb02f53e 100644 --- a/processing/src/test/java/io/druid/segment/IndexBuilder.java +++ b/processing/src/test/java/io/druid/segment/IndexBuilder.java @@ -45,7 +45,7 @@ import java.util.UUID; public class IndexBuilder { private static final int ROWS_PER_INDEX_FOR_MERGING = 1; - private static final int MAX_ROWS = 50_000; + private static final int DEFAULT_MAX_ROWS = Integer.MAX_VALUE; private IncrementalIndexSchema schema = new IncrementalIndexSchema.Builder().withMetrics(new AggregatorFactory[]{ new CountAggregatorFactory("count") @@ -53,6 +53,7 @@ public class IndexBuilder private IndexMerger indexMerger = TestHelper.getTestIndexMerger(); private File tmpDir; private IndexSpec indexSpec = new IndexSpec(); + private int maxRows = DEFAULT_MAX_ROWS; private final List rows = Lists.newArrayList(); @@ -90,20 +91,33 @@ public class IndexBuilder return this; } + public IndexBuilder maxRows(int maxRows) + { + this.maxRows = maxRows; + return this; + } + public IndexBuilder add(InputRow... rows) { return add(Arrays.asList(rows)); } - public IndexBuilder add(List rows) + public IndexBuilder add(Iterable rows) { - this.rows.addAll(rows); + Iterables.addAll(this.rows, rows); + return this; + } + + public IndexBuilder rows(Iterable rows) + { + this.rows.clear(); + Iterables.addAll(this.rows, rows); return this; } public IncrementalIndex buildIncrementalIndex() { - return buildIncrementalIndexWithRows(schema, rows); + return buildIncrementalIndexWithRows(schema, maxRows, rows); } public QueryableIndex buildMMappedIndex() @@ -138,6 +152,7 @@ public class IndexBuilder indexMerger.persist( buildIncrementalIndexWithRows( schema, + maxRows, rows.subList(i, Math.min(rows.size(), i + ROWS_PER_INDEX_FOR_MERGING)) ), new File(tmpDir, String.format("testIndex-%s", UUID.randomUUID().toString())), @@ -190,6 +205,7 @@ public class IndexBuilder private static IncrementalIndex buildIncrementalIndexWithRows( IncrementalIndexSchema schema, + int maxRows, Iterable rows ) { @@ -197,7 +213,7 @@ public class IndexBuilder final IncrementalIndex incrementalIndex = new OnheapIncrementalIndex( schema, true, - MAX_ROWS + maxRows ); for (InputRow row : rows) { try {