diff --git a/server/src/main/java/org/elasticsearch/action/search/TransportSearchAction.java b/server/src/main/java/org/elasticsearch/action/search/TransportSearchAction.java index 2255681e275b9..5a6c01aaad2fa 100644 --- a/server/src/main/java/org/elasticsearch/action/search/TransportSearchAction.java +++ b/server/src/main/java/org/elasticsearch/action/search/TransportSearchAction.java @@ -453,7 +453,12 @@ void executeRequest( } } }); - Rewriteable.rewriteAndFetch(original, searchService.getRewriteContext(timeProvider::absoluteStartMillis), rewriteListener); + + Rewriteable.rewriteAndFetch( + original, + searchService.getRewriteContext(timeProvider::absoluteStartMillis, original.source().timeout()), + rewriteListener + ); } static void adjustSearchType(SearchRequest searchRequest, boolean singleShard) { diff --git a/server/src/main/java/org/elasticsearch/index/query/QueryRewriteContext.java b/server/src/main/java/org/elasticsearch/index/query/QueryRewriteContext.java index fd8d3794cf2d8..8bbccddcb3569 100644 --- a/server/src/main/java/org/elasticsearch/index/query/QueryRewriteContext.java +++ b/server/src/main/java/org/elasticsearch/index/query/QueryRewriteContext.java @@ -13,6 +13,7 @@ import org.elasticsearch.common.io.stream.NamedWriteableRegistry; import org.elasticsearch.common.regex.Regex; import org.elasticsearch.common.util.concurrent.CountDown; +import org.elasticsearch.core.TimeValue; import org.elasticsearch.index.Index; import org.elasticsearch.index.IndexSettings; import org.elasticsearch.index.analysis.IndexAnalyzers; @@ -32,6 +33,7 @@ import java.util.List; import java.util.Map; import java.util.Objects; +import java.util.Optional; import java.util.Set; import java.util.function.BiConsumer; import java.util.function.BooleanSupplier; @@ -60,6 +62,43 @@ public class QueryRewriteContext { protected boolean allowUnmappedFields; protected boolean mapUnmappedFieldAsString; protected Predicate allowedFields; + private final TimeValue timeout; + + public QueryRewriteContext( + final XContentParserConfiguration parserConfiguration, + final Client client, + final LongSupplier nowInMillis, + final MapperService mapperService, + final MappingLookup mappingLookup, + final Map runtimeMappings, + final Predicate allowedFields, + final IndexSettings indexSettings, + final Index fullyQualifiedIndex, + final Predicate indexNameMatcher, + final NamedWriteableRegistry namedWriteableRegistry, + final ValuesSourceRegistry valuesSourceRegistry, + final BooleanSupplier allowExpensiveQueries, + final ScriptCompiler scriptService, + final TimeValue timeout + ) { + + this.parserConfiguration = parserConfiguration; + this.client = client; + this.nowInMillis = nowInMillis; + this.mapperService = mapperService; + this.mappingLookup = Objects.requireNonNull(mappingLookup); + this.allowUnmappedFields = indexSettings == null || indexSettings.isDefaultAllowUnmappedFields(); + this.runtimeMappings = runtimeMappings; + this.allowedFields = allowedFields; + this.indexSettings = indexSettings; + this.fullyQualifiedIndex = fullyQualifiedIndex; + this.indexNameMatcher = indexNameMatcher; + this.writeableRegistry = namedWriteableRegistry; + this.valuesSourceRegistry = valuesSourceRegistry; + this.allowExpensiveQueries = allowExpensiveQueries; + this.scriptService = scriptService; + this.timeout = timeout; + } public QueryRewriteContext( final XContentParserConfiguration parserConfiguration, @@ -93,6 +132,7 @@ public QueryRewriteContext( this.valuesSourceRegistry = valuesSourceRegistry; this.allowExpensiveQueries = allowExpensiveQueries; this.scriptService = scriptService; + this.timeout = null; } public QueryRewriteContext(final XContentParserConfiguration parserConfiguration, final Client client, final LongSupplier nowInMillis) { @@ -110,10 +150,31 @@ public QueryRewriteContext(final XContentParserConfiguration parserConfiguration null, null, null, + null, null ); } + public QueryRewriteContext(final XContentParserConfiguration parserConfiguration, final Client client, final LongSupplier nowInMillis, final TimeValue timeout) { + this( + parserConfiguration, + client, + nowInMillis, + null, + MappingLookup.EMPTY, + Collections.emptyMap(), + null, + null, + null, + null, + null, + null, + null, + null, + timeout + ); + } + /** * The registry used to build new {@link XContentParser}s. Contains registered named parsers needed to parse the query. * @@ -360,4 +421,8 @@ public Iterable> getAllFields() { // runtime mappings and non-runtime fields don't overlap, so we can simply concatenate the iterables here return () -> Iterators.concat(allEntrySet.iterator(), runtimeEntrySet.iterator()); } + + public Optional timeout() { + return Optional.ofNullable(timeout); + } } diff --git a/server/src/main/java/org/elasticsearch/indices/IndicesService.java b/server/src/main/java/org/elasticsearch/indices/IndicesService.java index 026a20415aa91..37526eed698c8 100644 --- a/server/src/main/java/org/elasticsearch/indices/IndicesService.java +++ b/server/src/main/java/org/elasticsearch/indices/IndicesService.java @@ -1723,8 +1723,8 @@ public AliasFilter buildAliasFilter(ClusterState state, String index, Set { @@ -177,7 +182,7 @@ protected QueryBuilder doRewrite(QueryRewriteContext queryRewriteContext) { List.of(modelText), TextExpansionConfigUpdate.EMPTY_UPDATE, false, - InferModelAction.Request.DEFAULT_TIMEOUT_FOR_API + queryRewriteContext.timeout().orElse(InferModelAction.Request.DEFAULT_TIMEOUT_FOR_API) ); inferRequest.setHighPriority(true); inferRequest.setPrefixType(TrainedModelPrefixStrings.PrefixType.SEARCH);