Skip to content

Commit

Permalink
add size param
Browse files Browse the repository at this point in the history
  • Loading branch information
talevy committed Oct 29, 2020
1 parent 285ffd9 commit baaea0e
Show file tree
Hide file tree
Showing 4 changed files with 23 additions and 6 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@ public class GeoLineAggregationBuilder
static final ParseField SORT_FIELD = new ParseField("sort");
static final ParseField ORDER_FIELD = new ParseField("sort_order");
static final ParseField INCLUDE_SORT_FIELD = new ParseField("include_sort");
static final ParseField SIZE_FIELD = new ParseField("size");

public static final String NAME = "geo_line";

Expand All @@ -50,10 +51,12 @@ public class GeoLineAggregationBuilder
MultiValuesSourceParseHelper.declareField(SORT_FIELD.getPreferredName(), PARSER, true, false, false);
PARSER.declareString((builder, order) -> builder.sortOrder(SortOrder.fromString(order)), ORDER_FIELD);
PARSER.declareBoolean(GeoLineAggregationBuilder::includeSort, INCLUDE_SORT_FIELD);
PARSER.declareInt(GeoLineAggregationBuilder::size, SIZE_FIELD);
}

private boolean includeSort;
private SortOrder sortOrder = SortOrder.ASC;
private int size = GeoLineAggregator.MAX_PATH_SIZE;

public static void registerUsage(ValuesSourceRegistry.Builder builder) {
builder.registerUsage(NAME, CoreValuesSourceType.GEOPOINT);
Expand Down Expand Up @@ -87,6 +90,14 @@ public GeoLineAggregationBuilder sortOrder(SortOrder sortOrder) {
return this;
}

public GeoLineAggregationBuilder size(int size) {
if (size > GeoLineAggregator.MAX_PATH_SIZE) {
throw new IllegalArgumentException("invalid [size] value [" + size + "] must be <= " + GeoLineAggregator.MAX_PATH_SIZE);
}
this.size = size;
return this;
}

@Override
protected AggregationBuilder shallowCopy(AggregatorFactories.Builder factoriesBuilder, Map<String, Object> metaData) {
return new GeoLineAggregationBuilder(this, factoriesBuilder, metaData);
Expand Down Expand Up @@ -116,7 +127,7 @@ protected MultiValuesSourceAggregatorFactory innerBuild(AggregationContext aggre
AggregatorFactory parent,
AggregatorFactories.Builder subFactoriesBuilder) throws IOException {
return new GeoLineAggregatorFactory(name, configs, format, aggregationContext, parent, subFactoriesBuilder, metadata,
includeSort, sortOrder);
includeSort, sortOrder, size);
}

public GeoLineAggregationBuilder value(MultiValuesSourceFieldConfig valueConfig) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -32,9 +32,11 @@ final class GeoLineAggregator extends MetricsAggregator {
private final GeoLineBucketedSort.Extra extra;
private final boolean includeSorts;
private final SortOrder sortOrder;
private final int size;

GeoLineAggregator(String name, MultiValuesSource.AnyMultiValuesSource valuesSources, SearchContext context,
Aggregator parent, Map<String,Object> metaData, boolean includeSorts, SortOrder sortOrder) throws IOException {
Aggregator parent, Map<String,Object> metaData, boolean includeSorts, SortOrder sortOrder,
int size) throws IOException {
super(name, context, parent, metaData);
this.valuesSources = valuesSources;
if (valuesSources != null) {
Expand All @@ -46,6 +48,7 @@ final class GeoLineAggregator extends MetricsAggregator {
}
this.includeSorts = includeSorts;
this.sortOrder = sortOrder;
this.size = size;
}

@Override
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -24,22 +24,24 @@ final class GeoLineAggregatorFactory extends MultiValuesSourceAggregatorFactory

private boolean includeSort;
private SortOrder sortOrder;
private int size;

GeoLineAggregatorFactory(String name,
Map<String, ValuesSourceConfig> configs,
DocValueFormat format, AggregationContext aggregationContext, AggregatorFactory parent,
AggregatorFactories.Builder subFactoriesBuilder,
Map<String, Object> metaData, boolean includeSort, SortOrder sortOrder) throws IOException {
Map<String, Object> metaData, boolean includeSort, SortOrder sortOrder, int size) throws IOException {
super(name, configs, format, aggregationContext, parent, subFactoriesBuilder, metaData);
this.includeSort = includeSort;
this.sortOrder = sortOrder;
this.size = size;
}

@Override
protected Aggregator createUnmapped(SearchContext searchContext,
Aggregator parent,
Map<String, Object> metaData) throws IOException {
return new GeoLineAggregator(name, null, searchContext, parent, metaData, includeSort, sortOrder);
return new GeoLineAggregator(name, null, searchContext, parent, metaData, includeSort, sortOrder, size);
}

@Override
Expand All @@ -51,7 +53,7 @@ protected Aggregator doCreateInternal(SearchContext searchContext,
Map<String, Object> metaData) throws IOException {
MultiValuesSource.AnyMultiValuesSource valuesSources =
new MultiValuesSource.AnyMultiValuesSource(configs, searchContext.getQueryShardContext());
return new GeoLineAggregator(name, valuesSources, searchContext, parent, metaData, includeSort, sortOrder);
return new GeoLineAggregator(name, valuesSources, searchContext, parent, metaData, includeSort, sortOrder, size);
}

@Override
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -88,8 +88,9 @@ public InternalAggregation reduce(List<InternalAggregation> aggregations, Reduce
idx += 1;
}
}
// the final reduce should always be in ascending order

new PathArraySorter(finalList, finalSortVals, sortOrder).sort();
new PathArraySorter(finalList, finalSortVals, SortOrder.ASC).sort();
long[] finalCappedList = Arrays.copyOf(finalList, Math.min(GeoLineAggregator.MAX_PATH_SIZE, mergedSize));
double[] finalCappedSortVals = Arrays.copyOf(finalSortVals, Math.min(GeoLineAggregator.MAX_PATH_SIZE, mergedSize));
return new InternalGeoLine(name, finalCappedList, finalCappedSortVals, getMetadata(), complete, includeSorts, sortOrder);
Expand Down

0 comments on commit baaea0e

Please sign in to comment.