value_count Aggregation optimization (backport of #54854)

docs/reference/release-notes.asciidoc

@@ -6,6 +6,7 @@
 
 This section summarizes the changes in each release.
 
+* <<release-notes-7.8.0>>
 * <<release-notes-7.7.0>>
 * <<release-notes-7.6.2>>
 * <<release-notes-7.6.1>>
@@ -32,6 +33,7 @@ This section summarizes the changes in each release.
 
 --
 
+include::release-notes/7.8.asciidoc[]
 include::release-notes/7.7.asciidoc[]
 include::release-notes/7.6.asciidoc[]
 include::release-notes/7.5.asciidoc[]

docs/reference/release-notes/7.8.asciidoc

@@ -0,0 +1,14 @@
+[[release-notes-7.8.0]]
+== {es} version 7.8.0
+
+coming[7.8.0]
+
+[[breaking-7.8.0]]
+[float]
+=== Breaking changes
+
+Search::
+* Scripts used in `value_count` will now receive a number if they are counting
+  a numeric field and a `GeoPoint` if they are counting a `geo_point` fields.
+  They used to always receive the `String` representation of those values.
+  {pull}54854[#54854]

server/src/main/java/org/elasticsearch/search/aggregations/metrics/ValueCountAggregator.java

@@ -19,9 +19,11 @@
 package org.elasticsearch.search.aggregations.metrics;
 
 import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.index.SortedNumericDocValues;
 import org.elasticsearch.common.lease.Releasables;
 import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.common.util.LongArray;
+import org.elasticsearch.index.fielddata.MultiGeoPointValues;
 import org.elasticsearch.index.fielddata.SortedBinaryDocValues;
 import org.elasticsearch.search.aggregations.Aggregator;
 import org.elasticsearch.search.aggregations.InternalAggregation;
@@ -62,6 +64,34 @@ public LeafBucketCollector getLeafCollector(LeafReaderContext ctx,
             return LeafBucketCollector.NO_OP_COLLECTOR;
         }
         final BigArrays bigArrays = context.bigArrays();
+
+        if (valuesSource instanceof ValuesSource.Numeric) {
+            final SortedNumericDocValues values = ((ValuesSource.Numeric)valuesSource).longValues(ctx);
+            return new LeafBucketCollectorBase(sub, values) {
+
+                @Override
+                public void collect(int doc, long bucket) throws IOException {
+                    counts = bigArrays.grow(counts, bucket + 1);
+                    if (values.advanceExact(doc)) {
+                        counts.increment(bucket, values.docValueCount());
+                    }
+                }
+            };
+        }
+        if (valuesSource instanceof ValuesSource.Bytes.GeoPoint) {
+            MultiGeoPointValues values = ((ValuesSource.GeoPoint)valuesSource).geoPointValues(ctx);
+            return new LeafBucketCollectorBase(sub, null) {
+
+                @Override
+                public void collect(int doc, long bucket) throws IOException {
+                    counts = bigArrays.grow(counts, bucket + 1);
+                    if (values.advanceExact(doc)) {
+                        counts.increment(bucket, values.docValueCount());
+                    }
+                }
+            };
+        }
+        // The following is default collector. Including the keyword FieldType
         final SortedBinaryDocValues values = valuesSource.bytesValues(ctx);
         return new LeafBucketCollectorBase(sub, values) {
 

server/src/test/java/org/elasticsearch/search/aggregations/metrics/ValueCountAggregatorTests.java

@@ -20,10 +20,14 @@
 package org.elasticsearch.search.aggregations.metrics;
 
 import org.apache.lucene.document.BinaryDocValuesField;
+import org.apache.lucene.document.Document;
+import org.apache.lucene.document.DoubleDocValuesField;
 import org.apache.lucene.document.IntPoint;
+import org.apache.lucene.document.LatLonDocValuesField;
 import org.apache.lucene.document.NumericDocValuesField;
 import org.apache.lucene.document.SortedDocValuesField;
 import org.apache.lucene.document.SortedNumericDocValuesField;
+import org.apache.lucene.document.SortedSetDocValuesField;
 import org.apache.lucene.index.DirectoryReader;
 import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.index.RandomIndexWriter;
@@ -35,6 +39,7 @@
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.common.CheckedConsumer;
 import org.elasticsearch.common.geo.GeoPoint;
+import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.index.mapper.BooleanFieldMapper;
 import org.elasticsearch.index.mapper.DateFieldMapper;
 import org.elasticsearch.index.mapper.GeoPointFieldMapper;
@@ -44,22 +49,103 @@
 import org.elasticsearch.index.mapper.NumberFieldMapper;
 import org.elasticsearch.index.mapper.RangeFieldMapper;
 import org.elasticsearch.index.mapper.RangeType;
+import org.elasticsearch.script.MockScriptEngine;
+import org.elasticsearch.script.Script;
+import org.elasticsearch.script.ScriptEngine;
+import org.elasticsearch.script.ScriptModule;
+import org.elasticsearch.script.ScriptService;
+import org.elasticsearch.script.ScriptType;
+import org.elasticsearch.search.aggregations.AggregationBuilder;
 import org.elasticsearch.search.aggregations.AggregatorTestCase;
 import org.elasticsearch.search.aggregations.support.AggregationInspectionHelper;
+import org.elasticsearch.search.aggregations.support.CoreValuesSourceType;
 import org.elasticsearch.search.aggregations.support.ValueType;
+import org.elasticsearch.search.aggregations.support.ValuesSourceType;
 
 import java.io.IOException;
 import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
 import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import java.util.function.Consumer;
+import java.util.function.Function;
 
 import static java.util.Collections.singleton;
 
 public class ValueCountAggregatorTests extends AggregatorTestCase {
 
     private static final String FIELD_NAME = "field";
 
+    private static final String STRING_VALUE_SCRIPT = "string_value";
+    private static final String NUMBER_VALUE_SCRIPT = "number_value";
+    private static final String SINGLE_SCRIPT = "single";
+
+    @Override
+    protected AggregationBuilder createAggBuilderForTypeTest(MappedFieldType fieldType, String fieldName) {
+        return new ValueCountAggregationBuilder("foo", null).field(fieldName);
+    }
+
+    @Override
+    protected List<ValuesSourceType> getSupportedValuesSourceTypes() {
+        return Arrays.asList(
+            CoreValuesSourceType.NUMERIC,
+            CoreValuesSourceType.BYTES,
+            CoreValuesSourceType.GEOPOINT,
+            CoreValuesSourceType.RANGE
+        );
+    }
+
+    @Override
+    protected ScriptService getMockScriptService() {
+        Map<String, Function<Map<String, Object>, Object>> scripts = new HashMap<>();
+
+        scripts.put(STRING_VALUE_SCRIPT, vars -> (Double.valueOf((String) vars.get("_value")) + 1));
+        scripts.put(NUMBER_VALUE_SCRIPT, vars -> (((Number) vars.get("_value")).doubleValue() + 1));
+        scripts.put(SINGLE_SCRIPT, vars -> 1);
+
+        MockScriptEngine scriptEngine = new MockScriptEngine(MockScriptEngine.NAME,
+            scripts,
+            Collections.emptyMap());
+        Map<String, ScriptEngine> engines = Collections.singletonMap(scriptEngine.getType(), scriptEngine);
+
+        return new ScriptService(Settings.EMPTY, engines, ScriptModule.CORE_CONTEXTS);
+    }
+
+
+    public void testGeoField() throws IOException {
+        testCase(new MatchAllDocsQuery(), ValueType.GEOPOINT, iw -> {
+            for (int i = 0; i < 10; i++) {
+                Document document = new Document();
+                document.add(new LatLonDocValuesField("field", 10, 10));
+                iw.addDocument(document);
+            }
+        }, count -> assertEquals(10L, count.getValue()));
+    }
+
+    public void testDoubleField() throws IOException {
+        testCase(new MatchAllDocsQuery(), ValueType.DOUBLE, iw -> {
+            for (int i = 0; i < 15; i++) {
+                Document document = new Document();
+                document.add(new DoubleDocValuesField(FIELD_NAME, 23D));
+                iw.addDocument(document);
+            }
+        }, count -> assertEquals(15L, count.getValue()));
+    }
+
+    public void testKeyWordField() throws IOException {
+        testCase(new MatchAllDocsQuery(), ValueType.STRING, iw -> {
+            for (int i = 0; i < 20; i++) {
+                Document document = new Document();
+                document.add(new SortedSetDocValuesField(FIELD_NAME, new BytesRef("stringValue")));
+                document.add(new SortedSetDocValuesField(FIELD_NAME, new BytesRef("string11Value")));
+                iw.addDocument(document);
+            }
+        }, count -> assertEquals(40L, count.getValue()));
+    }
+
     public void testNoDocs() throws IOException {
         for (ValueType valueType : ValueType.values()) {
             testCase(new MatchAllDocsQuery(), valueType, iw -> {
@@ -189,6 +275,105 @@ public void testRangeFieldValues() throws IOException {
         }, fieldType);
     }
 
+    public void testValueScriptNumber() throws IOException {
+        ValueCountAggregationBuilder aggregationBuilder = new ValueCountAggregationBuilder("name", null)
+            .field(FIELD_NAME)
+            .script(new Script(ScriptType.INLINE, MockScriptEngine.NAME, NUMBER_VALUE_SCRIPT, Collections.emptyMap()));
+
+        MappedFieldType fieldType = createMappedFieldType(ValueType.NUMERIC);
+        fieldType.setName(FIELD_NAME);
+        fieldType.setHasDocValues(true);
+
+        testCase(aggregationBuilder, new MatchAllDocsQuery(), iw -> {
+            iw.addDocument(singleton(new NumericDocValuesField(FIELD_NAME, 7)));
+            iw.addDocument(singleton(new NumericDocValuesField(FIELD_NAME, 8)));
+            iw.addDocument(singleton(new NumericDocValuesField(FIELD_NAME, 9)));
+        }, card -> {
+            assertEquals(3, card.getValue(), 0);
+            assertTrue(AggregationInspectionHelper.hasValue(card));
+        }, fieldType);
+    }
+
+    public void testSingleScriptNumber() throws IOException {
+        ValueCountAggregationBuilder aggregationBuilder = new ValueCountAggregationBuilder("name", null)
+            .field(FIELD_NAME);
+
+        MappedFieldType fieldType = createMappedFieldType(ValueType.NUMERIC);
+        fieldType.setName(FIELD_NAME);
+        fieldType.setHasDocValues(true);
+
+        testCase(aggregationBuilder, new MatchAllDocsQuery(), iw -> {
+            Document doc = new Document();
+            doc.add(new SortedNumericDocValuesField(FIELD_NAME, 7));
+            doc.add(new SortedNumericDocValuesField(FIELD_NAME, 7));
+            iw.addDocument(doc);
+
+            doc = new Document();
+            doc.add(new SortedNumericDocValuesField(FIELD_NAME, 8));
+            doc.add(new SortedNumericDocValuesField(FIELD_NAME, 8));
+            iw.addDocument(doc);
+
+            doc = new Document();
+            doc.add(new SortedNumericDocValuesField(FIELD_NAME, 1));
+            doc.add(new SortedNumericDocValuesField(FIELD_NAME, 1));
+            iw.addDocument(doc);
+        }, card -> {
+            // note: this is 6, even though the script returns a single value.  ValueCount does not de-dedupe
+            assertEquals(6, card.getValue(), 0);
+            assertTrue(AggregationInspectionHelper.hasValue(card));
+        }, fieldType);
+    }
+
+    public void testValueScriptString() throws IOException {
+        ValueCountAggregationBuilder aggregationBuilder = new ValueCountAggregationBuilder("name", null)
+            .field(FIELD_NAME)
+            .script(new Script(ScriptType.INLINE, MockScriptEngine.NAME, STRING_VALUE_SCRIPT, Collections.emptyMap()));
+
+        MappedFieldType fieldType = createMappedFieldType(ValueType.STRING);
+        fieldType.setName(FIELD_NAME);
+        fieldType.setHasDocValues(true);
+
+        testCase(aggregationBuilder, new MatchAllDocsQuery(), iw -> {
+            iw.addDocument(singleton(new SortedDocValuesField(FIELD_NAME, new BytesRef("1"))));
+            iw.addDocument(singleton(new SortedDocValuesField(FIELD_NAME, new BytesRef("2"))));
+            iw.addDocument(singleton(new SortedDocValuesField(FIELD_NAME, new BytesRef("3"))));
+        }, card -> {
+            assertEquals(3, card.getValue(), 0);
+            assertTrue(AggregationInspectionHelper.hasValue(card));
+        }, fieldType);
+    }
+
+    public void testSingleScriptString() throws IOException {
+        ValueCountAggregationBuilder aggregationBuilder = new ValueCountAggregationBuilder("name", null)
+            .field(FIELD_NAME);
+
+        MappedFieldType fieldType = createMappedFieldType(ValueType.STRING);
+        fieldType.setName(FIELD_NAME);
+        fieldType.setHasDocValues(true);
+
+        testCase(aggregationBuilder, new MatchAllDocsQuery(), iw -> {
+            Document doc = new Document();
+            // Note: unlike numerics, lucene de-dupes strings so we increment here
+            doc.add(new SortedSetDocValuesField(FIELD_NAME, new BytesRef("1")));
+            doc.add(new SortedSetDocValuesField(FIELD_NAME, new BytesRef("2")));
+            iw.addDocument(doc);
+
+            doc = new Document();
+            doc.add(new SortedSetDocValuesField(FIELD_NAME, new BytesRef("3")));
+            doc.add(new SortedSetDocValuesField(FIELD_NAME, new BytesRef("4")));
+            iw.addDocument(doc);
+
+            doc = new Document();
+            doc.add(new SortedSetDocValuesField(FIELD_NAME, new BytesRef("5")));
+            doc.add(new SortedSetDocValuesField(FIELD_NAME, new BytesRef("6")));
+            iw.addDocument(doc);
+        }, card -> {
+            // note: this is 6, even though the script returns a single value.  ValueCount does not de-dedupe
+            assertEquals(6, card.getValue(), 0);
+            assertTrue(AggregationInspectionHelper.hasValue(card));
+        }, fieldType);
+    }
+
     private void testCase(Query query,
                           ValueType valueType,
                           CheckedConsumer<RandomIndexWriter, IOException> indexer,