diff --git a/src/main/java/org/elasticsearch/index/fielddata/fieldcomparator/GeoDistanceComparator.java b/src/main/java/org/elasticsearch/index/fielddata/fieldcomparator/GeoDistanceComparator.java new file mode 100644 index 00000000000..b415a8cebf9 --- /dev/null +++ b/src/main/java/org/elasticsearch/index/fielddata/fieldcomparator/GeoDistanceComparator.java @@ -0,0 +1,133 @@ +/* + * Licensed to ElasticSearch and Shay Banon under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. ElasticSearch licenses this + * file to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.elasticsearch.index.fielddata.fieldcomparator; + +import org.apache.lucene.index.AtomicReaderContext; +import org.apache.lucene.search.FieldComparator; +import org.elasticsearch.common.unit.DistanceUnit; +import org.elasticsearch.index.fielddata.GeoPointValues; +import org.elasticsearch.index.fielddata.IndexGeoPointFieldData; +import org.elasticsearch.index.mapper.geo.GeoPoint; +import org.elasticsearch.index.search.geo.GeoDistance; + +import java.io.IOException; + +/** + */ +public class GeoDistanceComparator extends FieldComparator { + + protected final IndexGeoPointFieldData indexFieldData; + + protected final double lat; + protected final double lon; + protected final DistanceUnit unit; + protected final GeoDistance geoDistance; + protected final GeoDistance.FixedSourceDistance fixedSourceDistance; + + private final double[] values; + private double bottom; + + private GeoPointValues readerValues; + + public GeoDistanceComparator(int numHits, IndexGeoPointFieldData indexFieldData, double lat, double lon, DistanceUnit unit, GeoDistance geoDistance) { + this.values = new double[numHits]; + this.indexFieldData = indexFieldData; + this.lat = lat; + this.lon = lon; + this.unit = unit; + this.geoDistance = geoDistance; + this.fixedSourceDistance = geoDistance.fixedSourceDistance(lat, lon, unit); + } + + @Override + public FieldComparator setNextReader(AtomicReaderContext context) throws IOException { + this.readerValues = indexFieldData.load(context).getGeoPointValues(); + return this; + } + + @Override + public int compare(int slot1, int slot2) { + final double v1 = values[slot1]; + final double v2 = values[slot2]; + if (v1 > v2) { + return 1; + } else if (v1 < v2) { + return -1; + } else { + return 0; + } + } + + @Override + public int compareBottom(int doc) { + double distance; + GeoPoint geoPoint = readerValues.getValue(doc); + if (geoPoint == null) { + // is this true? push this to the "end" + distance = Double.MAX_VALUE; + } else { + distance = fixedSourceDistance.calculate(geoPoint.lat(), geoPoint.lon()); + } + final double v2 = distance; + if (bottom > v2) { + return 1; + } else if (bottom < v2) { + return -1; + } else { + return 0; + } + } + + @Override + public int compareDocToValue(int doc, Double distance2) throws IOException { + double distance1; + GeoPoint geoPoint = readerValues.getValue(doc); + if (geoPoint == null) { + // is this true? push this to the "end" + distance1 = Double.MAX_VALUE; + } else { + distance1 = fixedSourceDistance.calculate(geoPoint.lat(), geoPoint.lon()); + } + return (int) (distance1 - distance2); + } + + @Override + public void copy(int slot, int doc) { + double distance; + GeoPoint geoPoint = readerValues.getValue(doc); + if (geoPoint == null) { + // is this true? push this to the "end" + distance = Double.MAX_VALUE; + } else { + distance = fixedSourceDistance.calculate(geoPoint.lat(), geoPoint.lon()); + } + values[slot] = distance; + } + + @Override + public void setBottom(final int bottom) { + this.bottom = values[bottom]; + } + + @Override + public Double value(int slot) { + return values[slot]; + } +} diff --git a/src/main/java/org/elasticsearch/index/fielddata/fieldcomparator/GeoDistanceComparatorSource.java b/src/main/java/org/elasticsearch/index/fielddata/fieldcomparator/GeoDistanceComparatorSource.java new file mode 100644 index 00000000000..c4266121fc9 --- /dev/null +++ b/src/main/java/org/elasticsearch/index/fielddata/fieldcomparator/GeoDistanceComparatorSource.java @@ -0,0 +1,60 @@ +/* + * Licensed to ElasticSearch and Shay Banon under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. ElasticSearch licenses this + * file to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.elasticsearch.index.fielddata.fieldcomparator; + +import org.apache.lucene.search.FieldComparator; +import org.apache.lucene.search.SortField; +import org.elasticsearch.common.unit.DistanceUnit; +import org.elasticsearch.index.fielddata.IndexFieldData; +import org.elasticsearch.index.fielddata.IndexGeoPointFieldData; +import org.elasticsearch.index.search.geo.GeoDistance; + +import java.io.IOException; + +/** + */ +public class GeoDistanceComparatorSource extends IndexFieldData.XFieldComparatorSource { + + private final IndexGeoPointFieldData indexFieldData; + private final double lat; + private final double lon; + private final DistanceUnit unit; + private final GeoDistance geoDistance; + + public GeoDistanceComparatorSource(IndexGeoPointFieldData indexFieldData, double lat, double lon, DistanceUnit unit, GeoDistance geoDistance) { + this.indexFieldData = indexFieldData; + this.lat = lat; + this.lon = lon; + this.unit = unit; + this.geoDistance = geoDistance; + } + + @Override + public SortField.Type reducedType() { + return SortField.Type.DOUBLE; + } + + @Override + public FieldComparator newComparator(String fieldname, int numHits, int sortPos, boolean reversed) throws IOException { + assert indexFieldData.getFieldNames().indexName().equals(fieldname); + // TODO support multi value? + return new GeoDistanceComparator(numHits, indexFieldData, lat, lon, unit, geoDistance); + } +} diff --git a/src/main/java/org/elasticsearch/search/sort/GeoDistanceSortParser.java b/src/main/java/org/elasticsearch/search/sort/GeoDistanceSortParser.java index be41bc03f55..9065467645a 100644 --- a/src/main/java/org/elasticsearch/search/sort/GeoDistanceSortParser.java +++ b/src/main/java/org/elasticsearch/search/sort/GeoDistanceSortParser.java @@ -20,10 +20,17 @@ package org.elasticsearch.search.sort; import org.apache.lucene.search.SortField; +import org.elasticsearch.ElasticSearchIllegalArgumentException; import org.elasticsearch.common.unit.DistanceUnit; import org.elasticsearch.common.xcontent.XContentParser; +import org.elasticsearch.index.fielddata.IndexGeoPointFieldData; +import org.elasticsearch.index.fielddata.fieldcomparator.GeoDistanceComparatorSource; +import org.elasticsearch.index.mapper.FieldMapper; import org.elasticsearch.index.mapper.geo.GeoPointFieldMapper; -import org.elasticsearch.index.search.geo.*; +import org.elasticsearch.index.search.geo.GeoDistance; +import org.elasticsearch.index.search.geo.GeoHashUtils; +import org.elasticsearch.index.search.geo.GeoUtils; +import org.elasticsearch.index.search.geo.Point; import org.elasticsearch.search.internal.SearchContext; /** @@ -117,6 +124,12 @@ public class GeoDistanceSortParser implements SortParser { lon = point.lon; } - return new SortField(fieldName, GeoDistanceDataComparator.comparatorSource(fieldName, lat, lon, unit, geoDistance, context.fieldDataCache(), context.mapperService()), reverse); + FieldMapper mapper = context.smartNameFieldMapper(fieldName); + if (mapper == null) { + throw new ElasticSearchIllegalArgumentException("failed to find mapper for [" + fieldName + "] for geo distance based sort"); + } + IndexGeoPointFieldData indexFieldData = context.fieldData().getForField(mapper); + + return new SortField(fieldName, new GeoDistanceComparatorSource(indexFieldData, lat, lon, unit, geoDistance), reverse); } }