Stream result pages from sub plans for FORK

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

@@ -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/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/ForkIT.java

@@ -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/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java

@@ -978,9 +978,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()),
 
         /**
          * Does the usage information for ESQL contain a histogram of {@code took} values?

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/ProjectAwayColumns.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;
             }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/PlannerUtils.java

@@ -36,13 +36,15 @@
 import org.elasticsearch.xpack.esql.plan.QueryPlan;
 import org.elasticsearch.xpack.esql.plan.logical.EsRelation;
 import org.elasticsearch.xpack.esql.plan.logical.Filter;
+import org.elasticsearch.xpack.esql.plan.logical.Project;
 import org.elasticsearch.xpack.esql.plan.physical.AggregateExec;
 import org.elasticsearch.xpack.esql.plan.physical.EsSourceExec;
 import org.elasticsearch.xpack.esql.plan.physical.EstimatesRowSize;
 import org.elasticsearch.xpack.esql.plan.physical.ExchangeExec;
 import org.elasticsearch.xpack.esql.plan.physical.ExchangeSinkExec;
 import org.elasticsearch.xpack.esql.plan.physical.ExchangeSourceExec;
 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.planner.mapper.LocalMapper;
 import org.elasticsearch.xpack.esql.planner.mapper.Mapper;
@@ -67,6 +69,25 @@
 
 public class PlannerUtils {
 
+    public static Tuple<List<PhysicalPlan>, PhysicalPlan> breakPlanIntoSubPlansAndMainPlan(PhysicalPlan plan) {
+        var subplans = new Holder<List<PhysicalPlan>>();
+        PhysicalPlan mainPlan = plan.transformUp(MergeExec.class, me -> {
+            subplans.set(me.children().stream().map(child -> {
+                // TODO: we are adding a Project plan to force InsertFieldExtraction - we should remove this transformation
+                child = child.transformUp(FragmentExec.class, f -> {
+                    var logicalFragment = f.fragment();
+                    logicalFragment = new Project(logicalFragment.source(), logicalFragment, logicalFragment.output());
+                    return new FragmentExec(logicalFragment);
+                });
+
+                return (PhysicalPlan) new ExchangeSinkExec(child.source(), child.output(), false, child);
+            }).toList());
+            return new ExchangeSourceExec(me.source(), me.output(), false);
+        });
+
+        return new Tuple<>(subplans.get(), mainPlan);
+    }
+
     public static Tuple<PhysicalPlan, PhysicalPlan> breakPlanBetweenCoordinatorAndDataNode(PhysicalPlan plan, Configuration config) {
         var dataNodePlan = new Holder<PhysicalPlan>();
 

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/mapper/Mapper.java

@@ -41,7 +41,6 @@
 import org.elasticsearch.xpack.esql.plan.physical.UnaryExec;
 import org.elasticsearch.xpack.esql.plan.physical.inference.RerankExec;
 
-import java.util.ArrayList;
 import java.util.List;
 
 /**
@@ -235,12 +234,7 @@ private PhysicalPlan mapBinary(BinaryPlan bp) {
     }
 
     private PhysicalPlan mapFork(Fork fork) {
-        List<PhysicalPlan> physicalChildren = new ArrayList<>();
-        for (var child : fork.children()) {
-            var mappedChild = new FragmentExec(child);
-            physicalChildren.add(mappedChild);
-        }
-        return new MergeExec(fork.source(), physicalChildren, fork.output());
+        return new MergeExec(fork.source(), fork.children().stream().map(child -> map(child)).toList(), fork.output());
     }
 
     public static boolean isPipelineBreaker(LogicalPlan p) {

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

@@ -23,6 +23,8 @@
 import org.elasticsearch.compute.operator.DriverTaskRunner;
 import org.elasticsearch.compute.operator.FailureCollector;
 import org.elasticsearch.compute.operator.exchange.ExchangeService;
+import org.elasticsearch.compute.operator.exchange.ExchangeSink;
+import org.elasticsearch.compute.operator.exchange.ExchangeSinkHandler;
 import org.elasticsearch.compute.operator.exchange.ExchangeSourceHandler;
 import org.elasticsearch.core.Releasable;
 import org.elasticsearch.core.Releasables;
@@ -175,6 +177,89 @@ public void execute(
         FoldContext foldContext,
         EsqlExecutionInfo execInfo,
         ActionListener<Result> listener
+    ) {
+        Tuple<List<PhysicalPlan>, PhysicalPlan> subplansAndMainPlan = PlannerUtils.breakPlanIntoSubPlansAndMainPlan(physicalPlan);
+
+        List<PhysicalPlan> subplans = subplansAndMainPlan.v1();
+
+        // we have no sub plans, so we can just execute the given plan
+        if (subplans == null || subplans.size() == 0) {
+            executePlan(sessionId, rootTask, physicalPlan, configuration, foldContext, execInfo, listener, null);
+            return;
+        }
+
+        final List<Page> collectedPages = Collections.synchronizedList(new ArrayList<>());
+        PhysicalPlan mainPlan = new OutputExec(subplansAndMainPlan.v2(), collectedPages::add);
+
+        listener = listener.delegateResponse((l, e) -> {
+            collectedPages.forEach(p -> Releasables.closeExpectNoException(p::releaseBlocks));
+            l.onFailure(e);
+        });
+
+        var mainSessionId = newChildSession(sessionId);
+        QueryPragmas queryPragmas = configuration.pragmas();
+
+        ExchangeSourceHandler mainExchangeSource = new ExchangeSourceHandler(
+            queryPragmas.exchangeBufferSize(),
+            transportService.getThreadPool().executor(ThreadPool.Names.SEARCH)
+        );
+
+        exchangeService.addExchangeSourceHandler(mainSessionId, mainExchangeSource);
+        var finalListener = ActionListener.runBefore(listener, () -> exchangeService.removeExchangeSourceHandler(sessionId));
+
+        for (PhysicalPlan subplan : subplans) {
+            var childSessionId = newChildSession(sessionId);
+            ExchangeSinkHandler exchangeSink = exchangeService.createSinkHandler(childSessionId, queryPragmas.exchangeBufferSize());
+            // funnel sub plan pages into the main plan exchange source
+            mainExchangeSource.addRemoteSink(exchangeSink::fetchPageAsync, true, () -> {}, 1, ActionListener.noop());
+            executePlan(childSessionId, rootTask, subplan, configuration, foldContext, execInfo, ActionListener.wrap(result -> {
+                exchangeSink.addCompletionListener(
+                    ActionListener.running(() -> { exchangeService.finishSinkHandler(childSessionId, null); })
+                );
+            }, e -> {
+                exchangeService.finishSinkHandler(childSessionId, e);
+                finalListener.onFailure(e);
+            }), () -> exchangeSink.createExchangeSink(() -> {}));
+        }
+
+        var computeContext = new ComputeContext(
+            newChildSession(sessionId),
+            "single",
+            LOCAL_CLUSTER,
+            List.of(),
+            configuration,
+            foldContext,
+            mainExchangeSource::createExchangeSource,
+            null
+        );
+
+        Runnable cancelQueryOnFailure = cancelQueryOnFailure(rootTask);
+
+        PhysicalPlan finalMainPlan = mainPlan;
+
+        try (
+            ComputeListener localListener = new ComputeListener(
+                transportService.getThreadPool(),
+                cancelQueryOnFailure,
+                finalListener.map(profiles -> {
+                    execInfo.markEndQuery();
+                    return new Result(finalMainPlan.output(), collectedPages, profiles, execInfo);
+                })
+            )
+        ) {
+            runCompute(rootTask, computeContext, finalMainPlan, localListener.acquireCompute());
+        }
+    }
+
+    public void executePlan(
+        String sessionId,
+        CancellableTask rootTask,
+        PhysicalPlan physicalPlan,
+        Configuration configuration,
+        FoldContext foldContext,
+        EsqlExecutionInfo execInfo,
+        ActionListener<Result> listener,
+        Supplier<ExchangeSink> exchangeSinkSupplier
     ) {
         Tuple<PhysicalPlan, PhysicalPlan> coordinatorAndDataNodePlan = PlannerUtils.breakPlanBetweenCoordinatorAndDataNode(
             physicalPlan,
@@ -185,7 +270,12 @@ public void execute(
             collectedPages.forEach(p -> Releasables.closeExpectNoException(p::releaseBlocks));
             l.onFailure(e);
         });
-        PhysicalPlan coordinatorPlan = new OutputExec(coordinatorAndDataNodePlan.v1(), collectedPages::add);
+        PhysicalPlan coordinatorPlan = coordinatorAndDataNodePlan.v1();
+
+        if (exchangeSinkSupplier == null) {
+            coordinatorPlan = new OutputExec(coordinatorAndDataNodePlan.v1(), collectedPages::add);
+        }
+
         PhysicalPlan dataNodePlan = coordinatorAndDataNodePlan.v2();
         if (dataNodePlan != null && dataNodePlan instanceof ExchangeSinkExec == false) {
             assert false : "expected data node plan starts with an ExchangeSink; got " + dataNodePlan;
@@ -211,7 +301,7 @@ public void execute(
                 configuration,
                 foldContext,
                 null,
-                null
+                exchangeSinkSupplier
             );
             updateShardCountForCoordinatorOnlyQuery(execInfo);
             try (
@@ -287,7 +377,7 @@ public void execute(
                             configuration,
                             foldContext,
                             exchangeSource::createExchangeSource,
-                            null
+                            exchangeSinkSupplier
                         ),
                         coordinatorPlan,
                         localListener.acquireCompute()