Enable semantic search in FORK (#125960)

Author: ioanatia

x-pack/plugin/esql/qa/testFixtures/src/main/resources/fork.csv-spec

@@ -92,3 +92,23 @@ fork1         | 10052
 fork2         | 10099
 fork2         | 10100
 ;
+
+forkWithSemanticSearchAndScore
+required_capability: fork
+required_capability: semantic_text_field_caps
+required_capability: metadata_score
+
+FROM semantic_text METADATA _id, _score
+| FORK ( WHERE semantic_text_field:"something" | SORT _score DESC | LIMIT 2)
+       ( WHERE semantic_text_field:"something else" | SORT _score DESC | LIMIT 2)
+| EVAL _score = round(_score, 4)
+| SORT _fork, _score, _id
+| KEEP _fork, _score, _id, semantic_text_field
+;
+
+_fork:keyword | _score:double         | _id:keyword | semantic_text_field:text
+fork1         | 2.156063961865257E18  | 3           | be excellent to each other
+fork1         | 5.603396578413904E18  | 2           | all we have to decide is what to do with the time that is given to us
+fork2         | 2.3447541759648727E18 | 3           | be excellent to each other
+fork2         | 6.093784261960139E18  | 2           | all we have to decide is what to do with the time that is given to us
+;

x-pack/plugin/esql/qa/testFixtures/src/main/resources/rrf.csv-spec

@@ -109,3 +109,21 @@ _score:double | author:keyword                    | title:keyword        | _fork
 0.0161        | Ursula K. Le Guin                 | The Word For World i | fork2
 0.0159        | Ursula K. Le Guin                 | The Dispossessed     | fork2
 ;
+
+rrfWithSemanticSearch
+required_capability: rrf
+required_capability: semantic_text_field_caps
+required_capability: metadata_score
+
+FROM semantic_text METADATA _id, _score, _index
+| FORK ( WHERE semantic_text_field:"something" | SORT _score DESC | LIMIT 2)
+       ( WHERE semantic_text_field:"something else" | SORT _score DESC | LIMIT 2)
+| RRF
+| EVAL _score = round(_score, 4)
+| KEEP _fork, _score, _id, semantic_text_field
+;
+
+_fork:keyword  | _score:double | _id:keyword | semantic_text_field:keyword
+[fork1, fork2] | 0.0328        | 2           | all we have to decide is what to do with the time that is given to us
+[fork1, fork2] | 0.0323        | 3           | be excellent to each other
+;

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/fulltext/QueryBuilderResolver.java

@@ -12,8 +12,10 @@
 import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.index.query.QueryRewriteContext;
 import org.elasticsearch.index.query.Rewriteable;
+import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.util.Holder;
 import org.elasticsearch.xpack.esql.plan.logical.EsRelation;
+import org.elasticsearch.xpack.esql.plan.logical.Fork;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
 import org.elasticsearch.xpack.esql.planner.TranslatorHandler;
 import org.elasticsearch.xpack.esql.plugin.TransportActionServices;
@@ -22,6 +24,8 @@
 import java.io.IOException;
 import java.util.HashSet;
 import java.util.Set;
+import java.util.function.Function;
+import java.util.stream.Collectors;
 
 /**
  * Some {@link FullTextFunction} implementations such as {@link org.elasticsearch.xpack.esql.expression.function.fulltext.Match}
@@ -34,11 +38,7 @@ public final class QueryBuilderResolver {
     private QueryBuilderResolver() {}
 
     public static void resolveQueryBuilders(LogicalPlan plan, TransportActionServices services, ActionListener<LogicalPlan> listener) {
-        var hasFullTextFunctions = plan.anyMatch(p -> {
-            Holder<Boolean> hasFullTextFunction = new Holder<>(false);
-            p.forEachExpression(FullTextFunction.class, unused -> hasFullTextFunction.set(true));
-            return hasFullTextFunction.get();
-        });
+        var hasFullTextFunctions = hasFullTextFunctions(plan);
         if (hasFullTextFunctions) {
             Rewriteable.rewriteAndFetch(
                 new FullTextFunctionsRewritable(plan),
@@ -69,12 +69,29 @@ private static Set<String> indexNames(LogicalPlan plan) {
         return indexNames;
     }
 
+    private static boolean hasFullTextFunctions(LogicalPlan plan) {
+        return plan.anyMatch(p -> {
+            Holder<Boolean> hasFullTextFunction = new Holder<>(false);
+            p.forEachExpression(FullTextFunction.class, unused -> hasFullTextFunction.set(true));
+
+            if (p instanceof Fork fork) {
+                fork.subPlans().forEach(subPlan -> {
+                    if (hasFullTextFunctions(subPlan)) {
+                        hasFullTextFunction.set(true);
+                    }
+                });
+            }
+
+            return hasFullTextFunction.get();
+        });
+    }
+
     private record FullTextFunctionsRewritable(LogicalPlan plan) implements Rewriteable<QueryBuilderResolver.FullTextFunctionsRewritable> {
         @Override
         public FullTextFunctionsRewritable rewrite(QueryRewriteContext ctx) throws IOException {
             Holder<IOException> exceptionHolder = new Holder<>();
             Holder<Boolean> updated = new Holder<>(false);
-            LogicalPlan newPlan = plan.transformExpressionsDown(FullTextFunction.class, f -> {
+            LogicalPlan newPlan = transformPlan(plan, f -> {
                 QueryBuilder builder = f.queryBuilder(), initial = builder;
                 builder = builder == null ? f.asQuery(TranslatorHandler.TRANSLATOR_HANDLER).toQueryBuilder() : builder;
                 try {
@@ -91,5 +108,15 @@ public FullTextFunctionsRewritable rewrite(QueryRewriteContext ctx) throws IOExc
             }
             return updated.get() ? new FullTextFunctionsRewritable(newPlan) : this;
         }
+
+        private LogicalPlan transformPlan(LogicalPlan plan, Function<FullTextFunction, ? extends Expression> rule) {
+            return plan.transformExpressionsDown(FullTextFunction.class, rule).transformDown(Fork.class, fork -> {
+                var subPlans = fork.subPlans()
+                    .stream()
+                    .map(subPlan -> subPlan.transformExpressionsDown(FullTextFunction.class, rule))
+                    .collect(Collectors.toList());
+                return new Fork(fork.source(), fork.child(), subPlans);
+            });
+        }
     }
 }