Allow LuceneSourceOperator to early terminate (elastic#108820)

We can make the LuceneSourceOperator collector to throw a CollectionTerminatedException whenever it has collected enough documents.
carlosdelest · May 20, 2024 · 60777cf · 60777cf
1 parent 4e1b8bf
commit 60777cf
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 8 deletions.
diff --git a/docs/changelog/108820.yaml b/docs/changelog/108820.yaml
@@ -0,0 +1,5 @@
+pr: 108820
+summary: Allow `LuceneSourceOperator` to early terminate
+area: ES|QL
+type: enhancement
+issues: []
diff --git a/...gin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneSourceOperator.java b/...gin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneSourceOperator.java
@@ -7,6 +7,7 @@
 
 package org.elasticsearch.compute.lucene;
 
+import org.apache.lucene.search.CollectionTerminatedException;
 import org.apache.lucene.search.LeafCollector;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.Scorable;
@@ -90,6 +91,8 @@ public void collect(int doc) {
                     --remainingDocs;
                     docsBuilder.appendInt(doc);
                     currentPagePos++;
+                } else {
+                    throw new CollectionTerminatedException();
                 }
             }
         };
@@ -117,14 +120,19 @@ public Page getCheckedOutput() throws IOException {
             if (scorer == null) {
                 return null;
             }
-            scorer.scoreNextRange(
-                leafCollector,
-                scorer.leafReaderContext().reader().getLiveDocs(),
-                // Note: if (maxPageSize - currentPagePos) is a small "remaining" interval, this could lead to slow collection with a
-                // highly selective filter. Having a large "enough" difference between max- and minPageSize (and thus currentPagePos)
-                // alleviates this issue.
-                maxPageSize - currentPagePos
-            );
+            try {
+                scorer.scoreNextRange(
+                    leafCollector,
+                    scorer.leafReaderContext().reader().getLiveDocs(),
+                    // Note: if (maxPageSize - currentPagePos) is a small "remaining" interval, this could lead to slow collection with a
+                    // highly selective filter. Having a large "enough" difference between max- and minPageSize (and thus currentPagePos)
+                    // alleviates this issue.
+                    maxPageSize - currentPagePos
+                );
+            } catch (CollectionTerminatedException ex) {
+                // The leaf collector terminated the execution
+                scorer.markAsDone();
+            }
             Page page = null;
             if (currentPagePos >= minPageSize || remainingDocs <= 0 || scorer.isDone()) {
                 pagesEmitted++;