From eaa3054636bca45d09ef39bce115cb0037f483d1 Mon Sep 17 00:00:00 2001 From: Saikat Sarkar Date: Thu, 15 Feb 2024 16:26:39 -0700 Subject: [PATCH 01/21] Display error for text_expansion if the queried field does not have the right type --- .java-version | 1 + .../ml/queries/TextExpansionQueryBuilder.java | 51 ++++++++++++------- .../queries/WeightedTokensQueryBuilder.java | 7 +++ 3 files changed, 41 insertions(+), 18 deletions(-) create mode 100644 .java-version diff --git a/.java-version b/.java-version new file mode 100644 index 0000000000000..189d2e5f1483a --- /dev/null +++ b/.java-version @@ -0,0 +1 @@ +openjdk64-17.0.9 diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilder.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilder.java index 7d5197b9e9ba0..54bfee433fa05 100644 --- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilder.java +++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilder.java @@ -51,6 +51,31 @@ public class TextExpansionQueryBuilder extends AbstractQueryBuilder weightedTokensSupplier; private final TokenPruningConfig tokenPruningConfig; + public enum AllowedFieldTypesForTextExpansion { + RANK_FEATURES("rank_features"), + SPARSE_VECTOR("sparse_vector"); + + private final String typeName; + + AllowedFieldTypesForTextExpansion(String typeName) { + this.typeName = typeName; + } + + public String getTypeName() { + return typeName; + } + + public static boolean contains(String typeName) { + for (AllowedFieldTypesForTextExpansion fieldType : values()) { + if (fieldType.getTypeName().equals(typeName)) { + return true; + } + } + return false; + } + } + + public TextExpansionQueryBuilder(String fieldName, String modelText, String modelId) { this(fieldName, modelText, modelId, null); } @@ -198,24 +223,14 @@ protected QueryBuilder doRewrite(QueryRewriteContext queryRewriteContext) throws } private QueryBuilder weightedTokensToQuery(String fieldName, TextExpansionResults textExpansionResults) { - if (tokenPruningConfig != null) { - WeightedTokensQueryBuilder weightedTokensQueryBuilder = new WeightedTokensQueryBuilder( - fieldName, - textExpansionResults.getWeightedTokens(), - tokenPruningConfig - ); - weightedTokensQueryBuilder.queryName(queryName); - weightedTokensQueryBuilder.boost(boost); - return weightedTokensQueryBuilder; - } - var boolQuery = QueryBuilders.boolQuery(); - for (var weightedToken : textExpansionResults.getWeightedTokens()) { - boolQuery.should(QueryBuilders.termQuery(fieldName, weightedToken.token()).boost(weightedToken.weight())); - } - boolQuery.minimumShouldMatch(1); - boolQuery.boost(this.boost); - boolQuery.queryName(this.queryName); - return boolQuery; + WeightedTokensQueryBuilder weightedTokensQueryBuilder = new WeightedTokensQueryBuilder( + fieldName, + textExpansionResults.getWeightedTokens(), + tokenPruningConfig + ); + weightedTokensQueryBuilder.queryName(queryName); + weightedTokensQueryBuilder.boost(boost); + return weightedTokensQueryBuilder; } @Override diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java index a09bcadaacfc0..ff612a00e969b 100644 --- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java +++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java @@ -35,6 +35,7 @@ import java.util.Objects; import static org.elasticsearch.xpack.ml.queries.TextExpansionQueryBuilder.PRUNING_CONFIG; +import static org.elasticsearch.xpack.ml.queries.TextExpansionQueryBuilder.AllowedFieldTypesForTextExpansion; public class WeightedTokensQueryBuilder extends AbstractQueryBuilder { public static final String NAME = "weighted_tokens"; @@ -152,6 +153,12 @@ protected Query doToQuery(SearchExecutionContext context) throws IOException { if (ft == null) { return new MatchNoDocsQuery("The \"" + getName() + "\" query is against a field that does not exist"); } + + String fieldTypeName = ft.typeName(); + if (AllowedFieldTypesForTextExpansion.contains(fieldTypeName) == false) { + throw new ElasticsearchParseException("[" + fieldTypeName + "]" + " is not an appropriate field type for text expansion query"); + } + var qb = new BooleanQuery.Builder(); int fieldDocCount = context.getIndexReader().getDocCount(fieldName); float bestWeight = 0f; From 496aa0c868fc690643bf55e13f83c37bbe702aa4 Mon Sep 17 00:00:00 2001 From: Saikat Sarkar Date: Thu, 15 Feb 2024 16:46:44 -0700 Subject: [PATCH 02/21] Display error for text_expansion if the queried field does not have the right type --- .../xpack/ml/queries/TextExpansionQueryBuilder.java | 2 -- .../xpack/ml/queries/WeightedTokensQueryBuilder.java | 2 +- 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilder.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilder.java index 54bfee433fa05..cdd2ad8d09e0c 100644 --- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilder.java +++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilder.java @@ -18,7 +18,6 @@ import org.elasticsearch.core.Nullable; import org.elasticsearch.index.query.AbstractQueryBuilder; import org.elasticsearch.index.query.QueryBuilder; -import org.elasticsearch.index.query.QueryBuilders; import org.elasticsearch.index.query.QueryRewriteContext; import org.elasticsearch.index.query.SearchExecutionContext; import org.elasticsearch.xcontent.ParseField; @@ -75,7 +74,6 @@ public static boolean contains(String typeName) { } } - public TextExpansionQueryBuilder(String fieldName, String modelText, String modelId) { this(fieldName, modelText, modelId, null); } diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java index ff612a00e969b..91e486b6811d2 100644 --- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java +++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java @@ -34,8 +34,8 @@ import java.util.List; import java.util.Objects; -import static org.elasticsearch.xpack.ml.queries.TextExpansionQueryBuilder.PRUNING_CONFIG; import static org.elasticsearch.xpack.ml.queries.TextExpansionQueryBuilder.AllowedFieldTypesForTextExpansion; +import static org.elasticsearch.xpack.ml.queries.TextExpansionQueryBuilder.PRUNING_CONFIG; public class WeightedTokensQueryBuilder extends AbstractQueryBuilder { public static final String NAME = "weighted_tokens"; From 3601fa156622fb42641d54e9867a064bbc27ecdd Mon Sep 17 00:00:00 2001 From: Saikat Sarkar Date: Thu, 15 Feb 2024 19:53:01 -0700 Subject: [PATCH 03/21] Display error for text_expansion if the queried field does not have the right type --- .../xpack/ml/queries/WeightedTokensQueryBuilder.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java index 91e486b6811d2..9e1955c6bbe1f 100644 --- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java +++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java @@ -151,7 +151,7 @@ private boolean shouldKeepToken( protected Query doToQuery(SearchExecutionContext context) throws IOException { final MappedFieldType ft = context.getFieldType(fieldName); if (ft == null) { - return new MatchNoDocsQuery("The \"" + getName() + "\" query is against a field that does not exist"); + throw new ElasticsearchParseException("[" + fieldName + "]" + " is not a mapped field"); } String fieldTypeName = ft.typeName(); From d287e38ed097814f3abce91ec4adabaf65de09fd Mon Sep 17 00:00:00 2001 From: Saikat Sarkar Date: Thu, 15 Feb 2024 19:53:40 -0700 Subject: [PATCH 04/21] Display error for text_expansion if the queried field does not have the right type --- .../xpack/ml/queries/WeightedTokensQueryBuilder.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java index 9e1955c6bbe1f..2ce2d0d238e83 100644 --- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java +++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java @@ -154,7 +154,7 @@ protected Query doToQuery(SearchExecutionContext context) throws IOException { throw new ElasticsearchParseException("[" + fieldName + "]" + " is not a mapped field"); } - String fieldTypeName = ft.typeName(); + final String fieldTypeName = ft.typeName(); if (AllowedFieldTypesForTextExpansion.contains(fieldTypeName) == false) { throw new ElasticsearchParseException("[" + fieldTypeName + "]" + " is not an appropriate field type for text expansion query"); } From c29a17ef54547aad3c674f09f8d12cec3d0e711d Mon Sep 17 00:00:00 2001 From: Saikat Sarkar Date: Fri, 16 Feb 2024 13:14:14 -0700 Subject: [PATCH 05/21] Display error for text_expansion if the queried field does not have the right type --- .../resources/rest-api-spec/test/ml/text_expansion_search.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/ml/text_expansion_search.yml b/x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/ml/text_expansion_search.yml index 5e29d3cdf2ae6..2272e2eb6b301 100644 --- a/x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/ml/text_expansion_search.yml +++ b/x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/ml/text_expansion_search.yml @@ -245,7 +245,7 @@ setup: ml.tokens: tokens: [{"the": 1.0}, {"comforter":1.0}, {"smells":1.0}, {"bad": 1.0}] pruning_config: {} - - match: { hits.total.value: 5 } + - match: { hits.total.value: 7 } - match: { hits.hits.0._source.source_text: "the octopus comforter smells" } --- From df7d7ed91ac75e99f9e33d66bc84176934292062 Mon Sep 17 00:00:00 2001 From: Saikat Sarkar Date: Tue, 20 Feb 2024 19:10:06 -0700 Subject: [PATCH 06/21] Display error for text_expansion if the queried field does not have the right type --- .../xpack/ml/queries/TextExpansionQueryBuilder.java | 8 ++++---- .../xpack/ml/queries/WeightedTokensQueryBuilder.java | 4 ++-- .../rest-api-spec/test/ml/text_expansion_search.yml | 2 +- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilder.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilder.java index cdd2ad8d09e0c..313877155bce0 100644 --- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilder.java +++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilder.java @@ -50,13 +50,13 @@ public class TextExpansionQueryBuilder extends AbstractQueryBuilder weightedTokensSupplier; private final TokenPruningConfig tokenPruningConfig; - public enum AllowedFieldTypesForTextExpansion { + public enum AllowedFieldType { RANK_FEATURES("rank_features"), SPARSE_VECTOR("sparse_vector"); private final String typeName; - AllowedFieldTypesForTextExpansion(String typeName) { + AllowedFieldType(String typeName) { this.typeName = typeName; } @@ -64,8 +64,8 @@ public String getTypeName() { return typeName; } - public static boolean contains(String typeName) { - for (AllowedFieldTypesForTextExpansion fieldType : values()) { + public static boolean isTypeNameValid(String typeName) { + for (AllowedFieldType fieldType : values()) { if (fieldType.getTypeName().equals(typeName)) { return true; } diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java index 2ce2d0d238e83..58028ba718ed7 100644 --- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java +++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java @@ -34,7 +34,7 @@ import java.util.List; import java.util.Objects; -import static org.elasticsearch.xpack.ml.queries.TextExpansionQueryBuilder.AllowedFieldTypesForTextExpansion; +import static org.elasticsearch.xpack.ml.queries.TextExpansionQueryBuilder.AllowedFieldType; import static org.elasticsearch.xpack.ml.queries.TextExpansionQueryBuilder.PRUNING_CONFIG; public class WeightedTokensQueryBuilder extends AbstractQueryBuilder { @@ -155,7 +155,7 @@ protected Query doToQuery(SearchExecutionContext context) throws IOException { } final String fieldTypeName = ft.typeName(); - if (AllowedFieldTypesForTextExpansion.contains(fieldTypeName) == false) { + if (AllowedFieldType.isTypeNameValid(fieldTypeName) == false) { throw new ElasticsearchParseException("[" + fieldTypeName + "]" + " is not an appropriate field type for text expansion query"); } diff --git a/x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/ml/text_expansion_search.yml b/x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/ml/text_expansion_search.yml index 2272e2eb6b301..5e29d3cdf2ae6 100644 --- a/x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/ml/text_expansion_search.yml +++ b/x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/ml/text_expansion_search.yml @@ -245,7 +245,7 @@ setup: ml.tokens: tokens: [{"the": 1.0}, {"comforter":1.0}, {"smells":1.0}, {"bad": 1.0}] pruning_config: {} - - match: { hits.total.value: 7 } + - match: { hits.total.value: 5 } - match: { hits.hits.0._source.source_text: "the octopus comforter smells" } --- From 023a5c7d512e76fbdbe60ba14227daa50e8503b9 Mon Sep 17 00:00:00 2001 From: Saikat Sarkar Date: Tue, 20 Feb 2024 19:11:52 -0700 Subject: [PATCH 07/21] Display error for text_expansion if the queried field does not have the right type --- .java-version | 1 - 1 file changed, 1 deletion(-) delete mode 100644 .java-version diff --git a/.java-version b/.java-version deleted file mode 100644 index 189d2e5f1483a..0000000000000 --- a/.java-version +++ /dev/null @@ -1 +0,0 @@ -openjdk64-17.0.9 From fc6f437adda207e93198862029ad914be1fffb73 Mon Sep 17 00:00:00 2001 From: Saikat Sarkar Date: Wed, 21 Feb 2024 10:35:26 -0700 Subject: [PATCH 08/21] Display error for text_expansion if the queried field does not have the right type --- .../test/ml/text_expansion_search.yml | 34 +++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/ml/text_expansion_search.yml b/x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/ml/text_expansion_search.yml index 5e29d3cdf2ae6..570b0bb06f55d 100644 --- a/x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/ml/text_expansion_search.yml +++ b/x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/ml/text_expansion_search.yml @@ -287,3 +287,37 @@ setup: tokens_weight_threshold: 0.4 only_score_pruned_tokens: true - match: { hits.total.value: 0 } + +--- +"Test text-expansion that displays error for invalid queried field type": + - skip: + version: " - 8.13.99" + reason: "validation for invalid field type introduced in 8.14.0" + + - do: + catch: /\[keyword\] is not an appropriate field type for text expansion query/ + search: + index: index-with-rank-features + body: + query: + text_expansion: + source_text: + model_id: text_expansion_model + model_text: "octopus comforter smells" + +--- +"Test text-expansion that displays error for non-existent queried field": + - skip: + version: " - 8.13.99" + reason: "validation for non-existent queried field introduced in 8.14.0" + + - do: + catch: /\[non_existent_field\] is not a mapped field/ + search: + index: index-with-rank-features + body: + query: + text_expansion: + non_existent_field: + model_id: text_expansion_model + model_text: "octopus comforter smells" From 454f3c6fe045fd22123ed60b052a40e8d1d14be2 Mon Sep 17 00:00:00 2001 From: Saikat Sarkar Date: Wed, 21 Feb 2024 11:17:59 -0700 Subject: [PATCH 09/21] Display error for text_expansion if the queried field does not have the right type --- .../ml/queries/WeightedTokensQueryBuilder.java | 2 +- .../test/ml/text_expansion_search.yml | 17 ----------------- 2 files changed, 1 insertion(+), 18 deletions(-) diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java index 58028ba718ed7..5d063acd313c3 100644 --- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java +++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java @@ -151,7 +151,7 @@ private boolean shouldKeepToken( protected Query doToQuery(SearchExecutionContext context) throws IOException { final MappedFieldType ft = context.getFieldType(fieldName); if (ft == null) { - throw new ElasticsearchParseException("[" + fieldName + "]" + " is not a mapped field"); + return new MatchNoDocsQuery("The \"" + getName() + "\" query is against a field that does not exist"); } final String fieldTypeName = ft.typeName(); diff --git a/x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/ml/text_expansion_search.yml b/x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/ml/text_expansion_search.yml index 570b0bb06f55d..c35059654da8c 100644 --- a/x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/ml/text_expansion_search.yml +++ b/x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/ml/text_expansion_search.yml @@ -304,20 +304,3 @@ setup: source_text: model_id: text_expansion_model model_text: "octopus comforter smells" - ---- -"Test text-expansion that displays error for non-existent queried field": - - skip: - version: " - 8.13.99" - reason: "validation for non-existent queried field introduced in 8.14.0" - - - do: - catch: /\[non_existent_field\] is not a mapped field/ - search: - index: index-with-rank-features - body: - query: - text_expansion: - non_existent_field: - model_id: text_expansion_model - model_text: "octopus comforter smells" From b2f7b0db1eca5a81292781c8d90af2462a48e28d Mon Sep 17 00:00:00 2001 From: Saikat Sarkar Date: Wed, 21 Feb 2024 11:27:30 -0700 Subject: [PATCH 10/21] Display error for text_expansion if the queried field does not have the right type --- .../xpack/ml/queries/TextExpansionQueryBuilder.java | 2 +- .../xpack/ml/queries/WeightedTokensQueryBuilder.java | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilder.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilder.java index 313877155bce0..f3a98e8b643e2 100644 --- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilder.java +++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilder.java @@ -64,7 +64,7 @@ public String getTypeName() { return typeName; } - public static boolean isTypeNameValid(String typeName) { + public static boolean isFieldTypeAllowed(String typeName) { for (AllowedFieldType fieldType : values()) { if (fieldType.getTypeName().equals(typeName)) { return true; diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java index 5d063acd313c3..de5ebb3415752 100644 --- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java +++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java @@ -155,7 +155,7 @@ protected Query doToQuery(SearchExecutionContext context) throws IOException { } final String fieldTypeName = ft.typeName(); - if (AllowedFieldType.isTypeNameValid(fieldTypeName) == false) { + if (AllowedFieldType.isFieldTypeAllowed(fieldTypeName) == false) { throw new ElasticsearchParseException("[" + fieldTypeName + "]" + " is not an appropriate field type for text expansion query"); } From f7b821c6e31bc54d63581964f5f117e4fa0f3935 Mon Sep 17 00:00:00 2001 From: Saikat Sarkar Date: Wed, 21 Feb 2024 12:01:56 -0700 Subject: [PATCH 11/21] Display error for text_expansion if the queried field does not have the right type --- .../xpack/ml/queries/TextExpansionQueryBuilder.java | 12 ++++++++++++ .../xpack/ml/queries/WeightedTokensQueryBuilder.java | 3 ++- 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilder.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilder.java index f3a98e8b643e2..482fbacc5cff5 100644 --- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilder.java +++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilder.java @@ -72,6 +72,18 @@ public static boolean isFieldTypeAllowed(String typeName) { } return false; } + + public static String getAllowedFieldTypesAsString() { + StringBuilder result = new StringBuilder(); + for (AllowedFieldType fieldType : values()) { + result.append(fieldType.getTypeName()).append(", "); + } + // Remove the trailing ", " if there are values + if (result.length() > 0) { + result.setLength(result.length() - 2); + } + return result.toString(); + } } public TextExpansionQueryBuilder(String fieldName, String modelText, String modelId) { diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java index de5ebb3415752..1a96b47f2a8d2 100644 --- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java +++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java @@ -156,7 +156,8 @@ protected Query doToQuery(SearchExecutionContext context) throws IOException { final String fieldTypeName = ft.typeName(); if (AllowedFieldType.isFieldTypeAllowed(fieldTypeName) == false) { - throw new ElasticsearchParseException("[" + fieldTypeName + "]" + " is not an appropriate field type for text expansion query"); + throw new ElasticsearchParseException("[" + fieldTypeName + "]" + " is not an appropriate field type for text expansion query. " + + "Allowed field types are " + AllowedFieldType.getAllowedFieldTypesAsString() + "."); } var qb = new BooleanQuery.Builder(); From d2568533c039ff36a619037d42dd46d2c566bb53 Mon Sep 17 00:00:00 2001 From: Saikat Sarkar Date: Wed, 21 Feb 2024 12:09:48 -0700 Subject: [PATCH 12/21] Display error for text_expansion if the queried field does not have the right type --- .../xpack/ml/queries/TextExpansionQueryBuilder.java | 3 ++- .../xpack/ml/queries/WeightedTokensQueryBuilder.java | 11 +++++++++-- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilder.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilder.java index 482fbacc5cff5..128740a13677b 100644 --- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilder.java +++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilder.java @@ -78,7 +78,8 @@ public static String getAllowedFieldTypesAsString() { for (AllowedFieldType fieldType : values()) { result.append(fieldType.getTypeName()).append(", "); } - // Remove the trailing ", " if there are values + + // Remove the trailing ", " if (result.length() > 0) { result.setLength(result.length() - 2); } diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java index 1a96b47f2a8d2..4422d1a25cde5 100644 --- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java +++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java @@ -156,8 +156,15 @@ protected Query doToQuery(SearchExecutionContext context) throws IOException { final String fieldTypeName = ft.typeName(); if (AllowedFieldType.isFieldTypeAllowed(fieldTypeName) == false) { - throw new ElasticsearchParseException("[" + fieldTypeName + "]" + " is not an appropriate field type for text expansion query. " - + "Allowed field types are " + AllowedFieldType.getAllowedFieldTypesAsString() + "."); + throw new ElasticsearchParseException( + "[" + + fieldTypeName + + "]" + + " is not an appropriate field type for text expansion query. " + + "Allowed field types are " + + AllowedFieldType.getAllowedFieldTypesAsString() + + "." + ); } var qb = new BooleanQuery.Builder(); From 762409df7b611d04e170d252fcb65afc0072423f Mon Sep 17 00:00:00 2001 From: Saikat Sarkar Date: Wed, 21 Feb 2024 13:03:02 -0700 Subject: [PATCH 13/21] Display error for text_expansion if the queried field does not have the right type --- .../xpack/ml/queries/TextExpansionQueryBuilder.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilder.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilder.java index 128740a13677b..f652bd9475be7 100644 --- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilder.java +++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilder.java @@ -78,7 +78,7 @@ public static String getAllowedFieldTypesAsString() { for (AllowedFieldType fieldType : values()) { result.append(fieldType.getTypeName()).append(", "); } - + // Remove the trailing ", " if (result.length() > 0) { result.setLength(result.length() - 2); From f4c2c33d9f25530d9b1d9a883476c89038650c68 Mon Sep 17 00:00:00 2001 From: Saikat Sarkar Date: Wed, 21 Feb 2024 13:56:07 -0700 Subject: [PATCH 14/21] Display error for text_expansion if the queried field does not have the right type --- .../xpack/ml/queries/TextExpansionQueryBuilderTests.java | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilderTests.java b/x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilderTests.java index 13f12f3cdc1e1..6ee78475fb300 100644 --- a/x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilderTests.java +++ b/x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilderTests.java @@ -260,10 +260,6 @@ public void testThatTokensAreCorrectlyPruned() { SearchExecutionContext searchExecutionContext = createSearchExecutionContext(); TextExpansionQueryBuilder queryBuilder = createTestQueryBuilder(); QueryBuilder rewrittenQueryBuilder = rewriteAndFetch(queryBuilder, searchExecutionContext); - if (queryBuilder.getTokenPruningConfig() == null) { - assertTrue(rewrittenQueryBuilder instanceof BoolQueryBuilder); - } else { - assertTrue(rewrittenQueryBuilder instanceof WeightedTokensQueryBuilder); - } + assertTrue(rewrittenQueryBuilder instanceof WeightedTokensQueryBuilder); } } From 91f0a733b1058472c361fa84811f5e9cb3dfd5bd Mon Sep 17 00:00:00 2001 From: Saikat Sarkar Date: Wed, 21 Feb 2024 14:10:17 -0700 Subject: [PATCH 15/21] Display error for text_expansion if the queried field does not have the right type --- .../xpack/ml/queries/TextExpansionQueryBuilderTests.java | 1 - 1 file changed, 1 deletion(-) diff --git a/x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilderTests.java b/x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilderTests.java index 6ee78475fb300..50561d92f5d37 100644 --- a/x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilderTests.java +++ b/x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilderTests.java @@ -25,7 +25,6 @@ import org.elasticsearch.common.compress.CompressedXContent; import org.elasticsearch.index.mapper.MapperService; import org.elasticsearch.index.mapper.extras.MapperExtrasPlugin; -import org.elasticsearch.index.query.BoolQueryBuilder; import org.elasticsearch.index.query.QueryBuilder; import org.elasticsearch.index.query.SearchExecutionContext; import org.elasticsearch.plugins.Plugin; From 6113c8cda4fc264e192f61517c00ce55bf1977fc Mon Sep 17 00:00:00 2001 From: Saikat Sarkar Date: Thu, 22 Feb 2024 11:27:25 -0700 Subject: [PATCH 16/21] Clean up the code --- .../xpack/ml/queries/TextExpansionQueryBuilder.java | 13 +++---------- .../ml/queries/WeightedTokensQueryBuilder.java | 4 ++-- 2 files changed, 5 insertions(+), 12 deletions(-) diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilder.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilder.java index f652bd9475be7..675d062fdb3af 100644 --- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilder.java +++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilder.java @@ -31,8 +31,10 @@ import org.elasticsearch.xpack.core.ml.inference.trainedmodel.TextExpansionConfigUpdate; import java.io.IOException; +import java.util.Arrays; import java.util.List; import java.util.Objects; +import java.util.stream.Collectors; import static org.elasticsearch.xpack.core.ClientHelper.ML_ORIGIN; import static org.elasticsearch.xpack.core.ClientHelper.executeAsyncWithOrigin; @@ -74,16 +76,7 @@ public static boolean isFieldTypeAllowed(String typeName) { } public static String getAllowedFieldTypesAsString() { - StringBuilder result = new StringBuilder(); - for (AllowedFieldType fieldType : values()) { - result.append(fieldType.getTypeName()).append(", "); - } - - // Remove the trailing ", " - if (result.length() > 0) { - result.setLength(result.length() - 2); - } - return result.toString(); + return Arrays.stream(values()).map(value -> value.typeName).collect(Collectors.joining(", ")); } } diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java index 4422d1a25cde5..12ad14aca40a2 100644 --- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java +++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java @@ -161,9 +161,9 @@ protected Query doToQuery(SearchExecutionContext context) throws IOException { + fieldTypeName + "]" + " is not an appropriate field type for text expansion query. " - + "Allowed field types are " + + "Allowed field types are [" + AllowedFieldType.getAllowedFieldTypesAsString() - + "." + + "]." ); } From d536e07d9f4c8629410de14f8f6198c2126f7fb5 Mon Sep 17 00:00:00 2001 From: Saikat Sarkar Date: Thu, 22 Feb 2024 14:32:55 -0700 Subject: [PATCH 17/21] Optimize the code for pruning config --- .../queries/WeightedTokensQueryBuilder.java | 34 ++++++++++++++----- 1 file changed, 25 insertions(+), 9 deletions(-) diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java index 12ad14aca40a2..a240306129b8b 100644 --- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java +++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java @@ -167,27 +167,43 @@ protected Query doToQuery(SearchExecutionContext context) throws IOException { ); } + return (this.tokenPruningConfig == null) ? queryBuilderWithAllTokens(tokens, ft, context) : queryBuilderWithPrunedTokens(tokens, ft, context); + } + + private Query queryBuilderWithAllTokens(List tokens, MappedFieldType ft, SearchExecutionContext context) { var qb = new BooleanQuery.Builder(); - int fieldDocCount = context.getIndexReader().getDocCount(fieldName); - float bestWeight = 0f; - for (var t : tokens) { - bestWeight = Math.max(t.weight(), bestWeight); + + for (var token : tokens) { + qb.add(new BoostQuery(ft.termQuery(token.token(), context), token.weight()), BooleanClause.Occur.SHOULD); } + return qb.setMinimumNumberShouldMatch(1).build(); + } + private Query queryBuilderWithPrunedTokens(List tokens, MappedFieldType ft, SearchExecutionContext context) throws IOException { + var qb = new BooleanQuery.Builder(); + int fieldDocCount = context.getIndexReader().getDocCount(fieldName); + float bestWeight = findBestWeightFor(tokens); float averageTokenFreqRatio = getAverageTokenFreqRatio(context.getIndexReader(), fieldDocCount); if (averageTokenFreqRatio == 0) { return new MatchNoDocsQuery("The \"" + getName() + "\" query is against an empty field"); } + for (var token : tokens) { boolean keep = shouldKeepToken(context.getIndexReader(), token, fieldDocCount, averageTokenFreqRatio, bestWeight); - if (this.tokenPruningConfig != null) { - keep ^= this.tokenPruningConfig.isOnlyScorePrunedTokens(); - } + keep ^= this.tokenPruningConfig.isOnlyScorePrunedTokens(); if (keep) { qb.add(new BoostQuery(ft.termQuery(token.token(), context), token.weight()), BooleanClause.Occur.SHOULD); } } - qb.setMinimumNumberShouldMatch(1); - return qb.build(); + + return qb.setMinimumNumberShouldMatch(1).build(); + } + + private float findBestWeightFor(List tokens) { + float bestWeight = 0f; + for (var t : tokens) { + bestWeight = Math.max(t.weight(), bestWeight); + } + return bestWeight; } @Override From 3fb95e727d68fa931700d016c7290ae8a37017f0 Mon Sep 17 00:00:00 2001 From: Saikat Sarkar Date: Thu, 22 Feb 2024 14:35:55 -0700 Subject: [PATCH 18/21] Optimize the code for pruning config --- .../xpack/ml/queries/WeightedTokensQueryBuilder.java | 1 + 1 file changed, 1 insertion(+) diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java index a240306129b8b..bf110a3bab8a7 100644 --- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java +++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java @@ -178,6 +178,7 @@ private Query queryBuilderWithAllTokens(List tokens, MappedFieldT } return qb.setMinimumNumberShouldMatch(1).build(); } + private Query queryBuilderWithPrunedTokens(List tokens, MappedFieldType ft, SearchExecutionContext context) throws IOException { var qb = new BooleanQuery.Builder(); int fieldDocCount = context.getIndexReader().getDocCount(fieldName); From cb071b6ea2813f9853f399cbb92c70e54c05452e Mon Sep 17 00:00:00 2001 From: Saikat Sarkar Date: Thu, 22 Feb 2024 14:45:14 -0700 Subject: [PATCH 19/21] Write findBestWeightFor for clear code --- .../xpack/ml/queries/WeightedTokensQueryBuilder.java | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java index bf110a3bab8a7..4820d44596456 100644 --- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java +++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java @@ -200,11 +200,9 @@ private Query queryBuilderWithPrunedTokens(List tokens, MappedFie } private float findBestWeightFor(List tokens) { - float bestWeight = 0f; - for (var t : tokens) { - bestWeight = Math.max(t.weight(), bestWeight); - } - return bestWeight; + return tokens.stream() + .map(WeightedToken::weight) + .reduce(0f, Math::max); } @Override From ee3233bc4284c5b736ae4ebe89081dd082ac2131 Mon Sep 17 00:00:00 2001 From: Saikat Sarkar Date: Thu, 22 Feb 2024 14:48:16 -0700 Subject: [PATCH 20/21] Run Spotless --- .../xpack/ml/queries/WeightedTokensQueryBuilder.java | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java index 4820d44596456..fc3f8fe56d152 100644 --- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java +++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java @@ -167,7 +167,9 @@ protected Query doToQuery(SearchExecutionContext context) throws IOException { ); } - return (this.tokenPruningConfig == null) ? queryBuilderWithAllTokens(tokens, ft, context) : queryBuilderWithPrunedTokens(tokens, ft, context); + return (this.tokenPruningConfig == null) + ? queryBuilderWithAllTokens(tokens, ft, context) + : queryBuilderWithPrunedTokens(tokens, ft, context); } private Query queryBuilderWithAllTokens(List tokens, MappedFieldType ft, SearchExecutionContext context) { @@ -179,7 +181,8 @@ private Query queryBuilderWithAllTokens(List tokens, MappedFieldT return qb.setMinimumNumberShouldMatch(1).build(); } - private Query queryBuilderWithPrunedTokens(List tokens, MappedFieldType ft, SearchExecutionContext context) throws IOException { + private Query queryBuilderWithPrunedTokens(List tokens, MappedFieldType ft, SearchExecutionContext context) + throws IOException { var qb = new BooleanQuery.Builder(); int fieldDocCount = context.getIndexReader().getDocCount(fieldName); float bestWeight = findBestWeightFor(tokens); @@ -200,9 +203,7 @@ private Query queryBuilderWithPrunedTokens(List tokens, MappedFie } private float findBestWeightFor(List tokens) { - return tokens.stream() - .map(WeightedToken::weight) - .reduce(0f, Math::max); + return tokens.stream().map(WeightedToken::weight).reduce(0f, Math::max); } @Override From 8969c3181c1997045dfa9b1589feedfb9479020d Mon Sep 17 00:00:00 2001 From: Saikat Sarkar Date: Fri, 23 Feb 2024 09:11:31 -0700 Subject: [PATCH 21/21] Remove findBestWeightFor method --- .../xpack/ml/queries/WeightedTokensQueryBuilder.java | 8 ++------ .../rest-api-spec/test/ml/text_expansion_search.yml | 2 +- 2 files changed, 3 insertions(+), 7 deletions(-) diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java index fc3f8fe56d152..51139881fc2e4 100644 --- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java +++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/WeightedTokensQueryBuilder.java @@ -160,7 +160,7 @@ protected Query doToQuery(SearchExecutionContext context) throws IOException { "[" + fieldTypeName + "]" - + " is not an appropriate field type for text expansion query. " + + " is not an appropriate field type for this query. " + "Allowed field types are [" + AllowedFieldType.getAllowedFieldTypesAsString() + "]." @@ -185,7 +185,7 @@ private Query queryBuilderWithPrunedTokens(List tokens, MappedFie throws IOException { var qb = new BooleanQuery.Builder(); int fieldDocCount = context.getIndexReader().getDocCount(fieldName); - float bestWeight = findBestWeightFor(tokens); + float bestWeight = tokens.stream().map(WeightedToken::weight).reduce(0f, Math::max); float averageTokenFreqRatio = getAverageTokenFreqRatio(context.getIndexReader(), fieldDocCount); if (averageTokenFreqRatio == 0) { return new MatchNoDocsQuery("The \"" + getName() + "\" query is against an empty field"); @@ -202,10 +202,6 @@ private Query queryBuilderWithPrunedTokens(List tokens, MappedFie return qb.setMinimumNumberShouldMatch(1).build(); } - private float findBestWeightFor(List tokens) { - return tokens.stream().map(WeightedToken::weight).reduce(0f, Math::max); - } - @Override protected boolean doEquals(WeightedTokensQueryBuilder other) { return Objects.equals(fieldName, other.fieldName) diff --git a/x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/ml/text_expansion_search.yml b/x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/ml/text_expansion_search.yml index c35059654da8c..dc4e1751ccdee 100644 --- a/x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/ml/text_expansion_search.yml +++ b/x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/ml/text_expansion_search.yml @@ -295,7 +295,7 @@ setup: reason: "validation for invalid field type introduced in 8.14.0" - do: - catch: /\[keyword\] is not an appropriate field type for text expansion query/ + catch: /\[keyword\] is not an appropriate field type for this query/ search: index: index-with-rank-features body: