better handling cancel extractor

iverase · Jan 16, 2025 · b869267 · b869267
1 parent 3a0ff94
commit b869267
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/...in/ml/src/main/java/org/elasticsearch/xpack/ml/dataframe/process/DataFrameRowsJoiner.java b/...in/ml/src/main/java/org/elasticsearch/xpack/ml/dataframe/process/DataFrameRowsJoiner.java
@@ -96,6 +96,9 @@ private void addResultAndJoinIfEndOfBatch(RowResults rowResults) {
     private void joinCurrentResults() {
         try (LimitAwareBulkIndexer bulkIndexer = new LimitAwareBulkIndexer(settings, this::executeBulkRequest)) {
             while (currentResults.isEmpty() == false) {
+                if (dataExtractor.isCancelled()) {
+                    break;
+                }
                 RowResults result = currentResults.pop();
                 DataFrameDataExtractor.Row row = dataFrameRowsIterator.next();
                 checkChecksumsMatch(row, result);
@@ -174,7 +177,7 @@ public boolean hasNext() {
         @Override
         public DataFrameDataExtractor.Row next() {
             DataFrameDataExtractor.Row row = null;
-            while (hasNoMatch(row) && hasNext() && dataExtractor.isCancelled() == false) {
+            while (hasNoMatch(row) && hasNext()) {
                 advanceToNextBatchIfNecessary();
                 row = dataExtractor.createRow(currentDataFrameRows[currentDataFrameRowsIndex++]);
             }