diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/TimeSeriesSourceOperator.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/TimeSeriesSourceOperator.java index 089846f9939ae..ba6da814542e4 100644 --- a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/TimeSeriesSourceOperator.java +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/TimeSeriesSourceOperator.java @@ -183,16 +183,20 @@ void readDocsForNextPage() throws IOException { for (LeafIterator leaf : oneTsidQueue) { leaf.reinitializeIfNeeded(executingThread); } - do { - PriorityQueue sub = subQueueForNextTsid(); - if (sub.size() == 0) { - break; - } - tsHashesBuilder.appendNewTsid(sub.top().timeSeriesHash); - if (readValuesForOneTsid(sub)) { - break; - } - } while (mainQueue.size() > 0); + if (mainQueue.size() + oneTsidQueue.size() == 1) { + readValuesFromSingleRemainingLeaf(); + } else { + do { + PriorityQueue sub = subQueueForNextTsid(); + if (sub.size() == 0) { + break; + } + tsHashesBuilder.appendNewTsid(sub.top().timeSeriesHash); + if (readValuesForOneTsid(sub)) { + break; + } + } while (mainQueue.size() > 0); + } } private boolean readValuesForOneTsid(PriorityQueue sub) throws IOException { @@ -236,6 +240,38 @@ private PriorityQueue subQueueForNextTsid() { return oneTsidQueue; } + private void readValuesFromSingleRemainingLeaf() throws IOException { + if (oneTsidQueue.size() == 0) { + oneTsidQueue.add(getMainQueue().pop()); + tsidsLoaded++; + } + final LeafIterator sub = oneTsidQueue.top(); + int lastTsid = -1; + do { + currentPagePos++; + remainingDocs--; + docCollector.collect(sub.segmentOrd, sub.docID); + if (lastTsid != sub.lastTsidOrd) { + tsHashesBuilder.appendNewTsid(sub.timeSeriesHash); + lastTsid = sub.lastTsidOrd; + } + tsHashesBuilder.appendOrdinal(); + timestampsBuilder.appendLong(sub.timestamp); + if (sub.nextDoc() == false) { + if (sub.docID == DocIdSetIterator.NO_MORE_DOCS) { + oneTsidQueue.clear(); + return; + } else { + ++tsidsLoaded; + } + } + } while (remainingDocs > 0 && currentPagePos < maxPageSize); + } + + private PriorityQueue getMainQueue() { + return mainQueue; + } + boolean completed() { return mainQueue.size() == 0 && oneTsidQueue.size() == 0; }