Improve handling of incompatible after key types in composite aggs (#70839) (#70939)

Composite aggs on a field that have the same name but different types in
different indices can produce confusing errors. This commit add more
information to the error message making debugging of such issues easier.

Fixes #70480

Author: imotov

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

@@ -32,6 +32,7 @@
 import org.apache.lucene.search.comparators.LongComparator;
 import org.apache.lucene.util.Bits;
 import org.apache.lucene.util.RoaringDocIdSet;
+import org.elasticsearch.ElasticsearchParseException;
 import org.elasticsearch.common.lease.Releasables;
 import org.elasticsearch.index.IndexSortConfig;
 import org.elasticsearch.search.DocValueFormat;
@@ -103,7 +104,15 @@ final class CompositeAggregator extends BucketsAggregator {
                 this::addRequestCircuitBreakerBytes
             );
         }
-        this.queue = new CompositeValuesCollectorQueue(context.bigArrays(), sources, size, rawAfterKey);
+        this.queue = new CompositeValuesCollectorQueue(context.bigArrays(), sources, size);
+        if (rawAfterKey != null) {
+            try {
+                this.queue.setAfterKey(rawAfterKey);
+            } catch (IllegalArgumentException ex) {
+                throw new ElasticsearchParseException("Cannot set after key in the composite aggregation [" + name + "] - " +
+                    ex.getMessage(), ex);
+            }
+        }
         this.rawAfterKey = rawAfterKey;
     }
 

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

@@ -62,22 +62,30 @@ public int hashCode() {
      *
      * @param sources The list of {@link CompositeValuesSourceConfig} to build the composite buckets.
      * @param size The number of composite buckets to keep.
-     * @param afterKey composite key
      */
-    CompositeValuesCollectorQueue(BigArrays bigArrays, SingleDimensionValuesSource<?>[] sources, int size, CompositeKey afterKey) {
+    CompositeValuesCollectorQueue(BigArrays bigArrays, SingleDimensionValuesSource<?>[] sources, int size) {
         super(size);
         this.bigArrays = bigArrays;
         this.maxSize = size;
         this.arrays = sources;
         this.map = new HashMap<>(size);
-        if (afterKey != null) {
-            assert afterKey.size() == sources.length;
-            afterKeyIsSet = true;
-            for (int i = 0; i < afterKey.size(); i++) {
-                sources[i].setAfter(afterKey.get(i));
+        this.docCounts = bigArrays.newLongArray(1, false);
+    }
+
+    /**
+     * Sets after key
+     * @param afterKey composite key
+     */
+    public void setAfterKey(CompositeKey afterKey) {
+        assert afterKey.size() == arrays.length;
+        afterKeyIsSet = true;
+        for (int i = 0; i < afterKey.size(); i++) {
+            try {
+                arrays[i].setAfter(afterKey.get(i));
+            } catch (IllegalArgumentException ex) {
+                throw new IllegalArgumentException("incompatible value in the position " + i + ": " + ex.getMessage(), ex);
             }
         }
-        this.docCounts = bigArrays.newLongArray(1, false);
     }
 
     @Override

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

@@ -938,6 +938,21 @@ public void testWithKeywordAndLong() throws Exception {
                 assertEquals(1L, result.getBuckets().get(1).getDocCount());
             }
         );
+
+        Exception exc = expectThrows(ElasticsearchParseException.class,
+            () -> testSearchCase(Arrays.asList(new MatchAllDocsQuery(), new DocValuesFieldExistsQuery("date")), Collections.emptyList(),
+                () -> new CompositeAggregationBuilder("test",
+                    Arrays.asList(
+                        new TermsValuesSourceBuilder("keyword").field("keyword"),
+                        new TermsValuesSourceBuilder("long").field("long")
+                    )
+                ).aggregateAfter(createAfterKey("keyword", 0L, "long", 100L)
+                ),
+                (result) -> {
+                }
+            ));
+        assertThat(exc.getMessage(), containsString("Cannot set after key in the composite aggregation [test] - incompatible value in " +
+            "the position 0: invalid value, expected string, got Long"));
     }
 
     public void testWithKeywordAndLongDesc() throws Exception {

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

@@ -288,7 +288,10 @@ private void testRandomCase(boolean forceMerge,
                 CompositeKey last = null;
                 while (pos < size) {
                     final CompositeValuesCollectorQueue queue =
-                        new CompositeValuesCollectorQueue(BigArrays.NON_RECYCLING_INSTANCE, sources, size, last);
+                        new CompositeValuesCollectorQueue(BigArrays.NON_RECYCLING_INSTANCE, sources, size);
+                    if (last != null) {
+                        queue.setAfterKey(last);
+                    }
                     final SortedDocsProducer docsProducer = sources[0].createSortedDocsProducerOrNull(reader, new MatchAllDocsQuery());
                     for (LeafReaderContext leafReaderContext : reader.leaves()) {
                         if (docsProducer != null && withProducer) {