Optimized path for single segment in TS source

Author: dnhatn

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<LeafIterator> 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<LeafIterator> sub = subQueueForNextTsid();
+                    if (sub.size() == 0) {
+                        break;
+                    }
+                    tsHashesBuilder.appendNewTsid(sub.top().timeSeriesHash);
+                    if (readValuesForOneTsid(sub)) {
+                        break;
+                    }
+                } while (mainQueue.size() > 0);
+            }
         }
 
         private boolean readValuesForOneTsid(PriorityQueue<LeafIterator> sub) throws IOException {
@@ -236,6 +240,38 @@ private PriorityQueue<LeafIterator> 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<LeafIterator> getMainQueue() {
+            return mainQueue;
+        }
+
         boolean completed() {
             return mainQueue.size() == 0 && oneTsidQueue.size() == 0;
         }