Facets: Script statistical facets, closes #227.

modules/elasticsearch/src/main/java/org/elasticsearch/client/action/search/SearchRequestBuilder.java

@@ -346,11 +346,49 @@ public class SearchRequestBuilder {
         return this;
     }
 
+    /**
+     * Adds a numeric statistical facet for the provided field name.
+     *
+     * @param name      The name of the facet
+     * @param fieldName The name of the <b>numeric</b> field
+     * @param filter    An optional filter to reduce the scope of the facet
+     * @see org.elasticsearch.search.facets.statistical.StatisticalFacet
+     */
     public SearchRequestBuilder addFacetStatistical(String name, String fieldName, @Nullable XContentFilterBuilder filter) {
         facetsBuilder().statisticalFacet(name, fieldName, filter);
         return this;
     }
 
+    public SearchRequestBuilder addFacetStatisticalScript(String name, String script) {
+        facetsBuilder().statisticalScriptFacet(name, script);
+        return this;
+    }
+
+    public SearchRequestBuilder addFacetStatisticalScript(String name, String script, @Nullable Map<String, Object> params) {
+        facetsBuilder().statisticalScriptFacet(name, script, params);
+        return this;
+    }
+
+    public SearchRequestBuilder addFacetStatisticalScript(String name, String script, @Nullable Map<String, Object> params, @Nullable XContentFilterBuilder filter) {
+        facetsBuilder().statisticalScriptFacet(name, script, params, filter);
+        return this;
+    }
+
+    public SearchRequestBuilder addFacetGlobalStatisticalScript(String name, String script) {
+        facetsBuilder().statisticalScriptFacetGlobal(name, script);
+        return this;
+    }
+
+    public SearchRequestBuilder addFacetGlobalStatisticalScript(String name, String script, @Nullable Map<String, Object> params) {
+        facetsBuilder().statisticalScriptFacetGlobal(name, script, params);
+        return this;
+    }
+
+    public SearchRequestBuilder addFacetGlobalStatisticalScript(String name, String script, @Nullable Map<String, Object> params, @Nullable XContentFilterBuilder filter) {
+        facetsBuilder().statisticalScriptFacetGlobal(name, script, params, filter);
+        return this;
+    }
+
     /**
      * Adds a numeric statistical <b>global</b> facet for the provided field name.
      *

modules/elasticsearch/src/main/java/org/elasticsearch/search/builder/SearchSourceFacetsBuilder.java

@@ -32,6 +32,7 @@ import org.elasticsearch.search.facets.terms.TermsFacetCollectorParser;
 import javax.annotation.Nullable;
 import java.io.IOException;
 import java.util.List;
+import java.util.Map;
 
 import static org.elasticsearch.common.collect.Lists.*;
 
@@ -46,6 +47,7 @@ public class SearchSourceFacetsBuilder implements ToXContent {
     private List<BuilderQueryFacet> queryFacets;
     private List<BuilderTermsFacet> termsFacets;
     private List<BuilderStatisticalFacet> statisticalFacets;
+    private List<BuilderScriptStatisticalFacet> scriptStatisticalFacets;
     private List<BuilderHistogramFacet> histogramFacets;
 
     public SearchSourceFacetsBuilder queryFacet(String name, XContentQueryBuilder query) {
@@ -133,6 +135,38 @@ public class SearchSourceFacetsBuilder implements ToXContent {
         return this;
     }
 
+    public SearchSourceFacetsBuilder statisticalScriptFacet(String name, String script) {
+        return statisticalScriptFacet(name, script, null, null);
+    }
+
+    public SearchSourceFacetsBuilder statisticalScriptFacet(String name, String script, @Nullable Map<String, Object> params) {
+        return statisticalScriptFacet(name, script, params, null);
+    }
+
+    public SearchSourceFacetsBuilder statisticalScriptFacet(String name, String script, @Nullable Map<String, Object> params, @Nullable XContentFilterBuilder filter) {
+        if (scriptStatisticalFacets == null) {
+            scriptStatisticalFacets = newArrayListWithCapacity(2);
+        }
+        scriptStatisticalFacets.add(new BuilderScriptStatisticalFacet(name, script, params, filter, false));
+        return this;
+    }
+
+    public SearchSourceFacetsBuilder statisticalScriptFacetGlobal(String name, String script) {
+        return statisticalScriptFacetGlobal(name, script, null, null);
+    }
+
+    public SearchSourceFacetsBuilder statisticalScriptFacetGlobal(String name, String script, @Nullable Map<String, Object> params) {
+        return statisticalScriptFacetGlobal(name, script, params, null);
+    }
+
+    public SearchSourceFacetsBuilder statisticalScriptFacetGlobal(String name, String script, @Nullable Map<String, Object> params, @Nullable XContentFilterBuilder filter) {
+        if (scriptStatisticalFacets == null) {
+            scriptStatisticalFacets = newArrayListWithCapacity(2);
+        }
+        scriptStatisticalFacets.add(new BuilderScriptStatisticalFacet(name, script, params, filter, true));
+        return this;
+    }
+
     public SearchSourceFacetsBuilder histogramFacet(String name, String fieldName, long interval) {
         return histogramFacet(name, fieldName, interval, HistogramFacet.ComparatorType.KEY, null);
     }
@@ -210,7 +244,7 @@ public class SearchSourceFacetsBuilder implements ToXContent {
     }
 
     @Override public void toXContent(XContentBuilder builder, Params params) throws IOException {
-        if (queryFacets == null && termsFacets == null && statisticalFacets == null && histogramFacets == null) {
+        if (queryFacets == null && termsFacets == null && statisticalFacets == null && histogramFacets == null && scriptStatisticalFacets == null) {
             return;
         }
         builder.field("facets");
@@ -276,6 +310,31 @@ public class SearchSourceFacetsBuilder implements ToXContent {
             }
         }
 
+        if (scriptStatisticalFacets != null) {
+            for (BuilderScriptStatisticalFacet statisticalFacet : scriptStatisticalFacets) {
+                builder.startObject(statisticalFacet.name());
+
+                builder.startObject(StatisticalFacetCollectorParser.NAME);
+                builder.field("script", statisticalFacet.script());
+                if (statisticalFacet.params() != null) {
+                    builder.field("params");
+                    builder.map(statisticalFacet.params());
+                }
+                builder.endObject();
+
+                if (statisticalFacet.filter() != null) {
+                    builder.field("filter");
+                    statisticalFacet.filter().toXContent(builder, params);
+                }
+
+                if (statisticalFacet.global() != null) {
+                    builder.field("global", statisticalFacet.global());
+                }
+
+                builder.endObject();
+            }
+        }
+
         if (histogramFacets != null) {
             for (BuilderHistogramFacet histogramFacet : histogramFacets) {
                 builder.startObject(histogramFacet.name());
@@ -373,6 +432,42 @@ public class SearchSourceFacetsBuilder implements ToXContent {
         }
     }
 
+    private static class BuilderScriptStatisticalFacet {
+        private final String name;
+        private final String script;
+        private final Map<String, Object> params;
+        private final XContentFilterBuilder filter;
+        private final Boolean global;
+
+        private BuilderScriptStatisticalFacet(String name, String script, Map<String, Object> params, XContentFilterBuilder filter, Boolean global) {
+            this.name = name;
+            this.script = script;
+            this.params = params;
+            this.filter = filter;
+            this.global = global;
+        }
+
+        public String name() {
+            return name;
+        }
+
+        public String script() {
+            return script;
+        }
+
+        public Map<String, Object> params() {
+            return params;
+        }
+
+        public XContentFilterBuilder filter() {
+            return filter;
+        }
+
+        public Boolean global() {
+            return global;
+        }
+    }
+
     private static class BuilderStatisticalFacet {
         private final String name;
         private final String fieldName;

modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/statistical/ScriptStatisticalFacetCollector.java

@@ -0,0 +1,79 @@
+/*
+ * 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.statistical;
+
+import org.apache.lucene.index.IndexReader;
+import org.elasticsearch.index.cache.field.data.FieldDataCache;
+import org.elasticsearch.index.field.function.FieldsFunction;
+import org.elasticsearch.index.field.function.script.ScriptFieldsFunction;
+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.Map;
+
+/**
+ * @author kimchy (shay.banon)
+ */
+public class ScriptStatisticalFacetCollector extends AbstractFacetCollector {
+
+    private final FieldsFunction function;
+
+    private final Map<String, Object> params;
+
+    private double min = Double.NaN;
+
+    private double max = Double.NaN;
+
+    private double total = 0;
+
+    private double sumOfSquares = 0.0;
+
+    private long count;
+
+    public ScriptStatisticalFacetCollector(String facetName, String script, Map<String, Object> params, ScriptService scriptService, FieldDataCache fieldDataCache, MapperService mapperService) {
+        super(facetName);
+        this.params = params;
+        this.function = new ScriptFieldsFunction(script, scriptService, mapperService, fieldDataCache);
+    }
+
+    @Override protected void doCollect(int doc) throws IOException {
+        double value = ((Number) function.execute(doc, params)).doubleValue();
+        if (value < min || Double.isNaN(min)) {
+            min = value;
+        }
+        if (value > max || Double.isNaN(max)) {
+            max = value;
+        }
+        sumOfSquares += value * value;
+        total += value;
+        count++;
+    }
+
+    @Override protected void doSetNextReader(IndexReader reader, int docBase) throws IOException {
+        function.setNextReader(reader);
+    }
+
+    @Override public Facet facet() {
+        return new InternalStatisticalFacet(facetName, "_na", min, max, total, sumOfSquares, count);
+    }
+}

modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/statistical/StatisticalFacetCollectorParser.java

@@ -19,18 +19,28 @@
 
 package org.elasticsearch.search.facets.statistical;
 
+import org.elasticsearch.common.thread.ThreadLocals;
 import org.elasticsearch.common.xcontent.XContentParser;
+import org.elasticsearch.search.facets.FacetPhaseExecutionException;
 import org.elasticsearch.search.facets.collector.FacetCollector;
 import org.elasticsearch.search.facets.collector.FacetCollectorParser;
 import org.elasticsearch.search.internal.SearchContext;
 
 import java.io.IOException;
+import java.util.HashMap;
+import java.util.Map;
 
 /**
  * @author kimchy (shay.banon)
  */
 public class StatisticalFacetCollectorParser implements FacetCollectorParser {
 
+    private static ThreadLocal<ThreadLocals.CleanableValue<Map<String, Object>>> cachedParams = new ThreadLocal<ThreadLocals.CleanableValue<Map<String, Object>>>() {
+        @Override protected ThreadLocals.CleanableValue<Map<String, Object>> initialValue() {
+            return new ThreadLocals.CleanableValue<Map<String, Object>>(new HashMap<String, Object>());
+        }
+    };
+
     public static final String NAME = "statistical";
 
     @Override public String name() {
@@ -40,17 +50,33 @@ public class StatisticalFacetCollectorParser implements FacetCollectorParser {
     @Override public FacetCollector parser(String facetName, XContentParser parser, SearchContext context) throws IOException {
         String field = null;
 
-        String fieldName = null;
+        String currentFieldName = null;
+        String script = null;
+        Map<String, Object> params = cachedParams.get().get();
+        params.clear();
         XContentParser.Token token;
         while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
             if (token == XContentParser.Token.FIELD_NAME) {
-                fieldName = parser.currentName();
+                currentFieldName = parser.currentName();
+            } else if (token == XContentParser.Token.START_OBJECT) {
+                if ("params".equals(currentFieldName)) {
+                    params = parser.map();
+                }
             } else if (token.isValue()) {
-                if ("field".equals(fieldName)) {
+                if ("field".equals(currentFieldName)) {
                     field = parser.text();
+                } else if ("script".equals(currentFieldName)) {
+                    script = parser.text();
                 }
             }
         }
-        return new StatisticalFacetCollector(facetName, field, context.fieldDataCache(), context.mapperService());
+        if (script == null && field == null) {
+            throw new FacetPhaseExecutionException(facetName, "statistical facet requires either [script] or [field] to be set");
+        }
+        if (field != null) {
+            return new StatisticalFacetCollector(facetName, field, context.fieldDataCache(), context.mapperService());
+        } else {
+            return new ScriptStatisticalFacetCollector(facetName, script, params, context.scriptService(), context.fieldDataCache(), context.mapperService());
+        }
     }
 }

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

@@ -20,6 +20,7 @@
 package org.elasticsearch.test.integration.search.facets;
 
 import org.elasticsearch.action.search.SearchResponse;
+import org.elasticsearch.action.search.ShardSearchFailure;
 import org.elasticsearch.client.Client;
 import org.elasticsearch.index.query.xcontent.QueryBuilders;
 import org.elasticsearch.search.facets.histogram.HistogramFacet;
@@ -135,8 +136,17 @@ public class SimpleFacetsTests extends AbstractNodesTests {
                 .setQuery(QueryBuilders.matchAllQuery())
                 .addFacetStatistical("stats1", "num")
                 .addFacetStatistical("stats2", "multi_num")
+                .addFacetStatisticalScript("stats3", "doc['num'].value * 2")
                 .execute().actionGet();
 
+        if (searchResponse.failedShards() > 0) {
+            logger.warn("Failed shards:");
+            for (ShardSearchFailure shardSearchFailure : searchResponse.shardFailures()) {
+                logger.warn("-> {}", shardSearchFailure);
+            }
+        }
+        assertThat(searchResponse.failedShards(), equalTo(0));
+
         StatisticalFacet facet = searchResponse.facets().facet(StatisticalFacet.class, "stats1");
         assertThat(facet.name(), equalTo(facet.name()));
         assertThat(facet.count(), equalTo(2l));
@@ -153,6 +163,15 @@ public class SimpleFacetsTests extends AbstractNodesTests {
         assertThat(facet.min(), equalTo(1d));
         assertThat(facet.max(), equalTo(4d));
         assertThat(facet.mean(), equalTo(2.5d));
+
+        facet = searchResponse.facets().facet(StatisticalFacet.class, "stats3");
+        assertThat(facet.name(), equalTo(facet.name()));
+        assertThat(facet.count(), equalTo(2l));
+        assertThat(facet.total(), equalTo(6d));
+        assertThat(facet.min(), equalTo(2d));
+        assertThat(facet.max(), equalTo(4d));
+        assertThat(facet.mean(), equalTo(3d));
+        assertThat(facet.sumOfSquares(), equalTo(20d));
     }
 
     @Test public void testHistoFacets() throws Exception {