Add script support to value_count aggregations.

Close #5001
elastic · Feb 4, 2014 · 6777be6 · 6777be6
1 parent 9333a3c
commit 6777be6
Show file tree

Hide file tree

Showing 4 changed files with 121 additions and 21 deletions.
diff --git a/docs/reference/search/aggregations/metrics/valuecount-aggregation.asciidoc b/docs/reference/search/aggregations/metrics/valuecount-aggregation.asciidoc
@@ -2,7 +2,7 @@
 === Value Count
 
 A `single-value` metrics aggregation that counts the number of values that are extracted from the aggregated documents.
-These values can be extracted either from specific fields in the documents. Typically,
+These values can be extracted either from specific fields in the documents, or be generated by a provided script. Typically,
 this aggregator will be used in conjunction with other single-value aggregations. For example, when computing the `avg`
 one might be interested in the number of values the average is computed over.
 
@@ -33,3 +33,17 @@ Response:
 The name of the aggregation (`grades_count` above) also serves as the key by which the aggregation result can be
 retrieved from the returned response.
 
+==== Script
+
+Counting the values generated by a script:
+
+[source,js]
+--------------------------------------------------
+{
+    ...,
+
+    "aggs" : {
+        "grades_count" : { "value_count" : { "script" : "doc['grade'].value" } }
+    }
+}
+--------------------------------------------------
diff --git a/...main/java/org/elasticsearch/search/aggregations/metrics/valuecount/ValueCountBuilder.java b/...main/java/org/elasticsearch/search/aggregations/metrics/valuecount/ValueCountBuilder.java
@@ -19,31 +19,15 @@
 
 package org.elasticsearch.search.aggregations.metrics.valuecount;
 
-import org.elasticsearch.common.xcontent.XContentBuilder;
-import org.elasticsearch.search.aggregations.metrics.MetricsAggregationBuilder;
-
-import java.io.IOException;
+import org.elasticsearch.search.aggregations.metrics.ValuesSourceMetricsAggregationBuilder;
 
 /**
  *
  */
-public class ValueCountBuilder extends MetricsAggregationBuilder<ValueCountBuilder> {
-
-    private String field;
+public class ValueCountBuilder extends ValuesSourceMetricsAggregationBuilder<ValueCountBuilder> {
 
     public ValueCountBuilder(String name) {
         super(name, InternalValueCount.TYPE.name());
     }
 
-    public ValueCountBuilder field(String field) {
-        this.field = field;
-        return this;
-    }
-
-    @Override
-    protected void internalXContent(XContentBuilder builder, Params params) throws IOException {
-        if (field != null) {
-            builder.field("field", field);
-        }
-    }
 }
diff --git a/src/main/java/org/elasticsearch/search/aggregations/metrics/valuecount/ValueCountParser.java b/src/main/java/org/elasticsearch/search/aggregations/metrics/valuecount/ValueCountParser.java
@@ -30,6 +30,7 @@
 import org.elasticsearch.search.internal.SearchContext;
 
 import java.io.IOException;
+import java.util.Map;
 
 /**
  *
@@ -47,6 +48,10 @@ public AggregatorFactory parse(String aggregationName, XContentParser parser, Se
         ValuesSourceConfig<BytesValuesSource> config = new ValuesSourceConfig<BytesValuesSource>(BytesValuesSource.class);
 
         String field = null;
+        String script = null;
+        String scriptLang = null;
+        Map<String, Object> scriptParams = null;
+        boolean assumeUnique = false;
 
         XContentParser.Token token;
         String currentFieldName = null;
@@ -56,14 +61,36 @@ public AggregatorFactory parse(String aggregationName, XContentParser parser, Se
             } else if (token == XContentParser.Token.VALUE_STRING) {
                 if ("field".equals(currentFieldName)) {
                     field = parser.text();
+                } else if ("script".equals(currentFieldName)) {
+                    script = parser.text();
+                } else if ("lang".equals(currentFieldName)) {
+                    scriptLang = parser.text();
                 } else {
                     throw new SearchParseException(context, "Unknown key for a " + token + " in [" + aggregationName + "]: [" + currentFieldName + "].");
                 }
+            } else if (token == XContentParser.Token.VALUE_BOOLEAN) {
+                if ("script_values_unique".equals(currentFieldName)) {
+                    assumeUnique = parser.booleanValue();
+                } else {
+                    throw new SearchParseException(context, "Unknown key for a " + token + " in [" + aggregationName + "]: [" + currentFieldName + "].");
+                }
+            } else if (token == XContentParser.Token.START_OBJECT) {
+                if ("params".equals(currentFieldName)) {
+                    scriptParams = parser.map();
+                }
             } else {
                 throw new SearchParseException(context, "Unexpected token " + token + " in [" + aggregationName + "].");
             }
         }
 
+        if (script != null) {
+            config.script(context.scriptService().search(context.lookup(), scriptLang, script, scriptParams));
+        }
+
+        if (!assumeUnique) {
+            config.ensureUnique(true);
+        }
+
         if (field == null) {
             return new ValueCountAggregator.Factory(aggregationName, config);
         }

diff --git a/src/test/java/org/elasticsearch/search/aggregations/metrics/ValueCountTests.java b/src/test/java/org/elasticsearch/search/aggregations/metrics/ValueCountTests.java
@@ -23,7 +23,6 @@
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.search.aggregations.metrics.valuecount.ValueCount;
 import org.elasticsearch.test.ElasticsearchIntegrationTest;
-import org.elasticsearch.test.junit.annotations.TestLogging;
 import org.junit.Before;
 import org.junit.Test;
 
@@ -37,7 +36,7 @@
  *
  */
 public class ValueCountTests extends ElasticsearchIntegrationTest {
-    
+
     @Override
     public Settings indexSettings() {
         return ImmutableSettings.builder()
@@ -124,4 +123,80 @@ public void multiValuedField() throws Exception {
         assertThat(valueCount.getName(), equalTo("count"));
         assertThat(valueCount.getValue(), equalTo(20l));
     }
+
+    @Test
+    public void singleValuedScript() throws Exception {
+        SearchResponse searchResponse = client().prepareSearch("idx")
+                .setQuery(matchAllQuery())
+                .addAggregation(count("count").script("doc['value'].value"))
+                .execute().actionGet();
+
+        assertThat(searchResponse.getHits().getTotalHits(), equalTo(10l));
+
+        ValueCount valueCount = searchResponse.getAggregations().get("count");
+        assertThat(valueCount, notNullValue());
+        assertThat(valueCount.getName(), equalTo("count"));
+        assertThat(valueCount.getValue(), equalTo(10l));
+    }
+
+    @Test
+    public void multiValuedScript() throws Exception {
+        SearchResponse searchResponse = client().prepareSearch("idx")
+                .setQuery(matchAllQuery())
+                .addAggregation(count("count").script("doc['values'].values"))
+                .execute().actionGet();
+
+        assertThat(searchResponse.getHits().getTotalHits(), equalTo(10l));
+
+        ValueCount valueCount = searchResponse.getAggregations().get("count");
+        assertThat(valueCount, notNullValue());
+        assertThat(valueCount.getName(), equalTo("count"));
+        assertThat(valueCount.getValue(), equalTo(20l));
+    }
+
+    @Test
+    public void singleValuedScriptWithParams() throws Exception {
+        SearchResponse searchResponse = client().prepareSearch("idx")
+                .setQuery(matchAllQuery())
+                .addAggregation(count("count").script("doc[s].value").param("s", "value"))
+                .execute().actionGet();
+
+        assertThat(searchResponse.getHits().getTotalHits(), equalTo(10l));
+
+        ValueCount valueCount = searchResponse.getAggregations().get("count");
+        assertThat(valueCount, notNullValue());
+        assertThat(valueCount.getName(), equalTo("count"));
+        assertThat(valueCount.getValue(), equalTo(10l));
+    }
+
+    @Test
+    public void multiValuedScriptWithParams() throws Exception {
+        SearchResponse searchResponse = client().prepareSearch("idx")
+                .setQuery(matchAllQuery())
+                .addAggregation(count("count").script("doc[s].values").param("s", "values"))
+                .execute().actionGet();
+
+        assertThat(searchResponse.getHits().getTotalHits(), equalTo(10l));
+
+        ValueCount valueCount = searchResponse.getAggregations().get("count");
+        assertThat(valueCount, notNullValue());
+        assertThat(valueCount.getName(), equalTo("count"));
+        assertThat(valueCount.getValue(), equalTo(20l));
+    }
+
+    @Test
+    public void deduplication() throws Exception {
+        SearchResponse searchResponse = client().prepareSearch("idx")
+                .setQuery(matchAllQuery())
+                .addAggregation(count("count").script("doc['values'].values + [5L]"))
+                .execute().actionGet();
+
+        assertThat(searchResponse.getHits().getTotalHits(), equalTo(10l));
+
+        ValueCount valueCount = searchResponse.getAggregations().get("count");
+        assertThat(valueCount, notNullValue());
+        assertThat(valueCount.getName(), equalTo("count"));
+        assertThat(valueCount.getValue(), equalTo(28l));
+    }
+
 }