CR comments

Author: GalLalouche

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

@@ -138,11 +138,7 @@ public final Page getOutput() {
     protected abstract Page getCheckedOutput() throws IOException;
 
     @Override
-    public void close() {
-        if (currentScorer != null) {
-            currentScorer.shardContext().decRef();
-        }
-    }
+    public void close() {}
 
     LuceneScorer getCurrentOrLoadNextScorer() {
         while (currentScorer == null || currentScorer.isDone()) {
@@ -165,11 +161,7 @@ LuceneScorer getCurrentOrLoadNextScorer() {
             ) {
                 final Weight weight = currentSlice.weight();
                 processedQueries.add(weight.getQuery());
-                var previousScorer = currentScorer;
                 currentScorer = new LuceneScorer(currentSlice.shardContext(), weight, currentSlice.tags(), leaf);
-                if (previousScorer != null) {
-                    previousScorer.shardContext().decRef();
-                }
             }
             assert currentScorer.maxPosition <= partialLeaf.maxDoc() : currentScorer.maxPosition + ">" + partialLeaf.maxDoc();
             currentScorer.maxPosition = partialLeaf.maxDoc();
@@ -196,7 +188,6 @@ static final class LuceneScorer {
         private Thread executingThread;
 
         LuceneScorer(ShardContext shardContext, Weight weight, List<Object> tags, LeafReaderContext leafReaderContext) {
-            shardContext.incRef();
             this.shardContext = shardContext;
             this.weight = weight;
             this.tags = tags;

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

@@ -11,6 +11,8 @@
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.ElementType;
+import org.elasticsearch.core.Nullable;
+import org.elasticsearch.core.RefCounted;
 import org.elasticsearch.core.Releasable;
 
 /**
@@ -33,6 +35,12 @@ interface ResultBuilder extends Releasable {
      */
     void decodeValue(BytesRef values);
 
+    /**
+     * Sets the RefCounted value, which was extracted by {@link ValueExtractor#getRefCountedForShard(int)}. By default, this is a no-op,
+     * since most builders do not the shard ref counter.
+     */
+    default void setNextRefCounted(@Nullable RefCounted nextRefCounted) { /* no-op */ }
+
     /**
      * Build the result block.
      */

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

@@ -42,7 +42,8 @@ public void decodeKey(BytesRef keys) {
         throw new AssertionError("_doc can't be a key");
     }
 
-    void setNextRefCounted(RefCounted nextRefCounted) {
+    @Override
+    public void setNextRefCounted(RefCounted nextRefCounted) {
         this.nextRefCounted = nextRefCounted;
         // Since rows can be closed before build is called, we need to increment the ref count to ensure the shard context isn't closed.
         this.nextRefCounted.mustIncRef();

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

@@ -215,8 +215,9 @@ 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().shardRefCounted().get(fd.vector().shards().getInt(position)));
+                var refCounted = e.getRefCountedForShard(position);
+                if (refCounted != null) {
+                    destination.setShardRefCountersAndShard(refCounted);
                 }
                 e.writeValue(destination.values, position);
             }
@@ -486,10 +487,7 @@ private Iterator<Page> toPages() {
 
                 BytesRef values = row.values.bytesRefView();
                 for (ResultBuilder builder : builders) {
-                    if (builder instanceof ResultBuilderForDoc fd) {
-                        assert row.shardRefCounter != null : "shardRefCounter must be set for ResultBuilderForDoc";
-                        fd.setNextRefCounted(row.shardRefCounter);
-                    }
+                    builder.setNextRefCounted(row.shardRefCounter);
                     builder.decodeValue(values);
                 }
                 if (values.length != 0) {

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

@@ -18,13 +18,24 @@
 import org.elasticsearch.compute.data.IntBlock;
 import org.elasticsearch.compute.data.LongBlock;
 import org.elasticsearch.compute.operator.BreakingBytesRefBuilder;
+import org.elasticsearch.core.Nullable;
+import org.elasticsearch.core.RefCounted;
 
 /**
  * Extracts values into a {@link BreakingBytesRefBuilder}.
  */
 interface ValueExtractor {
     void writeValue(BreakingBytesRefBuilder values, int position);
 
+    /**
+     * This should return a non-null value if the row is supposed to hold a temporary reference to a shard (including incrementing and
+     * decrementing it) in between encoding and decoding the row values.
+     */
+    @Nullable
+    default RefCounted getRefCountedForShard(int position) {
+        return null;
+    }
+
     static ValueExtractor extractorFor(ElementType elementType, TopNEncoder encoder, boolean inKey, Block block) {
         if (false == (elementType == block.elementType() || ElementType.NULL == block.elementType())) {
             // While this maybe should be an IllegalArgumentException, it's important to throw an exception that causes a 500 response.

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

@@ -9,10 +9,16 @@
 
 import org.elasticsearch.compute.data.DocVector;
 import org.elasticsearch.compute.operator.BreakingBytesRefBuilder;
+import org.elasticsearch.core.RefCounted;
 
 class ValueExtractorForDoc implements ValueExtractor {
     private final DocVector vector;
 
+    @Override
+    public RefCounted getRefCountedForShard(int position) {
+        return vector().shardRefCounted().get(vector().shards().getInt(position));
+    }
+
     ValueExtractorForDoc(TopNEncoder encoder, DocVector vector) {
         assert encoder == TopNEncoder.DEFAULT_UNSORTABLE;
         this.vector = vector;