Skip to content

Commit bbddbcd

Browse files
talevytalevy
committed
insilico
1 parent 9e87913 commit bbddbcd

File tree

7 files changed

+650
-0
lines changed

7 files changed

+650
-0
lines changed

‎server/src/main/java/org/elasticsearch/search/aggregations/support/MultiValuesSource.java‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
package org.elasticsearch.search.aggregations.support;
2020

2121
import org.apache.lucene.index.LeafReaderContext;
22+
import org.elasticsearch.index.fielddata.MultiGeoPointValues;
2223
import org.elasticsearch.index.fielddata.SortedNumericDoubleValues;
2324
import org.elasticsearch.index.query.QueryShardContext;
2425

@@ -51,6 +52,41 @@ public SortedNumericDoubleValues getField(String fieldName, LeafReaderContext ct
5152
}
5253
}
5354

55+
public static class AnyMultiValuesSource extends MultiValuesSource<ValuesSource> {
56+
public AnyMultiValuesSource(Map<String, ValuesSourceConfig<ValuesSource>> valuesSourceConfigs,
57+
QueryShardContext context) {
58+
values = new HashMap<>(valuesSourceConfigs.size());
59+
for (Map.Entry<String, ValuesSourceConfig<ValuesSource>> entry : valuesSourceConfigs.entrySet()) {
60+
values.put(entry.getKey(), entry.getValue().toValuesSource(context));
61+
}
62+
}
63+
64+
private ValuesSource getField(String fieldName) {
65+
ValuesSource valuesSource = values.get(fieldName);
66+
if (valuesSource == null) {
67+
throw new IllegalArgumentException("Could not find field name [" + fieldName + "] in multiValuesSource");
68+
}
69+
return valuesSource;
70+
}
71+
72+
public SortedNumericDoubleValues getNumericField(String fieldName, LeafReaderContext ctx) throws IOException {
73+
ValuesSource valuesSource = getField(fieldName);
74+
if (valuesSource instanceof ValuesSource.Numeric) {
75+
return ((ValuesSource.Numeric) valuesSource).doubleValues(ctx);
76+
}
77+
throw new IllegalArgumentException("field [" + fieldName + "] is not a numeric type");
78+
}
79+
80+
public MultiGeoPointValues getGeoPointField(String fieldName, LeafReaderContext ctx) {
81+
ValuesSource valuesSource = getField(fieldName);
82+
if (valuesSource instanceof ValuesSource.GeoPoint) {
83+
return ((ValuesSource.GeoPoint) valuesSource).geoPointValues(ctx);
84+
}
85+
throw new IllegalArgumentException("field [" + fieldName + "] is not a geo_point type");
86+
}
87+
88+
}
89+
5490
public boolean needsScores() {
5591
return values.values().stream().anyMatch(ValuesSource::needsScores);
5692
}

‎x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/search/aggregations/GeoLineAggregationBuilder.java‎

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,106 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License;
4+
* you may not use this file except in compliance with the Elastic License.
5+
*/
6+
package org.elasticsearch.xpack.spatial.search.aggregations;
7+
8+
import org.elasticsearch.common.ParseField;
9+
import org.elasticsearch.common.io.stream.StreamInput;
10+
import org.elasticsearch.common.io.stream.StreamOutput;
11+
import org.elasticsearch.common.xcontent.ObjectParser;
12+
import org.elasticsearch.common.xcontent.ToXContent;
13+
import org.elasticsearch.common.xcontent.XContentBuilder;
14+
import org.elasticsearch.common.xcontent.XContentParser;
15+
import org.elasticsearch.index.query.QueryShardContext;
16+
import org.elasticsearch.search.DocValueFormat;
17+
import org.elasticsearch.search.aggregations.AggregationBuilder;
18+
import org.elasticsearch.search.aggregations.AggregatorFactories;
19+
import org.elasticsearch.search.aggregations.AggregatorFactory;
20+
import org.elasticsearch.search.aggregations.support.MultiValuesSourceAggregationBuilder;
21+
import org.elasticsearch.search.aggregations.support.MultiValuesSourceAggregatorFactory;
22+
import org.elasticsearch.search.aggregations.support.MultiValuesSourceFieldConfig;
23+
import org.elasticsearch.search.aggregations.support.MultiValuesSourceParseHelper;
24+
import org.elasticsearch.search.aggregations.support.ValueType;
25+
import org.elasticsearch.search.aggregations.support.ValuesSource;
26+
import org.elasticsearch.search.aggregations.support.ValuesSourceConfig;
27+
28+
import java.io.IOException;
29+
import java.util.Map;
30+
import java.util.Objects;
31+
32+
public class GeoLineAggregationBuilder
33+
extends MultiValuesSourceAggregationBuilder<ValuesSource, GeoLineAggregationBuilder> {
34+
35+
static final ParseField GEO_POINT_FIELD = new ParseField("geo_point");
36+
static final ParseField SORT_FIELD = new ParseField("sort");
37+
38+
static final String NAME = "geo_line";
39+
40+
private static final ObjectParser<GeoLineAggregationBuilder, Void> PARSER;
41+
static {
42+
PARSER = new ObjectParser<>(NAME);
43+
MultiValuesSourceParseHelper.declareCommon(PARSER, true, ValueType.NUMERIC);
44+
MultiValuesSourceParseHelper.declareField(GEO_POINT_FIELD.getPreferredName(), PARSER, true, false);
45+
MultiValuesSourceParseHelper.declareField(SORT_FIELD.getPreferredName(), PARSER, true, false);
46+
}
47+
48+
GeoLineAggregationBuilder(String name) {
49+
super(name, null);
50+
}
51+
52+
private GeoLineAggregationBuilder(GeoLineAggregationBuilder clone,
53+
AggregatorFactories.Builder factoriesBuilder, Map<String, Object> metaData) {
54+
super(clone, factoriesBuilder, metaData);
55+
}
56+
57+
/**
58+
* Read from a stream.
59+
*/
60+
GeoLineAggregationBuilder(StreamInput in) throws IOException {
61+
super(in, null);
62+
}
63+
64+
static AggregationBuilder parse(String aggregationName, XContentParser parser) throws IOException {
65+
return PARSER.parse(parser, new GeoLineAggregationBuilder(aggregationName), null);
66+
}
67+
68+
@Override
69+
protected AggregationBuilder shallowCopy(AggregatorFactories.Builder factoriesBuilder, Map<String, Object> metaData) {
70+
return new GeoLineAggregationBuilder(this, factoriesBuilder, metaData);
71+
}
72+
73+
@Override
74+
protected void innerWriteTo(StreamOutput out) {
75+
// Do nothing, no extra state to write to stream
76+
}
77+
78+
@Override
79+
protected MultiValuesSourceAggregatorFactory<ValuesSource> innerBuild(QueryShardContext queryShardContext, Map<String,
80+
ValuesSourceConfig<ValuesSource>> configs, DocValueFormat format, AggregatorFactory parent,
81+
AggregatorFactories.Builder subFactoriesBuilder) throws IOException {
82+
return new GeoLineAggregatorFactory(name, configs, format, queryShardContext, parent, subFactoriesBuilder, metaData);
83+
}
84+
85+
public GeoLineAggregationBuilder value(MultiValuesSourceFieldConfig valueConfig) {
86+
valueConfig = Objects.requireNonNull(valueConfig, "Configuration for field [" + GEO_POINT_FIELD + "] cannot be null");
87+
field(GEO_POINT_FIELD.getPreferredName(), valueConfig);
88+
return this;
89+
}
90+
91+
public GeoLineAggregationBuilder sort(MultiValuesSourceFieldConfig sortConfig) {
92+
sortConfig = Objects.requireNonNull(sortConfig, "Configuration for field [" + SORT_FIELD + "] cannot be null");
93+
field(SORT_FIELD.getPreferredName(), sortConfig);
94+
return this;
95+
}
96+
97+
@Override
98+
public XContentBuilder doXContentBody(XContentBuilder builder, ToXContent.Params params) {
99+
return builder;
100+
}
101+
102+
@Override
103+
public String getType() {
104+
return NAME;
105+
}
106+
}

‎x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/search/aggregations/GeoLineAggregator.java‎

Lines changed: 151 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,151 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License;
4+
* you may not use this file except in compliance with the Elastic License.
5+
*/
6+
package org.elasticsearch.xpack.spatial.search.aggregations;
7+
8+
import org.apache.lucene.geo.GeoEncodingUtils;
9+
import org.apache.lucene.index.LeafReaderContext;
10+
import org.apache.lucene.search.ScoreMode;
11+
import org.apache.lucene.util.ArrayUtil;
12+
import org.apache.lucene.util.NumericUtils;
13+
import org.elasticsearch.common.geo.GeoPoint;
14+
import org.elasticsearch.common.lease.Releasables;
15+
import org.elasticsearch.common.util.BigArrays;
16+
import org.elasticsearch.common.util.IntArray;
17+
import org.elasticsearch.common.util.ObjectArray;
18+
import org.elasticsearch.index.fielddata.MultiGeoPointValues;
19+
import org.elasticsearch.index.fielddata.SortedNumericDoubleValues;
20+
import org.elasticsearch.search.aggregations.AggregationExecutionException;
21+
import org.elasticsearch.search.aggregations.Aggregator;
22+
import org.elasticsearch.search.aggregations.InternalAggregation;
23+
import org.elasticsearch.search.aggregations.LeafBucketCollector;
24+
import org.elasticsearch.search.aggregations.LeafBucketCollectorBase;
25+
import org.elasticsearch.search.aggregations.metrics.MetricsAggregator;
26+
import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator;
27+
import org.elasticsearch.search.aggregations.support.MultiValuesSource;
28+
import org.elasticsearch.search.internal.SearchContext;
29+
30+
import java.io.IOException;
31+
import java.util.List;
32+
import java.util.Map;
33+
34+
import static org.elasticsearch.xpack.spatial.search.aggregations.GeoLineAggregationBuilder.GEO_POINT_FIELD;
35+
import static org.elasticsearch.xpack.spatial.search.aggregations.GeoLineAggregationBuilder.SORT_FIELD;
36+
37+
/**
38+
* Metric Aggregation for computing the pearson product correlation coefficient between multiple fields
39+
**/
40+
final class GeoLineAggregator extends MetricsAggregator {
41+
/** Multiple ValuesSource with field names */
42+
private final MultiValuesSource.AnyMultiValuesSource valuesSources;
43+
44+
private ObjectArray<long[]> paths;
45+
private ObjectArray<double[]> sortValues;
46+
private IntArray idxs;
47+
48+
GeoLineAggregator(String name, MultiValuesSource.AnyMultiValuesSource valuesSources, SearchContext context,
49+
Aggregator parent, List<PipelineAggregator> pipelineAggregators,
50+
Map<String,Object> metaData) throws IOException {
51+
super(name, context, parent, pipelineAggregators, metaData);
52+
this.valuesSources = valuesSources;
53+
if (valuesSources != null) {
54+
paths = context.bigArrays().newObjectArray(1);
55+
sortValues = context.bigArrays().newObjectArray(1);
56+
idxs = context.bigArrays().newIntArray(1);
57+
}
58+
}
59+
60+
@Override
61+
public ScoreMode scoreMode() {
62+
if (valuesSources != null && valuesSources.needsScores()) {
63+
return ScoreMode.COMPLETE;
64+
}
65+
return super.scoreMode();
66+
}
67+
68+
@Override
69+
public LeafBucketCollector getLeafCollector(LeafReaderContext ctx,
70+
final LeafBucketCollector sub) throws IOException {
71+
if (valuesSources == null) {
72+
return LeafBucketCollector.NO_OP_COLLECTOR;
73+
}
74+
final BigArrays bigArrays = context.bigArrays();
75+
MultiGeoPointValues docGeoPointValues = valuesSources.getGeoPointField(GEO_POINT_FIELD.getPreferredName(), ctx);
76+
SortedNumericDoubleValues docSortValues = valuesSources.getNumericField(SORT_FIELD.getPreferredName(), ctx);
77+
78+
return new LeafBucketCollectorBase(sub, docGeoPointValues) {
79+
@Override
80+
public void collect(int doc, long bucket) throws IOException {
81+
paths = bigArrays.grow(paths, bucket + 1);
82+
if (docGeoPointValues.advanceExact(doc) && docSortValues.advanceExact(doc)) {
83+
if (docSortValues.docValueCount() > 1) {
84+
throw new AggregationExecutionException("Encountered more than one sort value for a " +
85+
"single document. Use a script to combine multiple sort-values-per-doc into a single value.");
86+
}
87+
if (docGeoPointValues.docValueCount() > 1) {
88+
throw new AggregationExecutionException("Encountered more than one geo_point value for a " +
89+
"single document. Use a script to combine multiple geo_point-values-per-doc into a single value.");
90+
}
91+
92+
// There should always be one weight if advanceExact lands us here, either
93+
// a real weight or a `missing` weight
94+
assert docSortValues.docValueCount() == 1;
95+
assert docGeoPointValues.docValueCount() == 1;
96+
final double sort = docSortValues.nextValue();
97+
final GeoPoint point = docGeoPointValues.nextValue();
98+
99+
int idx = idxs.get(bucket);
100+
long[] bucketLine = paths.get(bucket);
101+
double[] sortVals = sortValues.get(bucket);
102+
if (bucketLine == null) {
103+
bucketLine = new long[10];
104+
} else {
105+
bucketLine = ArrayUtil.grow(bucketLine, idx + 1);
106+
}
107+
108+
109+
if (sortVals == null) {
110+
sortVals = new double[10];
111+
} else {
112+
sortVals = ArrayUtil.grow(sortVals, idx + 1);
113+
}
114+
115+
int encodedLat = GeoEncodingUtils.encodeLatitude(point.lat());
116+
int encodedLon = GeoEncodingUtils.encodeLongitude(point.lon());
117+
long lonLat = (((long) encodedLon) << 32) | (encodedLat & 0xffffffffL);
118+
119+
sortVals[idx] = sort;
120+
bucketLine[idx] = lonLat;
121+
122+
paths.set(bucket, bucketLine);
123+
sortValues.set(bucket, sortVals);
124+
idxs.set(bucket, idx + 1);
125+
}
126+
}
127+
};
128+
}
129+
130+
@Override
131+
public InternalAggregation buildAggregation(long bucket) {
132+
if (valuesSources == null) {
133+
return buildEmptyAggregation();
134+
}
135+
long[] bucketLine = paths.get(bucket);
136+
double[] sortVals = sortValues.get(bucket);
137+
int length = idxs.get(bucket);
138+
new PathArraySorter(bucketLine, sortVals, length).sort();
139+
return new InternalGeoLine(name, bucketLine, sortVals, length, pipelineAggregators(), metaData());
140+
}
141+
142+
@Override
143+
public InternalAggregation buildEmptyAggregation() {
144+
return new InternalGeoLine(name, null, null, 0, pipelineAggregators(), metaData());
145+
}
146+
147+
@Override
148+
public void doClose() {
149+
Releasables.close(paths, idxs, sortValues);
150+
}
151+
}

‎x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/search/aggregations/GeoLineAggregatorFactory.java‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License;
4+
* you may not use this file except in compliance with the Elastic License.
5+
*/
6+
package org.elasticsearch.xpack.spatial.search.aggregations;
7+
8+
import org.elasticsearch.index.query.QueryShardContext;
9+
import org.elasticsearch.search.DocValueFormat;
10+
import org.elasticsearch.search.aggregations.Aggregator;
11+
import org.elasticsearch.search.aggregations.AggregatorFactories;
12+
import org.elasticsearch.search.aggregations.AggregatorFactory;
13+
import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator;
14+
import org.elasticsearch.search.aggregations.support.MultiValuesSource;
15+
import org.elasticsearch.search.aggregations.support.MultiValuesSourceAggregatorFactory;
16+
import org.elasticsearch.search.aggregations.support.ValuesSource;
17+
import org.elasticsearch.search.aggregations.support.ValuesSourceConfig;
18+
import org.elasticsearch.search.internal.SearchContext;
19+
20+
import java.io.IOException;
21+
import java.util.List;
22+
import java.util.Map;
23+
24+
final class GeoLineAggregatorFactory extends MultiValuesSourceAggregatorFactory<ValuesSource> {
25+
26+
GeoLineAggregatorFactory(String name,
27+
Map<String, ValuesSourceConfig<ValuesSource>> configs,
28+
DocValueFormat format, QueryShardContext queryShardContext, AggregatorFactory parent,
29+
AggregatorFactories.Builder subFactoriesBuilder,
30+
Map<String, Object> metaData) throws IOException {
31+
super(name, configs, format, queryShardContext, parent, subFactoriesBuilder, metaData);
32+
}
33+
34+
@Override
35+
protected Aggregator createUnmapped(SearchContext searchContext, Aggregator parent, List<PipelineAggregator> pipelineAggregators,
36+
Map<String, Object> metaData) throws IOException {
37+
return new GeoLineAggregator(name, null, searchContext, parent, pipelineAggregators, metaData);
38+
}
39+
40+
@Override
41+
protected Aggregator doCreateInternal(SearchContext searchContext, Map<String, ValuesSourceConfig<ValuesSource>> configs,
42+
DocValueFormat format, Aggregator parent, boolean collectsFromSingleBucket,
43+
List<PipelineAggregator> pipelineAggregators, Map<String, Object> metaData) throws IOException {
44+
MultiValuesSource.AnyMultiValuesSource valuesSources = new MultiValuesSource
45+
.AnyMultiValuesSource(configs, searchContext.getQueryShardContext());
46+
return new GeoLineAggregator(name, valuesSources, searchContext, parent, pipelineAggregators, metaData);
47+
}
48+
}

0 commit comments

Comments
 (0)