[8.17] Fix early termination in LuceneSourceOperator (#123197) (#123839)

The LuceneSourceOperator is supposed to terminate when it reaches the
limit; unfortunately, we don't have a test to cover this. Due to this
bug, we continue scanning all segments, even though we discard the
results as the limit was reached. This can cause performance issues for
simple queries like FROM .. | LIMIT 10, when Lucene indices are on the
warm or cold tier. I will submit a follow-up PR to ensure we only
collect up to the limit across multiple drivers.

Author: dnhatn

docs/changelog/123197.yaml

@@ -0,0 +1,5 @@
+pr: 123197
+summary: Fix early termination in `LuceneSourceOperator`
+area: ES|QL
+type: bug
+issues: []

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneSourceOperator.java

@@ -100,7 +100,7 @@ public void collect(int doc) {
 
     @Override
     public boolean isFinished() {
-        return doneCollecting;
+        return doneCollecting || remainingDocs <= 0;
     }
 
     @Override

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneSourceOperatorTests.java

@@ -25,6 +25,7 @@
 import org.elasticsearch.compute.operator.DriverContext;
 import org.elasticsearch.compute.operator.Operator;
 import org.elasticsearch.compute.operator.OperatorTestCase;
+import org.elasticsearch.compute.operator.SourceOperator;
 import org.elasticsearch.compute.operator.TestResultPageSinkOperator;
 import org.elasticsearch.core.IOUtils;
 import org.elasticsearch.index.cache.query.TrivialQueryCachingPolicy;
@@ -112,6 +113,27 @@ public void testShardDataPartitioning() {
         testSimple(driverContext(), size, limit);
     }
 
+    public void testEarlyTermination() {
+        int size = between(1_000, 20_000);
+        int limit = between(10, size);
+        LuceneSourceOperator.Factory factory = simple(randomFrom(DataPartitioning.values()), size, limit);
+        try (SourceOperator sourceOperator = factory.get(driverContext())) {
+            assertFalse(sourceOperator.isFinished());
+            int collected = 0;
+            while (sourceOperator.isFinished() == false) {
+                Page page = sourceOperator.getOutput();
+                if (page != null) {
+                    collected += page.getPositionCount();
+                    page.releaseBlocks();
+                }
+                if (collected >= limit) {
+                    assertTrue("source operator is not finished after reaching limit", sourceOperator.isFinished());
+                    assertThat(collected, equalTo(limit));
+                }
+            }
+        }
+    }
+
     public void testEmpty() {
         testSimple(driverContext(), 0, between(10, 10_000));
     }