ES|QL: refactor generative tests

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

@@ -12,6 +12,7 @@
 import org.elasticsearch.test.rest.ESRestTestCase;
 import org.elasticsearch.xpack.esql.CsvTestsDataLoader;
 import org.elasticsearch.xpack.esql.qa.rest.RestEsqlTestCase;
+import org.elasticsearch.xpack.esql.qa.rest.generative.command.CommandGenerator;
 import org.junit.AfterClass;
 import org.junit.Before;
 
@@ -47,11 +48,14 @@ public abstract class GenerativeRestTest extends ESRestTestCase {
         "Field '.*' shadowed by field at line .*",
         "evaluation of \\[.*\\] failed, treating result as null", // TODO investigate?
 
-        // Awaiting fixes
+        // Awaiting fixes for query failure
         "Unknown column \\[<all-fields-projected>\\]", // https://github.com/elastic/elasticsearch/issues/121741,
         "Plan \\[ProjectExec\\[\\[<no-fields>.* optimized incorrectly due to missing references", // https://github.com/elastic/elasticsearch/issues/125866
         "optimized incorrectly due to missing references", // https://github.com/elastic/elasticsearch/issues/116781
-        "The incoming YAML document exceeds the limit:" // still to investigate, but it seems to be specific to the test framework
+        "The incoming YAML document exceeds the limit:", // still to investigate, but it seems to be specific to the test framework
+
+        // Awaiting fixes for correctness
+        "Expecting the following columns \\[.*\\], got" // https://github.com/elastic/elasticsearch/issues/129000
     );
 
     public static final Set<Pattern> ALLOWED_ERROR_PATTERNS = ALLOWED_ERRORS.stream()
@@ -84,27 +88,73 @@ public void test() throws IOException {
         List<String> indices = availableIndices();
         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++) {
-            String command = EsqlQueryGenerator.sourceCommand(indices);
+            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;
+                break;
+            }
+            if (checkResults(List.of(), commandGenerator, desc, null, result).success() == false) {
+                break;
             }
+            previousResult = result;
+            previousCommands.add(desc);
             for (int j = 0; j < MAX_DEPTH; j++) {
                 if (result.outputSchema().isEmpty()) {
                     break;
                 }
-                command = EsqlQueryGenerator.pipeCommand(result.outputSchema(), policies, lookupIndices);
+                commandGenerator = EsqlQueryGenerator.randomPipeCommandGenerator();
+                desc = commandGenerator.generate(previousCommands, result.outputSchema(), mappingInfo);
+                if (desc == CommandGenerator.EMPTY_DESCRIPTION) {
+                    continue;
+                }
+                command = desc.commandString();
                 result = execute(result.query() + command, result.depth() + 1);
                 if (result.exception() != null) {
                     checkException(result);
                     break;
                 }
+                if (checkResults(previousCommands, commandGenerator, desc, previousResult, result).success() == false) {
+                    break;
+                }
+                previousCommands.add(desc);
+                previousResult = result;
             }
         }
     }
 
+    private static CommandGenerator.ValidationResult checkResults(
+        List<CommandGenerator.CommandDescription> previousCommands,
+        CommandGenerator commandGenerator,
+        CommandGenerator.CommandDescription commandDescription,
+        EsqlQueryGenerator.QueryExecuted previousResult,
+        EsqlQueryGenerator.QueryExecuted result
+    ) {
+        CommandGenerator.ValidationResult outputValidation = commandGenerator.validateOutput(
+            previousCommands,
+            commandDescription,
+            previousResult == null ? null : previousResult.outputSchema(),
+            previousResult == null ? null : previousResult.result(),
+            result.outputSchema(),
+            result.result()
+        );
+        if (outputValidation.success() == false) {
+            for (Pattern allowedError : ALLOWED_ERROR_PATTERNS) {
+                if (allowedError.matcher(outputValidation.errorMessage()).matches()) {
+                    return outputValidation;
+                }
+            }
+            fail("query: " + result.query() + "\nerror: " + outputValidation.errorMessage());
+        }
+        return outputValidation;
+    }
+
     private void checkException(EsqlQueryGenerator.QueryExecuted query) {
         for (Pattern allowedError : ALLOWED_ERROR_PATTERNS) {
             if (allowedError.matcher(query.exception().getMessage()).matches()) {
@@ -114,16 +164,18 @@ private void checkException(EsqlQueryGenerator.QueryExecuted query) {
         fail("query: " + query.query() + "\nexception: " + query.exception().getMessage());
     }
 
+    @SuppressWarnings("unchecked")
     private EsqlQueryGenerator.QueryExecuted execute(String command, int depth) {
         try {
             Map<String, Object> a = RestEsqlTestCase.runEsqlSync(new RestEsqlTestCase.RequestObjectBuilder().query(command).build());
             List<EsqlQueryGenerator.Column> outputSchema = outputSchema(a);
-            return new EsqlQueryGenerator.QueryExecuted(command, depth, outputSchema, null);
+            List<List<Object>> values = (List<List<Object>>) a.get("values");
+            return new EsqlQueryGenerator.QueryExecuted(command, depth, outputSchema, values, null);
         } catch (Exception e) {
-            return new EsqlQueryGenerator.QueryExecuted(command, depth, null, e);
+            return new EsqlQueryGenerator.QueryExecuted(command, depth, null, null, e);
         } catch (AssertionError ae) {
             // this is for ensureNoWarnings()
-            return new EsqlQueryGenerator.QueryExecuted(command, depth, null, new RuntimeException(ae.getMessage()));
+            return new EsqlQueryGenerator.QueryExecuted(command, depth, null, null, new RuntimeException(ae.getMessage()));
         }
 
     }
@@ -144,7 +196,7 @@ private List<String> availableIndices() throws IOException {
             .toList();
     }
 
-    record LookupIdx(String idxName, String key, String keyType) {}
+    public record LookupIdx(String idxName, String key, String keyType) {}
 
     private List<LookupIdx> lookupIndices() {
         List<LookupIdx> result = new ArrayList<>();

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

@@ -0,0 +1,43 @@
+= ES|QL Generative Tests
+
+These tests generate random queries and execute them.
+
+The intention is not to test the single commands, but rather to test how ES|QL query engine
+(parser, optimizers, query layout, compute) manages very complex queries.
+
+The test workflow is the following:
+
+1. Generate a source command (eg. `FROM idx`)
+2. Execute it
+3. Check the result
+4. Based on the previous query output, generate a pipe command (eg. `| EVAL foo = to_lower(bar))`
+5. Append the command to the query and execute it
+6. Check the result
+7. If the query is less than N commands (see `GenerativeRestTest.MAX_DEPTH)`, go to point `4`
+
+This workflow is executed M times (see `GenerativeRestTest.ITERATIONS`)
+
+The result check happens at two levels:
+
+* query success/failure - If the query fails:
+  ** If the error is in `GenerativeRestTest.ALLOWED_ERRORS`, ignore it and start with next iteration.
+  ** Otherwise throw an assertion error
+* check result correctness - this is delegated to last executed command generator
+
+== Implementing your own command generator
+
+If you implement a new command, and you want it to be tested by the generative tests, you can add a command generator here.
+
+All you have to do is:
+
+* add a class in `org.elasticsearch.xpack.esql.qa.rest.generative.command.source` (if it's a source command) or in `org.elasticsearch.xpack.esql.qa.rest.generative.command.pipe` (if it's a pipe command)
+* Implement `CommandGenerator` interface (see its javadoc, it should be explicative. Or just have a look at one of the existing commands, eg. `SortGenerator`)
+** Implement `CommandGenerator.generate()` method, that will return the command.
+*** Have a look at `EsqlQueryGenerator`, it contains many utility methods that will help you generate random expressions.
+** Implement `CommandGenerator.validateOutput()` to validate the output of the query.
+* Add your class to `EsqlQueryGenerator.SOURCE_COMMANDS` (if it's a source command) or `EsqlQueryGenerator.PIPE_COMMANDS` (if it's a pipe command).
+* Run `GenerativeIT` at least a couple of times: these tests can be pretty noisy.
+* If you get unexpected errors (real bugs in ES|QL), please open an issue and add the error to `GenerativeRestTest.ALLOWED_ERRORS`. Run tests again until everything works fine.
+
+
+IMPORTANT: be careful when validating the output (Eg. the row count), as ES|QL can be quite non-deterministic when there are no SORTs

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

@@ -0,0 +1,142 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.esql.qa.rest.generative.command;
+
+import org.elasticsearch.xpack.esql.CsvTestsDataLoader;
+import org.elasticsearch.xpack.esql.qa.rest.generative.EsqlQueryGenerator;
+import org.elasticsearch.xpack.esql.qa.rest.generative.GenerativeRestTest;
+
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Implement this if you want to your command to be tested by the random query generator.
+ * Then add it to the right list in {@link EsqlQueryGenerator}
+ * <p>
+ * The i
+ */
+public interface CommandGenerator {
+
+    /**
+     * @param commandName   the name of the command that is being generated
+     * @param commandString the full command string, including the "|"
+     * @param context       additional information that could be useful for output validation.
+     *                      This will be passed to validateOutput after the query execution, together with the query output
+     */
+    record CommandDescription(String commandName, CommandGenerator generator, String commandString, Map<String, Object> context) {}
+
+    record QuerySchema(
+        List<String> baseIndices,
+        List<GenerativeRestTest.LookupIdx> lookupIndices,
+        List<CsvTestsDataLoader.EnrichConfig> enrichPolicies
+    ) {}
+
+    record ValidationResult(boolean success, String errorMessage) {}
+
+    CommandDescription EMPTY_DESCRIPTION = new CommandDescription("<empty>", new CommandGenerator() {
+        @Override
+        public CommandDescription generate(
+            List<CommandDescription> previousCommands,
+            List<EsqlQueryGenerator.Column> previousOutput,
+            QuerySchema schema
+        ) {
+            return EMPTY_DESCRIPTION;
+        }
+
+        @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;
+        }
+    }, "", Map.of());
+
+    ValidationResult VALIDATION_OK = new ValidationResult(true, null);
+
+    /**
+     * Implement this method to generate a command, that will be appended to an existing query and then executed.
+     * See also {@link CommandDescription}
+     *
+     * @param previousCommands the list of the previous commands in the query
+     * @param previousOutput   the output returned by the query so far.
+     * @param schema           The columns returned by the query so far. It contains name and type information for each column.
+     * @return All the details about the generated command. See {@link CommandDescription}.
+     * If something goes wrong and for some reason you can't generate a command, you should return {@link CommandGenerator#EMPTY_DESCRIPTION}
+     */
+    CommandDescription generate(
+        List<CommandDescription> previousCommands,
+        List<EsqlQueryGenerator.Column> previousOutput,
+        QuerySchema schema
+    );
+
+    /**
+     * This will be invoked after the query execution.
+     * You are expected to put validation logic in here.
+     *
+     * @param previousCommands The list of commands before the last generated one.
+     * @param command          The description of the command you just generated.
+     *                         It also contains the context information you stored during command generation.
+     * @param previousColumns  The output schema of the original query (without last generated command).
+     *                         It contains name and type information for each column, see {@link EsqlQueryGenerator.Column}
+     * @param previousOutput   The output of the original query (without last generated command), as a list (rows) of lists (columns) of values
+     * @param columns          The output schema of the full query (WITH last generated command).
+     * @param output           The output of the full query (WITH last generated command), as a list (rows) of lists (columns) of values
+     * @return The result of the output validation. If the validation succeeds, you should return {@link CommandGenerator#VALIDATION_OK}.
+     * Also, if for some reason you can't validate the output, just return {@link CommandGenerator#VALIDATION_OK}; for a command, having a generator without
+     * validation is much better than having no generator at all.
+     */
+    ValidationResult validateOutput(
+        List<CommandDescription> previousCommands,
+        CommandDescription command,
+        List<EsqlQueryGenerator.Column> previousColumns,
+        List<List<Object>> previousOutput,
+        List<EsqlQueryGenerator.Column> columns,
+        List<List<Object>> output
+    );
+
+    static ValidationResult expectSameRowCount(
+        List<CommandDescription> previousCommands,
+        List<List<Object>> previousOutput,
+        List<List<Object>> output
+    ) {
+
+        // ES|QL is quite non-deterministic in this sense, we can't guarantee it for now
+        // if (output.size() != previousOutput.size()) {
+        // return new ValidationResult(false, "Expecting [" + previousOutput.size() + "] rows, but got [" + output.size() + "]");
+        // }
+
+        return VALIDATION_OK;
+    }
+
+    static ValidationResult expectSameColumns(List<EsqlQueryGenerator.Column> previousColumns, List<EsqlQueryGenerator.Column> columns) {
+
+        if (previousColumns.stream().anyMatch(x -> x.name().contains("<all-fields-projected>"))) {
+            return VALIDATION_OK; // known bug
+        }
+
+        if (previousColumns.size() != columns.size()) {
+            return new ValidationResult(false, "Expecting [" + previousColumns.size() + "] columns, got [" + columns.size() + "]");
+        }
+
+        List<String> prevColNames = previousColumns.stream().map(EsqlQueryGenerator.Column::name).toList();
+        List<String> newColNames = columns.stream().map(EsqlQueryGenerator.Column::name).toList();
+        if (prevColNames.equals(newColNames) == false) {
+            return new ValidationResult(
+                false,
+                "Expecting the following columns [" + String.join(", ", prevColNames) + "], got [" + String.join(", ", newColNames) + "]"
+            );
+        }
+
+        return VALIDATION_OK;
+    }
+}

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

@@ -0,0 +1,78 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+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.command.CommandGenerator;
+
+import java.util.List;
+import java.util.Map;
+
+import static org.elasticsearch.test.ESTestCase.randomBoolean;
+import static org.elasticsearch.test.ESTestCase.randomIntBetween;
+
+public class DissectGenerator implements CommandGenerator {
+
+    public static final String DISSECT = "dissect";
+    public static final CommandGenerator INSTANCE = new DissectGenerator();
+
+    @Override
+    public CommandDescription generate(
+        List<CommandDescription> previousCommands,
+        List<EsqlQueryGenerator.Column> previousOutput,
+        QuerySchema schema
+    ) {
+        String field = EsqlQueryGenerator.randomStringField(previousOutput);
+        if (field == null) {
+            return EMPTY_DESCRIPTION;// no strings to dissect, just skip
+        }
+        StringBuilder result = new StringBuilder(" | dissect ");
+        result.append(field);
+        result.append(" \"");
+        for (int i = 0; i < randomIntBetween(1, 3); i++) {
+            if (i > 0) {
+                result.append(" ");
+            }
+            result.append("%{");
+            String fieldName;
+            if (randomBoolean()) {
+                fieldName = EsqlQueryGenerator.randomIdentifier();
+            } else {
+                fieldName = EsqlQueryGenerator.randomRawName(previousOutput);
+                if (fieldName == null) {
+                    fieldName = EsqlQueryGenerator.randomIdentifier();
+                }
+            }
+            result.append(fieldName);
+            result.append("}");
+        }
+        result.append("\"");
+        String cmdString = result.toString();
+        return new CommandDescription(DISSECT, this, cmdString, Map.of());
+    }
+
+    @Override
+    public ValidationResult validateOutput(
+        List<CommandDescription> previousCommands,
+        CommandDescription commandDescription,
+        List<EsqlQueryGenerator.Column> previousColumns,
+        List<List<Object>> previousOutput,
+        List<EsqlQueryGenerator.Column> columns,
+        List<List<Object>> output
+    ) {
+        if (commandDescription == EMPTY_DESCRIPTION) {
+            return VALIDATION_OK;
+        }
+
+        if (previousColumns.size() > columns.size()) {
+            return new ValidationResult(false, "Expecting at least [" + previousColumns.size() + "] columns, got [" + columns.size() + "]");
+        }
+
+        return CommandGenerator.expectSameRowCount(previousCommands, previousOutput, output);
+    }
+}