Facets: Script Histogram facet, closes #228.

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

@@ -490,6 +490,38 @@ public class SearchRequestBuilder {
         return this;
     }
 
+    public SearchRequestBuilder addFacetHistogramScript(String name, String keyScript, String valueScript) {
+        facetsBuilder().histogramScriptFacet(name, keyScript, valueScript);
+        return this;
+    }
+
+    public SearchRequestBuilder addFacetHistogramScript(String name, String keyScript, String valueScript, HistogramFacet.ComparatorType comparatorType) {
+        facetsBuilder().histogramScriptFacet(name, keyScript, valueScript, comparatorType);
+        return this;
+    }
+
+    public SearchRequestBuilder addFacetHistogramScript(String name, String keyScript, String valueScript, @Nullable Map<String, Object> params, long interval, HistogramFacet.ComparatorType comparatorType,
+                                                        @Nullable XContentFilterBuilder filter) {
+        facetsBuilder().histogramScriptFacet(name, keyScript, valueScript, params, interval, comparatorType, filter);
+        return this;
+    }
+
+    public SearchRequestBuilder addFacetHistogramScriptGlobal(String name, String keyScript, String valueScript) {
+        facetsBuilder().histogramScriptFacetGlobal(name, keyScript, valueScript);
+        return this;
+    }
+
+    public SearchRequestBuilder addFacetHistogramScriptGlobal(String name, String keyScript, String valueScript, HistogramFacet.ComparatorType comparatorType) {
+        facetsBuilder().histogramScriptFacetGlobal(name, keyScript, valueScript, comparatorType);
+        return this;
+    }
+
+    public SearchRequestBuilder addFacetHistogramScriptGlobal(String name, String keyScript, String valueScript, @Nullable Map<String, Object> params, long interval, HistogramFacet.ComparatorType comparatorType,
+                                                              @Nullable XContentFilterBuilder filter) {
+        facetsBuilder().histogramScriptFacetGlobal(name, keyScript, valueScript, params, interval, comparatorType, filter);
+        return this;
+    }
+
     /**
      * Adds a field to be highlighted with default fragment size of 100 characters, and
      * default number of fragments of 5.

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

@@ -49,6 +49,7 @@ public class SearchSourceFacetsBuilder implements ToXContent {
     private List<BuilderStatisticalFacet> statisticalFacets;
     private List<BuilderScriptStatisticalFacet> scriptStatisticalFacets;
     private List<BuilderHistogramFacet> histogramFacets;
+    private List<BuilderScriptHistogramFacet> scriptHistogramFacets;
 
     public SearchSourceFacetsBuilder queryFacet(String name, XContentQueryBuilder query) {
         return queryFacet(name, query, null);
@@ -243,8 +244,43 @@ public class SearchSourceFacetsBuilder implements ToXContent {
         return this;
     }
 
+    public SearchSourceFacetsBuilder histogramScriptFacet(String name, String keyScript, String valueScript) {
+        return histogramScriptFacet(name, keyScript, valueScript, null, -1, HistogramFacet.ComparatorType.KEY, null);
+    }
+
+    public SearchSourceFacetsBuilder histogramScriptFacet(String name, String keyScript, String valueScript, HistogramFacet.ComparatorType comparatorType) {
+        return histogramScriptFacet(name, keyScript, valueScript, null, -1, comparatorType, null);
+    }
+
+    public SearchSourceFacetsBuilder histogramScriptFacet(String name, String keyScript, String valueScript, @Nullable Map<String, Object> params, long interval, HistogramFacet.ComparatorType comparatorType,
+                                                          @Nullable XContentFilterBuilder filter) {
+        if (scriptHistogramFacets == null) {
+            scriptHistogramFacets = newArrayListWithCapacity(2);
+        }
+        scriptHistogramFacets.add(new BuilderScriptHistogramFacet(name, keyScript, valueScript, params, interval, comparatorType, filter, false));
+        return this;
+    }
+
+    public SearchSourceFacetsBuilder histogramScriptFacetGlobal(String name, String keyScript, String valueScript) {
+        return histogramScriptFacetGlobal(name, keyScript, valueScript, null, -1, HistogramFacet.ComparatorType.KEY, null);
+    }
+
+    public SearchSourceFacetsBuilder histogramScriptFacetGlobal(String name, String keyScript, String valueScript, HistogramFacet.ComparatorType comparatorType) {
+        return histogramScriptFacetGlobal(name, keyScript, valueScript, null, -1, comparatorType, null);
+    }
+
+    public SearchSourceFacetsBuilder histogramScriptFacetGlobal(String name, String keyScript, String valueScript, @Nullable Map<String, Object> params, long interval, HistogramFacet.ComparatorType comparatorType,
+                                                                @Nullable XContentFilterBuilder filter) {
+        if (scriptHistogramFacets == null) {
+            scriptHistogramFacets = newArrayListWithCapacity(2);
+        }
+        scriptHistogramFacets.add(new BuilderScriptHistogramFacet(name, keyScript, valueScript, params, interval, comparatorType, filter, true));
+        return this;
+    }
+
+
     @Override public void toXContent(XContentBuilder builder, Params params) throws IOException {
-        if (queryFacets == null && termsFacets == null && statisticalFacets == null && histogramFacets == null && scriptStatisticalFacets == null) {
+        if (queryFacets == null && termsFacets == null && statisticalFacets == null && histogramFacets == null && scriptStatisticalFacets == null && scriptHistogramFacets == null) {
             return;
         }
         builder.field("facets");
@@ -363,6 +399,36 @@ public class SearchSourceFacetsBuilder implements ToXContent {
             }
         }
 
+        if (scriptHistogramFacets != null) {
+            for (BuilderScriptHistogramFacet histogramFacet : scriptHistogramFacets) {
+                builder.startObject(histogramFacet.name());
+
+                builder.startObject(HistogramFacetCollectorParser.NAME);
+                builder.field("key_script", histogramFacet.keyScript());
+                builder.field("value_script", histogramFacet.valueScript());
+                if (histogramFacet.interval() > 0) {
+                    builder.field("interval", histogramFacet.interval());
+                }
+                if (histogramFacet.params() != null) {
+                    builder.field("params");
+                    builder.map(histogramFacet.params());
+                }
+                builder.field("comparator", histogramFacet.comparatorType().description());
+                builder.endObject();
+
+                if (histogramFacet.filter() != null) {
+                    builder.field("filter");
+                    histogramFacet.filter().toXContent(builder, params);
+                }
+
+                if (histogramFacet.global() != null) {
+                    builder.field("global", histogramFacet.global());
+                }
+
+                builder.endObject();
+            }
+        }
+
         builder.endObject();
     }
 
@@ -498,6 +564,60 @@ public class SearchSourceFacetsBuilder implements ToXContent {
         }
     }
 
+    private static class BuilderScriptHistogramFacet {
+        private final String name;
+        private final String keyScript;
+        private final String valueScript;
+        private final Map<String, Object> params;
+        private final long interval;
+        private final HistogramFacet.ComparatorType comparatorType;
+        private final XContentFilterBuilder filter;
+        private final Boolean global;
+
+        private BuilderScriptHistogramFacet(String name, String keyScript, String valueScript, Map<String, Object> params, long interval, HistogramFacet.ComparatorType comparatorType, XContentFilterBuilder filter, Boolean global) {
+            this.name = name;
+            this.keyScript = keyScript;
+            this.valueScript = valueScript;
+            this.params = params;
+            this.interval = interval;
+            this.comparatorType = comparatorType;
+            this.filter = filter;
+            this.global = global;
+        }
+
+        public String name() {
+            return name;
+        }
+
+        public Map<String, Object> params() {
+            return params;
+        }
+
+        public String keyScript() {
+            return keyScript;
+        }
+
+        public String valueScript() {
+            return valueScript;
+        }
+
+        public long interval() {
+            return interval;
+        }
+
+        public HistogramFacet.ComparatorType comparatorType() {
+            return comparatorType;
+        }
+
+        public XContentFilterBuilder filter() {
+            return filter;
+        }
+
+        public Boolean global() {
+            return global;
+        }
+    }
+
     private static class BuilderHistogramFacet {
         private final String name;
         private final String keyFieldName;

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

@@ -19,6 +19,7 @@
 
 package org.elasticsearch.search.facets.histogram;
 
+import org.elasticsearch.common.thread.ThreadLocals;
 import org.elasticsearch.common.unit.TimeValue;
 import org.elasticsearch.common.xcontent.XContentParser;
 import org.elasticsearch.search.facets.FacetPhaseExecutionException;
@@ -27,12 +28,20 @@ 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 HistogramFacetCollectorParser 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 = "histogram";
 
     @Override public String name() {
@@ -42,13 +51,20 @@ public class HistogramFacetCollectorParser implements FacetCollectorParser {
     @Override public FacetCollector parser(String facetName, XContentParser parser, SearchContext context) throws IOException {
         String keyField = null;
         String valueField = null;
-        long interval = -1;
+        String keyScript = null;
+        String valueScript = null;
+        Map<String, Object> params = cachedParams.get().get();
+        long interval = 0;
         HistogramFacet.ComparatorType comparatorType = HistogramFacet.ComparatorType.KEY;
         XContentParser.Token token;
         String fieldName = 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.isValue()) {
                 if ("field".equals(fieldName)) {
                     keyField = parser.text();
@@ -60,17 +76,25 @@ public class HistogramFacetCollectorParser implements FacetCollectorParser {
                     interval = parser.longValue();
                 } else if ("time_interval".equals(fieldName)) {
                     interval = TimeValue.parseTimeValue(parser.text(), null).millis();
+                } else if ("key_script".equals(fieldName) || "keyScript".equals(fieldName)) {
+                    keyScript = parser.text();
+                } else if ("value_script".equals(fieldName) || "valueScript".equals(fieldName)) {
+                    valueScript = parser.text();
                 } else if ("comparator".equals(fieldName)) {
                     comparatorType = HistogramFacet.ComparatorType.fromString(parser.text());
                 }
             }
         }
 
+        if (keyScript != null && valueScript != null) {
+            return new ScriptHistogramFacetCollector(facetName, keyScript, valueScript, params, interval, comparatorType, context.scriptService(), context.fieldDataCache(), context.mapperService());
+        }
+
         if (keyField == null) {
             throw new FacetPhaseExecutionException(facetName, "key field is required to be set for histogram facet, either using [field] or using [key_field]");
         }
 
-        if (interval == -1) {
+        if (interval <= 0) {
             throw new FacetPhaseExecutionException(facetName, "[interval] is required to be set for histogram facet");
         }
 

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

@@ -0,0 +1,89 @@
+/*
+ * 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.histogram;
+
+import org.apache.lucene.index.IndexReader;
+import org.elasticsearch.common.trove.TLongDoubleHashMap;
+import org.elasticsearch.common.trove.TLongLongHashMap;
+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 ScriptHistogramFacetCollector extends AbstractFacetCollector {
+
+    private final FieldsFunction keyFunction;
+
+    private final FieldsFunction valueFunction;
+
+    private final Map<String, Object> params;
+
+    private final long interval;
+
+    private final HistogramFacet.ComparatorType comparatorType;
+
+    private final TLongLongHashMap counts = new TLongLongHashMap();
+
+    private final TLongDoubleHashMap totals = new TLongDoubleHashMap();
+
+    public ScriptHistogramFacetCollector(String facetName, String keyScript, String valueScript, Map<String, Object> params, long interval, HistogramFacet.ComparatorType comparatorType, ScriptService scriptService, FieldDataCache fieldDataCache, MapperService mapperService) {
+        super(facetName);
+        this.keyFunction = new ScriptFieldsFunction(keyScript, scriptService, mapperService, fieldDataCache);
+        this.valueFunction = new ScriptFieldsFunction(valueScript, scriptService, mapperService, fieldDataCache);
+        this.interval = interval > 0 ? interval : 0;
+        this.params = params;
+        this.comparatorType = comparatorType;
+    }
+
+    @Override protected void doCollect(int doc) throws IOException {
+        Number keyValue = (Number) keyFunction.execute(doc, params);
+        long bucket;
+        if (interval == 0) {
+            bucket = keyValue.longValue();
+        } else {
+            bucket = bucket(keyValue.doubleValue(), interval);
+        }
+        double value = ((Number) valueFunction.execute(doc, params)).doubleValue();
+        counts.adjustOrPutValue(bucket, 1, 1);
+        totals.adjustOrPutValue(bucket, value, value);
+    }
+
+    @Override protected void doSetNextReader(IndexReader reader, int docBase) throws IOException {
+        keyFunction.setNextReader(reader);
+        valueFunction.setNextReader(reader);
+    }
+
+    @Override public Facet facet() {
+        return new InternalHistogramFacet(facetName, "_na", "_na", -1, comparatorType, counts, totals);
+    }
+
+    public static long bucket(double value, long interval) {
+        return (((long) (value / interval)) * interval);
+    }
+}

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

@@ -22,7 +22,6 @@ 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;
 import org.elasticsearch.search.facets.statistical.StatisticalFacet;
 import org.elasticsearch.search.facets.terms.TermsFacet;
@@ -133,7 +132,7 @@ public class SimpleFacetsTests extends AbstractNodesTests {
         client.admin().indices().prepareRefresh().execute().actionGet();
 
         SearchResponse searchResponse = client.prepareSearch()
-                .setQuery(QueryBuilders.matchAllQuery())
+                .setQuery(matchAllQuery())
                 .addFacetStatistical("stats1", "num")
                 .addFacetStatistical("stats2", "multi_num")
                 .addFacetStatisticalScript("stats3", "doc['num'].value * 2")
@@ -185,29 +184,41 @@ public class SimpleFacetsTests extends AbstractNodesTests {
 
         client.prepareIndex("test", "type1").setSource(jsonBuilder().startObject()
                 .field("num", 1055)
+                .field("date", "1970-01-01T00:00:00")
                 .startArray("multi_num").value(13.0f).value(23.f).endArray()
                 .endObject()).execute().actionGet();
         client.admin().indices().prepareFlush().setRefresh(true).execute().actionGet();
 
         client.prepareIndex("test", "type1").setSource(jsonBuilder().startObject()
                 .field("num", 1065)
+                .field("date", "1970-01-01T00:00:25")
                 .startArray("multi_num").value(15.0f).value(31.0f).endArray()
                 .endObject()).execute().actionGet();
         client.admin().indices().prepareRefresh().execute().actionGet();
 
         client.prepareIndex("test", "type1").setSource(jsonBuilder().startObject()
                 .field("num", 1175)
+                .field("date", "1970-01-01T00:02:00")
                 .startArray("multi_num").value(17.0f).value(25.0f).endArray()
                 .endObject()).execute().actionGet();
         client.admin().indices().prepareRefresh().execute().actionGet();
 
         SearchResponse searchResponse = client.prepareSearch()
-                .setQuery(QueryBuilders.matchAllQuery())
+                .setQuery(matchAllQuery())
                 .addFacetHistogram("stats1", "num", 100)
                 .addFacetHistogram("stats2", "multi_num", 10)
                 .addFacetHistogram("stats3", "num", "multi_num", 100)
+                .addFacetHistogramScript("stats4", "doc['date'].date.minuteOfHour", "doc['num'].value")
                 .execute().actionGet();
 
+        if (searchResponse.failedShards() > 0) {
+            logger.warn("Failed shards:");
+            for (ShardSearchFailure shardSearchFailure : searchResponse.shardFailures()) {
+                logger.warn("-> {}", shardSearchFailure);
+            }
+        }
+        assertThat(searchResponse.failedShards(), equalTo(0));
+
         HistogramFacet facet = searchResponse.facets().facet(HistogramFacet.class, "stats1");
         assertThat(facet.name(), equalTo("stats1"));
         assertThat(facet.entries().size(), equalTo(2));
@@ -247,5 +258,17 @@ public class SimpleFacetsTests extends AbstractNodesTests {
         assertThat(facet.entries().get(1).count(), equalTo(2l));
         assertThat(facet.entries().get(1).total(), equalTo(42d));
         assertThat(facet.entries().get(1).mean(), equalTo(21d));
+
+        facet = searchResponse.facets().facet(HistogramFacet.class, "stats4");
+        assertThat(facet.name(), equalTo("stats4"));
+        assertThat(facet.entries().size(), equalTo(2));
+        assertThat(facet.entries().get(0).key(), equalTo(0l));
+        assertThat(facet.entries().get(0).count(), equalTo(2l));
+        assertThat(facet.entries().get(0).total(), equalTo(2120d));
+        assertThat(facet.entries().get(0).mean(), equalTo(1060d));
+        assertThat(facet.entries().get(1).key(), equalTo(2l));
+        assertThat(facet.entries().get(1).count(), equalTo(1l));
+        assertThat(facet.entries().get(1).total(), equalTo(1175d));
+        assertThat(facet.entries().get(1).mean(), equalTo(1175d));
     }
 }