Stream result pages from sub plans for FORK (#126705)

Author: ioanatia

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;
             }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/ProjectAwayColumns.java

@@ -31,7 +31,6 @@
 import org.elasticsearch.compute.operator.LimitOperator;
 import org.elasticsearch.compute.operator.LocalSourceOperator;
 import org.elasticsearch.compute.operator.LocalSourceOperator.LocalSourceFactory;
-import org.elasticsearch.compute.operator.MergeOperator;
 import org.elasticsearch.compute.operator.MvExpandOperator;
 import org.elasticsearch.compute.operator.Operator;
 import org.elasticsearch.compute.operator.Operator.OperatorFactory;
@@ -103,7 +102,6 @@
 import org.elasticsearch.xpack.esql.plan.physical.LimitExec;
 import org.elasticsearch.xpack.esql.plan.physical.LocalSourceExec;
 import org.elasticsearch.xpack.esql.plan.physical.LookupJoinExec;
-import org.elasticsearch.xpack.esql.plan.physical.MergeExec;
 import org.elasticsearch.xpack.esql.plan.physical.MvExpandExec;
 import org.elasticsearch.xpack.esql.plan.physical.OutputExec;
 import org.elasticsearch.xpack.esql.plan.physical.PhysicalPlan;
@@ -281,8 +279,6 @@ else if (node instanceof OutputExec outputExec) {
             return planOutput(outputExec, context);
         } else if (node instanceof ExchangeSinkExec exchangeSink) {
             return planExchangeSink(exchangeSink, context);
-        } else if (node instanceof MergeExec mergeExec) {
-            return planMerge(mergeExec, context);
         } else if (node instanceof RrfScoreEvalExec rrf) {
             return planRrfScoreEvalExec(rrf, context);
         }
@@ -804,13 +800,6 @@ private PhysicalOperation planChangePoint(ChangePointExec changePoint, LocalExec
         );
     }
 
-    private PhysicalOperation planMerge(MergeExec mergeExec, LocalExecutionPlannerContext context) {
-        Layout.Builder layout = new Layout.Builder();
-        layout.append(mergeExec.output());
-        MergeOperator.BlockSuppliers suppliers = () -> mergeExec.suppliers().stream().map(s -> s.get()).toList();
-        return PhysicalOperation.fromSource(new MergeOperator.MergeOperatorFactory(suppliers), layout.build());
-    }
-
     /**
      * Immutable physical operation.
      */

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/LocalExecutionPlanner.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,34 @@
 
 public class PlannerUtils {
 
+    /**
+     * When the plan contains children like {@code MergeExec} resulted from the planning of commands such as FORK,
+     * we need to break the plan into sub plans and a main coordinator plan.
+     * The result pages from each sub plan will be funneled to the main coordinator plan.
+     * To achieve this, we wire each sub plan with a {@code ExchangeSinkExec} and add a {@code ExchangeSourceExec}
+     * to the main coordinator plan.
+     * There is an additional split of each sub plan into a data node plan and coordinator plan.
+     * This split is not done here, but as part of {@code PlannerUtils#breakPlanBetweenCoordinatorAndDataNode}.
+     */
+    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>();