From 9e646a9c92bf9d416b7c62142c547e6019275174 Mon Sep 17 00:00:00 2001
From: Nhat Nguyen <nhat.nguyen@elastic.co>
Date: Fri, 21 Feb 2025 14:05:18 -0800
Subject: [PATCH 1/3] Fix early termination in LuceneSourceOperator

---
 .../compute/lucene/LuceneSourceOperator.java  |  2 +-
 .../lucene/LuceneSourceOperatorTests.java     | 22 +++++++++++++++++++
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneSourceOperator.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneSourceOperator.java
index 3d34067e1a839..37b7091ac540f 100644
--- a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneSourceOperator.java
+++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneSourceOperator.java
@@ -140,7 +140,7 @@ public void collect(int doc) throws IOException {
 
     @Override
     public boolean isFinished() {
-        return doneCollecting;
+        return doneCollecting || remainingDocs == 0;
     }
 
     @Override
diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneSourceOperatorTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneSourceOperatorTests.java
index 574f9b25ff146..42c9f49a2db7c 100644
--- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneSourceOperatorTests.java
+++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneSourceOperatorTests.java
@@ -25,6 +25,7 @@
 import org.elasticsearch.compute.operator.Driver;
 import org.elasticsearch.compute.operator.DriverContext;
 import org.elasticsearch.compute.operator.Operator;
+import org.elasticsearch.compute.operator.SourceOperator;
 import org.elasticsearch.compute.test.AnyOperatorTestCase;
 import org.elasticsearch.compute.test.OperatorTestCase;
 import org.elasticsearch.compute.test.TestResultPageSinkOperator;
@@ -117,6 +118,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, scoring);
+        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));
     }

From 1c0dd56cc48f204489ed667d6cbe5e8fedc27bf9 Mon Sep 17 00:00:00 2001
From: Nhat Nguyen <nhat.nguyen@elastic.co>
Date: Fri, 21 Feb 2025 16:02:50 -0800
Subject: [PATCH 2/3] Update docs/changelog/123197.yaml

---
 docs/changelog/123197.yaml | 5 +++++
 1 file changed, 5 insertions(+)
 create mode 100644 docs/changelog/123197.yaml

diff --git a/docs/changelog/123197.yaml b/docs/changelog/123197.yaml
new file mode 100644
index 0000000000000..ffb4bab79fe8c
--- /dev/null
+++ b/docs/changelog/123197.yaml
@@ -0,0 +1,5 @@
+pr: 123197
+summary: Fix early termination in `LuceneSourceOperator`
+area: ES|QL
+type: bug
+issues: []

From 02ebb8ae63768bb1b594598a70fc3dbde74717b8 Mon Sep 17 00:00:00 2001
From: Nhat Nguyen <nhat.nguyen@elastic.co>
Date: Fri, 21 Feb 2025 16:05:38 -0800
Subject: [PATCH 3/3] harden

---
 .../org/elasticsearch/compute/lucene/LuceneSourceOperator.java  | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneSourceOperator.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneSourceOperator.java
index 37b7091ac540f..61a7cbad3e8af 100644
--- a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneSourceOperator.java
+++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneSourceOperator.java
@@ -140,7 +140,7 @@ public void collect(int doc) throws IOException {
 
     @Override
     public boolean isFinished() {
-        return doneCollecting || remainingDocs == 0;
+        return doneCollecting || remainingDocs <= 0;
     }
 
     @Override