Merge branch 'main' into dense_vector_stats

Author: ChrisHegarty

test/framework/src/main/java/org/elasticsearch/node/MockNode.java

@@ -9,6 +9,8 @@
 
 package org.elasticsearch.node;
 
+import org.elasticsearch.action.search.OnlinePrewarmingService;
+import org.elasticsearch.action.search.OnlinePrewarmingServiceProvider;
 import org.elasticsearch.client.internal.node.NodeClient;
 import org.elasticsearch.cluster.ClusterInfoService;
 import org.elasticsearch.cluster.MockInternalClusterInfoService;
@@ -117,6 +119,10 @@ SearchService newSearchService(
                     tracer
                 );
             }
+            OnlinePrewarmingService onlinePrewarmingService = pluginsService.loadSingletonServiceProvider(
+                OnlinePrewarmingServiceProvider.class,
+                () -> OnlinePrewarmingServiceProvider.DEFAULT
+            ).create(clusterService.getSettings(), threadPool, clusterService);
             return new MockSearchService(
                 clusterService,
                 indicesService,
@@ -126,7 +132,8 @@ SearchService newSearchService(
                 fetchPhase,
                 circuitBreakerService,
                 executorSelector,
-                tracer
+                tracer,
+                onlinePrewarmingService
             );
         }
 

test/framework/src/main/java/org/elasticsearch/search/MockSearchService.java

@@ -85,7 +85,8 @@ public MockSearchService(
         FetchPhase fetchPhase,
         CircuitBreakerService circuitBreakerService,
         ExecutorSelector executorSelector,
-        Tracer tracer
+        Tracer tracer,
+        OnlinePrewarmingService onlinePrewarmingService
     ) {
         super(
             clusterService,
@@ -97,7 +98,7 @@ public MockSearchService(
             circuitBreakerService,
             executorSelector,
             tracer,
-            OnlinePrewarmingService.NOOP
+            onlinePrewarmingService
         );
     }
 

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/MergeOperator.java

@@ -68,8 +68,11 @@ protected Page process(Page page) {
         for (int i = 0; i < page.getBlockCount() - 1; i++) {
             projections[i] = i == scorePosition ? page.getBlockCount() - 1 : i;
         }
-
-        return page.projectBlocks(projections);
+        try {
+            return page.projectBlocks(projections);
+        } finally {
+            page.releaseBlocks();
+        }
     }
 
     @Override

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/RrfScoreEvalOperator.java

@@ -3,7 +3,7 @@
 //
 
 simpleFork
-required_capability: fork
+required_capability: fork_v3
 
 FROM employees
 | FORK ( WHERE emp_no == 10001 )
@@ -18,7 +18,7 @@ emp_no:integer | _fork:keyword
 ;
 
 forkWithWhereSortAndLimit
-required_capability: fork
+required_capability: fork_v3
 
 FROM employees
 | FORK ( WHERE hire_date < "1985-03-01T00:00:00Z" | SORT first_name | LIMIT 5 )
@@ -38,7 +38,7 @@ emp_no:integer | first_name:keyword | _fork:keyword
 ;
 
 fiveFork
-required_capability: fork
+required_capability: fork_v3
 
 FROM employees
 | FORK ( WHERE emp_no == 10005 )
@@ -59,7 +59,7 @@ fork5         | 10001
 ;
 
 forkWithWhereSortDescAndLimit
-required_capability: fork
+required_capability: fork_v3
 
 FROM employees
 | FORK ( WHERE hire_date < "1985-03-01T00:00:00Z" | SORT first_name DESC | LIMIT 2 )
@@ -76,7 +76,7 @@ fork2         | 10087          | Xinglin
 ;
 
 forkWithCommonPrefilter
-required_capability: fork
+required_capability: fork_v3
 
 FROM employees
 | WHERE emp_no > 10050
@@ -94,7 +94,7 @@ fork2         | 10100
 ;
 
 forkWithSemanticSearchAndScore
-required_capability: fork
+required_capability: fork_v3
 required_capability: semantic_text_field_caps
 required_capability: metadata_score
 
@@ -114,7 +114,7 @@ fork2         | 6.093784261960139E18  | 2           | all we have to decide is w
 ;
 
 forkWithEvals
-required_capability: fork_v2
+required_capability: fork_v3
 
 FROM employees
 | FORK (WHERE emp_no == 10048 OR emp_no == 10081 | EVAL x = "abc" | EVAL y = 1)
@@ -131,7 +131,7 @@ fork2         | 10087          | def       | null      | 2
 ;
 
 forkWithStats
-required_capability: fork_v2
+required_capability: fork_v3
 
 FROM employees
 | FORK (WHERE emp_no == 10048 OR emp_no == 10081)
@@ -152,7 +152,7 @@ fork4         | null           | 100    | 10001     | null
 ;
 
 forkWithDissect
-required_capability: fork_v2
+required_capability: fork_v3
 
 FROM employees
 | WHERE emp_no == 10048 OR emp_no == 10081
@@ -172,7 +172,7 @@ fork2         | 10081          | Rosen     | 10081     | null      | Zhongwei
 ;
 
 forkWithMixOfCommands
-required_capability: fork_v2
+required_capability: fork_v3
 
 FROM employees
 | WHERE emp_no == 10048 OR emp_no == 10081

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

@@ -76,6 +76,7 @@ FROM books METADATA _id, _index, _score
        ( WHERE author:"Ursula K. Le Guin" AND title:"short stories" | SORT _score, _id DESC | LIMIT 3)
 | RRF
 | STATS count_fork=COUNT(*) BY _fork
+| SORT _fork
 ;
 
 count_fork:long | _fork:keyword
@@ -120,6 +121,7 @@ FROM semantic_text METADATA _id, _score, _index
        ( WHERE semantic_text_field:"something else" | SORT _score DESC | LIMIT 2)
 | RRF
 | EVAL _score = round(_score, 4)
+| EVAL _fork = mv_sort(_fork)
 | KEEP _fork, _score, _id, semantic_text_field
 ;
 

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

@@ -72,6 +72,29 @@ private void testSimpleImpl(String query) {
         }
     }
 
+    public void testRow() {
+        var query = """
+            ROW a = [1, 2, 3, 4], b = 100
+            | MV_EXPAND a
+            | FORK (WHERE a % 2 == 1)
+                   (WHERE a % 2 == 0)
+            | SORT _fork, a
+            """;
+
+        try (var resp = run(query)) {
+            assertColumnNames(resp.columns(), List.of("a", "b", "_fork"));
+            assertColumnTypes(resp.columns(), List.of("integer", "integer", "keyword"));
+
+            Iterable<Iterable<Object>> expectedValues = List.of(
+                List.of(1, 100, "fork1"),
+                List.of(3, 100, "fork1"),
+                List.of(2, 100, "fork2"),
+                List.of(4, 100, "fork2")
+            );
+            assertValues(resp.values(), expectedValues);
+        }
+    }
+
     public void testSortAndLimitInFirstSubQuery() {
         var query = """
             FROM test
@@ -216,13 +239,15 @@ public void testWhereSortOnlyInFork() {
                ( WHERE content:"fox" | SORT id )
                ( WHERE content:"dog" | SORT id )
             | KEEP _fork, id, content
+            | SORT _fork, id
             """;
         var queryWithMatchFunction = """
             FROM test
             | FORK
                ( WHERE match(content, "fox") | SORT id )
                ( WHERE match(content, "dog") | SORT id )
             | KEEP _fork, id, content
+            | SORT _fork, id
             """;
         for (var query : List.of(queryWithMatchOperator, queryWithMatchFunction)) {
             try (var resp = run(query)) {
@@ -509,6 +534,7 @@ public void testWithEvalSimple() {
                 | FORK ( EVAL a = 1 )
                        ( EVAL a = 2 )
                 | KEEP a, _fork, id, content
+                | SORT _fork
             """;
 
         try (var resp = run(query)) {

x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/ForkIT.java

@@ -40,6 +40,7 @@ public void testRrf() {
                ( WHERE content:"fox" | SORT _score, _id DESC )
                ( WHERE content:"dog" | SORT _score, _id DESC )
             | RRF
+            | EVAL _fork = mv_sort(_fork)
             | EVAL _score = round(_score, 4)
             | KEEP id, content, _score, _fork
             """;

x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/RrfIT.java

@@ -983,9 +983,9 @@ public enum Cap {
         MAX_OVER_TIME(Build.current().isSnapshot()),
 
         /**
-         * Support STATS/EVAL/DISSECT in Fork branches
+         * Support streaming of sub plan results
          */
-        FORK_V2(Build.current().isSnapshot()),
+        FORK_V3(Build.current().isSnapshot()),
 
         /**
          * Support for the {@code leading_zeros} named parameter.

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java

@@ -19,6 +19,7 @@
 import org.elasticsearch.xpack.esql.plan.logical.Project;
 import org.elasticsearch.xpack.esql.plan.physical.ExchangeExec;
 import org.elasticsearch.xpack.esql.plan.physical.FragmentExec;
+import org.elasticsearch.xpack.esql.plan.physical.MergeExec;
 import org.elasticsearch.xpack.esql.plan.physical.PhysicalPlan;
 import org.elasticsearch.xpack.esql.rule.Rule;
 
@@ -45,6 +46,10 @@ public PhysicalPlan apply(PhysicalPlan plan) {
 
         // This will require updating should we choose to have non-unary execution plans in the future.
         return plan.transformDown(currentPlanNode -> {
+            if (currentPlanNode instanceof MergeExec) {
+                keepTraversing.set(FALSE);
+            }
+
             if (keepTraversing.get() == false) {
                 return currentPlanNode;
             }