ESQL: Use a little less memory in DocVector (#133319)

Prevously we made a unique `ShardRefCounted` for each `DocVector`. This
shares the `ShardRefCounted`s between all `DocVector`s that are in the
same shard.

Author: nik9000

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

@@ -70,6 +70,10 @@ public abstract class LuceneOperator extends SourceOperator {
     private int sliceIndex;
 
     private LuceneScorer currentScorer;
+    /**
+     * The {@link ShardRefCounted} for the current scorer.
+     */
+    private ShardRefCounted.Single currentScorerShardRefCounted;
 
     long processingNanos;
     int pagesEmitted;
@@ -172,6 +176,10 @@ LuceneScorer getCurrentOrLoadNextScorer() {
                 }
                 processedSlices++;
                 processedShards.add(currentSlice.shardContext().shardIdentifier());
+                int shardId = currentSlice.shardContext().index();
+                if (currentScorerShardRefCounted == null || currentScorerShardRefCounted.index() != shardId) {
+                    currentScorerShardRefCounted = new ShardRefCounted.Single(shardId, shardContextCounters.get(shardId));
+                }
             }
             final PartialLeafReaderContext partialLeaf = currentSlice.getLeaf(sliceIndex++);
             logger.trace("Starting {}", partialLeaf);
@@ -194,6 +202,13 @@ LuceneScorer getCurrentOrLoadNextScorer() {
         return currentScorer;
     }
 
+    /**
+     * The {@link ShardRefCounted} for the current scorer.
+     */
+    ShardRefCounted currentScorerShardRefCounted() {
+        return currentScorerShardRefCounted;
+    }
+
     /**
      * Wraps a {@link BulkScorer} with shard information
      */

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

@@ -334,8 +334,7 @@ public Page getCheckedOutput() throws IOException {
                     docs = buildDocsVector(currentPagePos);
                     docsBuilder = blockFactory.newIntVectorBuilder(Math.min(remainingDocs, maxPageSize));
                     int b = 0;
-                    ShardRefCounted refCounted = ShardRefCounted.single(shardId, shardContextCounters.get(shardId));
-                    blocks[b++] = new DocVector(refCounted, shard, leaf, docs, true).asBlock();
+                    blocks[b++] = new DocVector(currentScorerShardRefCounted(), shard, leaf, docs, true).asBlock();
                     shard = null;
                     leaf = null;
                     docs = null;

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

@@ -114,7 +114,6 @@ public String describe() {
     private int offset = 0;
 
     private PerShardCollector perShardCollector;
-    private final List<? extends ShardContext> contexts;
     private final List<SortBuilder<?>> sorts;
     private final int limit;
     private final boolean needsScore;
@@ -129,7 +128,6 @@ public LuceneTopNSourceOperator(
         boolean needsScore
     ) {
         super(contexts, blockFactory, maxPageSize, sliceQueue);
-        this.contexts = contexts;
         this.sorts = sorts;
         this.limit = limit;
         this.needsScore = needsScore;
@@ -241,7 +239,7 @@ private Page emit(boolean startEmitting) {
             shard = blockFactory.newConstantIntBlockWith(shardId, size);
             segments = currentSegmentBuilder.build();
             docs = currentDocsBuilder.build();
-            ShardRefCounted shardRefCounted = ShardRefCounted.single(shardId, contexts.get(shardId));
+            ShardRefCounted shardRefCounted = new ShardRefCounted.Single(shardId, shardContextCounters.get(shardId));
             docBlock = new DocVector(shardRefCounted, shard.asVector(), segments, docs, null).asBlock();
             shard = null;
             segments = null;

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

@@ -24,17 +24,18 @@ static ShardRefCounted fromList(List<? extends RefCounted> refCounters) {
     }
 
     static ShardRefCounted fromShardContext(ShardContext shardContext) {
-        return single(shardContext.index(), shardContext);
+        return new Single(shardContext.index(), shardContext);
     }
 
-    static ShardRefCounted single(int index, RefCounted refCounted) {
-        return shardId -> {
+    ShardRefCounted ALWAYS_REFERENCED = shardId -> RefCounted.ALWAYS_REFERENCED;
+
+    record Single(int index, RefCounted refCounted) implements ShardRefCounted {
+        @Override
+        public RefCounted get(int shardId) {
             if (shardId != index) {
                 throw new IllegalArgumentException("Invalid shardId: " + shardId + ", expected: " + index);
             }
             return refCounted;
-        };
+        }
     }
-
-    ShardRefCounted ALWAYS_REFERENCED = shardId -> RefCounted.ALWAYS_REFERENCED;
 }

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneSourceOperatorTests.java

@@ -63,6 +63,7 @@
 import static org.hamcrest.Matchers.lessThan;
 import static org.hamcrest.Matchers.lessThanOrEqualTo;
 import static org.hamcrest.Matchers.matchesRegex;
+import static org.hamcrest.Matchers.sameInstance;
 
 public class LuceneSourceOperatorTests extends SourceOperatorTestCase {
     private static final MappedFieldType S_FIELD = new NumberFieldMapper.NumberFieldType("s", NumberFieldMapper.NumberType.LONG);
@@ -373,6 +374,7 @@ private void testSimple(DriverContext ctx, DataPartitioning partitioning, int nu
         for (Page page : results) {
             assertThat(page.getPositionCount(), lessThanOrEqualTo(factory.maxPageSize()));
         }
+        assertAllRefCountedSameInstance(results);
 
         for (Page page : results) {
             LongBlock sBlock = page.getBlock(initialBlockIndex(page));
@@ -483,4 +485,17 @@ public boolean hasReferences() {
             return true;
         }
     }
+
+    static void assertAllRefCountedSameInstance(List<Page> results) {
+        ShardRefCounted firstRefCounted = null;
+        for (Page page : results) {
+            DocBlock docs = page.getBlock(0);
+            ShardRefCounted refCounted = docs.asVector().shardRefCounted();
+            if (firstRefCounted == null) {
+                firstRefCounted = refCounted;
+            } else {
+                assertThat(refCounted, sameInstance(firstRefCounted));
+            }
+        }
+    }
 }

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneTopNSourceOperatorTests.java

@@ -48,6 +48,7 @@
 import java.util.Optional;
 import java.util.function.Function;
 
+import static org.elasticsearch.compute.lucene.LuceneSourceOperatorTests.assertAllRefCountedSameInstance;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.hasSize;
 import static org.hamcrest.Matchers.matchesRegex;
@@ -207,6 +208,7 @@ protected void testSimple(DriverContext ctx, int size, int limit) {
                 assertThat(sBlock.getLong(sBlock.getFirstValueIndex(p)), equalTo(expectedS++));
             }
         }
+        assertAllRefCountedSameInstance(results);
         int pages = (int) Math.ceil((float) Math.min(size, limit) / factory.maxPageSize());
         assertThat(results, hasSize(pages));
     }