Optimize the composite aggregation for match_all and range queries

docs/reference/aggregations/bucket/composite-aggregation.asciidoc

@@ -545,88 +545,3 @@ GET /_search
 }
 --------------------------------------------------
 // TESTRESPONSE[s/\.\.\.//]
-
-==== Index sorting
-
-By default this aggregation runs on every document that match the query.
-Though if the index sort matches the composite sort this aggregation can optimize
-the execution and can skip documents that contain composite buckets that would not
-be part of the response.
-
-For instance the following aggregations:
-
-[source,js]
---------------------------------------------------
-GET /_search
-{
-    "aggs" : {
-        "my_buckets": {
-            "composite" : {
-                "size": 2,
-                "sources" : [
-                    { "date": { "date_histogram": { "field": "timestamp", "interval": "1d", "order": "asc" } } },
-                    { "product": { "terms": { "field": "product", "order": "asc" } } }
-                ]
-            }
-        }
-    }
-}
---------------------------------------------------
-// CONSOLE
-
-\... is much faster on an index that uses the following sort:
-
-[source,js]
---------------------------------------------------
-PUT twitter
-{
-    "settings" : {
-        "index" : {
-            "sort.field" : ["timestamp", "product"],
-            "sort.order" : ["asc", "asc"]
-        }
-    },
-    "mappings": {
-        "sales": {
-            "properties": {
-                "timestamp": {
-                    "type": "date"
-                },
-                "product": {
-                    "type": "keyword"
-                }
-            }
-        }
-    }
-}
---------------------------------------------------
-// CONSOLE
-
-WARNING: The optimization takes effect only if the fields used for sorting are single-valued and follow
-the same order as the aggregation (`desc` or `asc`).
-
-If only the aggregation results are needed it is also better to set the size of the query to 0
-and `track_total_hits` to false in order to remove other slowing factors:
-
-[source,js]
---------------------------------------------------
-GET /_search
-{
-    "size": 0,
-    "track_total_hits": false,
-    "aggs" : {
-        "my_buckets": {
-            "composite" : {
-                "size": 2,
-                "sources" : [
-                    { "date": { "date_histogram": { "field": "timestamp", "interval": "1d" } } },
-                    { "product": { "terms": { "field": "product" } } }
-                ]
-            }
-        }
-    }
-}
---------------------------------------------------
-// CONSOLE
-
-See <<index-modules-index-sorting, index sorting>> for more details.

rest-api-spec/src/main/resources/rest-api-spec/test/search.aggregation/230_composite.yml

@@ -99,6 +99,7 @@ setup:
   - do:
       search:
         index: test
+        allow_partial_search_results: false
         body:
           aggregations:
             test:

server/src/main/java/org/elasticsearch/search/aggregations/bucket/composite/BinaryValuesSource.java

@@ -0,0 +1,137 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch 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.aggregations.bucket.composite;
+
+import org.apache.lucene.index.IndexReader;
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.search.MatchAllDocsQuery;
+import org.apache.lucene.search.Query;
+import org.apache.lucene.util.BytesRef;
+import org.elasticsearch.common.CheckedFunction;
+import org.elasticsearch.index.fielddata.SortedBinaryDocValues;
+import org.elasticsearch.index.mapper.MappedFieldType;
+import org.elasticsearch.search.aggregations.LeafBucketCollector;
+
+import java.io.IOException;
+
+/**
+ * A {@link SingleDimensionValuesSource} for binary source ({@link BytesRef}).
+ */
+class BinaryValuesSource extends SingleDimensionValuesSource<BytesRef> {
+    private final CheckedFunction<LeafReaderContext, SortedBinaryDocValues, IOException> docValuesFunc;
+    private final BytesRef[] values;
+    private BytesRef currentValue;
+
+    BinaryValuesSource(MappedFieldType fieldType, CheckedFunction<LeafReaderContext, SortedBinaryDocValues, IOException> docValuesFunc,
+                       int size, int reverseMul) {
+        super(fieldType, size, reverseMul);
+        this.docValuesFunc = docValuesFunc;
+        this.values = new BytesRef[size];
+    }
+
+    @Override
+    String type() {
+        return "binary";
+    }
+
+    @Override
+    public void copyCurrent(int slot) {
+        values[slot] = BytesRef.deepCopyOf(currentValue);
+    }
+
+    @Override
+    public int compare(int from, int to) {
+        return compareValues(values[from], values[to]);
+    }
+
+    @Override
+    int compareCurrent(int slot) {
+        return compareValues(currentValue, values[slot]);
+    }
+
+    @Override
+    int compareCurrentWithAfter() {
+        return compareValues(currentValue, afterValue);
+    }
+
+    int compareValues(BytesRef v1, BytesRef v2) {
+        return v1.compareTo(v2) * reverseMul;
+    }
+
+    @Override
+    void setAfter(Comparable<?> value) {
+        if (value.getClass() == BytesRef.class) {
+            afterValue = (BytesRef) value;
+        } else if (value.getClass() == String.class) {
+            afterValue = new BytesRef((String) value);
+        } else {
+            throw new IllegalArgumentException("invalid value, expected string, got " + value.getClass().getSimpleName());
+        }
+    }
+
+    @Override
+    BytesRef toComparable(int slot) {
+        return values[slot];
+    }
+
+    @Override
+    LeafBucketCollector getLeafCollector(LeafReaderContext context, LeafBucketCollector next) throws IOException {
+        final SortedBinaryDocValues dvs = docValuesFunc.apply(context);
+        return new LeafBucketCollector() {
+            @Override
+            public void collect(int doc, long bucket) throws IOException {
+                if (dvs.advanceExact(doc)) {
+                    int num = dvs.docValueCount();
+                    for (int i = 0; i < num; i++) {
+                        currentValue = dvs.nextValue();
+                        next.collect(doc, bucket);
+                    }
+                }
+            }
+        };
+    }
+
+    @Override
+    LeafBucketCollector getLeafCollector(Comparable<?> value, LeafReaderContext context, LeafBucketCollector next) {
+        if (value.getClass() != BytesRef.class) {
+            throw new IllegalArgumentException("Expected BytesRef, got " + value.getClass());
+        }
+        final BytesRef filterValue = (BytesRef) value;
+        return new LeafBucketCollector() {
+            @Override
+            public void collect(int doc, long bucket) throws IOException {
+                currentValue = filterValue;
+                next.collect(doc, bucket);
+            }
+        };
+    }
+
+    @Override
+    SortedDocsProducer createSortedDocsProducerOrNull(IndexReader reader, Query query) {
+        if (checkIfSortedDocsIsApplicable(reader, fieldType) == false ||
+                (query != null && query.getClass() != MatchAllDocsQuery.class)) {
+            return null;
+        }
+        return new TermsSortedDocsProducer(fieldType.name());
+    }
+
+    @Override
+    public void close() {}
+}

server/src/main/java/org/elasticsearch/search/aggregations/bucket/composite/CompositeAggregationBuilder.java

@@ -19,16 +19,12 @@
 
 package org.elasticsearch.search.aggregations.bucket.composite;
 
-import org.apache.lucene.search.Sort;
-import org.apache.lucene.search.SortField;
 import org.elasticsearch.common.ParseField;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.xcontent.ObjectParser;
 import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.common.xcontent.XContentParser;
-import org.elasticsearch.index.IndexSortConfig;
-import org.elasticsearch.index.query.QueryShardContext;
 import org.elasticsearch.search.aggregations.AbstractAggregationBuilder;
 import org.elasticsearch.search.aggregations.AggregationBuilder;
 import org.elasticsearch.search.aggregations.AggregatorFactories;
@@ -154,16 +150,9 @@ protected AggregatorFactory<?> doBuild(SearchContext context, AggregatorFactory<
         if (parent != null) {
             throw new IllegalArgumentException("[composite] aggregation cannot be used with a parent aggregation");
         }
-        final QueryShardContext shardContext = context.getQueryShardContext();
         CompositeValuesSourceConfig[] configs = new CompositeValuesSourceConfig[sources.size()];
-        SortField[] sortFields = new SortField[configs.length];
-        IndexSortConfig indexSortConfig = shardContext.getIndexSettings().getIndexSortConfig();
-        if (indexSortConfig.hasIndexSort()) {
-            Sort sort = indexSortConfig.buildIndexSort(shardContext::fieldMapper, shardContext::getForField);
-            System.arraycopy(sort.getSort(), 0, sortFields, 0, sortFields.length);
-        }
         for (int i = 0; i < configs.length; i++) {
-            configs[i] = sources.get(i).build(context, i, configs.length, sortFields[i]);
+            configs[i] = sources.get(i).build(context);
             if (configs[i].valuesSource().needsScores()) {
                 throw new IllegalArgumentException("[sources] cannot access _score");
             }