Skip to content

Commit e35c2cc

Browse files
kderussoMikep86
authored and
Max Hniebergall
committed
Add match support for semantic_text fields (elastic#117839) (elastic#118587)
* Added query name to inference field metadata * Fix build error * Added query builder service * Add query builder service to query rewrite context * Updated match query to support querying semantic text fields * Fix build error * Fix NPE * Update the POC to rewrite to a bool query when combined inference and non-inference fields * Separate clause for each inference index (to avoid inference ID clashes) * Simplify query builder service concept to a single default inference query * Rename QueryBuilderService, remove query name from inference metadata * Fix too many rewrite rounds error by injecting booleans in constructors for match query builder and semantic text * Fix test compilation errors * Fix tests * Add yaml test for semantic match * Add NodeFeature * Fix license headers * Spotless * Updated getClass comparison in MatchQueryBuilder * Cleanup * Add Mock Inference Query Builder Service * Spotless * Cleanup * Update docs/changelog/117839.yaml * Update changelog * Replace the default inference query builder with a query rewrite interceptor * Cleanup * Some more cleanup/renames * Some more cleanup/renames * Spotless * Checkstyle * Convert List<QueryRewriteInterceptor> to Map keyed on query name, error on query name collisions * PR feedback - remove check on QueryRewriteContext class only * PR feedback * Remove intercept flag from MatchQueryBuilder and replace with wrapper * Move feature to test feature * Ensure interception happens only once * Rename InterceptedQueryBuilderWrapper to AbstractQueryBuilderWrapper * Add lenient field to SemanticQueryBuilder * Clean up yaml test * Add TODO comment * Add comment * Spotless * Rename AbstractQueryBuilderWrapper back to InterceptedQueryBuilderWrapper * Spotless * Didn't mean to commit that * Remove static class wrapping the InterceptedQueryBuilderWrapper * Make InterceptedQueryBuilderWrapper part of QueryRewriteInterceptor * Refactor the interceptor to be an internal plugin that cannot be used outside inference plugin * Fix tests * Spotless * Minor cleanup * C'mon spotless * Test spotless * Cleanup InternalQueryRewriter * Change if statement to assert * Simplify template of InterceptedQueryBuilderWrapper * Change constructor of InterceptedQueryBuilderWrapper * Refactor InterceptedQueryBuilderWrapper to extend QueryBuilder * Cleanup * Add test * Spotless * Rename rewrite to interceptAndRewrite in QueryRewriteInterceptor * DOESN'T WORK - for testing * Add comment * Getting closer - match on single typed fields works now * Deleted line by mistake * Checkstyle * Fix over-aggressive IntelliJ Refactor/Rename * And another one * Move SemanticMatchQueryRewriteInterceptor.SEMANTIC_MATCH_QUERY_REWRITE_INTERCEPTION_SUPPORTED to Test feature * PR feedback * Require query name with no default * PR feedback & update test * Add rewrite test * Update server/src/main/java/org/elasticsearch/index/query/InnerHitContextBuilder.java Co-authored-by: Mike Pellegrini <[email protected]> --------- Co-authored-by: Mike Pellegrini <[email protected]> (cherry picked from commit c9a6a2c)
1 parent cfb31da commit e35c2cc

File tree

31 files changed

+890
-32
lines changed

31 files changed

+890
-32
lines changed

‎docs/changelog/117839.yaml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
pr: 117839
2+
summary: Add match support for `semantic_text` fields
3+
area: "Search"
4+
type: enhancement
5+
issues: []

‎server/src/main/java/module-info.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -474,5 +474,5 @@
474474
exports org.elasticsearch.lucene.spatial;
475475
exports org.elasticsearch.inference.configuration;
476476
exports org.elasticsearch.monitor.metrics;
477-
477+
exports org.elasticsearch.plugins.internal.rewriter to org.elasticsearch.inference;
478478
}

‎server/src/main/java/org/elasticsearch/TransportVersions.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -144,6 +144,7 @@ static TransportVersion def(int id) {
144144
public static final TransportVersion RETRIES_AND_OPERATIONS_IN_BLOBSTORE_STATS = def(8_804_00_0);
145145
public static final TransportVersion ADD_DATA_STREAM_OPTIONS_TO_TEMPLATES = def(8_805_00_0);
146146
public static final TransportVersion KNN_QUERY_RESCORE_OVERSAMPLE = def(8_806_00_0);
147+
public static final TransportVersion SEMANTIC_QUERY_LENIENT = def(8_807_00_0);
147148

148149
/*
149150
* STOP! READ THIS FIRST! No, really,

‎server/src/main/java/org/elasticsearch/index/IndexModule.java‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,7 @@
5959
import org.elasticsearch.indices.fielddata.cache.IndicesFieldDataCache;
6060
import org.elasticsearch.indices.recovery.RecoveryState;
6161
import org.elasticsearch.plugins.IndexStorePlugin;
62+
import org.elasticsearch.plugins.internal.rewriter.QueryRewriteInterceptor;
6263
import org.elasticsearch.script.ScriptService;
6364
import org.elasticsearch.search.aggregations.support.ValuesSourceRegistry;
6465
import org.elasticsearch.threadpool.ThreadPool;
@@ -483,7 +484,8 @@ public IndexService newIndexService(
483484
IdFieldMapper idFieldMapper,
484485
ValuesSourceRegistry valuesSourceRegistry,
485486
IndexStorePlugin.IndexFoldersDeletionListener indexFoldersDeletionListener,
486-
Map<String, IndexStorePlugin.SnapshotCommitSupplier> snapshotCommitSuppliers
487+
Map<String, IndexStorePlugin.SnapshotCommitSupplier> snapshotCommitSuppliers,
488+
QueryRewriteInterceptor queryRewriteInterceptor
487489
) throws IOException {
488490
final IndexEventListener eventListener = freeze();
489491
Function<IndexService, CheckedFunction<DirectoryReader, DirectoryReader, IOException>> readerWrapperFactory = indexReaderWrapper
@@ -545,7 +547,8 @@ public IndexService newIndexService(
545547
indexFoldersDeletionListener,
546548
snapshotCommitSupplier,
547549
indexCommitListener.get(),
548-
mapperMetrics
550+
mapperMetrics,
551+
queryRewriteInterceptor
549552
);
550553
success = true;
551554
return indexService;

‎server/src/main/java/org/elasticsearch/index/IndexService.java‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,7 @@
8585
import org.elasticsearch.indices.fielddata.cache.IndicesFieldDataCache;
8686
import org.elasticsearch.indices.recovery.RecoveryState;
8787
import org.elasticsearch.plugins.IndexStorePlugin;
88+
import org.elasticsearch.plugins.internal.rewriter.QueryRewriteInterceptor;
8889
import org.elasticsearch.script.ScriptService;
8990
import org.elasticsearch.search.aggregations.support.ValuesSourceRegistry;
9091
import org.elasticsearch.threadpool.ThreadPool;
@@ -162,6 +163,7 @@ public class IndexService extends AbstractIndexComponent implements IndicesClust
162163
private final Supplier<Sort> indexSortSupplier;
163164
private final ValuesSourceRegistry valuesSourceRegistry;
164165
private final MapperMetrics mapperMetrics;
166+
private final QueryRewriteInterceptor queryRewriteInterceptor;
165167

166168
@SuppressWarnings("this-escape")
167169
public IndexService(
@@ -196,7 +198,8 @@ public IndexService(
196198
IndexStorePlugin.IndexFoldersDeletionListener indexFoldersDeletionListener,
197199
IndexStorePlugin.SnapshotCommitSupplier snapshotCommitSupplier,
198200
Engine.IndexCommitListener indexCommitListener,
199-
MapperMetrics mapperMetrics
201+
MapperMetrics mapperMetrics,
202+
QueryRewriteInterceptor queryRewriteInterceptor
200203
) {
201204
super(indexSettings);
202205
assert indexCreationContext != IndexCreationContext.RELOAD_ANALYZERS
@@ -271,6 +274,7 @@ public IndexService(
271274
this.indexingOperationListeners = Collections.unmodifiableList(indexingOperationListeners);
272275
this.indexCommitListener = indexCommitListener;
273276
this.mapperMetrics = mapperMetrics;
277+
this.queryRewriteInterceptor = queryRewriteInterceptor;
274278
try (var ignored = threadPool.getThreadContext().clearTraceContext()) {
275279
// kick off async ops for the first shard in this index
276280
this.refreshTask = new AsyncRefreshTask(this);
@@ -802,6 +806,7 @@ public QueryRewriteContext newQueryRewriteContext(
802806
allowExpensiveQueries,
803807
scriptService,
804808
null,
809+
null,
805810
null
806811
);
807812
}

‎server/src/main/java/org/elasticsearch/index/query/AbstractQueryBuilder.java‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import org.elasticsearch.common.lucene.BytesRefs;
2222
import org.elasticsearch.common.settings.Settings;
2323
import org.elasticsearch.common.xcontent.SuggestingErrorOnUnknown;
24+
import org.elasticsearch.plugins.internal.rewriter.QueryRewriteInterceptor;
2425
import org.elasticsearch.xcontent.AbstractObjectParser;
2526
import org.elasticsearch.xcontent.FilterXContentParser;
2627
import org.elasticsearch.xcontent.FilterXContentParserWrapper;
@@ -278,6 +279,14 @@ protected static List<QueryBuilder> readQueries(StreamInput in) throws IOExcepti
278279

279280
@Override
280281
public final QueryBuilder rewrite(QueryRewriteContext queryRewriteContext) throws IOException {
282+
QueryRewriteInterceptor queryRewriteInterceptor = queryRewriteContext.getQueryRewriteInterceptor();
283+
if (queryRewriteInterceptor != null) {
284+
var rewritten = queryRewriteInterceptor.interceptAndRewrite(queryRewriteContext, this);
285+
if (rewritten != this) {
286+
return new InterceptedQueryBuilderWrapper(rewritten);
287+
}
288+
}
289+
281290
QueryBuilder rewritten = doRewrite(queryRewriteContext);
282291
if (rewritten == this) {
283292
return rewritten;

‎server/src/main/java/org/elasticsearch/index/query/CoordinatorRewriteContext.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -104,6 +104,7 @@ public CoordinatorRewriteContext(
104104
null,
105105
null,
106106
null,
107+
null,
107108
null
108109
);
109110
this.dateFieldRangeInfo = dateFieldRangeInfo;

‎server/src/main/java/org/elasticsearch/index/query/InnerHitContextBuilder.java‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,9 @@ public InnerHitBuilder innerHitBuilder() {
6666
public static void extractInnerHits(QueryBuilder query, Map<String, InnerHitContextBuilder> innerHitBuilders) {
6767
if (query instanceof AbstractQueryBuilder) {
6868
((AbstractQueryBuilder<?>) query).extractInnerHitBuilders(innerHitBuilders);
69+
} else if (query instanceof InterceptedQueryBuilderWrapper interceptedQuery) {
70+
// Unwrap an intercepted query here
71+
extractInnerHits(interceptedQuery.queryBuilder, innerHitBuilders);
6972
} else {
7073
throw new IllegalStateException(
7174
"provided query builder [" + query.getClass() + "] class should inherit from AbstractQueryBuilder, but it doesn't"

‎server/src/main/java/org/elasticsearch/index/query/InterceptedQueryBuilderWrapper.java‎

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,109 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the "Elastic License
4+
* 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
5+
* Public License v 1"; you may not use this file except in compliance with, at
6+
* your election, the "Elastic License 2.0", the "GNU Affero General Public
7+
* License v3.0 only", or the "Server Side Public License, v 1".
8+
*/
9+
10+
package org.elasticsearch.index.query;
11+
12+
import org.apache.lucene.search.Query;
13+
import org.elasticsearch.TransportVersion;
14+
import org.elasticsearch.common.io.stream.StreamOutput;
15+
import org.elasticsearch.plugins.internal.rewriter.QueryRewriteInterceptor;
16+
import org.elasticsearch.xcontent.XContentBuilder;
17+
18+
import java.io.IOException;
19+
import java.util.Objects;
20+
21+
/**
22+
* Wrapper for instances of {@link QueryBuilder} that have been intercepted using the {@link QueryRewriteInterceptor} to
23+
* break out of the rewrite phase. These instances are unwrapped on serialization.
24+
*/
25+
class InterceptedQueryBuilderWrapper implements QueryBuilder {
26+
27+
protected final QueryBuilder queryBuilder;
28+
29+
InterceptedQueryBuilderWrapper(QueryBuilder queryBuilder) {
30+
super();
31+
this.queryBuilder = queryBuilder;
32+
}
33+
34+
@Override
35+
public QueryBuilder rewrite(QueryRewriteContext queryRewriteContext) throws IOException {
36+
QueryRewriteInterceptor queryRewriteInterceptor = queryRewriteContext.getQueryRewriteInterceptor();
37+
try {
38+
queryRewriteContext.setQueryRewriteInterceptor(null);
39+
QueryBuilder rewritten = queryBuilder.rewrite(queryRewriteContext);
40+
return rewritten != queryBuilder ? new InterceptedQueryBuilderWrapper(rewritten) : this;
41+
} finally {
42+
queryRewriteContext.setQueryRewriteInterceptor(queryRewriteInterceptor);
43+
}
44+
}
45+
46+
@Override
47+
public String getWriteableName() {
48+
return queryBuilder.getWriteableName();
49+
}
50+
51+
@Override
52+
public TransportVersion getMinimalSupportedVersion() {
53+
return queryBuilder.getMinimalSupportedVersion();
54+
}
55+
56+
@Override
57+
public Query toQuery(SearchExecutionContext context) throws IOException {
58+
return queryBuilder.toQuery(context);
59+
}
60+
61+
@Override
62+
public QueryBuilder queryName(String queryName) {
63+
queryBuilder.queryName(queryName);
64+
return this;
65+
}
66+
67+
@Override
68+
public String queryName() {
69+
return queryBuilder.queryName();
70+
}
71+
72+
@Override
73+
public float boost() {
74+
return queryBuilder.boost();
75+
}
76+
77+
@Override
78+
public QueryBuilder boost(float boost) {
79+
queryBuilder.boost(boost);
80+
return this;
81+
}
82+
83+
@Override
84+
public String getName() {
85+
return queryBuilder.getName();
86+
}
87+
88+
@Override
89+
public void writeTo(StreamOutput out) throws IOException {
90+
queryBuilder.writeTo(out);
91+
}
92+
93+
@Override
94+
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
95+
return queryBuilder.toXContent(builder, params);
96+
}
97+
98+
@Override
99+
public boolean equals(Object o) {
100+
if (this == o) return true;
101+
if (o instanceof InterceptedQueryBuilderWrapper == false) return false;
102+
return Objects.equals(queryBuilder, ((InterceptedQueryBuilderWrapper) o).queryBuilder);
103+
}
104+
105+
@Override
106+
public int hashCode() {
107+
return Objects.hashCode(queryBuilder);
108+
}
109+
}

‎server/src/main/java/org/elasticsearch/index/query/QueryRewriteContext.java‎

Lines changed: 19 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
import org.elasticsearch.index.mapper.MappingLookup;
2929
import org.elasticsearch.index.mapper.SourceFieldMapper;
3030
import org.elasticsearch.index.mapper.TextFieldMapper;
31+
import org.elasticsearch.plugins.internal.rewriter.QueryRewriteInterceptor;
3132
import org.elasticsearch.script.ScriptCompiler;
3233
import org.elasticsearch.search.aggregations.support.ValuesSourceRegistry;
3334
import org.elasticsearch.search.builder.PointInTimeBuilder;
@@ -70,6 +71,7 @@ public class QueryRewriteContext {
7071
protected Predicate<String> allowedFields;
7172
private final ResolvedIndices resolvedIndices;
7273
private final PointInTimeBuilder pit;
74+
private QueryRewriteInterceptor queryRewriteInterceptor;
7375

7476
public QueryRewriteContext(
7577
final XContentParserConfiguration parserConfiguration,
@@ -86,7 +88,8 @@ public QueryRewriteContext(
8688
final BooleanSupplier allowExpensiveQueries,
8789
final ScriptCompiler scriptService,
8890
final ResolvedIndices resolvedIndices,
89-
final PointInTimeBuilder pit
91+
final PointInTimeBuilder pit,
92+
final QueryRewriteInterceptor queryRewriteInterceptor
9093
) {
9194

9295
this.parserConfiguration = parserConfiguration;
@@ -105,6 +108,7 @@ public QueryRewriteContext(
105108
this.scriptService = scriptService;
106109
this.resolvedIndices = resolvedIndices;
107110
this.pit = pit;
111+
this.queryRewriteInterceptor = queryRewriteInterceptor;
108112
}
109113

110114
public QueryRewriteContext(final XContentParserConfiguration parserConfiguration, final Client client, final LongSupplier nowInMillis) {
@@ -123,6 +127,7 @@ public QueryRewriteContext(final XContentParserConfiguration parserConfiguration
123127
null,
124128
null,
125129
null,
130+
null,
126131
null
127132
);
128133
}
@@ -132,7 +137,8 @@ public QueryRewriteContext(
132137
final Client client,
133138
final LongSupplier nowInMillis,
134139
final ResolvedIndices resolvedIndices,
135-
final PointInTimeBuilder pit
140+
final PointInTimeBuilder pit,
141+
final QueryRewriteInterceptor queryRewriteInterceptor
136142
) {
137143
this(
138144
parserConfiguration,
@@ -149,7 +155,8 @@ public QueryRewriteContext(
149155
null,
150156
null,
151157
resolvedIndices,
152-
pit
158+
pit,
159+
queryRewriteInterceptor
153160
);
154161
}
155162

@@ -428,4 +435,13 @@ public String getTierPreference() {
428435
// It was decided we should only test the first of these potentially multiple preferences.
429436
return value.split(",")[0].trim();
430437
}
438+
439+
public QueryRewriteInterceptor getQueryRewriteInterceptor() {
440+
return queryRewriteInterceptor;
441+
}
442+
443+
public void setQueryRewriteInterceptor(QueryRewriteInterceptor queryRewriteInterceptor) {
444+
this.queryRewriteInterceptor = queryRewriteInterceptor;
445+
}
446+
431447
}

0 commit comments

Comments
 (0)