Skip to content

Commit 134813a

Browse files
ES|QL: refactor generative tests (#129028)
1 parent 155c0da commit 134813a

18 files changed

+1333
-385
lines changed

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

Lines changed: 73 additions & 376 deletions
Large diffs are not rendered by default.

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

Lines changed: 61 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
import org.elasticsearch.test.rest.ESRestTestCase;
1313
import org.elasticsearch.xpack.esql.CsvTestsDataLoader;
1414
import org.elasticsearch.xpack.esql.qa.rest.RestEsqlTestCase;
15+
import org.elasticsearch.xpack.esql.qa.rest.generative.command.CommandGenerator;
1516
import org.junit.AfterClass;
1617
import org.junit.Before;
1718

@@ -47,11 +48,14 @@ public abstract class GenerativeRestTest extends ESRestTestCase {
4748
"Field '.*' shadowed by field at line .*",
4849
"evaluation of \\[.*\\] failed, treating result as null", // TODO investigate?
4950

50-
// Awaiting fixes
51+
// Awaiting fixes for query failure
5152
"Unknown column \\[<all-fields-projected>\\]", // https://github.com/elastic/elasticsearch/issues/121741,
5253
"Plan \\[ProjectExec\\[\\[<no-fields>.* optimized incorrectly due to missing references", // https://github.com/elastic/elasticsearch/issues/125866
5354
"optimized incorrectly due to missing references", // https://github.com/elastic/elasticsearch/issues/116781
54-
"The incoming YAML document exceeds the limit:" // still to investigate, but it seems to be specific to the test framework
55+
"The incoming YAML document exceeds the limit:", // still to investigate, but it seems to be specific to the test framework
56+
57+
// Awaiting fixes for correctness
58+
"Expecting the following columns \\[.*\\], got" // https://github.com/elastic/elasticsearch/issues/129000
5559
);
5660

5761
public static final Set<Pattern> ALLOWED_ERROR_PATTERNS = ALLOWED_ERRORS.stream()
@@ -84,27 +88,73 @@ public void test() throws IOException {
8488
List<String> indices = availableIndices();
8589
List<LookupIdx> lookupIndices = lookupIndices();
8690
List<CsvTestsDataLoader.EnrichConfig> policies = availableEnrichPolicies();
91+
CommandGenerator.QuerySchema mappingInfo = new CommandGenerator.QuerySchema(indices, lookupIndices, policies);
92+
EsqlQueryGenerator.QueryExecuted previousResult = null;
8793
for (int i = 0; i < ITERATIONS; i++) {
88-
String command = EsqlQueryGenerator.sourceCommand(indices);
94+
List<CommandGenerator.CommandDescription> previousCommands = new ArrayList<>();
95+
CommandGenerator commandGenerator = EsqlQueryGenerator.sourceCommand();
96+
CommandGenerator.CommandDescription desc = commandGenerator.generate(List.of(), List.of(), mappingInfo);
97+
String command = desc.commandString();
8998
EsqlQueryGenerator.QueryExecuted result = execute(command, 0);
9099
if (result.exception() != null) {
91100
checkException(result);
92-
continue;
101+
break;
102+
}
103+
if (checkResults(List.of(), commandGenerator, desc, null, result).success() == false) {
104+
break;
93105
}
106+
previousResult = result;
107+
previousCommands.add(desc);
94108
for (int j = 0; j < MAX_DEPTH; j++) {
95109
if (result.outputSchema().isEmpty()) {
96110
break;
97111
}
98-
command = EsqlQueryGenerator.pipeCommand(result.outputSchema(), policies, lookupIndices);
112+
commandGenerator = EsqlQueryGenerator.randomPipeCommandGenerator();
113+
desc = commandGenerator.generate(previousCommands, result.outputSchema(), mappingInfo);
114+
if (desc == CommandGenerator.EMPTY_DESCRIPTION) {
115+
continue;
116+
}
117+
command = desc.commandString();
99118
result = execute(result.query() + command, result.depth() + 1);
100119
if (result.exception() != null) {
101120
checkException(result);
102121
break;
103122
}
123+
if (checkResults(previousCommands, commandGenerator, desc, previousResult, result).success() == false) {
124+
break;
125+
}
126+
previousCommands.add(desc);
127+
previousResult = result;
104128
}
105129
}
106130
}
107131

132+
private static CommandGenerator.ValidationResult checkResults(
133+
List<CommandGenerator.CommandDescription> previousCommands,
134+
CommandGenerator commandGenerator,
135+
CommandGenerator.CommandDescription commandDescription,
136+
EsqlQueryGenerator.QueryExecuted previousResult,
137+
EsqlQueryGenerator.QueryExecuted result
138+
) {
139+
CommandGenerator.ValidationResult outputValidation = commandGenerator.validateOutput(
140+
previousCommands,
141+
commandDescription,
142+
previousResult == null ? null : previousResult.outputSchema(),
143+
previousResult == null ? null : previousResult.result(),
144+
result.outputSchema(),
145+
result.result()
146+
);
147+
if (outputValidation.success() == false) {
148+
for (Pattern allowedError : ALLOWED_ERROR_PATTERNS) {
149+
if (allowedError.matcher(outputValidation.errorMessage()).matches()) {
150+
return outputValidation;
151+
}
152+
}
153+
fail("query: " + result.query() + "\nerror: " + outputValidation.errorMessage());
154+
}
155+
return outputValidation;
156+
}
157+
108158
private void checkException(EsqlQueryGenerator.QueryExecuted query) {
109159
for (Pattern allowedError : ALLOWED_ERROR_PATTERNS) {
110160
if (allowedError.matcher(query.exception().getMessage()).matches()) {
@@ -114,16 +164,18 @@ private void checkException(EsqlQueryGenerator.QueryExecuted query) {
114164
fail("query: " + query.query() + "\nexception: " + query.exception().getMessage());
115165
}
116166

167+
@SuppressWarnings("unchecked")
117168
private EsqlQueryGenerator.QueryExecuted execute(String command, int depth) {
118169
try {
119170
Map<String, Object> a = RestEsqlTestCase.runEsqlSync(new RestEsqlTestCase.RequestObjectBuilder().query(command).build());
120171
List<EsqlQueryGenerator.Column> outputSchema = outputSchema(a);
121-
return new EsqlQueryGenerator.QueryExecuted(command, depth, outputSchema, null);
172+
List<List<Object>> values = (List<List<Object>>) a.get("values");
173+
return new EsqlQueryGenerator.QueryExecuted(command, depth, outputSchema, values, null);
122174
} catch (Exception e) {
123-
return new EsqlQueryGenerator.QueryExecuted(command, depth, null, e);
175+
return new EsqlQueryGenerator.QueryExecuted(command, depth, null, null, e);
124176
} catch (AssertionError ae) {
125177
// this is for ensureNoWarnings()
126-
return new EsqlQueryGenerator.QueryExecuted(command, depth, null, new RuntimeException(ae.getMessage()));
178+
return new EsqlQueryGenerator.QueryExecuted(command, depth, null, null, new RuntimeException(ae.getMessage()));
127179
}
128180

129181
}
@@ -144,7 +196,7 @@ private List<String> availableIndices() throws IOException {
144196
.toList();
145197
}
146198

147-
record LookupIdx(String idxName, String key, String keyType) {}
199+
public record LookupIdx(String idxName, String key, String keyType) {}
148200

149201
private List<LookupIdx> lookupIndices() {
150202
List<LookupIdx> result = new ArrayList<>();
Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
= ES|QL Generative Tests
2+
3+
These tests generate random queries and execute them.
4+
5+
The intention is not to test the single commands, but rather to test how ES|QL query engine
6+
(parser, optimizers, query layout, compute) manages very complex queries.
7+
8+
The test workflow is the following:
9+
10+
1. Generate a source command (eg. `FROM idx`)
11+
2. Execute it
12+
3. Check the result
13+
4. Based on the previous query output, generate a pipe command (eg. `| EVAL foo = to_lower(bar))`
14+
5. Append the command to the query and execute it
15+
6. Check the result
16+
7. If the query is less than N commands (see `GenerativeRestTest.MAX_DEPTH)`, go to point `4`
17+
18+
This workflow is executed M times (see `GenerativeRestTest.ITERATIONS`)
19+
20+
The result check happens at two levels:
21+
22+
* query success/failure - If the query fails:
23+
** If the error is in `GenerativeRestTest.ALLOWED_ERRORS`, ignore it and start with next iteration.
24+
** Otherwise throw an assertion error
25+
* check result correctness - this is delegated to last executed command generator
26+
27+
== Implementing your own command generator
28+
29+
If you implement a new command, and you want it to be tested by the generative tests, you can add a command generator here.
30+
31+
All you have to do is:
32+
33+
* 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)
34+
* Implement `CommandGenerator` interface (see its javadoc, it should be explicative. Or just have a look at one of the existing commands, eg. `SortGenerator`)
35+
** Implement `CommandGenerator.generate()` method, that will return the command.
36+
*** Have a look at `EsqlQueryGenerator`, it contains many utility methods that will help you generate random expressions.
37+
** Implement `CommandGenerator.validateOutput()` to validate the output of the query.
38+
* Add your class to `EsqlQueryGenerator.SOURCE_COMMANDS` (if it's a source command) or `EsqlQueryGenerator.PIPE_COMMANDS` (if it's a pipe command).
39+
* Run `GenerativeIT` at least a couple of times: these tests can be pretty noisy.
40+
* 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.
41+
42+
43+
IMPORTANT: be careful when validating the output (Eg. the row count), as ES|QL can be quite non-deterministic when there are no SORTs
Lines changed: 142 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,142 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
package org.elasticsearch.xpack.esql.qa.rest.generative.command;
9+
10+
import org.elasticsearch.xpack.esql.CsvTestsDataLoader;
11+
import org.elasticsearch.xpack.esql.qa.rest.generative.EsqlQueryGenerator;
12+
import org.elasticsearch.xpack.esql.qa.rest.generative.GenerativeRestTest;
13+
14+
import java.util.List;
15+
import java.util.Map;
16+
17+
/**
18+
* Implement this if you want to your command to be tested by the random query generator.
19+
* Then add it to the right list in {@link EsqlQueryGenerator}
20+
* <p>
21+
* The i
22+
*/
23+
public interface CommandGenerator {
24+
25+
/**
26+
* @param commandName the name of the command that is being generated
27+
* @param commandString the full command string, including the "|"
28+
* @param context additional information that could be useful for output validation.
29+
* This will be passed to validateOutput after the query execution, together with the query output
30+
*/
31+
record CommandDescription(String commandName, CommandGenerator generator, String commandString, Map<String, Object> context) {}
32+
33+
record QuerySchema(
34+
List<String> baseIndices,
35+
List<GenerativeRestTest.LookupIdx> lookupIndices,
36+
List<CsvTestsDataLoader.EnrichConfig> enrichPolicies
37+
) {}
38+
39+
record ValidationResult(boolean success, String errorMessage) {}
40+
41+
CommandDescription EMPTY_DESCRIPTION = new CommandDescription("<empty>", new CommandGenerator() {
42+
@Override
43+
public CommandDescription generate(
44+
List<CommandDescription> previousCommands,
45+
List<EsqlQueryGenerator.Column> previousOutput,
46+
QuerySchema schema
47+
) {
48+
return EMPTY_DESCRIPTION;
49+
}
50+
51+
@Override
52+
public ValidationResult validateOutput(
53+
List<CommandDescription> previousCommands,
54+
CommandDescription command,
55+
List<EsqlQueryGenerator.Column> previousColumns,
56+
List<List<Object>> previousOutput,
57+
List<EsqlQueryGenerator.Column> columns,
58+
List<List<Object>> output
59+
) {
60+
return VALIDATION_OK;
61+
}
62+
}, "", Map.of());
63+
64+
ValidationResult VALIDATION_OK = new ValidationResult(true, null);
65+
66+
/**
67+
* Implement this method to generate a command, that will be appended to an existing query and then executed.
68+
* See also {@link CommandDescription}
69+
*
70+
* @param previousCommands the list of the previous commands in the query
71+
* @param previousOutput the output returned by the query so far.
72+
* @param schema The columns returned by the query so far. It contains name and type information for each column.
73+
* @return All the details about the generated command. See {@link CommandDescription}.
74+
* If something goes wrong and for some reason you can't generate a command, you should return {@link CommandGenerator#EMPTY_DESCRIPTION}
75+
*/
76+
CommandDescription generate(
77+
List<CommandDescription> previousCommands,
78+
List<EsqlQueryGenerator.Column> previousOutput,
79+
QuerySchema schema
80+
);
81+
82+
/**
83+
* This will be invoked after the query execution.
84+
* You are expected to put validation logic in here.
85+
*
86+
* @param previousCommands The list of commands before the last generated one.
87+
* @param command The description of the command you just generated.
88+
* It also contains the context information you stored during command generation.
89+
* @param previousColumns The output schema of the original query (without last generated command).
90+
* It contains name and type information for each column, see {@link EsqlQueryGenerator.Column}
91+
* @param previousOutput The output of the original query (without last generated command), as a list (rows) of lists (columns) of values
92+
* @param columns The output schema of the full query (WITH last generated command).
93+
* @param output The output of the full query (WITH last generated command), as a list (rows) of lists (columns) of values
94+
* @return The result of the output validation. If the validation succeeds, you should return {@link CommandGenerator#VALIDATION_OK}.
95+
* Also, if for some reason you can't validate the output, just return {@link CommandGenerator#VALIDATION_OK}; for a command, having a generator without
96+
* validation is much better than having no generator at all.
97+
*/
98+
ValidationResult validateOutput(
99+
List<CommandDescription> previousCommands,
100+
CommandDescription command,
101+
List<EsqlQueryGenerator.Column> previousColumns,
102+
List<List<Object>> previousOutput,
103+
List<EsqlQueryGenerator.Column> columns,
104+
List<List<Object>> output
105+
);
106+
107+
static ValidationResult expectSameRowCount(
108+
List<CommandDescription> previousCommands,
109+
List<List<Object>> previousOutput,
110+
List<List<Object>> output
111+
) {
112+
113+
// ES|QL is quite non-deterministic in this sense, we can't guarantee it for now
114+
// if (output.size() != previousOutput.size()) {
115+
// return new ValidationResult(false, "Expecting [" + previousOutput.size() + "] rows, but got [" + output.size() + "]");
116+
// }
117+
118+
return VALIDATION_OK;
119+
}
120+
121+
static ValidationResult expectSameColumns(List<EsqlQueryGenerator.Column> previousColumns, List<EsqlQueryGenerator.Column> columns) {
122+
123+
if (previousColumns.stream().anyMatch(x -> x.name().contains("<all-fields-projected>"))) {
124+
return VALIDATION_OK; // known bug
125+
}
126+
127+
if (previousColumns.size() != columns.size()) {
128+
return new ValidationResult(false, "Expecting [" + previousColumns.size() + "] columns, got [" + columns.size() + "]");
129+
}
130+
131+
List<String> prevColNames = previousColumns.stream().map(EsqlQueryGenerator.Column::name).toList();
132+
List<String> newColNames = columns.stream().map(EsqlQueryGenerator.Column::name).toList();
133+
if (prevColNames.equals(newColNames) == false) {
134+
return new ValidationResult(
135+
false,
136+
"Expecting the following columns [" + String.join(", ", prevColNames) + "], got [" + String.join(", ", newColNames) + "]"
137+
);
138+
}
139+
140+
return VALIDATION_OK;
141+
}
142+
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
package org.elasticsearch.xpack.esql.qa.rest.generative.command.pipe;
9+
10+
import org.elasticsearch.xpack.esql.qa.rest.generative.EsqlQueryGenerator;
11+
import org.elasticsearch.xpack.esql.qa.rest.generative.command.CommandGenerator;
12+
13+
import java.util.List;
14+
import java.util.Map;
15+
16+
import static org.elasticsearch.test.ESTestCase.randomBoolean;
17+
import static org.elasticsearch.test.ESTestCase.randomIntBetween;
18+
19+
public class DissectGenerator implements CommandGenerator {
20+
21+
public static final String DISSECT = "dissect";
22+
public static final CommandGenerator INSTANCE = new DissectGenerator();
23+
24+
@Override
25+
public CommandDescription generate(
26+
List<CommandDescription> previousCommands,
27+
List<EsqlQueryGenerator.Column> previousOutput,
28+
QuerySchema schema
29+
) {
30+
String field = EsqlQueryGenerator.randomStringField(previousOutput);
31+
if (field == null) {
32+
return EMPTY_DESCRIPTION;// no strings to dissect, just skip
33+
}
34+
StringBuilder result = new StringBuilder(" | dissect ");
35+
result.append(field);
36+
result.append(" \"");
37+
for (int i = 0; i < randomIntBetween(1, 3); i++) {
38+
if (i > 0) {
39+
result.append(" ");
40+
}
41+
result.append("%{");
42+
String fieldName;
43+
if (randomBoolean()) {
44+
fieldName = EsqlQueryGenerator.randomIdentifier();
45+
} else {
46+
fieldName = EsqlQueryGenerator.randomRawName(previousOutput);
47+
if (fieldName == null) {
48+
fieldName = EsqlQueryGenerator.randomIdentifier();
49+
}
50+
}
51+
result.append(fieldName);
52+
result.append("}");
53+
}
54+
result.append("\"");
55+
String cmdString = result.toString();
56+
return new CommandDescription(DISSECT, this, cmdString, Map.of());
57+
}
58+
59+
@Override
60+
public ValidationResult validateOutput(
61+
List<CommandDescription> previousCommands,
62+
CommandDescription commandDescription,
63+
List<EsqlQueryGenerator.Column> previousColumns,
64+
List<List<Object>> previousOutput,
65+
List<EsqlQueryGenerator.Column> columns,
66+
List<List<Object>> output
67+
) {
68+
if (commandDescription == EMPTY_DESCRIPTION) {
69+
return VALIDATION_OK;
70+
}
71+
72+
if (previousColumns.size() > columns.size()) {
73+
return new ValidationResult(false, "Expecting at least [" + previousColumns.size() + "] columns, got [" + columns.size() + "]");
74+
}
75+
76+
return CommandGenerator.expectSameRowCount(previousCommands, previousOutput, output);
77+
}
78+
}

0 commit comments

Comments
 (0)