Fix composite agg sort bug

When an composite aggregation is run against an index with a sort that
*starts* with the "source" fields from the composite but has additional
fields it'd blow up in while trying to decide if it could use the sort.
This changes it to decide that it *can* use the sort.

Closes elastic#52480

Author: nik9000

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

@@ -850,3 +850,83 @@ setup:
   - length: { aggregations.test.buckets: 1 }
   - match: { aggregations.test.buckets.0.key.date: "2017-10-21T00:00:00.000-02:00" }
   - match: { aggregations.test.buckets.0.doc_count: 2 }
+
+---
+"Terms source from sorted":
+  - do:
+        indices.create:
+          index: sorted_test
+          body:
+              settings:
+                sort.field: keyword
+              mappings:
+                properties:
+                  keyword:
+                    type: keyword
+                  long:
+                    type: long
+
+
+  - do:
+      index:
+        index:   sorted_test
+        id:      2
+        refresh: true
+        body:  { "keyword": "foo", "long": 1 }
+
+  - do:
+      search:
+        index: sorted_test
+        body:
+          aggregations:
+            test:
+              composite:
+                sources:
+                  - keyword:
+                      terms:
+                        field: keyword
+
+  - match: {hits.total.value: 1}
+  - length: { aggregations.test.buckets: 1 }
+  - match: { aggregations.test.buckets.0.key.keyword: "foo" }
+  - match: { aggregations.test.buckets.0.doc_count: 1 }
+
+---
+"Terms source from part of sorted":
+  - do:
+        indices.create:
+          index: sorted_test
+          body:
+              settings:
+                sort.field: [keyword, long]
+              mappings:
+                properties:
+                  keyword:
+                    type: keyword
+                  long:
+                    type: long
+
+
+  - do:
+      index:
+        index:   sorted_test
+        id:      2
+        refresh: true
+        body:  { "keyword": "foo", "long": 1 }
+
+  - do:
+      search:
+        index: sorted_test
+        body:
+          aggregations:
+            test:
+              composite:
+                sources:
+                  - keyword:
+                      terms:
+                        field: keyword
+
+  - match: {hits.total.value: 1}
+  - length: { aggregations.test.buckets: 1 }
+  - match: { aggregations.test.buckets.0.key.keyword: "foo" }
+  - match: { aggregations.test.buckets.0.doc_count: 1 }

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

@@ -202,7 +202,8 @@ private Sort buildIndexSortPrefix(LeafReaderContext context) throws IOException
             return null;
         }
         List<SortField> sortFields = new ArrayList<>();
-        for (int i = 0; i < indexSort.getSort().length; i++) {
+        int max = Math.min(indexSort.getSort().length, sourceConfigs.length);
+        for (int i = 0; i < max; i++) {
             CompositeValuesSourceConfig sourceConfig = sourceConfigs[i];
             SingleDimensionValuesSource<?> source = sources[i];
             SortField indexSortField = indexSort.getSort()[i];

server/src/test/java/org/elasticsearch/search/aggregations/bucket/composite/CompositeAggregatorTests.java

@@ -2077,40 +2077,52 @@ private static long asLong(String dateTime) {
 
     private static Sort buildIndexSort(List<CompositeValuesSourceBuilder<?>> sources, Map<String, MappedFieldType> fieldTypes) {
         List<SortField> sortFields = new ArrayList<>();
+        Map<String, MappedFieldType> remainingFieldTypes = new HashMap<>(fieldTypes);
         for (CompositeValuesSourceBuilder<?> source : sources) {
-            MappedFieldType type = fieldTypes.get(source.field());
-            if (type instanceof KeywordFieldMapper.KeywordFieldType) {
-                sortFields.add(new SortedSetSortField(type.name(), false));
-            } else if (type instanceof DateFieldMapper.DateFieldType) {
-                sortFields.add(new SortedNumericSortField(type.name(), SortField.Type.LONG, false));
-            } else if (type instanceof NumberFieldMapper.NumberFieldType) {
-                boolean comp = false;
-                switch (type.typeName()) {
-                    case "byte":
-                    case "short":
-                    case "integer":
-                        comp = true;
-                        sortFields.add(new SortedNumericSortField(type.name(), SortField.Type.INT, false));
-                        break;
-
-                    case "long":
-                        sortFields.add(new SortedNumericSortField(type.name(), SortField.Type.LONG, false));
-                        break;
-
-                    case "float":
-                    case "double":
-                        comp = true;
-                        sortFields.add(new SortedNumericSortField(type.name(), SortField.Type.DOUBLE, false));
-                        break;
-
-                    default:
-                        break;
-                }
-                if (comp == false) {
-                    break;
-                }
+            MappedFieldType type = fieldTypes.remove(source.field());
+            remainingFieldTypes.remove(source.field());
+            SortField sortField = sortFieldFrom(type);
+            if (sortField == null) {
+                break;
+            }
+            sortFields.add(sortField);
+            if (sortField instanceof SortedNumericSortField && ((SortedNumericSortField) sortField).getType() == SortField.Type.LONG) {
+                break;
+            }
+        }
+        while (remainingFieldTypes.size() > 0 && randomBoolean()) {
+            // Add extra unused sorts
+            List<String> fields = new ArrayList<>(remainingFieldTypes.keySet());
+            Collections.sort(fields);
+            String field = fields.get(between(0, fields.size() - 1));
+            SortField sortField = sortFieldFrom(remainingFieldTypes.remove(field));
+            if (sortField != null) {
+                sortFields.add(sortField);
+            }
+        }
+        return sortFields.size() > 0 ? new Sort(sortFields.toArray(SortField[]::new)) : null;
+    }
+    
+    private static SortField sortFieldFrom(MappedFieldType type) {
+        if (type instanceof KeywordFieldMapper.KeywordFieldType) {
+            return new SortedSetSortField(type.name(), false);
+        } else if (type instanceof DateFieldMapper.DateFieldType) {
+            return new SortedNumericSortField(type.name(), SortField.Type.LONG, false);
+        } else if (type instanceof NumberFieldMapper.NumberFieldType) {
+            switch (type.typeName()) {
+                case "byte":
+                case "short":
+                case "integer":
+                    return new SortedNumericSortField(type.name(), SortField.Type.INT, false);
+                case "long":
+                    return new SortedNumericSortField(type.name(), SortField.Type.LONG, false);
+                case "float":
+                case "double":
+                    return new SortedNumericSortField(type.name(), SortField.Type.DOUBLE, false);
+                default:
+                    return null;
             }
         }
-        return sortFields.size() > 0 ? new Sort(sortFields.toArray(new SortField[0])) : null;
+        return null;
     }
 }