Term Facet: Use a script that provides the facets to index, closes #410.

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) {
+        if (fieldName == null && fieldsNames == null && script == null) {
-            throw new SearchSourceBuilderException("field must be set on terms facet for facet [" + name + "]");
+            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);

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);
     }
 }

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<String> excluded;
+
+    private final TObjectIntHashMap<String> facets;
+
+    public TermsScriptFieldFacetCollector(String facetName, int size, InternalTermsFacet.ComparatorType comparatorType, SearchContext context,
+                                          ImmutableSet<String> excluded, Pattern pattern, String scriptLang, String script, Map<String, Object> 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.<InternalTermsFacet.Entry>of());
+        } else {
+            // we need to fetch facets of "size * numberOfShards" because of problems in how they are distributed across shards
+            BoundedTreeSet<InternalTermsFacet.Entry> ordered = new BoundedTreeSet<InternalTermsFacet.Entry>(InternalTermsFacet.ComparatorType.COUNT.comparator(), size * numberOfShards);
+            for (TObjectIntIterator<String> 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<String> excluded;
+
+        private final Matcher matcher;
+
+        private final SearchScript script;
+
+        private final Map<String, Object> scriptParams;
+
+        public AggregatorValueProc(TObjectIntHashMap<String> facets, ImmutableSet<String> 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<String> facets;
+
+        public StaticAggregatorValueProc(TObjectIntHashMap<String> facets) {
+            this.facets = facets;
+        }
+
+        @Override public void onValue(int docId, String value) {
+            facets.adjustOrPutValue(value, 1, 1);
+        }
+
+        public final TObjectIntHashMap<String> facets() {
+            return facets;
+        }
+    }
+}

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 {