From b8b21a3363ef531efc6bd9fe9bf32d98a23212ae Mon Sep 17 00:00:00 2001 From: kimchy Date: Sun, 1 Aug 2010 13:31:03 +0300 Subject: [PATCH] Geo: `geo_distance` facet, closes #286. --- .../common/unit/DistanceUnit.java | 43 ++- .../index/field/function/FieldsFunction.java | 2 +- .../function/script/ScriptFieldsFunction.java | 12 +- .../elasticsearch/search/facets/Facet.java | 9 +- .../search/facets/FacetBuilders.java | 5 + .../search/facets/FacetsParseElement.java | 17 +- .../collector/FacetCollectorParser.java | 2 +- .../facets/geodistance/GeoDistanceFacet.java | 120 +++++++++ .../geodistance/GeoDistanceFacetBuilder.java | 250 ++++++++++++++++++ .../GeoDistanceFacetCollector.java | 127 +++++++++ .../GeoDistanceFacetCollectorParser.java | 171 ++++++++++++ .../geodistance/InternalGeoDistanceFacet.java | 185 +++++++++++++ .../ScriptGeoDistanceFacetCollector.java | 86 ++++++ .../ValueGeoDistanceFacetCollector.java | 111 ++++++++ .../HistogramFacetCollectorParser.java | 14 +- .../facets/internal/InternalFacets.java | 3 + .../query/QueryFacetCollectorParser.java | 4 +- .../StatisticalFacetCollectorParser.java | 4 +- .../terms/TermsFacetCollectorParser.java | 4 +- .../search/geo/GeoDistanceFacetTests.java | 222 ++++++++++++++++ 20 files changed, 1364 insertions(+), 27 deletions(-) create mode 100644 modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/geodistance/GeoDistanceFacet.java create mode 100644 modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/geodistance/GeoDistanceFacetBuilder.java create mode 100644 modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/geodistance/GeoDistanceFacetCollector.java create mode 100644 modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/geodistance/GeoDistanceFacetCollectorParser.java create mode 100644 modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/geodistance/InternalGeoDistanceFacet.java create mode 100644 modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/geodistance/ScriptGeoDistanceFacetCollector.java create mode 100644 modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/geodistance/ValueGeoDistanceFacetCollector.java create mode 100644 modules/test/integration/src/test/java/org/elasticsearch/test/integration/search/geo/GeoDistanceFacetTests.java diff --git a/modules/elasticsearch/src/main/java/org/elasticsearch/common/unit/DistanceUnit.java b/modules/elasticsearch/src/main/java/org/elasticsearch/common/unit/DistanceUnit.java index 795bf171075..9304b6ffee9 100644 --- a/modules/elasticsearch/src/main/java/org/elasticsearch/common/unit/DistanceUnit.java +++ b/modules/elasticsearch/src/main/java/org/elasticsearch/common/unit/DistanceUnit.java @@ -20,13 +20,19 @@ package org.elasticsearch.common.unit; import org.elasticsearch.ElasticSearchIllegalArgumentException; +import org.elasticsearch.common.io.stream.StreamInput; +import org.elasticsearch.common.io.stream.StreamOutput; + +import java.io.IOException; /** * @author kimchy (shay.banon) */ public enum DistanceUnit { MILES(3959, 24902) { - @Override public double toMiles(double distance) { + @Override public String toString() { + return "miles"; + }@Override public double toMiles(double distance) { return distance; }@Override public double toKilometers(double distance) { return distance / MILES_KILOMETRES_RATIO; @@ -35,7 +41,9 @@ public enum DistanceUnit { return distance + "mi"; }}, KILOMETERS(6371, 40076) { - @Override public double toMiles(double distance) { + @Override public String toString() { + return "km"; + }@Override public double toMiles(double distance) { return distance * MILES_KILOMETRES_RATIO; }@Override public double toKilometers(double distance) { return distance; @@ -73,6 +81,18 @@ public enum DistanceUnit { } } + public static DistanceUnit parseUnit(String distance, DistanceUnit defaultUnit) { + if (distance.endsWith("mi")) { + return MILES; + } else if (distance.endsWith("miles")) { + return MILES; + } else if (distance.endsWith("km")) { + return KILOMETERS; + } else { + return defaultUnit; + } + } + protected final double earthCircumference; protected final double earthRadius; @@ -97,4 +117,23 @@ public enum DistanceUnit { } throw new ElasticSearchIllegalArgumentException("No distance unit match [" + unit + "]"); } + + public static void writeDistanceUnit(StreamOutput out, DistanceUnit unit) throws IOException { + if (unit == MILES) { + out.writeByte((byte) 0); + } else if (unit == KILOMETERS) { + out.writeByte((byte) 1); + } + } + + public static DistanceUnit readDistanceUnit(StreamInput in) throws IOException { + byte b = in.readByte(); + if (b == 0) { + return MILES; + } else if (b == 1) { + return KILOMETERS; + } else { + throw new ElasticSearchIllegalArgumentException("No type for distance unit matching [" + b + "]"); + } + } } diff --git a/modules/elasticsearch/src/main/java/org/elasticsearch/index/field/function/FieldsFunction.java b/modules/elasticsearch/src/main/java/org/elasticsearch/index/field/function/FieldsFunction.java index 1c04f2840f4..1d9f45a49ff 100644 --- a/modules/elasticsearch/src/main/java/org/elasticsearch/index/field/function/FieldsFunction.java +++ b/modules/elasticsearch/src/main/java/org/elasticsearch/index/field/function/FieldsFunction.java @@ -35,5 +35,5 @@ public interface FieldsFunction { * @param vars The vars providing additional parameters, should be reused and has values added to it in execute * @return */ - Object execute(int docId, Map vars); + Object execute(int docId, Map vars); } diff --git a/modules/elasticsearch/src/main/java/org/elasticsearch/index/field/function/script/ScriptFieldsFunction.java b/modules/elasticsearch/src/main/java/org/elasticsearch/index/field/function/script/ScriptFieldsFunction.java index 98fff88b87e..68603d5760d 100644 --- a/modules/elasticsearch/src/main/java/org/elasticsearch/index/field/function/script/ScriptFieldsFunction.java +++ b/modules/elasticsearch/src/main/java/org/elasticsearch/index/field/function/script/ScriptFieldsFunction.java @@ -47,6 +47,12 @@ public class ScriptFieldsFunction implements FieldsFunction, Map { } }; + private static ThreadLocal>> cachedVars = new ThreadLocal>>() { + @Override protected ThreadLocals.CleanableValue> initialValue() { + return new ThreadLocals.CleanableValue>(new HashMap()); + } + }; + final Object script; final MapperService mapperService; @@ -73,8 +79,12 @@ public class ScriptFieldsFunction implements FieldsFunction, Map { localCacheFieldData.clear(); } - @Override public Object execute(int docId, Map vars) { + @Override public Object execute(int docId, Map vars) { this.docId = docId; + if (vars == null) { + vars = cachedVars.get().get(); + vars.clear(); + } vars.put("doc", this); return scriptService.execute(script, vars); } diff --git a/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/Facet.java b/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/Facet.java index 0ff4dad9071..d1ce75dc63a 100644 --- a/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/Facet.java +++ b/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/Facet.java @@ -20,6 +20,7 @@ package org.elasticsearch.search.facets; import org.elasticsearch.ElasticSearchIllegalArgumentException; +import org.elasticsearch.search.facets.geodistance.GeoDistanceFacet; import org.elasticsearch.search.facets.histogram.HistogramFacet; import org.elasticsearch.search.facets.query.QueryFacet; import org.elasticsearch.search.facets.statistical.StatisticalFacet; @@ -51,7 +52,11 @@ public interface Facet { /** * Histogram facet type, matching {@link HistogramFacet}. */ - HISTOGRAM(3, HistogramFacet.class); + HISTOGRAM(3, HistogramFacet.class), + /** + * Geo Distance facet type, matching {@link GeoDistanceFacet}. + */ + GEO_DISTANCE(4, GeoDistanceFacet.class); private int id; @@ -89,6 +94,8 @@ public interface Facet { return STATISTICAL; } else if (id == 3) { return HISTOGRAM; + } else if (id == 4) { + return GEO_DISTANCE; } else { throw new ElasticSearchIllegalArgumentException("No match for id [" + id + "]"); } diff --git a/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/FacetBuilders.java b/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/FacetBuilders.java index f6288018c48..07f1d35c019 100644 --- a/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/FacetBuilders.java +++ b/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/FacetBuilders.java @@ -20,6 +20,7 @@ package org.elasticsearch.search.facets; import org.elasticsearch.index.query.xcontent.XContentQueryBuilder; +import org.elasticsearch.search.facets.geodistance.GeoDistanceFacetBuilder; import org.elasticsearch.search.facets.histogram.HistogramFacetBuilder; import org.elasticsearch.search.facets.histogram.HistogramScriptFacetBuilder; import org.elasticsearch.search.facets.query.QueryFacetBuilder; @@ -59,4 +60,8 @@ public class FacetBuilders { public static HistogramScriptFacetBuilder histogramScriptFacet(String facetName) { return new HistogramScriptFacetBuilder(facetName); } + + public static GeoDistanceFacetBuilder geoDistanceFacet(String facetName) { + return new GeoDistanceFacetBuilder(facetName); + } } diff --git a/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/FacetsParseElement.java b/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/FacetsParseElement.java index 718399559d4..7d70d812caf 100644 --- a/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/FacetsParseElement.java +++ b/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/FacetsParseElement.java @@ -29,6 +29,7 @@ import org.elasticsearch.search.SearchParseElement; import org.elasticsearch.search.SearchParseException; import org.elasticsearch.search.facets.collector.FacetCollector; import org.elasticsearch.search.facets.collector.FacetCollectorParser; +import org.elasticsearch.search.facets.geodistance.GeoDistanceFacetCollectorParser; import org.elasticsearch.search.facets.histogram.HistogramFacetCollectorParser; import org.elasticsearch.search.facets.query.QueryFacetCollectorParser; import org.elasticsearch.search.facets.statistical.StatisticalFacetCollectorParser; @@ -64,13 +65,21 @@ public class FacetsParseElement implements SearchParseElement { public FacetsParseElement() { MapBuilder builder = newMapBuilder(); - builder.put(TermsFacetCollectorParser.NAME, new TermsFacetCollectorParser()); - builder.put(QueryFacetCollectorParser.NAME, new QueryFacetCollectorParser()); - builder.put(StatisticalFacetCollectorParser.NAME, new StatisticalFacetCollectorParser()); - builder.put(HistogramFacetCollectorParser.NAME, new HistogramFacetCollectorParser()); + addFacetParser(builder, new TermsFacetCollectorParser()); + addFacetParser(builder, new QueryFacetCollectorParser()); + addFacetParser(builder, new StatisticalFacetCollectorParser()); + addFacetParser(builder, new HistogramFacetCollectorParser()); + addFacetParser(builder, new GeoDistanceFacetCollectorParser()); + addFacetParser(builder, new GeoDistanceFacetCollectorParser()); this.facetCollectorParsers = builder.immutableMap(); } + private void addFacetParser(MapBuilder builder, FacetCollectorParser facetCollectorParser) { + for (String s : facetCollectorParser.names()) { + builder.put(s, facetCollectorParser); + } + } + @Override public void parse(XContentParser parser, SearchContext context) throws Exception { XContentParser.Token token; diff --git a/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/collector/FacetCollectorParser.java b/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/collector/FacetCollectorParser.java index b7c971cd9cb..d08a716bae8 100644 --- a/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/collector/FacetCollectorParser.java +++ b/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/collector/FacetCollectorParser.java @@ -29,7 +29,7 @@ import java.io.IOException; */ public interface FacetCollectorParser { - String name(); + String[] names(); FacetCollector parser(String facetName, XContentParser parser, SearchContext context) throws IOException; } diff --git a/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/geodistance/GeoDistanceFacet.java b/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/geodistance/GeoDistanceFacet.java new file mode 100644 index 00000000000..c2c32516e16 --- /dev/null +++ b/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/geodistance/GeoDistanceFacet.java @@ -0,0 +1,120 @@ +/* + * Licensed to Elastic Search and Shay Banon under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. Elastic Search 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.search.facets.geodistance; + +import org.elasticsearch.common.unit.DistanceUnit; +import org.elasticsearch.search.facets.Facet; + +import java.util.List; + +/** + * @author kimchy (shay.banon) + */ +public interface GeoDistanceFacet extends Facet, Iterable { + + String fieldName(); + + String getFieldName(); + + String valueFieldName(); + + String getValueFieldName(); + + DistanceUnit unit(); + + DistanceUnit getUnit(); + + /** + * An ordered list of histogram facet entries. + */ + List entries(); + + /** + * An ordered list of histogram facet entries. + */ + List getEntries(); + + public class Entry { + + double from = Double.NEGATIVE_INFINITY; + + double to = Double.POSITIVE_INFINITY; + + long count; + + double total; + + Entry() { + } + + public Entry(double from, double to, long count, double total) { + this.from = from; + this.to = to; + this.count = count; + this.total = total; + } + + public double from() { + return this.from; + } + + public double getFrom() { + return from(); + } + + public double to() { + return this.to; + } + + public double getTo() { + return to(); + } + + public long count() { + return this.count; + } + + public long getCount() { + return count(); + } + + public double total() { + return this.total; + } + + public double getTotal() { + return total(); + } + + /** + * The mean of this facet interval. + */ + public double mean() { + return total / count; + } + + /** + * The mean of this facet interval. + */ + public double getMean() { + return mean(); + } + } +} diff --git a/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/geodistance/GeoDistanceFacetBuilder.java b/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/geodistance/GeoDistanceFacetBuilder.java new file mode 100644 index 00000000000..afefe9f36a6 --- /dev/null +++ b/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/geodistance/GeoDistanceFacetBuilder.java @@ -0,0 +1,250 @@ +/* + * Licensed to Elastic Search and Shay Banon under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. Elastic Search 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.search.facets.geodistance; + +import org.elasticsearch.common.collect.Lists; +import org.elasticsearch.common.collect.Maps; +import org.elasticsearch.common.lucene.geo.GeoDistance; +import org.elasticsearch.common.unit.DistanceUnit; +import org.elasticsearch.common.xcontent.builder.XContentBuilder; +import org.elasticsearch.search.builder.SearchSourceBuilderException; +import org.elasticsearch.search.facets.AbstractFacetBuilder; + +import java.io.IOException; +import java.util.List; +import java.util.Map; + +/** + * A geo distance builder allowing to create a facet of distances from a specific location including the + * number of hits within each distance range, and aggregated data (like totals of either the distance or + * cusotm value fields). + * + * @author kimchy (shay.banon) + */ +public class GeoDistanceFacetBuilder extends AbstractFacetBuilder { + + private String fieldName; + + private String valueFieldName; + + private double lat; + + private double lon; + + private String geohash; + + private GeoDistance geoDistance; + + private DistanceUnit unit; + + private Map params; + + private String valueScript; + + private List entries = Lists.newArrayList(); + + /** + * Constructs a new geo distance with the provided facet name. + */ + public GeoDistanceFacetBuilder(String name) { + super(name); + } + + /** + * The geo point field that will be used to extract the document location(s). + */ + public GeoDistanceFacetBuilder field(String fieldName) { + this.fieldName = fieldName; + return this; + } + + /** + * A custom value field (numeric) that will be used to provide aggregated data for each facet (for example, total). + */ + public GeoDistanceFacetBuilder valueField(String valueFieldName) { + this.valueFieldName = valueFieldName; + return this; + } + + /** + * A custom value script (result is numeric) that will be used to provide aggregated data for each facet (for example, total). + */ + public GeoDistanceFacetBuilder valueScript(String valueScript) { + this.valueScript = valueScript; + return this; + } + + /** + * Parameters for {@link #valueScript(String)} to improve performance when executing the same script with different parameters. + */ + public GeoDistanceFacetBuilder scriptParam(String name, Object value) { + if (params == null) { + params = Maps.newHashMap(); + } + params.put(name, value); + return this; + } + + /** + * The point to create the range distance facets from. + * + * @param lat latitude. + * @param lon longitude. + */ + public GeoDistanceFacetBuilder point(double lat, double lon) { + this.lat = lat; + this.lon = lon; + return this; + } + + /** + * The latitude to create the range distance facets from. + */ + public GeoDistanceFacetBuilder lat(double lat) { + this.lat = lat; + return this; + } + + /** + * The longitude to create the range distance facets from. + */ + public GeoDistanceFacetBuilder lon(double lon) { + this.lon = lon; + return this; + } + + /** + * The geohash of the geo point to create the range distance facets from. + */ + public GeoDistanceFacetBuilder geohash(String geohash) { + this.geohash = geohash; + return this; + } + + /** + * The geo distance type used to compute the distnace. + */ + public GeoDistanceFacetBuilder geoDistance(GeoDistance geoDistance) { + this.geoDistance = geoDistance; + return this; + } + + /** + * Adds a range entry with explicit from and to. + * + * @param from The from distance limit + * @param to The to distance limit + */ + public GeoDistanceFacetBuilder addRange(double from, double to) { + entries.add(new Entry(from, to)); + return this; + } + + /** + * Adds a range entry with explicit from and unbounded to. + * + * @param from the from distance limit, to is unbounded. + */ + public GeoDistanceFacetBuilder addUnboundedTo(double from) { + entries.add(new Entry(from, Double.POSITIVE_INFINITY)); + return this; + } + + /** + * Adds a range entry with explicit to and unbounded from. + * + * @param to the to distance limit, from is unbounded. + */ + public GeoDistanceFacetBuilder addUnboundedFrom(double to) { + entries.add(new Entry(Double.NEGATIVE_INFINITY, to)); + return this; + } + + /** + * The distance unit to use. Defaults to {@link org.elasticsearch.common.unit.DistanceUnit#KILOMETERS} + */ + public GeoDistanceFacetBuilder unit(DistanceUnit unit) { + this.unit = unit; + return this; + } + + @Override public void toXContent(XContentBuilder builder, Params params) throws IOException { + if (fieldName == null) { + throw new SearchSourceBuilderException("field must be set on geo_distance facet for facet [" + name + "]"); + } + if (entries.isEmpty()) { + throw new SearchSourceBuilderException("at least one range must be defined for geo_distance face [" + name + "]"); + } + + builder.startObject(name); + builder.startObject(GeoDistanceFacetCollectorParser.NAME); + + if (geohash != null) { + builder.field(fieldName, geohash); + } else { + builder.startArray(fieldName).value(lat).value(lon).endArray(); + } + + if (valueFieldName != null) { + builder.field("value_field", valueFieldName); + } + + if (valueScript != null) { + builder.field("value_script", valueScript); + if (this.params != null) { + builder.field("params"); + builder.map(this.params); + } + } + + builder.startArray("ranges"); + for (Entry entry : entries) { + builder.startObject(); + if (!Double.isInfinite(entry.from)) { + builder.field("from", entry.from); + } + if (!Double.isInfinite(entry.to)) { + builder.field("to", entry.to); + } + builder.endObject(); + } + builder.endArray(); + + if (unit != null) { + builder.field("unit", unit); + } + if (geoDistance != null) { + builder.field("distance_type", geoDistance.name().toLowerCase()); + } + + builder.endObject(); + builder.endObject(); + } + + private static class Entry { + final double from; + final double to; + + private Entry(double from, double to) { + this.from = from; + this.to = to; + } + } +} diff --git a/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/geodistance/GeoDistanceFacetCollector.java b/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/geodistance/GeoDistanceFacetCollector.java new file mode 100644 index 00000000000..e7564eb2051 --- /dev/null +++ b/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/geodistance/GeoDistanceFacetCollector.java @@ -0,0 +1,127 @@ +/* + * Licensed to Elastic Search and Shay Banon under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. Elastic Search 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.search.facets.geodistance; + +import org.apache.lucene.index.IndexReader; +import org.elasticsearch.common.lucene.geo.GeoDistance; +import org.elasticsearch.common.unit.DistanceUnit; +import org.elasticsearch.index.cache.field.data.FieldDataCache; +import org.elasticsearch.index.field.data.FieldData; +import org.elasticsearch.index.field.data.NumericFieldData; +import org.elasticsearch.index.mapper.FieldMapper; +import org.elasticsearch.index.mapper.MapperService; +import org.elasticsearch.index.mapper.xcontent.XContentGeoPointFieldMapper; +import org.elasticsearch.search.facets.Facet; +import org.elasticsearch.search.facets.FacetPhaseExecutionException; +import org.elasticsearch.search.facets.support.AbstractFacetCollector; + +import java.io.IOException; + +/** + * @author kimchy (shay.banon) + */ +public class GeoDistanceFacetCollector extends AbstractFacetCollector { + + protected final String fieldName; + + protected final String indexLatFieldName; + + protected final String indexLonFieldName; + + protected final double lat; + + protected final double lon; + + protected final DistanceUnit unit; + + protected final GeoDistance geoDistance; + + protected final FieldDataCache fieldDataCache; + + protected final FieldData.Type fieldDataType; + + protected NumericFieldData latFieldData; + + protected NumericFieldData lonFieldData; + + protected final GeoDistanceFacet.Entry[] entries; + + public GeoDistanceFacetCollector(String facetName, String fieldName, double lat, double lon, DistanceUnit unit, GeoDistance geoDistance, + GeoDistanceFacet.Entry[] entries, FieldDataCache fieldDataCache, MapperService mapperService) { + super(facetName); + this.fieldName = fieldName; + this.lat = lat; + this.lon = lon; + this.unit = unit; + this.entries = entries; + this.geoDistance = geoDistance; + this.fieldDataCache = fieldDataCache; + + FieldMapper mapper = mapperService.smartNameFieldMapper(fieldName + XContentGeoPointFieldMapper.Names.LAT_SUFFIX); + if (mapper == null) { + throw new FacetPhaseExecutionException(facetName, "No mapping found for field [" + fieldName + "]"); + } + this.indexLatFieldName = mapper.names().indexName(); + + mapper = mapperService.smartNameFieldMapper(fieldName + XContentGeoPointFieldMapper.Names.LON_SUFFIX); + if (mapper == null) { + throw new FacetPhaseExecutionException(facetName, "No mapping found for field [" + fieldName + "]"); + } + this.indexLonFieldName = mapper.names().indexName(); + this.fieldDataType = mapper.fieldDataType(); + } + + @Override protected void doSetNextReader(IndexReader reader, int docBase) throws IOException { + latFieldData = (NumericFieldData) fieldDataCache.cache(fieldDataType, reader, indexLatFieldName); + lonFieldData = (NumericFieldData) fieldDataCache.cache(fieldDataType, reader, indexLonFieldName); + } + + @Override protected void doCollect(int doc) throws IOException { + if (!latFieldData.hasValue(doc) || !lonFieldData.hasValue(doc)) { + return; + } + + if (latFieldData.multiValued()) { + double[] lats = latFieldData.doubleValues(doc); + double[] lons = latFieldData.doubleValues(doc); + for (int i = 0; i < lats.length; i++) { + double distance = geoDistance.calculate(lat, lon, lats[i], lons[i], unit); + for (GeoDistanceFacet.Entry entry : entries) { + if (distance >= entry.getFrom() && distance < entry.getTo()) { + entry.count++; + entry.total += distance; + } + } + } + } else { + double distance = geoDistance.calculate(lat, lon, latFieldData.doubleValue(doc), lonFieldData.doubleValue(doc), unit); + for (GeoDistanceFacet.Entry entry : entries) { + if (distance >= entry.getFrom() && distance < entry.getTo()) { + entry.count++; + entry.total += distance; + } + } + } + } + + @Override public Facet facet() { + return new InternalGeoDistanceFacet(facetName, fieldName, fieldName, unit, entries); + } +} diff --git a/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/geodistance/GeoDistanceFacetCollectorParser.java b/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/geodistance/GeoDistanceFacetCollectorParser.java new file mode 100644 index 00000000000..7debd890f60 --- /dev/null +++ b/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/geodistance/GeoDistanceFacetCollectorParser.java @@ -0,0 +1,171 @@ +/* + * Licensed to Elastic Search and Shay Banon under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. Elastic Search 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.search.facets.geodistance; + +import org.elasticsearch.common.collect.Lists; +import org.elasticsearch.common.lucene.geo.GeoDistance; +import org.elasticsearch.common.lucene.geo.GeoHashUtils; +import org.elasticsearch.common.thread.ThreadLocals; +import org.elasticsearch.common.unit.DistanceUnit; +import org.elasticsearch.common.xcontent.XContentParser; +import org.elasticsearch.index.mapper.xcontent.XContentGeoPointFieldMapper; +import org.elasticsearch.search.facets.FacetPhaseExecutionException; +import org.elasticsearch.search.facets.collector.FacetCollector; +import org.elasticsearch.search.facets.collector.FacetCollectorParser; +import org.elasticsearch.search.internal.SearchContext; + +import java.io.IOException; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * @author kimchy (shay.banon) + */ +public class GeoDistanceFacetCollectorParser implements FacetCollectorParser { + + private static ThreadLocal>> cachedParams = new ThreadLocal>>() { + @Override protected ThreadLocals.CleanableValue> initialValue() { + return new ThreadLocals.CleanableValue>(new HashMap()); + } + }; + + public static final String NAME = "geo_distance"; + + @Override public String[] names() { + return new String[]{NAME, "geoDistance"}; + } + + @Override public FacetCollector parser(String facetName, XContentParser parser, SearchContext context) throws IOException { + String fieldName = null; + String valueFieldName = null; + String valueScript = null; + Map params = null; + double lat = Double.NaN; + double lon = Double.NaN; + DistanceUnit unit = DistanceUnit.KILOMETERS; + GeoDistance geoDistance = GeoDistance.ARC; + List entries = Lists.newArrayList(); + + XContentParser.Token token; + String currentName = parser.currentName(); + + while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) { + if (token == XContentParser.Token.FIELD_NAME) { + currentName = parser.currentName(); + } else if (token == XContentParser.Token.START_ARRAY) { + if ("ranges".equals(currentName) || "entries".equals(currentName)) { + // "ranges" : [ + // { "from" : "0', to : "12.5" } + // { "from" : "12.5" } + // ] + while ((token = parser.nextToken()) != XContentParser.Token.END_ARRAY) { + double from = Double.NEGATIVE_INFINITY; + double to = Double.POSITIVE_INFINITY; + while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) { + if (token == XContentParser.Token.FIELD_NAME) { + currentName = parser.currentName(); + } else if (token.isValue()) { + if ("from".equals(currentName)) { + from = parser.doubleValue(); + } else if ("to".equals(currentName)) { + to = parser.doubleValue(); + } + } + } + entries.add(new GeoDistanceFacet.Entry(from, to, 0, 0)); + } + } else { + token = parser.nextToken(); + lat = parser.doubleValue(); + token = parser.nextToken(); + lon = parser.doubleValue(); + while ((token = parser.nextToken()) != XContentParser.Token.END_ARRAY) { + + } + fieldName = currentName; + } + } else if (token == XContentParser.Token.START_OBJECT) { + if ("params".equals(currentName)) { + params = parser.map(); + } else { + // the json in the format of -> field : { lat : 30, lon : 12 } + fieldName = currentName; + while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) { + if (token == XContentParser.Token.FIELD_NAME) { + currentName = parser.currentName(); + } else if (token.isValue()) { + if (currentName.equals(XContentGeoPointFieldMapper.Names.LAT)) { + lat = parser.doubleValue(); + } else if (currentName.equals(XContentGeoPointFieldMapper.Names.LON)) { + lon = parser.doubleValue(); + } else if (currentName.equals(XContentGeoPointFieldMapper.Names.GEOHASH)) { + double[] values = GeoHashUtils.decode(parser.text()); + lat = values[0]; + lon = values[1]; + } + } + } + } + } else if (token.isValue()) { + if (currentName.equals("unit")) { + unit = DistanceUnit.fromString(parser.text()); + } else if (currentName.equals("distance_type") || currentName.equals("distanceType")) { + geoDistance = GeoDistance.fromString(parser.text()); + } else if ("value_field".equals(currentName) || "valueName".equals(currentName)) { + valueFieldName = parser.text(); + } else if ("value_script".equals(currentName) || "valueScript".equals(currentName)) { + valueScript = parser.text(); + } else { + // assume the value is the actual value + String value = parser.text(); + int comma = value.indexOf(','); + if (comma != -1) { + lat = Double.parseDouble(value.substring(0, comma).trim()); + lon = Double.parseDouble(value.substring(comma + 1).trim()); + } else { + double[] values = GeoHashUtils.decode(value); + lat = values[0]; + lon = values[1]; + } + + fieldName = currentName; + } + } + } + + if (Double.isNaN(lat) || Double.isNaN(lon)) { + throw new FacetPhaseExecutionException(facetName, "lat/lon not set for geo_distance facet"); + } + + if (valueFieldName != null) { + return new ValueGeoDistanceFacetCollector(facetName, fieldName, lat, lon, unit, geoDistance, entries.toArray(new GeoDistanceFacet.Entry[entries.size()]), + context.fieldDataCache(), context.mapperService(), valueFieldName); + } + + if (valueScript != null) { + return new ScriptGeoDistanceFacetCollector(facetName, fieldName, lat, lon, unit, geoDistance, entries.toArray(new GeoDistanceFacet.Entry[entries.size()]), + context.fieldDataCache(), context.mapperService(), valueScript, params, context.scriptService()); + } + + return new GeoDistanceFacetCollector(facetName, fieldName, lat, lon, unit, geoDistance, entries.toArray(new GeoDistanceFacet.Entry[entries.size()]), + context.fieldDataCache(), context.mapperService()); + } +} diff --git a/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/geodistance/InternalGeoDistanceFacet.java b/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/geodistance/InternalGeoDistanceFacet.java new file mode 100644 index 00000000000..3df9ba99c05 --- /dev/null +++ b/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/geodistance/InternalGeoDistanceFacet.java @@ -0,0 +1,185 @@ +/* + * Licensed to Elastic Search and Shay Banon under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. Elastic Search 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.search.facets.geodistance; + +import org.elasticsearch.common.collect.ImmutableList; +import org.elasticsearch.common.io.stream.StreamInput; +import org.elasticsearch.common.io.stream.StreamOutput; +import org.elasticsearch.common.unit.DistanceUnit; +import org.elasticsearch.common.xcontent.builder.XContentBuilder; +import org.elasticsearch.search.facets.Facet; +import org.elasticsearch.search.facets.internal.InternalFacet; + +import java.io.IOException; +import java.util.Iterator; +import java.util.List; + +/** + * @author kimchy (shay.banon) + */ +public class InternalGeoDistanceFacet implements GeoDistanceFacet, InternalFacet { + + private String name; + + private String fieldName; + + private String valueFieldName; + + private DistanceUnit unit; + + private Entry[] entries; + + InternalGeoDistanceFacet() { + } + + public InternalGeoDistanceFacet(String name, String fieldName, String valueFieldName, DistanceUnit unit, Entry[] entries) { + this.name = name; + this.fieldName = fieldName; + this.valueFieldName = valueFieldName; + this.unit = unit; + this.entries = entries; + } + + @Override public String name() { + return this.name; + } + + @Override public String getName() { + return name(); + } + + @Override public Type type() { + return Type.GEO_DISTANCE; + } + + @Override public Type getType() { + return type(); + } + + @Override public String fieldName() { + return this.fieldName; + } + + @Override public String getFieldName() { + return fieldName(); + } + + @Override public String valueFieldName() { + return this.valueFieldName; + } + + @Override public String getValueFieldName() { + return valueFieldName(); + } + + @Override public DistanceUnit unit() { + return this.unit; + } + + @Override public DistanceUnit getUnit() { + return unit(); + } + + @Override public List entries() { + return ImmutableList.copyOf(entries); + } + + @Override public List getEntries() { + return entries(); + } + + @Override public Iterator iterator() { + return entries().iterator(); + } + + @Override public Facet aggregate(Iterable facets) { + InternalGeoDistanceFacet agg = null; + for (Facet facet : facets) { + if (!facet.name().equals(name)) { + continue; + } + InternalGeoDistanceFacet geoDistanceFacet = (InternalGeoDistanceFacet) facet; + if (agg == null) { + agg = geoDistanceFacet; + } else { + for (int i = 0; i < geoDistanceFacet.entries.length; i++) { + agg.entries[i].count += geoDistanceFacet.entries[i].count; + agg.entries[i].total += geoDistanceFacet.entries[i].total; + } + } + } + return agg; + } + + public static InternalGeoDistanceFacet readGeoDistanceFacet(StreamInput in) throws IOException { + InternalGeoDistanceFacet facet = new InternalGeoDistanceFacet(); + facet.readFrom(in); + return facet; + } + + @Override public void readFrom(StreamInput in) throws IOException { + name = in.readUTF(); + fieldName = in.readUTF(); + valueFieldName = in.readUTF(); + unit = DistanceUnit.readDistanceUnit(in); + entries = new Entry[in.readVInt()]; + for (int i = 0; i < entries.length; i++) { + entries[i] = new Entry(in.readDouble(), in.readDouble(), in.readVLong(), in.readDouble()); + } + } + + @Override public void writeTo(StreamOutput out) throws IOException { + out.writeUTF(name); + out.writeUTF(fieldName); + out.writeUTF(valueFieldName); + DistanceUnit.writeDistanceUnit(out, unit); + out.writeVInt(entries.length); + for (Entry entry : entries) { + out.writeDouble(entry.from); + out.writeDouble(entry.to); + out.writeVLong(entry.count); + out.writeDouble(entry.total); + } + } + + @Override public void toXContent(XContentBuilder builder, Params params) throws IOException { + builder.startObject(name); + builder.field("_type", "histogram"); + builder.field("_field", fieldName); + builder.field("_value_field", valueFieldName); + builder.field("_unit", unit); + builder.startArray("ranges"); + for (Entry entry : entries) { + builder.startObject(); + if (!Double.isInfinite(entry.from)) { + builder.field("from", entry.from); + } + if (!Double.isInfinite(entry.to)) { + builder.field("to", entry.to); + } + builder.field("count", entry.count()); + builder.field("total", entry.total()); + builder.field("mean", entry.mean()); + builder.endObject(); + } + builder.endArray(); + builder.endObject(); + } +} diff --git a/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/geodistance/ScriptGeoDistanceFacetCollector.java b/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/geodistance/ScriptGeoDistanceFacetCollector.java new file mode 100644 index 00000000000..50398e3b3b2 --- /dev/null +++ b/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/geodistance/ScriptGeoDistanceFacetCollector.java @@ -0,0 +1,86 @@ +/* + * Licensed to Elastic Search and Shay Banon under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. Elastic Search 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.search.facets.geodistance; + +import org.apache.lucene.index.IndexReader; +import org.elasticsearch.common.lucene.geo.GeoDistance; +import org.elasticsearch.common.unit.DistanceUnit; +import org.elasticsearch.index.cache.field.data.FieldDataCache; +import org.elasticsearch.index.field.function.FieldsFunction; +import org.elasticsearch.index.field.function.script.ScriptFieldsFunction; +import org.elasticsearch.index.mapper.MapperService; +import org.elasticsearch.script.ScriptService; + +import java.io.IOException; +import java.util.Map; + +/** + * @author kimchy (shay.banon) + */ +public class ScriptGeoDistanceFacetCollector extends GeoDistanceFacetCollector { + + private final FieldsFunction valueFunction; + + private final Map params; + + public ScriptGeoDistanceFacetCollector(String facetName, String fieldName, double lat, double lon, DistanceUnit unit, GeoDistance geoDistance, + GeoDistanceFacet.Entry[] entries, FieldDataCache fieldDataCache, MapperService mapperService, + String script, Map params, ScriptService scriptService) { + super(facetName, fieldName, lat, lon, unit, geoDistance, entries, fieldDataCache, mapperService); + this.params = params; + + this.valueFunction = new ScriptFieldsFunction(script, scriptService, mapperService, fieldDataCache); + } + + @Override protected void doSetNextReader(IndexReader reader, int docBase) throws IOException { + super.doSetNextReader(reader, docBase); + valueFunction.setNextReader(reader); + } + + @Override protected void doCollect(int doc) throws IOException { + if (!latFieldData.hasValue(doc) || !lonFieldData.hasValue(doc)) { + return; + } + + double value = ((Number) valueFunction.execute(doc, params)).doubleValue(); + + if (latFieldData.multiValued()) { + double[] lats = latFieldData.doubleValues(doc); + double[] lons = latFieldData.doubleValues(doc); + for (int i = 0; i < lats.length; i++) { + double distance = geoDistance.calculate(lat, lon, lats[i], lons[i], unit); + for (GeoDistanceFacet.Entry entry : entries) { + if (distance >= entry.getFrom() && distance < entry.getTo()) { + entry.count++; + entry.total += value; + } + } + } + } else { + double distance = geoDistance.calculate(lat, lon, latFieldData.doubleValue(doc), lonFieldData.doubleValue(doc), unit); + for (GeoDistanceFacet.Entry entry : entries) { + if (distance >= entry.getFrom() && distance < entry.getTo()) { + entry.count++; + entry.total += value; + } + } + } + } +} diff --git a/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/geodistance/ValueGeoDistanceFacetCollector.java b/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/geodistance/ValueGeoDistanceFacetCollector.java new file mode 100644 index 00000000000..8ac62f80036 --- /dev/null +++ b/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/geodistance/ValueGeoDistanceFacetCollector.java @@ -0,0 +1,111 @@ +/* + * Licensed to Elastic Search and Shay Banon under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. Elastic Search 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.search.facets.geodistance; + +import org.apache.lucene.index.IndexReader; +import org.elasticsearch.common.lucene.geo.GeoDistance; +import org.elasticsearch.common.unit.DistanceUnit; +import org.elasticsearch.index.cache.field.data.FieldDataCache; +import org.elasticsearch.index.field.data.FieldData; +import org.elasticsearch.index.field.data.NumericFieldData; +import org.elasticsearch.index.mapper.FieldMapper; +import org.elasticsearch.index.mapper.MapperService; +import org.elasticsearch.search.facets.Facet; +import org.elasticsearch.search.facets.FacetPhaseExecutionException; + +import java.io.IOException; + +/** + * @author kimchy (shay.banon) + */ +public class ValueGeoDistanceFacetCollector extends GeoDistanceFacetCollector { + + private final String valueFieldName; + + private final String indexValueFieldName; + + private final FieldData.Type valueFieldDataType; + + private NumericFieldData valueFieldData; + + public ValueGeoDistanceFacetCollector(String facetName, String fieldName, double lat, double lon, DistanceUnit unit, GeoDistance geoDistance, + GeoDistanceFacet.Entry[] entries, FieldDataCache fieldDataCache, MapperService mapperService, String valueFieldName) { + super(facetName, fieldName, lat, lon, unit, geoDistance, entries, fieldDataCache, mapperService); + this.valueFieldName = valueFieldName; + + FieldMapper mapper = mapperService.smartNameFieldMapper(valueFieldName); + if (mapper == null) { + throw new FacetPhaseExecutionException(facetName, "No mapping found for field [" + valueFieldName + "]"); + } + this.indexValueFieldName = valueFieldName; + this.valueFieldDataType = mapper.fieldDataType(); + } + + @Override protected void doSetNextReader(IndexReader reader, int docBase) throws IOException { + super.doSetNextReader(reader, docBase); + valueFieldData = (NumericFieldData) fieldDataCache.cache(valueFieldDataType, reader, indexValueFieldName); + } + + @Override protected void doCollect(int doc) throws IOException { + if (!latFieldData.hasValue(doc) || !lonFieldData.hasValue(doc)) { + return; + } + + if (latFieldData.multiValued()) { + double[] lats = latFieldData.doubleValues(doc); + double[] lons = latFieldData.doubleValues(doc); + double[] values = valueFieldData.multiValued() ? valueFieldData.doubleValues(doc) : null; + for (int i = 0; i < lats.length; i++) { + double distance = geoDistance.calculate(lat, lon, lats[i], lons[i], unit); + for (GeoDistanceFacet.Entry entry : entries) { + if (distance >= entry.getFrom() && distance < entry.getTo()) { + entry.count++; + if (values != null) { + if (i < values.length) { + entry.total += values[i]; + } + } else if (valueFieldData.hasValue(doc)) { + entry.total += valueFieldData.doubleValue(doc); + } + } + } + } + } else { + double distance = geoDistance.calculate(lat, lon, latFieldData.doubleValue(doc), lonFieldData.doubleValue(doc), unit); + for (GeoDistanceFacet.Entry entry : entries) { + if (distance >= entry.getFrom() && distance < entry.getTo()) { + entry.count++; + if (valueFieldData.multiValued()) { + double[] values = valueFieldData.doubleValues(doc); + for (double value : values) { + entry.total += value; + } + } else if (valueFieldData.hasValue(doc)) { + entry.total += valueFieldData.doubleValue(doc); + } + } + } + } + } + + @Override public Facet facet() { + return new InternalGeoDistanceFacet(facetName, fieldName, valueFieldName, unit, entries); + } +} diff --git a/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/histogram/HistogramFacetCollectorParser.java b/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/histogram/HistogramFacetCollectorParser.java index 51995cc58fc..eb23c880da1 100644 --- a/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/histogram/HistogramFacetCollectorParser.java +++ b/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/histogram/HistogramFacetCollectorParser.java @@ -19,7 +19,6 @@ package org.elasticsearch.search.facets.histogram; -import org.elasticsearch.common.thread.ThreadLocals; import org.elasticsearch.common.unit.TimeValue; import org.elasticsearch.common.xcontent.XContentParser; import org.elasticsearch.search.facets.FacetPhaseExecutionException; @@ -28,7 +27,6 @@ import org.elasticsearch.search.facets.collector.FacetCollectorParser; import org.elasticsearch.search.internal.SearchContext; import java.io.IOException; -import java.util.HashMap; import java.util.Map; /** @@ -36,16 +34,10 @@ import java.util.Map; */ public class HistogramFacetCollectorParser implements FacetCollectorParser { - private static ThreadLocal>> cachedParams = new ThreadLocal>>() { - @Override protected ThreadLocals.CleanableValue> initialValue() { - return new ThreadLocals.CleanableValue>(new HashMap()); - } - }; - public static final String NAME = "histogram"; - @Override public String name() { - return NAME; + @Override public String[] names() { + return new String[]{NAME}; } @Override public FacetCollector parser(String facetName, XContentParser parser, SearchContext context) throws IOException { @@ -53,7 +45,7 @@ public class HistogramFacetCollectorParser implements FacetCollectorParser { String valueField = null; String keyScript = null; String valueScript = null; - Map params = cachedParams.get().get(); + Map params = null; long interval = 0; HistogramFacet.ComparatorType comparatorType = HistogramFacet.ComparatorType.KEY; XContentParser.Token token; diff --git a/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/internal/InternalFacets.java b/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/internal/InternalFacets.java index 9d147f03663..5a98373488f 100644 --- a/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/internal/InternalFacets.java +++ b/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/internal/InternalFacets.java @@ -29,6 +29,7 @@ import org.elasticsearch.common.xcontent.ToXContent; import org.elasticsearch.common.xcontent.builder.XContentBuilder; import org.elasticsearch.search.facets.Facet; import org.elasticsearch.search.facets.Facets; +import org.elasticsearch.search.facets.geodistance.InternalGeoDistanceFacet; import org.elasticsearch.search.facets.histogram.InternalHistogramFacet; import org.elasticsearch.search.facets.query.InternalQueryFacet; import org.elasticsearch.search.facets.statistical.InternalStatisticalFacet; @@ -142,6 +143,8 @@ public class InternalFacets implements Facets, Streamable, ToXContent, Iterable< facets.add(InternalStatisticalFacet.readStatisticalFacet(in)); } else if (id == Facet.Type.HISTOGRAM.id()) { facets.add(InternalHistogramFacet.readHistogramFacet(in)); + } else if (id == Facet.Type.GEO_DISTANCE.id()) { + facets.add(InternalGeoDistanceFacet.readGeoDistanceFacet(in)); } else { throw new IOException("Can't handle facet type with id [" + id + "]"); } diff --git a/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/query/QueryFacetCollectorParser.java b/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/query/QueryFacetCollectorParser.java index e7f25671acf..6bf8ea8b2a8 100644 --- a/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/query/QueryFacetCollectorParser.java +++ b/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/query/QueryFacetCollectorParser.java @@ -33,8 +33,8 @@ public class QueryFacetCollectorParser implements FacetCollectorParser { public static final String NAME = "query"; - @Override public String name() { - return "query"; + @Override public String[] names() { + return new String[]{"query"}; } @Override public FacetCollector parser(String facetName, XContentParser parser, SearchContext context) { diff --git a/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/statistical/StatisticalFacetCollectorParser.java b/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/statistical/StatisticalFacetCollectorParser.java index 52c90b5ff7f..5ef3ea52586 100644 --- a/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/statistical/StatisticalFacetCollectorParser.java +++ b/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/statistical/StatisticalFacetCollectorParser.java @@ -43,8 +43,8 @@ public class StatisticalFacetCollectorParser implements FacetCollectorParser { public static final String NAME = "statistical"; - @Override public String name() { - return NAME; + @Override public String[] names() { + return new String[]{NAME}; } @Override public FacetCollector parser(String facetName, XContentParser parser, SearchContext context) throws IOException { diff --git a/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/terms/TermsFacetCollectorParser.java b/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/terms/TermsFacetCollectorParser.java index b32bc2eec2d..308f9deb17b 100644 --- a/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/terms/TermsFacetCollectorParser.java +++ b/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/terms/TermsFacetCollectorParser.java @@ -36,8 +36,8 @@ public class TermsFacetCollectorParser implements FacetCollectorParser { public static final String NAME = "terms"; - @Override public String name() { - return NAME; + @Override public String[] names() { + return new String[]{NAME}; } @Override public FacetCollector parser(String facetName, XContentParser parser, SearchContext context) throws IOException { diff --git a/modules/test/integration/src/test/java/org/elasticsearch/test/integration/search/geo/GeoDistanceFacetTests.java b/modules/test/integration/src/test/java/org/elasticsearch/test/integration/search/geo/GeoDistanceFacetTests.java new file mode 100644 index 00000000000..ecedc5f8884 --- /dev/null +++ b/modules/test/integration/src/test/java/org/elasticsearch/test/integration/search/geo/GeoDistanceFacetTests.java @@ -0,0 +1,222 @@ +/* + * Licensed to Elastic Search and Shay Banon under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. Elastic Search 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.test.integration.search.geo; + +import org.elasticsearch.action.search.SearchResponse; +import org.elasticsearch.client.Client; +import org.elasticsearch.common.unit.DistanceUnit; +import org.elasticsearch.search.facets.geodistance.GeoDistanceFacet; +import org.elasticsearch.test.integration.AbstractNodesTests; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; + +import static org.elasticsearch.common.xcontent.XContentFactory.*; +import static org.elasticsearch.index.query.xcontent.QueryBuilders.*; +import static org.elasticsearch.search.facets.FacetBuilders.*; +import static org.hamcrest.MatcherAssert.*; +import static org.hamcrest.Matchers.*; + +/** + * @author kimchy (shay.banon) + */ +public class GeoDistanceFacetTests extends AbstractNodesTests { + + private Client client; + + @BeforeClass public void createNodes() throws Exception { + startNode("server1"); + startNode("server2"); + client = getClient(); + } + + @AfterClass public void closeNodes() { + client.close(); + closeAllNodes(); + } + + protected Client getClient() { + return client("server1"); + } + + @Test public void simpleGeoFacetTests() throws Exception { + try { + client.admin().indices().prepareDelete("test").execute().actionGet(); + } catch (Exception e) { + // ignore + } + client.admin().indices().prepareCreate("test").execute().actionGet(); + client.admin().cluster().prepareHealth().setWaitForGreenStatus().execute().actionGet(); + + // to NY: 0 + client.prepareIndex("test", "type1", "1").setSource(jsonBuilder().startObject() + .field("name", "New York") + .field("num", 1) + .startObject("location").field("lat", 40.7143528).field("lon", -74.0059731).endObject() + .endObject()).execute().actionGet(); + + // to NY: 5.286 km + client.prepareIndex("test", "type1", "2").setSource(jsonBuilder().startObject() + .field("name", "Times Square") + .field("num", 2) + .startObject("location").field("lat", 40.759011).field("lon", -73.9844722).endObject() + .endObject()).execute().actionGet(); + + // to NY: 0.4621 km + client.prepareIndex("test", "type1", "3").setSource(jsonBuilder().startObject() + .field("name", "Tribeca") + .field("num", 3) + .startObject("location").field("lat", 40.718266).field("lon", -74.007819).endObject() + .endObject()).execute().actionGet(); + + // to NY: 1.055 km + client.prepareIndex("test", "type1", "4").setSource(jsonBuilder().startObject() + .field("name", "Wall Street") + .field("num", 4) + .startObject("location").field("lat", 40.7051157).field("lon", -74.0088305).endObject() + .endObject()).execute().actionGet(); + + // to NY: 1.258 km + client.prepareIndex("test", "type1", "5").setSource(jsonBuilder().startObject() + .field("name", "Soho") + .field("num", 5) + .startObject("location").field("lat", 40.7247222).field("lon", -74).endObject() + .endObject()).execute().actionGet(); + + // to NY: 2.029 km + client.prepareIndex("test", "type1", "6").setSource(jsonBuilder().startObject() + .field("name", "Greenwich Village") + .field("num", 6) + .startObject("location").field("lat", 40.731033).field("lon", -73.9962255).endObject() + .endObject()).execute().actionGet(); + + // to NY: 8.572 km + client.prepareIndex("test", "type1", "7").setSource(jsonBuilder().startObject() + .field("name", "Brooklyn") + .field("num", 7) + .startObject("location").field("lat", 40.65).field("lon", -73.95).endObject() + .endObject()).execute().actionGet(); + + client.admin().indices().prepareRefresh().execute().actionGet(); + + SearchResponse searchResponse = client.prepareSearch() // from NY + .setQuery(matchAllQuery()) + .addFacet(geoDistanceFacet("geo1").field("location").point(40.7143528, -74.0059731).unit(DistanceUnit.KILOMETERS) + .addUnboundedFrom(2) + .addRange(0, 1) + .addRange(0.5, 2.5) + .addUnboundedTo(1) + ) + .execute().actionGet(); + + assertThat(searchResponse.hits().totalHits(), equalTo(7l)); + GeoDistanceFacet facet = searchResponse.facets().facet("geo1"); + assertThat(facet.fieldName(), equalTo("location")); + assertThat(facet.unit(), equalTo(DistanceUnit.KILOMETERS)); + assertThat(facet.entries().size(), equalTo(4)); + + assertThat(facet.entries().get(0).to(), closeTo(2, 0.000001)); + assertThat(facet.entries().get(0).count(), equalTo(4l)); + assertThat(facet.entries().get(0).total(), not(closeTo(0, 0.00001))); + + assertThat(facet.entries().get(1).from(), closeTo(0, 0.000001)); + assertThat(facet.entries().get(1).to(), closeTo(1, 0.000001)); + assertThat(facet.entries().get(1).count(), equalTo(2l)); + assertThat(facet.entries().get(1).total(), not(closeTo(0, 0.00001))); + + assertThat(facet.entries().get(2).from(), closeTo(0.5, 0.000001)); + assertThat(facet.entries().get(2).to(), closeTo(2.5, 0.000001)); + assertThat(facet.entries().get(2).count(), equalTo(3l)); + assertThat(facet.entries().get(2).total(), not(closeTo(0, 0.00001))); + + assertThat(facet.entries().get(3).from(), closeTo(1, 0.000001)); + assertThat(facet.entries().get(3).count(), equalTo(5l)); + assertThat(facet.entries().get(3).total(), not(closeTo(0, 0.00001))); + + + searchResponse = client.prepareSearch() // from NY + .setQuery(matchAllQuery()) + .addFacet(geoDistanceFacet("geo1").field("location").point(40.7143528, -74.0059731).unit(DistanceUnit.KILOMETERS).valueField("num") + .addUnboundedFrom(2) + .addRange(0, 1) + .addRange(0.5, 2.5) + .addUnboundedTo(1) + ) + .execute().actionGet(); + + assertThat(searchResponse.hits().totalHits(), equalTo(7l)); + facet = searchResponse.facets().facet("geo1"); + assertThat(facet.fieldName(), equalTo("location")); + assertThat(facet.unit(), equalTo(DistanceUnit.KILOMETERS)); + assertThat(facet.entries().size(), equalTo(4)); + + assertThat(facet.entries().get(0).to(), closeTo(2, 0.000001)); + assertThat(facet.entries().get(0).count(), equalTo(4l)); + assertThat(facet.entries().get(0).total(), closeTo(13, 0.00001)); + + assertThat(facet.entries().get(1).from(), closeTo(0, 0.000001)); + assertThat(facet.entries().get(1).to(), closeTo(1, 0.000001)); + assertThat(facet.entries().get(1).count(), equalTo(2l)); + assertThat(facet.entries().get(1).total(), closeTo(4, 0.00001)); + + assertThat(facet.entries().get(2).from(), closeTo(0.5, 0.000001)); + assertThat(facet.entries().get(2).to(), closeTo(2.5, 0.000001)); + assertThat(facet.entries().get(2).count(), equalTo(3l)); + assertThat(facet.entries().get(2).total(), closeTo(15, 0.00001)); + + assertThat(facet.entries().get(3).from(), closeTo(1, 0.000001)); + assertThat(facet.entries().get(3).count(), equalTo(5l)); + assertThat(facet.entries().get(3).total(), closeTo(24, 0.00001)); + + searchResponse = client.prepareSearch() // from NY + .setQuery(matchAllQuery()) + .addFacet(geoDistanceFacet("geo1").field("location").point(40.7143528, -74.0059731).unit(DistanceUnit.KILOMETERS).valueScript("doc['num'].value") + .addUnboundedFrom(2) + .addRange(0, 1) + .addRange(0.5, 2.5) + .addUnboundedTo(1) + ) + .execute().actionGet(); + + assertThat(searchResponse.hits().totalHits(), equalTo(7l)); + facet = searchResponse.facets().facet("geo1"); + assertThat(facet.fieldName(), equalTo("location")); + assertThat(facet.unit(), equalTo(DistanceUnit.KILOMETERS)); + assertThat(facet.entries().size(), equalTo(4)); + + assertThat(facet.entries().get(0).to(), closeTo(2, 0.000001)); + assertThat(facet.entries().get(0).count(), equalTo(4l)); + assertThat(facet.entries().get(0).total(), closeTo(13, 0.00001)); + + assertThat(facet.entries().get(1).from(), closeTo(0, 0.000001)); + assertThat(facet.entries().get(1).to(), closeTo(1, 0.000001)); + assertThat(facet.entries().get(1).count(), equalTo(2l)); + assertThat(facet.entries().get(1).total(), closeTo(4, 0.00001)); + + assertThat(facet.entries().get(2).from(), closeTo(0.5, 0.000001)); + assertThat(facet.entries().get(2).to(), closeTo(2.5, 0.000001)); + assertThat(facet.entries().get(2).count(), equalTo(3l)); + assertThat(facet.entries().get(2).total(), closeTo(15, 0.00001)); + + assertThat(facet.entries().get(3).from(), closeTo(1, 0.000001)); + assertThat(facet.entries().get(3).count(), equalTo(5l)); + assertThat(facet.entries().get(3).total(), closeTo(24, 0.00001)); + } +}