diff --git a/docs/changelog/117839.yaml b/docs/changelog/117839.yaml new file mode 100644 index 000000000000..98c97b5078c0 --- /dev/null +++ b/docs/changelog/117839.yaml @@ -0,0 +1,5 @@ +pr: 117839 +summary: Add match support for `semantic_text` fields +area: "Search" +type: enhancement +issues: [] diff --git a/server/src/main/java/module-info.java b/server/src/main/java/module-info.java index ff902dbede00..51896fb80a62 100644 --- a/server/src/main/java/module-info.java +++ b/server/src/main/java/module-info.java @@ -479,5 +479,5 @@ module org.elasticsearch.server { exports org.elasticsearch.lucene.spatial; exports org.elasticsearch.inference.configuration; exports org.elasticsearch.monitor.metrics; - + exports org.elasticsearch.plugins.internal.rewriter to org.elasticsearch.inference; } diff --git a/server/src/main/java/org/elasticsearch/TransportVersions.java b/server/src/main/java/org/elasticsearch/TransportVersions.java index ac083862357d..4135b1f0b8e9 100644 --- a/server/src/main/java/org/elasticsearch/TransportVersions.java +++ b/server/src/main/java/org/elasticsearch/TransportVersions.java @@ -137,6 +137,7 @@ public class TransportVersions { public static final TransportVersion RETRIES_AND_OPERATIONS_IN_BLOBSTORE_STATS = def(8_804_00_0); public static final TransportVersion ADD_DATA_STREAM_OPTIONS_TO_TEMPLATES = def(8_805_00_0); public static final TransportVersion KNN_QUERY_RESCORE_OVERSAMPLE = def(8_806_00_0); + public static final TransportVersion SEMANTIC_QUERY_LENIENT = def(8_807_00_0); /* * STOP! READ THIS FIRST! No, really, diff --git a/server/src/main/java/org/elasticsearch/index/IndexModule.java b/server/src/main/java/org/elasticsearch/index/IndexModule.java index 64182b000827..2168ad1df5d2 100644 --- a/server/src/main/java/org/elasticsearch/index/IndexModule.java +++ b/server/src/main/java/org/elasticsearch/index/IndexModule.java @@ -58,6 +58,7 @@ import org.elasticsearch.indices.breaker.CircuitBreakerService; import org.elasticsearch.indices.fielddata.cache.IndicesFieldDataCache; import org.elasticsearch.indices.recovery.RecoveryState; import org.elasticsearch.plugins.IndexStorePlugin; +import org.elasticsearch.plugins.internal.rewriter.QueryRewriteInterceptor; import org.elasticsearch.script.ScriptService; import org.elasticsearch.search.aggregations.support.ValuesSourceRegistry; import org.elasticsearch.threadpool.ThreadPool; @@ -478,7 +479,8 @@ public final class IndexModule { IdFieldMapper idFieldMapper, ValuesSourceRegistry valuesSourceRegistry, IndexStorePlugin.IndexFoldersDeletionListener indexFoldersDeletionListener, - Map snapshotCommitSuppliers + Map snapshotCommitSuppliers, + QueryRewriteInterceptor queryRewriteInterceptor ) throws IOException { final IndexEventListener eventListener = freeze(); Function> readerWrapperFactory = indexReaderWrapper @@ -540,7 +542,8 @@ public final class IndexModule { indexFoldersDeletionListener, snapshotCommitSupplier, indexCommitListener.get(), - mapperMetrics + mapperMetrics, + queryRewriteInterceptor ); success = true; return indexService; diff --git a/server/src/main/java/org/elasticsearch/index/IndexService.java b/server/src/main/java/org/elasticsearch/index/IndexService.java index 571bbd76a49d..a5b3991d89bc 100644 --- a/server/src/main/java/org/elasticsearch/index/IndexService.java +++ b/server/src/main/java/org/elasticsearch/index/IndexService.java @@ -85,6 +85,7 @@ import org.elasticsearch.indices.cluster.IndicesClusterStateService; import org.elasticsearch.indices.fielddata.cache.IndicesFieldDataCache; import org.elasticsearch.indices.recovery.RecoveryState; import org.elasticsearch.plugins.IndexStorePlugin; +import org.elasticsearch.plugins.internal.rewriter.QueryRewriteInterceptor; import org.elasticsearch.script.ScriptService; import org.elasticsearch.search.aggregations.support.ValuesSourceRegistry; import org.elasticsearch.threadpool.ThreadPool; @@ -162,6 +163,7 @@ public class IndexService extends AbstractIndexComponent implements IndicesClust private final Supplier indexSortSupplier; private final ValuesSourceRegistry valuesSourceRegistry; private final MapperMetrics mapperMetrics; + private final QueryRewriteInterceptor queryRewriteInterceptor; @SuppressWarnings("this-escape") public IndexService( @@ -196,7 +198,8 @@ public class IndexService extends AbstractIndexComponent implements IndicesClust IndexStorePlugin.IndexFoldersDeletionListener indexFoldersDeletionListener, IndexStorePlugin.SnapshotCommitSupplier snapshotCommitSupplier, Engine.IndexCommitListener indexCommitListener, - MapperMetrics mapperMetrics + MapperMetrics mapperMetrics, + QueryRewriteInterceptor queryRewriteInterceptor ) { super(indexSettings); assert indexCreationContext != IndexCreationContext.RELOAD_ANALYZERS @@ -271,6 +274,7 @@ public class IndexService extends AbstractIndexComponent implements IndicesClust this.indexingOperationListeners = Collections.unmodifiableList(indexingOperationListeners); this.indexCommitListener = indexCommitListener; this.mapperMetrics = mapperMetrics; + this.queryRewriteInterceptor = queryRewriteInterceptor; try (var ignored = threadPool.getThreadContext().clearTraceContext()) { // kick off async ops for the first shard in this index this.refreshTask = new AsyncRefreshTask(this); @@ -802,6 +806,7 @@ public class IndexService extends AbstractIndexComponent implements IndicesClust allowExpensiveQueries, scriptService, null, + null, null ); } diff --git a/server/src/main/java/org/elasticsearch/index/query/AbstractQueryBuilder.java b/server/src/main/java/org/elasticsearch/index/query/AbstractQueryBuilder.java index f00e6904feac..05262798bac2 100644 --- a/server/src/main/java/org/elasticsearch/index/query/AbstractQueryBuilder.java +++ b/server/src/main/java/org/elasticsearch/index/query/AbstractQueryBuilder.java @@ -21,6 +21,7 @@ import org.elasticsearch.common.io.stream.StreamOutput; import org.elasticsearch.common.lucene.BytesRefs; import org.elasticsearch.common.settings.Settings; import org.elasticsearch.common.xcontent.SuggestingErrorOnUnknown; +import org.elasticsearch.plugins.internal.rewriter.QueryRewriteInterceptor; import org.elasticsearch.xcontent.AbstractObjectParser; import org.elasticsearch.xcontent.FilterXContentParser; import org.elasticsearch.xcontent.FilterXContentParserWrapper; @@ -278,6 +279,14 @@ public abstract class AbstractQueryBuilder> @Override public final QueryBuilder rewrite(QueryRewriteContext queryRewriteContext) throws IOException { + QueryRewriteInterceptor queryRewriteInterceptor = queryRewriteContext.getQueryRewriteInterceptor(); + if (queryRewriteInterceptor != null) { + var rewritten = queryRewriteInterceptor.interceptAndRewrite(queryRewriteContext, this); + if (rewritten != this) { + return new InterceptedQueryBuilderWrapper(rewritten); + } + } + QueryBuilder rewritten = doRewrite(queryRewriteContext); if (rewritten == this) { return rewritten; diff --git a/server/src/main/java/org/elasticsearch/index/query/CoordinatorRewriteContext.java b/server/src/main/java/org/elasticsearch/index/query/CoordinatorRewriteContext.java index e054f17ef64d..a84455ef09bf 100644 --- a/server/src/main/java/org/elasticsearch/index/query/CoordinatorRewriteContext.java +++ b/server/src/main/java/org/elasticsearch/index/query/CoordinatorRewriteContext.java @@ -104,6 +104,7 @@ public class CoordinatorRewriteContext extends QueryRewriteContext { null, null, null, + null, null ); this.dateFieldRangeInfo = dateFieldRangeInfo; diff --git a/server/src/main/java/org/elasticsearch/index/query/InnerHitContextBuilder.java b/server/src/main/java/org/elasticsearch/index/query/InnerHitContextBuilder.java index aacb4b4129c7..31bc7dddacb7 100644 --- a/server/src/main/java/org/elasticsearch/index/query/InnerHitContextBuilder.java +++ b/server/src/main/java/org/elasticsearch/index/query/InnerHitContextBuilder.java @@ -66,6 +66,9 @@ public abstract class InnerHitContextBuilder { public static void extractInnerHits(QueryBuilder query, Map innerHitBuilders) { if (query instanceof AbstractQueryBuilder) { ((AbstractQueryBuilder) query).extractInnerHitBuilders(innerHitBuilders); + } else if (query instanceof InterceptedQueryBuilderWrapper interceptedQuery) { + // Unwrap an intercepted query here + extractInnerHits(interceptedQuery.queryBuilder, innerHitBuilders); } else { throw new IllegalStateException( "provided query builder [" + query.getClass() + "] class should inherit from AbstractQueryBuilder, but it doesn't" diff --git a/server/src/main/java/org/elasticsearch/index/query/InterceptedQueryBuilderWrapper.java b/server/src/main/java/org/elasticsearch/index/query/InterceptedQueryBuilderWrapper.java new file mode 100644 index 000000000000..b1030e4a76d9 --- /dev/null +++ b/server/src/main/java/org/elasticsearch/index/query/InterceptedQueryBuilderWrapper.java @@ -0,0 +1,109 @@ +/* + * 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.query; + +import org.apache.lucene.search.Query; +import org.elasticsearch.TransportVersion; +import org.elasticsearch.common.io.stream.StreamOutput; +import org.elasticsearch.plugins.internal.rewriter.QueryRewriteInterceptor; +import org.elasticsearch.xcontent.XContentBuilder; + +import java.io.IOException; +import java.util.Objects; + +/** + * Wrapper for instances of {@link QueryBuilder} that have been intercepted using the {@link QueryRewriteInterceptor} to + * break out of the rewrite phase. These instances are unwrapped on serialization. + */ +class InterceptedQueryBuilderWrapper implements QueryBuilder { + + protected final QueryBuilder queryBuilder; + + InterceptedQueryBuilderWrapper(QueryBuilder queryBuilder) { + super(); + this.queryBuilder = queryBuilder; + } + + @Override + public QueryBuilder rewrite(QueryRewriteContext queryRewriteContext) throws IOException { + QueryRewriteInterceptor queryRewriteInterceptor = queryRewriteContext.getQueryRewriteInterceptor(); + try { + queryRewriteContext.setQueryRewriteInterceptor(null); + QueryBuilder rewritten = queryBuilder.rewrite(queryRewriteContext); + return rewritten != queryBuilder ? new InterceptedQueryBuilderWrapper(rewritten) : this; + } finally { + queryRewriteContext.setQueryRewriteInterceptor(queryRewriteInterceptor); + } + } + + @Override + public String getWriteableName() { + return queryBuilder.getWriteableName(); + } + + @Override + public TransportVersion getMinimalSupportedVersion() { + return queryBuilder.getMinimalSupportedVersion(); + } + + @Override + public Query toQuery(SearchExecutionContext context) throws IOException { + return queryBuilder.toQuery(context); + } + + @Override + public QueryBuilder queryName(String queryName) { + queryBuilder.queryName(queryName); + return this; + } + + @Override + public String queryName() { + return queryBuilder.queryName(); + } + + @Override + public float boost() { + return queryBuilder.boost(); + } + + @Override + public QueryBuilder boost(float boost) { + queryBuilder.boost(boost); + return this; + } + + @Override + public String getName() { + return queryBuilder.getName(); + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + queryBuilder.writeTo(out); + } + + @Override + public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { + return queryBuilder.toXContent(builder, params); + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o instanceof InterceptedQueryBuilderWrapper == false) return false; + return Objects.equals(queryBuilder, ((InterceptedQueryBuilderWrapper) o).queryBuilder); + } + + @Override + public int hashCode() { + return Objects.hashCode(queryBuilder); + } +} 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 fce74aa60ab1..265a0c52593b 100644 --- a/server/src/main/java/org/elasticsearch/index/query/QueryRewriteContext.java +++ b/server/src/main/java/org/elasticsearch/index/query/QueryRewriteContext.java @@ -28,6 +28,7 @@ import org.elasticsearch.index.mapper.MapperService; import org.elasticsearch.index.mapper.MappingLookup; import org.elasticsearch.index.mapper.SourceFieldMapper; import org.elasticsearch.index.mapper.TextFieldMapper; +import org.elasticsearch.plugins.internal.rewriter.QueryRewriteInterceptor; import org.elasticsearch.script.ScriptCompiler; import org.elasticsearch.search.aggregations.support.ValuesSourceRegistry; import org.elasticsearch.search.builder.PointInTimeBuilder; @@ -70,6 +71,7 @@ public class QueryRewriteContext { protected Predicate allowedFields; private final ResolvedIndices resolvedIndices; private final PointInTimeBuilder pit; + private QueryRewriteInterceptor queryRewriteInterceptor; public QueryRewriteContext( final XContentParserConfiguration parserConfiguration, @@ -86,7 +88,8 @@ public class QueryRewriteContext { final BooleanSupplier allowExpensiveQueries, final ScriptCompiler scriptService, final ResolvedIndices resolvedIndices, - final PointInTimeBuilder pit + final PointInTimeBuilder pit, + final QueryRewriteInterceptor queryRewriteInterceptor ) { this.parserConfiguration = parserConfiguration; @@ -105,6 +108,7 @@ public class QueryRewriteContext { this.scriptService = scriptService; this.resolvedIndices = resolvedIndices; this.pit = pit; + this.queryRewriteInterceptor = queryRewriteInterceptor; } public QueryRewriteContext(final XContentParserConfiguration parserConfiguration, final Client client, final LongSupplier nowInMillis) { @@ -123,6 +127,7 @@ public class QueryRewriteContext { null, null, null, + null, null ); } @@ -132,7 +137,8 @@ public class QueryRewriteContext { final Client client, final LongSupplier nowInMillis, final ResolvedIndices resolvedIndices, - final PointInTimeBuilder pit + final PointInTimeBuilder pit, + final QueryRewriteInterceptor queryRewriteInterceptor ) { this( parserConfiguration, @@ -149,7 +155,8 @@ public class QueryRewriteContext { null, null, resolvedIndices, - pit + pit, + queryRewriteInterceptor ); } @@ -428,4 +435,13 @@ public class QueryRewriteContext { // It was decided we should only test the first of these potentially multiple preferences. return value.split(",")[0].trim(); } + + public QueryRewriteInterceptor getQueryRewriteInterceptor() { + return queryRewriteInterceptor; + } + + public void setQueryRewriteInterceptor(QueryRewriteInterceptor queryRewriteInterceptor) { + this.queryRewriteInterceptor = queryRewriteInterceptor; + } + } 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 fbc3696d4022..b2ee6134a772 100644 --- a/server/src/main/java/org/elasticsearch/index/query/SearchExecutionContext.java +++ b/server/src/main/java/org/elasticsearch/index/query/SearchExecutionContext.java @@ -271,6 +271,7 @@ public class SearchExecutionContext extends QueryRewriteContext { allowExpensiveQueries, scriptService, null, + null, null ); this.shardId = shardId; diff --git a/server/src/main/java/org/elasticsearch/indices/IndicesService.java b/server/src/main/java/org/elasticsearch/indices/IndicesService.java index 818bf2036e3b..a5765a1a707d 100644 --- a/server/src/main/java/org/elasticsearch/indices/IndicesService.java +++ b/server/src/main/java/org/elasticsearch/indices/IndicesService.java @@ -135,6 +135,7 @@ import org.elasticsearch.node.Node; import org.elasticsearch.plugins.FieldPredicate; import org.elasticsearch.plugins.IndexStorePlugin; import org.elasticsearch.plugins.PluginsService; +import org.elasticsearch.plugins.internal.rewriter.QueryRewriteInterceptor; import org.elasticsearch.repositories.RepositoriesService; import org.elasticsearch.script.ScriptService; import org.elasticsearch.search.aggregations.support.ValuesSourceRegistry; @@ -262,6 +263,7 @@ public class IndicesService extends AbstractLifecycleComponent private final MapperMetrics mapperMetrics; private final PostRecoveryMerger postRecoveryMerger; private final List searchOperationListeners; + private final QueryRewriteInterceptor queryRewriteInterceptor; @Override protected void doStart() { @@ -330,6 +332,7 @@ public class IndicesService extends AbstractLifecycleComponent this.indexFoldersDeletionListeners = new CompositeIndexFoldersDeletionListener(builder.indexFoldersDeletionListeners); this.snapshotCommitSuppliers = builder.snapshotCommitSuppliers; this.requestCacheKeyDifferentiator = builder.requestCacheKeyDifferentiator; + this.queryRewriteInterceptor = builder.queryRewriteInterceptor; this.mapperMetrics = builder.mapperMetrics; // doClose() is called when shutting down a node, yet there might still be ongoing requests // that we need to wait for before closing some resources such as the caches. In order to @@ -779,7 +782,8 @@ public class IndicesService extends AbstractLifecycleComponent idFieldMappers.apply(idxSettings.getMode()), valuesSourceRegistry, indexFoldersDeletionListeners, - snapshotCommitSuppliers + snapshotCommitSuppliers, + queryRewriteInterceptor ); } @@ -1766,7 +1770,7 @@ public class IndicesService extends AbstractLifecycleComponent * Returns a new {@link QueryRewriteContext} with the given {@code now} provider */ public QueryRewriteContext getRewriteContext(LongSupplier nowInMillis, ResolvedIndices resolvedIndices, PointInTimeBuilder pit) { - return new QueryRewriteContext(parserConfig, client, nowInMillis, resolvedIndices, pit); + return new QueryRewriteContext(parserConfig, client, nowInMillis, resolvedIndices, pit, queryRewriteInterceptor); } public DataRewriteContext getDataRewriteContext(LongSupplier nowInMillis) { diff --git a/server/src/main/java/org/elasticsearch/indices/IndicesServiceBuilder.java b/server/src/main/java/org/elasticsearch/indices/IndicesServiceBuilder.java index 08d1b5ce3a96..f0f0f453e3be 100644 --- a/server/src/main/java/org/elasticsearch/indices/IndicesServiceBuilder.java +++ b/server/src/main/java/org/elasticsearch/indices/IndicesServiceBuilder.java @@ -32,6 +32,8 @@ import org.elasticsearch.indices.breaker.CircuitBreakerService; import org.elasticsearch.plugins.EnginePlugin; import org.elasticsearch.plugins.IndexStorePlugin; import org.elasticsearch.plugins.PluginsService; +import org.elasticsearch.plugins.SearchPlugin; +import org.elasticsearch.plugins.internal.rewriter.QueryRewriteInterceptor; import org.elasticsearch.script.ScriptService; import org.elasticsearch.search.aggregations.support.ValuesSourceRegistry; import org.elasticsearch.search.internal.ShardSearchRequest; @@ -76,6 +78,7 @@ public class IndicesServiceBuilder { CheckedBiConsumer requestCacheKeyDifferentiator; MapperMetrics mapperMetrics; List searchOperationListener = List.of(); + QueryRewriteInterceptor queryRewriteInterceptor = null; public IndicesServiceBuilder settings(Settings settings) { this.settings = settings; @@ -239,6 +242,27 @@ public class IndicesServiceBuilder { .flatMap(m -> m.entrySet().stream()) .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); + var queryRewriteInterceptors = pluginsService.filterPlugins(SearchPlugin.class) + .map(SearchPlugin::getQueryRewriteInterceptors) + .flatMap(List::stream) + .collect(Collectors.toMap(QueryRewriteInterceptor::getQueryName, interceptor -> { + if (interceptor.getQueryName() == null) { + throw new IllegalArgumentException("QueryRewriteInterceptor [" + interceptor.getClass().getName() + "] requires name"); + } + return interceptor; + }, (a, b) -> { + throw new IllegalStateException( + "Conflicting rewrite interceptors [" + + a.getQueryName() + + "] found in [" + + a.getClass().getName() + + "] and [" + + b.getClass().getName() + + "]" + ); + })); + queryRewriteInterceptor = QueryRewriteInterceptor.multi(queryRewriteInterceptors); + return new IndicesService(this); } } diff --git a/server/src/main/java/org/elasticsearch/plugins/SearchPlugin.java b/server/src/main/java/org/elasticsearch/plugins/SearchPlugin.java index f5670ebd8a54..e87e9ee85b29 100644 --- a/server/src/main/java/org/elasticsearch/plugins/SearchPlugin.java +++ b/server/src/main/java/org/elasticsearch/plugins/SearchPlugin.java @@ -23,6 +23,7 @@ import org.elasticsearch.index.query.QueryBuilder; import org.elasticsearch.index.query.QueryParser; import org.elasticsearch.index.query.functionscore.ScoreFunctionBuilder; import org.elasticsearch.index.query.functionscore.ScoreFunctionParser; +import org.elasticsearch.plugins.internal.rewriter.QueryRewriteInterceptor; import org.elasticsearch.search.SearchExtBuilder; import org.elasticsearch.search.aggregations.Aggregation; import org.elasticsearch.search.aggregations.AggregationBuilder; @@ -128,6 +129,14 @@ public interface SearchPlugin { return emptyList(); } + /** + * @return Applicable {@link QueryRewriteInterceptor}s configured for this plugin. + * Note: This is internal to Elasticsearch's API and not extensible by external plugins. + */ + default List getQueryRewriteInterceptors() { + return emptyList(); + } + /** * The new {@link Aggregation}s added by this plugin. */ diff --git a/server/src/main/java/org/elasticsearch/plugins/internal/rewriter/QueryRewriteInterceptor.java b/server/src/main/java/org/elasticsearch/plugins/internal/rewriter/QueryRewriteInterceptor.java new file mode 100644 index 000000000000..8f4fb2ce7491 --- /dev/null +++ b/server/src/main/java/org/elasticsearch/plugins/internal/rewriter/QueryRewriteInterceptor.java @@ -0,0 +1,75 @@ +/* + * 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.plugins.internal.rewriter; + +import org.elasticsearch.index.query.QueryBuilder; +import org.elasticsearch.index.query.QueryRewriteContext; + +import java.util.Map; + +/** + * Enables modules and plugins to intercept and rewrite queries during the query rewrite phase on the coordinator node. + */ +public interface QueryRewriteInterceptor { + + /** + * Intercepts and returns a rewritten query if modifications are required; otherwise, + * returns the same provided {@link QueryBuilder} instance unchanged. + * + * @param context the {@link QueryRewriteContext} providing the context for the rewrite operation + * @param queryBuilder the original {@link QueryBuilder} to potentially rewrite + * @return the rewritten {@link QueryBuilder}, or the original instance if no rewrite was needed + */ + QueryBuilder interceptAndRewrite(QueryRewriteContext context, QueryBuilder queryBuilder); + + /** + * Name of the query to be intercepted and rewritten. + */ + String getQueryName(); + + static QueryRewriteInterceptor multi(Map interceptors) { + return interceptors.isEmpty() ? new NoOpQueryRewriteInterceptor() : new CompositeQueryRewriteInterceptor(interceptors); + } + + class CompositeQueryRewriteInterceptor implements QueryRewriteInterceptor { + final String NAME = "composite"; + private final Map interceptors; + + private CompositeQueryRewriteInterceptor(Map interceptors) { + this.interceptors = interceptors; + } + + @Override + public String getQueryName() { + return NAME; + } + + @Override + public QueryBuilder interceptAndRewrite(QueryRewriteContext context, QueryBuilder queryBuilder) { + QueryRewriteInterceptor interceptor = interceptors.get(queryBuilder.getName()); + if (interceptor != null) { + return interceptor.interceptAndRewrite(context, queryBuilder); + } + return queryBuilder; + } + } + + class NoOpQueryRewriteInterceptor implements QueryRewriteInterceptor { + @Override + public QueryBuilder interceptAndRewrite(QueryRewriteContext context, QueryBuilder queryBuilder) { + return queryBuilder; + } + + @Override + public String getQueryName() { + return null; + } + } +} diff --git a/server/src/test/java/org/elasticsearch/action/search/TransportSearchActionTests.java b/server/src/test/java/org/elasticsearch/action/search/TransportSearchActionTests.java index a94427cd0df6..0707f1356516 100644 --- a/server/src/test/java/org/elasticsearch/action/search/TransportSearchActionTests.java +++ b/server/src/test/java/org/elasticsearch/action/search/TransportSearchActionTests.java @@ -1744,7 +1744,9 @@ public class TransportSearchActionTests extends ESTestCase { NodeClient client = new NodeClient(settings, threadPool); SearchService searchService = mock(SearchService.class); - when(searchService.getRewriteContext(any(), any(), any())).thenReturn(new QueryRewriteContext(null, null, null, null, null)); + when(searchService.getRewriteContext(any(), any(), any())).thenReturn( + new QueryRewriteContext(null, null, null, null, null, null) + ); ClusterService clusterService = new ClusterService( settings, new ClusterSettings(settings, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS), diff --git a/server/src/test/java/org/elasticsearch/cluster/metadata/IndexMetadataTests.java b/server/src/test/java/org/elasticsearch/cluster/metadata/IndexMetadataTests.java index 8036a964071d..4abd0c4a9d46 100644 --- a/server/src/test/java/org/elasticsearch/cluster/metadata/IndexMetadataTests.java +++ b/server/src/test/java/org/elasticsearch/cluster/metadata/IndexMetadataTests.java @@ -690,7 +690,12 @@ public class IndexMetadataTests extends ESTestCase { } private static InferenceFieldMetadata randomInferenceFieldMetadata(String name) { - return new InferenceFieldMetadata(name, randomIdentifier(), randomSet(1, 5, ESTestCase::randomIdentifier).toArray(String[]::new)); + return new InferenceFieldMetadata( + name, + randomIdentifier(), + randomIdentifier(), + randomSet(1, 5, ESTestCase::randomIdentifier).toArray(String[]::new) + ); } private IndexMetadataStats randomIndexStats(int numberOfShards) { diff --git a/server/src/test/java/org/elasticsearch/index/IndexModuleTests.java b/server/src/test/java/org/elasticsearch/index/IndexModuleTests.java index 49a4d519c0ea..c519d4834148 100644 --- a/server/src/test/java/org/elasticsearch/index/IndexModuleTests.java +++ b/server/src/test/java/org/elasticsearch/index/IndexModuleTests.java @@ -86,6 +86,7 @@ import org.elasticsearch.indices.cluster.IndicesClusterStateService.AllocatedInd import org.elasticsearch.indices.fielddata.cache.IndicesFieldDataCache; import org.elasticsearch.indices.recovery.RecoveryState; import org.elasticsearch.plugins.IndexStorePlugin; +import org.elasticsearch.plugins.internal.rewriter.MockQueryRewriteInterceptor; import org.elasticsearch.script.ScriptService; import org.elasticsearch.search.internal.ReaderContext; import org.elasticsearch.test.ClusterServiceUtils; @@ -223,7 +224,8 @@ public class IndexModuleTests extends ESTestCase { module.indexSettings().getMode().idFieldMapperWithoutFieldData(), null, indexDeletionListener, - emptyMap() + emptyMap(), + new MockQueryRewriteInterceptor() ); } diff --git a/server/src/test/java/org/elasticsearch/index/mapper/MappingLookupInferenceFieldMapperTests.java b/server/src/test/java/org/elasticsearch/index/mapper/MappingLookupInferenceFieldMapperTests.java index 809fb161fcbe..b1470c1ee5b3 100644 --- a/server/src/test/java/org/elasticsearch/index/mapper/MappingLookupInferenceFieldMapperTests.java +++ b/server/src/test/java/org/elasticsearch/index/mapper/MappingLookupInferenceFieldMapperTests.java @@ -15,7 +15,6 @@ import org.elasticsearch.index.query.SearchExecutionContext; import org.elasticsearch.plugins.MapperPlugin; import org.elasticsearch.plugins.Plugin; -import java.io.IOException; import java.util.Collection; import java.util.List; import java.util.Map; @@ -94,6 +93,7 @@ public class MappingLookupInferenceFieldMapperTests extends MapperServiceTestCas public static final TypeParser PARSER = new TypeParser((n, c) -> new Builder(n)); public static final String INFERENCE_ID = "test_inference_id"; + public static final String SEARCH_INFERENCE_ID = "test_search_inference_id"; public static final String CONTENT_TYPE = "test_inference_field"; TestInferenceFieldMapper(String simpleName) { @@ -102,7 +102,7 @@ public class MappingLookupInferenceFieldMapperTests extends MapperServiceTestCas @Override public InferenceFieldMetadata getMetadata(Set sourcePaths) { - return new InferenceFieldMetadata(fullPath(), INFERENCE_ID, sourcePaths.toArray(new String[0])); + return new InferenceFieldMetadata(fullPath(), INFERENCE_ID, SEARCH_INFERENCE_ID, sourcePaths.toArray(new String[0])); } @Override @@ -111,7 +111,7 @@ public class MappingLookupInferenceFieldMapperTests extends MapperServiceTestCas } @Override - protected void parseCreateField(DocumentParserContext context) throws IOException {} + protected void parseCreateField(DocumentParserContext context) {} @Override public Builder getMergeBuilder() { diff --git a/server/src/test/java/org/elasticsearch/index/query/InterceptedQueryBuilderWrapperTests.java b/server/src/test/java/org/elasticsearch/index/query/InterceptedQueryBuilderWrapperTests.java new file mode 100644 index 000000000000..6c570e0e7172 --- /dev/null +++ b/server/src/test/java/org/elasticsearch/index/query/InterceptedQueryBuilderWrapperTests.java @@ -0,0 +1,92 @@ +/* + * 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.query; + +import org.elasticsearch.plugins.internal.rewriter.QueryRewriteInterceptor; +import org.elasticsearch.test.ESTestCase; +import org.elasticsearch.test.client.NoOpClient; +import org.elasticsearch.threadpool.TestThreadPool; +import org.junit.After; +import org.junit.Before; + +import java.io.IOException; + +public class InterceptedQueryBuilderWrapperTests extends ESTestCase { + + private TestThreadPool threadPool; + private NoOpClient client; + + @Before + public void setup() { + threadPool = createThreadPool(); + client = new NoOpClient(threadPool); + } + + @After + public void cleanup() { + threadPool.close(); + } + + public void testQueryNameReturnsWrappedQueryBuilder() { + MatchAllQueryBuilder matchAllQueryBuilder = new MatchAllQueryBuilder(); + InterceptedQueryBuilderWrapper interceptedQueryBuilderWrapper = new InterceptedQueryBuilderWrapper(matchAllQueryBuilder); + String queryName = randomAlphaOfLengthBetween(5, 10); + QueryBuilder namedQuery = interceptedQueryBuilderWrapper.queryName(queryName); + assertTrue(namedQuery instanceof InterceptedQueryBuilderWrapper); + assertEquals(queryName, namedQuery.queryName()); + } + + public void testQueryBoostReturnsWrappedQueryBuilder() { + MatchAllQueryBuilder matchAllQueryBuilder = new MatchAllQueryBuilder(); + InterceptedQueryBuilderWrapper interceptedQueryBuilderWrapper = new InterceptedQueryBuilderWrapper(matchAllQueryBuilder); + float boost = randomFloat(); + QueryBuilder boostedQuery = interceptedQueryBuilderWrapper.boost(boost); + assertTrue(boostedQuery instanceof InterceptedQueryBuilderWrapper); + assertEquals(boost, boostedQuery.boost(), 0.0001f); + } + + public void testRewrite() throws IOException { + QueryRewriteContext context = new QueryRewriteContext(null, client, null); + context.setQueryRewriteInterceptor(myMatchInterceptor); + + // Queries that are not intercepted behave normally + TermQueryBuilder termQueryBuilder = new TermQueryBuilder("field", "value"); + QueryBuilder rewritten = termQueryBuilder.rewrite(context); + assertTrue(rewritten instanceof TermQueryBuilder); + + // Queries that should be intercepted are and the right thing happens + MatchQueryBuilder matchQueryBuilder = new MatchQueryBuilder("field", "value"); + rewritten = matchQueryBuilder.rewrite(context); + assertTrue(rewritten instanceof InterceptedQueryBuilderWrapper); + assertTrue(((InterceptedQueryBuilderWrapper) rewritten).queryBuilder instanceof MatchQueryBuilder); + MatchQueryBuilder rewrittenMatchQueryBuilder = (MatchQueryBuilder) ((InterceptedQueryBuilderWrapper) rewritten).queryBuilder; + assertEquals("intercepted", rewrittenMatchQueryBuilder.value()); + + // An additional rewrite on an already intercepted query returns the same query + QueryBuilder rewrittenAgain = rewritten.rewrite(context); + assertTrue(rewrittenAgain instanceof InterceptedQueryBuilderWrapper); + assertEquals(rewritten, rewrittenAgain); + } + + private final QueryRewriteInterceptor myMatchInterceptor = new QueryRewriteInterceptor() { + @Override + public QueryBuilder interceptAndRewrite(QueryRewriteContext context, QueryBuilder queryBuilder) { + if (queryBuilder instanceof MatchQueryBuilder matchQueryBuilder) { + return new MatchQueryBuilder(matchQueryBuilder.fieldName(), "intercepted"); + } + return queryBuilder; + } + + @Override + public String getQueryName() { + return MatchQueryBuilder.NAME; + } + }; +} diff --git a/server/src/test/java/org/elasticsearch/index/query/QueryRewriteContextTests.java b/server/src/test/java/org/elasticsearch/index/query/QueryRewriteContextTests.java index d07bcf54fdf0..5dd231ab9788 100644 --- a/server/src/test/java/org/elasticsearch/index/query/QueryRewriteContextTests.java +++ b/server/src/test/java/org/elasticsearch/index/query/QueryRewriteContextTests.java @@ -52,6 +52,7 @@ public class QueryRewriteContextTests extends ESTestCase { null, null, null, + null, null ); @@ -79,6 +80,7 @@ public class QueryRewriteContextTests extends ESTestCase { null, null, null, + null, null ); diff --git a/test/framework/src/main/java/org/elasticsearch/plugins/internal/rewriter/MockQueryRewriteInterceptor.java b/test/framework/src/main/java/org/elasticsearch/plugins/internal/rewriter/MockQueryRewriteInterceptor.java new file mode 100644 index 000000000000..196e5bd4f4a2 --- /dev/null +++ b/test/framework/src/main/java/org/elasticsearch/plugins/internal/rewriter/MockQueryRewriteInterceptor.java @@ -0,0 +1,26 @@ +/* + * 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.plugins.internal.rewriter; + +import org.elasticsearch.index.query.QueryBuilder; +import org.elasticsearch.index.query.QueryRewriteContext; + +public class MockQueryRewriteInterceptor implements QueryRewriteInterceptor { + + @Override + public QueryBuilder interceptAndRewrite(QueryRewriteContext context, QueryBuilder queryBuilder) { + return queryBuilder; + } + + @Override + public String getQueryName() { + return this.getClass().getSimpleName(); + } +} diff --git a/test/framework/src/main/java/org/elasticsearch/test/AbstractBuilderTestCase.java b/test/framework/src/main/java/org/elasticsearch/test/AbstractBuilderTestCase.java index bdf323afb8d9..20cb66affdde 100644 --- a/test/framework/src/main/java/org/elasticsearch/test/AbstractBuilderTestCase.java +++ b/test/framework/src/main/java/org/elasticsearch/test/AbstractBuilderTestCase.java @@ -71,6 +71,8 @@ import org.elasticsearch.plugins.Plugin; import org.elasticsearch.plugins.PluginsService; import org.elasticsearch.plugins.ScriptPlugin; import org.elasticsearch.plugins.SearchPlugin; +import org.elasticsearch.plugins.internal.rewriter.MockQueryRewriteInterceptor; +import org.elasticsearch.plugins.internal.rewriter.QueryRewriteInterceptor; import org.elasticsearch.plugins.scanners.StablePluginsRegistry; import org.elasticsearch.script.MockScriptEngine; import org.elasticsearch.script.MockScriptService; @@ -629,7 +631,8 @@ public abstract class AbstractBuilderTestCase extends ESTestCase { () -> true, scriptService, createMockResolvedIndices(), - null + null, + createMockQueryRewriteInterceptor() ); } @@ -670,5 +673,9 @@ public abstract class AbstractBuilderTestCase extends ESTestCase { Map.of(index, indexMetadata) ); } + + private QueryRewriteInterceptor createMockQueryRewriteInterceptor() { + return new MockQueryRewriteInterceptor(); + } } } diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferenceFeatures.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferenceFeatures.java index 67892dfe7862..3b7613b8b0e1 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferenceFeatures.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferenceFeatures.java @@ -10,6 +10,7 @@ package org.elasticsearch.xpack.inference; import org.elasticsearch.features.FeatureSpecification; import org.elasticsearch.features.NodeFeature; import org.elasticsearch.xpack.inference.mapper.SemanticTextFieldMapper; +import org.elasticsearch.xpack.inference.queries.SemanticMatchQueryRewriteInterceptor; import org.elasticsearch.xpack.inference.queries.SemanticQueryBuilder; import org.elasticsearch.xpack.inference.rank.random.RandomRankRetrieverBuilder; import org.elasticsearch.xpack.inference.rank.textsimilarity.TextSimilarityRankRetrieverBuilder; @@ -43,7 +44,8 @@ public class InferenceFeatures implements FeatureSpecification { SemanticTextFieldMapper.SEMANTIC_TEXT_DELETE_FIX, SemanticTextFieldMapper.SEMANTIC_TEXT_ZERO_SIZE_FIX, SemanticTextFieldMapper.SEMANTIC_TEXT_ALWAYS_EMIT_INFERENCE_ID_FIX, - SEMANTIC_TEXT_HIGHLIGHTER + SEMANTIC_TEXT_HIGHLIGHTER, + SemanticMatchQueryRewriteInterceptor.SEMANTIC_MATCH_QUERY_REWRITE_INTERCEPTION_SUPPORTED ); } } diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferencePlugin.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferencePlugin.java index eef07aefb30c..93743a5485c2 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferencePlugin.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferencePlugin.java @@ -36,6 +36,7 @@ import org.elasticsearch.plugins.MapperPlugin; import org.elasticsearch.plugins.Plugin; import org.elasticsearch.plugins.SearchPlugin; import org.elasticsearch.plugins.SystemIndexPlugin; +import org.elasticsearch.plugins.internal.rewriter.QueryRewriteInterceptor; import org.elasticsearch.rest.RestController; import org.elasticsearch.rest.RestHandler; import org.elasticsearch.search.fetch.subphase.highlight.Highlighter; @@ -77,6 +78,7 @@ import org.elasticsearch.xpack.inference.highlight.SemanticTextHighlighter; import org.elasticsearch.xpack.inference.logging.ThrottlerManager; import org.elasticsearch.xpack.inference.mapper.OffsetSourceFieldMapper; import org.elasticsearch.xpack.inference.mapper.SemanticTextFieldMapper; +import org.elasticsearch.xpack.inference.queries.SemanticMatchQueryRewriteInterceptor; import org.elasticsearch.xpack.inference.queries.SemanticQueryBuilder; import org.elasticsearch.xpack.inference.rank.random.RandomRankBuilder; import org.elasticsearch.xpack.inference.rank.random.RandomRankRetrieverBuilder; @@ -436,6 +438,11 @@ public class InferencePlugin extends Plugin implements ActionPlugin, ExtensibleP return List.of(new QuerySpec<>(SemanticQueryBuilder.NAME, SemanticQueryBuilder::new, SemanticQueryBuilder::fromXContent)); } + @Override + public List getQueryRewriteInterceptors() { + return List.of(new SemanticMatchQueryRewriteInterceptor()); + } + @Override public List> getRetrievers() { return List.of( diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticMatchQueryRewriteInterceptor.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticMatchQueryRewriteInterceptor.java new file mode 100644 index 000000000000..a4a8123935c3 --- /dev/null +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticMatchQueryRewriteInterceptor.java @@ -0,0 +1,95 @@ +/* + * 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; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.inference.queries; + +import org.elasticsearch.action.ResolvedIndices; +import org.elasticsearch.cluster.metadata.IndexMetadata; +import org.elasticsearch.cluster.metadata.InferenceFieldMetadata; +import org.elasticsearch.features.NodeFeature; +import org.elasticsearch.index.mapper.IndexFieldMapper; +import org.elasticsearch.index.query.BoolQueryBuilder; +import org.elasticsearch.index.query.MatchQueryBuilder; +import org.elasticsearch.index.query.QueryBuilder; +import org.elasticsearch.index.query.QueryRewriteContext; +import org.elasticsearch.index.query.TermQueryBuilder; +import org.elasticsearch.index.query.TermsQueryBuilder; +import org.elasticsearch.plugins.internal.rewriter.QueryRewriteInterceptor; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.List; + +public class SemanticMatchQueryRewriteInterceptor implements QueryRewriteInterceptor { + + public static final NodeFeature SEMANTIC_MATCH_QUERY_REWRITE_INTERCEPTION_SUPPORTED = new NodeFeature( + "search.semantic_match_query_rewrite_interception_supported" + ); + + public SemanticMatchQueryRewriteInterceptor() {} + + @Override + public QueryBuilder interceptAndRewrite(QueryRewriteContext context, QueryBuilder queryBuilder) { + assert (queryBuilder instanceof MatchQueryBuilder); + MatchQueryBuilder matchQueryBuilder = (MatchQueryBuilder) queryBuilder; + QueryBuilder rewritten = queryBuilder; + ResolvedIndices resolvedIndices = context.getResolvedIndices(); + if (resolvedIndices != null) { + Collection indexMetadataCollection = resolvedIndices.getConcreteLocalIndicesMetadata().values(); + List inferenceIndices = new ArrayList<>(); + List nonInferenceIndices = new ArrayList<>(); + for (IndexMetadata indexMetadata : indexMetadataCollection) { + String indexName = indexMetadata.getIndex().getName(); + InferenceFieldMetadata inferenceFieldMetadata = indexMetadata.getInferenceFields().get(matchQueryBuilder.fieldName()); + if (inferenceFieldMetadata != null) { + inferenceIndices.add(indexName); + } else { + nonInferenceIndices.add(indexName); + } + } + + if (inferenceIndices.isEmpty()) { + return rewritten; + } else if (nonInferenceIndices.isEmpty() == false) { + BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); + for (String inferenceIndexName : inferenceIndices) { + // Add a separate clause for each semantic query, because they may be using different inference endpoints + // TODO - consolidate this to a single clause once the semantic query supports multiple inference endpoints + boolQueryBuilder.should( + createSemanticSubQuery(inferenceIndexName, matchQueryBuilder.fieldName(), (String) matchQueryBuilder.value()) + ); + } + boolQueryBuilder.should(createMatchSubQuery(nonInferenceIndices, matchQueryBuilder)); + rewritten = boolQueryBuilder; + } else { + rewritten = new SemanticQueryBuilder(matchQueryBuilder.fieldName(), (String) matchQueryBuilder.value(), false); + } + } + + return rewritten; + + } + + @Override + public String getQueryName() { + return MatchQueryBuilder.NAME; + } + + private QueryBuilder createSemanticSubQuery(String indexName, String fieldName, String value) { + BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); + boolQueryBuilder.must(new SemanticQueryBuilder(fieldName, value, true)); + boolQueryBuilder.filter(new TermQueryBuilder(IndexFieldMapper.NAME, indexName)); + return boolQueryBuilder; + } + + private QueryBuilder createMatchSubQuery(List indices, MatchQueryBuilder matchQueryBuilder) { + BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); + boolQueryBuilder.must(matchQueryBuilder); + boolQueryBuilder.filter(new TermsQueryBuilder(IndexFieldMapper.NAME, indices)); + return boolQueryBuilder; + } +} diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryBuilder.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryBuilder.java index d648db2fbfdb..2a34651efcd9 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryBuilder.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryBuilder.java @@ -46,6 +46,7 @@ import java.util.Map; import java.util.Objects; import static org.elasticsearch.xcontent.ConstructingObjectParser.constructorArg; +import static org.elasticsearch.xcontent.ConstructingObjectParser.optionalConstructorArg; import static org.elasticsearch.xpack.core.ClientHelper.ML_ORIGIN; import static org.elasticsearch.xpack.core.ClientHelper.executeAsyncWithOrigin; @@ -57,16 +58,18 @@ public class SemanticQueryBuilder extends AbstractQueryBuilder PARSER = new ConstructingObjectParser<>( NAME, false, - args -> new SemanticQueryBuilder((String) args[0], (String) args[1]) + args -> new SemanticQueryBuilder((String) args[0], (String) args[1], (Boolean) args[2]) ); static { PARSER.declareString(constructorArg(), FIELD_FIELD); PARSER.declareString(constructorArg(), QUERY_FIELD); + PARSER.declareBoolean(optionalConstructorArg(), LENIENT_FIELD); declareStandardFields(PARSER); } @@ -75,8 +78,13 @@ public class SemanticQueryBuilder extends AbstractQueryBuilder inferenceResultsSupplier; private final InferenceResults inferenceResults; private final boolean noInferenceResults; + private final Boolean lenient; public SemanticQueryBuilder(String fieldName, String query) { + this(fieldName, query, null); + } + + public SemanticQueryBuilder(String fieldName, String query, Boolean lenient) { if (fieldName == null) { throw new IllegalArgumentException("[" + NAME + "] requires a " + FIELD_FIELD.getPreferredName() + " value"); } @@ -88,6 +96,7 @@ public class SemanticQueryBuilder extends AbstractQueryBuilder + { + "service": "test_service", + "service_settings": { + "model": "my_model", + "api_key": "abc64" + }, + "task_settings": { + } + } + + - do: + inference.put: + task_type: sparse_embedding + inference_id: sparse-inference-id-2 + body: > + { + "service": "test_service", + "service_settings": { + "model": "my_model", + "api_key": "abc64" + }, + "task_settings": { + } + } + + - do: + inference.put: + task_type: text_embedding + inference_id: dense-inference-id + body: > + { + "service": "text_embedding_test_service", + "service_settings": { + "model": "my_model", + "dimensions": 10, + "api_key": "abc64", + "similarity": "COSINE" + }, + "task_settings": { + } + } + + - do: + indices.create: + index: test-sparse-index + body: + mappings: + properties: + inference_field: + type: semantic_text + inference_id: sparse-inference-id + non_inference_field: + type: text + + - do: + indices.create: + index: test-dense-index + body: + mappings: + properties: + inference_field: + type: semantic_text + inference_id: dense-inference-id + non_inference_field: + type: text + + - do: + indices.create: + index: test-text-only-index + body: + mappings: + properties: + inference_field: + type: text + non_inference_field: + type: text + +--- +"Query using a sparse embedding model": + - skip: + features: [ "headers", "close_to" ] + + - do: + index: + index: test-sparse-index + id: doc_1 + body: + inference_field: [ "inference test", "another inference test" ] + non_inference_field: "non inference test" + refresh: true + + - do: + headers: + # Force JSON content type so that we use a parser that interprets the floating-point score as a double + Content-Type: application/json + search: + index: test-sparse-index + body: + query: + match: + inference_field: + query: "inference test" + + - match: { hits.total.value: 1 } + - match: { hits.hits.0._id: "doc_1" } + +--- +"Query using a dense embedding model": + - skip: + features: [ "headers", "close_to" ] + + - do: + index: + index: test-dense-index + id: doc_1 + body: + inference_field: [ "inference test", "another inference test" ] + non_inference_field: "non inference test" + refresh: true + + - do: + headers: + # Force JSON content type so that we use a parser that interprets the floating-point score as a double + Content-Type: application/json + search: + index: test-dense-index + body: + query: + match: + inference_field: + query: "inference test" + + - match: { hits.total.value: 1 } + - match: { hits.hits.0._id: "doc_1" } + +--- +"Query an index alias": + - skip: + features: [ "headers", "close_to" ] + + - do: + index: + index: test-sparse-index + id: doc_1 + body: + inference_field: [ "inference test", "another inference test" ] + non_inference_field: "non inference test" + refresh: true + + - do: + indices.put_alias: + index: test-sparse-index + name: my-alias + + - do: + headers: + # Force JSON content type so that we use a parser that interprets the floating-point score as a double + Content-Type: application/json + search: + index: my-alias + body: + query: + match: + inference_field: + query: "inference test" + + - match: { hits.total.value: 1 } + - match: { hits.hits.0._id: "doc_1" } + +--- +"Query indices with both semantic_text and regular text content": + + - do: + index: + index: test-sparse-index + id: doc_1 + body: + inference_field: [ "inference test", "another inference test" ] + non_inference_field: "non inference test" + refresh: true + + - do: + index: + index: test-text-only-index + id: doc_2 + body: + inference_field: [ "inference test", "not an inference field" ] + non_inference_field: "non inference test" + refresh: true + + - do: + search: + index: + - test-sparse-index + - test-text-only-index + body: + query: + match: + inference_field: + query: "inference test" + + - match: { hits.total.value: 2 } + - match: { hits.hits.0._id: "doc_1" } + - match: { hits.hits.1._id: "doc_2" } + + # Test querying multiple indices that either use the same inference ID or combine semantic_text with lexical search + - do: + indices.create: + index: test-sparse-index-2 + body: + mappings: + properties: + inference_field: + type: semantic_text + inference_id: sparse-inference-id + non_inference_field: + type: text + + - do: + index: + index: test-sparse-index-2 + id: doc_3 + body: + inference_field: "another inference test" + refresh: true + + - do: + search: + index: + - test-sparse-index* + - test-text-only-index + body: + query: + match: + inference_field: + query: "inference test" + + - match: { hits.total.value: 3 } + - match: { hits.hits.0._id: "doc_1" } + - match: { hits.hits.1._id: "doc_3" } + - match: { hits.hits.2._id: "doc_2" } + +--- +"Query a field that has no indexed inference results": + - skip: + features: [ "headers" ] + + - do: + headers: + # Force JSON content type so that we use a parser that interprets the floating-point score as a double + Content-Type: application/json + search: + index: test-sparse-index + body: + query: + match: + inference_field: + query: "inference test" + + - match: { hits.total.value: 0 } + + - do: + headers: + # Force JSON content type so that we use a parser that interprets the floating-point score as a double + Content-Type: application/json + search: + index: test-dense-index + body: + query: + match: + inference_field: + query: "inference test" + + - match: { hits.total.value: 0 } diff --git a/x-pack/plugin/rank-rrf/src/test/java/org/elasticsearch/xpack/rank/rrf/RRFRetrieverBuilderTests.java b/x-pack/plugin/rank-rrf/src/test/java/org/elasticsearch/xpack/rank/rrf/RRFRetrieverBuilderTests.java index d20f0f88aeb1..bdd6d73ec0fb 100644 --- a/x-pack/plugin/rank-rrf/src/test/java/org/elasticsearch/xpack/rank/rrf/RRFRetrieverBuilderTests.java +++ b/x-pack/plugin/rank-rrf/src/test/java/org/elasticsearch/xpack/rank/rrf/RRFRetrieverBuilderTests.java @@ -54,7 +54,9 @@ public class RRFRetrieverBuilderTests extends ESTestCase { IllegalArgumentException iae = expectThrows( IllegalArgumentException.class, () -> ssb.parseXContent(parser, true, nf -> true) - .rewrite(new QueryRewriteContext(parserConfig(), null, null, null, new PointInTimeBuilder(new BytesArray("pitid")))) + .rewrite( + new QueryRewriteContext(parserConfig(), null, null, null, new PointInTimeBuilder(new BytesArray("pitid")), null) + ) ); assertEquals("[search_after] cannot be used in children of compound retrievers", iae.getMessage()); } @@ -70,7 +72,9 @@ public class RRFRetrieverBuilderTests extends ESTestCase { IllegalArgumentException iae = expectThrows( IllegalArgumentException.class, () -> ssb.parseXContent(parser, true, nf -> true) - .rewrite(new QueryRewriteContext(parserConfig(), null, null, null, new PointInTimeBuilder(new BytesArray("pitid")))) + .rewrite( + new QueryRewriteContext(parserConfig(), null, null, null, new PointInTimeBuilder(new BytesArray("pitid")), null) + ) ); assertEquals("[terminate_after] cannot be used in children of compound retrievers", iae.getMessage()); }