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 386d0373d00..9d9503942e5 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 @@ -89,6 +89,15 @@ public class TermsFacetBuilder extends AbstractFacetBuilder { return this; } + /** + * Define a script field that will control the terms that will be used (and not filtered, as is the + * case when the script is provided on top of field / fields). + */ + public TermsFacetBuilder scriptField(String scriptField) { + this.script = scriptField; + return this; + } + /** * A set of terms that will be excluded. */ @@ -160,8 +169,8 @@ public class TermsFacetBuilder extends AbstractFacetBuilder { } @Override public void toXContent(XContentBuilder builder, Params params) throws IOException { - if (fieldName == null && fieldsNames == null) { - throw new SearchSourceBuilderException("field must be set on terms facet for facet [" + name + "]"); + if (fieldName == null && fieldsNames == null && script == null) { + throw new SearchSourceBuilderException("field/fields/script must be set on terms facet for facet [" + name + "]"); } builder.startObject(name); @@ -172,7 +181,7 @@ public class TermsFacetBuilder extends AbstractFacetBuilder { } else { builder.field("fields", fieldsNames); } - } else { + } else if (fieldName != null) { builder.field("field", fieldName); } builder.field("size", size); 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 4126cfaf80e..9a68746a01a 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 @@ -81,6 +81,8 @@ public class TermsFacetCollectorParser implements FacetCollectorParser { } else if (token.isValue()) { if ("field".equals(fieldName)) { field = parser.text(); + } else if ("script_field".equals(fieldName)) { + script = parser.text(); } else if ("size".equals(fieldName)) { size = parser.intValue(); } else if ("regex".equals(fieldName)) { @@ -108,6 +110,9 @@ public class TermsFacetCollectorParser implements FacetCollectorParser { if (fieldsNames != null) { return new TermsFieldsFacetCollector(facetName, fieldsNames, size, comparatorType, context, excluded, pattern, scriptLang, script, params); } + if (field == null && fieldsNames == null && script != null) { + return new TermsScriptFieldFacetCollector(facetName, 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/TermsScriptFieldFacetCollector.java b/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/terms/TermsScriptFieldFacetCollector.java new file mode 100644 index 00000000000..f4e94d687a2 --- /dev/null +++ b/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/terms/TermsScriptFieldFacetCollector.java @@ -0,0 +1,196 @@ +/* + * 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.field.data.FieldData; +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; + +/** + * @author kimchy (shay.banon) + */ +public class TermsScriptFieldFacetCollector extends AbstractFacetCollector { + + private final InternalTermsFacet.ComparatorType comparatorType; + + private final int size; + + private final int numberOfShards; + + private final String sScript; + + private final SearchScript script; + + private final Matcher matcher; + + private final ImmutableSet excluded; + + private final TObjectIntHashMap facets; + + public TermsScriptFieldFacetCollector(String facetName, int size, InternalTermsFacet.ComparatorType comparatorType, SearchContext context, + ImmutableSet excluded, Pattern pattern, String scriptLang, String script, Map params) { + super(facetName); + this.size = size; + this.comparatorType = comparatorType; + this.numberOfShards = context.numberOfShards(); + this.sScript = script; + this.script = new SearchScript(context.scriptSearchLookup(), scriptLang, script, params, context.scriptService()); + + this.excluded = excluded; + this.matcher = pattern != null ? pattern.matcher("") : null; + + this.facets = TermsFacetCollector.popFacets(); + } + + @Override protected void doSetNextReader(IndexReader reader, int docBase) throws IOException { + script.setNextReader(reader); + } + + @Override protected void doCollect(int doc) throws IOException { + Object o = script.execute(doc); + if (o == null) { + return; + } + if (o instanceof Iterable) { + for (Object o1 : ((Iterable) o)) { + String value = o1.toString(); + if (match(value)) { + facets.adjustOrPutValue(value, 1, 1); + } + } + } else if (o instanceof Object[]) { + for (Object o1 : ((Object[]) o)) { + String value = o1.toString(); + if (match(value)) { + facets.adjustOrPutValue(value, 1, 1); + } + } + } else { + String value = o.toString(); + if (match(value)) { + facets.adjustOrPutValue(value, 1, 1); + } + } + } + + private boolean match(String value) { + if (excluded != null && excluded.contains(value)) { + return false; + } + if (matcher != null && !matcher.reset(value).matches()) { + return false; + } + return true; + } + + @Override public Facet facet() { + if (facets.isEmpty()) { + TermsFacetCollector.pushFacets(facets); + return new InternalTermsFacet(facetName, sScript, 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, sScript, 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 41587f84f81..dae5eeb12af 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 @@ -369,7 +369,7 @@ public class SimpleFacetsTests extends AbstractNodesTests { assertThat(facet.entries().get(1).term(), equalTo("zzz")); assertThat(facet.entries().get(1).count(), equalTo(1)); - // Fields + // Fields Facets searchResponse = client.prepareSearch() .setQuery(matchAllQuery()) @@ -387,6 +387,45 @@ public class SimpleFacetsTests extends AbstractNodesTests { 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)); + + // Script Field + + searchResponse = client.prepareSearch() + .setQuery(matchAllQuery()) + .addFacet(termsFacet("facet1").scriptField("_source.stag").size(10)) + .addFacet(termsFacet("facet2").scriptField("_source.tag").size(10)) + .execute().actionGet(); + + facet = searchResponse.facets().facet("facet1"); + assertThat(facet.name(), equalTo("facet1")); + assertThat(facet.entries().size(), equalTo(1)); + assertThat(facet.entries().get(0).term(), equalTo("111")); + assertThat(facet.entries().get(0).count(), equalTo(2)); + + facet = searchResponse.facets().facet("facet2"); + assertThat(facet.name(), equalTo("facet2")); + assertThat(facet.entries().size(), equalTo(3)); + assertThat(facet.entries().get(0).term(), equalTo("yyy")); + assertThat(facet.entries().get(0).count(), equalTo(2)); + + // test joining two scripts using the same facet name :) + searchResponse = client.prepareSearch() + .setQuery(matchAllQuery()) + .addFacet(termsFacet("facet1").scriptField("_source.stag").size(10)) + .addFacet(termsFacet("facet1").scriptField("_source.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 {