Skip to content

Commit 9137fcc

Browse files
kimchykimchy
committed
move geo distance sorting to use new field data
1 parent d5e70a2 commit 9137fcc

File tree

3 files changed

+208
-2
lines changed

3 files changed

+208
-2
lines changed

‎src/main/java/org/elasticsearch/index/fielddata/fieldcomparator/GeoDistanceComparator.java

Lines changed: 133 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,133 @@
1+
/*
2+
* Licensed to ElasticSearch and Shay Banon under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. ElasticSearch licenses this
6+
* file to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package org.elasticsearch.index.fielddata.fieldcomparator;
21+
22+
import org.apache.lucene.index.AtomicReaderContext;
23+
import org.apache.lucene.search.FieldComparator;
24+
import org.elasticsearch.common.unit.DistanceUnit;
25+
import org.elasticsearch.index.fielddata.GeoPointValues;
26+
import org.elasticsearch.index.fielddata.IndexGeoPointFieldData;
27+
import org.elasticsearch.index.mapper.geo.GeoPoint;
28+
import org.elasticsearch.index.search.geo.GeoDistance;
29+
30+
import java.io.IOException;
31+
32+
/**
33+
*/
34+
public class GeoDistanceComparator extends FieldComparator<Double> {
35+
36+
protected final IndexGeoPointFieldData indexFieldData;
37+
38+
protected final double lat;
39+
protected final double lon;
40+
protected final DistanceUnit unit;
41+
protected final GeoDistance geoDistance;
42+
protected final GeoDistance.FixedSourceDistance fixedSourceDistance;
43+
44+
private final double[] values;
45+
private double bottom;
46+
47+
private GeoPointValues readerValues;
48+
49+
public GeoDistanceComparator(int numHits, IndexGeoPointFieldData indexFieldData, double lat, double lon, DistanceUnit unit, GeoDistance geoDistance) {
50+
this.values = new double[numHits];
51+
this.indexFieldData = indexFieldData;
52+
this.lat = lat;
53+
this.lon = lon;
54+
this.unit = unit;
55+
this.geoDistance = geoDistance;
56+
this.fixedSourceDistance = geoDistance.fixedSourceDistance(lat, lon, unit);
57+
}
58+
59+
@Override
60+
public FieldComparator<Double> setNextReader(AtomicReaderContext context) throws IOException {
61+
this.readerValues = indexFieldData.load(context).getGeoPointValues();
62+
return this;
63+
}
64+
65+
@Override
66+
public int compare(int slot1, int slot2) {
67+
final double v1 = values[slot1];
68+
final double v2 = values[slot2];
69+
if (v1 > v2) {
70+
return 1;
71+
} else if (v1 < v2) {
72+
return -1;
73+
} else {
74+
return 0;
75+
}
76+
}
77+
78+
@Override
79+
public int compareBottom(int doc) {
80+
double distance;
81+
GeoPoint geoPoint = readerValues.getValue(doc);
82+
if (geoPoint == null) {
83+
// is this true? push this to the "end"
84+
distance = Double.MAX_VALUE;
85+
} else {
86+
distance = fixedSourceDistance.calculate(geoPoint.lat(), geoPoint.lon());
87+
}
88+
final double v2 = distance;
89+
if (bottom > v2) {
90+
return 1;
91+
} else if (bottom < v2) {
92+
return -1;
93+
} else {
94+
return 0;
95+
}
96+
}
97+
98+
@Override
99+
public int compareDocToValue(int doc, Double distance2) throws IOException {
100+
double distance1;
101+
GeoPoint geoPoint = readerValues.getValue(doc);
102+
if (geoPoint == null) {
103+
// is this true? push this to the "end"
104+
distance1 = Double.MAX_VALUE;
105+
} else {
106+
distance1 = fixedSourceDistance.calculate(geoPoint.lat(), geoPoint.lon());
107+
}
108+
return (int) (distance1 - distance2);
109+
}
110+
111+
@Override
112+
public void copy(int slot, int doc) {
113+
double distance;
114+
GeoPoint geoPoint = readerValues.getValue(doc);
115+
if (geoPoint == null) {
116+
// is this true? push this to the "end"
117+
distance = Double.MAX_VALUE;
118+
} else {
119+
distance = fixedSourceDistance.calculate(geoPoint.lat(), geoPoint.lon());
120+
}
121+
values[slot] = distance;
122+
}
123+
124+
@Override
125+
public void setBottom(final int bottom) {
126+
this.bottom = values[bottom];
127+
}
128+
129+
@Override
130+
public Double value(int slot) {
131+
return values[slot];
132+
}
133+
}

‎src/main/java/org/elasticsearch/index/fielddata/fieldcomparator/GeoDistanceComparatorSource.java

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
/*
2+
* Licensed to ElasticSearch and Shay Banon under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. ElasticSearch licenses this
6+
* file to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package org.elasticsearch.index.fielddata.fieldcomparator;
21+
22+
import org.apache.lucene.search.FieldComparator;
23+
import org.apache.lucene.search.SortField;
24+
import org.elasticsearch.common.unit.DistanceUnit;
25+
import org.elasticsearch.index.fielddata.IndexFieldData;
26+
import org.elasticsearch.index.fielddata.IndexGeoPointFieldData;
27+
import org.elasticsearch.index.search.geo.GeoDistance;
28+
29+
import java.io.IOException;
30+
31+
/**
32+
*/
33+
public class GeoDistanceComparatorSource extends IndexFieldData.XFieldComparatorSource {
34+
35+
private final IndexGeoPointFieldData indexFieldData;
36+
private final double lat;
37+
private final double lon;
38+
private final DistanceUnit unit;
39+
private final GeoDistance geoDistance;
40+
41+
public GeoDistanceComparatorSource(IndexGeoPointFieldData indexFieldData, double lat, double lon, DistanceUnit unit, GeoDistance geoDistance) {
42+
this.indexFieldData = indexFieldData;
43+
this.lat = lat;
44+
this.lon = lon;
45+
this.unit = unit;
46+
this.geoDistance = geoDistance;
47+
}
48+
49+
@Override
50+
public SortField.Type reducedType() {
51+
return SortField.Type.DOUBLE;
52+
}
53+
54+
@Override
55+
public FieldComparator<?> newComparator(String fieldname, int numHits, int sortPos, boolean reversed) throws IOException {
56+
assert indexFieldData.getFieldNames().indexName().equals(fieldname);
57+
// TODO support multi value?
58+
return new GeoDistanceComparator(numHits, indexFieldData, lat, lon, unit, geoDistance);
59+
}
60+
}

‎src/main/java/org/elasticsearch/search/sort/GeoDistanceSortParser.java

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,10 +20,17 @@
2020
package org.elasticsearch.search.sort;
2121

2222
import org.apache.lucene.search.SortField;
23+
import org.elasticsearch.ElasticSearchIllegalArgumentException;
2324
import org.elasticsearch.common.unit.DistanceUnit;
2425
import org.elasticsearch.common.xcontent.XContentParser;
26+
import org.elasticsearch.index.fielddata.IndexGeoPointFieldData;
27+
import org.elasticsearch.index.fielddata.fieldcomparator.GeoDistanceComparatorSource;
28+
import org.elasticsearch.index.mapper.FieldMapper;
2529
import org.elasticsearch.index.mapper.geo.GeoPointFieldMapper;
26-
import org.elasticsearch.index.search.geo.*;
30+
import org.elasticsearch.index.search.geo.GeoDistance;
31+
import org.elasticsearch.index.search.geo.GeoHashUtils;
32+
import org.elasticsearch.index.search.geo.GeoUtils;
33+
import org.elasticsearch.index.search.geo.Point;
2734
import org.elasticsearch.search.internal.SearchContext;
2835

2936
/**
@@ -117,6 +124,12 @@ public SortField parse(XContentParser parser, SearchContext context) throws Exce
117124
lon = point.lon;
118125
}
119126

120-
return new SortField(fieldName, GeoDistanceDataComparator.comparatorSource(fieldName, lat, lon, unit, geoDistance, context.fieldDataCache(), context.mapperService()), reverse);
127+
FieldMapper mapper = context.smartNameFieldMapper(fieldName);
128+
if (mapper == null) {
129+
throw new ElasticSearchIllegalArgumentException("failed to find mapper for [" + fieldName + "] for geo distance based sort");
130+
}
131+
IndexGeoPointFieldData indexFieldData = context.fieldData().getForField(mapper);
132+
133+
return new SortField(fieldName, new GeoDistanceComparatorSource(indexFieldData, lat, lon, unit, geoDistance), reverse);
121134
}
122135
}

0 commit comments

Comments
 (0)