diff --git a/modules/elasticsearch/src/main/java/org/elasticsearch/client/action/search/SearchRequestBuilder.java b/modules/elasticsearch/src/main/java/org/elasticsearch/client/action/search/SearchRequestBuilder.java index be4319f22fc..c039ef8e6c5 100644 --- a/modules/elasticsearch/src/main/java/org/elasticsearch/client/action/search/SearchRequestBuilder.java +++ b/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 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 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. diff --git a/modules/elasticsearch/src/main/java/org/elasticsearch/search/builder/SearchSourceFacetsBuilder.java b/modules/elasticsearch/src/main/java/org/elasticsearch/search/builder/SearchSourceFacetsBuilder.java index 9d6130cac1f..9e16e336794 100644 --- a/modules/elasticsearch/src/main/java/org/elasticsearch/search/builder/SearchSourceFacetsBuilder.java +++ b/modules/elasticsearch/src/main/java/org/elasticsearch/search/builder/SearchSourceFacetsBuilder.java @@ -49,6 +49,7 @@ public class SearchSourceFacetsBuilder implements ToXContent { private List statisticalFacets; private List scriptStatisticalFacets; private List histogramFacets; + private List 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 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 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 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 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 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; diff --git a/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/histogram/HistogramFacetCollectorParser.java b/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/histogram/HistogramFacetCollectorParser.java index d71e5b23a2f..51995cc58fc 100644 --- a/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/histogram/HistogramFacetCollectorParser.java +++ b/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>> cachedParams = new ThreadLocal>>() { + @Override protected ThreadLocals.CleanableValue> initialValue() { + return new ThreadLocals.CleanableValue>(new HashMap()); + } + }; + 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 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"); } diff --git a/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/histogram/ScriptHistogramFacetCollector.java b/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/histogram/ScriptHistogramFacetCollector.java new file mode 100644 index 00000000000..6f5d789c2c8 --- /dev/null +++ b/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 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 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); + } +} \ No newline at end of file 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 a410bbd55aa..9333c1e77dc 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 @@ -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)); } }