add script support for terms facets, allowing to control if a certain term will be included in the facets result (boolean result), or muging of the term

modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/histogram/HistogramFacetCollectorParser.java

@@ -72,7 +72,7 @@ public class HistogramFacetCollectorParser implements FacetCollectorParser {
                     keyScript = parser.text();
                 } else if ("value_script".equals(fieldName) || "valueScript".equals(fieldName)) {
                     valueScript = parser.text();
-                } else if ("comparator".equals(fieldName)) {
+                } else if ("order".equals(fieldName) || "comparator".equals(fieldName)) {
                     comparatorType = HistogramFacet.ComparatorType.fromString(parser.text());
                 }
             }

modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/terms/TermsFacetBuilder.java

@@ -19,6 +19,7 @@
 
 package org.elasticsearch.search.facets.terms;
 
+import org.elasticsearch.common.collect.Maps;
 import org.elasticsearch.common.regex.Regex;
 import org.elasticsearch.common.xcontent.builder.XContentBuilder;
 import org.elasticsearch.index.query.xcontent.XContentFilterBuilder;
@@ -26,6 +27,7 @@ import org.elasticsearch.search.builder.SearchSourceBuilderException;
 import org.elasticsearch.search.facets.AbstractFacetBuilder;
 
 import java.io.IOException;
+import java.util.Map;
 
 /**
  * @author kimchy (shay.banon)
@@ -37,6 +39,8 @@ public class TermsFacetBuilder extends AbstractFacetBuilder {
     private String regex;
     private int regexFlags = 0;
     private TermsFacet.ComparatorType comparatorType;
+    private String script;
+    private Map<String, Object> params;
 
     public TermsFacetBuilder(String name) {
         super(name);
@@ -82,6 +86,19 @@ public class TermsFacetBuilder extends AbstractFacetBuilder {
         return this;
     }
 
+    public TermsFacetBuilder script(String script) {
+        this.script = script;
+        return this;
+    }
+
+    public TermsFacetBuilder param(String name, Object value) {
+        if (params == null) {
+            params = Maps.newHashMap();
+        }
+        params.put(name, value);
+        return this;
+    }
+
     @Override public void toXContent(XContentBuilder builder, Params params) throws IOException {
         if (fieldName == null) {
             throw new SearchSourceBuilderException("field must be set on terms facet for facet [" + name + "]");
@@ -107,6 +124,15 @@ public class TermsFacetBuilder extends AbstractFacetBuilder {
         if (comparatorType != null) {
             builder.field("order", comparatorType.name().toLowerCase());
         }
+
+        if (script != null) {
+            builder.field("script", script);
+            if (this.params != null) {
+                builder.field("params");
+                builder.map(this.params);
+            }
+        }
+
         builder.endObject();
 
         addFilterFacetAndGlobal(builder, params);

modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/terms/TermsFacetCollector.java

@@ -23,19 +23,24 @@ 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.thread.ThreadLocals;
 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.field.function.FieldsFunction;
+import org.elasticsearch.index.field.function.script.ScriptFieldsFunction;
 import org.elasticsearch.index.mapper.FieldMapper;
 import org.elasticsearch.index.mapper.MapperService;
+import org.elasticsearch.script.ScriptService;
 import org.elasticsearch.search.facets.Facet;
 import org.elasticsearch.search.facets.support.AbstractFacetCollector;
 
 import java.io.IOException;
 import java.util.ArrayDeque;
 import java.util.Deque;
+import java.util.Map;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
@@ -73,7 +78,12 @@ public class TermsFacetCollector extends AbstractFacetCollector {
 
     private final Pattern pattern;
 
-    public TermsFacetCollector(String facetName, String fieldName, int size, InternalTermsFacet.ComparatorType comparatorType, int numberOfShards, FieldDataCache fieldDataCache, MapperService mapperService, ImmutableSet<String> excluded, Pattern pattern) {
+    private final FieldsFunction scriptFunction;
+
+    private final Map<String, Object> params;
+
+    public TermsFacetCollector(String facetName, String fieldName, int size, InternalTermsFacet.ComparatorType comparatorType, int numberOfShards, FieldDataCache fieldDataCache, MapperService mapperService, ScriptService scriptService,
+                               ImmutableSet<String> excluded, Pattern pattern, String script, Map<String, Object> params) {
         super(facetName);
         this.fieldDataCache = fieldDataCache;
         this.size = size;
@@ -91,15 +101,31 @@ public class TermsFacetCollector extends AbstractFacetCollector {
             this.indexFieldName = fieldName;
             this.fieldDataType = FieldData.Type.STRING;
         }
-        if (excluded.isEmpty() && pattern == null) {
+
+        if (script != null) {
+            scriptFunction = new ScriptFieldsFunction(script, scriptService, mapperService, fieldDataCache);
+            if (params == null) {
+                this.params = Maps.newHashMapWithExpectedSize(1);
+            } else {
+                this.params = params;
+            }
+        } else {
+            this.params = null;
+            scriptFunction = null;
+        }
+
+        if (excluded.isEmpty() && pattern == null && scriptFunction == null) {
             aggregator = new StaticAggregatorValueProc(popFacets());
         } else {
-            aggregator = new AggregatorValueProc(popFacets(), excluded, pattern);
+            aggregator = new AggregatorValueProc(popFacets(), excluded, pattern, scriptFunction, params);
         }
     }
 
     @Override protected void doSetNextReader(IndexReader reader, int docBase) throws IOException {
         fieldData = fieldDataCache.cache(fieldDataType, reader, indexFieldName);
+        if (scriptFunction != null) {
+            scriptFunction.setNextReader(reader);
+        }
     }
 
     @Override protected void doCollect(int doc) throws IOException {
@@ -147,10 +173,17 @@ public class TermsFacetCollector extends AbstractFacetCollector {
 
         private final Matcher matcher;
 
-        public AggregatorValueProc(TObjectIntHashMap<String> facets, ImmutableSet<String> excluded, Pattern pattern) {
+        private final FieldsFunction scriptFunction;
+
+        private final Map<String, Object> params;
+
+        public AggregatorValueProc(TObjectIntHashMap<String> facets, ImmutableSet<String> excluded, Pattern pattern,
+                                   FieldsFunction scriptFunction, Map<String, Object> params) {
             super(facets);
             this.excluded = excluded;
             this.matcher = pattern != null ? pattern.matcher("") : null;
+            this.scriptFunction = scriptFunction;
+            this.params = params;
         }
 
         @Override public void onValue(int docId, String value) {
@@ -160,6 +193,17 @@ public class TermsFacetCollector extends AbstractFacetCollector {
             if (matcher != null && !matcher.reset(value).matches()) {
                 return;
             }
+            if (scriptFunction != null) {
+                params.put("term", value);
+                Object scriptValue = scriptFunction.execute(docId, params);
+                if (scriptValue instanceof Boolean) {
+                    if (!((Boolean) scriptValue)) {
+                        return;
+                    }
+                } else {
+                    value = scriptValue.toString();
+                }
+            }
             super.onValue(docId, value);
         }
     }

modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/terms/TermsFacetCollectorParser.java

@@ -27,6 +27,7 @@ import org.elasticsearch.search.facets.collector.FacetCollectorParser;
 import org.elasticsearch.search.internal.SearchContext;
 
 import java.io.IOException;
+import java.util.Map;
 import java.util.regex.Pattern;
 
 /**
@@ -50,9 +51,15 @@ public class TermsFacetCollectorParser implements FacetCollectorParser {
         String regex = null;
         String regexFlags = null;
         TermsFacet.ComparatorType comparatorType = TermsFacet.ComparatorType.COUNT;
+        String script = null;
+        Map<String, Object> params = null;
         while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
             if (token == XContentParser.Token.FIELD_NAME) {
                 fieldName = parser.currentName();
+            } else if (token == XContentParser.Token.START_OBJECT) {
+                if ("params".equals(fieldName)) {
+                    params = parser.map();
+                }
             } else if (token == XContentParser.Token.START_ARRAY) {
                 if ("exclude".equals(fieldName)) {
                     ImmutableSet.Builder<String> builder = ImmutableSet.builder();
@@ -72,6 +79,8 @@ public class TermsFacetCollectorParser implements FacetCollectorParser {
                     regexFlags = parser.text();
                 } else if ("order".equals(fieldName) || "comparator".equals(field)) {
                     comparatorType = TermsFacet.ComparatorType.fromString(parser.text());
+                } else if ("script".equals(fieldName)) {
+                    script = parser.text();
                 }
             }
         }
@@ -79,6 +88,7 @@ public class TermsFacetCollectorParser implements FacetCollectorParser {
         if (regex != null) {
             pattern = Regex.compile(regex, regexFlags);
         }
-        return new TermsFacetCollector(facetName, field, size, comparatorType, context.numberOfShards(), context.fieldDataCache(), context.mapperService(), excluded, pattern);
+        return new TermsFacetCollector(facetName, field, size, comparatorType, context.numberOfShards(), context.fieldDataCache(), context.mapperService(), context.scriptService(),
+                excluded, pattern, script, params);
     }
 }

modules/test/integration/src/test/java/org/elasticsearch/test/integration/search/facets/SimpleFacetsTests.java

@@ -261,6 +261,23 @@ public class SimpleFacetsTests extends AbstractNodesTests {
         assertThat(facet.entries().get(1).count(), equalTo(2));
         assertThat(facet.entries().get(2).term(), equalTo("zzz"));
         assertThat(facet.entries().get(2).count(), equalTo(1));
+
+        // Script
+
+        searchResponse = client.prepareSearch()
+                .setQuery(matchAllQuery())
+                .addFacet(termsFacet("facet1").field("tag").size(10).script("term + param1").param("param1", "a").order(TermsFacet.ComparatorType.TERM))
+                .execute().actionGet();
+
+        facet = searchResponse.facets().facet("facet1");
+        assertThat(facet.name(), equalTo("facet1"));
+        assertThat(facet.entries().size(), equalTo(3));
+        assertThat(facet.entries().get(0).term(), equalTo("xxxa"));
+        assertThat(facet.entries().get(0).count(), equalTo(1));
+        assertThat(facet.entries().get(1).term(), equalTo("yyya"));
+        assertThat(facet.entries().get(1).count(), equalTo(2));
+        assertThat(facet.entries().get(2).term(), equalTo("zzza"));
+        assertThat(facet.entries().get(2).count(), equalTo(1));
     }
 
     @Test public void testTermFacetWithEqualTermDistribution() throws Exception {