using weighted tokens as criteria for rewriting text expansion query

Author: henriquepaes1

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/results/TextExpansionResults.java

@@ -96,4 +96,8 @@ public Map<String, Object> asMap(String outputField) {
         map.put(outputField, weightedTokens.stream().collect(Collectors.toMap(WeightedToken::token, WeightedToken::weight)));
         return map;
     }
+
+    public boolean hasWeightedTokens() {
+        return this.getWeightedTokens().isEmpty() == false;
+    }
 }

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilder.java

@@ -38,7 +38,6 @@
 import java.io.IOException;
 import java.util.List;
 import java.util.Objects;
-import java.util.Set;
 
 import static org.elasticsearch.xpack.core.ClientHelper.ML_ORIGIN;
 import static org.elasticsearch.xpack.core.ClientHelper.executeAsyncWithOrigin;
@@ -60,8 +59,6 @@ public class TextExpansionQueryBuilder extends AbstractQueryBuilder<TextExpansio
     private SetOnce<TextExpansionResults> weightedTokensSupplier;
     private final TokenPruningConfig tokenPruningConfig;
 
-    private static final Set<String> ALLOWED_FIELD_TYPES = Set.of("sparse_vector", "rank_features");
-
     private static final DeprecationLogger deprecationLogger = DeprecationLogger.getLogger(ParseField.class);
     public static final String TEXT_EXPANSION_DEPRECATION_MESSAGE = NAME + " is deprecated. Use sparse_vector instead.";
 
@@ -161,8 +158,6 @@ protected QueryBuilder doRewrite(QueryRewriteContext queryRewriteContext) {
             return weightedTokensToQuery(fieldName, weightedTokensSupplier.get());
         }
 
-        // Do field type check if query won't be rewritten as a WeightedTokensQuery
-
         CoordinatedInferenceAction.Request inferRequest = CoordinatedInferenceAction.Request.forTextInput(
             modelId,
             List.of(modelText),
@@ -213,7 +208,7 @@ protected QueryBuilder doRewrite(QueryRewriteContext queryRewriteContext) {
     }
 
     private QueryBuilder weightedTokensToQuery(String fieldName, TextExpansionResults textExpansionResults) {
-        if (tokenPruningConfig != null) {
+        if (tokenPruningConfig != null || textExpansionResults.hasWeightedTokens()) {
             WeightedTokensQueryBuilder weightedTokensQueryBuilder = new WeightedTokensQueryBuilder(
                 fieldName,
                 textExpansionResults.getWeightedTokens(),