Move shardRefCounter logic to the super class

Author: GalLalouche

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

@@ -35,7 +35,6 @@ public class LuceneCountOperator extends LuceneOperator {
 
     private static final int PAGE_SIZE = 1;
 
-    private final List<? extends RefCounted> shardRefCounters;
     private int totalHits = 0;
     private int remainingDocs;
 
@@ -81,9 +80,7 @@ public LuceneCountOperator(
         LuceneSliceQueue sliceQueue,
         int limit
     ) {
-        super(blockFactory, PAGE_SIZE, sliceQueue);
-        this.shardRefCounters = shardRefCounters;
-        shardRefCounters.forEach(RefCounted::mustIncRef);
+        super(shardRefCounters, blockFactory, PAGE_SIZE, sliceQueue);
         this.remainingDocs = limit;
         this.leafCollector = new LeafCollector() {
             @Override
@@ -182,9 +179,4 @@ protected Page getCheckedOutput() throws IOException {
     protected void describe(StringBuilder sb) {
         sb.append(", remainingDocs=").append(remainingDocs);
     }
-
-    @Override
-    public void close() {
-        shardRefCounters.forEach(RefCounted::decRef);
-    }
 }

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

@@ -108,6 +108,7 @@ public final long evaluate(long value1, long value2) {
         abstract long bytesToLong(byte[] bytes);
     }
 
+    private final List<? extends ShardContext> contexts;
     private final String fieldName;
     private final NumberType numberType;
 
@@ -130,13 +131,14 @@ public LuceneMaxFactory(
             false,
             ScoreMode.COMPLETE_NO_SCORES
         );
+        this.contexts = contexts;
         this.fieldName = fieldName;
         this.numberType = numberType;
     }
 
     @Override
     public SourceOperator get(DriverContext driverContext) {
-        return new LuceneMinMaxOperator(driverContext.blockFactory(), sliceQueue, fieldName, numberType, limit, Long.MIN_VALUE);
+        return new LuceneMinMaxOperator(contexts, driverContext.blockFactory(), sliceQueue, fieldName, numberType, limit, Long.MIN_VALUE);
     }
 
     @Override

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

@@ -16,6 +16,7 @@
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.operator.DriverContext;
 import org.elasticsearch.compute.operator.SourceOperator;
+import org.elasticsearch.core.RefCounted;
 import org.elasticsearch.search.MultiValueMode;
 
 import java.io.IOException;
@@ -108,6 +109,7 @@ public final long evaluate(long value1, long value2) {
         abstract long bytesToLong(byte[] bytes);
     }
 
+    private final List<? extends RefCounted> shardRefCounters;
     private final String fieldName;
     private final NumberType numberType;
 
@@ -130,13 +132,22 @@ public LuceneMinFactory(
             false,
             ScoreMode.COMPLETE_NO_SCORES
         );
+        this.shardRefCounters = contexts;
         this.fieldName = fieldName;
         this.numberType = numberType;
     }
 
     @Override
     public SourceOperator get(DriverContext driverContext) {
-        return new LuceneMinMaxOperator(driverContext.blockFactory(), sliceQueue, fieldName, numberType, limit, Long.MAX_VALUE);
+        return new LuceneMinMaxOperator(
+            shardRefCounters,
+            driverContext.blockFactory(),
+            sliceQueue,
+            fieldName,
+            numberType,
+            limit,
+            Long.MAX_VALUE
+        );
     }
 
     @Override

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

@@ -20,10 +20,12 @@
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.BooleanBlock;
 import org.elasticsearch.compute.data.Page;
+import org.elasticsearch.core.RefCounted;
 import org.elasticsearch.core.Releasables;
 import org.elasticsearch.search.MultiValueMode;
 
 import java.io.IOException;
+import java.util.List;
 
 /**
  * Operator that finds the min or max value of a field using Lucene searches
@@ -65,14 +67,15 @@ sealed interface NumberType permits LuceneMinFactory.NumberType, LuceneMaxFactor
     private final String fieldName;
 
     LuceneMinMaxOperator(
+        List<? extends RefCounted> shardRefCounters,
         BlockFactory blockFactory,
         LuceneSliceQueue sliceQueue,
         String fieldName,
         NumberType numberType,
         int limit,
         long initialResult
     ) {
-        super(blockFactory, PAGE_SIZE, sliceQueue);
+        super(shardRefCounters, blockFactory, PAGE_SIZE, sliceQueue);
         this.remainingDocs = limit;
         this.numberType = numberType;
         this.fieldName = fieldName;

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

@@ -25,6 +25,7 @@
 import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.Operator;
 import org.elasticsearch.compute.operator.SourceOperator;
+import org.elasticsearch.core.RefCounted;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.logging.LogManager;
 import org.elasticsearch.logging.Logger;
@@ -52,6 +53,7 @@ public abstract class LuceneOperator extends SourceOperator {
 
     public static final int NO_LIMIT = Integer.MAX_VALUE;
 
+    protected final List<? extends RefCounted> shardContextCounters;
     protected final BlockFactory blockFactory;
 
     /**
@@ -77,7 +79,14 @@ public abstract class LuceneOperator extends SourceOperator {
      */
     long rowsEmitted;
 
-    protected LuceneOperator(BlockFactory blockFactory, int maxPageSize, LuceneSliceQueue sliceQueue) {
+    protected LuceneOperator(
+        List<? extends RefCounted> shardContextCounters,
+        BlockFactory blockFactory,
+        int maxPageSize,
+        LuceneSliceQueue sliceQueue
+    ) {
+        this.shardContextCounters = shardContextCounters;
+        shardContextCounters.forEach(RefCounted::incRef);
         this.blockFactory = blockFactory;
         this.maxPageSize = maxPageSize;
         this.sliceQueue = sliceQueue;
@@ -138,7 +147,12 @@ public final Page getOutput() {
     protected abstract Page getCheckedOutput() throws IOException;
 
     @Override
-    public void close() {}
+    public final void close() {
+        shardContextCounters.forEach(RefCounted::decRef);
+        additionalClose();
+    }
+
+    protected void additionalClose() { /* Override this method to add any additional cleanup logic if needed */ }
 
     LuceneScorer getCurrentOrLoadNextScorer() {
         while (currentScorer == null || currentScorer.isDone()) {

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

@@ -49,7 +49,6 @@
  */
 public class LuceneSourceOperator extends LuceneOperator {
     private static final Logger log = LogManager.getLogger(LuceneSourceOperator.class);
-    private final List<? extends RefCounted> shardContextCounters;
 
     private int currentPagePos = 0;
     private int remainingDocs;
@@ -227,8 +226,7 @@ public LuceneSourceOperator(
         Limiter limiter,
         boolean needsScore
     ) {
-        super(blockFactory, maxPageSize, sliceQueue);
-        this.shardContextCounters = shardContextCounters;
+        super(shardContextCounters, blockFactory, maxPageSize, sliceQueue);
         shardContextCounters.forEach(RefCounted::mustIncRef);
         this.minPageSize = Math.max(1, maxPageSize / 2);
         this.remainingDocs = limit;
@@ -330,12 +328,14 @@ public Page getCheckedOutput() throws IOException {
                 Block[] blocks = new Block[1 + (scoreBuilder == null ? 0 : 1) + scorer.tags().size()];
                 currentPagePos -= discardedDocs;
                 try {
-                    shard = blockFactory.newConstantIntVector(scorer.shardContext().index(), currentPagePos);
+                    int shardId = scorer.shardContext().index();
+                    shard = blockFactory.newConstantIntVector(shardId, currentPagePos);
                     leaf = blockFactory.newConstantIntVector(scorer.leafReaderContext().ord, currentPagePos);
                     docs = buildDocsVector(currentPagePos);
                     docsBuilder = blockFactory.newIntVectorBuilder(Math.min(remainingDocs, maxPageSize));
                     int b = 0;
-                    blocks[b++] = new DocVector(ShardRefCounted.fromList(shardContextCounters), shard, leaf, docs, true).asBlock();
+                    ShardRefCounted refCounted = ShardRefCounted.single(shardId, shardContextCounters.get(shardId));
+                    blocks[b++] = new DocVector(refCounted, shard, leaf, docs, true).asBlock();
                     shard = null;
                     leaf = null;
                     docs = null;
@@ -393,12 +393,8 @@ private DoubleVector buildScoresVector(int upToPositions) {
     }
 
     @Override
-    public void close() {
-        Releasables.close(
-            docsBuilder,
-            scoreBuilder,
-            Releasables.wrap(shardContextCounters.stream().map(Releasables::fromRefCounted).toList())
-        );
+    public void additionalClose() {
+        Releasables.close(docsBuilder, scoreBuilder);
     }
 
     @Override

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

@@ -132,9 +132,8 @@ public LuceneTopNSourceOperator(
         LuceneSliceQueue sliceQueue,
         boolean needsScore
     ) {
-        super(blockFactory, maxPageSize, sliceQueue);
+        super(contexts, blockFactory, maxPageSize, sliceQueue);
         this.contexts = contexts;
-        contexts.forEach(ShardContext::mustIncRef);
         this.sorts = sorts;
         this.limit = limit;
         this.needsScore = needsScore;
@@ -370,10 +369,4 @@ private static PerShardCollector newPerShardCollector(ShardContext context, List
         sort = new Sort(l.toArray(SortField[]::new));
         return new ScoringPerShardCollector(context, new TopFieldCollectorManager(sort, limit, null, 0).newCollector());
     }
-
-    @Override
-    public void close() {
-        super.close();
-        contexts.forEach(ShardContext::decRef);
-    }
 }

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

@@ -43,7 +43,6 @@
 import java.util.Set;
 
 public final class TimeSeriesSourceOperator extends LuceneOperator {
-    private final List<? extends ShardContext> contexts;
     private final int maxPageSize;
     private final BlockFactory blockFactory;
     private final LuceneSliceQueue sliceQueue;
@@ -64,9 +63,8 @@ public final class TimeSeriesSourceOperator extends LuceneOperator {
         int maxPageSize,
         int limit
     ) {
-        super(blockFactory, maxPageSize, sliceQueue);
+        super(contexts, blockFactory, maxPageSize, sliceQueue);
         contexts.forEach(RefCounted::mustIncRef);
-        this.contexts = contexts;
         this.maxPageSize = maxPageSize;
         this.blockFactory = blockFactory;
         this.remainingDocs = limit;
@@ -110,7 +108,7 @@ public Page getCheckedOutput() throws IOException {
                     throw new UnsupportedOperationException("tags not supported by " + getClass());
                 }
                 iterator = new SegmentsIterator(slice);
-                docCollector = new DocIdCollector(contexts, blockFactory, slice.shardContext());
+                docCollector = new DocIdCollector(blockFactory, slice.shardContext());
             }
             iterator.readDocsForNextPage();
             if (currentPagePos > 0) {
@@ -141,13 +139,8 @@ public Page getCheckedOutput() throws IOException {
     }
 
     @Override
-    public void close() {
-        Releasables.closeExpectNoException(
-            timestampsBuilder,
-            tsHashesBuilder,
-            docCollector,
-            Releasables.wrap(contexts.stream().map(Releasables::fromRefCounted).toList())
-        );
+    public void additionalClose() {
+        Releasables.closeExpectNoException(timestampsBuilder, tsHashesBuilder, docCollector);
     }
 
     class SegmentsIterator {
@@ -365,14 +358,12 @@ OrdinalBytesRefVector build() throws IOException {
     }
 
     static final class DocIdCollector implements Releasable {
-        private final List<? extends ShardContext> contexts;
         private final BlockFactory blockFactory;
         private final ShardContext shardContext;
         private IntVector.Builder docsBuilder;
         private IntVector.Builder segmentsBuilder;
 
-        DocIdCollector(List<? extends ShardContext> contexts, BlockFactory blockFactory, ShardContext shardContext) {
-            this.contexts = contexts;
+        DocIdCollector(BlockFactory blockFactory, ShardContext shardContext) {
             this.blockFactory = blockFactory;
             this.shardContext = shardContext;
         }