From 7f649e04a678abdb2fcf0913ce974ae36ea8f4c9 Mon Sep 17 00:00:00 2001 From: Samiul Monir Date: Thu, 24 Jul 2025 21:47:22 -0400 Subject: [PATCH 1/8] update fieldNames to accomodate per field boosting from now --- .../SemanticKnnVectorQueryRewriteInterceptor.java | 4 ++-- .../SemanticMatchQueryRewriteInterceptor.java | 6 ++++-- .../queries/SemanticQueryRewriteInterceptor.java | 12 +++++++++--- .../SemanticSparseVectorQueryRewriteInterceptor.java | 4 ++-- 4 files changed, 17 insertions(+), 9 deletions(-) diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticKnnVectorQueryRewriteInterceptor.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticKnnVectorQueryRewriteInterceptor.java index b1f5c240371f8..581ebb45b4dd1 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticKnnVectorQueryRewriteInterceptor.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticKnnVectorQueryRewriteInterceptor.java @@ -33,10 +33,10 @@ public class SemanticKnnVectorQueryRewriteInterceptor extends SemanticQueryRewri public SemanticKnnVectorQueryRewriteInterceptor() {} @Override - protected String getFieldName(QueryBuilder queryBuilder) { + protected Map getFieldNamesWithWeights(QueryBuilder queryBuilder) { assert (queryBuilder instanceof KnnVectorQueryBuilder); KnnVectorQueryBuilder knnVectorQueryBuilder = (KnnVectorQueryBuilder) queryBuilder; - return knnVectorQueryBuilder.getFieldName(); + return Map.of(knnVectorQueryBuilder.getFieldName(), 1.0f); } @Override 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 index a6599afc66c3f..d9e34f76bba8d 100644 --- 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 @@ -12,6 +12,8 @@ import org.elasticsearch.index.query.MatchQueryBuilder; import org.elasticsearch.index.query.QueryBuilder; +import java.util.Map; + public class SemanticMatchQueryRewriteInterceptor extends SemanticQueryRewriteInterceptor { public static final NodeFeature SEMANTIC_MATCH_QUERY_REWRITE_INTERCEPTION_SUPPORTED = new NodeFeature( @@ -21,10 +23,10 @@ public class SemanticMatchQueryRewriteInterceptor extends SemanticQueryRewriteIn public SemanticMatchQueryRewriteInterceptor() {} @Override - protected String getFieldName(QueryBuilder queryBuilder) { + protected Map getFieldNamesWithWeights(QueryBuilder queryBuilder) { assert (queryBuilder instanceof MatchQueryBuilder); MatchQueryBuilder matchQueryBuilder = (MatchQueryBuilder) queryBuilder; - return matchQueryBuilder.fieldName(); + return Map.of(matchQueryBuilder.fieldName(), 1.0f); } @Override diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryRewriteInterceptor.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryRewriteInterceptor.java index bb76ef0be24e9..dd39549b54e70 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryRewriteInterceptor.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryRewriteInterceptor.java @@ -33,7 +33,12 @@ public SemanticQueryRewriteInterceptor() {} @Override public QueryBuilder interceptAndRewrite(QueryRewriteContext context, QueryBuilder queryBuilder) { - String fieldName = getFieldName(queryBuilder); + Map fieldNamesWithWeights = getFieldNamesWithWeights(queryBuilder); + if (fieldNamesWithWeights.size() > 1) { + // Multi-field query, so return the original query, and an exception will be thrown eventually + return queryBuilder; + } + String fieldName = fieldNamesWithWeights.keySet().iterator().next(); ResolvedIndices resolvedIndices = context.getResolvedIndices(); if (resolvedIndices == null) { @@ -59,9 +64,10 @@ public QueryBuilder interceptAndRewrite(QueryRewriteContext context, QueryBuilde /** * @param queryBuilder {@link QueryBuilder} - * @return The singular field name requested by the provided query builder. + * @return Map of field names with their weights for multi-field queries. + * For single-field queries, return a map with one entry. */ - protected abstract String getFieldName(QueryBuilder queryBuilder); + protected abstract Map getFieldNamesWithWeights(QueryBuilder queryBuilder); /** * @param queryBuilder {@link QueryBuilder} diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticSparseVectorQueryRewriteInterceptor.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticSparseVectorQueryRewriteInterceptor.java index c85a21f10301d..1e20aa240d887 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticSparseVectorQueryRewriteInterceptor.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticSparseVectorQueryRewriteInterceptor.java @@ -27,10 +27,10 @@ public class SemanticSparseVectorQueryRewriteInterceptor extends SemanticQueryRe public SemanticSparseVectorQueryRewriteInterceptor() {} @Override - protected String getFieldName(QueryBuilder queryBuilder) { + protected Map getFieldNamesWithWeights(QueryBuilder queryBuilder) { assert (queryBuilder instanceof SparseVectorQueryBuilder); SparseVectorQueryBuilder sparseVectorQueryBuilder = (SparseVectorQueryBuilder) queryBuilder; - return sparseVectorQueryBuilder.getFieldName(); + return Map.of(sparseVectorQueryBuilder.getFieldName(), 1.0f); } @Override From e663f7004412c97e0c5c798f62525b71aaf7c195 Mon Sep 17 00:00:00 2001 From: Samiul Monir Date: Fri, 25 Jul 2025 13:11:11 -0400 Subject: [PATCH 2/8] Introducing multimatch intercepter --- .../xpack/inference/InferencePlugin.java | 4 +- ...nticMultiMatchQueryRewriteInterceptor.java | 106 ++++++++++++++++++ .../SemanticQueryRewriteInterceptor.java | 35 +++--- 3 files changed, 126 insertions(+), 19 deletions(-) create mode 100644 x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticMultiMatchQueryRewriteInterceptor.java 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 bbb1bd1a2fec2..642cd6f76421c 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 @@ -95,6 +95,7 @@ import org.elasticsearch.xpack.inference.mapper.SemanticTextFieldMapper; import org.elasticsearch.xpack.inference.queries.SemanticKnnVectorQueryRewriteInterceptor; import org.elasticsearch.xpack.inference.queries.SemanticMatchQueryRewriteInterceptor; +import org.elasticsearch.xpack.inference.queries.SemanticMultiMatchQueryRewriteInterceptor; import org.elasticsearch.xpack.inference.queries.SemanticQueryBuilder; import org.elasticsearch.xpack.inference.queries.SemanticSparseVectorQueryRewriteInterceptor; import org.elasticsearch.xpack.inference.rank.random.RandomRankBuilder; @@ -556,7 +557,8 @@ public List getQueryRewriteInterceptors() { return List.of( new SemanticKnnVectorQueryRewriteInterceptor(), new SemanticMatchQueryRewriteInterceptor(), - new SemanticSparseVectorQueryRewriteInterceptor() + new SemanticSparseVectorQueryRewriteInterceptor(), + new SemanticMultiMatchQueryRewriteInterceptor() ); } diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticMultiMatchQueryRewriteInterceptor.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticMultiMatchQueryRewriteInterceptor.java new file mode 100644 index 0000000000000..b758fc8c6a75c --- /dev/null +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticMultiMatchQueryRewriteInterceptor.java @@ -0,0 +1,106 @@ +/* + * 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.index.query.BoolQueryBuilder; +import org.elasticsearch.index.query.MultiMatchQueryBuilder; +import org.elasticsearch.index.query.QueryBuilder; + +import java.util.Map; + +public class SemanticMultiMatchQueryRewriteInterceptor extends SemanticQueryRewriteInterceptor { + @Override + protected Map getFieldNamesWithWeights(QueryBuilder queryBuilder) { + assert (queryBuilder instanceof MultiMatchQueryBuilder); + MultiMatchQueryBuilder multiMatchQueryBuilder = (MultiMatchQueryBuilder) queryBuilder; + return multiMatchQueryBuilder.fields(); + } + + @Override + protected String getQuery(QueryBuilder queryBuilder) { + assert (queryBuilder instanceof MultiMatchQueryBuilder); + MultiMatchQueryBuilder multiMatchQueryBuilder = (MultiMatchQueryBuilder) queryBuilder; + return (String) multiMatchQueryBuilder.value(); + } + + @Override + protected QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation) { + SemanticQueryBuilder semanticQueryBuilder = new SemanticQueryBuilder( + indexInformation.fieldName(), + getQuery(queryBuilder), + false + ); + // TODO:: add boost + semanticQueryBuilder.queryName(queryBuilder.queryName()); + return semanticQueryBuilder; + } + + @Override + protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation) { + assert (queryBuilder instanceof MultiMatchQueryBuilder); + MultiMatchQueryBuilder originalMultiMatchQueryBuilder = (MultiMatchQueryBuilder) queryBuilder; + + // Create a copy for non-inference fields with only this specific field + MultiMatchQueryBuilder multiMatchQueryBuilder = createSingleFieldMultiMatch( + originalMultiMatchQueryBuilder, + indexInformation.fieldName() + ); + + BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); + + // Add semantic query for inference indices + boolQueryBuilder.should( + createSemanticSubQuery( + indexInformation.getInferenceIndices(), + indexInformation.fieldName(), + getQuery(queryBuilder) + ) + ); + + // Add regular query for non-inference indices + boolQueryBuilder.should( + createSubQueryForIndices(indexInformation.nonInferenceIndices(), multiMatchQueryBuilder) + ); + + // TODO:: add boost + boolQueryBuilder.queryName(queryBuilder.queryName()); + return boolQueryBuilder; + } + + @Override + public String getQueryName() { + return MultiMatchQueryBuilder.NAME; + } + + /** + * Create a MultiMatchQueryBuilder with only a single field for non-inference indices + */ + private MultiMatchQueryBuilder createSingleFieldMultiMatch(MultiMatchQueryBuilder original, String fieldName) { + MultiMatchQueryBuilder singleFieldQuery = new MultiMatchQueryBuilder(original.value()); + + // Copy all properties from original query + singleFieldQuery.type(original.type()); + singleFieldQuery.operator(original.operator()); + singleFieldQuery.analyzer(original.analyzer()); + singleFieldQuery.fuzziness(original.fuzziness()); + singleFieldQuery.prefixLength(original.prefixLength()); + singleFieldQuery.maxExpansions(original.maxExpansions()); + singleFieldQuery.minimumShouldMatch(original.minimumShouldMatch()); + singleFieldQuery.fuzzyRewrite(original.fuzzyRewrite()); + singleFieldQuery.tieBreaker(original.tieBreaker()); + singleFieldQuery.lenient(original.lenient()); + singleFieldQuery.zeroTermsQuery(original.zeroTermsQuery()); + singleFieldQuery.autoGenerateSynonymsPhraseQuery(original.autoGenerateSynonymsPhraseQuery()); + singleFieldQuery.fuzzyTranspositions(original.fuzzyTranspositions()); + + // Add only the specific field (without boost for now) + singleFieldQuery.field(fieldName); + + return singleFieldQuery; + } +} diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryRewriteInterceptor.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryRewriteInterceptor.java index dd39549b54e70..02a02b85a8e42 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryRewriteInterceptor.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryRewriteInterceptor.java @@ -34,11 +34,6 @@ public SemanticQueryRewriteInterceptor() {} @Override public QueryBuilder interceptAndRewrite(QueryRewriteContext context, QueryBuilder queryBuilder) { Map fieldNamesWithWeights = getFieldNamesWithWeights(queryBuilder); - if (fieldNamesWithWeights.size() > 1) { - // Multi-field query, so return the original query, and an exception will be thrown eventually - return queryBuilder; - } - String fieldName = fieldNamesWithWeights.keySet().iterator().next(); ResolvedIndices resolvedIndices = context.getResolvedIndices(); if (resolvedIndices == null) { @@ -46,20 +41,24 @@ public QueryBuilder interceptAndRewrite(QueryRewriteContext context, QueryBuilde return queryBuilder; } - InferenceIndexInformationForField indexInformation = resolveIndicesForField(fieldName, resolvedIndices); - if (indexInformation.getInferenceIndices().isEmpty()) { - // No inference fields were identified, so return the original query. - return queryBuilder; - } else if (indexInformation.nonInferenceIndices().isEmpty() == false) { - // Combined case where the field name requested by this query contains both - // semantic_text and non-inference fields, so we have to combine queries per index - // containing each field type. - return buildCombinedInferenceAndNonInferenceQuery(queryBuilder, indexInformation); - } else { - // The only fields we've identified are inference fields (e.g. semantic_text), - // so rewrite the entire query to work on a semantic_text field. - return buildInferenceQuery(queryBuilder, indexInformation); + BoolQueryBuilder finalQueryBuilder = new BoolQueryBuilder(); + for (String fieldName : fieldNamesWithWeights.keySet()) { + InferenceIndexInformationForField indexInformation = resolveIndicesForField(fieldName, resolvedIndices); + if (indexInformation.getInferenceIndices().isEmpty()) { + // No inference fields were identified, so return the original query. + finalQueryBuilder.should(queryBuilder); + } else if (indexInformation.nonInferenceIndices().isEmpty() == false) { + // Combined case where the field name requested by this query contains both + // semantic_text and non-inference fields, so we have to combine queries per index + // containing each field type. + return finalQueryBuilder.should(buildCombinedInferenceAndNonInferenceQuery(queryBuilder, indexInformation)); + } else { + // The only fields we've identified are inference fields (e.g. semantic_text), + // so rewrite the entire query to work on a semantic_text field. + return finalQueryBuilder.should(buildInferenceQuery(queryBuilder, indexInformation)); + } } + return finalQueryBuilder; } /** From 84161fd64bdc5fd068b42148c1a4f082cb46a47a Mon Sep 17 00:00:00 2001 From: Samiul Monir Date: Fri, 25 Jul 2025 13:52:45 -0400 Subject: [PATCH 3/8] Supply boosting to multi_match --- ...SemanticKnnVectorQueryRewriteInterceptor.java | 14 ++++++++++++-- .../SemanticMatchQueryRewriteInterceptor.java | 14 ++++++++++++-- ...emanticMultiMatchQueryRewriteInterceptor.java | 16 +++++++++++++--- .../queries/SemanticQueryRewriteInterceptor.java | 11 +++++++---- ...anticSparseVectorQueryRewriteInterceptor.java | 14 ++++++++++++-- 5 files changed, 56 insertions(+), 13 deletions(-) diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticKnnVectorQueryRewriteInterceptor.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticKnnVectorQueryRewriteInterceptor.java index 581ebb45b4dd1..75deba9e35ad8 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticKnnVectorQueryRewriteInterceptor.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticKnnVectorQueryRewriteInterceptor.java @@ -47,8 +47,7 @@ protected String getQuery(QueryBuilder queryBuilder) { return queryVectorBuilder != null ? queryVectorBuilder.getModelText() : null; } - @Override - protected QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation) { + private QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation) { assert (queryBuilder instanceof KnnVectorQueryBuilder); KnnVectorQueryBuilder knnVectorQueryBuilder = (KnnVectorQueryBuilder) queryBuilder; Map> inferenceIdsIndices = indexInformation.getInferenceIdsIndices(); @@ -68,6 +67,17 @@ protected QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceI return finalQueryBuilder; } + @Override + protected QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation, Float fieldWeight) { + QueryBuilder inferenceQuery = buildInferenceQuery(queryBuilder, indexInformation); + + if (fieldWeight != null && fieldWeight.equals(1.0f) == false) { + inferenceQuery.boost(fieldWeight); + } + + return inferenceQuery; + } + private QueryBuilder buildInferenceQueryWithMultipleInferenceIds( KnnVectorQueryBuilder queryBuilder, Map> inferenceIdsIndices 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 index d9e34f76bba8d..0cf0b19c6bfc5 100644 --- 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 @@ -36,14 +36,24 @@ protected String getQuery(QueryBuilder queryBuilder) { return (String) matchQueryBuilder.value(); } - @Override - protected QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation) { + private QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation) { SemanticQueryBuilder semanticQueryBuilder = new SemanticQueryBuilder(indexInformation.fieldName(), getQuery(queryBuilder), false); semanticQueryBuilder.boost(queryBuilder.boost()); semanticQueryBuilder.queryName(queryBuilder.queryName()); return semanticQueryBuilder; } + @Override + protected QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation, Float fieldWeight) { + QueryBuilder inferenceQuery = buildInferenceQuery(queryBuilder, indexInformation); + + if (fieldWeight != null && fieldWeight.equals(1.0f) == false) { + inferenceQuery.boost(fieldWeight); + } + + return inferenceQuery; + } + @Override protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery( QueryBuilder queryBuilder, diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticMultiMatchQueryRewriteInterceptor.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticMultiMatchQueryRewriteInterceptor.java index b758fc8c6a75c..a636ebb4041ec 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticMultiMatchQueryRewriteInterceptor.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticMultiMatchQueryRewriteInterceptor.java @@ -28,18 +28,28 @@ protected String getQuery(QueryBuilder queryBuilder) { return (String) multiMatchQueryBuilder.value(); } - @Override - protected QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation) { + private QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation) { SemanticQueryBuilder semanticQueryBuilder = new SemanticQueryBuilder( indexInformation.fieldName(), getQuery(queryBuilder), false ); - // TODO:: add boost + semanticQueryBuilder.boost(queryBuilder.boost()); semanticQueryBuilder.queryName(queryBuilder.queryName()); return semanticQueryBuilder; } + @Override + protected QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation, Float fieldWeight) { + QueryBuilder inferenceQuery = buildInferenceQuery(queryBuilder, indexInformation); + + if (fieldWeight != null && fieldWeight.equals(1.0f) == false) { + inferenceQuery.boost(fieldWeight); + } + + return inferenceQuery; + } + @Override protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation) { assert (queryBuilder instanceof MultiMatchQueryBuilder); diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryRewriteInterceptor.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryRewriteInterceptor.java index 02a02b85a8e42..32ba719487ecd 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryRewriteInterceptor.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryRewriteInterceptor.java @@ -42,7 +42,9 @@ public QueryBuilder interceptAndRewrite(QueryRewriteContext context, QueryBuilde } BoolQueryBuilder finalQueryBuilder = new BoolQueryBuilder(); - for (String fieldName : fieldNamesWithWeights.keySet()) { + for (Map.Entry fieldSet : fieldNamesWithWeights.entrySet()) { + String fieldName = fieldSet.getKey(); + Float fieldWeight = fieldSet.getValue(); InferenceIndexInformationForField indexInformation = resolveIndicesForField(fieldName, resolvedIndices); if (indexInformation.getInferenceIndices().isEmpty()) { // No inference fields were identified, so return the original query. @@ -51,11 +53,11 @@ public QueryBuilder interceptAndRewrite(QueryRewriteContext context, QueryBuilde // Combined case where the field name requested by this query contains both // semantic_text and non-inference fields, so we have to combine queries per index // containing each field type. - return finalQueryBuilder.should(buildCombinedInferenceAndNonInferenceQuery(queryBuilder, indexInformation)); + finalQueryBuilder.should(buildCombinedInferenceAndNonInferenceQuery(queryBuilder, indexInformation)); } else { // The only fields we've identified are inference fields (e.g. semantic_text), // so rewrite the entire query to work on a semantic_text field. - return finalQueryBuilder.should(buildInferenceQuery(queryBuilder, indexInformation)); + finalQueryBuilder.should(buildInferenceQuery(queryBuilder, indexInformation, fieldWeight)); } } return finalQueryBuilder; @@ -79,9 +81,10 @@ public QueryBuilder interceptAndRewrite(QueryRewriteContext context, QueryBuilde * * @param queryBuilder {@link QueryBuilder} * @param indexInformation {@link InferenceIndexInformationForField} + * @param weight {@link Float} * @return {@link QueryBuilder} */ - protected abstract QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation); + protected abstract QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation, Float weight); /** * Builds a combined inference and non-inference query, diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticSparseVectorQueryRewriteInterceptor.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticSparseVectorQueryRewriteInterceptor.java index 1e20aa240d887..f40368f4db3a5 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticSparseVectorQueryRewriteInterceptor.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticSparseVectorQueryRewriteInterceptor.java @@ -40,8 +40,7 @@ protected String getQuery(QueryBuilder queryBuilder) { return sparseVectorQueryBuilder.getQuery(); } - @Override - protected QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation) { + private QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation) { Map> inferenceIdsIndices = indexInformation.getInferenceIdsIndices(); QueryBuilder finalQueryBuilder; if (inferenceIdsIndices.size() == 1) { @@ -57,6 +56,17 @@ protected QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceI return finalQueryBuilder; } + @Override + protected QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation, Float fieldWeight) { + QueryBuilder inferenceQuery = buildInferenceQuery(queryBuilder, indexInformation); + + if (fieldWeight != null && fieldWeight.equals(1.0f) == false) { + inferenceQuery.boost(fieldWeight); + } + + return inferenceQuery; + } + private QueryBuilder buildInferenceQueryWithMultipleInferenceIds( QueryBuilder queryBuilder, Map> inferenceIdsIndices From 1e21f9606c8a1fde6af3c4f9c0fa94316e6196a5 Mon Sep 17 00:00:00 2001 From: Samiul Monir Date: Fri, 25 Jul 2025 16:34:41 -0400 Subject: [PATCH 4/8] working commit --- ...anticKnnVectorQueryRewriteInterceptor.java | 18 ++++++++- .../SemanticMatchQueryRewriteInterceptor.java | 16 +++++++- ...nticMultiMatchQueryRewriteInterceptor.java | 40 ++++++++++++++++++- .../SemanticQueryRewriteInterceptor.java | 10 ++++- ...icSparseVectorQueryRewriteInterceptor.java | 15 ++++++- 5 files changed, 91 insertions(+), 8 deletions(-) diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticKnnVectorQueryRewriteInterceptor.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticKnnVectorQueryRewriteInterceptor.java index 75deba9e35ad8..74da674be452c 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticKnnVectorQueryRewriteInterceptor.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticKnnVectorQueryRewriteInterceptor.java @@ -94,8 +94,7 @@ private QueryBuilder buildInferenceQueryWithMultipleInferenceIds( return boolQueryBuilder; } - @Override - protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery( + private QueryBuilder buildCombinedInferenceAndNonInferenceQuery( QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation ) { @@ -121,6 +120,21 @@ protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery( return boolQueryBuilder; } + @Override + protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery( + QueryBuilder queryBuilder, + InferenceIndexInformationForField indexInformation, + Float fieldWeight + ) { + QueryBuilder inferenceQuery = buildCombinedInferenceAndNonInferenceQuery(queryBuilder, indexInformation); + + if (fieldWeight != null && fieldWeight.equals(1.0f) == false) { + inferenceQuery.boost(fieldWeight); + } + + return inferenceQuery; + } + private QueryBuilder buildNestedQueryFromKnnVectorQuery( KnnVectorQueryBuilder knnVectorQueryBuilder, List indices, 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 index 0cf0b19c6bfc5..95764ca09d0e6 100644 --- 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 @@ -54,7 +54,6 @@ protected QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceI return inferenceQuery; } - @Override protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery( QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation @@ -78,6 +77,21 @@ protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery( return boolQueryBuilder; } + @Override + protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery( + QueryBuilder queryBuilder, + InferenceIndexInformationForField indexInformation, + Float fieldWeight + ) { + QueryBuilder inferenceQuery = buildCombinedInferenceAndNonInferenceQuery(queryBuilder, indexInformation); + + if (fieldWeight != null && fieldWeight.equals(1.0f) == false) { + inferenceQuery.boost(fieldWeight); + } + + return inferenceQuery; + } + @Override public String getQueryName() { return MatchQueryBuilder.NAME; diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticMultiMatchQueryRewriteInterceptor.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticMultiMatchQueryRewriteInterceptor.java index a636ebb4041ec..b75c6b897051f 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticMultiMatchQueryRewriteInterceptor.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticMultiMatchQueryRewriteInterceptor.java @@ -8,6 +8,7 @@ package org.elasticsearch.xpack.inference.queries; import org.elasticsearch.index.query.BoolQueryBuilder; +import org.elasticsearch.index.query.MatchQueryBuilder; import org.elasticsearch.index.query.MultiMatchQueryBuilder; import org.elasticsearch.index.query.QueryBuilder; @@ -50,8 +51,7 @@ protected QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceI return inferenceQuery; } - @Override - protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation) { + private QueryBuilder buildCombinedInferenceAndNonInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation) { assert (queryBuilder instanceof MultiMatchQueryBuilder); MultiMatchQueryBuilder originalMultiMatchQueryBuilder = (MultiMatchQueryBuilder) queryBuilder; @@ -82,6 +82,42 @@ protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery(QueryBuilder q return boolQueryBuilder; } + @Override + protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation, Float fieldWeight) { + assert (queryBuilder instanceof MultiMatchQueryBuilder); + MultiMatchQueryBuilder originalMultiMatchQueryBuilder = (MultiMatchQueryBuilder) queryBuilder; + + MatchQueryBuilder matchQueryBuilder = new MatchQueryBuilder(indexInformation.fieldName(), getQuery(queryBuilder)); + + // Create a copy for non-inference fields with only this specific field +// MultiMatchQueryBuilder multiMatchQueryBuilder = createSingleFieldMultiMatch( +// originalMultiMatchQueryBuilder, +// indexInformation.fieldName() +// ); + + BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); + + // Add semantic query for inference indices + boolQueryBuilder.should( + createSemanticSubQuery( + indexInformation.getInferenceIndices(), + indexInformation.fieldName(), + getQuery(matchQueryBuilder) + ) + ); + + // Add regular query for non-inference indices + boolQueryBuilder.should( + createSubQueryForIndices(indexInformation.nonInferenceIndices(), matchQueryBuilder) + ); + + if (fieldWeight != null && fieldWeight.equals(1.0f) == false) { + boolQueryBuilder.boost(fieldWeight); + } + boolQueryBuilder.queryName(queryBuilder.queryName()); + return boolQueryBuilder; + } + @Override public String getQueryName() { return MultiMatchQueryBuilder.NAME; diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryRewriteInterceptor.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryRewriteInterceptor.java index 32ba719487ecd..971e965969391 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryRewriteInterceptor.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryRewriteInterceptor.java @@ -53,7 +53,7 @@ public QueryBuilder interceptAndRewrite(QueryRewriteContext context, QueryBuilde // Combined case where the field name requested by this query contains both // semantic_text and non-inference fields, so we have to combine queries per index // containing each field type. - finalQueryBuilder.should(buildCombinedInferenceAndNonInferenceQuery(queryBuilder, indexInformation)); + finalQueryBuilder.should(buildCombinedInferenceAndNonInferenceQuery(queryBuilder, indexInformation, fieldWeight)); } else { // The only fields we've identified are inference fields (e.g. semantic_text), // so rewrite the entire query to work on a semantic_text field. @@ -93,9 +93,15 @@ public QueryBuilder interceptAndRewrite(QueryRewriteContext context, QueryBuilde * @param indexInformation {@link InferenceIndexInformationForField} * @return {@link QueryBuilder} */ +// protected abstract QueryBuilder buildCombinedInferenceAndNonInferenceQuery( +// QueryBuilder queryBuilder, +// InferenceIndexInformationForField indexInformation +// ); + protected abstract QueryBuilder buildCombinedInferenceAndNonInferenceQuery( QueryBuilder queryBuilder, - InferenceIndexInformationForField indexInformation + InferenceIndexInformationForField indexInformation, + Float fieldWeight ); private InferenceIndexInformationForField resolveIndicesForField(String fieldName, ResolvedIndices resolvedIndices) { diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticSparseVectorQueryRewriteInterceptor.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticSparseVectorQueryRewriteInterceptor.java index f40368f4db3a5..2d28de8414428 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticSparseVectorQueryRewriteInterceptor.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticSparseVectorQueryRewriteInterceptor.java @@ -83,7 +83,6 @@ private QueryBuilder buildInferenceQueryWithMultipleInferenceIds( return boolQueryBuilder; } - @Override protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery( QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation @@ -121,6 +120,20 @@ protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery( return boolQueryBuilder; } + protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery( + QueryBuilder queryBuilder, + InferenceIndexInformationForField indexInformation, + Float fieldWeight + ) { + QueryBuilder inferenceQuery = buildCombinedInferenceAndNonInferenceQuery(queryBuilder, indexInformation); + + if (fieldWeight != null && fieldWeight.equals(1.0f) == false) { + inferenceQuery.boost(fieldWeight); + } + + return inferenceQuery; + } + private QueryBuilder buildNestedQueryFromSparseVectorQuery(QueryBuilder queryBuilder, String searchInferenceId) { assert (queryBuilder instanceof SparseVectorQueryBuilder); SparseVectorQueryBuilder sparseVectorQueryBuilder = (SparseVectorQueryBuilder) queryBuilder; From 742fbdc647b8ebe0e122072cdb47eada73f425c9 Mon Sep 17 00:00:00 2001 From: elasticsearchmachine Date: Fri, 25 Jul 2025 20:45:22 +0000 Subject: [PATCH 5/8] [CI] Auto commit changes from spotless --- ...anticKnnVectorQueryRewriteInterceptor.java | 6 ++- .../SemanticMatchQueryRewriteInterceptor.java | 6 ++- ...nticMultiMatchQueryRewriteInterceptor.java | 53 +++++++++---------- .../SemanticQueryRewriteInterceptor.java | 14 +++-- ...icSparseVectorQueryRewriteInterceptor.java | 6 ++- 5 files changed, 48 insertions(+), 37 deletions(-) diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticKnnVectorQueryRewriteInterceptor.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticKnnVectorQueryRewriteInterceptor.java index 74da674be452c..71ba43259c257 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticKnnVectorQueryRewriteInterceptor.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticKnnVectorQueryRewriteInterceptor.java @@ -68,7 +68,11 @@ private QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceInd } @Override - protected QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation, Float fieldWeight) { + protected QueryBuilder buildInferenceQuery( + QueryBuilder queryBuilder, + InferenceIndexInformationForField indexInformation, + Float fieldWeight + ) { QueryBuilder inferenceQuery = buildInferenceQuery(queryBuilder, indexInformation); if (fieldWeight != null && fieldWeight.equals(1.0f) == false) { 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 index 95764ca09d0e6..06554a53c0888 100644 --- 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 @@ -44,7 +44,11 @@ private QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceInd } @Override - protected QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation, Float fieldWeight) { + protected QueryBuilder buildInferenceQuery( + QueryBuilder queryBuilder, + InferenceIndexInformationForField indexInformation, + Float fieldWeight + ) { QueryBuilder inferenceQuery = buildInferenceQuery(queryBuilder, indexInformation); if (fieldWeight != null && fieldWeight.equals(1.0f) == false) { diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticMultiMatchQueryRewriteInterceptor.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticMultiMatchQueryRewriteInterceptor.java index b75c6b897051f..26ef9b85bdf0a 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticMultiMatchQueryRewriteInterceptor.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticMultiMatchQueryRewriteInterceptor.java @@ -17,7 +17,7 @@ public class SemanticMultiMatchQueryRewriteInterceptor extends SemanticQueryRewriteInterceptor { @Override protected Map getFieldNamesWithWeights(QueryBuilder queryBuilder) { - assert (queryBuilder instanceof MultiMatchQueryBuilder); + assert (queryBuilder instanceof MultiMatchQueryBuilder); MultiMatchQueryBuilder multiMatchQueryBuilder = (MultiMatchQueryBuilder) queryBuilder; return multiMatchQueryBuilder.fields(); } @@ -30,18 +30,18 @@ protected String getQuery(QueryBuilder queryBuilder) { } private QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation) { - SemanticQueryBuilder semanticQueryBuilder = new SemanticQueryBuilder( - indexInformation.fieldName(), - getQuery(queryBuilder), - false - ); + SemanticQueryBuilder semanticQueryBuilder = new SemanticQueryBuilder(indexInformation.fieldName(), getQuery(queryBuilder), false); semanticQueryBuilder.boost(queryBuilder.boost()); semanticQueryBuilder.queryName(queryBuilder.queryName()); return semanticQueryBuilder; } @Override - protected QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation, Float fieldWeight) { + protected QueryBuilder buildInferenceQuery( + QueryBuilder queryBuilder, + InferenceIndexInformationForField indexInformation, + Float fieldWeight + ) { QueryBuilder inferenceQuery = buildInferenceQuery(queryBuilder, indexInformation); if (fieldWeight != null && fieldWeight.equals(1.0f) == false) { @@ -51,7 +51,10 @@ protected QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceI return inferenceQuery; } - private QueryBuilder buildCombinedInferenceAndNonInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation) { + private QueryBuilder buildCombinedInferenceAndNonInferenceQuery( + QueryBuilder queryBuilder, + InferenceIndexInformationForField indexInformation + ) { assert (queryBuilder instanceof MultiMatchQueryBuilder); MultiMatchQueryBuilder originalMultiMatchQueryBuilder = (MultiMatchQueryBuilder) queryBuilder; @@ -65,17 +68,11 @@ private QueryBuilder buildCombinedInferenceAndNonInferenceQuery(QueryBuilder que // Add semantic query for inference indices boolQueryBuilder.should( - createSemanticSubQuery( - indexInformation.getInferenceIndices(), - indexInformation.fieldName(), - getQuery(queryBuilder) - ) + createSemanticSubQuery(indexInformation.getInferenceIndices(), indexInformation.fieldName(), getQuery(queryBuilder)) ); // Add regular query for non-inference indices - boolQueryBuilder.should( - createSubQueryForIndices(indexInformation.nonInferenceIndices(), multiMatchQueryBuilder) - ); + boolQueryBuilder.should(createSubQueryForIndices(indexInformation.nonInferenceIndices(), multiMatchQueryBuilder)); // TODO:: add boost boolQueryBuilder.queryName(queryBuilder.queryName()); @@ -83,33 +80,31 @@ private QueryBuilder buildCombinedInferenceAndNonInferenceQuery(QueryBuilder que } @Override - protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation, Float fieldWeight) { + protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery( + QueryBuilder queryBuilder, + InferenceIndexInformationForField indexInformation, + Float fieldWeight + ) { assert (queryBuilder instanceof MultiMatchQueryBuilder); MultiMatchQueryBuilder originalMultiMatchQueryBuilder = (MultiMatchQueryBuilder) queryBuilder; MatchQueryBuilder matchQueryBuilder = new MatchQueryBuilder(indexInformation.fieldName(), getQuery(queryBuilder)); // Create a copy for non-inference fields with only this specific field -// MultiMatchQueryBuilder multiMatchQueryBuilder = createSingleFieldMultiMatch( -// originalMultiMatchQueryBuilder, -// indexInformation.fieldName() -// ); + // MultiMatchQueryBuilder multiMatchQueryBuilder = createSingleFieldMultiMatch( + // originalMultiMatchQueryBuilder, + // indexInformation.fieldName() + // ); BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); // Add semantic query for inference indices boolQueryBuilder.should( - createSemanticSubQuery( - indexInformation.getInferenceIndices(), - indexInformation.fieldName(), - getQuery(matchQueryBuilder) - ) + createSemanticSubQuery(indexInformation.getInferenceIndices(), indexInformation.fieldName(), getQuery(matchQueryBuilder)) ); // Add regular query for non-inference indices - boolQueryBuilder.should( - createSubQueryForIndices(indexInformation.nonInferenceIndices(), matchQueryBuilder) - ); + boolQueryBuilder.should(createSubQueryForIndices(indexInformation.nonInferenceIndices(), matchQueryBuilder)); if (fieldWeight != null && fieldWeight.equals(1.0f) == false) { boolQueryBuilder.boost(fieldWeight); diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryRewriteInterceptor.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryRewriteInterceptor.java index 971e965969391..a5a739cdd6206 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryRewriteInterceptor.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryRewriteInterceptor.java @@ -84,7 +84,11 @@ public QueryBuilder interceptAndRewrite(QueryRewriteContext context, QueryBuilde * @param weight {@link Float} * @return {@link QueryBuilder} */ - protected abstract QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation, Float weight); + protected abstract QueryBuilder buildInferenceQuery( + QueryBuilder queryBuilder, + InferenceIndexInformationForField indexInformation, + Float weight + ); /** * Builds a combined inference and non-inference query, @@ -93,10 +97,10 @@ public QueryBuilder interceptAndRewrite(QueryRewriteContext context, QueryBuilde * @param indexInformation {@link InferenceIndexInformationForField} * @return {@link QueryBuilder} */ -// protected abstract QueryBuilder buildCombinedInferenceAndNonInferenceQuery( -// QueryBuilder queryBuilder, -// InferenceIndexInformationForField indexInformation -// ); + // protected abstract QueryBuilder buildCombinedInferenceAndNonInferenceQuery( + // QueryBuilder queryBuilder, + // InferenceIndexInformationForField indexInformation + // ); protected abstract QueryBuilder buildCombinedInferenceAndNonInferenceQuery( QueryBuilder queryBuilder, diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticSparseVectorQueryRewriteInterceptor.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticSparseVectorQueryRewriteInterceptor.java index 2d28de8414428..2513465546edf 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticSparseVectorQueryRewriteInterceptor.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticSparseVectorQueryRewriteInterceptor.java @@ -57,7 +57,11 @@ private QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceInd } @Override - protected QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation, Float fieldWeight) { + protected QueryBuilder buildInferenceQuery( + QueryBuilder queryBuilder, + InferenceIndexInformationForField indexInformation, + Float fieldWeight + ) { QueryBuilder inferenceQuery = buildInferenceQuery(queryBuilder, indexInformation); if (fieldWeight != null && fieldWeight.equals(1.0f) == false) { From 344e2de92cc92bb20b7bf30d0747acb7d54012f2 Mon Sep 17 00:00:00 2001 From: Samiul Monir Date: Mon, 28 Jul 2025 11:22:37 -0400 Subject: [PATCH 6/8] handle combined queries for inference and non-inference fields --- ...nticMultiMatchQueryRewriteInterceptor.java | 32 ++++--- .../SemanticQueryRewriteInterceptor.java | 93 ++++++++++++++++--- 2 files changed, 100 insertions(+), 25 deletions(-) diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticMultiMatchQueryRewriteInterceptor.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticMultiMatchQueryRewriteInterceptor.java index 26ef9b85bdf0a..17ce709cd455d 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticMultiMatchQueryRewriteInterceptor.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticMultiMatchQueryRewriteInterceptor.java @@ -86,25 +86,31 @@ protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery( Float fieldWeight ) { assert (queryBuilder instanceof MultiMatchQueryBuilder); - MultiMatchQueryBuilder originalMultiMatchQueryBuilder = (MultiMatchQueryBuilder) queryBuilder; - - MatchQueryBuilder matchQueryBuilder = new MatchQueryBuilder(indexInformation.fieldName(), getQuery(queryBuilder)); - - // Create a copy for non-inference fields with only this specific field - // MultiMatchQueryBuilder multiMatchQueryBuilder = createSingleFieldMultiMatch( - // originalMultiMatchQueryBuilder, - // indexInformation.fieldName() - // ); + MultiMatchQueryBuilder multiMatchQueryBuilder = (MultiMatchQueryBuilder) queryBuilder; BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); + boolQueryBuilder.should( + createSemanticSubQuery(indexInformation.getInferenceIndices(), indexInformation.fieldName(), (String) multiMatchQueryBuilder.value()) + ); - // Add semantic query for inference indices boolQueryBuilder.should( - createSemanticSubQuery(indexInformation.getInferenceIndices(), indexInformation.fieldName(), getQuery(matchQueryBuilder)) + createMatchSubQuery( + indexInformation.nonInferenceIndices(), + indexInformation.fieldName(), + (String) multiMatchQueryBuilder.value() + ) ); - // Add regular query for non-inference indices - boolQueryBuilder.should(createSubQueryForIndices(indexInformation.nonInferenceIndices(), matchQueryBuilder)); + + QueryBuilder nonSemanticFieldQuery = buildNonSemanticFieldQuery( + queryBuilder, + indexInformation.fieldName(), + fieldWeight + ); +// boolQueryBuilder.should( +// createSubQueryForIndices(indexInformation.nonInferenceIndices(), nonSemanticFieldQuery) +// ); +// boolQueryBuilder.should(createSubQueryForIndices(indexInformation.nonInferenceIndices(), multiMatchQueryBuilder)); if (fieldWeight != null && fieldWeight.equals(1.0f) == false) { boolQueryBuilder.boost(fieldWeight); diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryRewriteInterceptor.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryRewriteInterceptor.java index a5a739cdd6206..426c0f639df18 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryRewriteInterceptor.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryRewriteInterceptor.java @@ -12,6 +12,7 @@ import org.elasticsearch.cluster.metadata.InferenceFieldMetadata; 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.TermsQueryBuilder; @@ -41,25 +42,71 @@ public QueryBuilder interceptAndRewrite(QueryRewriteContext context, QueryBuilde return queryBuilder; } + if (fieldNamesWithWeights.size() > 1) { + // Multi-field query, so return the original query. + return handleMultiFieldQuery(queryBuilder, fieldNamesWithWeights, resolvedIndices); + } + + String fieldName = fieldNamesWithWeights.keySet().iterator().next(); + Float weight = fieldNamesWithWeights.get(fieldName); + InferenceIndexInformationForField indexInformation = resolveIndicesForField(fieldName, resolvedIndices); + if (indexInformation.getInferenceIndices().isEmpty()) { + // No inference fields were identified, so return the original query. + return queryBuilder; + } else if (indexInformation.nonInferenceIndices().isEmpty() == false) { + // Combined case where the field name requested by this query contains both + // semantic_text and non-inference fields, so we have to combine queries per index + // containing each field type. + return buildCombinedInferenceAndNonInferenceQuery(queryBuilder, indexInformation, weight); + } else { + // The only fields we've identified are inference fields (e.g. semantic_text), + // so rewrite the entire query to work on a semantic_text field. + return buildInferenceQuery(queryBuilder, indexInformation, weight); + } + } + + /** + * Handle multi-field queries (new logic) + */ + private QueryBuilder handleMultiFieldQuery( + QueryBuilder queryBuilder, + Map fieldNamesWithWeights, + ResolvedIndices resolvedIndices + ) { BoolQueryBuilder finalQueryBuilder = new BoolQueryBuilder(); - for (Map.Entry fieldSet : fieldNamesWithWeights.entrySet()) { - String fieldName = fieldSet.getKey(); - Float fieldWeight = fieldSet.getValue(); + boolean hasAnySemanticFields = false; + + for (Map.Entry fieldEntry : fieldNamesWithWeights.entrySet()) { + String fieldName = fieldEntry.getKey(); + Float fieldWeight = fieldEntry.getValue(); InferenceIndexInformationForField indexInformation = resolveIndicesForField(fieldName, resolvedIndices); + if (indexInformation.getInferenceIndices().isEmpty()) { - // No inference fields were identified, so return the original query. - finalQueryBuilder.should(queryBuilder); + // Pure non-semantic field - create individual match query + QueryBuilder nonSemanticQuery = createMatchSubQuery( + indexInformation.nonInferenceIndices(), + fieldName, + getQuery(queryBuilder)); + finalQueryBuilder.should(nonSemanticQuery); } else if (indexInformation.nonInferenceIndices().isEmpty() == false) { - // Combined case where the field name requested by this query contains both - // semantic_text and non-inference fields, so we have to combine queries per index - // containing each field type. - finalQueryBuilder.should(buildCombinedInferenceAndNonInferenceQuery(queryBuilder, indexInformation, fieldWeight)); + // Mixed semantic/non-semantic field - use combined approach + QueryBuilder combinedQuery = buildCombinedInferenceAndNonInferenceQuery(queryBuilder, indexInformation, fieldWeight); + finalQueryBuilder.should(combinedQuery); + hasAnySemanticFields = true; } else { - // The only fields we've identified are inference fields (e.g. semantic_text), - // so rewrite the entire query to work on a semantic_text field. - finalQueryBuilder.should(buildInferenceQuery(queryBuilder, indexInformation, fieldWeight)); + // Pure semantic field - create semantic query + QueryBuilder semanticQuery = buildInferenceQuery(queryBuilder, indexInformation, fieldWeight); + finalQueryBuilder.should(semanticQuery); + hasAnySemanticFields = true; } } + + // If no semantic fields were found, return original query + if (hasAnySemanticFields == false) { + return queryBuilder; + } + +// finalQueryBuilder.minimumShouldMatch(1); return finalQueryBuilder; } @@ -125,6 +172,28 @@ private InferenceIndexInformationForField resolveIndicesForField(String fieldNam return new InferenceIndexInformationForField(fieldName, inferenceIndicesMetadata, nonInferenceIndices); } + /** + * Build a non-semantic field query (for multi-field scenarios) + */ + protected QueryBuilder buildNonSemanticFieldQuery(QueryBuilder queryBuilder, String fieldName, Float fieldWeight) { + // Default implementation - subclasses can override for specific query types + String query = getQuery(queryBuilder); + MatchQueryBuilder matchQueryBuilder = new MatchQueryBuilder(fieldName, query); + matchQueryBuilder.boost(fieldWeight); + return matchQueryBuilder; + } + + protected QueryBuilder createMatchSubQuery(Collection indices, String fieldName, String queryText) { + BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); + MatchQueryBuilder matchQuery = new MatchQueryBuilder(fieldName, queryText); +// if (fieldWeight != null && !fieldWeight.equals(1.0f)) { +// matchQuery.boost(fieldWeight); +// } + boolQueryBuilder.must(matchQuery); + boolQueryBuilder.filter(new TermsQueryBuilder(IndexFieldMapper.NAME, indices)); + return boolQueryBuilder; + } + protected QueryBuilder createSubQueryForIndices(Collection indices, QueryBuilder queryBuilder) { BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); boolQueryBuilder.must(queryBuilder); From 791552b367cca77714367816ce25abb2d332b358 Mon Sep 17 00:00:00 2001 From: Samiul Monir Date: Mon, 28 Jul 2025 16:32:07 -0400 Subject: [PATCH 7/8] simplified query generation --- ...anticKnnVectorQueryRewriteInterceptor.java | 49 +++------- .../SemanticMatchQueryRewriteInterceptor.java | 43 ++------- ...nticMultiMatchQueryRewriteInterceptor.java | 95 ++----------------- .../SemanticQueryRewriteInterceptor.java | 50 ++++------ ...icSparseVectorQueryRewriteInterceptor.java | 45 +++------ 5 files changed, 58 insertions(+), 224 deletions(-) diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticKnnVectorQueryRewriteInterceptor.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticKnnVectorQueryRewriteInterceptor.java index 71ba43259c257..05821effa2f3a 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticKnnVectorQueryRewriteInterceptor.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticKnnVectorQueryRewriteInterceptor.java @@ -33,7 +33,7 @@ public class SemanticKnnVectorQueryRewriteInterceptor extends SemanticQueryRewri public SemanticKnnVectorQueryRewriteInterceptor() {} @Override - protected Map getFieldNamesWithWeights(QueryBuilder queryBuilder) { + protected Map getFieldNamesWithBoosts(QueryBuilder queryBuilder) { assert (queryBuilder instanceof KnnVectorQueryBuilder); KnnVectorQueryBuilder knnVectorQueryBuilder = (KnnVectorQueryBuilder) queryBuilder; return Map.of(knnVectorQueryBuilder.getFieldName(), 1.0f); @@ -47,7 +47,12 @@ protected String getQuery(QueryBuilder queryBuilder) { return queryVectorBuilder != null ? queryVectorBuilder.getModelText() : null; } - private QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation) { + @Override + protected QueryBuilder buildInferenceQuery( + QueryBuilder queryBuilder, + InferenceIndexInformationForField indexInformation, + Float fieldBoost + ) { assert (queryBuilder instanceof KnnVectorQueryBuilder); KnnVectorQueryBuilder knnVectorQueryBuilder = (KnnVectorQueryBuilder) queryBuilder; Map> inferenceIdsIndices = indexInformation.getInferenceIdsIndices(); @@ -62,26 +67,11 @@ private QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceInd // Multiple inference IDs, construct a boolean query finalQueryBuilder = buildInferenceQueryWithMultipleInferenceIds(knnVectorQueryBuilder, inferenceIdsIndices); } - finalQueryBuilder.boost(queryBuilder.boost()); + finalQueryBuilder.boost(queryBuilder.boost() * fieldBoost); finalQueryBuilder.queryName(queryBuilder.queryName()); return finalQueryBuilder; } - @Override - protected QueryBuilder buildInferenceQuery( - QueryBuilder queryBuilder, - InferenceIndexInformationForField indexInformation, - Float fieldWeight - ) { - QueryBuilder inferenceQuery = buildInferenceQuery(queryBuilder, indexInformation); - - if (fieldWeight != null && fieldWeight.equals(1.0f) == false) { - inferenceQuery.boost(fieldWeight); - } - - return inferenceQuery; - } - private QueryBuilder buildInferenceQueryWithMultipleInferenceIds( KnnVectorQueryBuilder queryBuilder, Map> inferenceIdsIndices @@ -98,9 +88,11 @@ private QueryBuilder buildInferenceQueryWithMultipleInferenceIds( return boolQueryBuilder; } - private QueryBuilder buildCombinedInferenceAndNonInferenceQuery( + @Override + protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery( QueryBuilder queryBuilder, - InferenceIndexInformationForField indexInformation + InferenceIndexInformationForField indexInformation, + Float fieldBoost ) { assert (queryBuilder instanceof KnnVectorQueryBuilder); KnnVectorQueryBuilder knnVectorQueryBuilder = (KnnVectorQueryBuilder) queryBuilder; @@ -119,26 +111,11 @@ private QueryBuilder buildCombinedInferenceAndNonInferenceQuery( ) ); } - boolQueryBuilder.boost(queryBuilder.boost()); + boolQueryBuilder.boost(queryBuilder.boost() * fieldBoost); boolQueryBuilder.queryName(queryBuilder.queryName()); return boolQueryBuilder; } - @Override - protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery( - QueryBuilder queryBuilder, - InferenceIndexInformationForField indexInformation, - Float fieldWeight - ) { - QueryBuilder inferenceQuery = buildCombinedInferenceAndNonInferenceQuery(queryBuilder, indexInformation); - - if (fieldWeight != null && fieldWeight.equals(1.0f) == false) { - inferenceQuery.boost(fieldWeight); - } - - return inferenceQuery; - } - private QueryBuilder buildNestedQueryFromKnnVectorQuery( KnnVectorQueryBuilder knnVectorQueryBuilder, List indices, 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 index 06554a53c0888..915234f4a5a66 100644 --- 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 @@ -23,7 +23,7 @@ public class SemanticMatchQueryRewriteInterceptor extends SemanticQueryRewriteIn public SemanticMatchQueryRewriteInterceptor() {} @Override - protected Map getFieldNamesWithWeights(QueryBuilder queryBuilder) { + protected Map getFieldNamesWithBoosts(QueryBuilder queryBuilder) { assert (queryBuilder instanceof MatchQueryBuilder); MatchQueryBuilder matchQueryBuilder = (MatchQueryBuilder) queryBuilder; return Map.of(matchQueryBuilder.fieldName(), 1.0f); @@ -36,31 +36,23 @@ protected String getQuery(QueryBuilder queryBuilder) { return (String) matchQueryBuilder.value(); } - private QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation) { - SemanticQueryBuilder semanticQueryBuilder = new SemanticQueryBuilder(indexInformation.fieldName(), getQuery(queryBuilder), false); - semanticQueryBuilder.boost(queryBuilder.boost()); - semanticQueryBuilder.queryName(queryBuilder.queryName()); - return semanticQueryBuilder; - } - @Override protected QueryBuilder buildInferenceQuery( QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation, - Float fieldWeight + Float fieldBoost ) { - QueryBuilder inferenceQuery = buildInferenceQuery(queryBuilder, indexInformation); - - if (fieldWeight != null && fieldWeight.equals(1.0f) == false) { - inferenceQuery.boost(fieldWeight); - } - - return inferenceQuery; + SemanticQueryBuilder semanticQueryBuilder = new SemanticQueryBuilder(indexInformation.fieldName(), getQuery(queryBuilder), false); + semanticQueryBuilder.boost(queryBuilder.boost() * fieldBoost); + semanticQueryBuilder.queryName(queryBuilder.queryName()); + return semanticQueryBuilder; } + @Override protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery( QueryBuilder queryBuilder, - InferenceIndexInformationForField indexInformation + InferenceIndexInformationForField indexInformation, + Float fieldBoost ) { assert (queryBuilder instanceof MatchQueryBuilder); MatchQueryBuilder originalMatchQueryBuilder = (MatchQueryBuilder) queryBuilder; @@ -76,26 +68,11 @@ protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery( ) ); boolQueryBuilder.should(createSubQueryForIndices(indexInformation.nonInferenceIndices(), matchQueryBuilder)); - boolQueryBuilder.boost(queryBuilder.boost()); + boolQueryBuilder.boost(queryBuilder.boost() * fieldBoost); boolQueryBuilder.queryName(queryBuilder.queryName()); return boolQueryBuilder; } - @Override - protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery( - QueryBuilder queryBuilder, - InferenceIndexInformationForField indexInformation, - Float fieldWeight - ) { - QueryBuilder inferenceQuery = buildCombinedInferenceAndNonInferenceQuery(queryBuilder, indexInformation); - - if (fieldWeight != null && fieldWeight.equals(1.0f) == false) { - inferenceQuery.boost(fieldWeight); - } - - return inferenceQuery; - } - @Override public String getQueryName() { return MatchQueryBuilder.NAME; diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticMultiMatchQueryRewriteInterceptor.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticMultiMatchQueryRewriteInterceptor.java index 17ce709cd455d..36edabee26203 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticMultiMatchQueryRewriteInterceptor.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticMultiMatchQueryRewriteInterceptor.java @@ -8,7 +8,6 @@ package org.elasticsearch.xpack.inference.queries; import org.elasticsearch.index.query.BoolQueryBuilder; -import org.elasticsearch.index.query.MatchQueryBuilder; import org.elasticsearch.index.query.MultiMatchQueryBuilder; import org.elasticsearch.index.query.QueryBuilder; @@ -16,7 +15,7 @@ public class SemanticMultiMatchQueryRewriteInterceptor extends SemanticQueryRewriteInterceptor { @Override - protected Map getFieldNamesWithWeights(QueryBuilder queryBuilder) { + protected Map getFieldNamesWithBoosts(QueryBuilder queryBuilder) { assert (queryBuilder instanceof MultiMatchQueryBuilder); MultiMatchQueryBuilder multiMatchQueryBuilder = (MultiMatchQueryBuilder) queryBuilder; return multiMatchQueryBuilder.fields(); @@ -29,61 +28,23 @@ protected String getQuery(QueryBuilder queryBuilder) { return (String) multiMatchQueryBuilder.value(); } - private QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation) { - SemanticQueryBuilder semanticQueryBuilder = new SemanticQueryBuilder(indexInformation.fieldName(), getQuery(queryBuilder), false); - semanticQueryBuilder.boost(queryBuilder.boost()); - semanticQueryBuilder.queryName(queryBuilder.queryName()); - return semanticQueryBuilder; - } - @Override protected QueryBuilder buildInferenceQuery( QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation, - Float fieldWeight + Float fieldWBoost ) { - QueryBuilder inferenceQuery = buildInferenceQuery(queryBuilder, indexInformation); - - if (fieldWeight != null && fieldWeight.equals(1.0f) == false) { - inferenceQuery.boost(fieldWeight); - } - - return inferenceQuery; - } - - private QueryBuilder buildCombinedInferenceAndNonInferenceQuery( - QueryBuilder queryBuilder, - InferenceIndexInformationForField indexInformation - ) { - assert (queryBuilder instanceof MultiMatchQueryBuilder); - MultiMatchQueryBuilder originalMultiMatchQueryBuilder = (MultiMatchQueryBuilder) queryBuilder; - - // Create a copy for non-inference fields with only this specific field - MultiMatchQueryBuilder multiMatchQueryBuilder = createSingleFieldMultiMatch( - originalMultiMatchQueryBuilder, - indexInformation.fieldName() - ); - - BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); - - // Add semantic query for inference indices - boolQueryBuilder.should( - createSemanticSubQuery(indexInformation.getInferenceIndices(), indexInformation.fieldName(), getQuery(queryBuilder)) - ); - - // Add regular query for non-inference indices - boolQueryBuilder.should(createSubQueryForIndices(indexInformation.nonInferenceIndices(), multiMatchQueryBuilder)); - - // TODO:: add boost - boolQueryBuilder.queryName(queryBuilder.queryName()); - return boolQueryBuilder; + SemanticQueryBuilder semanticQueryBuilder = new SemanticQueryBuilder(indexInformation.fieldName(), getQuery(queryBuilder), false); + semanticQueryBuilder.boost(queryBuilder.boost() * fieldWBoost); + semanticQueryBuilder.queryName(queryBuilder.queryName()); + return semanticQueryBuilder; } @Override protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery( QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation, - Float fieldWeight + Float fieldBoost ) { assert (queryBuilder instanceof MultiMatchQueryBuilder); MultiMatchQueryBuilder multiMatchQueryBuilder = (MultiMatchQueryBuilder) queryBuilder; @@ -101,20 +62,7 @@ protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery( ) ); - - QueryBuilder nonSemanticFieldQuery = buildNonSemanticFieldQuery( - queryBuilder, - indexInformation.fieldName(), - fieldWeight - ); -// boolQueryBuilder.should( -// createSubQueryForIndices(indexInformation.nonInferenceIndices(), nonSemanticFieldQuery) -// ); -// boolQueryBuilder.should(createSubQueryForIndices(indexInformation.nonInferenceIndices(), multiMatchQueryBuilder)); - - if (fieldWeight != null && fieldWeight.equals(1.0f) == false) { - boolQueryBuilder.boost(fieldWeight); - } + boolQueryBuilder.boost(queryBuilder.boost() * fieldBoost); boolQueryBuilder.queryName(queryBuilder.queryName()); return boolQueryBuilder; } @@ -123,31 +71,4 @@ protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery( public String getQueryName() { return MultiMatchQueryBuilder.NAME; } - - /** - * Create a MultiMatchQueryBuilder with only a single field for non-inference indices - */ - private MultiMatchQueryBuilder createSingleFieldMultiMatch(MultiMatchQueryBuilder original, String fieldName) { - MultiMatchQueryBuilder singleFieldQuery = new MultiMatchQueryBuilder(original.value()); - - // Copy all properties from original query - singleFieldQuery.type(original.type()); - singleFieldQuery.operator(original.operator()); - singleFieldQuery.analyzer(original.analyzer()); - singleFieldQuery.fuzziness(original.fuzziness()); - singleFieldQuery.prefixLength(original.prefixLength()); - singleFieldQuery.maxExpansions(original.maxExpansions()); - singleFieldQuery.minimumShouldMatch(original.minimumShouldMatch()); - singleFieldQuery.fuzzyRewrite(original.fuzzyRewrite()); - singleFieldQuery.tieBreaker(original.tieBreaker()); - singleFieldQuery.lenient(original.lenient()); - singleFieldQuery.zeroTermsQuery(original.zeroTermsQuery()); - singleFieldQuery.autoGenerateSynonymsPhraseQuery(original.autoGenerateSynonymsPhraseQuery()); - singleFieldQuery.fuzzyTranspositions(original.fuzzyTranspositions()); - - // Add only the specific field (without boost for now) - singleFieldQuery.field(fieldName); - - return singleFieldQuery; - } } diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryRewriteInterceptor.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryRewriteInterceptor.java index 426c0f639df18..d78eb848d6eed 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryRewriteInterceptor.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryRewriteInterceptor.java @@ -34,7 +34,7 @@ public SemanticQueryRewriteInterceptor() {} @Override public QueryBuilder interceptAndRewrite(QueryRewriteContext context, QueryBuilder queryBuilder) { - Map fieldNamesWithWeights = getFieldNamesWithWeights(queryBuilder); + Map fieldsWithBoosts = getFieldNamesWithBoosts(queryBuilder); ResolvedIndices resolvedIndices = context.getResolvedIndices(); if (resolvedIndices == null) { @@ -42,13 +42,13 @@ public QueryBuilder interceptAndRewrite(QueryRewriteContext context, QueryBuilde return queryBuilder; } - if (fieldNamesWithWeights.size() > 1) { + if (fieldsWithBoosts.size() > 1) { // Multi-field query, so return the original query. - return handleMultiFieldQuery(queryBuilder, fieldNamesWithWeights, resolvedIndices); + return handleMultiFieldQuery(queryBuilder, fieldsWithBoosts, resolvedIndices); } - String fieldName = fieldNamesWithWeights.keySet().iterator().next(); - Float weight = fieldNamesWithWeights.get(fieldName); + String fieldName = fieldsWithBoosts.keySet().iterator().next(); + Float weight = fieldsWithBoosts.get(fieldName); InferenceIndexInformationForField indexInformation = resolveIndicesForField(fieldName, resolvedIndices); if (indexInformation.getInferenceIndices().isEmpty()) { // No inference fields were identified, so return the original query. @@ -106,16 +106,16 @@ private QueryBuilder handleMultiFieldQuery( return queryBuilder; } -// finalQueryBuilder.minimumShouldMatch(1); return finalQueryBuilder; } /** - * @param queryBuilder {@link QueryBuilder} - * @return Map of field names with their weights for multi-field queries. - * For single-field queries, return a map with one entry. + * Extracts field names and their associated boost values from the query builder. + * + * @param queryBuilder the query builder to extract field information from + * @return a map where keys are field names and values are their boost multipliers */ - protected abstract Map getFieldNamesWithWeights(QueryBuilder queryBuilder); + protected abstract Map getFieldNamesWithBoosts(QueryBuilder queryBuilder); /** * @param queryBuilder {@link QueryBuilder} @@ -128,13 +128,13 @@ private QueryBuilder handleMultiFieldQuery( * * @param queryBuilder {@link QueryBuilder} * @param indexInformation {@link InferenceIndexInformationForField} - * @param weight {@link Float} + * @param fieldBoost per field boost value * @return {@link QueryBuilder} */ protected abstract QueryBuilder buildInferenceQuery( QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation, - Float weight + Float fieldBoost ); /** @@ -142,17 +142,13 @@ protected abstract QueryBuilder buildInferenceQuery( * which separates the different queries into appropriate indices based on field type. * @param queryBuilder {@link QueryBuilder} * @param indexInformation {@link InferenceIndexInformationForField} + * @param fieldBoost per field boost value * @return {@link QueryBuilder} */ - // protected abstract QueryBuilder buildCombinedInferenceAndNonInferenceQuery( - // QueryBuilder queryBuilder, - // InferenceIndexInformationForField indexInformation - // ); - protected abstract QueryBuilder buildCombinedInferenceAndNonInferenceQuery( QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation, - Float fieldWeight + Float fieldBoost ); private InferenceIndexInformationForField resolveIndicesForField(String fieldName, ResolvedIndices resolvedIndices) { @@ -172,23 +168,9 @@ private InferenceIndexInformationForField resolveIndicesForField(String fieldNam return new InferenceIndexInformationForField(fieldName, inferenceIndicesMetadata, nonInferenceIndices); } - /** - * Build a non-semantic field query (for multi-field scenarios) - */ - protected QueryBuilder buildNonSemanticFieldQuery(QueryBuilder queryBuilder, String fieldName, Float fieldWeight) { - // Default implementation - subclasses can override for specific query types - String query = getQuery(queryBuilder); - MatchQueryBuilder matchQueryBuilder = new MatchQueryBuilder(fieldName, query); - matchQueryBuilder.boost(fieldWeight); - return matchQueryBuilder; - } - - protected QueryBuilder createMatchSubQuery(Collection indices, String fieldName, String queryText) { + protected QueryBuilder createMatchSubQuery(Collection indices, String fieldName, String value) { BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); - MatchQueryBuilder matchQuery = new MatchQueryBuilder(fieldName, queryText); -// if (fieldWeight != null && !fieldWeight.equals(1.0f)) { -// matchQuery.boost(fieldWeight); -// } + MatchQueryBuilder matchQuery = new MatchQueryBuilder(fieldName, value); boolQueryBuilder.must(matchQuery); boolQueryBuilder.filter(new TermsQueryBuilder(IndexFieldMapper.NAME, indices)); return boolQueryBuilder; diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticSparseVectorQueryRewriteInterceptor.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticSparseVectorQueryRewriteInterceptor.java index 2513465546edf..8230356190a7f 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticSparseVectorQueryRewriteInterceptor.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticSparseVectorQueryRewriteInterceptor.java @@ -27,7 +27,7 @@ public class SemanticSparseVectorQueryRewriteInterceptor extends SemanticQueryRe public SemanticSparseVectorQueryRewriteInterceptor() {} @Override - protected Map getFieldNamesWithWeights(QueryBuilder queryBuilder) { + protected Map getFieldNamesWithBoosts(QueryBuilder queryBuilder) { assert (queryBuilder instanceof SparseVectorQueryBuilder); SparseVectorQueryBuilder sparseVectorQueryBuilder = (SparseVectorQueryBuilder) queryBuilder; return Map.of(sparseVectorQueryBuilder.getFieldName(), 1.0f); @@ -40,7 +40,12 @@ protected String getQuery(QueryBuilder queryBuilder) { return sparseVectorQueryBuilder.getQuery(); } - private QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceIndexInformationForField indexInformation) { + @Override + protected QueryBuilder buildInferenceQuery( + QueryBuilder queryBuilder, + InferenceIndexInformationForField indexInformation, + Float fieldBoost + ) { Map> inferenceIdsIndices = indexInformation.getInferenceIdsIndices(); QueryBuilder finalQueryBuilder; if (inferenceIdsIndices.size() == 1) { @@ -52,25 +57,10 @@ private QueryBuilder buildInferenceQuery(QueryBuilder queryBuilder, InferenceInd finalQueryBuilder = buildInferenceQueryWithMultipleInferenceIds(queryBuilder, inferenceIdsIndices); } finalQueryBuilder.queryName(queryBuilder.queryName()); - finalQueryBuilder.boost(queryBuilder.boost()); + finalQueryBuilder.boost(queryBuilder.boost() * fieldBoost); return finalQueryBuilder; } - @Override - protected QueryBuilder buildInferenceQuery( - QueryBuilder queryBuilder, - InferenceIndexInformationForField indexInformation, - Float fieldWeight - ) { - QueryBuilder inferenceQuery = buildInferenceQuery(queryBuilder, indexInformation); - - if (fieldWeight != null && fieldWeight.equals(1.0f) == false) { - inferenceQuery.boost(fieldWeight); - } - - return inferenceQuery; - } - private QueryBuilder buildInferenceQueryWithMultipleInferenceIds( QueryBuilder queryBuilder, Map> inferenceIdsIndices @@ -89,7 +79,8 @@ private QueryBuilder buildInferenceQueryWithMultipleInferenceIds( protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery( QueryBuilder queryBuilder, - InferenceIndexInformationForField indexInformation + InferenceIndexInformationForField indexInformation, + Float fieldBoost ) { assert (queryBuilder instanceof SparseVectorQueryBuilder); SparseVectorQueryBuilder sparseVectorQueryBuilder = (SparseVectorQueryBuilder) queryBuilder; @@ -119,25 +110,11 @@ protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery( ) ); } - boolQueryBuilder.boost(queryBuilder.boost()); + boolQueryBuilder.boost(queryBuilder.boost() * fieldBoost); boolQueryBuilder.queryName(queryBuilder.queryName()); return boolQueryBuilder; } - protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery( - QueryBuilder queryBuilder, - InferenceIndexInformationForField indexInformation, - Float fieldWeight - ) { - QueryBuilder inferenceQuery = buildCombinedInferenceAndNonInferenceQuery(queryBuilder, indexInformation); - - if (fieldWeight != null && fieldWeight.equals(1.0f) == false) { - inferenceQuery.boost(fieldWeight); - } - - return inferenceQuery; - } - private QueryBuilder buildNestedQueryFromSparseVectorQuery(QueryBuilder queryBuilder, String searchInferenceId) { assert (queryBuilder instanceof SparseVectorQueryBuilder); SparseVectorQueryBuilder sparseVectorQueryBuilder = (SparseVectorQueryBuilder) queryBuilder; From b93b16aa1fb842015f8ce14eb7f97e06e0b452ec Mon Sep 17 00:00:00 2001 From: Samiul Monir Date: Mon, 28 Jul 2025 16:32:52 -0400 Subject: [PATCH 8/8] linting --- .../queries/SemanticMultiMatchQueryRewriteInterceptor.java | 6 +++++- .../inference/queries/SemanticQueryRewriteInterceptor.java | 3 ++- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticMultiMatchQueryRewriteInterceptor.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticMultiMatchQueryRewriteInterceptor.java index 36edabee26203..39d4486232a9c 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticMultiMatchQueryRewriteInterceptor.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticMultiMatchQueryRewriteInterceptor.java @@ -51,7 +51,11 @@ protected QueryBuilder buildCombinedInferenceAndNonInferenceQuery( BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); boolQueryBuilder.should( - createSemanticSubQuery(indexInformation.getInferenceIndices(), indexInformation.fieldName(), (String) multiMatchQueryBuilder.value()) + createSemanticSubQuery( + indexInformation.getInferenceIndices(), + indexInformation.fieldName(), + (String) multiMatchQueryBuilder.value() + ) ); boolQueryBuilder.should( diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryRewriteInterceptor.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryRewriteInterceptor.java index d78eb848d6eed..85d5c0feee4b9 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryRewriteInterceptor.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryRewriteInterceptor.java @@ -86,7 +86,8 @@ private QueryBuilder handleMultiFieldQuery( QueryBuilder nonSemanticQuery = createMatchSubQuery( indexInformation.nonInferenceIndices(), fieldName, - getQuery(queryBuilder)); + getQuery(queryBuilder) + ); finalQueryBuilder.should(nonSemanticQuery); } else if (indexInformation.nonInferenceIndices().isEmpty() == false) { // Mixed semantic/non-semantic field - use combined approach