ESQL: Load script fields row-by-row

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

@@ -35,7 +35,9 @@ public void advanceTo(int docId) {
 
     private void advanceIfNeeded() throws IOException {
         if (loaderDocId != docId) {
-            loader.advanceTo(docId);
+            if (loader != null) {
+                loader.advanceTo(docId);
+            }
             loaderDocId = docId;
         }
     }

server/src/main/java/org/elasticsearch/index/mapper/blockloader/script/BlockScriptReader.java

@@ -0,0 +1,97 @@
+/*
+ * 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", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.index.mapper.blockloader.script;
+
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.index.SortedSetDocValues;
+import org.apache.lucene.util.IOFunction;
+import org.elasticsearch.common.breaker.CircuitBreaker;
+import org.elasticsearch.common.unit.ByteSizeValue;
+import org.elasticsearch.index.mapper.BlockLoader;
+import org.elasticsearch.search.fetch.StoredFieldsSpec;
+
+import java.io.IOException;
+
+/**
+ * A reader that supports reading doc-values from a Lucene segment in Block fashion.
+ */
+public abstract class BlockScriptReader implements BlockLoader.RowStrideReader {
+    protected final CircuitBreaker breaker;
+    private final long byteSize;
+    private final Thread creationThread;
+
+    public BlockScriptReader(CircuitBreaker breaker, long byteSize) {
+        this.breaker = breaker;
+        this.byteSize = byteSize;
+        this.creationThread = Thread.currentThread();
+    }
+
+    protected abstract int docId();
+
+    /**
+     * Checks if the reader can be used to read a range documents starting with the given docID by the current thread.
+     */
+    @Override
+    public final boolean canReuse(int startingDocID) {
+        return creationThread == Thread.currentThread() && docId() <= startingDocID;
+    }
+
+    @Override
+    public final void close() {
+        breaker.addWithoutBreaking(-byteSize);
+    }
+
+    @Override
+    public abstract String toString();
+
+    public abstract static class ScriptBlockLoader implements BlockLoader {
+        private final long byteSize;
+
+        protected ScriptBlockLoader(ByteSizeValue byteSize) {
+            this.byteSize = byteSize.getBytes();
+        }
+
+        public abstract BlockScriptReader reader(CircuitBreaker breaker, LeafReaderContext context) throws IOException;
+
+        @Override
+        public final IOFunction<CircuitBreaker, ColumnAtATimeReader> columnAtATimeReader(LeafReaderContext context) {
+            return null;
+        }
+
+        @Override
+        public final RowStrideReader rowStrideReader(CircuitBreaker breaker, LeafReaderContext context) throws IOException {
+            breaker.addEstimateBytesAndMaybeBreak(byteSize, "load blocks");
+            RowStrideReader reader = null;
+            try {
+                reader = reader(breaker, context);
+                return reader;
+            } finally {
+                if (reader == null) {
+                    breaker.addWithoutBreaking(-byteSize);
+                }
+            }
+        }
+
+        @Override
+        public final StoredFieldsSpec rowStrideStoredFieldSpec() {
+            return StoredFieldsSpec.NO_REQUIREMENTS;
+        }
+
+        @Override
+        public boolean supportsOrdinals() {
+            return false;
+        }
+
+        @Override
+        public SortedSetDocValues ordinals(LeafReaderContext context) throws IOException {
+            throw new UnsupportedOperationException();
+        }
+    }
+}

server/src/main/java/org/elasticsearch/index/mapper/blockloader/script/BooleanScriptBlockDocValuesReader.java

@@ -21,12 +21,13 @@
 /**
  * {@link BlockDocValuesReader} implementation for {@code boolean} scripts.
  */
-public class BooleanScriptBlockDocValuesReader extends BlockDocValuesReader {
-    public static class BooleanScriptBlockLoader extends DocValuesBlockLoader {
+public class BooleanScriptBlockDocValuesReader extends BlockScriptReader {
+    public static class BooleanScriptBlockLoader extends ScriptBlockLoader {
         private final BooleanFieldScript.LeafFactory factory;
         private final long byteSize;
 
         public BooleanScriptBlockLoader(BooleanFieldScript.LeafFactory factory, ByteSizeValue byteSize) {
+            super(byteSize);
             this.factory = factory;
             this.byteSize = byteSize.getBytes();
         }
@@ -37,28 +38,17 @@ public Builder builder(BlockFactory factory, int expectedCount) {
         }
 
         @Override
-        public AllReader reader(CircuitBreaker breaker, LeafReaderContext context) throws IOException {
-            breaker.addEstimateBytesAndMaybeBreak(byteSize, "load blocks");
-            BooleanFieldScript script = null;
-            try {
-                script = factory.newInstance(context);
-                return new BooleanScriptBlockDocValuesReader(breaker, script, byteSize);
-            } finally {
-                if (script == null) {
-                    breaker.addWithoutBreaking(-byteSize);
-                }
-            }
+        public BlockScriptReader reader(CircuitBreaker breaker, LeafReaderContext context) throws IOException {
+            return new BooleanScriptBlockDocValuesReader(breaker, factory.newInstance(context), byteSize);
         }
     }
 
     private final BooleanFieldScript script;
-    private final long byteSize;
     private int docId;
 
     BooleanScriptBlockDocValuesReader(CircuitBreaker breaker, BooleanFieldScript script, long byteSize) {
-        super(breaker);
+        super(breaker, byteSize);
         this.script = script;
-        this.byteSize = byteSize;
     }
 
     @Override
@@ -67,24 +57,9 @@ public int docId() {
     }
 
     @Override
-    public BlockLoader.Block read(BlockLoader.BlockFactory factory, BlockLoader.Docs docs, int offset, boolean nullsFiltered)
-        throws IOException {
-        // Note that we don't emit falses before trues so we conform to the doc values contract and can use booleansFromDocValues
-        try (BlockLoader.BooleanBuilder builder = factory.booleans(docs.count() - offset)) {
-            for (int i = offset; i < docs.count(); i++) {
-                read(docs.get(i), builder);
-            }
-            return builder.build();
-        }
-    }
-
-    @Override
-    public void read(int docId, BlockLoader.StoredFields storedFields, BlockLoader.Builder builder) throws IOException {
+    public void read(int docId, BlockLoader.StoredFields storedFields, BlockLoader.Builder b) throws IOException {
+        BlockLoader.BooleanBuilder builder = (BlockLoader.BooleanBuilder) b;
         this.docId = docId;
-        read(docId, (BlockLoader.BooleanBuilder) builder);
-    }
-
-    private void read(int docId, BlockLoader.BooleanBuilder builder) {
         script.runForDoc(docId);
         int total = script.falses() + script.trues();
         switch (total) {
@@ -107,9 +82,4 @@ private void read(int docId, BlockLoader.BooleanBuilder builder) {
     public String toString() {
         return "ScriptBooleans";
     }
-
-    @Override
-    public void close() {
-        breaker.addWithoutBreaking(-byteSize);
-    }
 }

server/src/main/java/org/elasticsearch/index/mapper/blockloader/script/DateScriptBlockDocValuesReader.java

@@ -21,12 +21,13 @@
 /**
  * {@link BlockDocValuesReader} implementation for date scripts.
  */
-public class DateScriptBlockDocValuesReader extends BlockDocValuesReader {
-    public static class DateScriptBlockLoader extends DocValuesBlockLoader {
+public class DateScriptBlockDocValuesReader extends BlockScriptReader {
+    public static class DateScriptBlockLoader extends ScriptBlockLoader {
         private final DateFieldScript.LeafFactory factory;
         private final long byteSize;
 
         public DateScriptBlockLoader(DateFieldScript.LeafFactory factory, ByteSizeValue byteSize) {
+            super(byteSize);
             this.factory = factory;
             this.byteSize = byteSize.getBytes();
         }
@@ -37,28 +38,17 @@ public Builder builder(BlockFactory factory, int expectedCount) {
         }
 
         @Override
-        public AllReader reader(CircuitBreaker breaker, LeafReaderContext context) throws IOException {
-            breaker.addEstimateBytesAndMaybeBreak(byteSize, "load blocks");
-            DateFieldScript script = null;
-            try {
-                script = factory.newInstance(context);
-                return new DateScriptBlockDocValuesReader(breaker, script, byteSize);
-            } finally {
-                if (script == null) {
-                    breaker.addWithoutBreaking(-byteSize);
-                }
-            }
+        public BlockScriptReader reader(CircuitBreaker breaker, LeafReaderContext context) throws IOException {
+            return new DateScriptBlockDocValuesReader(breaker, factory.newInstance(context), byteSize);
         }
     }
 
     private final DateFieldScript script;
-    private final long byteSize;
     private int docId;
 
     DateScriptBlockDocValuesReader(CircuitBreaker breaker, DateFieldScript script, long byteSize) {
-        super(breaker);
+        super(breaker, byteSize);
         this.script = script;
-        this.byteSize = byteSize;
     }
 
     @Override
@@ -67,24 +57,9 @@ public int docId() {
     }
 
     @Override
-    public BlockLoader.Block read(BlockLoader.BlockFactory factory, BlockLoader.Docs docs, int offset, boolean nullsFiltered)
-        throws IOException {
-        // Note that we don't sort the values sort, so we can't use factory.longsFromDocValues
-        try (BlockLoader.LongBuilder builder = factory.longs(docs.count() - offset)) {
-            for (int i = offset; i < docs.count(); i++) {
-                read(docs.get(i), builder);
-            }
-            return builder.build();
-        }
-    }
-
-    @Override
-    public void read(int docId, BlockLoader.StoredFields storedFields, BlockLoader.Builder builder) throws IOException {
+    public void read(int docId, BlockLoader.StoredFields storedFields, BlockLoader.Builder b) throws IOException {
+        BlockLoader.LongBuilder builder = (BlockLoader.LongBuilder) b;
         this.docId = docId;
-        read(docId, (BlockLoader.LongBuilder) builder);
-    }
-
-    private void read(int docId, BlockLoader.LongBuilder builder) {
         script.runForDoc(docId);
         switch (script.count()) {
             case 0 -> builder.appendNull();
@@ -103,9 +78,4 @@ private void read(int docId, BlockLoader.LongBuilder builder) {
     public String toString() {
         return "ScriptDates";
     }
-
-    @Override
-    public void close() {
-        breaker.addWithoutBreaking(-byteSize);
-    }
 }

server/src/main/java/org/elasticsearch/index/mapper/blockloader/script/DoubleScriptBlockDocValuesReader.java

@@ -21,12 +21,13 @@
 /**
  * {@link BlockDocValuesReader} implementation for {@code double} scripts.
  */
-public class DoubleScriptBlockDocValuesReader extends BlockDocValuesReader {
-    public static class DoubleScriptBlockLoader extends DocValuesBlockLoader {
+public class DoubleScriptBlockDocValuesReader extends BlockScriptReader {
+    public static class DoubleScriptBlockLoader extends ScriptBlockLoader {
         private final DoubleFieldScript.LeafFactory factory;
         private final long byteSize;
 
         public DoubleScriptBlockLoader(DoubleFieldScript.LeafFactory factory, ByteSizeValue byteSize) {
+            super(byteSize);
             this.factory = factory;
             this.byteSize = byteSize.getBytes();
         }
@@ -37,29 +38,17 @@ public Builder builder(BlockFactory factory, int expectedCount) {
         }
 
         @Override
-        public AllReader reader(CircuitBreaker breaker, LeafReaderContext context) throws IOException {
-            breaker.addEstimateBytesAndMaybeBreak(byteSize, "load blocks");
-            DoubleFieldScript script = null;
-            try {
-                script = factory.newInstance(context);
-                return new DoubleScriptBlockDocValuesReader(breaker, script, byteSize);
-            } finally {
-                if (script == null) {
-                    breaker.addWithoutBreaking(-byteSize);
-                }
-            }
+        public BlockScriptReader reader(CircuitBreaker breaker, LeafReaderContext context) throws IOException {
+            return new DoubleScriptBlockDocValuesReader(breaker, factory.newInstance(context), byteSize);
         }
     }
 
     private final DoubleFieldScript script;
-    private final long byteSize;
     private int docId;
 
     DoubleScriptBlockDocValuesReader(CircuitBreaker breaker, DoubleFieldScript script, long byteSize) {
-        super(breaker);
+        super(breaker, byteSize);
         this.script = script;
-        this.byteSize = byteSize;
-
     }
 
     @Override
@@ -68,24 +57,9 @@ public int docId() {
     }
 
     @Override
-    public BlockLoader.Block read(BlockLoader.BlockFactory factory, BlockLoader.Docs docs, int offset, boolean nullsFiltered)
-        throws IOException {
-        // Note that we don't sort the values sort, so we can't use factory.doublesFromDocValues
-        try (BlockLoader.DoubleBuilder builder = factory.doubles(docs.count() - offset)) {
-            for (int i = offset; i < docs.count(); i++) {
-                read(docs.get(i), builder);
-            }
-            return builder.build();
-        }
-    }
-
-    @Override
-    public void read(int docId, BlockLoader.StoredFields storedFields, BlockLoader.Builder builder) throws IOException {
+    public void read(int docId, BlockLoader.StoredFields storedFields, BlockLoader.Builder b) throws IOException {
+        BlockLoader.DoubleBuilder builder = (BlockLoader.DoubleBuilder) b;
         this.docId = docId;
-        read(docId, (BlockLoader.DoubleBuilder) builder);
-    }
-
-    private void read(int docId, BlockLoader.DoubleBuilder builder) {
         script.runForDoc(docId);
         switch (script.count()) {
             case 0 -> builder.appendNull();
@@ -104,9 +78,4 @@ private void read(int docId, BlockLoader.DoubleBuilder builder) {
     public String toString() {
         return "ScriptDoubles";
     }
-
-    @Override
-    public void close() {
-        breaker.addWithoutBreaking(-byteSize);
-    }
 }