FORK - allow EVAL/DISSECT/STATS in branches

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

@@ -112,3 +112,86 @@ fork1         | 5.603396578413904E18  | 2           | all we have to decide is w
 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
 ;
+
+forkWithEvals
+required_capability: fork_v2
+
+FROM employees
+| FORK (WHERE emp_no == 10048 OR emp_no == 10081 | EVAL x = "abc" | EVAL y = 1)
+       (WHERE emp_no == 10081 OR emp_no == 10087 | EVAL x = "def" | EVAL z = 2)
+| KEEP _fork, emp_no, x, y, z
+| SORT _fork, emp_no
+;
+
+_fork:keyword | emp_no:integer | x:keyword | y:integer | z:integer
+fork1         | 10048          | abc       | 1         | null
+fork1         | 10081          | abc       | 1         | null
+fork2         | 10081          | def       | null      | 2
+fork2         | 10087          | def       | null      | 2
+;
+
+forkWithStats
+required_capability: fork_v2
+
+FROM employees
+| FORK (WHERE emp_no == 10048 OR emp_no == 10081)
+       (WHERE emp_no == 10081 OR emp_no == 10087)
+       (STATS x = COUNT(*), y = MAX(emp_no), z = MIN(emp_no))
+       (STATS x = COUNT(*), y = MIN(emp_no))
+| KEEP _fork, emp_no, x, y, z
+| SORT _fork, emp_no
+;
+
+_fork:keyword | emp_no:integer | x:long | y:integer | z:integer
+fork1         | 10048          | null   | null      | null
+fork1         | 10081          | null   | null      | null
+fork2         | 10081          | null   | null      | null
+fork2         | 10087          | null   | null      | null
+fork3         | null           | 100    | 10100     | 10001
+fork4         | null           | 100    | 10001     | null
+;
+
+forkWithDissect
+required_capability: fork_v2
+
+FROM employees
+| WHERE emp_no == 10048 OR emp_no == 10081
+| FORK (EVAL a = CONCAT(first_name, " ", emp_no::keyword, " ", last_name)
+        | DISSECT a "%{x} %{y} %{z}" )
+       (EVAL b = CONCAT(last_name, " ", emp_no::keyword, " ", first_name)
+        | DISSECT b "%{x} %{y} %{w}" )
+| KEEP _fork, emp_no, x, y, z, w
+| SORT _fork, emp_no
+;
+
+_fork:keyword | emp_no:integer | x:keyword | y:keyword | z:keyword | w:keyword
+fork1         | 10048          | Florian   | 10048     | Syrotiuk  | null
+fork1         | 10081          | Zhongwei  | 10081     | Rosen     | null
+fork2         | 10048          | Syrotiuk  | 10048     | null      | Florian
+fork2         | 10081          | Rosen     | 10081     | null      | Zhongwei
+;
+
+forkWithMixOfCommands
+required_capability: fork_v2
+
+FROM employees
+| WHERE emp_no == 10048 OR emp_no == 10081
+| FORK ( EVAL a = CONCAT(first_name, " ", emp_no::keyword, " ", last_name)
+        | DISSECT a "%{x} %{y} %{z}"
+        | EVAL y = y::keyword )
+       ( STATS x = COUNT(*)::keyword, y = MAX(emp_no)::keyword, z = MIN(emp_no)::keyword )
+       ( SORT emp_no ASC | LIMIT 2 | EVAL x = last_name )
+       ( EVAL x = "abc" | EVAL y = "aaa" )
+| KEEP _fork, emp_no, x, y, z, a
+| SORT _fork, emp_no
+;
+
+_fork:keyword | emp_no:integer | x:keyword | y:keyword | z:keyword | a:keyword
+fork1         | 10048          | Florian   | 10048     | Syrotiuk  | Florian 10048 Syrotiuk
+fork1         | 10081          | Zhongwei  | 10081     | Rosen     | Zhongwei 10081 Rosen
+fork2         | null           | 2         | 10081     | 10048     | null
+fork3         | 10048          | Syrotiuk  | null      | null      | null
+fork3         | 10081          | Rosen     | null      | null      | null
+fork4         | 10048          | abc       | aaa       | null      | null
+fork4         | 10081          | abc       | aaa       | null      | null
+;

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

@@ -14,6 +14,7 @@
 import org.elasticsearch.xpack.esql.parser.ParsingException;
 import org.junit.Before;
 
+import java.util.Arrays;
 import java.util.Iterator;
 import java.util.List;
 import java.util.function.Predicate;
@@ -501,6 +502,75 @@ public void testSubqueryWithoutLimitOnly() {   // this should
         }
     }
 
+    public void testWithEvalSimple() {
+        var query = """
+                FROM test
+                | WHERE content:"cat"
+                | FORK ( EVAL a = 1 )
+                       ( EVAL a = 2 )
+                | KEEP a, _fork, id, content
+            """;
+
+        try (var resp = run(query)) {
+            assertColumnNames(resp.columns(), List.of("a", "_fork", "id", "content"));
+
+            Iterable<Iterable<Object>> expectedValues = List.of(
+                List.of(1, "fork1", 5, "There is also a white cat"),
+                List.of(2, "fork2", 5, "There is also a white cat")
+            );
+            assertValues(resp.values(), expectedValues);
+        }
+    }
+
+    public void testWithEvalDifferentOutputs() {
+        var query = """
+                FROM test
+                | WHERE id == 2
+                | FORK ( EVAL a = 1 )
+                       ( EVAL b = 2 )
+                | KEEP a, b, _fork
+                | SORT _fork, a
+            """;
+        try (var resp = run(query)) {
+            assertColumnNames(resp.columns(), List.of("a", "b", "_fork"));
+            Iterable<Iterable<Object>> expectedValues = List.of(
+                Arrays.stream(new Object[] { 1, null, "fork1" }).toList(),
+                Arrays.stream(new Object[] { null, 2, "fork2" }).toList()
+            );
+            assertValues(resp.values(), expectedValues);
+        }
+    }
+
+    public void testWithStatsSimple() {
+        var query = """
+                FROM test
+                | FORK (STATS x=COUNT(*), y=VALUES(id))
+                       (WHERE id == 2)
+                | KEEP _fork, x, y, id
+                | SORT _fork, id
+            """;
+        try (var resp = run(query)) {
+            assertColumnNames(resp.columns(), List.of("_fork", "x", "y", "id"));
+            Iterable<Iterable<Object>> expectedValues = List.of(
+                Arrays.stream(new Object[] { "fork1", 6L, List.of(1, 2, 3, 4, 5, 6), null }).toList(),
+                Arrays.stream(new Object[] { "fork2", null, null, 2 }).toList()
+            );
+            assertValues(resp.values(), expectedValues);
+        }
+    }
+
+    public void testWithEvalWithConflictingTypes() {
+        var query = """
+                FROM test
+                | FORK ( EVAL a = 1 )
+                       ( EVAL a = "aaaa" )
+                | KEEP a, _fork
+            """;
+
+        var e = expectThrows(VerificationException.class, () -> run(query));
+        assertTrue(e.getMessage().contains("Column [a] has conflicting data types"));
+    }
+
     public void testSubqueryWithUnknownField() {
         var query = """
             FROM test
@@ -565,6 +635,19 @@ public void testSubqueryWithUnknownFieldInSort() {
         assertTrue(e.getMessage().contains("Unknown column [bar]"));
     }
 
+    public void testSubqueryWithUnknownFieldInEval() {
+        var query = """
+            FROM test
+            | FORK
+               ( EVAL x = baz + 1)
+               ( WHERE content:"cat" )
+            | KEEP _fork, id, content
+            | SORT _fork, id
+            """;
+        var e = expectThrows(VerificationException.class, () -> run(query));
+        assertTrue(e.getMessage().contains("Unknown column [baz]"));
+    }
+
     public void testOneSubQuery() {
         var query = """
             FROM test

x-pack/plugin/esql/src/main/antlr/EsqlBaseParser.g4

@@ -281,9 +281,12 @@ forkSubQueryCommand
     ;
 
 forkSubQueryProcessingCommand
-    : whereCommand
-    | sortCommand
+    : evalCommand
+    | whereCommand
     | limitCommand
+    | statsCommand
+    | sortCommand
+    | dissectCommand
     ;
 
 rrfCommand

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

@@ -964,7 +964,12 @@ public enum Cap {
         /**
          * Support max_over_time aggregation
          */
-        MAX_OVER_TIME(Build.current().isSnapshot());
+        MAX_OVER_TIME(Build.current().isSnapshot()),
+
+        /**
+         * Support STATS/EVAL/DISSECT in Fork branches
+         */
+        FORK_V2(Build.current().isSnapshot());
 
         private final boolean enabled;
 

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

@@ -153,8 +153,9 @@
 public class Analyzer extends ParameterizedRuleExecutor<LogicalPlan, AnalyzerContext> {
     // marker list of attributes for plans that do not have any concrete fields to return, but have other computed columns to return
     // ie from test | stats c = count(*)
+    public static final String NO_FIELDS_NAME = "<no-fields>";
     public static final List<Attribute> NO_FIELDS = List.of(
-        new ReferenceAttribute(Source.EMPTY, "<no-fields>", DataType.NULL, Nullability.TRUE, null, true)
+        new ReferenceAttribute(Source.EMPTY, NO_FIELDS_NAME, DataType.NULL, Nullability.TRUE, null, true)
     );
 
     private static final List<Batch<LogicalPlan>> RULES = List.of(
@@ -499,6 +500,10 @@ protected LogicalPlan rule(LogicalPlan plan, AnalyzerContext context) {
                 return resolveKeep(p, childrenOutput);
             }
 
+            if (plan instanceof Fork f) {
+                return resolveFork(f, context);
+            }
+
             if (plan instanceof Eval p) {
                 return resolveEval(p, childrenOutput);
             }
@@ -714,6 +719,56 @@ private Join resolveLookupJoin(LookupJoin join) {
             return join;
         }
 
+        private LogicalPlan resolveFork(Fork fork, AnalyzerContext context) {
+            // we align the outputs of the sub plans such that they have the same columns
+            boolean changed = false;
+            List<LogicalPlan> newSubPlans = new ArrayList<>();
+            Set<String> forkColumns = fork.outputSet().names();
+
+            for (LogicalPlan logicalPlan : fork.children()) {
+                Source source = logicalPlan.source();
+
+                // find the missing columns
+                List<Attribute> missing = new ArrayList<>();
+                Set<String> currentNames = logicalPlan.outputSet().names();
+                for (Attribute attr : fork.outputSet()) {
+                    if (currentNames.contains(attr.name()) == false) {
+                        missing.add(attr);
+                    }
+                }
+
+                List<Alias> aliases = missing.stream().map(attr -> new Alias(source, attr.name(), Literal.of(attr, null))).toList();
+
+                // add the missing columns
+                if (aliases.size() > 0) {
+                    logicalPlan = new Eval(source, logicalPlan, aliases);
+                    changed = true;
+                }
+
+                List<String> subPlanColumns = logicalPlan.output().stream().map(Attribute::name).toList();
+                // We need to add an explicit Keep even if the outputs align
+                // This is because at the moment the sub plans are executed and optimized separately and the output might change
+                // during optimizations. Once we add streaming we might not need to add a Keep when the outputs already align.
+                // Note that until we add explicit support for KEEP in FORK branches, this condition will always be true.
+                if (logicalPlan instanceof Keep == false || subPlanColumns.equals(forkColumns) == false) {
+                    changed = true;
+                    List<Attribute> newOutput = new ArrayList<>();
+                    for (String attrName : forkColumns) {
+                        for (Attribute subAttr : logicalPlan.output()) {
+                            if (attrName.equals(subAttr.name())) {
+                                newOutput.add(subAttr);
+                            }
+                        }
+                    }
+                    logicalPlan = new Keep(logicalPlan.source(), logicalPlan, newOutput);
+                }
+
+                newSubPlans.add(logicalPlan);
+            }
+
+            return changed ? new Fork(fork.source(), newSubPlans) : fork;
+        }
+
         private LogicalPlan resolveRerank(Rerank rerank, List<Attribute> childrenOutput) {
             List<Alias> newFields = new ArrayList<>();
             boolean changed = false;

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/PruneColumns.java

@@ -14,9 +14,11 @@
 import org.elasticsearch.xpack.esql.core.expression.EmptyAttribute;
 import org.elasticsearch.xpack.esql.core.expression.Expressions;
 import org.elasticsearch.xpack.esql.core.expression.NamedExpression;
+import org.elasticsearch.xpack.esql.core.util.Holder;
 import org.elasticsearch.xpack.esql.plan.logical.Aggregate;
 import org.elasticsearch.xpack.esql.plan.logical.EsRelation;
 import org.elasticsearch.xpack.esql.plan.logical.Eval;
+import org.elasticsearch.xpack.esql.plan.logical.Fork;
 import org.elasticsearch.xpack.esql.plan.logical.Limit;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
 import org.elasticsearch.xpack.esql.plan.logical.local.LocalRelation;
@@ -36,6 +38,8 @@ public final class PruneColumns extends Rule<LogicalPlan, LogicalPlan> {
     public LogicalPlan apply(LogicalPlan plan) {
         // track used references
         var used = plan.outputSet().asBuilder();
+        Holder<Boolean> forkPresent = new Holder<>(false);
+
         // while going top-to-bottom (upstream)
         var pl = plan.transformDown(p -> {
             // Note: It is NOT required to do anything special for binary plans like JOINs. It is perfectly fine that transformDown descends
@@ -50,6 +54,14 @@ public LogicalPlan apply(LogicalPlan plan) {
                 return p;
             }
 
+            if (p instanceof Fork) {
+                forkPresent.set(true);
+            }
+            // pruning columns for Fork branches can have the side effect of having misaligned outputs
+            if (forkPresent.get()) {
+                return p;
+            }
+
             // remember used
             boolean recheck;
             // analyze the unused items against dedicated 'producer' nodes such as Eval and Aggregate