ESQL: Late materialization after TopN (Node level) (elastic#132757)

This PR adds a late(r) materialization for TopN queries, such that the materialization happes in the "node_reduce" phase instead of during the "data" phase.

For example, if the limit is 20, and each data node spawns 10 workers, we would only read 20 additional columns (i.e., ones not needed for the TopN) filters, instead of 200. To support this, the reducer node maintains a global list of all shard contexts used by its individual data workers (although some of those might be closed if they are no longer needed, thanks to elastic#129454).

There is some additional book-keeping involved, since previously, every data node held a local list of shard contexts, and used its local indices to access it. To avoid changing too much (this local-index logic is spread throughout much of the code!), a new global index is introduced, which replaces the local index after all the rows are merged together in the reduce phase's TopN.

Author: GalLalouche

benchmarks/src/main/java/org/elasticsearch/benchmark/_nightly/esql/ValuesSourceReaderBenchmark.java

@@ -41,8 +41,9 @@
 import org.elasticsearch.compute.data.LongBlock;
 import org.elasticsearch.compute.data.LongVector;
 import org.elasticsearch.compute.data.Page;
+import org.elasticsearch.compute.lucene.AlwaysReferencedIndexedByShardId;
+import org.elasticsearch.compute.lucene.IndexedByShardIdFromSingleton;
 import org.elasticsearch.compute.lucene.LuceneSourceOperator;
-import org.elasticsearch.compute.lucene.ShardRefCounted;
 import org.elasticsearch.compute.lucene.read.ValuesSourceReaderOperator;
 import org.elasticsearch.compute.lucene.read.ValuesSourceReaderOperatorStatus;
 import org.elasticsearch.compute.operator.topn.TopNOperator;
@@ -368,7 +369,7 @@ public void benchmark() {
             blockFactory,
             ByteSizeValue.ofMb(1).getBytes(),
             fields(name),
-            List.of(new ValuesSourceReaderOperator.ShardContext(reader, () -> {
+            new IndexedByShardIdFromSingleton<>(new ValuesSourceReaderOperator.ShardContext(reader, () -> {
                 throw new UnsupportedOperationException("can't load _source here");
             }, EsqlPlugin.STORED_FIELDS_SEQUENTIAL_PROPORTION.getDefault(Settings.EMPTY))),
             0
@@ -538,7 +539,7 @@ private void setupPages() {
                         pages.add(
                             new Page(
                                 new DocVector(
-                                    ShardRefCounted.ALWAYS_REFERENCED,
+                                    AlwaysReferencedIndexedByShardId.INSTANCE,
                                     blockFactory.newConstantIntBlockWith(0, end - begin).asVector(),
                                     blockFactory.newConstantIntBlockWith(ctx.ord, end - begin).asVector(),
                                     docs.build(),
@@ -575,8 +576,7 @@ record ItrAndOrd(PrimitiveIterator.OfInt itr, int ord) {}
                             pages.add(
                                 new Page(
                                     new DocVector(
-
-                                        ShardRefCounted.ALWAYS_REFERENCED,
+                                        AlwaysReferencedIndexedByShardId.INSTANCE,
                                         blockFactory.newConstantIntVector(0, size),
                                         leafs.build(),
                                         docs.build(),
@@ -594,7 +594,7 @@ record ItrAndOrd(PrimitiveIterator.OfInt itr, int ord) {}
                     pages.add(
                         new Page(
                             new DocVector(
-                                ShardRefCounted.ALWAYS_REFERENCED,
+                                AlwaysReferencedIndexedByShardId.INSTANCE,
                                 blockFactory.newConstantIntBlockWith(0, size).asVector(),
                                 leafs.build().asBlock().asVector(),
                                 docs.build(),
@@ -621,8 +621,7 @@ record ItrAndOrd(PrimitiveIterator.OfInt itr, int ord) {}
                         pages.add(
                             new Page(
                                 new DocVector(
-
-                                    ShardRefCounted.ALWAYS_REFERENCED,
+                                    AlwaysReferencedIndexedByShardId.INSTANCE,
                                     blockFactory.newConstantIntVector(0, 1),
                                     blockFactory.newConstantIntVector(next.ord, 1),
                                     blockFactory.newConstantIntVector(next.itr.nextInt(), 1),

docs/changelog/132757.yaml

@@ -0,0 +1,5 @@
+pr: 132757
+summary: Late materialization after TopN (Node level)
+area: ES|QL
+type: feature
+issues: []

server/src/main/resources/transport/definitions/referable/esql_reduce_late_materialization.csv

@@ -0,0 +1 @@
+9194000

server/src/main/resources/transport/upper_bounds/9.3.csv

@@ -1 +1 @@
-available_processors_in_os_stats,9193000
+esql_reduce_late_materialization,9194000

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/esql/action/EsqlQueryRequestBuilder.java

@@ -41,4 +41,5 @@ public final ActionType<Response> action() {
 
     public abstract EsqlQueryRequestBuilder<Request, Response> allowPartialResults(boolean allowPartialResults);
 
+    public abstract EsqlQueryRequestBuilder<Request, Response> profile(boolean profile);
 }

x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/util/CollectionUtils.java

@@ -68,4 +68,26 @@ public static <T> List<T> combine(Collection<? extends T> left, T... entries) {
         }
         return list;
     }
+
+    /**
+     * Creates a copy of the given collection with the given element prepended.
+     *
+     * @param collection collection to copy
+     * @param element    element to prepend
+     */
+    @SuppressWarnings("unchecked")
+    public static <T> List<T> prependToCopy(T element, Collection<T> collection) {
+        T[] result = (T[]) new Object[collection.size() + 1];
+        result[0] = element;
+        if (collection instanceof ArrayList<T> arrayList && arrayList.size() <= 1_000_000) {
+            // Creating an array out of a relatively small ArrayList and copying it is faster than iterating.
+            System.arraycopy(arrayList.toArray(), 0, result, 1, result.length - 1);
+        } else {
+            var i = 1;
+            for (T t : collection) {
+                result[i++] = t;
+            }
+        }
+        return List.of(result);
+    }
 }

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

@@ -9,7 +9,8 @@
 
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.unit.ByteSizeValue;
-import org.elasticsearch.compute.lucene.ShardRefCounted;
+import org.elasticsearch.compute.lucene.AlwaysReferencedIndexedByShardId;
+import org.elasticsearch.compute.lucene.IndexedByShardId;
 import org.elasticsearch.core.RefCounted;
 import org.elasticsearch.core.ReleasableIterator;
 import org.elasticsearch.core.Releasables;
@@ -92,6 +93,14 @@ public void closeInternal() {
         Releasables.closeExpectNoException(vector);
     }
 
+    @Override
+    public String toString() {
+        final StringBuffer sb = new StringBuffer("DocBlock[");
+        sb.append("vector=").append(vector);
+        sb.append(']');
+        return sb.toString();
+    }
+
     /**
      * A builder the for {@link DocBlock}.
      */
@@ -103,9 +112,9 @@ public static class Builder implements Block.Builder {
         private final IntVector.Builder shards;
         private final IntVector.Builder segments;
         private final IntVector.Builder docs;
-        private ShardRefCounted shardRefCounters = ShardRefCounted.ALWAYS_REFERENCED;
+        private IndexedByShardId<? extends RefCounted> shardRefCounters = null;
 
-        public Builder setShardRefCounted(ShardRefCounted shardRefCounters) {
+        public Builder shardRefCounters(IndexedByShardId<? extends RefCounted> shardRefCounters) {
             this.shardRefCounters = shardRefCounters;
             return this;
         }
@@ -196,7 +205,13 @@ public DocBlock build() {
                 shards = this.shards.build();
                 segments = this.segments.build();
                 docs = this.docs.build();
-                result = new DocVector(shardRefCounters, shards, segments, docs, null);
+                result = new DocVector(
+                    shardRefCounters == null ? AlwaysReferencedIndexedByShardId.INSTANCE : shardRefCounters,
+                    shards,
+                    segments,
+                    docs,
+                    null
+                );
                 return result.asBlock();
             } finally {
                 if (result == null) {

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

@@ -10,7 +10,7 @@
 import org.apache.lucene.util.IntroSorter;
 import org.apache.lucene.util.RamUsageEstimator;
 import org.elasticsearch.common.unit.ByteSizeValue;
-import org.elasticsearch.compute.lucene.ShardRefCounted;
+import org.elasticsearch.compute.lucene.IndexedByShardId;
 import org.elasticsearch.core.RefCounted;
 import org.elasticsearch.core.ReleasableIterator;
 import org.elasticsearch.core.Releasables;
@@ -31,6 +31,11 @@ public final class DocVector extends AbstractVector implements Vector {
      */
     public static final int SHARD_SEGMENT_DOC_MAP_PER_ROW_OVERHEAD = Integer.BYTES * 2;
 
+    /**
+     * The shard IDs for each position. Note that these shard IDs are shared between all doc vectors running in the same node, but a given
+     * doc vector might only reference a subset of the shard IDs (Which is the subset is also the one exposed by {@link #refCounteds}).
+     * These shard IDs are sliced up by DataNodeComputeHandler, and depend on the MAX_CONCURRENT_SHARDS_PER_NODE setting.
+     */
     private final IntVector shards;
     private final IntVector segments;
     private final IntVector docs;
@@ -51,21 +56,41 @@ public final class DocVector extends AbstractVector implements Vector {
      */
     private int[] shardSegmentDocMapBackwards;
 
-    private final ShardRefCounted shardRefCounters;
+    private final IndexedByShardId<? extends RefCounted> refCounteds;
 
-    public ShardRefCounted shardRefCounted() {
-        return shardRefCounters;
+    public RefCounted shardRefCounted(int position) {
+        return refCounteds.get(shards.getInt(position));
     }
 
     public DocVector(
-        ShardRefCounted shardRefCounters,
+        IndexedByShardId<? extends RefCounted> refCounteds,
         IntVector shards,
         IntVector segments,
         IntVector docs,
         Boolean singleSegmentNonDecreasing
+    ) {
+        this(refCounteds, shards, segments, docs, singleSegmentNonDecreasing, true);
+    }
+
+    public static DocVector withoutIncrementingShardRefCounts(
+        IndexedByShardId<? extends RefCounted> refCounteds,
+        IntVector shards,
+        IntVector segments,
+        IntVector docs
+    ) {
+        return new DocVector(refCounteds, shards, segments, docs, null, false);
+    }
+
+    private DocVector(
+        IndexedByShardId<? extends RefCounted> refCounteds,
+        IntVector shards,
+        IntVector segments,
+        IntVector docs,
+        Boolean singleSegmentNonDecreasing,
+        boolean incrementShardRefCounts
     ) {
         super(shards.getPositionCount(), shards.blockFactory());
-        this.shardRefCounters = shardRefCounters;
+        this.refCounteds = refCounteds;
         this.shards = shards;
         this.segments = segments;
         this.docs = docs;
@@ -82,18 +107,20 @@ public DocVector(
         }
         blockFactory().adjustBreaker(BASE_RAM_BYTES_USED);
 
-        forEachShardRefCounter(RefCounted::mustIncRef);
+        if (incrementShardRefCounts) {
+            forEachShardRefCounter(RefCounted::mustIncRef);
+        }
     }
 
     public DocVector(
-        ShardRefCounted shardRefCounters,
+        IndexedByShardId<? extends RefCounted> refCounteds,
         IntVector shards,
         IntVector segments,
         IntVector docs,
         int[] docMapForwards,
         int[] docMapBackwards
     ) {
-        this(shardRefCounters, shards, segments, docs, null);
+        this(refCounteds, shards, segments, docs, null);
         this.shardSegmentDocMapForwards = docMapForwards;
         this.shardSegmentDocMapBackwards = docMapBackwards;
     }
@@ -269,7 +296,7 @@ public DocVector filter(int... positions) {
             filteredShards = shards.filter(positions);
             filteredSegments = segments.filter(positions);
             filteredDocs = docs.filter(positions);
-            result = new DocVector(shardRefCounters, filteredShards, filteredSegments, filteredDocs, null);
+            result = new DocVector(refCounteds, filteredShards, filteredSegments, filteredDocs, null);
             return result;
         } finally {
             if (result == null) {
@@ -288,7 +315,7 @@ public DocVector deepCopy(BlockFactory blockFactory) {
             filteredShards = shards.deepCopy(blockFactory);
             filteredSegments = segments.deepCopy(blockFactory);
             filteredDocs = docs.deepCopy(blockFactory);
-            result = new DocVector(shardRefCounters, filteredShards, filteredSegments, filteredDocs, null);
+            result = new DocVector(refCounteds, filteredShards, filteredSegments, filteredDocs, null);
             return result;
         } finally {
             if (result == null) {
@@ -331,6 +358,16 @@ public boolean equals(Object obj) {
         return shards.equals(other.shards) && segments.equals(other.segments) && docs.equals(other.docs);
     }
 
+    @Override
+    public String toString() {
+        final StringBuffer sb = new StringBuffer("DocVector[");
+        sb.append("shards=").append(shards);
+        sb.append(", segments=").append(segments);
+        sb.append(", docs=").append(docs);
+        sb.append(']');
+        return sb.toString();
+    }
+
     private static long ramBytesOrZero(int[] array) {
         return array == null ? 0 : RamUsageEstimator.shallowSizeOf(array);
     }
@@ -372,13 +409,13 @@ public void closeInternal() {
 
     private void forEachShardRefCounter(Consumer<RefCounted> consumer) {
         switch (shards) {
-            case ConstantIntVector constantIntVector -> consumer.accept(shardRefCounters.get(constantIntVector.getInt(0)));
+            case ConstantIntVector constantIntVector -> consumer.accept(refCounteds.get(constantIntVector.getInt(0)));
             case ConstantNullVector ignored -> {
                 // Noop
             }
             default -> {
                 for (int i = 0; i < shards.getPositionCount(); i++) {
-                    consumer.accept(shardRefCounters.get(shards.getInt(i)));
+                    consumer.accept(refCounteds.get(shards.getInt(i)));
                 }
             }
         }

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

@@ -0,0 +1,39 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.compute.lucene;
+
+import org.elasticsearch.core.RefCounted;
+
+import java.util.Collection;
+import java.util.List;
+import java.util.function.Function;
+
+/**
+ * An implementation which always returns {@link RefCounted#ALWAYS_REFERENCED} for any shard ID. Used by tests, but defined here so it could
+ * also be used by the benchmarks.
+ */
+public class AlwaysReferencedIndexedByShardId implements IndexedByShardId<RefCounted> {
+    public static final AlwaysReferencedIndexedByShardId INSTANCE = new AlwaysReferencedIndexedByShardId();
+
+    private AlwaysReferencedIndexedByShardId() {}
+
+    @Override
+    public RefCounted get(int shardId) {
+        return RefCounted.ALWAYS_REFERENCED;
+    }
+
+    @Override
+    public Collection<? extends RefCounted> collection() {
+        return List.of(RefCounted.ALWAYS_REFERENCED);
+    }
+
+    @Override
+    public <S> IndexedByShardId<S> map(Function<RefCounted, S> mapper) {
+        throw new UnsupportedOperationException();
+    }
+}

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

@@ -0,0 +1,36 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.compute.lucene;
+
+import java.util.Collection;
+import java.util.function.Function;
+
+public class EmptyIndexedByShardId {
+    @SuppressWarnings("unchecked")
+    public static <T> IndexedByShardId<T> instance() {
+        return (IndexedByShardId<T>) EMPTY;
+    }
+
+    private static IndexedByShardId<?> EMPTY = new IndexedByShardId<>() {
+        @Override
+        public Object get(int shardId) {
+            throw new IndexOutOfBoundsException("no shards");
+        }
+
+        @Override
+        public Collection<?> collection() {
+            return java.util.List.of();
+        }
+
+        @SuppressWarnings("unchecked")
+        @Override
+        public <S> IndexedByShardId<S> map(Function<Object, S> mapper) {
+            return (IndexedByShardId<S>) this;
+        }
+    };
+}