Skip to content

Commit e663f70

Browse files
Samiul-TheSoccerFanSamiul-TheSoccerFan
committed
Introducing multimatch intercepter
1 parent 7f649e0 commit e663f70

File tree

3 files changed

+126
-19
lines changed

3 files changed

+126
-19
lines changed

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferencePlugin.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,7 @@
9595
import org.elasticsearch.xpack.inference.mapper.SemanticTextFieldMapper;
9696
import org.elasticsearch.xpack.inference.queries.SemanticKnnVectorQueryRewriteInterceptor;
9797
import org.elasticsearch.xpack.inference.queries.SemanticMatchQueryRewriteInterceptor;
98+
import org.elasticsearch.xpack.inference.queries.SemanticMultiMatchQueryRewriteInterceptor;
9899
import org.elasticsearch.xpack.inference.queries.SemanticQueryBuilder;
99100
import org.elasticsearch.xpack.inference.queries.SemanticSparseVectorQueryRewriteInterceptor;
100101
import org.elasticsearch.xpack.inference.rank.random.RandomRankBuilder;
@@ -556,7 +557,8 @@ public List<QueryRewriteInterceptor> getQueryRewriteInterceptors() {
556557
return List.of(
557558
new SemanticKnnVectorQueryRewriteInterceptor(),
558559
new SemanticMatchQueryRewriteInterceptor(),
559-
new SemanticSparseVectorQueryRewriteInterceptor()
560+
new SemanticSparseVectorQueryRewriteInterceptor(),
561+
new SemanticMultiMatchQueryRewriteInterceptor()
560562
);
561563
}
562564

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticMultiMatchQueryRewriteInterceptor.java‎

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,106 @@
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; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
package org.elasticsearch.xpack.inference.queries;
9+
10+
import org.elasticsearch.index.query.BoolQueryBuilder;
11+
import org.elasticsearch.index.query.MultiMatchQueryBuilder;
12+
import org.elasticsearch.index.query.QueryBuilder;
13+
14+
import java.util.Map;
15+
16+
public class SemanticMultiMatchQueryRewriteInterceptor extends SemanticQueryRewriteInterceptor {
17+
@Override
18+
protected Map<String, Float> getFieldNamesWithWeights(QueryBuilder queryBuilder) {
19+
assert (queryBuilder instanceof MultiMatchQueryBuilder);
20+
MultiMatchQueryBuilder multiMatchQueryBuilder = (MultiMatchQueryBuilder) queryBuilder;
21+
return multiMatchQueryBuilder.fields();
22+
}
23+
24+
@Override
25+
protected String getQuery(QueryBuilder queryBuilder) {
26+
assert (queryBuilder instanceof MultiMatchQueryBuilder);
27+
MultiMatchQueryBuilder multiMatchQueryBuilder = (MultiMatchQueryBuilder) queryBuilder;
28+
return (String) multiMatchQueryBuilder.value();
29+
}
30+
31+
@Override
32+
protected QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation) {
33+
SemanticQueryBuilder semanticQueryBuilder = new SemanticQueryBuilder(
34+
indexInformation.fieldName(),
35+
getQuery(queryBuilder),
36+
false
37+
);
38+
// TODO:: add boost
39+
semanticQueryBuilder.queryName(queryBuilder.queryName());
40+
return semanticQueryBuilder;
41+
}
42+
43+
@Override
44+
protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation) {
45+
assert (queryBuilder instanceof MultiMatchQueryBuilder);
46+
MultiMatchQueryBuilder originalMultiMatchQueryBuilder = (MultiMatchQueryBuilder) queryBuilder;
47+
48+
// Create a copy for non-inference fields with only this specific field
49+
MultiMatchQueryBuilder multiMatchQueryBuilder = createSingleFieldMultiMatch(
50+
originalMultiMatchQueryBuilder,
51+
indexInformation.fieldName()
52+
);
53+
54+
BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder();
55+
56+
// Add semantic query for inference indices
57+
boolQueryBuilder.should(
58+
createSemanticSubQuery(
59+
indexInformation.getInferenceIndices(),
60+
indexInformation.fieldName(),
61+
getQuery(queryBuilder)
62+
)
63+
);
64+
65+
// Add regular query for non-inference indices
66+
boolQueryBuilder.should(
67+
createSubQueryForIndices(indexInformation.nonInferenceIndices(), multiMatchQueryBuilder)
68+
);
69+
70+
// TODO:: add boost
71+
boolQueryBuilder.queryName(queryBuilder.queryName());
72+
return boolQueryBuilder;
73+
}
74+
75+
@Override
76+
public String getQueryName() {
77+
return MultiMatchQueryBuilder.NAME;
78+
}
79+
80+
/**
81+
* Create a MultiMatchQueryBuilder with only a single field for non-inference indices
82+
*/
83+
private MultiMatchQueryBuilder createSingleFieldMultiMatch(MultiMatchQueryBuilder original, String fieldName) {
84+
MultiMatchQueryBuilder singleFieldQuery = new MultiMatchQueryBuilder(original.value());
85+
86+
// Copy all properties from original query
87+
singleFieldQuery.type(original.type());
88+
singleFieldQuery.operator(original.operator());
89+
singleFieldQuery.analyzer(original.analyzer());
90+
singleFieldQuery.fuzziness(original.fuzziness());
91+
singleFieldQuery.prefixLength(original.prefixLength());
92+
singleFieldQuery.maxExpansions(original.maxExpansions());
93+
singleFieldQuery.minimumShouldMatch(original.minimumShouldMatch());
94+
singleFieldQuery.fuzzyRewrite(original.fuzzyRewrite());
95+
singleFieldQuery.tieBreaker(original.tieBreaker());
96+
singleFieldQuery.lenient(original.lenient());
97+
singleFieldQuery.zeroTermsQuery(original.zeroTermsQuery());
98+
singleFieldQuery.autoGenerateSynonymsPhraseQuery(original.autoGenerateSynonymsPhraseQuery());
99+
singleFieldQuery.fuzzyTranspositions(original.fuzzyTranspositions());
100+
101+
// Add only the specific field (without boost for now)
102+
singleFieldQuery.field(fieldName);
103+
104+
return singleFieldQuery;
105+
}
106+
}

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryRewriteInterceptor.java‎

Lines changed: 17 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -34,32 +34,31 @@ public SemanticQueryRewriteInterceptor() {}
3434
@Override
3535
public QueryBuilder interceptAndRewrite(QueryRewriteContext context, QueryBuilder queryBuilder) {
3636
Map<String, Float> fieldNamesWithWeights = getFieldNamesWithWeights(queryBuilder);
37-
if (fieldNamesWithWeights.size() > 1) {
38-
// Multi-field query, so return the original query, and an exception will be thrown eventually
39-
return queryBuilder;
40-
}
41-
String fieldName = fieldNamesWithWeights.keySet().iterator().next();
4237
ResolvedIndices resolvedIndices = context.getResolvedIndices();
4338

4439
if (resolvedIndices == null) {
4540
// No resolved indices, so return the original query.
4641
return queryBuilder;
4742
}
4843

49-
InferenceIndexInformationForField indexInformation = resolveIndicesForField(fieldName, resolvedIndices);
50-
if (indexInformation.getInferenceIndices().isEmpty()) {
51-
// No inference fields were identified, so return the original query.
52-
return queryBuilder;
53-
} else if (indexInformation.nonInferenceIndices().isEmpty() == false) {
54-
// Combined case where the field name requested by this query contains both
55-
// semantic_text and non-inference fields, so we have to combine queries per index
56-
// containing each field type.
57-
return buildCombinedInferenceAndNonInferenceQuery(queryBuilder, indexInformation);
58-
} else {
59-
// The only fields we've identified are inference fields (e.g. semantic_text),
60-
// so rewrite the entire query to work on a semantic_text field.
61-
return buildInferenceQuery(queryBuilder, indexInformation);
44+
BoolQueryBuilder finalQueryBuilder = new BoolQueryBuilder();
45+
for (String fieldName : fieldNamesWithWeights.keySet()) {
46+
InferenceIndexInformationForField indexInformation = resolveIndicesForField(fieldName, resolvedIndices);
47+
if (indexInformation.getInferenceIndices().isEmpty()) {
48+
// No inference fields were identified, so return the original query.
49+
finalQueryBuilder.should(queryBuilder);
50+
} else if (indexInformation.nonInferenceIndices().isEmpty() == false) {
51+
// Combined case where the field name requested by this query contains both
52+
// semantic_text and non-inference fields, so we have to combine queries per index
53+
// containing each field type.
54+
return finalQueryBuilder.should(buildCombinedInferenceAndNonInferenceQuery(queryBuilder, indexInformation));
55+
} else {
56+
// The only fields we've identified are inference fields (e.g. semantic_text),
57+
// so rewrite the entire query to work on a semantic_text field.
58+
return finalQueryBuilder.should(buildInferenceQuery(queryBuilder, indexInformation));
59+
}
6260
}
61+
return finalQueryBuilder;
6362
}
6463

6564
/**

0 commit comments

Comments
 (0)