Add ref counter to OrdinalsGroupingOperator

Author: GalLalouche

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/DocVector.java

@@ -51,7 +51,7 @@ public final class DocVector extends AbstractVector implements Vector {
 
     private final ShardRefCounted shardRefCounters;
 
-    public ShardRefCounted shardRefCounters() {
+    public ShardRefCounted shardRefCounted() {
         return shardRefCounters;
     }
 

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/OrdinalsGroupingOperator.java

@@ -33,6 +33,7 @@
 import org.elasticsearch.compute.data.IntVector;
 import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.lucene.ValuesSourceReaderOperator;
+import org.elasticsearch.core.RefCounted;
 import org.elasticsearch.core.Releasable;
 import org.elasticsearch.core.Releasables;
 import org.elasticsearch.index.mapper.BlockLoader;
@@ -136,6 +137,7 @@ public void addInput(Page page) {
         requireNonNull(page, "page is null");
         DocVector docVector = page.<DocBlock>getBlock(docChannel).asVector();
         final int shardIndex = docVector.shards().getInt(0);
+        RefCounted shardRefCounter = docVector.shardRefCounted().get(shardIndex);
         final var blockLoader = blockLoaders.apply(shardIndex);
         boolean pagePassed = false;
         try {
@@ -150,7 +152,8 @@ public void addInput(Page page) {
                                 driverContext.blockFactory(),
                                 this::createGroupingAggregators,
                                 () -> blockLoader.ordinals(shardContexts.get(k.shardIndex).reader().leaves().get(k.segmentIndex)),
-                                driverContext.bigArrays()
+                                driverContext.bigArrays(),
+                                shardRefCounter
                             );
                         } catch (IOException e) {
                             throw new UncheckedIOException(e);
@@ -343,15 +346,19 @@ static final class OrdinalSegmentAggregator implements Releasable, SeenGroupIds
         private final List<GroupingAggregator> aggregators;
         private final CheckedSupplier<SortedSetDocValues, IOException> docValuesSupplier;
         private final BitArray visitedOrds;
+        private final RefCounted shardRefCounted;
         private BlockOrdinalsReader currentReader;
 
         OrdinalSegmentAggregator(
             BlockFactory blockFactory,
             Supplier<List<GroupingAggregator>> aggregatorsSupplier,
             CheckedSupplier<SortedSetDocValues, IOException> docValuesSupplier,
-            BigArrays bigArrays
+            BigArrays bigArrays,
+            RefCounted shardRefCounted
         ) throws IOException {
             boolean success = false;
+            shardRefCounted.mustIncRef();
+            this.shardRefCounted = shardRefCounted;
             List<GroupingAggregator> groupingAggregators = null;
             BitArray bitArray = null;
             try {
@@ -363,11 +370,13 @@ static final class OrdinalSegmentAggregator implements Releasable, SeenGroupIds
                 this.docValuesSupplier = docValuesSupplier;
                 this.aggregators = groupingAggregators;
                 this.visitedOrds = bitArray;
+                this.shardRefCounted.mustIncRef();
                 success = true;
             } finally {
                 if (success == false) {
                     if (bitArray != null) Releasables.close(bitArray);
                     if (groupingAggregators != null) Releasables.close(groupingAggregators);
+                    shardRefCounted.decRef();
                 }
             }
         }
@@ -447,7 +456,7 @@ public BitArray seenGroupIds(BigArrays bigArrays) {
 
         @Override
         public void close() {
-            Releasables.close(visitedOrds, () -> Releasables.close(aggregators));
+            Releasables.close(visitedOrds, () -> Releasables.close(aggregators), Releasables.fromRefCounted(shardRefCounted));
         }
     }
 

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/topn/TopNOperator.java

@@ -216,7 +216,7 @@ private void writeKey(int position, Row row) {
         private void writeValues(int position, Row destination) {
             for (ValueExtractor e : valueExtractors) {
                 if (e instanceof ValueExtractorForDoc fd) {
-                    destination.setShardRefCountersAndShard(fd.vector().shardRefCounters().get(fd.vector().shards().getInt(position)));
+                    destination.setShardRefCountersAndShard(fd.vector().shardRefCounted().get(fd.vector().shards().getInt(position)));
                 }
                 e.writeValue(destination.values, position);
             }