ESQL: Speed loading stored fields (#127348)

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.EsqlPlugin;
 import org.openjdk.jmh.annotations.Benchmark;
 import org.openjdk.jmh.annotations.BenchmarkMode;
 import org.openjdk.jmh.annotations.Fork;
@@ -296,7 +298,7 @@ public void benchmark() {
  fields(name),
  List.of(new ValuesSourceReaderOperator.ShardContext(reader, () -> {
  throw new UnsupportedOperationException("can't load _source here");
- })),
+ }, EsqlPlugin.STORED_FIELDS_SEQUENTIAL_PROPORTION.getDefault(Settings.EMPTY))),
  0
  );
  long sum = 0;

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

@@ -105,7 +105,7 @@ public String describe() {
  */
  public record FieldInfo(String name, ElementType type, IntFunction<BlockLoader> blockLoader) {}
 
- public record ShardContext(IndexReader reader, Supplier<SourceLoader> newSourceLoader) {}
+ public record ShardContext(IndexReader reader, Supplier<SourceLoader> newSourceLoader, double storedFieldsSequentialProportion) {}
 
  private final FieldWork[] fields;
  private final List<ShardContext> shardContexts;
@@ -241,8 +241,9 @@ private void loadFromSingleLeaf(Block[] blocks, int shard, int segment, BlockLoa
  }
 
  SourceLoader sourceLoader = null;
+ ShardContext shardContext = shardContexts.get(shard);
  if (storedFieldsSpec.requiresSource()) {
- sourceLoader = shardContexts.get(shard).newSourceLoader.get();
+ sourceLoader = shardContext.newSourceLoader.get();
  storedFieldsSpec = storedFieldsSpec.merge(new StoredFieldsSpec(true, false, sourceLoader.requiredStoredFields()));
  }
 
@@ -255,7 +256,7 @@ private void loadFromSingleLeaf(Block[] blocks, int shard, int segment, BlockLoa
  );
  }
  StoredFieldLoader storedFieldLoader;
- if (useSequentialStoredFieldsReader(docs)) {
+ if (useSequentialStoredFieldsReader(docs, shardContext.storedFieldsSequentialProportion())) {
  storedFieldLoader = StoredFieldLoader.fromSpecSequential(storedFieldsSpec);
  trackStoredFields(storedFieldsSpec, true);
  } else {
@@ -432,9 +433,13 @@ public void close() {
  * Is it more efficient to use a sequential stored field reader
  * when reading stored fields for the documents contained in {@code docIds}?
  */
- private boolean useSequentialStoredFieldsReader(BlockLoader.Docs docs) {
+ private boolean useSequentialStoredFieldsReader(BlockLoader.Docs docs, double storedFieldsSequentialProportion) {
  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;
  }
 
  private void trackStoredFields(StoredFieldsSpec spec, boolean sequential) {

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

@@ -193,7 +193,7 @@ public String toString() {
  operators.add(
  new OrdinalsGroupingOperator(
  shardIdx -> new KeywordFieldMapper.KeywordFieldType("g").blockLoader(null),
- List.of(new ValuesSourceReaderOperator.ShardContext(reader, () -> SourceLoader.FROM_STORED_SOURCE)),
+ List.of(new ValuesSourceReaderOperator.ShardContext(reader, () -> SourceLoader.FROM_STORED_SOURCE, 0.2)),
  ElementType.BYTES_REF,
  0,
  gField,

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

@@ -207,7 +207,7 @@ private List<Page> runQuery(Set<String> values, Query query, boolean shuffleDocs
  ),
  List.of(new ValuesSourceReaderOperator.ShardContext(reader, () -> {
  throw new UnsupportedOperationException();
- })),
+ }, 0.2)),
  0
  )
  );