Handle IndexOrDocValuesQuery in composite aggregation (#35392)

The `composite` aggregation can optimize its execution when the query
is a `match_all` or a `range` over the field that is used in the first source
of the aggregation. However we only check for instances of `PointRangeQuery` whereas
the range query builder creates an  `IndexOrDocValuesQuery`. This means that
today the optimization does not apply to `range` query even if the code could handle it.
This change fixes this issue by extracting the index query inside `IndexOrDocValuesQuery`.

Author: jimczi

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

@@ -24,6 +24,8 @@
 import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.SortedNumericDocValues;
+import org.apache.lucene.search.BoostQuery;
+import org.apache.lucene.search.IndexOrDocValuesQuery;
 import org.apache.lucene.search.MatchAllDocsQuery;
 import org.apache.lucene.search.PointRangeQuery;
 import org.apache.lucene.search.Query;
@@ -178,8 +180,19 @@ public void collect(int doc, long bucket) throws IOException {
         };
     }
 
+    static Query extractQuery(Query query) {
+        if (query instanceof BoostQuery) {
+            return extractQuery(((BoostQuery) query).getQuery());
+        } else if (query instanceof IndexOrDocValuesQuery) {
+            return extractQuery(((IndexOrDocValuesQuery) query).getIndexQuery());
+        } else {
+            return query;
+        }
+    }
+
     @Override
     SortedDocsProducer createSortedDocsProducerOrNull(IndexReader reader, Query query) {
+        query = extractQuery(query);
         if (checkIfSortedDocsIsApplicable(reader, fieldType) == false ||
                 (query != null &&
                     query.getClass() != MatchAllDocsQuery.class &&

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

@@ -22,6 +22,8 @@
 import org.apache.lucene.document.LongPoint;
 import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.index.Term;
+import org.apache.lucene.search.BoostQuery;
+import org.apache.lucene.search.IndexOrDocValuesQuery;
 import org.apache.lucene.search.MatchAllDocsQuery;
 import org.apache.lucene.search.TermQuery;
 import org.elasticsearch.common.util.BigArrays;
@@ -175,6 +177,10 @@ public void testNumericSorted() {
                 assertNotNull(source.createSortedDocsProducerOrNull(reader, new MatchAllDocsQuery()));
                 assertNotNull(source.createSortedDocsProducerOrNull(reader, null));
                 assertNotNull(source.createSortedDocsProducerOrNull(reader, LongPoint.newRangeQuery("number", 0, 1)));
+                assertNotNull(source.createSortedDocsProducerOrNull(reader, new IndexOrDocValuesQuery(
+                    LongPoint.newRangeQuery("number", 0, 1), new MatchAllDocsQuery())));
+                assertNotNull(source.createSortedDocsProducerOrNull(reader, new BoostQuery(new IndexOrDocValuesQuery(
+                    LongPoint.newRangeQuery("number", 0, 1), new MatchAllDocsQuery()), 2.0f)));
                 assertNull(source.createSortedDocsProducerOrNull(reader, new TermQuery(new Term("keyword", "toto)"))));
 
                 LongValuesSource sourceWithMissing = new LongValuesSource(