From 9b06d06a79a57676cbb2002ce59866cef8243414 Mon Sep 17 00:00:00 2001 From: Aurelien FOUCRET Date: Wed, 17 Jul 2024 16:26:34 +0200 Subject: [PATCH] Additional refactoring. --- .../search/function/ScriptScoreQuery.java | 3 +- .../index/query/SearchExecutionContext.java | 4 +- .../org/elasticsearch/script/ScoreScript.java | 17 ++++++--- .../elasticsearch/script/TermStatsReader.java | 7 +--- .../search/lookup/SearchLookup.java | 37 ++++++++----------- 5 files changed, 32 insertions(+), 36 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/common/lucene/search/function/ScriptScoreQuery.java b/server/src/main/java/org/elasticsearch/common/lucene/search/function/ScriptScoreQuery.java index 90e6bae244e0b..1fcff2ee32227 100644 --- a/server/src/main/java/org/elasticsearch/common/lucene/search/function/ScriptScoreQuery.java +++ b/server/src/main/java/org/elasticsearch/common/lucene/search/function/ScriptScoreQuery.java @@ -30,6 +30,7 @@ import org.elasticsearch.script.ScoreScript; import org.elasticsearch.script.ScoreScript.ExplanationHolder; import org.elasticsearch.script.Script; +import org.elasticsearch.script.TermStatsReader; import org.elasticsearch.search.lookup.SearchLookup; import java.io.IOException; @@ -180,7 +181,7 @@ private ScoreScript makeScoreScript(LeafReaderContext context) throws IOExceptio scoreScript._setIndexName(indexName); scoreScript._setShard(shardId); if (needsTermStatistics) { - scoreScript.get_termStatistics().setTerms(terms); + scoreScript._setTermStats(new TermStatsReader(searcher, scoreScript::_getDocId, context, terms)); } return scoreScript; } diff --git a/server/src/main/java/org/elasticsearch/index/query/SearchExecutionContext.java b/server/src/main/java/org/elasticsearch/index/query/SearchExecutionContext.java index 553ab858d3418..4b16a5833cf41 100644 --- a/server/src/main/java/org/elasticsearch/index/query/SearchExecutionContext.java +++ b/server/src/main/java/org/elasticsearch/index/query/SearchExecutionContext.java @@ -52,7 +52,6 @@ import org.elasticsearch.script.ScriptCompiler; import org.elasticsearch.script.ScriptContext; import org.elasticsearch.script.ScriptFactory; -import org.elasticsearch.script.TermStatsReader; import org.elasticsearch.search.NestedDocuments; import org.elasticsearch.search.aggregations.support.ValuesSourceRegistry; import org.elasticsearch.search.lookup.LeafFieldLookupProvider; @@ -525,8 +524,7 @@ public void setLookupProviders( ) ), sourceProvider, - fieldLookupProvider, - (leafReaderContext, docIdSupplier) -> new TermStatsReader(searcher, docIdSupplier, leafReaderContext) + fieldLookupProvider ); } diff --git a/server/src/main/java/org/elasticsearch/script/ScoreScript.java b/server/src/main/java/org/elasticsearch/script/ScoreScript.java index 524c1c0792e9a..126ac399034d4 100644 --- a/server/src/main/java/org/elasticsearch/script/ScoreScript.java +++ b/server/src/main/java/org/elasticsearch/script/ScoreScript.java @@ -27,7 +27,6 @@ * A script used for adjusting the score on a per document basis. */ public abstract class ScoreScript extends DocBasedScript { - /** A helper to take in an explanation from a script and turn it into an {@link org.apache.lucene.search.Explanation} */ public static class ExplanationHolder { private String description; @@ -83,7 +82,7 @@ public Explanation get(double score, Explanation subQueryExplanation) { private int shardId = -1; private String indexName = null; - private final TermStatsReader termStatsReader; + private TermStatsReader termStatsReader = null; public ScoreScript(Map params, SearchLookup searchLookup, DocReader docReader) { // searchLookup parameter is ignored but part of the ScriptFactory contract. It is part of that contract because it's required @@ -101,8 +100,6 @@ public ScoreScript(Map params, SearchLookup searchLookup, DocRea this.params = new DynamicMap(params, PARAMS_FUNCTIONS); LeafReaderContext leafReaderContext = ((DocValuesDocReader) docReader).getLeafReaderContext(); this.docBase = leafReaderContext.docBase; - this.termStatsReader = searchLookup.getTermStatsReader(leafReaderContext, this::_getDocId); - } } @@ -195,12 +192,22 @@ public void _setIndexName(String indexName) { this.indexName = indexName; } + /** + * Starting a name with underscore, so that the user cannot access this function directly through a script. + */ + public void _setTermStats(TermStatsReader termStatsReader) { + this.termStatsReader = termStatsReader; + } + + /** + * Accessed as _termStatistics in the painless script. + */ public TermStatsReader get_termStatistics() { if (termStatsReader != null) { return termStatsReader; } - throw new IllegalArgumentException("_termStats is not available"); + throw new IllegalArgumentException("_termStatistics is not available"); } /** A factory to construct {@link ScoreScript} instances. */ diff --git a/server/src/main/java/org/elasticsearch/script/TermStatsReader.java b/server/src/main/java/org/elasticsearch/script/TermStatsReader.java index 330b81e5c8b2f..bfe882352c254 100644 --- a/server/src/main/java/org/elasticsearch/script/TermStatsReader.java +++ b/server/src/main/java/org/elasticsearch/script/TermStatsReader.java @@ -31,15 +31,12 @@ public class TermStatsReader { private final Supplier docIdSupplier; private final Map termContexts = new HashMap<>(); private final Map postings = new HashMap<>(); - private Set terms = Set.of(); + private final Set terms; - public TermStatsReader(IndexSearcher searcher, Supplier docIdSupplier, LeafReaderContext leafReaderContext) { + public TermStatsReader(IndexSearcher searcher, Supplier docIdSupplier, LeafReaderContext leafReaderContext, Set terms) { this.searcher = searcher; this.docIdSupplier = docIdSupplier; this.leafReaderContext = leafReaderContext; - } - - public void setTerms(Set terms) { this.terms = terms; } diff --git a/server/src/main/java/org/elasticsearch/search/lookup/SearchLookup.java b/server/src/main/java/org/elasticsearch/search/lookup/SearchLookup.java index a0ce10363873c..5ed40e509976a 100644 --- a/server/src/main/java/org/elasticsearch/search/lookup/SearchLookup.java +++ b/server/src/main/java/org/elasticsearch/search/lookup/SearchLookup.java @@ -12,14 +12,12 @@ import org.elasticsearch.common.TriFunction; import org.elasticsearch.index.fielddata.IndexFieldData; import org.elasticsearch.index.mapper.MappedFieldType; -import org.elasticsearch.script.TermStatsReader; import java.io.IOException; import java.util.Collections; import java.util.LinkedHashSet; import java.util.Objects; import java.util.Set; -import java.util.function.BiFunction; import java.util.function.Function; import java.util.function.Supplier; @@ -54,50 +52,49 @@ public class SearchLookup implements SourceProvider { IndexFieldData> fieldDataLookup; private final Function fieldLookupProvider; - private final BiFunction, TermStatsReader> termStatsReaderFactory; - /** * Create a new SearchLookup, using the default stored fields provider - * @param fieldTypeLookup defines how to look up field types - * @param fieldDataLookup defines how to look up field data - * @param sourceProvider defines how to look up the source + * + * @param fieldTypeLookup defines how to look up field types + * @param fieldDataLookup defines how to look up field data + * @param sourceProvider defines how to look up the source */ public SearchLookup( Function fieldTypeLookup, TriFunction, MappedFieldType.FielddataOperation, IndexFieldData> fieldDataLookup, SourceProvider sourceProvider ) { - this(fieldTypeLookup, fieldDataLookup, sourceProvider, LeafFieldLookupProvider.fromStoredFields(), null); + this(fieldTypeLookup, fieldDataLookup, sourceProvider, LeafFieldLookupProvider.fromStoredFields()); } /** * Create a new SearchLookup, using the default stored fields provider - * @param fieldTypeLookup defines how to look up field types - * @param fieldDataLookup defines how to look up field data - * @param sourceProvider defines how to look up the source - * @param fieldLookupProvider defines how to look up stored fields + * + * @param fieldTypeLookup defines how to look up field types + * @param fieldDataLookup defines how to look up field data + * @param sourceProvider defines how to look up the source + * @param fieldLookupProvider defines how to look up stored fields */ public SearchLookup( Function fieldTypeLookup, TriFunction, MappedFieldType.FielddataOperation, IndexFieldData> fieldDataLookup, SourceProvider sourceProvider, - Function fieldLookupProvider, - BiFunction, TermStatsReader> termStatsReaderFactory + Function fieldLookupProvider ) { this.fieldTypeLookup = fieldTypeLookup; this.fieldChain = Collections.emptySet(); this.sourceProvider = sourceProvider; this.fieldDataLookup = fieldDataLookup; this.fieldLookupProvider = fieldLookupProvider; - this.termStatsReaderFactory = termStatsReaderFactory; } /** * Create a new {@link SearchLookup} that looks fields up the same as the one provided as argument, * while also tracking field references starting from the provided field name. It detects cycles * and prevents resolving fields that depend on more than {@link #MAX_FIELD_CHAIN_DEPTH} fields. + * * @param searchLookup the existing lookup to create a new one from - * @param fieldChain the chain of fields that required the field currently being loaded + * @param fieldChain the chain of fields that required the field currently being loaded */ private SearchLookup(SearchLookup searchLookup, Set fieldChain) { this.fieldChain = Collections.unmodifiableSet(fieldChain); @@ -105,16 +102,16 @@ private SearchLookup(SearchLookup searchLookup, Set fieldChain) { this.fieldTypeLookup = searchLookup.fieldTypeLookup; this.fieldDataLookup = searchLookup.fieldDataLookup; this.fieldLookupProvider = searchLookup.fieldLookupProvider; - this.termStatsReaderFactory = searchLookup.termStatsReaderFactory; } /** * Creates a copy of the current {@link SearchLookup} that looks fields up in the same way, but also tracks field references * in order to detect cycles and prevent resolving fields that depend on more than {@link #MAX_FIELD_CHAIN_DEPTH} other fields. + * * @param field the field being referred to, for which fielddata needs to be loaded * @return the new lookup * @throws IllegalArgumentException if a cycle is detected in the fields required to build doc values, or if the field - * being resolved depends on more than {@link #MAX_FIELD_CHAIN_DEPTH} + * being resolved depends on more than {@link #MAX_FIELD_CHAIN_DEPTH} */ public final SearchLookup forkAndTrackFieldReferences(String field) { Objects.requireNonNull(field, "field cannot be null"); @@ -150,8 +147,4 @@ public IndexFieldData getForField(MappedFieldType fieldType, MappedFieldType. public Source getSource(LeafReaderContext ctx, int doc) throws IOException { return sourceProvider.getSource(ctx, doc); } - - public TermStatsReader getTermStatsReader(LeafReaderContext leafReaderContext, Supplier docIdSupplier) { - return termStatsReaderFactory != null ? termStatsReaderFactory.apply(leafReaderContext, docIdSupplier) : null; - } }