Skip to content

Commit 9db4361

Browse files
svilen-mihaylov-elasticsvilen-mihaylov-elastic
authored
ES|QL: Improve generative tests for FORK [130015] (#131206)
Addresses #130015
1 parent 0eca703 commit 9db4361

File tree

4 files changed

+198
-50
lines changed

4 files changed

+198
-50
lines changed

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

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,46 @@ public static CommandGenerator randomPipeCommandGenerator() {
7676
return randomFrom(PIPE_COMMANDS);
7777
}
7878

79+
public interface Executor {
80+
void run(CommandGenerator generator, CommandGenerator.CommandDescription current);
81+
82+
List<CommandGenerator.CommandDescription> previousCommands();
83+
84+
boolean continueExecuting();
85+
86+
List<EsqlQueryGenerator.Column> currentSchema();
87+
88+
}
89+
90+
public static void generatePipeline(
91+
final int depth,
92+
CommandGenerator commandGenerator,
93+
final CommandGenerator.QuerySchema schema,
94+
Executor executor
95+
) {
96+
CommandGenerator.CommandDescription desc = commandGenerator.generate(List.of(), List.of(), schema);
97+
executor.run(commandGenerator, desc);
98+
if (executor.continueExecuting() == false) {
99+
return;
100+
}
101+
102+
for (int j = 0; j < depth; j++) {
103+
if (executor.currentSchema().isEmpty()) {
104+
break;
105+
}
106+
commandGenerator = EsqlQueryGenerator.randomPipeCommandGenerator();
107+
desc = commandGenerator.generate(executor.previousCommands(), executor.currentSchema(), schema);
108+
if (desc == CommandGenerator.EMPTY_DESCRIPTION) {
109+
continue;
110+
}
111+
112+
executor.run(commandGenerator, desc);
113+
if (executor.continueExecuting() == false) {
114+
break;
115+
}
116+
}
117+
}
118+
79119
public static String booleanExpression(List<Column> previousOutput) {
80120
// TODO LIKE, RLIKE, functions etc.
81121
return switch (randomIntBetween(0, 3)) {

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

Lines changed: 44 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -87,43 +87,53 @@ public void test() throws IOException {
8787
List<LookupIdx> lookupIndices = lookupIndices();
8888
List<CsvTestsDataLoader.EnrichConfig> policies = availableEnrichPolicies();
8989
CommandGenerator.QuerySchema mappingInfo = new CommandGenerator.QuerySchema(indices, lookupIndices, policies);
90-
EsqlQueryGenerator.QueryExecuted previousResult = null;
90+
9191
for (int i = 0; i < ITERATIONS; i++) {
92-
List<CommandGenerator.CommandDescription> previousCommands = new ArrayList<>();
93-
CommandGenerator commandGenerator = EsqlQueryGenerator.sourceCommand();
94-
CommandGenerator.CommandDescription desc = commandGenerator.generate(List.of(), List.of(), mappingInfo);
95-
String command = desc.commandString();
96-
EsqlQueryGenerator.QueryExecuted result = execute(command, 0);
97-
if (result.exception() != null) {
98-
checkException(result);
99-
continue;
100-
}
101-
if (checkResults(List.of(), commandGenerator, desc, null, result).success() == false) {
102-
continue;
103-
}
104-
previousResult = result;
105-
previousCommands.add(desc);
106-
for (int j = 0; j < MAX_DEPTH; j++) {
107-
if (result.outputSchema().isEmpty()) {
108-
break;
92+
var exec = new EsqlQueryGenerator.Executor() {
93+
@Override
94+
public void run(CommandGenerator generator, CommandGenerator.CommandDescription current) {
95+
previousCommands.add(current);
96+
final String command = current.commandString();
97+
98+
final EsqlQueryGenerator.QueryExecuted result = previousResult == null
99+
? execute(command, 0)
100+
: execute(previousResult.query() + command, previousResult.depth());
101+
previousResult = result;
102+
103+
final boolean hasException = result.exception() != null;
104+
if (hasException || checkResults(List.of(), generator, current, previousResult, result).success() == false) {
105+
if (hasException) {
106+
checkException(result);
107+
}
108+
continueExecuting = false;
109+
currentSchema = List.of();
110+
} else {
111+
continueExecuting = true;
112+
currentSchema = result.outputSchema();
113+
}
109114
}
110-
commandGenerator = EsqlQueryGenerator.randomPipeCommandGenerator();
111-
desc = commandGenerator.generate(previousCommands, result.outputSchema(), mappingInfo);
112-
if (desc == CommandGenerator.EMPTY_DESCRIPTION) {
113-
continue;
115+
116+
@Override
117+
public List<CommandGenerator.CommandDescription> previousCommands() {
118+
return previousCommands;
114119
}
115-
command = desc.commandString();
116-
result = execute(result.query() + command, result.depth() + 1);
117-
if (result.exception() != null) {
118-
checkException(result);
119-
break;
120+
121+
@Override
122+
public boolean continueExecuting() {
123+
return continueExecuting;
120124
}
121-
if (checkResults(previousCommands, commandGenerator, desc, previousResult, result).success() == false) {
122-
break;
125+
126+
@Override
127+
public List<EsqlQueryGenerator.Column> currentSchema() {
128+
return currentSchema;
123129
}
124-
previousCommands.add(desc);
125-
previousResult = result;
126-
}
130+
131+
boolean continueExecuting;
132+
List<EsqlQueryGenerator.Column> currentSchema;
133+
final List<CommandGenerator.CommandDescription> previousCommands = new ArrayList<>();
134+
EsqlQueryGenerator.QueryExecuted previousResult;
135+
};
136+
EsqlQueryGenerator.generatePipeline(MAX_DEPTH, EsqlQueryGenerator.sourceCommand(), mappingInfo, exec);
127137
}
128138
}
129139

@@ -163,7 +173,7 @@ private void checkException(EsqlQueryGenerator.QueryExecuted query) {
163173
}
164174

165175
@SuppressWarnings("unchecked")
166-
private EsqlQueryGenerator.QueryExecuted execute(String command, int depth) {
176+
public static EsqlQueryGenerator.QueryExecuted execute(String command, int depth) {
167177
try {
168178
Map<String, Object> a = RestEsqlTestCase.runEsql(
169179
new RestEsqlTestCase.RequestObjectBuilder().query(command).build(),
@@ -183,7 +193,7 @@ private EsqlQueryGenerator.QueryExecuted execute(String command, int depth) {
183193
}
184194

185195
@SuppressWarnings("unchecked")
186-
private List<EsqlQueryGenerator.Column> outputSchema(Map<String, Object> a) {
196+
private static List<EsqlQueryGenerator.Column> outputSchema(Map<String, Object> a) {
187197
List<Map<String, String>> cols = (List<Map<String, String>>) a.get("columns");
188198
if (cols == null) {
189199
return null;

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

Lines changed: 95 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,10 +8,13 @@
88
package org.elasticsearch.xpack.esql.qa.rest.generative.command.pipe;
99

1010
import org.elasticsearch.xpack.esql.qa.rest.generative.EsqlQueryGenerator;
11+
import org.elasticsearch.xpack.esql.qa.rest.generative.GenerativeRestTest;
1112
import org.elasticsearch.xpack.esql.qa.rest.generative.command.CommandGenerator;
1213

14+
import java.util.ArrayList;
1315
import java.util.List;
1416
import java.util.Map;
17+
import java.util.stream.Collectors;
1518

1619
import static org.elasticsearch.test.ESTestCase.randomIntBetween;
1720

@@ -28,17 +31,105 @@ public CommandDescription generate(
2831
) {
2932
// FORK can only be allowed once - so we skip adding another FORK if we already have one
3033
// otherwise, most generated queries would only result in a validation error
34+
StringBuilder completeCommand = new StringBuilder();
3135
for (CommandDescription command : previousCommands) {
3236
if (command.commandName().equals(FORK)) {
33-
return new CommandDescription(FORK, this, " ", Map.of());
37+
return EMPTY_DESCRIPTION;
3438
}
39+
40+
completeCommand.append(command.commandString());
3541
}
3642

37-
int n = randomIntBetween(2, 3);
43+
final int branchCount = randomIntBetween(2, 3);
44+
final int branchToRetain = randomIntBetween(1, branchCount);
45+
46+
StringBuilder forkCmd = new StringBuilder(" | FORK ");
47+
for (int i = 0; i < branchCount; i++) {
48+
var expr = WhereGenerator.randomExpression(randomIntBetween(1, 2), previousOutput);
49+
if (expr == null) {
50+
expr = "true";
51+
}
52+
forkCmd.append(" (").append("where ").append(expr);
53+
54+
var exec = new EsqlQueryGenerator.Executor() {
55+
@Override
56+
public void run(CommandGenerator generator, CommandDescription current) {
57+
final String command = current.commandString();
58+
59+
// Try appending new command to parent of Fork. If we successfully execute (without exception) AND still retain the same
60+
// schema, we append the command. Enforcing the same schema is stricter than the Fork needs (it only needs types to be
61+
// the same on columns which are present), but given we currently generate independent sub-pipelines, this way we can
62+
// generate more valid Fork queries.
63+
final EsqlQueryGenerator.QueryExecuted result = previousResult == null
64+
? GenerativeRestTest.execute(command, 0)
65+
: GenerativeRestTest.execute(previousResult.query() + command, previousResult.depth());
66+
previousResult = result;
67+
68+
continueExecuting = result.exception() == null && result.outputSchema().equals(previousOutput);
69+
if (continueExecuting) {
70+
previousCommands.add(current);
71+
}
72+
}
73+
74+
@Override
75+
public List<CommandDescription> previousCommands() {
76+
return previousCommands;
77+
}
78+
79+
@Override
80+
public boolean continueExecuting() {
81+
return continueExecuting;
82+
}
3883

39-
String cmd = " | FORK " + "( WHERE true ) ".repeat(n) + " | WHERE _fork == \"fork" + randomIntBetween(1, n) + "\" | DROP _fork";
84+
@Override
85+
public List<EsqlQueryGenerator.Column> currentSchema() {
86+
return previousOutput;
87+
}
88+
89+
final List<CommandGenerator.CommandDescription> previousCommands = new ArrayList<>();
90+
boolean continueExecuting;
91+
EsqlQueryGenerator.QueryExecuted previousResult;
92+
};
93+
94+
var gen = new CommandGenerator() {
95+
@Override
96+
public CommandDescription generate(
97+
List<CommandDescription> previousCommands,
98+
List<EsqlQueryGenerator.Column> previousOutput,
99+
QuerySchema schema
100+
) {
101+
return new CommandDescription(FORK, this, completeCommand.toString(), Map.of());
102+
}
103+
104+
@Override
105+
public ValidationResult validateOutput(
106+
List<CommandDescription> previousCommands,
107+
CommandDescription command,
108+
List<EsqlQueryGenerator.Column> previousColumns,
109+
List<List<Object>> previousOutput,
110+
List<EsqlQueryGenerator.Column> columns,
111+
List<List<Object>> output
112+
) {
113+
return VALIDATION_OK;
114+
}
115+
};
116+
117+
EsqlQueryGenerator.generatePipeline(3, gen, schema, exec);
118+
if (exec.previousCommands().size() > 1) {
119+
String previousCmd = exec.previousCommands()
120+
.stream()
121+
.skip(1)
122+
.map(CommandDescription::commandString)
123+
.collect(Collectors.joining(" "));
124+
forkCmd.append(previousCmd);
125+
}
126+
127+
forkCmd.append(")");
128+
}
129+
forkCmd.append(" | WHERE _fork == \"fork").append(branchToRetain).append("\" | DROP _fork");
40130

41-
return new CommandDescription(FORK, this, cmd, Map.of());
131+
// System.out.println("Generated fork command: " + forkCmd);
132+
return new CommandDescription(FORK, this, forkCmd.toString(), Map.of());
42133
}
43134

44135
@Override

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

Lines changed: 19 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -21,20 +21,15 @@ public class WhereGenerator implements CommandGenerator {
2121
public static final String WHERE = "where";
2222
public static final CommandGenerator INSTANCE = new WhereGenerator();
2323

24-
@Override
25-
public CommandDescription generate(
26-
List<CommandDescription> previousCommands,
27-
List<EsqlQueryGenerator.Column> previousOutput,
28-
QuerySchema schema
29-
) {
24+
public static String randomExpression(final int nConditions, List<EsqlQueryGenerator.Column> previousOutput) {
3025
// TODO more complex conditions
31-
StringBuilder result = new StringBuilder(" | where ");
32-
int nConditions = randomIntBetween(1, 5);
26+
var result = new StringBuilder();
27+
3328
for (int i = 0; i < nConditions; i++) {
3429
String exp = EsqlQueryGenerator.booleanExpression(previousOutput);
3530
if (exp == null) {
36-
// cannot generate expressions, just skip
37-
return EMPTY_DESCRIPTION;
31+
// Cannot generate expressions, just skip.
32+
return null;
3833
}
3934
if (i > 0) {
4035
result.append(randomBoolean() ? " AND " : " OR ");
@@ -45,8 +40,20 @@ public CommandDescription generate(
4540
result.append(exp);
4641
}
4742

48-
String cmd = result.toString();
49-
return new CommandDescription(WHERE, this, cmd, Map.of());
43+
return result.toString();
44+
}
45+
46+
@Override
47+
public CommandDescription generate(
48+
List<CommandDescription> previousCommands,
49+
List<EsqlQueryGenerator.Column> previousOutput,
50+
QuerySchema schema
51+
) {
52+
String expression = randomExpression(randomIntBetween(1, 5), previousOutput);
53+
if (expression == null) {
54+
return EMPTY_DESCRIPTION;
55+
}
56+
return new CommandDescription(WHERE, this, " | where " + expression, Map.of());
5057
}
5158

5259
@Override

0 commit comments

Comments
 (0)