ES|QL: Improve generative tests for FORK [130015]

x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/EsqlQueryGenerator.java

@@ -76,6 +76,46 @@ public static CommandGenerator randomPipeCommandGenerator() {
         return randomFrom(PIPE_COMMANDS);
     }
 
+    public interface Executor {
+        void run(CommandGenerator generator, CommandGenerator.CommandDescription current);
+
+        List<CommandGenerator.CommandDescription> previousCommands();
+
+        boolean continueExecuting();
+
+        List<EsqlQueryGenerator.Column> currentSchema();
+
+    }
+
+    public static void generatePipeline(
+        final int depth,
+        CommandGenerator commandGenerator,
+        final CommandGenerator.QuerySchema schema,
+        Executor executor
+    ) {
+        CommandGenerator.CommandDescription desc = commandGenerator.generate(List.of(), List.of(), schema);
+        executor.run(commandGenerator, desc);
+        if (executor.continueExecuting() == false) {
+            return;
+        }
+
+        for (int j = 0; j < depth; j++) {
+            if (executor.currentSchema().isEmpty()) {
+                break;
+            }
+            commandGenerator = EsqlQueryGenerator.randomPipeCommandGenerator();
+            desc = commandGenerator.generate(executor.previousCommands(), executor.currentSchema(), schema);
+            if (desc == CommandGenerator.EMPTY_DESCRIPTION) {
+                continue;
+            }
+
+            executor.run(commandGenerator, desc);
+            if (executor.continueExecuting() == false) {
+                break;
+            }
+        }
+    }
+
     public static String booleanExpression(List<Column> previousOutput) {
         // TODO LIKE, RLIKE, functions etc.
         return switch (randomIntBetween(0, 3)) {

x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/GenerativeRestTest.java

@@ -87,43 +87,53 @@ public void test() throws IOException {
         List<LookupIdx> lookupIndices = lookupIndices();
         List<CsvTestsDataLoader.EnrichConfig> policies = availableEnrichPolicies();
         CommandGenerator.QuerySchema mappingInfo = new CommandGenerator.QuerySchema(indices, lookupIndices, policies);
-        EsqlQueryGenerator.QueryExecuted previousResult = null;
+
         for (int i = 0; i < ITERATIONS; i++) {
-            List<CommandGenerator.CommandDescription> previousCommands = new ArrayList<>();
-            CommandGenerator commandGenerator = EsqlQueryGenerator.sourceCommand();
-            CommandGenerator.CommandDescription desc = commandGenerator.generate(List.of(), List.of(), mappingInfo);
-            String command = desc.commandString();
-            EsqlQueryGenerator.QueryExecuted result = execute(command, 0);
-            if (result.exception() != null) {
-                checkException(result);
-                continue;
-            }
-            if (checkResults(List.of(), commandGenerator, desc, null, result).success() == false) {
-                continue;
-            }
-            previousResult = result;
-            previousCommands.add(desc);
-            for (int j = 0; j < MAX_DEPTH; j++) {
-                if (result.outputSchema().isEmpty()) {
-                    break;
+            var exec = new EsqlQueryGenerator.Executor() {
+                @Override
+                public void run(CommandGenerator generator, CommandGenerator.CommandDescription current) {
+                    previousCommands.add(current);
+                    final String command = current.commandString();
+
+                    final EsqlQueryGenerator.QueryExecuted result = previousResult == null
+                        ? execute(command, 0)
+                        : execute(previousResult.query() + command, previousResult.depth());
+                    previousResult = result;
+
+                    final boolean hasException = result.exception() != null;
+                    if (hasException || checkResults(List.of(), generator, current, previousResult, result).success() == false) {
+                        if (hasException) {
+                            checkException(result);
+                        }
+                        continueExecuting = false;
+                        currentSchema = List.of();
+                    } else {
+                        continueExecuting = true;
+                        currentSchema = result.outputSchema();
+                    }
                 }
-                commandGenerator = EsqlQueryGenerator.randomPipeCommandGenerator();
-                desc = commandGenerator.generate(previousCommands, result.outputSchema(), mappingInfo);
-                if (desc == CommandGenerator.EMPTY_DESCRIPTION) {
-                    continue;
+
+                @Override
+                public List<CommandGenerator.CommandDescription> previousCommands() {
+                    return previousCommands;
                 }
-                command = desc.commandString();
-                result = execute(result.query() + command, result.depth() + 1);
-                if (result.exception() != null) {
-                    checkException(result);
-                    break;
+
+                @Override
+                public boolean continueExecuting() {
+                    return continueExecuting;
                 }
-                if (checkResults(previousCommands, commandGenerator, desc, previousResult, result).success() == false) {
-                    break;
+
+                @Override
+                public List<EsqlQueryGenerator.Column> currentSchema() {
+                    return currentSchema;
                 }
-                previousCommands.add(desc);
-                previousResult = result;
-            }
+
+                boolean continueExecuting;
+                List<EsqlQueryGenerator.Column> currentSchema;
+                final List<CommandGenerator.CommandDescription> previousCommands = new ArrayList<>();
+                EsqlQueryGenerator.QueryExecuted previousResult;
+            };
+            EsqlQueryGenerator.generatePipeline(MAX_DEPTH, EsqlQueryGenerator.sourceCommand(), mappingInfo, exec);
         }
     }
 
@@ -163,7 +173,7 @@ private void checkException(EsqlQueryGenerator.QueryExecuted query) {
     }
 
     @SuppressWarnings("unchecked")
-    private EsqlQueryGenerator.QueryExecuted execute(String command, int depth) {
+    public static EsqlQueryGenerator.QueryExecuted execute(String command, int depth) {
         try {
             Map<String, Object> a = RestEsqlTestCase.runEsql(
                 new RestEsqlTestCase.RequestObjectBuilder().query(command).build(),
@@ -183,7 +193,7 @@ private EsqlQueryGenerator.QueryExecuted execute(String command, int depth) {
     }
 
     @SuppressWarnings("unchecked")
-    private List<EsqlQueryGenerator.Column> outputSchema(Map<String, Object> a) {
+    private static List<EsqlQueryGenerator.Column> outputSchema(Map<String, Object> a) {
         List<Map<String, String>> cols = (List<Map<String, String>>) a.get("columns");
         if (cols == null) {
             return null;

x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/pipe/ForkGenerator.java

@@ -8,10 +8,13 @@
 package org.elasticsearch.xpack.esql.qa.rest.generative.command.pipe;
 
 import org.elasticsearch.xpack.esql.qa.rest.generative.EsqlQueryGenerator;
+import org.elasticsearch.xpack.esql.qa.rest.generative.GenerativeRestTest;
 import org.elasticsearch.xpack.esql.qa.rest.generative.command.CommandGenerator;
 
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
+import java.util.stream.Collectors;
 
 import static org.elasticsearch.test.ESTestCase.randomIntBetween;
 
@@ -28,17 +31,105 @@ public CommandDescription generate(
     ) {
         // FORK can only be allowed once - so we skip adding another FORK if we already have one
         // otherwise, most generated queries would only result in a validation error
+        StringBuilder completeCommand = new StringBuilder();
         for (CommandDescription command : previousCommands) {
             if (command.commandName().equals(FORK)) {
-                return new CommandDescription(FORK, this, " ", Map.of());
+                return EMPTY_DESCRIPTION;
             }
+
+            completeCommand.append(command.commandString());
         }
 
-        int n = randomIntBetween(2, 3);
+        final int branchCount = randomIntBetween(2, 3);
+        final int branchToRetain = randomIntBetween(1, branchCount);
+
+        StringBuilder forkCmd = new StringBuilder(" | FORK ");
+        for (int i = 0; i < branchCount; i++) {
+            var expr = WhereGenerator.randomExpression(randomIntBetween(1, 2), previousOutput);
+            if (expr == null) {
+                expr = "true";
+            }
+            forkCmd.append(" (").append("where ").append(expr);
+
+            var exec = new EsqlQueryGenerator.Executor() {
+                @Override
+                public void run(CommandGenerator generator, CommandDescription current) {
+                    final String command = current.commandString();
+
+                    // Try appending new command to parent of Fork. If we successfully execute (without exception) AND still retain the same
+                    // schema, we append the command. Enforcing the same schema is stricter than the Fork needs (it only needs types to be
+                    // the same on columns which are present), but given we currently generate independent sub-pipelines, this way we can
+                    // generate more valid Fork queries.
+                    final EsqlQueryGenerator.QueryExecuted result = previousResult == null
+                        ? GenerativeRestTest.execute(command, 0)
+                        : GenerativeRestTest.execute(previousResult.query() + command, previousResult.depth());
+                    previousResult = result;
+
+                    continueExecuting = result.exception() == null && result.outputSchema().equals(previousOutput);
+                    if (continueExecuting) {
+                        previousCommands.add(current);
+                    }
+                }
+
+                @Override
+                public List<CommandDescription> previousCommands() {
+                    return previousCommands;
+                }
+
+                @Override
+                public boolean continueExecuting() {
+                    return continueExecuting;
+                }
 
-        String cmd = " | FORK " + "( WHERE true ) ".repeat(n) + " | WHERE _fork == \"fork" + randomIntBetween(1, n) + "\" | DROP _fork";
+                @Override
+                public List<EsqlQueryGenerator.Column> currentSchema() {
+                    return previousOutput;
+                }
+
+                final List<CommandGenerator.CommandDescription> previousCommands = new ArrayList<>();
+                boolean continueExecuting;
+                EsqlQueryGenerator.QueryExecuted previousResult;
+            };
+
+            var gen = new CommandGenerator() {
+                @Override
+                public CommandDescription generate(
+                    List<CommandDescription> previousCommands,
+                    List<EsqlQueryGenerator.Column> previousOutput,
+                    QuerySchema schema
+                ) {
+                    return new CommandDescription(FORK, this, completeCommand.toString(), Map.of());
+                }
+
+                @Override
+                public ValidationResult validateOutput(
+                    List<CommandDescription> previousCommands,
+                    CommandDescription command,
+                    List<EsqlQueryGenerator.Column> previousColumns,
+                    List<List<Object>> previousOutput,
+                    List<EsqlQueryGenerator.Column> columns,
+                    List<List<Object>> output
+                ) {
+                    return VALIDATION_OK;
+                }
+            };
+
+            EsqlQueryGenerator.generatePipeline(3, gen, schema, exec);
+            if (exec.previousCommands().size() > 1) {
+                String previousCmd = exec.previousCommands()
+                    .stream()
+                    .skip(1)
+                    .map(CommandDescription::commandString)
+                    .collect(Collectors.joining(" "));
+                forkCmd.append(previousCmd);
+            }
+
+            forkCmd.append(")");
+        }
+        forkCmd.append(" | WHERE _fork == \"fork").append(branchToRetain).append("\" | DROP _fork");
 
-        return new CommandDescription(FORK, this, cmd, Map.of());
+        // System.out.println("Generated fork command: " + forkCmd);
+        return new CommandDescription(FORK, this, forkCmd.toString(), Map.of());
     }
 
     @Override

x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/command/pipe/WhereGenerator.java

@@ -21,20 +21,15 @@ public class WhereGenerator implements CommandGenerator {
     public static final String WHERE = "where";
     public static final CommandGenerator INSTANCE = new WhereGenerator();
 
-    @Override
-    public CommandDescription generate(
-        List<CommandDescription> previousCommands,
-        List<EsqlQueryGenerator.Column> previousOutput,
-        QuerySchema schema
-    ) {
+    public static String randomExpression(final int nConditions, List<EsqlQueryGenerator.Column> previousOutput) {
         // TODO more complex conditions
-        StringBuilder result = new StringBuilder(" | where ");
-        int nConditions = randomIntBetween(1, 5);
+        var result = new StringBuilder();
+
         for (int i = 0; i < nConditions; i++) {
             String exp = EsqlQueryGenerator.booleanExpression(previousOutput);
             if (exp == null) {
-                // cannot generate expressions, just skip
-                return EMPTY_DESCRIPTION;
+                // Cannot generate expressions, just skip.
+                return null;
             }
             if (i > 0) {
                 result.append(randomBoolean() ? " AND " : " OR ");
@@ -45,8 +40,20 @@ public CommandDescription generate(
             result.append(exp);
         }
 
-        String cmd = result.toString();
-        return new CommandDescription(WHERE, this, cmd, Map.of());
+        return result.toString();
+    }
+
+    @Override
+    public CommandDescription generate(
+        List<CommandDescription> previousCommands,
+        List<EsqlQueryGenerator.Column> previousOutput,
+        QuerySchema schema
+    ) {
+        String expression = randomExpression(randomIntBetween(1, 5), previousOutput);
+        if (expression == null) {
+            return EMPTY_DESCRIPTION;
+        }
+        return new CommandDescription(WHERE, this, " | where " + expression, Map.of());
     }
 
     @Override