Handles exists query in composite aggs (#35758)

This commit adds the support for exists query in the sorted execution mode
of the `composite` aggregation. We'll execute the aggregation from the sorted
points and use early termination if the main query is an `exists` query over the
first source of the `composite` aggregation.

Author: jimczi

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

@@ -25,6 +25,8 @@
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.SortedNumericDocValues;
 import org.apache.lucene.search.BoostQuery;
+import org.apache.lucene.search.ConstantScoreQuery;
+import org.apache.lucene.search.DocValuesFieldExistsQuery;
 import org.apache.lucene.search.IndexOrDocValuesQuery;
 import org.apache.lucene.search.MatchAllDocsQuery;
 import org.apache.lucene.search.PointRangeQuery;
@@ -180,24 +182,42 @@ public void collect(int doc, long bucket) throws IOException {
         };
     }
 
-    static Query extractQuery(Query query) {
+    private 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 if (query instanceof ConstantScoreQuery){
+            return extractQuery(((ConstantScoreQuery) query).getQuery());
         } else {
             return query;
         }
     }
 
+    /**
+     * Returns true if we can use <code>query</code> with a {@link SortedDocsProducer} on <code>fieldName</code>.
+     */
+    private static boolean checkMatchAllOrRangeQuery(Query query, String fieldName) {
+        if (query == null) {
+            return true;
+        } else if (query.getClass() == MatchAllDocsQuery.class) {
+            return true;
+        } else if (query instanceof PointRangeQuery) {
+            PointRangeQuery pointQuery = (PointRangeQuery) query;
+            return fieldName.equals(pointQuery.getField());
+        } else if (query instanceof DocValuesFieldExistsQuery) {
+            DocValuesFieldExistsQuery existsQuery = (DocValuesFieldExistsQuery) query;
+            return fieldName.equals(existsQuery.getField());
+        } else {
+            return false;
+        }
+    }
+
     @Override
     SortedDocsProducer createSortedDocsProducerOrNull(IndexReader reader, Query query) {
         query = extractQuery(query);
         if (checkIfSortedDocsIsApplicable(reader, fieldType) == false ||
-                (query != null &&
-                    query.getClass() != MatchAllDocsQuery.class &&
-                    // if the query is a range query over the same field
-                    (query instanceof PointRangeQuery && fieldType.name().equals((((PointRangeQuery) query).getField()))) == false)) {
+                checkMatchAllOrRangeQuery(query, fieldType.name()) == false) {
             return null;
         }
         final byte[] lowerPoint;

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

@@ -23,6 +23,8 @@
 import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.index.Term;
 import org.apache.lucene.search.BoostQuery;
+import org.apache.lucene.search.ConstantScoreQuery;
+import org.apache.lucene.search.DocValuesFieldExistsQuery;
 import org.apache.lucene.search.IndexOrDocValuesQuery;
 import org.apache.lucene.search.MatchAllDocsQuery;
 import org.apache.lucene.search.TermQuery;
@@ -55,6 +57,7 @@ public void testBinarySorted() {
         IndexReader reader = mockIndexReader(1, 1);
         assertNotNull(source.createSortedDocsProducerOrNull(reader, new MatchAllDocsQuery()));
         assertNotNull(source.createSortedDocsProducerOrNull(reader, null));
+        assertNull(source.createSortedDocsProducerOrNull(reader, new TermQuery(new Term("foo", "bar"))));
         assertNull(source.createSortedDocsProducerOrNull(reader,
             new TermQuery(new Term("keyword", "toto)"))));
 
@@ -112,6 +115,7 @@ public void testGlobalOrdinalsSorted() {
         IndexReader reader = mockIndexReader(1, 1);
         assertNotNull(source.createSortedDocsProducerOrNull(reader, new MatchAllDocsQuery()));
         assertNotNull(source.createSortedDocsProducerOrNull(reader, null));
+        assertNull(source.createSortedDocsProducerOrNull(reader, new TermQuery(new Term("foo", "bar"))));
         assertNull(source.createSortedDocsProducerOrNull(reader,
             new TermQuery(new Term("keyword", "toto)"))));
 
@@ -126,6 +130,7 @@ public void testGlobalOrdinalsSorted() {
         );
         assertNull(source.createSortedDocsProducerOrNull(reader, new MatchAllDocsQuery()));
         assertNull(source.createSortedDocsProducerOrNull(reader, null));
+        assertNull(source.createSortedDocsProducerOrNull(reader, new TermQuery(new Term("foo", "bar"))));
 
         source = new GlobalOrdinalValuesSource(
             BigArrays.NON_RECYCLING_INSTANCE,
@@ -137,6 +142,7 @@ public void testGlobalOrdinalsSorted() {
             -1
         );
         assertNull(source.createSortedDocsProducerOrNull(reader, null));
+        assertNull(source.createSortedDocsProducerOrNull(reader, new TermQuery(new Term("foo", "bar"))));
 
         final MappedFieldType ip = new IpFieldMapper.IpFieldType();
         ip.setName("ip");
@@ -150,6 +156,7 @@ public void testGlobalOrdinalsSorted() {
             1
         );
         assertNull(source.createSortedDocsProducerOrNull(reader, null));
+        assertNull(source.createSortedDocsProducerOrNull(reader, new TermQuery(new Term("foo", "bar"))));
     }
 
     public void testNumericSorted() {
@@ -179,6 +186,9 @@ public void testNumericSorted() {
                 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 DocValuesFieldExistsQuery("number")));
+                assertNotNull(source.createSortedDocsProducerOrNull(reader,
+                    new ConstantScoreQuery(new DocValuesFieldExistsQuery("number"))));
                 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)"))));
@@ -195,6 +205,9 @@ public void testNumericSorted() {
                 assertNull(sourceWithMissing.createSortedDocsProducerOrNull(reader, new MatchAllDocsQuery()));
                 assertNull(sourceWithMissing.createSortedDocsProducerOrNull(reader, null));
                 assertNull(sourceWithMissing.createSortedDocsProducerOrNull(reader, new TermQuery(new Term("keyword", "toto)"))));
+                assertNull(sourceWithMissing.createSortedDocsProducerOrNull(reader, new DocValuesFieldExistsQuery("number")));
+                assertNull(sourceWithMissing.createSortedDocsProducerOrNull(reader,
+                    new ConstantScoreQuery(new DocValuesFieldExistsQuery("number"))));
 
                 LongValuesSource sourceRev = new LongValuesSource(
                     BigArrays.NON_RECYCLING_INSTANCE,
@@ -207,6 +220,10 @@ public void testNumericSorted() {
                     -1
                 );
                 assertNull(sourceRev.createSortedDocsProducerOrNull(reader, null));
+                assertNull(sourceRev.createSortedDocsProducerOrNull(reader, new DocValuesFieldExistsQuery("number")));
+                assertNull(sourceRev.createSortedDocsProducerOrNull(reader,
+                    new ConstantScoreQuery(new DocValuesFieldExistsQuery("number"))));
+                assertNull(sourceWithMissing.createSortedDocsProducerOrNull(reader, new TermQuery(new Term("keyword", "toto)"))));
             } else if (numberType == NumberFieldMapper.NumberType.HALF_FLOAT ||
                             numberType == NumberFieldMapper.NumberType.FLOAT ||
                             numberType == NumberFieldMapper.NumberType.DOUBLE) {
@@ -221,6 +238,10 @@ public void testNumericSorted() {
                 );
                 IndexReader reader = mockIndexReader(1, 1);
                 assertNull(source.createSortedDocsProducerOrNull(reader, null));
+                assertNull(source.createSortedDocsProducerOrNull(reader, new DocValuesFieldExistsQuery("number")));
+                assertNull(source.createSortedDocsProducerOrNull(reader, new TermQuery(new Term("keyword", "toto)"))));
+                assertNull(source.createSortedDocsProducerOrNull(reader,
+                    new ConstantScoreQuery(new DocValuesFieldExistsQuery("number"))));
             } else{
                 throw new AssertionError ("missing type:" + numberType.typeName());
             }