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 d2ba340d567..be4319f22fc 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 @@ -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 numeric 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 params) { + facetsBuilder().statisticalScriptFacet(name, script, params); + return this; + } + + public SearchRequestBuilder addFacetStatisticalScript(String name, String script, @Nullable Map 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 params) { + facetsBuilder().statisticalScriptFacetGlobal(name, script, params); + return this; + } + + public SearchRequestBuilder addFacetGlobalStatisticalScript(String name, String script, @Nullable Map params, @Nullable XContentFilterBuilder filter) { + facetsBuilder().statisticalScriptFacetGlobal(name, script, params, filter); + return this; + } + /** * Adds a numeric statistical global facet for the provided field name. * 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 798738a4305..9d6130cac1f 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 @@ -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 queryFacets; private List termsFacets; private List statisticalFacets; + private List scriptStatisticalFacets; private List 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 params) { + return statisticalScriptFacet(name, script, params, null); + } + + public SearchSourceFacetsBuilder statisticalScriptFacet(String name, String script, @Nullable Map 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 params) { + return statisticalScriptFacetGlobal(name, script, params, null); + } + + public SearchSourceFacetsBuilder statisticalScriptFacetGlobal(String name, String script, @Nullable Map 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 params; + private final XContentFilterBuilder filter; + private final Boolean global; + + private BuilderScriptStatisticalFacet(String name, String script, Map 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 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; diff --git a/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/statistical/ScriptStatisticalFacetCollector.java b/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/statistical/ScriptStatisticalFacetCollector.java new file mode 100644 index 00000000000..2580f236f9d --- /dev/null +++ b/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 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 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); + } +} \ No newline at end of file diff --git a/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/statistical/StatisticalFacetCollectorParser.java b/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/statistical/StatisticalFacetCollectorParser.java index 4a95e523aa9..52c90b5ff7f 100644 --- a/modules/elasticsearch/src/main/java/org/elasticsearch/search/facets/statistical/StatisticalFacetCollectorParser.java +++ b/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>> cachedParams = new ThreadLocal>>() { + @Override protected ThreadLocals.CleanableValue> initialValue() { + return new ThreadLocals.CleanableValue>(new HashMap()); + } + }; + 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 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()); + } } } \ 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 5e549c0fd2b..a410bbd55aa 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 @@ -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 {