ESQL: Load many fields column-at-a-time

server/src/main/java/org/elasticsearch/index/mapper/BlockLoader.java

@@ -39,6 +39,11 @@
  *     spent the past 40 years making them really really good at running tight loops over
  *     arrays of data. So we play along with the CPU and make arrays.
  * </p>
+ * <p>
+ *     Implementers will return non-null from either {@link #columnAtATimeReader} or
+ *     {@link #rowStrideReader}. As of 2026-2-4 many will return non-null from both
+ *     but that's deprecated and will be removed.
+ * </p>
  * <h2>How to implement</h2>
  * <p>
  *     There are a lot of interesting choices hiding in here to make getting those arrays
@@ -253,6 +258,11 @@ interface RowStrideReader extends Reader {
         void read(int docId, StoredFields storedFields, Builder builder) throws IOException;
     }
 
+    /**
+     * @deprecated we no longer need to implement {@link RowStrideReader} when
+     *             storage prefers column-at-a-time
+     */
+    @Deprecated
     interface AllReader extends ColumnAtATimeReader, RowStrideReader {}
 
     interface StoredFields {
@@ -293,16 +303,17 @@ interface StoredFields {
      * Build a column-at-a-time reader. <strong>May</strong> return {@code null}
      * if the underlying storage needs to be loaded row-by-row. Callers should try
      * this first, only falling back to {@link #rowStrideReader} if this returns
-     * {@code null} or if they can't load column-at-a-time themselves.
+     * {@code null}. If this returns null then {@link #rowStrideReader} may not.
      */
     @Nullable
     IOSupplier<ColumnAtATimeReader> columnAtATimeReader(LeafReaderContext context) throws IOException;
 
     /**
-     * Build a row-by-row reader. Must <strong>never</strong> return {@code null},
-     * evan if the underlying storage prefers to be loaded column-at-a-time. Some
-     * callers simply can't load column-at-a-time so all implementations must support
-     * this method.
+     * Build a row-by-row reader. <strong>May</strong> return {@code null} if the
+     * underlying storage prefers to be loaded column-at-a-time. Callers should try
+     * {@link #columnAtATimeReader} first, only falling back to this if
+     * {@link #columnAtATimeReader} returns null. This may not return null if
+     * {@link #columnAtATimeReader} does.
      */
     RowStrideReader rowStrideReader(LeafReaderContext context) throws IOException;
 

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

@@ -22,6 +22,11 @@ class ComputeBlockLoaderFactory extends DelegatingBlockLoaderFactory implements
     public Block constantNulls(int count) {
         if (nullBlock == null) {
             nullBlock = factory.newConstantNullBlock(count);
+        } else {
+            if (nullBlock.getPositionCount() != count) {
+                nullBlock.close();
+                nullBlock = factory.newConstantNullBlock(count);
+            }
         }
         nullBlock.incRef();
         return nullBlock;

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

@@ -22,6 +22,8 @@
 import org.elasticsearch.search.fetch.StoredFieldsSpec;
 
 import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
 
 /**
  * Loads values from a many leaves. Much less efficient than {@link ValuesFromSingleReader}.
@@ -31,15 +33,13 @@ class ValuesFromManyReader extends ValuesReader {
 
     private final int[] forwards;
     private final int[] backwards;
-    private final BlockLoader.RowStrideReader[] rowStride;
 
     private BlockLoaderStoredFieldsFromLeafLoader storedFields;
 
     ValuesFromManyReader(ValuesSourceReaderOperator operator, DocVector docs) {
         super(operator, docs);
         forwards = docs.shardSegmentDocMapForwards();
         backwards = docs.shardSegmentDocMapBackwards();
-        rowStride = new BlockLoader.RowStrideReader[operator.fields.length];
         log.debug("initializing {} positions", docs.getPositionCount());
     }
 
@@ -71,12 +71,17 @@ class Run implements Releasable {
          * </p>
          */
         private final CurrentWork[] current;
+
+        private final List<CurrentWork> columnAtATime;
+        private final List<CurrentWork> rowStride;
         private int currentShard = -1;
 
         Run(Block[] target) {
             this.target = target;
             finalBuilders = new Block.Builder[target.length];
             current = new CurrentWork[target.length];
+            columnAtATime = new ArrayList<>(target.length);
+            rowStride = new ArrayList<>(target.length);
         }
 
         void run(int offset) throws IOException {
@@ -98,8 +103,9 @@ void run(int offset) throws IOException {
             operator.positionFieldWork(shard, segment, firstDoc);
             LeafReaderContext ctx = operator.ctx(shard, segment);
             fieldsMoved(ctx, shard);
-            read(firstDoc);
+            readRowStride(firstDoc);
 
+            int segmentStart = offset;
             int i = offset + 1;
             long estimated = estimatedRamBytesUsed();
             long dangerZoneBytes = Long.MAX_VALUE; // TODO danger_zone if ascending
@@ -109,14 +115,17 @@ void run(int offset) throws IOException {
                 segment = docs.segments().getInt(p);
                 boolean changedSegment = operator.positionFieldWorkDocGuaranteedAscending(shard, segment);
                 if (changedSegment) {
+                    readColumnAtATime(segmentStart, i);
+                    segmentStart = i;
                     ctx = operator.ctx(shard, segment);
                     fieldsMoved(ctx, shard);
                 }
-                read(docs.docs().getInt(p));
+                readRowStride(docs.docs().getInt(p));
                 i++;
                 estimated = estimatedRamBytesUsed();
                 log.trace("{}: bytes loaded {}/{}", p, estimated, dangerZoneBytes);
             }
+            readColumnAtATime(segmentStart, i);
             buildBlocks();
             if (log.isDebugEnabled()) {
                 long actual = 0;
@@ -131,19 +140,36 @@ private void buildBlocks() {
             convertAndAccumulate();
             for (int f = 0; f < target.length; f++) {
                 try (Block targetBlock = finalBuilders[f].build()) {
+                    assert targetBlock.getPositionCount() == backwards.length
+                        : targetBlock.getPositionCount() + " == " + backwards.length + " " + targetBlock;
                     target[f] = targetBlock.filter(backwards);
                 }
-                operator.sanityCheckBlock(rowStride[f], backwards.length, target[f], f);
+                operator.sanityCheckBlock(current[f].rowStride, backwards.length, target[f], f);
             }
             if (target[0].getPositionCount() != docs.getPositionCount()) {
                 throw new IllegalStateException("partial pages not yet supported");
             }
         }
 
-        private void read(int doc) throws IOException {
+        private void readRowStride(int doc) throws IOException {
             storedFields.advanceTo(doc);
-            for (int f = 0; f < current.length; f++) {
-                rowStride[f].read(doc, storedFields, current[f].builder);
+            for (CurrentWork r : rowStride) {
+                assert r.columnAtATime == null;
+                r.rowStride.read(doc, storedFields, r.builder);
+            }
+        }
+
+        private void readColumnAtATime(int segmentStart, int segmentEnd) throws IOException {
+            ValuesReaderDocs readerDocs = new ValuesReaderDocs(docs).mapped(forwards);
+            readerDocs.setCount(segmentEnd);
+            for (CurrentWork c : columnAtATime) {
+                assert c.rowStride == null;
+                try (Block read = (Block) c.columnAtATime.read(blockFactory, readerDocs, segmentStart, c.field.info.nullsFiltered())) {
+                    // TODO add a `read(builder, docs, offset, nullsFiltered)` override
+                    assert read.getPositionCount() == segmentEnd - segmentStart
+                        : read.getPositionCount() + " == " + segmentEnd + " - " + segmentStart + " " + read;
+                    c.builder.copyFrom(read, 0, read.getPositionCount());
+                }
             }
         }
 
@@ -162,11 +188,25 @@ private long estimatedRamBytesUsed() {
         }
 
         private void fieldsMoved(LeafReaderContext ctx, int shard) throws IOException {
+            if (currentShard != shard) {
+                if (currentShard >= 0) {
+                    convertAndAccumulate();
+                }
+                moveBuildersAndLoadersToShard();
+                currentShard = shard;
+            }
             StoredFieldsSpec storedFieldsSpec = StoredFieldsSpec.NO_REQUIREMENTS;
-            for (int f = 0; f < operator.fields.length; f++) {
-                ValuesSourceReaderOperator.FieldWork field = operator.fields[f];
-                rowStride[f] = field.rowStride(ctx);
-                storedFieldsSpec = storedFieldsSpec.merge(field.loader.rowStrideStoredFieldSpec());
+            columnAtATime.clear();
+            rowStride.clear();
+            for (CurrentWork field : current) {
+                field.columnAtATime = field.field.columnAtATime(ctx);
+                if (field.columnAtATime != null) {
+                    columnAtATime.add(field);
+                } else {
+                    field.rowStride = field.field.rowStride(ctx);
+                    storedFieldsSpec = storedFieldsSpec.merge(field.field.loader.rowStrideStoredFieldSpec());
+                    rowStride.add(field);
+                }
             }
             SourceLoader sourceLoader = null;
             if (storedFieldsSpec.requiresSource()) {
@@ -180,14 +220,6 @@ private void fieldsMoved(LeafReaderContext ctx, int shard) throws IOException {
             if (false == storedFieldsSpec.equals(StoredFieldsSpec.NO_REQUIREMENTS)) {
                 operator.trackStoredFields(storedFieldsSpec, false);
             }
-
-            if (currentShard != shard) {
-                if (currentShard >= 0) {
-                    convertAndAccumulate();
-                }
-                moveBuildersAndLoadersToShard();
-                currentShard = shard;
-            }
         }
 
         private void convertAndAccumulate() {
@@ -228,17 +260,28 @@ private void moveBuildersAndLoadersToShard() {
      * the {@link #finalBuilder} immediately.
      */
     private static class CurrentWork implements Releasable {
-        private final Block.Builder builder;
+        private final ValuesSourceReaderOperator.FieldWork field;
+        /**
+         * Converter for the field at the current segment. It's a copy of
+         * {@code field.converter} at the time of construction. By the time we actually
+         * go to use the converter, the field has moved onto another shard, changing
+         * the value of {@code field.converter}.
+         */
         @Nullable
         private final ValuesSourceReaderOperator.ConverterEvaluator converter;
+        private final Block.Builder builder;
         private final Block.Builder finalBuilder;
 
+        private BlockLoader.ColumnAtATimeReader columnAtATime;
+        private BlockLoader.RowStrideReader rowStride;
+
         CurrentWork(
             ComputeBlockLoaderFactory blockFactory,
             DocVector docs,
             ValuesSourceReaderOperator.FieldWork field,
             Block.Builder finalBuilder
         ) {
+            this.field = field;
             this.converter = field.converter;
             this.builder = converter == null ? finalBuilder : (Block.Builder) field.loader.builder(blockFactory, docs.getPositionCount());
             this.finalBuilder = finalBuilder;

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

@@ -53,7 +53,7 @@ public int get(int i) {
         return docs.docs().getInt(i);
     }
 
-    private class Mapped extends ValuesReaderDocs {
+    private static class Mapped extends ValuesReaderDocs {
         private final int[] forwards;
 
         private Mapped(DocVector docs, int[] forwards) {

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

@@ -274,6 +274,8 @@ protected class FieldWork {
         private final int fieldIdx;
 
         BlockLoader loader;
+        // TODO rework this bit of mutable state into something harder to forget
+        // Seriously, I've tripped over this twice.
         @Nullable
         ConverterEvaluator converter;
         @Nullable