ESQL: Fix bug in topn

docs/changelog/133601.yaml

@@ -0,0 +1,5 @@
+pr: 133601
+summary: Fix bug in topn
+area: ES|QL
+type: bug
+issues: []

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

@@ -108,6 +108,12 @@ public String describe() {
      * Collected docs. {@code null} until we're {@link #emit(boolean)}.
      */
     private ScoreDoc[] scoreDocs;
+
+    /**
+     * {@link ShardRefCounted} for collected docs.
+     */
+    private ShardRefCounted shardRefCounted;
+
     /**
      * The offset in {@link #scoreDocs} of the next page.
      */
@@ -142,6 +148,7 @@ public boolean isFinished() {
     public void finish() {
         doneCollecting = true;
         scoreDocs = null;
+        shardRefCounted = null;
         assert isFinished();
     }
 
@@ -202,6 +209,8 @@ private Page emit(boolean startEmitting) {
             offset = 0;
             if (perShardCollector != null) {
                 scoreDocs = perShardCollector.collector.topDocs().scoreDocs;
+                int shardId = perShardCollector.shardContext.index();
+                shardRefCounted = new ShardRefCounted.Single(shardId, shardContextCounters.get(shardId));
             } else {
                 scoreDocs = new ScoreDoc[0];
             }
@@ -239,7 +248,6 @@ private Page emit(boolean startEmitting) {
             shard = blockFactory.newConstantIntBlockWith(shardId, size);
             segments = currentSegmentBuilder.build();
             docs = currentDocsBuilder.build();
-            ShardRefCounted shardRefCounted = new ShardRefCounted.Single(shardId, shardContextCounters.get(shardId));
             docBlock = new DocVector(shardRefCounted, shard.asVector(), segments, docs, null).asBlock();
             shard = null;
             segments = null;