Skip to content

Commit 053ab13

Browse files
working commit
1 parent a1d3ef3 commit 053ab13

File tree

2 files changed

+144
-8
lines changed

2 files changed

+144
-8
lines changed

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

Lines changed: 63 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,10 @@
77

88
package org.elasticsearch.xpack.inference.queries;
99

10+
import org.elasticsearch.index.query.BoolQueryBuilder;
1011
import org.elasticsearch.index.query.MultiMatchQueryBuilder;
1112
import org.elasticsearch.index.query.QueryBuilder;
13+
import org.elasticsearch.index.query.QueryBuilders;
1214

1315
import java.util.Map;
1416

@@ -33,7 +35,14 @@ protected QueryBuilder buildInferenceQuery(
3335
InferenceIndexInformationForField indexInformation,
3436
Float fieldBoost
3537
) {
36-
return queryBuilder;
38+
SemanticQueryBuilder semanticQueryBuilder = new SemanticQueryBuilder(
39+
indexInformation.fieldName(),
40+
getQuery(queryBuilder),
41+
false
42+
);
43+
semanticQueryBuilder.boost(queryBuilder.boost() * fieldBoost);
44+
semanticQueryBuilder.queryName(queryBuilder.queryName());
45+
return semanticQueryBuilder;
3746
}
3847

3948
@Override
@@ -42,11 +51,63 @@ protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery(
4251
InferenceIndexInformationForField indexInformation,
4352
Float fieldBoost
4453
) {
45-
return queryBuilder;
54+
BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder();
55+
56+
// Add the semantic part for inference indices
57+
boolQueryBuilder.should(
58+
createSemanticSubQuery(
59+
indexInformation.getInferenceIndices(),
60+
indexInformation.fieldName(),
61+
getQuery(queryBuilder)
62+
)
63+
);
64+
65+
// Add the non-semantic part for non-inference indices
66+
boolQueryBuilder.should(
67+
createSubQueryForIndices(
68+
indexInformation.nonInferenceIndices(),
69+
QueryBuilders.matchQuery(indexInformation.fieldName(), getQuery(queryBuilder))
70+
)
71+
);
72+
73+
// Apply the field boost
74+
boolQueryBuilder.boost(queryBuilder.boost() * fieldBoost);
75+
boolQueryBuilder.queryName(queryBuilder.queryName());
76+
77+
return boolQueryBuilder;
4678
}
4779

4880
@Override
4981
public String getQueryName() {
5082
return MultiMatchQueryBuilder.NAME;
5183
}
84+
85+
86+
public static void copyMultiMatchConfiguration(MultiMatchQueryBuilder source, MultiMatchQueryBuilder target) {
87+
target.type(source.type());
88+
target.operator(source.operator());
89+
if (source.analyzer() != null) {
90+
target.analyzer(source.analyzer());
91+
}
92+
target.slop(source.slop());
93+
if (source.fuzziness() != null) {
94+
target.fuzziness(source.fuzziness());
95+
}
96+
target.prefixLength(source.prefixLength());
97+
target.maxExpansions(source.maxExpansions());
98+
if (source.minimumShouldMatch() != null) {
99+
target.minimumShouldMatch(source.minimumShouldMatch());
100+
}
101+
if (source.fuzzyRewrite() != null) {
102+
target.fuzzyRewrite(source.fuzzyRewrite());
103+
}
104+
if (source.tieBreaker() != null) {
105+
target.tieBreaker(source.tieBreaker());
106+
}
107+
target.lenient(source.lenient());
108+
target.zeroTermsQuery(source.zeroTermsQuery());
109+
target.autoGenerateSynonymsPhraseQuery(source.autoGenerateSynonymsPhraseQuery());
110+
target.fuzzyTranspositions(source.fuzzyTranspositions());
111+
}
112+
52113
}

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

Lines changed: 81 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
import org.elasticsearch.index.query.BoolQueryBuilder;
1515
import org.elasticsearch.index.query.MultiMatchQueryBuilder;
1616
import org.elasticsearch.index.query.QueryBuilder;
17+
import org.elasticsearch.index.query.QueryBuilders;
1718
import org.elasticsearch.index.query.QueryRewriteContext;
1819
import org.elasticsearch.index.query.TermsQueryBuilder;
1920
import org.elasticsearch.plugins.internal.rewriter.QueryRewriteInterceptor;
@@ -43,9 +44,7 @@ public QueryBuilder interceptAndRewrite(QueryRewriteContext context, QueryBuilde
4344
}
4445

4546
if (queryBuilder instanceof MultiMatchQueryBuilder) {
46-
// This is a placeholder for the actual multi-field handling logic.
47-
// Going forward, we will also send fieldNamesWithWeights and resolvedIndices
48-
return handleMultiFieldQuery(queryBuilder);
47+
return handleMultiFieldQuery(queryBuilder, fieldsWithBoosts, resolvedIndices);
4948
}
5049

5150
String fieldName = fieldsWithBoosts.keySet().iterator().next();
@@ -68,10 +67,86 @@ public QueryBuilder interceptAndRewrite(QueryRewriteContext context, QueryBuilde
6867
}
6968

7069
/**
71-
* Handle multi-field queries
70+
* Handle multi-field queries by analyzing each field for semantic_text fields
71+
* and creating appropriate queries
7272
*/
73-
private QueryBuilder handleMultiFieldQuery(QueryBuilder queryBuilder) {
74-
return queryBuilder;
73+
private QueryBuilder handleMultiFieldQuery(
74+
QueryBuilder queryBuilder,
75+
Map<String, Float> fieldsWithBoosts,
76+
ResolvedIndices resolvedIndices
77+
) {
78+
assert (queryBuilder instanceof MultiMatchQueryBuilder);
79+
MultiMatchQueryBuilder multiMatchQueryBuilder = (MultiMatchQueryBuilder) queryBuilder;
80+
String queryText = getQuery(queryBuilder);
81+
82+
boolean hasSemanticField = false;
83+
Map<String, InferenceIndexInformationForField> fieldInfoMap = new HashMap<>();
84+
85+
// Analyze each field to determine if it's semantic or not
86+
for (Map.Entry<String, Float> fieldEntry : fieldsWithBoosts.entrySet()) {
87+
String fieldName = fieldEntry.getKey();
88+
InferenceIndexInformationForField indexInfo = resolveIndicesForField(fieldName, resolvedIndices);
89+
fieldInfoMap.put(fieldName, indexInfo);
90+
91+
if (indexInfo.getInferenceIndices().isEmpty() ==false) {
92+
hasSemanticField = true;
93+
}
94+
}
95+
96+
// If no semantic fields were found, return the original query
97+
if (hasSemanticField == false) {
98+
return queryBuilder;
99+
}
100+
101+
// Create a combined query
102+
BoolQueryBuilder combinedQuery = new BoolQueryBuilder();
103+
104+
// Apply the MultiMatch type and tie-breaker
105+
MultiMatchQueryBuilder.Type type = multiMatchQueryBuilder.type();
106+
Float tieBreaker = multiMatchQueryBuilder.tieBreaker();
107+
boolean shouldUseTieBreaker = (tieBreaker != null);
108+
109+
110+
Map<String, Float> nonSemanticFields = new HashMap<>();
111+
for (Map.Entry<String, Float> fieldEntry : fieldsWithBoosts.entrySet()) {
112+
String fieldName = fieldEntry.getKey();
113+
Float fieldBoost = fieldEntry.getValue();
114+
InferenceIndexInformationForField indexInfo = fieldInfoMap.get(fieldName);
115+
116+
if (indexInfo.getInferenceIndices().isEmpty()) {
117+
nonSemanticFields.put(fieldName, fieldBoost);
118+
} else if (indexInfo.nonInferenceIndices().isEmpty()) {
119+
QueryBuilder semanticQuery = buildInferenceQuery(queryBuilder, indexInfo, fieldBoost);
120+
if (shouldUseTieBreaker) {
121+
semanticQuery.boost(semanticQuery.boost() * (type == MultiMatchQueryBuilder.Type.MOST_FIELDS ? 1.0f : tieBreaker));
122+
}
123+
combinedQuery.should(semanticQuery);
124+
} else {
125+
QueryBuilder mixedQuery = buildCombinedInferenceAndNonInferenceQuery(queryBuilder, indexInfo, fieldBoost);
126+
if (shouldUseTieBreaker) {
127+
mixedQuery.boost(mixedQuery.boost() * (type == MultiMatchQueryBuilder.Type.MOST_FIELDS ? 1.0f : tieBreaker));
128+
}
129+
combinedQuery.should(mixedQuery);
130+
}
131+
}
132+
133+
if (nonSemanticFields.isEmpty() == false) {
134+
MultiMatchQueryBuilder nonSemanticQuery = QueryBuilders.multiMatchQuery(queryText);
135+
nonSemanticQuery.fields(nonSemanticFields);
136+
137+
SemanticMultiMatchQueryRewriteInterceptor.copyMultiMatchConfiguration(multiMatchQueryBuilder, nonSemanticQuery);
138+
139+
if (shouldUseTieBreaker) {
140+
nonSemanticQuery.boost(nonSemanticQuery.boost() * (type == MultiMatchQueryBuilder.Type.MOST_FIELDS ? 1.0f : tieBreaker));
141+
}
142+
143+
combinedQuery.should(nonSemanticQuery);
144+
}
145+
146+
combinedQuery.boost(queryBuilder.boost());
147+
combinedQuery.queryName(queryBuilder.queryName());
148+
149+
return combinedQuery;
75150
}
76151

77152
/**

0 commit comments

Comments
 (0)