From 8249e44463e2ac26d9dbd669a9a9f52955d85a7f Mon Sep 17 00:00:00 2001 From: Marcin Gornik Date: Tue, 22 Oct 2013 18:22:25 +0200 Subject: [PATCH] Add support for additional grouping Added parameter "additional_grouping" for geohashcell facet. When this parameter is provided, the region counts are additionaly broken down according to terms in this field. --- .../geohashcellfacet/GeoHashCellEntry.java | 21 +++--- .../geohashcellfacet/GeoHashCellFacet.java | 34 +++++----- .../GeoHashCellFacetEntry.java | 6 +- .../GeoHashCellFacetExecutor.java | 54 +++++++++++---- .../GeoHashCellFacetResult.java | 67 +++++++++++++++++++ .../GeoHashCellWithGroupingEntry.java | 12 ++++ .../GeoHashCellWithGroupings.java | 65 ++++++++++++++++++ .../MissingGeoValueEntry.java | 10 ++- .../MissingGeoValueWithGroupings.java | 49 ++++++++++++++ .../geohashcellfacet/ResultWithGroupings.java | 29 ++++++++ 10 files changed, 301 insertions(+), 46 deletions(-) create mode 100644 src/main/java/org/elasticsearch/plugin/geohashcellfacet/GeoHashCellFacetResult.java create mode 100644 src/main/java/org/elasticsearch/plugin/geohashcellfacet/GeoHashCellWithGroupings.java create mode 100644 src/main/java/org/elasticsearch/plugin/geohashcellfacet/MissingGeoValueWithGroupings.java create mode 100644 src/main/java/org/elasticsearch/plugin/geohashcellfacet/ResultWithGroupings.java diff --git a/src/main/java/org/elasticsearch/plugin/geohashcellfacet/GeoHashCellEntry.java b/src/main/java/org/elasticsearch/plugin/geohashcellfacet/GeoHashCellEntry.java index de3bec2..fd7c622 100644 --- a/src/main/java/org/elasticsearch/plugin/geohashcellfacet/GeoHashCellEntry.java +++ b/src/main/java/org/elasticsearch/plugin/geohashcellfacet/GeoHashCellEntry.java @@ -8,7 +8,7 @@ public class GeoHashCellEntry implements GeoHashCellFacetEntry { - private final GeoHashCell cell; + protected final GeoHashCell cell; public GeoHashCellEntry(GeoHashCell cell) { this.cell = cell; @@ -29,6 +29,10 @@ public boolean equals(Object o) { return cell.equals(that.cell); } + public static GeoHashCellEntry createEntry(GeoHashCell cell) { + return createEntry(cell, null); + } + public static GeoHashCellEntry createEntry(GeoHashCell cell, String additionalGroupingValue) { return additionalGroupingValue == null || additionalGroupingValue.isEmpty() ? new GeoHashCellEntry(cell) @@ -36,13 +40,12 @@ public static GeoHashCellEntry createEntry(GeoHashCell cell, String additionalGr } @Override - public XContentBuilder toXContent(XContentBuilder builder, AtomicLong value) throws IOException { - GeoPoint center = cell.getCenter(); - builder.startObject(); - builder.field("lat", center.lat()); - builder.field("lon", center.lon()); - builder.field("count", value); - builder.endObject(); - return builder; + public GeoHashCellWithGroupings createCellWithGroupings(AtomicLong value) { + return new GeoHashCellWithGroupings(cell, value.get()); + } + + @Override + public String getKey() { + return cell.toString(); } } \ No newline at end of file diff --git a/src/main/java/org/elasticsearch/plugin/geohashcellfacet/GeoHashCellFacet.java b/src/main/java/org/elasticsearch/plugin/geohashcellfacet/GeoHashCellFacet.java index ed9a4a8..30295ad 100644 --- a/src/main/java/org/elasticsearch/plugin/geohashcellfacet/GeoHashCellFacet.java +++ b/src/main/java/org/elasticsearch/plugin/geohashcellfacet/GeoHashCellFacet.java @@ -33,33 +33,36 @@ public class GeoHashCellFacet extends InternalFacet { private Map counts = Maps.newHashMap(); private final String fieldName; private final int userLevel; - private final String additionalGrouping; + private final String additionalGroupingField; + + private GeoHashCellFacetResult results; /** * Creates the facet instance. * @param facetName Name of the facet. - * @param counts Map of the group counts. + * @param results * @param fieldName Name of the field used for grouping. * @param mapBox Current map viewport. * @param userLevel Level for geo hash grouping. - * @param additionalGrouping + * @param additionalGroupingField */ - public GeoHashCellFacet(String facetName, Map counts, + public GeoHashCellFacet(String facetName, GeoHashCellFacetResult results, String fieldName, MapBox mapBox, - int userLevel, String additionalGrouping) { + int userLevel, String additionalGroupingField) { super(facetName); - this.counts = counts; + this.results = results; this.fieldName = fieldName; this.mapBox = mapBox; this.userLevel = userLevel; - this.additionalGrouping = additionalGrouping; + this.additionalGroupingField = additionalGroupingField; } private GeoHashCellFacet() { this.fieldName = null; this.mapBox = null; this.userLevel = 0; - this.additionalGrouping = null; + this.results = new GeoHashCellFacetResult(); + this.additionalGroupingField = null; } @@ -92,14 +95,7 @@ public Facet reduce(ReduceContext reduceContext) { if (facet instanceof GeoHashCellFacet) { GeoHashCellFacet hashCellFacet = (GeoHashCellFacet) facet; - for (Map.Entry entry : hashCellFacet.counts.entrySet()) { - if (geoHashCellFacet.counts.containsKey(entry.getKey())) { - geoHashCellFacet.counts.get(entry.getKey()).addAndGet(entry.getValue().longValue()); - } - else { - geoHashCellFacet.counts.put(entry.getKey(), entry.getValue()); - } - } + geoHashCellFacet.results = geoHashCellFacet.results.reduce(hashCellFacet.results); } } @@ -121,9 +117,9 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws builder.field(GeoHashCellFacetParser.ParamName.LEVEL, mapBox.getLevel(userLevel)); builder.field("base_map_level", mapBox.getBaseLevel()); builder.startArray("regions"); - for (Map.Entry entry: counts.entrySet()) { - entry.getKey().toXContent(builder, entry.getValue()); - } + + results.toXContent(builder); + builder.endArray(); builder.endObject(); return builder; diff --git a/src/main/java/org/elasticsearch/plugin/geohashcellfacet/GeoHashCellFacetEntry.java b/src/main/java/org/elasticsearch/plugin/geohashcellfacet/GeoHashCellFacetEntry.java index 77af19b..ab68aba 100644 --- a/src/main/java/org/elasticsearch/plugin/geohashcellfacet/GeoHashCellFacetEntry.java +++ b/src/main/java/org/elasticsearch/plugin/geohashcellfacet/GeoHashCellFacetEntry.java @@ -1,10 +1,14 @@ package org.elasticsearch.plugin.geohashcellfacet; +import com.sun.tools.corba.se.idl.InterfaceState; import org.elasticsearch.common.xcontent.XContentBuilder; import java.io.IOException; +import java.util.Set; import java.util.concurrent.atomic.AtomicLong; public interface GeoHashCellFacetEntry { - XContentBuilder toXContent(XContentBuilder builder, AtomicLong value) throws IOException; + ResultWithGroupings createCellWithGroupings(AtomicLong value); + + String getKey(); } diff --git a/src/main/java/org/elasticsearch/plugin/geohashcellfacet/GeoHashCellFacetExecutor.java b/src/main/java/org/elasticsearch/plugin/geohashcellfacet/GeoHashCellFacetExecutor.java index 5150036..3dfb064 100644 --- a/src/main/java/org/elasticsearch/plugin/geohashcellfacet/GeoHashCellFacetExecutor.java +++ b/src/main/java/org/elasticsearch/plugin/geohashcellfacet/GeoHashCellFacetExecutor.java @@ -1,8 +1,16 @@ package org.elasticsearch.plugin.geohashcellfacet; +import org.apache.lucene.index.AtomicReader; import org.apache.lucene.index.AtomicReaderContext; +import org.apache.lucene.index.IndexableField; +import org.apache.lucene.index.Terms; +import org.apache.lucene.util.BytesRef; +import org.elasticsearch.common.bytes.BytesArray; import org.elasticsearch.common.collect.Maps; import org.elasticsearch.common.geo.GeoPoint; +import org.elasticsearch.common.text.BytesText; +import org.elasticsearch.common.text.Text; +import org.elasticsearch.index.fielddata.BytesValues; import org.elasticsearch.index.fielddata.GeoPointValues; import org.elasticsearch.index.fielddata.IndexFieldData; import org.elasticsearch.index.fielddata.IndexGeoPointFieldData; @@ -11,35 +19,36 @@ import org.elasticsearch.search.internal.SearchContext; import java.io.IOException; +import java.util.List; import java.util.Map; import java.util.concurrent.atomic.AtomicLong; public class GeoHashCellFacetExecutor extends FacetExecutor { private final String fieldName; private final int userLevel; - private final String additionalGrouping; + private final String additionalGroupingField; + private final IndexFieldData additionalGroupingFieldData; private Map counts = Maps.newHashMap(); private final IndexGeoPointFieldData indexFieldData; - private final IndexFieldData additionalGroupingFieldData; private final MapBox mapBox; public GeoHashCellFacetExecutor(String fieldName, SearchContext searchContext, MapBox mapBox, - int userLevel, String additionalGrouping) { + int userLevel, String additionalGroupingField) { this.fieldName = fieldName; this.mapBox = mapBox; this.userLevel = userLevel; - this.additionalGrouping = additionalGrouping; + this.additionalGroupingField = additionalGroupingField; this.indexFieldData = searchContext .fieldData() .getForField(searchContext.smartNameFieldMapper(fieldName)); - if (additionalGrouping != null && !additionalGrouping.isEmpty()) { + if (additionalGroupingField != null && !additionalGroupingField.isEmpty()) { this.additionalGroupingFieldData = searchContext .fieldData() - .getForField(searchContext.smartNameFieldMapper(additionalGrouping)); + .getForField(searchContext.smartNameFieldMapper(additionalGroupingField)); } else { this.additionalGroupingFieldData = null; } @@ -47,9 +56,10 @@ public GeoHashCellFacetExecutor(String fieldName, @Override public InternalFacet buildFacet(String facetName) { + GeoHashCellFacetResult results = GeoHashCellFacetResult.createFromEntryCounts(counts); return new GeoHashCellFacet( - facetName, counts, fieldName, - mapBox, userLevel, additionalGrouping); + facetName, results, fieldName, + mapBox, userLevel, additionalGroupingField); } @Override @@ -60,15 +70,19 @@ public Collector collector() { final class Collector extends FacetExecutor.Collector { protected GeoPointValues values; - protected String additionalGroupingValue; + private BytesValues additionalValues; @Override public void setNextReader(AtomicReaderContext context) throws IOException { values = indexFieldData.load(context).getGeoPointValues(); - if (additionalGroupingFieldData != null) - additionalGroupingValue = additionalGroupingFieldData.load(context).toString(); + if (additionalGroupingFieldData != null) { + additionalValues = additionalGroupingFieldData.load(context).getBytesValues(); + + } else { + additionalValues = null; + } } @Override @@ -87,14 +101,26 @@ public void collect(int docId) throws IOException { continue; GeoHashCell cell = new GeoHashCell(point, mapBox.getLevel(userLevel)); - GeoHashCellEntry entry = GeoHashCellEntry.createEntry(cell, additionalGroupingValue); - increment(entry, 1L); + + if (additionalValues != null && additionalValues.hasValue(docId)) { + final BytesValues.Iter iter = additionalValues.getIter(docId); + while (iter.hasNext()) { + + BytesRef bytesRef = iter.next(); + Text text = new BytesText(new BytesArray(bytesRef)); + + increment(GeoHashCellEntry.createEntry(cell, text.string()), 1L); + } + } + else { + increment(GeoHashCellEntry.createEntry(cell), 1L); + } } } @Override public void postCollection() { - // do nothing + } private void increment(GeoHashCellFacetEntry entry, Long value) { diff --git a/src/main/java/org/elasticsearch/plugin/geohashcellfacet/GeoHashCellFacetResult.java b/src/main/java/org/elasticsearch/plugin/geohashcellfacet/GeoHashCellFacetResult.java new file mode 100644 index 0000000..c9eddd8 --- /dev/null +++ b/src/main/java/org/elasticsearch/plugin/geohashcellfacet/GeoHashCellFacetResult.java @@ -0,0 +1,67 @@ +package org.elasticsearch.plugin.geohashcellfacet; + +import org.elasticsearch.common.collect.Maps; +import org.elasticsearch.common.xcontent.XContentBuilder; + +import java.io.IOException; +import java.util.Map; +import java.util.concurrent.atomic.AtomicLong; + +public class GeoHashCellFacetResult { + + private final Map counts; + + public GeoHashCellFacetResult() { + this.counts = Maps.newHashMap(); + } + + public GeoHashCellFacetResult(GeoHashCellFacetEntry entry, AtomicLong value) { + Map counts = Maps.newHashMap(); + counts.put(entry.getKey(), entry.createCellWithGroupings(value)); + + this.counts = counts; + } + + private GeoHashCellFacetResult(Map counts) { + this.counts = Maps.newHashMap(counts); + } + + public GeoHashCellFacetResult reduce(GeoHashCellFacetResult other) { + + Map countsCopy = Maps.newHashMap(this.counts); + + for (Map.Entry pair : other.counts.entrySet()) { + String key = pair.getKey(); + ResultWithGroupings value = pair.getValue(); + + if (countsCopy.containsKey(key)) { + ResultWithGroupings oldValue = countsCopy.get(key); + countsCopy.remove(key); + countsCopy.put(key, oldValue.merge(value)); + } + else { + countsCopy.put(key, value); + } + } + + return new GeoHashCellFacetResult(countsCopy); + } + + public static GeoHashCellFacetResult createFromEntryCounts(Map counts) { + GeoHashCellFacetResult result = new GeoHashCellFacetResult(); + + for (Map.Entry entry : counts.entrySet()) { + result = result.reduce(new GeoHashCellFacetResult(entry.getKey(), entry.getValue())); + } + + return result; + } + + public void toXContent(XContentBuilder builder) throws IOException { + for (Map.Entry entry : counts.entrySet()) { + entry.getValue().toXContent(builder); + } + } +} + + diff --git a/src/main/java/org/elasticsearch/plugin/geohashcellfacet/GeoHashCellWithGroupingEntry.java b/src/main/java/org/elasticsearch/plugin/geohashcellfacet/GeoHashCellWithGroupingEntry.java index d128ec7..0af8bcb 100644 --- a/src/main/java/org/elasticsearch/plugin/geohashcellfacet/GeoHashCellWithGroupingEntry.java +++ b/src/main/java/org/elasticsearch/plugin/geohashcellfacet/GeoHashCellWithGroupingEntry.java @@ -1,5 +1,10 @@ package org.elasticsearch.plugin.geohashcellfacet; +import org.elasticsearch.common.collect.Maps; + +import java.util.Map; +import java.util.concurrent.atomic.AtomicLong; + public class GeoHashCellWithGroupingEntry extends GeoHashCellEntry { private final String additionalGroupingValue; @@ -20,6 +25,13 @@ public boolean equals(Object o) { return additionalGroupingValue.equals(that.additionalGroupingValue); } + @Override + public GeoHashCellWithGroupings createCellWithGroupings(AtomicLong value) { + Map counts = Maps.newHashMap(); + counts.put(additionalGroupingValue, value); + return new GeoHashCellWithGroupings(cell, value.get(), counts); + } + @Override public int hashCode() { int result = super.hashCode(); diff --git a/src/main/java/org/elasticsearch/plugin/geohashcellfacet/GeoHashCellWithGroupings.java b/src/main/java/org/elasticsearch/plugin/geohashcellfacet/GeoHashCellWithGroupings.java new file mode 100644 index 0000000..e380d68 --- /dev/null +++ b/src/main/java/org/elasticsearch/plugin/geohashcellfacet/GeoHashCellWithGroupings.java @@ -0,0 +1,65 @@ +package org.elasticsearch.plugin.geohashcellfacet; + +import org.elasticsearch.common.collect.Maps; +import org.elasticsearch.common.geo.GeoPoint; +import org.elasticsearch.common.xcontent.XContentBuilder; + +import java.io.IOException; +import java.util.Map; +import java.util.concurrent.atomic.AtomicLong; + +public class GeoHashCellWithGroupings extends ResultWithGroupings { + + private final GeoHashCell cell; + private final long total; + private final Map groupings; + + public GeoHashCellWithGroupings(GeoHashCell cell, long total) { + + this.cell = cell; + this.total = total; + this.groupings = Maps.newHashMap(); + } + + public GeoHashCellWithGroupings(GeoHashCell cell, long total, Map groupings) { + this.cell = cell; + this.total = total; + this.groupings = Maps.newHashMap(groupings); + } + + @Override + public ResultWithGroupings merge(ResultWithGroupings other) { + if (!(other instanceof GeoHashCellWithGroupings)) + throw new IllegalArgumentException("other"); + + GeoHashCellWithGroupings cellWithGroupings = (GeoHashCellWithGroupings) other; + + return new GeoHashCellWithGroupings(cell, + total + cellWithGroupings.total, + mergeGroupings(this.groupings, cellWithGroupings.groupings)); + } + + @Override + public void toXContent(XContentBuilder builder) throws IOException { + GeoPoint center = cell.getCenter(); + builder.startObject(); + builder.field("lat", center.lat()); + builder.field("lon", center.lon()); + + if (groupings.size() > 0) { + builder.startObject("count"); + + for (Map.Entry entry : groupings.entrySet()) { + builder.field(entry.getKey(), entry.getValue().get()); + } + + builder.endObject(); + } else { + builder.field("count", total); + } + + builder.endObject(); + } +} + + diff --git a/src/main/java/org/elasticsearch/plugin/geohashcellfacet/MissingGeoValueEntry.java b/src/main/java/org/elasticsearch/plugin/geohashcellfacet/MissingGeoValueEntry.java index 3404855..6260b2a 100644 --- a/src/main/java/org/elasticsearch/plugin/geohashcellfacet/MissingGeoValueEntry.java +++ b/src/main/java/org/elasticsearch/plugin/geohashcellfacet/MissingGeoValueEntry.java @@ -21,8 +21,12 @@ public int hashCode() { } @Override - public XContentBuilder toXContent(XContentBuilder builder, AtomicLong value) throws IOException { - builder.field(MISSING_ENTRY, value); - return builder; + public ResultWithGroupings createCellWithGroupings(AtomicLong value) { + return new MissingGeoValueWithGroupings(value); + } + + @Override + public String getKey() { + return MISSING_ENTRY; } } diff --git a/src/main/java/org/elasticsearch/plugin/geohashcellfacet/MissingGeoValueWithGroupings.java b/src/main/java/org/elasticsearch/plugin/geohashcellfacet/MissingGeoValueWithGroupings.java new file mode 100644 index 0000000..4da6c04 --- /dev/null +++ b/src/main/java/org/elasticsearch/plugin/geohashcellfacet/MissingGeoValueWithGroupings.java @@ -0,0 +1,49 @@ +package org.elasticsearch.plugin.geohashcellfacet; + +import org.elasticsearch.common.collect.Maps; +import org.elasticsearch.common.xcontent.XContentBuilder; + +import java.io.IOException; +import java.util.Map; +import java.util.concurrent.atomic.AtomicLong; + +public class MissingGeoValueWithGroupings extends ResultWithGroupings { + + private final AtomicLong total; + private final Map groupings; + + public MissingGeoValueWithGroupings(AtomicLong total) { + this.total = total; + this.groupings = Maps.newHashMap(); + } + + public MissingGeoValueWithGroupings(AtomicLong total, Map groupings) { + this.total = total; + this.groupings = groupings; + } + + @Override + public ResultWithGroupings merge(ResultWithGroupings other) { + if (!(other instanceof MissingGeoValueWithGroupings)) + throw new IllegalArgumentException("other"); + + MissingGeoValueWithGroupings missingWithGroupings = (MissingGeoValueWithGroupings) other; + + return new MissingGeoValueWithGroupings( + new AtomicLong(total.get() + missingWithGroupings.total.get()), + mergeGroupings(this.groupings, missingWithGroupings.groupings)); + } + + @Override + public void toXContent(XContentBuilder builder) throws IOException { + builder.startObject(); + if (groupings.size() > 0) { + for (Map.Entry entry : groupings.entrySet()) { + builder.field(entry.getKey(), entry.getValue().get()); + } + } else { + builder.field("_missing", total.get()); + } + builder.endObject(); + } +} diff --git a/src/main/java/org/elasticsearch/plugin/geohashcellfacet/ResultWithGroupings.java b/src/main/java/org/elasticsearch/plugin/geohashcellfacet/ResultWithGroupings.java new file mode 100644 index 0000000..b845e52 --- /dev/null +++ b/src/main/java/org/elasticsearch/plugin/geohashcellfacet/ResultWithGroupings.java @@ -0,0 +1,29 @@ +package org.elasticsearch.plugin.geohashcellfacet; + +import org.elasticsearch.common.collect.Maps; +import org.elasticsearch.common.xcontent.XContentBuilder; + +import java.io.IOException; +import java.util.Map; +import java.util.concurrent.atomic.AtomicLong; + +public abstract class ResultWithGroupings { + + public abstract ResultWithGroupings merge(ResultWithGroupings other); + + protected Map mergeGroupings( + Map firstGroupings, Map secondGroupings) { + Map result = Maps.newHashMap(firstGroupings); + + for (Map.Entry entry : secondGroupings.entrySet()) { + if (result.containsKey(entry.getKey())) + result.get(entry.getKey()).addAndGet(entry.getValue().get()); + else + result.put(entry.getKey(), entry.getValue()); + } + + return result; + } + + public abstract void toXContent(XContentBuilder builder) throws IOException; +}