org.elasticsearch.index.fieldvisitor.StoredFieldLoader Maven / Gradle / Ivy

Show more of this group Show more artifacts with this name
Show all versions of elasticsearch Show documentation
Elasticsearch subproject :server
There is a newer version: 8.14.0
/*
 * 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 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 or the Server
 * Side Public License, v 1.
 */

package org.elasticsearch.index.fieldvisitor;

import org.apache.lucene.index.LeafReader;
import org.apache.lucene.index.LeafReaderContext;
import org.elasticsearch.common.CheckedBiConsumer;
import org.elasticsearch.common.bytes.BytesReference;
import org.elasticsearch.common.lucene.index.SequentialStoredFieldsLeafReader;

import java.io.IOException;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

/**
 * Generates a {@link LeafStoredFieldLoader} for a given lucene segment to load stored fields.
 */
public abstract class StoredFieldLoader {

    /**
     * Return a {@link LeafStoredFieldLoader} for the given segment and document set
     *
     * The loader will use an internal lucene merge reader if the document set is of
     * sufficient size and is contiguous.  Callers may pass {@code null} if the set
     * is not known up front or if the merge reader optimisation will not apply.
     */
    public abstract LeafStoredFieldLoader getLoader(LeafReaderContext ctx, int[] docs);

    /**
     * @return a list of fields that will be loaded for each document
     */
    public abstract List fieldsToLoad();

    /**
     * Creates a new StoredFieldLoader
     * @param loadSource should this loader load the _source field
     * @param fields     a set of additional fields the loader should load
     */
    public static StoredFieldLoader create(boolean loadSource, Set fields) {
        List fieldsToLoad = fieldsToLoad(loadSource, fields);
        return new StoredFieldLoader() {
            @Override
            public LeafStoredFieldLoader getLoader(LeafReaderContext ctx, int[] docs) {
                return new ReaderStoredFieldLoader(reader(ctx, docs), loadSource, fields);
            }

            @Override
            public List fieldsToLoad() {
                return fieldsToLoad;
            }
        };
    }

    /**
     * Creates a no-op StoredFieldLoader that will not load any fields from disk
     */
    public static StoredFieldLoader empty() {
        return new StoredFieldLoader() {
            @Override
            public LeafStoredFieldLoader getLoader(LeafReaderContext ctx, int[] docs) {
                return new EmptyStoredFieldLoader();
            }

            @Override
            public List fieldsToLoad() {
                return List.of();
            }
        };
    }

    private static CheckedBiConsumer reader(LeafReaderContext ctx, int[] docs) {
        LeafReader leafReader = ctx.reader();
        if (docs == null) {
            return leafReader::document;
        }
        if (leafReader instanceof SequentialStoredFieldsLeafReader lf && docs.length > 10 && hasSequentialDocs(docs)) {
            return lf.getSequentialStoredFieldsReader()::visitDocument;
        }
        return leafReader::document;
    }

    private static List fieldsToLoad(boolean loadSource, Set fields) {
        Set fieldsToLoad = new HashSet<>();
        fieldsToLoad.add("_id");
        fieldsToLoad.add("_routing");
        if (loadSource) {
            fieldsToLoad.add("_source");
        }
        fieldsToLoad.addAll(fields);
        return fieldsToLoad.stream().sorted().toList();
    }

    private static boolean hasSequentialDocs(int[] docs) {
        return docs.length > 0 && docs[docs.length - 1] - docs[0] == docs.length - 1;
    }

    private static class EmptyStoredFieldLoader implements LeafStoredFieldLoader {

        @Override
        public void advanceTo(int doc) throws IOException {}

        @Override
        public BytesReference source() {
            return null;
        }

        @Override
        public String id() {
            return null;
        }

        @Override
        public String routing() {
            return null;
        }

        @Override
        public Map> storedFields() {
            return Collections.emptyMap();
        }
    }

    private static class ReaderStoredFieldLoader implements LeafStoredFieldLoader {

        private final CheckedBiConsumer reader;
        private final CustomFieldsVisitor visitor;
        private int doc = -1;

        ReaderStoredFieldLoader(CheckedBiConsumer reader, boolean loadSource, Set fields) {
            this.reader = reader;
            this.visitor = new CustomFieldsVisitor(fields, loadSource);
        }

        @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 visitor.source();
        }

        @Override
        public String id() {
            return visitor.id();
        }

        @Override
        public String routing() {
            return visitor.routing();
        }

        @Override
        public Map> storedFields() {
            return visitor.fields();
        }
    }

}