diff --git a/modules/elasticsearch/src/main/java/org/elasticsearch/common/xcontent/XContentBuilder.java b/modules/elasticsearch/src/main/java/org/elasticsearch/common/xcontent/XContentBuilder.java index df6d4016456..248f2f4fcfa 100644 --- a/modules/elasticsearch/src/main/java/org/elasticsearch/common/xcontent/XContentBuilder.java +++ b/modules/elasticsearch/src/main/java/org/elasticsearch/common/xcontent/XContentBuilder.java @@ -239,6 +239,15 @@ public final class XContentBuilder { return this; } + public XContentBuilder field(String name, String... value) throws IOException { + startArray(name); + for (String o : value) { + value(o); + } + endArray(); + return this; + } + public XContentBuilder field(String name, Object... value) throws IOException { startArray(name); for (Object o : value) { diff --git a/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/terms/TermsFacetBuilder.java b/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/terms/TermsFacetBuilder.java index 223b31f7fae..5934d348340 100644 --- a/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/terms/TermsFacetBuilder.java +++ b/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/terms/TermsFacetBuilder.java @@ -34,6 +34,7 @@ import java.util.Map; */ public class TermsFacetBuilder extends AbstractFacetBuilder { private String fieldName; + private String[] fieldsNames; private int size = 10; private String[] exclude; private String regex; @@ -62,6 +63,11 @@ public class TermsFacetBuilder extends AbstractFacetBuilder { return this; } + public TermsFacetBuilder fields(String... fields) { + this.fieldsNames = fields; + return this; + } + public TermsFacetBuilder exclude(String... exclude) { this.exclude = exclude; return this; @@ -109,13 +115,21 @@ public class TermsFacetBuilder extends AbstractFacetBuilder { } @Override public void toXContent(XContentBuilder builder, Params params) throws IOException { - if (fieldName == null) { + if (fieldName == null && fieldsNames == null) { throw new SearchSourceBuilderException("field must be set on terms facet for facet [" + name + "]"); } builder.startObject(name); builder.startObject(TermsFacetCollectorParser.NAME); - builder.field("field", fieldName); + if (fieldsNames != null) { + if (fieldsNames.length == 1) { + builder.field("field", fieldsNames[0]); + } else { + builder.field("fields", fieldsNames); + } + } else { + builder.field("field", fieldName); + } builder.field("size", size); if (exclude != null) { builder.startArray("exclude"); diff --git a/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/terms/TermsFacetCollector.java b/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/terms/TermsFacetCollector.java index 37adf8dc049..a23c4e01267 100644 --- a/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/terms/TermsFacetCollector.java +++ b/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/terms/TermsFacetCollector.java @@ -47,7 +47,7 @@ import java.util.regex.Pattern; */ public class TermsFacetCollector extends AbstractFacetCollector { - private static ThreadLocal>>> cache = new ThreadLocal>>>() { + static ThreadLocal>>> cache = new ThreadLocal>>>() { @Override protected ThreadLocals.CleanableValue>> initialValue() { return new ThreadLocals.CleanableValue>>(new ArrayDeque>()); } @@ -139,7 +139,7 @@ public class TermsFacetCollector extends AbstractFacetCollector { } } - private TObjectIntHashMap popFacets() { + static TObjectIntHashMap popFacets() { Deque> deque = cache.get().get(); if (deque.isEmpty()) { deque.add(new TObjectIntHashMap()); @@ -149,7 +149,7 @@ public class TermsFacetCollector extends AbstractFacetCollector { return facets; } - private void pushFacets(TObjectIntHashMap facets) { + static void pushFacets(TObjectIntHashMap facets) { facets.clear(); Deque> deque = cache.get().get(); if (deque != null) { 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 8cfed555fe1..4126cfaf80e 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 @@ -20,6 +20,7 @@ package org.elasticsearch.search.facets.terms; import org.elasticsearch.common.collect.ImmutableSet; +import org.elasticsearch.common.collect.Lists; import org.elasticsearch.common.regex.Regex; import org.elasticsearch.common.xcontent.XContentParser; import org.elasticsearch.search.facets.collector.FacetCollector; @@ -27,6 +28,7 @@ import org.elasticsearch.search.facets.collector.FacetCollectorParser; import org.elasticsearch.search.internal.SearchContext; import java.io.IOException; +import java.util.List; import java.util.Map; import java.util.regex.Pattern; @@ -46,6 +48,7 @@ public class TermsFacetCollectorParser implements FacetCollectorParser { int size = 10; String fieldName = null; + String[] fieldsNames = null; XContentParser.Token token; ImmutableSet excluded = ImmutableSet.of(); String regex = null; @@ -68,6 +71,12 @@ public class TermsFacetCollectorParser implements FacetCollectorParser { builder.add(parser.text()); } excluded = builder.build(); + } else if ("fields".equals(fieldName)) { + List fields = Lists.newArrayListWithCapacity(4); + while ((token = parser.nextToken()) != XContentParser.Token.END_ARRAY) { + fields.add(parser.text()); + } + fieldsNames = fields.toArray(new String[fields.size()]); } } else if (token.isValue()) { if ("field".equals(fieldName)) { @@ -96,6 +105,9 @@ public class TermsFacetCollectorParser implements FacetCollectorParser { if (regex != null) { pattern = Regex.compile(regex, regexFlags); } + if (fieldsNames != null) { + return new TermsFieldsFacetCollector(facetName, fieldsNames, size, comparatorType, context, excluded, pattern, scriptLang, script, params); + } return new TermsFacetCollector(facetName, field, size, comparatorType, context, excluded, pattern, scriptLang, script, params); } } diff --git a/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/terms/TermsFieldsFacetCollector.java b/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/terms/TermsFieldsFacetCollector.java new file mode 100644 index 00000000000..65fe9ebe266 --- /dev/null +++ b/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/terms/TermsFieldsFacetCollector.java @@ -0,0 +1,203 @@ +/* + * 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.terms; + +import org.apache.lucene.index.IndexReader; +import org.elasticsearch.common.collect.BoundedTreeSet; +import org.elasticsearch.common.collect.ImmutableList; +import org.elasticsearch.common.collect.ImmutableSet; +import org.elasticsearch.common.collect.Maps; +import org.elasticsearch.common.trove.TObjectIntHashMap; +import org.elasticsearch.common.trove.TObjectIntIterator; +import org.elasticsearch.index.cache.field.data.FieldDataCache; +import org.elasticsearch.index.field.data.FieldData; +import org.elasticsearch.index.mapper.MapperService; +import org.elasticsearch.script.search.SearchScript; +import org.elasticsearch.search.facets.Facet; +import org.elasticsearch.search.facets.support.AbstractFacetCollector; +import org.elasticsearch.search.internal.SearchContext; + +import java.io.IOException; +import java.util.Map; +import java.util.regex.Matcher; +import java.util.regex.Pattern; + +import static org.elasticsearch.common.Strings.*; + +/** + * @author kimchy (shay.banon) + */ +public class TermsFieldsFacetCollector extends AbstractFacetCollector { + + private final FieldDataCache fieldDataCache; + + private final String[] fieldsNames; + + private final String[] indexFieldsNames; + + private final InternalTermsFacet.ComparatorType comparatorType; + + private final int size; + + private final int numberOfShards; + + private final FieldData.Type[] fieldsDataType; + + private FieldData[] fieldsData; + + private final StaticAggregatorValueProc aggregator; + + private final SearchScript script; + + public TermsFieldsFacetCollector(String facetName, String[] fieldsNames, int size, InternalTermsFacet.ComparatorType comparatorType, SearchContext context, + ImmutableSet excluded, Pattern pattern, String scriptLang, String script, Map params) { + super(facetName); + this.fieldDataCache = context.fieldDataCache(); + this.size = size; + this.comparatorType = comparatorType; + this.numberOfShards = context.numberOfShards(); + + this.fieldsNames = fieldsNames; + + fieldsDataType = new FieldData.Type[fieldsNames.length]; + fieldsData = new FieldData[fieldsNames.length]; + indexFieldsNames = new String[fieldsNames.length]; + + for (int i = 0; i < fieldsNames.length; i++) { + MapperService.SmartNameFieldMappers smartMappers = context.mapperService().smartName(fieldsNames[i]); + if (smartMappers == null || !smartMappers.hasMapper()) { + this.indexFieldsNames[i] = fieldsNames[i]; + this.fieldsDataType[i] = FieldData.Type.STRING; + } else { + this.indexFieldsNames[i] = smartMappers.mapper().names().indexName(); + this.fieldsDataType[i] = smartMappers.mapper().fieldDataType(); + } + + } + + if (script != null) { + this.script = new SearchScript(context.scriptSearchLookup(), scriptLang, script, params, context.scriptService()); + } else { + this.script = null; + } + + if (excluded.isEmpty() && pattern == null && this.script == null) { + aggregator = new StaticAggregatorValueProc(TermsFacetCollector.popFacets()); + } else { + aggregator = new AggregatorValueProc(TermsFacetCollector.popFacets(), excluded, pattern, this.script); + } + } + + @Override protected void doSetNextReader(IndexReader reader, int docBase) throws IOException { + for (int i = 0; i < fieldsNames.length; i++) { + fieldsData[i] = fieldDataCache.cache(fieldsDataType[i], reader, indexFieldsNames[i]); + } + if (script != null) { + script.setNextReader(reader); + } + } + + @Override protected void doCollect(int doc) throws IOException { + for (FieldData fieldData : fieldsData) { + fieldData.forEachValueInDoc(doc, aggregator); + } + } + + @Override public Facet facet() { + TObjectIntHashMap facets = aggregator.facets(); + if (facets.isEmpty()) { + TermsFacetCollector.pushFacets(facets); + return new InternalTermsFacet(facetName, arrayToCommaDelimitedString(fieldsNames), comparatorType, size, ImmutableList.of()); + } else { + // we need to fetch facets of "size * numberOfShards" because of problems in how they are distributed across shards + BoundedTreeSet ordered = new BoundedTreeSet(InternalTermsFacet.ComparatorType.COUNT.comparator(), size * numberOfShards); + for (TObjectIntIterator it = facets.iterator(); it.hasNext();) { + it.advance(); + ordered.add(new InternalTermsFacet.Entry(it.key(), it.value())); + } + TermsFacetCollector.pushFacets(facets); + return new InternalTermsFacet(facetName, arrayToCommaDelimitedString(fieldsNames), comparatorType, size, ordered); + } + } + + public static class AggregatorValueProc extends StaticAggregatorValueProc { + + private final ImmutableSet excluded; + + private final Matcher matcher; + + private final SearchScript script; + + private final Map scriptParams; + + public AggregatorValueProc(TObjectIntHashMap facets, ImmutableSet excluded, Pattern pattern, SearchScript script) { + super(facets); + this.excluded = excluded; + this.matcher = pattern != null ? pattern.matcher("") : null; + this.script = script; + if (script != null) { + scriptParams = Maps.newHashMapWithExpectedSize(4); + } else { + scriptParams = null; + } + } + + @Override public void onValue(int docId, String value) { + if (excluded != null && excluded.contains(value)) { + return; + } + if (matcher != null && !matcher.reset(value).matches()) { + return; + } + if (script != null) { + scriptParams.put("term", value); + Object scriptValue = script.execute(docId, scriptParams); + if (scriptValue == null) { + return; + } + if (scriptValue instanceof Boolean) { + if (!((Boolean) scriptValue)) { + return; + } + } else { + value = scriptValue.toString(); + } + } + super.onValue(docId, value); + } + } + + public static class StaticAggregatorValueProc implements FieldData.StringValueInDocProc { + + private final TObjectIntHashMap facets; + + public StaticAggregatorValueProc(TObjectIntHashMap facets) { + this.facets = facets; + } + + @Override public void onValue(int docId, String value) { + facets.adjustOrPutValue(value, 1, 1); + } + + public final TObjectIntHashMap facets() { + return facets; + } + } +} diff --git a/modules/test/integration/src/test/java/org/elasticsearch/test/integration/search/facets/SimpleFacetsTests.java b/modules/test/integration/src/test/java/org/elasticsearch/test/integration/search/facets/SimpleFacetsTests.java index ce72b6250ce..41587f84f81 100644 --- a/modules/test/integration/src/test/java/org/elasticsearch/test/integration/search/facets/SimpleFacetsTests.java +++ b/modules/test/integration/src/test/java/org/elasticsearch/test/integration/search/facets/SimpleFacetsTests.java @@ -368,6 +368,25 @@ public class SimpleFacetsTests extends AbstractNodesTests { assertThat(facet.entries().get(0).count(), equalTo(2)); assertThat(facet.entries().get(1).term(), equalTo("zzz")); assertThat(facet.entries().get(1).count(), equalTo(1)); + + // Fields + + searchResponse = client.prepareSearch() + .setQuery(matchAllQuery()) + .addFacet(termsFacet("facet1").fields("stag", "tag").size(10)) + .execute().actionGet(); + + facet = searchResponse.facets().facet("facet1"); + assertThat(facet.name(), equalTo("facet1")); + assertThat(facet.entries().size(), equalTo(4)); + assertThat(facet.entries().get(0).term(), anyOf(equalTo("111"), equalTo("yyy"))); + assertThat(facet.entries().get(0).count(), equalTo(2)); + assertThat(facet.entries().get(1).term(), anyOf(equalTo("111"), equalTo("yyy"))); + assertThat(facet.entries().get(1).count(), equalTo(2)); + assertThat(facet.entries().get(2).term(), anyOf(equalTo("zzz"), equalTo("xxx"))); + assertThat(facet.entries().get(2).count(), equalTo(1)); + assertThat(facet.entries().get(3).term(), anyOf(equalTo("zzz"), equalTo("xxx"))); + assertThat(facet.entries().get(3).count(), equalTo(1)); } @Test public void testTermFacetWithEqualTermDistribution() throws Exception {