diff --git a/processing/src/main/java/io/druid/query/search/SearchBinaryFn.java b/processing/src/main/java/io/druid/query/search/SearchBinaryFn.java index 99093911739..1c5a0b5668e 100644 --- a/processing/src/main/java/io/druid/query/search/SearchBinaryFn.java +++ b/processing/src/main/java/io/druid/query/search/SearchBinaryFn.java @@ -19,6 +19,7 @@ package io.druid.query.search; +import com.google.common.collect.Iterables; import com.google.common.collect.Lists; import com.google.common.collect.Sets; import com.metamx.common.guava.nary.BinaryFn; @@ -32,18 +33,22 @@ import java.util.TreeSet; /** */ -public class SearchBinaryFn implements BinaryFn, Result, Result> +public class SearchBinaryFn + implements BinaryFn, Result, Result> { private final SearchSortSpec searchSortSpec; private final QueryGranularity gran; + private final int limit; public SearchBinaryFn( SearchSortSpec searchSortSpec, - QueryGranularity granularity + QueryGranularity granularity, + int limit ) { this.searchSortSpec = searchSortSpec; this.gran = granularity; + this.limit = limit; } @Override @@ -65,7 +70,13 @@ public class SearchBinaryFn implements BinaryFn, Resul results.addAll(Lists.newArrayList(arg2Vals)); return (gran instanceof AllGranularity) - ? new Result(arg1.getTimestamp(), new SearchResultValue(Lists.newArrayList(results))) + ? new Result( + arg1.getTimestamp(), new SearchResultValue( + Lists.newArrayList( + Iterables.limit(results, limit) + ) + ) + ) : new Result( gran.toDateTime(gran.truncate(arg1.getTimestamp().getMillis())), new SearchResultValue(Lists.newArrayList(results)) diff --git a/processing/src/main/java/io/druid/query/search/SearchQueryQueryToolChest.java b/processing/src/main/java/io/druid/query/search/SearchQueryQueryToolChest.java index 07f51a0f743..166ed71da8a 100644 --- a/processing/src/main/java/io/druid/query/search/SearchQueryQueryToolChest.java +++ b/processing/src/main/java/io/druid/query/search/SearchQueryQueryToolChest.java @@ -99,7 +99,7 @@ public class SearchQueryQueryToolChest extends QueryToolChest actual = new SearchBinaryFn(new LexicographicSearchSortSpec(), QueryGranularity.ALL).apply(r1, r2); + Result actual = new SearchBinaryFn(new LexicographicSearchSortSpec(), QueryGranularity.ALL, Integer.MAX_VALUE).apply(r1, r2); Assert.assertEquals(expected.getTimestamp(), actual.getTimestamp()); assertSearchMergeResult(expected.getValue(), actual.getValue()); } @@ -138,7 +138,7 @@ public class SearchBinaryFnTest ) ); - Result actual = new SearchBinaryFn(new LexicographicSearchSortSpec(), QueryGranularity.DAY).apply(r1, r2); + Result actual = new SearchBinaryFn(new LexicographicSearchSortSpec(), QueryGranularity.DAY, Integer.MAX_VALUE).apply(r1, r2); Assert.assertEquals(expected.getTimestamp(), actual.getTimestamp()); assertSearchMergeResult(expected.getValue(), actual.getValue()); } @@ -162,7 +162,7 @@ public class SearchBinaryFnTest Result expected = r1; - Result actual = new SearchBinaryFn(new LexicographicSearchSortSpec(), QueryGranularity.ALL).apply(r1, r2); + Result actual = new SearchBinaryFn(new LexicographicSearchSortSpec(), QueryGranularity.ALL, Integer.MAX_VALUE).apply(r1, r2); Assert.assertEquals(expected.getTimestamp(), actual.getTimestamp()); assertSearchMergeResult(expected.getValue(), actual.getValue()); } @@ -210,7 +210,7 @@ public class SearchBinaryFnTest ) ); - Result actual = new SearchBinaryFn(new LexicographicSearchSortSpec(), QueryGranularity.ALL).apply(r1, r2); + Result actual = new SearchBinaryFn(new LexicographicSearchSortSpec(), QueryGranularity.ALL, Integer.MAX_VALUE).apply(r1, r2); Assert.assertEquals(expected.getTimestamp(), actual.getTimestamp()); assertSearchMergeResult(expected.getValue(), actual.getValue()); } @@ -258,7 +258,7 @@ public class SearchBinaryFnTest ) ); - Result actual = new SearchBinaryFn(new StrlenSearchSortSpec(), QueryGranularity.ALL).apply(r1, r2); + Result actual = new SearchBinaryFn(new StrlenSearchSortSpec(), QueryGranularity.ALL, Integer.MAX_VALUE).apply(r1, r2); Assert.assertEquals(expected.getTimestamp(), actual.getTimestamp()); assertSearchMergeResult(expected.getValue(), actual.getValue()); } @@ -282,7 +282,38 @@ public class SearchBinaryFnTest Result expected = r1; - Result actual = new SearchBinaryFn(new LexicographicSearchSortSpec(), QueryGranularity.ALL).apply(r1, r2); + Result actual = new SearchBinaryFn(new LexicographicSearchSortSpec(), QueryGranularity.ALL, Integer.MAX_VALUE).apply(r1, r2); + Assert.assertEquals(expected.getTimestamp(), actual.getTimestamp()); + assertSearchMergeResult(expected.getValue(), actual.getValue()); + } + + @Test + public void testMergeLimit(){ + Result r1 = new Result( + currTime, + new SearchResultValue( + ImmutableList.of( + new SearchHit( + "blah", + "foo" + ) + ) + ) + ); + + Result r2 = new Result( + currTime, + new SearchResultValue( + ImmutableList.of( + new SearchHit( + "blah2", + "foo2" + ) + ) + ) + ); + Result expected = r1; + Result actual = new SearchBinaryFn(new LexicographicSearchSortSpec(), QueryGranularity.ALL, 1).apply(r1, r2); Assert.assertEquals(expected.getTimestamp(), actual.getTimestamp()); assertSearchMergeResult(expected.getValue(), actual.getValue()); }