Improve reading ignored source in esql

martijnvg · martijnvg · commit 08f59c3c7cbd · 2025-07-06T21:11:20.000+02:00
* Parse sub field fields once in IgnoredSourceRowStrideReader instead of for each doc that gets read.
* Introduce a dedicated StoredFieldLoader for ignored source (IgnoredSourceFieldLoader). Which optimizes reading stored fields just for ignored source by avoiding relatively expensive set stuff (see CustomFieldsVisitor) and make use aborting loading stored fields (`Status.STOP`) when ignored source is read.
diff --git a/server/src/main/java/org/elasticsearch/index/fieldvisitor/IgnoredSourceFieldLoader.java b/server/src/main/java/org/elasticsearch/index/fieldvisitor/IgnoredSourceFieldLoader.java
@@ -0,0 +1,118 @@
+/*
+ * 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.fieldvisitor;
+
+import org.apache.lucene.index.FieldInfo;
+import org.apache.lucene.index.LeafReader;
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.index.StoredFieldVisitor;
+import org.apache.lucene.util.BytesRef;
+import org.elasticsearch.common.CheckedBiConsumer;
+import org.elasticsearch.common.bytes.BytesReference;
+import org.elasticsearch.common.lucene.index.SequentialStoredFieldsLeafReader;
+import org.elasticsearch.index.mapper.IgnoredSourceFieldMapper;
+import org.elasticsearch.search.fetch.StoredFieldsSpec;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+
+class IgnoredSourceFieldLoader extends StoredFieldLoader {
+
+    @Override
+    public LeafStoredFieldLoader getLoader(LeafReaderContext ctx, int[] docs) throws IOException {
+        var reader = sequentialReader(ctx);
+        var visitor = new SFV();
+        return new LeafStoredFieldLoader() {
+
+            private int doc = -1;
+
+            @Override
+            public void advanceTo(int doc) throws IOException {
+                if (doc != this.doc) {
+                    visitor.reset();
+                    reader.accept(doc, visitor);
+                    this.doc = doc;
+                }
+            }
+
+            @Override
+            public BytesReference source() {
+                return null;
+            }
+
+            @Override
+            public String id() {
+                return null;
+            }
+
+            @Override
+            public String routing() {
+                return null;
+            }
+
+            @Override
+            public Map<String, List<Object>> storedFields() {
+                return Map.of(IgnoredSourceFieldMapper.NAME, visitor.values);
+            }
+        };
+    }
+
+    @Override
+    public List<String> fieldsToLoad() {
+        return List.of(IgnoredSourceFieldMapper.NAME);
+    }
+
+    static class SFV extends StoredFieldVisitor {
+
+        boolean processing;
+        final List<Object> values = new ArrayList<>();
+
+        @Override
+        public Status needsField(FieldInfo fieldInfo) throws IOException {
+            if (IgnoredSourceFieldMapper.NAME.equals(fieldInfo.name)) {
+                processing = true;
+                return Status.YES;
+            } else if (processing) {
+                return Status.STOP;
+            }
+
+            return Status.NO;
+        }
+
+        @Override
+        public void binaryField(FieldInfo fieldInfo, byte[] value) throws IOException {
+            values.add(new BytesRef(value));
+        }
+
+        void reset() {
+            values.clear();
+            processing = false;
+        }
+
+    }
+
+    static boolean supports(StoredFieldsSpec spec) {
+        return spec.requiresSource() == false
+            && spec.requiresMetadata() == false
+            && spec.requiredStoredFields().size() == 1
+            && spec.requiredStoredFields().contains(IgnoredSourceFieldMapper.NAME);
+    }
+
+    // TODO: use provided one
+    private static CheckedBiConsumer<Integer, StoredFieldVisitor, IOException> sequentialReader(LeafReaderContext ctx) throws IOException {
+        LeafReader leafReader = ctx.reader();
+        if (leafReader instanceof SequentialStoredFieldsLeafReader lf) {
+            return lf.getSequentialStoredFieldsReader()::document;
+        }
+        return leafReader.storedFields()::document;
+    }
+}
diff --git a/server/src/main/java/org/elasticsearch/index/fieldvisitor/StoredFieldLoader.java b/server/src/main/java/org/elasticsearch/index/fieldvisitor/StoredFieldLoader.java
@@ -50,6 +50,9 @@ public static StoredFieldLoader fromSpec(StoredFieldsSpec spec) {
         if (spec.noRequirements()) {
             return StoredFieldLoader.empty();
         }
+        if (IgnoredSourceFieldLoader.supports(spec)) {
+            return new IgnoredSourceFieldLoader();
+        }
         return create(spec.requiresSource(), spec.requiredStoredFields());
     }
 
@@ -91,6 +94,10 @@ public static StoredFieldLoader fromSpecSequential(StoredFieldsSpec spec) {
         if (spec.noRequirements()) {
             return StoredFieldLoader.empty();
         }
+        if (IgnoredSourceFieldLoader.supports(spec)) {
+            return new IgnoredSourceFieldLoader();
+        }
+
         List<String> fieldsToLoad = fieldsToLoad(spec.requiresSource(), spec.requiredStoredFields());
         return new StoredFieldLoader() {
             @Override
diff --git a/server/src/main/java/org/elasticsearch/index/mapper/FallbackSyntheticSourceBlockLoader.java b/server/src/main/java/org/elasticsearch/index/mapper/FallbackSyntheticSourceBlockLoader.java
@@ -21,6 +21,7 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
 import java.util.Stack;
@@ -39,10 +40,12 @@
 public abstract class FallbackSyntheticSourceBlockLoader implements BlockLoader {
     private final Reader<?> reader;
     private final String fieldName;
+    private final Set<String> fieldPaths;
 
     protected FallbackSyntheticSourceBlockLoader(Reader<?> reader, String fieldName) {
         this.reader = reader;
         this.fieldName = fieldName;
+        this.fieldPaths = splitIntoFieldPaths(fieldName);
     }
 
     @Override
@@ -52,7 +55,7 @@ public ColumnAtATimeReader columnAtATimeReader(LeafReaderContext context) throws
 
     @Override
     public RowStrideReader rowStrideReader(LeafReaderContext context) throws IOException {
-        return new IgnoredSourceRowStrideReader<>(fieldName, reader);
+        return new IgnoredSourceRowStrideReader<>(fieldName, reader, fieldPaths);
     }
 
     @Override
@@ -70,7 +73,31 @@ public SortedSetDocValues ordinals(LeafReaderContext context) throws IOException
         throw new UnsupportedOperationException();
     }
 
-    private record IgnoredSourceRowStrideReader<T>(String fieldName, Reader<T> reader) implements RowStrideReader {
+    static Set<String> splitIntoFieldPaths(String fieldName) {
+        var paths = new HashSet<String>();
+        paths.add("_doc");
+        var current = new StringBuilder();
+        for (var part : fieldName.split("\\.")) {
+            if (current.isEmpty() == false) {
+                current.append('.');
+            }
+            current.append(part);
+            paths.add(current.toString());
+        }
+        return paths;
+    }
+
+    private static final class IgnoredSourceRowStrideReader<T> implements RowStrideReader {
+        private final String fieldName;
+        private final Reader<T> reader;
+        private final Set<String> fieldPaths;
+
+        private IgnoredSourceRowStrideReader(String fieldName, Reader<T> reader, Set<String> fieldPaths) {
+            this.fieldName = fieldName;
+            this.reader = reader;
+            this.fieldPaths = fieldPaths;
+        }
+
         @Override
         public void read(int docId, StoredFields storedFields, Builder builder) throws IOException {
             var ignoredSource = storedFields.storedFields().get(IgnoredSourceFieldMapper.NAME);
@@ -79,28 +106,11 @@ public void read(int docId, StoredFields storedFields, Builder builder) throws I
                 return;
             }
 
-            Map<String, List<IgnoredSourceFieldMapper.NameValue>> valuesForFieldAndParents = new HashMap<>();
-
-            // Contains name of the field and all its parents
-            Set<String> fieldNames = new HashSet<>() {
-                {
-                    add("_doc");
-                }
-            };
-
-            var current = new StringBuilder();
-            for (String part : fieldName.split("\\.")) {
-                if (current.isEmpty() == false) {
-                    current.append('.');
-                }
-                current.append(part);
-                fieldNames.add(current.toString());
-            }
-
+            Map<String, IgnoredSourceFieldMapper.NameValue> valuesForFieldAndParents = new HashMap<>();
             for (Object value : ignoredSource) {
                 IgnoredSourceFieldMapper.NameValue nameValue = IgnoredSourceFieldMapper.decode(value);
-                if (fieldNames.contains(nameValue.name())) {
-                    valuesForFieldAndParents.computeIfAbsent(nameValue.name(), k -> new ArrayList<>()).add(nameValue);
+                if (fieldPaths.contains(nameValue.name())) {
+                    valuesForFieldAndParents.putIfAbsent(nameValue.name(), nameValue);
                 }
             }
 
@@ -130,41 +140,35 @@ public void read(int docId, StoredFields storedFields, Builder builder) throws I
             }
         }
 
-        private void readFromFieldValue(List<IgnoredSourceFieldMapper.NameValue> nameValues, List<T> blockValues) throws IOException {
-            if (nameValues.isEmpty()) {
+        private void readFromFieldValue(IgnoredSourceFieldMapper.NameValue nameValue, List<T> blockValues) throws IOException {
+            // Leaf field is stored directly (not as a part of a parent object), let's try to decode it.
+            Optional<Object> singleValue = XContentDataHelper.decode(nameValue.value());
+            if (singleValue.isPresent()) {
+                reader.convertValue(singleValue.get(), blockValues);
                 return;
             }
 
-            for (var nameValue : nameValues) {
-                // Leaf field is stored directly (not as a part of a parent object), let's try to decode it.
-                Optional<Object> singleValue = XContentDataHelper.decode(nameValue.value());
-                if (singleValue.isPresent()) {
-                    reader.convertValue(singleValue.get(), blockValues);
-                    continue;
-                }
+            // We have a value for this field but it's an array or an object
+            var type = XContentDataHelper.decodeType(nameValue.value());
+            assert type.isPresent();
 
-                // We have a value for this field but it's an array or an object
-                var type = XContentDataHelper.decodeType(nameValue.value());
-                assert type.isPresent();
-
-                try (
-                    XContentParser parser = type.get()
-                        .xContent()
-                        .createParser(
-                            XContentParserConfiguration.EMPTY,
-                            nameValue.value().bytes,
-                            nameValue.value().offset + 1,
-                            nameValue.value().length - 1
-                        )
-                ) {
-                    parser.nextToken();
-                    parseWithReader(parser, blockValues);
-                }
+            try (
+                XContentParser parser = type.get()
+                    .xContent()
+                    .createParser(
+                        XContentParserConfiguration.EMPTY,
+                        nameValue.value().bytes,
+                        nameValue.value().offset + 1,
+                        nameValue.value().length - 1
+                    )
+            ) {
+                parser.nextToken();
+                parseWithReader(parser, blockValues);
             }
         }
 
         private void readFromParentValue(
-            Map<String, List<IgnoredSourceFieldMapper.NameValue>> valuesForFieldAndParents,
+            Map<String, IgnoredSourceFieldMapper.NameValue> valuesForFieldAndParents,
             List<T> blockValues
         ) throws IOException {
             if (valuesForFieldAndParents.isEmpty()) {
@@ -174,11 +178,8 @@ private void readFromParentValue(
             // If a parent object is stored at a particular level its children won't be stored.
             // So we should only ever have one parent here.
             assert valuesForFieldAndParents.size() == 1 : "_ignored_source field contains multiple levels of the same object";
-            var parentValues = valuesForFieldAndParents.values().iterator().next();
-
-            for (var nameValue : parentValues) {
-                parseFieldFromParent(nameValue, blockValues);
-            }
+            var parentValue = valuesForFieldAndParents.values().iterator().next();
+            parseFieldFromParent(parentValue, blockValues);
         }
 
         private void parseFieldFromParent(IgnoredSourceFieldMapper.NameValue nameValue, List<T> blockValues) throws IOException {
@@ -243,6 +244,33 @@ private void parseWithReader(XContentParser parser, List<T> blockValues) throws
         public boolean canReuse(int startingDocID) {
             return true;
         }
+
+        public String fieldName() {
+            return fieldName;
+        }
+
+        public Reader<T> reader() {
+            return reader;
+        }
+
+        @Override
+        public boolean equals(Object obj) {
+            if (obj == this) return true;
+            if (obj == null || obj.getClass() != this.getClass()) return false;
+            var that = (IgnoredSourceRowStrideReader) obj;
+            return Objects.equals(this.fieldName, that.fieldName) && Objects.equals(this.reader, that.reader);
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(fieldName, reader);
+        }
+
+        @Override
+        public String toString() {
+            return "IgnoredSourceRowStrideReader[" + "fieldName=" + fieldName + ", " + "reader=" + reader + ']';
+        }
+
     }
 
     /**

Original file line number	Diff line number	Diff line change
`@@ -50,6 +50,9 @@ public static StoredFieldLoader fromSpec(StoredFieldsSpec spec) {`
`50`	`50`	`if (spec.noRequirements()) {`
`51`	`51`	`return StoredFieldLoader.empty();`
`52`	`52`	`}`
	`53`	`+ if (IgnoredSourceFieldLoader.supports(spec)) {`
	`54`	`+ return new IgnoredSourceFieldLoader();`
	`55`	`+ }`
`53`	`56`	`return create(spec.requiresSource(), spec.requiredStoredFields());`
`54`	`57`	`}`
`55`	`58`
`@@ -91,6 +94,10 @@ public static StoredFieldLoader fromSpecSequential(StoredFieldsSpec spec) {`
`91`	`94`	`if (spec.noRequirements()) {`
`92`	`95`	`return StoredFieldLoader.empty();`
`93`	`96`	`}`
	`97`	`+ if (IgnoredSourceFieldLoader.supports(spec)) {`
	`98`	`+ return new IgnoredSourceFieldLoader();`
	`99`	`+ }`
	`100`	`+`
`94`	`101`	`List<String> fieldsToLoad = fieldsToLoad(spec.requiresSource(), spec.requiredStoredFields());`
`95`	`102`	`return new StoredFieldLoader() {`
`96`	`103`	`@Override`