ESQL: Speed loading stored fields

benchmarks/src/main/java/org/elasticsearch/benchmark/compute/operator/ValuesSourceReaderBenchmark.java

@@ -25,6 +25,7 @@
 import org.apache.lucene.util.NumericUtils;
 import org.elasticsearch.common.breaker.NoopCircuitBreaker;
 import org.elasticsearch.common.lucene.Lucene;
+import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.BytesRefBlock;
@@ -50,6 +51,7 @@
 import org.elasticsearch.index.mapper.MappedFieldType;
 import org.elasticsearch.index.mapper.NumberFieldMapper;
 import org.elasticsearch.search.lookup.SearchLookup;
+import org.elasticsearch.xpack.esql.plugin.QueryPragmas;
 import org.openjdk.jmh.annotations.Benchmark;
 import org.openjdk.jmh.annotations.BenchmarkMode;
 import org.openjdk.jmh.annotations.Fork;
@@ -336,7 +338,8 @@ public void benchmark() {
             List.of(new ValuesSourceReaderOperator.ShardContext(reader, () -> {
                 throw new UnsupportedOperationException("can't load _source here");
             })),
-            0
+            0,
+            QueryPragmas.STORED_FIELDS_SEQUENTIAL_PROPORTION.getDefault(Settings.EMPTY)
         );
         long sum = 0;
         for (Page page : pages) {

docs/changelog/127348.yaml

@@ -0,0 +1,5 @@
+pr: 127348
+summary: Speed loading stored fields
+area: ES|QL
+type: enhancement
+issues: []

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

@@ -72,10 +72,18 @@ public class ValuesSourceReaderOperator extends AbstractPageMappingOperator {
      * @param shardContexts per-shard loading information
      * @param docChannel the channel containing the shard, leaf/segment and doc id
      */
-    public record Factory(List<FieldInfo> fields, List<ShardContext> shardContexts, int docChannel) implements OperatorFactory {
+    public record Factory(List<FieldInfo> fields, List<ShardContext> shardContexts, int docChannel, double storedFieldsSequentialProportion)
+        implements
+            OperatorFactory {
         @Override
         public Operator get(DriverContext driverContext) {
-            return new ValuesSourceReaderOperator(driverContext.blockFactory(), fields, shardContexts, docChannel);
+            return new ValuesSourceReaderOperator(
+                driverContext.blockFactory(),
+                fields,
+                shardContexts,
+                docChannel,
+                storedFieldsSequentialProportion
+            );
         }
 
         @Override
@@ -113,6 +121,7 @@ public record ShardContext(IndexReader reader, Supplier<SourceLoader> newSourceL
     private final List<ShardContext> shardContexts;
     private final int docChannel;
     private final BlockFactory blockFactory;
+    private final double storedFieldsSequentialProportion;
 
     private final Map<String, Integer> readersBuilt = new TreeMap<>();
     private long valuesLoaded;
@@ -125,11 +134,18 @@ public record ShardContext(IndexReader reader, Supplier<SourceLoader> newSourceL
      * @param fields fields to load
      * @param docChannel the channel containing the shard, leaf/segment and doc id
      */
-    public ValuesSourceReaderOperator(BlockFactory blockFactory, List<FieldInfo> fields, List<ShardContext> shardContexts, int docChannel) {
+    public ValuesSourceReaderOperator(
+        BlockFactory blockFactory,
+        List<FieldInfo> fields,
+        List<ShardContext> shardContexts,
+        int docChannel,
+        double storedFieldsSequentialProportion
+    ) {
         this.fields = fields.stream().map(f -> new FieldWork(f)).toArray(FieldWork[]::new);
         this.shardContexts = shardContexts;
         this.docChannel = docChannel;
         this.blockFactory = blockFactory;
+        this.storedFieldsSequentialProportion = storedFieldsSequentialProportion;
     }
 
     @Override
@@ -440,7 +456,11 @@ public void close() {
      */
     private boolean useSequentialStoredFieldsReader(BlockLoader.Docs docs) {
         int count = docs.count();
-        return count >= SEQUENTIAL_BOUNDARY && docs.get(count - 1) - docs.get(0) == count - 1;
+        if (count < SEQUENTIAL_BOUNDARY) {
+            return false;
+        }
+        int range = docs.get(count - 1) - docs.get(0);
+        return range * storedFieldsSequentialProportion < count - 1;
     }
 
     private void trackStoredFields(StoredFieldsSpec spec, boolean sequential) {

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

@@ -63,7 +63,8 @@ public record OrdinalsGroupingOperatorFactory(
         int docChannel,
         String groupingField,
         List<Factory> aggregators,
-        int maxPageSize
+        int maxPageSize,
+        double storedFieldsSequentialProportion
     ) implements OperatorFactory {
 
         @Override
@@ -76,6 +77,7 @@ public Operator get(DriverContext driverContext) {
                 groupingField,
                 aggregators,
                 maxPageSize,
+                storedFieldsSequentialProportion,
                 driverContext
             );
         }
@@ -94,6 +96,7 @@ public String describe() {
     private final List<Factory> aggregatorFactories;
     private final ElementType groupingElementType;
     private final Map<SegmentID, OrdinalSegmentAggregator> ordinalAggregators;
+    private final double storedFieldsSequentialProportion;
 
     private final DriverContext driverContext;
 
@@ -111,6 +114,7 @@ public OrdinalsGroupingOperator(
         String groupingField,
         List<GroupingAggregator.Factory> aggregatorFactories,
         int maxPageSize,
+        double storedFieldsSequentialProportion,
         DriverContext driverContext
     ) {
         Objects.requireNonNull(aggregatorFactories);
@@ -122,6 +126,7 @@ public OrdinalsGroupingOperator(
         this.aggregatorFactories = aggregatorFactories;
         this.ordinalAggregators = new HashMap<>();
         this.maxPageSize = maxPageSize;
+        this.storedFieldsSequentialProportion = storedFieldsSequentialProportion;
         this.driverContext = driverContext;
     }
 
@@ -171,6 +176,7 @@ public void addInput(Page page) {
                         channelIndex,
                         aggregatorFactories,
                         maxPageSize,
+                        storedFieldsSequentialProportion,
                         driverContext
                     );
                 }
@@ -485,6 +491,7 @@ boolean next() throws IOException {
     private static class ValuesAggregator implements Releasable {
         private final ValuesSourceReaderOperator extractor;
         private final HashAggregationOperator aggregator;
+        private final double storedFieldsSequentialProportion;
 
         ValuesAggregator(
             IntFunction<BlockLoader> blockLoaders,
@@ -495,13 +502,15 @@ private static class ValuesAggregator implements Releasable {
             int channelIndex,
             List<GroupingAggregator.Factory> aggregatorFactories,
             int maxPageSize,
+            double storedFieldsSequentialProportion,
             DriverContext driverContext
         ) {
             this.extractor = new ValuesSourceReaderOperator(
                 driverContext.blockFactory(),
                 List.of(new ValuesSourceReaderOperator.FieldInfo(groupingField, groupingElementType, blockLoaders)),
                 shardContexts,
-                docChannel
+                docChannel,
+                storedFieldsSequentialProportion
             );
             this.aggregator = new HashAggregationOperator(
                 aggregatorFactories,
@@ -513,6 +522,7 @@ private static class ValuesAggregator implements Releasable {
                 ),
                 driverContext
             );
+            this.storedFieldsSequentialProportion = storedFieldsSequentialProportion;
         }
 
         void addInput(Page page) {

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

@@ -204,6 +204,7 @@ public String toString() {
                         gField,
                         List.of(CountAggregatorFunction.supplier().groupingAggregatorFactory(INITIAL, List.of(1))),
                         randomPageSize(),
+                        0.1,
                         driverContext
                     )
                 );

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

@@ -209,7 +209,8 @@ private List<Page> runQuery(Set<String> values, Query query, boolean shuffleDocs
                     List.of(new ValuesSourceReaderOperator.ShardContext(reader, () -> {
                         throw new UnsupportedOperationException();
                     })),
-                    0
+                    0,
+                    0.1
                 )
             );
             LuceneQueryEvaluator.ShardConfig[] shards = new LuceneQueryEvaluator.ShardConfig[] {

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

@@ -108,6 +108,7 @@
 import java.util.stream.IntStream;
 import java.util.stream.LongStream;
 
+import static org.elasticsearch.compute.lucene.ValuesSourceReaderOperatorTests.STORED_FIELDS_SEQUENTIAL_PROPORTIONS;
 import static org.elasticsearch.test.MapMatcher.assertMap;
 import static org.elasticsearch.test.MapMatcher.matchesMap;
 import static org.elasticsearch.xpack.esql.core.type.DataType.IP;
@@ -239,7 +240,7 @@ private static Operator.OperatorFactory factory(
                 fail("unexpected shardIdx [" + shardIdx + "]");
             }
             return loader;
-        })), shardContexts, 0);
+        })), shardContexts, 0, STORED_FIELDS_SEQUENTIAL_PROPORTIONS);
     }
 
     protected SourceOperator simpleInput(DriverContext context, int size) {
@@ -488,7 +489,8 @@ public void testManySingleDocPages() {
             new ValuesSourceReaderOperator.Factory(
                 List.of(testCase.info, fieldInfo(mapperService(indexKey).fieldType("key"), ElementType.INT)),
                 shardContexts,
-                0
+                0,
+                STORED_FIELDS_SEQUENTIAL_PROPORTIONS
             ).get(driverContext)
         );
         List<Page> results = drive(operators, input.iterator(), driverContext);
@@ -598,7 +600,8 @@ private void loadSimpleAndAssert(
                     fieldInfo(mapperService("index1").fieldType("indexKey"), ElementType.BYTES_REF)
                 ),
                 shardContexts,
-                0
+                0,
+                STORED_FIELDS_SEQUENTIAL_PROPORTIONS
             ).get(driverContext)
         );
         List<FieldCase> tests = new ArrayList<>();
@@ -607,7 +610,12 @@ private void loadSimpleAndAssert(
             cases.removeAll(b);
             tests.addAll(b);
             operators.add(
-                new ValuesSourceReaderOperator.Factory(b.stream().map(i -> i.info).toList(), shardContexts, 0).get(driverContext)
+                new ValuesSourceReaderOperator.Factory(
+                    b.stream().map(i -> i.info).toList(),
+                    shardContexts,
+                    0,
+                    STORED_FIELDS_SEQUENTIAL_PROPORTIONS
+                ).get(driverContext)
             );
         }
         List<Page> results = drive(operators, input.iterator(), driverContext);
@@ -709,7 +717,11 @@ private void testLoadAllStatus(boolean allInOnePage) {
             Block.MvOrdering.DEDUPLICATED_AND_SORTED_ASCENDING
         );
         List<Operator> operators = cases.stream()
-            .map(i -> new ValuesSourceReaderOperator.Factory(List.of(i.info), shardContexts, 0).get(driverContext))
+            .map(
+                i -> new ValuesSourceReaderOperator.Factory(List.of(i.info), shardContexts, 0, STORED_FIELDS_SEQUENTIAL_PROPORTIONS).get(
+                    driverContext
+                )
+            )
             .toList();
         if (allInOnePage) {
             input = List.of(CannedSourceOperator.mergePages(input));
@@ -1385,7 +1397,8 @@ public void testNullsShared() {
                             new ValuesSourceReaderOperator.FieldInfo("null2", ElementType.NULL, shardIdx -> BlockLoader.CONSTANT_NULLS)
                         ),
                         shardContexts,
-                        0
+                        0,
+                        STORED_FIELDS_SEQUENTIAL_PROPORTIONS
                     ).get(driverContext)
                 ),
                 new PageConsumerOperator(page -> {
@@ -1416,7 +1429,8 @@ public void testDescriptionOfMany() throws IOException {
         ValuesSourceReaderOperator.Factory factory = new ValuesSourceReaderOperator.Factory(
             cases.stream().map(c -> c.info).toList(),
             List.of(new ValuesSourceReaderOperator.ShardContext(reader(indexKey), () -> SourceLoader.FROM_STORED_SOURCE)),
-            0
+            0,
+            STORED_FIELDS_SEQUENTIAL_PROPORTIONS
         );
         assertThat(factory.describe(), equalTo("ValuesSourceReaderOperator[fields = [" + cases.size() + " fields]]"));
         try (Operator op = factory.get(driverContext())) {
@@ -1462,7 +1476,8 @@ public void testManyShards() throws IOException {
                     return ft.blockLoader(blContext());
                 })),
                 readerShardContexts,
-                0
+                0,
+                STORED_FIELDS_SEQUENTIAL_PROPORTIONS
             );
             DriverContext driverContext = driverContext();
             List<Page> results = drive(

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

@@ -114,6 +114,8 @@ public class ValuesSourceReaderOperatorTests extends OperatorTestCase {
         { false, true, true },
         { false, false, true, true } };
 
+    static final double STORED_FIELDS_SEQUENTIAL_PROPORTIONS = 0.1;
+
     private Directory directory = newDirectory();
     private MapperService mapperService;
     private IndexReader reader;
@@ -147,7 +149,11 @@ static Operator.OperatorFactory factory(IndexReader reader, String name, Element
                 fail("unexpected shardIdx [" + shardIdx + "]");
             }
             return loader;
-        })), List.of(new ValuesSourceReaderOperator.ShardContext(reader, () -> SourceLoader.FROM_STORED_SOURCE)), 0);
+        })),
+            List.of(new ValuesSourceReaderOperator.ShardContext(reader, () -> SourceLoader.FROM_STORED_SOURCE)),
+            0,
+            STORED_FIELDS_SEQUENTIAL_PROPORTIONS
+        );
     }
 
     @Override
@@ -444,7 +450,8 @@ public void testManySingleDocPages() {
             new ValuesSourceReaderOperator.Factory(
                 List.of(testCase.info, fieldInfo(mapperService.fieldType("key"), ElementType.INT)),
                 List.of(new ValuesSourceReaderOperator.ShardContext(reader, () -> SourceLoader.FROM_STORED_SOURCE)),
-                0
+                0,
+                STORED_FIELDS_SEQUENTIAL_PROPORTIONS
             ).get(driverContext)
         );
         List<Page> results = drive(operators, input.iterator(), driverContext);
@@ -550,7 +557,8 @@ private void loadSimpleAndAssert(
             new ValuesSourceReaderOperator.Factory(
                 List.of(fieldInfo(mapperService.fieldType("key"), ElementType.INT)),
                 List.of(new ValuesSourceReaderOperator.ShardContext(reader, () -> SourceLoader.FROM_STORED_SOURCE)),
-                0
+                0,
+                STORED_FIELDS_SEQUENTIAL_PROPORTIONS
             ).get(driverContext)
         );
         List<FieldCase> tests = new ArrayList<>();
@@ -562,7 +570,8 @@ private void loadSimpleAndAssert(
                 new ValuesSourceReaderOperator.Factory(
                     b.stream().map(i -> i.info).toList(),
                     List.of(new ValuesSourceReaderOperator.ShardContext(reader, () -> SourceLoader.FROM_STORED_SOURCE)),
-                    0
+                    0,
+                    STORED_FIELDS_SEQUENTIAL_PROPORTIONS
                 ).get(driverContext)
             );
         }
@@ -652,7 +661,8 @@ private void testLoadAllStatus(boolean allInOnePage) {
                 i -> new ValuesSourceReaderOperator.Factory(
                     List.of(i.info),
                     List.of(new ValuesSourceReaderOperator.ShardContext(reader, () -> SourceLoader.FROM_STORED_SOURCE)),
-                    0
+                    0,
+                    STORED_FIELDS_SEQUENTIAL_PROPORTIONS
                 ).get(driverContext)
             )
             .toList();
@@ -1418,7 +1428,8 @@ public void testNullsShared() {
                             new ValuesSourceReaderOperator.FieldInfo("null2", ElementType.NULL, shardIdx -> BlockLoader.CONSTANT_NULLS)
                         ),
                         List.of(new ValuesSourceReaderOperator.ShardContext(reader, () -> SourceLoader.FROM_STORED_SOURCE)),
-                        0
+                        0,
+                        STORED_FIELDS_SEQUENTIAL_PROPORTIONS
                     ).get(driverContext)
                 ),
                 new PageConsumerOperator(page -> {
@@ -1463,7 +1474,8 @@ private void testSequentialStoredFields(boolean sequential, int docCount) throws
                 fieldInfo(storedTextField("stored_text"), ElementType.BYTES_REF)
             ),
             List.of(new ValuesSourceReaderOperator.ShardContext(reader, () -> SourceLoader.FROM_STORED_SOURCE)),
-            0
+            0,
+            STORED_FIELDS_SEQUENTIAL_PROPORTIONS
         ).get(driverContext);
         List<Page> results = drive(op, source.iterator(), driverContext);
         Checks checks = new Checks(Block.MvOrdering.UNORDERED, Block.MvOrdering.UNORDERED);
@@ -1491,7 +1503,8 @@ public void testDescriptionOfMany() throws IOException {
         ValuesSourceReaderOperator.Factory factory = new ValuesSourceReaderOperator.Factory(
             cases.stream().map(c -> c.info).toList(),
             List.of(new ValuesSourceReaderOperator.ShardContext(reader, () -> SourceLoader.FROM_STORED_SOURCE)),
-            0
+            0,
+            STORED_FIELDS_SEQUENTIAL_PROPORTIONS
         );
         assertThat(factory.describe(), equalTo("ValuesSourceReaderOperator[fields = [" + cases.size() + " fields]]"));
         try (Operator op = factory.get(driverContext())) {
@@ -1535,7 +1548,8 @@ public void testManyShards() throws IOException {
                     return ft.blockLoader(blContext());
                 })),
                 readerShardContexts,
-                0
+                0,
+                STORED_FIELDS_SEQUENTIAL_PROPORTIONS
             );
             DriverContext driverContext = driverContext();
             List<Page> results = drive(