CR fixes

Author: GalLalouche

libs/core/src/main/java/org/elasticsearch/core/Releasables.java

@@ -202,6 +202,7 @@ public boolean equals(Object obj) {
         }
     }
 
+    /** Creates a {@link Releasable} that calls {@link RefCounted#decRef()} when closed. */
     public static Releasable fromRefCounted(RefCounted refCounted) {
         return () -> refCounted.decRef();
     }

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

@@ -11,10 +11,13 @@
 import org.apache.lucene.util.RamUsageEstimator;
 import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.compute.lucene.ShardRefCounted;
+import org.elasticsearch.core.RefCounted;
 import org.elasticsearch.core.ReleasableIterator;
 import org.elasticsearch.core.Releasables;
 
+import java.util.BitSet;
 import java.util.Objects;
+import java.util.function.Consumer;
 
 /**
  * {@link Vector} where each entry references a lucene document.
@@ -30,6 +33,7 @@ public final class DocVector extends AbstractVector implements Vector {
     public static final int SHARD_SEGMENT_DOC_MAP_PER_ROW_OVERHEAD = Integer.BYTES * 2;
 
     private final IntVector shards;
+    private final IntVector uniqueShards;
     private final IntVector segments;
     private final IntVector docs;
 
@@ -80,7 +84,32 @@ public DocVector(
         }
         blockFactory().adjustBreaker(BASE_RAM_BYTES_USED);
 
-        forEachShardRefCounter(DecOrInc.INC);
+        forEachShardRefCounter(RefCounted::mustIncRef);
+        this.uniqueShards = computeUniqueShards(shards);
+    }
+
+    private static IntVector computeUniqueShards(IntVector shards) {
+        switch (shards) {
+            case ConstantIntVector constantIntVector -> {
+                return shards.blockFactory().newConstantIntVector(constantIntVector.getInt(0), 1);
+            }
+            case ConstantNullVector unused -> {
+                return shards.blockFactory().newConstantIntVector(0, 0);
+            }
+            default -> {
+                var seen = new BitSet(128);
+                try (IntVector.Builder uniqueShardsBuilder = shards.blockFactory().newIntVectorBuilder(shards.getPositionCount())) {
+                    for (int p = 0; p < shards.getPositionCount(); p++) {
+                        int shardId = shards.getInt(p);
+                        if (seen.get(shardId) == false) {
+                            seen.set(shardId);
+                            uniqueShardsBuilder.appendInt(shardId);
+                        }
+                    }
+                    return uniqueShardsBuilder.build();
+                }
+            }
+        }
     }
 
     public DocVector(
@@ -310,18 +339,20 @@ private static long ramBytesOrZero(int[] array) {
 
     public static long ramBytesEstimated(
         IntVector shards,
+        IntVector uniqueShards,
         IntVector segments,
         IntVector docs,
         int[] shardSegmentDocMapForwards,
         int[] shardSegmentDocMapBackwards
     ) {
-        return BASE_RAM_BYTES_USED + RamUsageEstimator.sizeOf(shards) + RamUsageEstimator.sizeOf(segments) + RamUsageEstimator.sizeOf(docs)
-            + ramBytesOrZero(shardSegmentDocMapForwards) + ramBytesOrZero(shardSegmentDocMapBackwards);
+        return BASE_RAM_BYTES_USED + RamUsageEstimator.sizeOf(shards) + RamUsageEstimator.sizeOf(uniqueShards) + RamUsageEstimator.sizeOf(
+            segments
+        ) + RamUsageEstimator.sizeOf(docs) + ramBytesOrZero(shardSegmentDocMapForwards) + ramBytesOrZero(shardSegmentDocMapBackwards);
     }
 
     @Override
     public long ramBytesUsed() {
-        return ramBytesEstimated(shards, segments, docs, shardSegmentDocMapForwards, shardSegmentDocMapBackwards);
+        return ramBytesEstimated(shards, uniqueShards, segments, docs, shardSegmentDocMapForwards, shardSegmentDocMapBackwards);
     }
 
     @Override
@@ -337,33 +368,22 @@ public void closeInternal() {
         Releasables.closeExpectNoException(
             () -> blockFactory().adjustBreaker(-BASE_RAM_BYTES_USED - (shardSegmentDocMapForwards == null ? 0 : sizeOfSegmentDocMap())),
             shards,
+            uniqueShards,
             segments,
             docs
         );
-        forEachShardRefCounter(DecOrInc.DEC);
-    }
-
-    private enum DecOrInc {
-        DEC,
-        INC;
-
-        void apply(ShardRefCounted counters, int shardId) {
-            switch (this) {
-                case DEC -> counters.get(shardId).decRef();
-                case INC -> counters.get(shardId).mustIncRef();
-            }
-        }
+        forEachShardRefCounter(RefCounted::decRef);
     }
 
-    private void forEachShardRefCounter(DecOrInc mode) {
+    private void forEachShardRefCounter(Consumer<RefCounted> consumer) {
         switch (shards) {
-            case ConstantIntVector constantIntVector -> mode.apply(shardRefCounters, constantIntVector.getInt(0));
+            case ConstantIntVector constantIntVector -> consumer.accept(shardRefCounters.get(constantIntVector.getInt(0)));
             case ConstantNullVector ignored -> {
                 // Noop
             }
             default -> {
-                for (int i = 0; i < shards.getPositionCount(); i++) {
-                    mode.apply(shardRefCounters, shards.getInt(i));
+                for (int i = 0; i < uniqueShards.getPositionCount(); i++) {
+                    consumer.accept(shardRefCounters.get(uniqueShards.getInt(i)));
                 }
             }
         }

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

@@ -13,6 +13,10 @@
 
 /** Manages reference counting for {@link ShardContext}. */
 public interface ShardRefCounted {
+    /**
+     * @param shardId The shard index used by {@link org.elasticsearch.compute.data.DocVector}.
+     * @return the {@link RefCounted} for the given shard. In production, this will almost always be a {@link ShardContext}.
+     */
     RefCounted get(int shardId);
 
     static ShardRefCounted fromList(List<? extends RefCounted> refCounters) {

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/EsPhysicalOperationProviders.java

@@ -92,7 +92,8 @@ public class EsPhysicalOperationProviders extends AbstractPhysicalOperationProvi
     private static final Logger logger = LogManager.getLogger(EsPhysicalOperationProviders.class);
 
     /**
-     * Context of each shard we're operating against.
+     * Context of each shard we're operating against. Note these objects are shared across multiple operators as
+     * {@link org.elasticsearch.core.RefCounted}.
      */
     public abstract static class ShardContext implements org.elasticsearch.compute.lucene.ShardContext, Releasable {
         private final AbstractRefCounted refCounted = new AbstractRefCounted() {
@@ -379,6 +380,10 @@ public Operator.OperatorFactory timeSeriesAggregatorOperatorFactory(
 
     public static class DefaultShardContext extends ShardContext {
         private final int index;
+        /**
+         * In production, this will be a {@link org.elasticsearch.search.internal.SearchContext}, but we don't want to drag that huge
+         * dependency here.
+         */
         private final Releasable releasable;
         private final SearchExecutionContext ctx;
         private final AliasFilter aliasFilter;

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/ComputeService.java

@@ -590,6 +590,8 @@ public SourceProvider createSourceProvider() {
                 LOGGER.debug("Local execution plan:\n{}", localExecutionPlan.describe());
             }
             drivers = localExecutionPlan.createDrivers(context.sessionId());
+            // After creating the drivers (and therefore, the operators), we can safely decrement the reference count since the operators
+            // will hold a reference to the contexts where relevant.
             contexts.forEach(RefCounted::decRef);
             if (drivers.isEmpty()) {
                 throw new IllegalStateException("no drivers created");