Skip to content

Commit 82296eb

Browse files
committed
[ES|QL] Add a TS variation of GenerativeIT
This commit adds a variation of the GenerativeRestTest in ES|QL that queries indices marked as time series indices and runs time series aggregations on them (amongst all the other commands already supported in the Generative tests)
1 parent 9fe97ab commit 82296eb

File tree

7 files changed

+353
-9
lines changed

7 files changed

+353
-9
lines changed
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
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.single_node;
9+
10+
import com.carrotsearch.randomizedtesting.annotations.ThreadLeakFilters;
11+
12+
import org.elasticsearch.test.TestClustersThreadFilter;
13+
import org.elasticsearch.test.cluster.ElasticsearchCluster;
14+
import org.elasticsearch.xpack.esql.qa.rest.generative.EsqlQueryGenerator;
15+
import org.elasticsearch.xpack.esql.qa.rest.generative.GenerativeRestTest;
16+
import org.elasticsearch.xpack.esql.qa.rest.generative.command.CommandGenerator;
17+
import org.junit.ClassRule;
18+
19+
@ThreadLeakFilters(filters = TestClustersThreadFilter.class)
20+
public class GenerativeMetricsIT extends GenerativeRestTest {
21+
@ClassRule
22+
public static ElasticsearchCluster cluster = Clusters.testCluster();
23+
24+
@Override
25+
protected String getTestRestCluster() {
26+
return cluster.getHttpAddresses();
27+
}
28+
29+
@Override
30+
protected boolean supportsSourceFieldMapping() {
31+
return cluster.getNumNodes() == 1;
32+
}
33+
34+
@Override
35+
protected CommandGenerator sourceCommand() {
36+
return EsqlQueryGenerator.timeSeriesSourceCommand();
37+
}
38+
39+
@Override
40+
protected boolean requiresTimeSeries() {
41+
return true;
42+
}
43+
}

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

Lines changed: 117 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,12 +19,14 @@
1919
import org.elasticsearch.xpack.esql.qa.rest.generative.command.pipe.KeepGenerator;
2020
import org.elasticsearch.xpack.esql.qa.rest.generative.command.pipe.LimitGenerator;
2121
import org.elasticsearch.xpack.esql.qa.rest.generative.command.pipe.LookupJoinGenerator;
22+
import org.elasticsearch.xpack.esql.qa.rest.generative.command.pipe.MetricsStatsGenerator;
2223
import org.elasticsearch.xpack.esql.qa.rest.generative.command.pipe.MvExpandGenerator;
2324
import org.elasticsearch.xpack.esql.qa.rest.generative.command.pipe.RenameGenerator;
2425
import org.elasticsearch.xpack.esql.qa.rest.generative.command.pipe.SortGenerator;
2526
import org.elasticsearch.xpack.esql.qa.rest.generative.command.pipe.StatsGenerator;
2627
import org.elasticsearch.xpack.esql.qa.rest.generative.command.pipe.WhereGenerator;
2728
import org.elasticsearch.xpack.esql.qa.rest.generative.command.source.FromGenerator;
29+
import org.elasticsearch.xpack.esql.qa.rest.generative.command.source.MetricGenerator;
2830

2931
import java.util.List;
3032
import java.util.Set;
@@ -51,6 +53,11 @@ public record QueryExecuted(String query, int depth, List<Column> outputSchema,
5153
*/
5254
static List<CommandGenerator> SOURCE_COMMANDS = List.of(FromGenerator.INSTANCE);
5355

56+
/**
57+
* Commands at the beginning of queries that begin queries on time series indices, eg. TS
58+
*/
59+
static List<CommandGenerator> TIME_SERIES_SOURCE_COMMANDS = List.of(MetricGenerator.INSTANCE);
60+
5461
/**
5562
* These are downstream commands, ie. that cannot appear as the first command in a query
5663
*/
@@ -72,14 +79,42 @@ public record QueryExecuted(String query, int depth, List<Column> outputSchema,
7279
WhereGenerator.INSTANCE
7380
);
7481

82+
static List<CommandGenerator> TIME_SERIES_PIPE_COMMANDS = List.of(
83+
ChangePointGenerator.INSTANCE,
84+
DissectGenerator.INSTANCE,
85+
DropGenerator.INSTANCE,
86+
EnrichGenerator.INSTANCE,
87+
EvalGenerator.INSTANCE,
88+
ForkGenerator.INSTANCE,
89+
GrokGenerator.INSTANCE,
90+
KeepGenerator.INSTANCE,
91+
LimitGenerator.INSTANCE,
92+
LookupJoinGenerator.INSTANCE,
93+
MetricsStatsGenerator.INSTANCE,
94+
MvExpandGenerator.INSTANCE,
95+
RenameGenerator.INSTANCE,
96+
SortGenerator.INSTANCE,
97+
StatsGenerator.INSTANCE,
98+
WhereGenerator.INSTANCE
99+
);
100+
75101
public static CommandGenerator sourceCommand() {
76102
return randomFrom(SOURCE_COMMANDS);
77103
}
78104

105+
public static CommandGenerator timeSeriesSourceCommand() {
106+
return randomFrom(TIME_SERIES_SOURCE_COMMANDS);
107+
}
108+
79109
public static CommandGenerator randomPipeCommandGenerator() {
80110
return randomFrom(PIPE_COMMANDS);
81111
}
82112

113+
public static CommandGenerator randomMetricsPipeCommandGenerator() {
114+
// todo better way
115+
return randomFrom(TIME_SERIES_PIPE_COMMANDS);
116+
}
117+
83118
public interface Executor {
84119
void run(CommandGenerator generator, CommandGenerator.CommandDescription current);
85120

@@ -95,7 +130,8 @@ public static void generatePipeline(
95130
final int depth,
96131
CommandGenerator commandGenerator,
97132
final CommandGenerator.QuerySchema schema,
98-
Executor executor
133+
Executor executor,
134+
boolean isTimeSeries
99135
) {
100136
CommandGenerator.CommandDescription desc = commandGenerator.generate(List.of(), List.of(), schema);
101137
executor.run(commandGenerator, desc);
@@ -107,7 +143,7 @@ public static void generatePipeline(
107143
if (executor.currentSchema().isEmpty()) {
108144
break;
109145
}
110-
commandGenerator = EsqlQueryGenerator.randomPipeCommandGenerator();
146+
commandGenerator = isTimeSeries ? randomMetricsPipeCommandGenerator() : EsqlQueryGenerator.randomPipeCommandGenerator();
111147
desc = commandGenerator.generate(executor.previousCommands(), executor.currentSchema(), schema);
112148
if (desc == CommandGenerator.EMPTY_DESCRIPTION) {
113149
continue;
@@ -217,6 +253,61 @@ public static boolean sortable(Column col) {
217253
|| col.type.equals("version");
218254
}
219255

256+
public static String metricsAgg(List<Column> previousOutput) {
257+
String outerCommand = randomFrom("min", "max", "sum", "count", "avg");
258+
String innerCommand = switch (randomIntBetween(0, 3)) {
259+
case 0 -> {
260+
// input can be numerics + aggregate_metric_double
261+
String numericPlusAggMetricFieldName = randomMetricsNumericField(previousOutput);
262+
if (numericPlusAggMetricFieldName == null) {
263+
yield null;
264+
}
265+
yield switch ((randomIntBetween(0, 3))) {
266+
case 0 -> "max_over_time(" + numericPlusAggMetricFieldName + ")";
267+
case 1 -> "min_over_time(" + numericPlusAggMetricFieldName + ")";
268+
case 2 -> "sum_over_time(" + numericPlusAggMetricFieldName + ")";
269+
default -> "avg_over_time(" + numericPlusAggMetricFieldName + ")";
270+
};
271+
}
272+
case 1 -> {
273+
// input can be a counter
274+
String counterField = randomCounterField(previousOutput);
275+
if (counterField == null) {
276+
yield null;
277+
}
278+
yield "rate(" + counterField + ")";
279+
}
280+
case 2 -> {
281+
// numerics except aggregate_metric_double
282+
// TODO: move to case 0 when support for aggregate_metric_double is added to these functions
283+
String numericFieldName = randomNumericField(previousOutput);
284+
if (numericFieldName == null) {
285+
yield null;
286+
}
287+
yield (randomBoolean() ? "first_over_time(" : "last_over_time(") + numericFieldName + ")";
288+
}
289+
default -> {
290+
// TODO: add other types that count_over_time supports
291+
String otherFieldName = randomBoolean() ? randomStringField(previousOutput) : randomNumericOrDateField(previousOutput);
292+
if (otherFieldName == null) {
293+
yield null;
294+
}
295+
if (randomBoolean()) {
296+
yield "count_over_time(" + otherFieldName + ")";
297+
} else {
298+
yield "count_distinct_over_time(" + otherFieldName + ")";
299+
// TODO: replace with the below
300+
// yield "count_distinct_over_time(" + otherFieldName + (randomBoolean() ? ", " + randomNonNegativeInt() : "") + ")";
301+
}
302+
}
303+
};
304+
if (innerCommand == null) {
305+
// TODO: figure out a default that maybe makes more sense than using a timestamp field
306+
innerCommand = "count_over_time(" + randomDateField(previousOutput) + ")";
307+
}
308+
return outerCommand + "(" + innerCommand + ")";
309+
}
310+
220311
public static String agg(List<Column> previousOutput) {
221312
String name = randomNumericOrDateField(previousOutput);
222313
if (name != null && randomBoolean()) {
@@ -251,6 +342,30 @@ public static String randomNumericField(List<Column> previousOutput) {
251342
return randomName(previousOutput, Set.of("long", "integer", "double"));
252343
}
253344

345+
public static String randomMetricsNumericField(List<Column> previousOutput) {
346+
Set<String> allowedTypes = Set.of("double", "long", "unsigned_long", "integer", "aggregate_metric_double");
347+
List<String> items = previousOutput.stream()
348+
.filter(
349+
x -> allowedTypes.contains(x.type())
350+
|| (x.type().equals("unsupported") && canBeCastedToAggregateMetricDouble(x.originalTypes()))
351+
)
352+
.map(Column::name)
353+
.toList();
354+
if (items.isEmpty()) {
355+
return null;
356+
}
357+
return items.get(randomIntBetween(0, items.size() - 1));
358+
}
359+
360+
public static String randomCounterField(List<Column> previousOutput) {
361+
return randomName(previousOutput, Set.of("counter_long", "counter_double", "counter_integer"));
362+
}
363+
364+
private static boolean canBeCastedToAggregateMetricDouble(List<String> types) {
365+
return types.contains("aggregate_metric_double")
366+
&& Set.of("double", "long", "unsigned_long", "integer", "aggregate_metric_double").containsAll(types);
367+
}
368+
254369
public static String randomStringField(List<Column> previousOutput) {
255370
return randomName(previousOutput, Set.of("text", "keyword"));
256371
}

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

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,10 @@ public void setup() throws IOException {
7979

8080
protected abstract boolean supportsSourceFieldMapping();
8181

82+
protected boolean requiresTimeSeries() {
83+
return false;
84+
}
85+
8286
@AfterClass
8387
public static void wipeTestData() throws IOException {
8488
try {
@@ -142,10 +146,14 @@ public List<EsqlQueryGenerator.Column> currentSchema() {
142146
final List<CommandGenerator.CommandDescription> previousCommands = new ArrayList<>();
143147
EsqlQueryGenerator.QueryExecuted previousResult;
144148
};
145-
EsqlQueryGenerator.generatePipeline(MAX_DEPTH, EsqlQueryGenerator.sourceCommand(), mappingInfo, exec);
149+
EsqlQueryGenerator.generatePipeline(MAX_DEPTH, sourceCommand(), mappingInfo, exec, requiresTimeSeries());
146150
}
147151
}
148152

153+
protected CommandGenerator sourceCommand() {
154+
return EsqlQueryGenerator.sourceCommand();
155+
}
156+
149157
private static CommandGenerator.ValidationResult checkResults(
150158
List<CommandGenerator.CommandDescription> previousCommands,
151159
CommandGenerator commandGenerator,
@@ -234,7 +242,7 @@ private static List<String> originalTypes(Map<String, ?> x) {
234242
}
235243

236244
private List<String> availableIndices() throws IOException {
237-
return availableDatasetsForEs(true, supportsSourceFieldMapping(), false).stream()
245+
return availableDatasetsForEs(true, supportsSourceFieldMapping(), false, requiresTimeSeries()).stream()
238246
.filter(x -> x.requiresInferenceEndpoint() == false)
239247
.map(x -> x.indexName())
240248
.toList();

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -114,7 +114,7 @@ public ValidationResult validateOutput(
114114
}
115115
};
116116

117-
EsqlQueryGenerator.generatePipeline(3, gen, schema, exec);
117+
EsqlQueryGenerator.generatePipeline(3, gen, schema, exec, false);
118118
if (exec.previousCommands().size() > 1) {
119119
String previousCmd = exec.previousCommands()
120120
.stream()
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
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+
import java.util.stream.Collectors;
16+
17+
import static org.elasticsearch.test.ESTestCase.randomBoolean;
18+
import static org.elasticsearch.test.ESTestCase.randomIntBetween;
19+
import static org.elasticsearch.xpack.esql.qa.rest.generative.EsqlQueryGenerator.randomDateField;
20+
21+
public class MetricsStatsGenerator implements CommandGenerator {
22+
23+
public static final String STATS = "stats";
24+
public static final CommandGenerator INSTANCE = new MetricsStatsGenerator();
25+
26+
@Override
27+
public CommandDescription generate(
28+
List<CommandDescription> previousCommands,
29+
List<EsqlQueryGenerator.Column> previousOutput,
30+
QuerySchema schema
31+
) {
32+
// generates stats in the form of:
33+
// `STATS some_aggregation(some_field) by optional_grouping_field, non_optional = bucket(time_field, 5minute)`
34+
// where `some_aggregation` can be a time series aggregation, or a regular aggregation
35+
// There is a variable number of aggregations per command
36+
37+
List<EsqlQueryGenerator.Column> nonNull = previousOutput.stream()
38+
.filter(EsqlQueryGenerator::fieldCanBeUsed)
39+
.filter(x -> x.type().equals("null") == false)
40+
.collect(Collectors.toList());
41+
if (nonNull.isEmpty()) {
42+
return EMPTY_DESCRIPTION;
43+
}
44+
String timestamp = randomDateField(nonNull);
45+
// if there's no timestamp field left, there's nothing to bucket on
46+
if (timestamp == null) {
47+
return EMPTY_DESCRIPTION;
48+
}
49+
50+
StringBuilder cmd = new StringBuilder(" | stats ");
51+
52+
// TODO: increase range max to 5
53+
int nStats = randomIntBetween(1, 2);
54+
for (int i = 0; i < nStats; i++) {
55+
String name;
56+
if (randomBoolean()) {
57+
name = EsqlQueryGenerator.randomIdentifier();
58+
} else {
59+
name = EsqlQueryGenerator.randomName(previousOutput);
60+
if (name == null) {
61+
name = EsqlQueryGenerator.randomIdentifier();
62+
}
63+
}
64+
// generate the aggregation
65+
String expression = randomBoolean() ? EsqlQueryGenerator.metricsAgg(nonNull) : EsqlQueryGenerator.agg(nonNull);
66+
if (i > 0) {
67+
cmd.append(",");
68+
}
69+
cmd.append(" ");
70+
cmd.append(name);
71+
cmd.append(" = ");
72+
cmd.append(expression);
73+
}
74+
75+
cmd.append(" by ");
76+
if (randomBoolean()) {
77+
var col = EsqlQueryGenerator.randomGroupableName(nonNull);
78+
if (col != null) {
79+
cmd.append(col + ", ");
80+
}
81+
}
82+
// TODO: add alternative time buckets
83+
cmd.append(
84+
(randomBoolean() ? EsqlQueryGenerator.randomIdentifier() : EsqlQueryGenerator.randomName(previousOutput))
85+
+ " = bucket("
86+
+ timestamp
87+
+ ",1hour)"
88+
);
89+
return new CommandDescription(STATS, this, cmd.toString(), Map.of());
90+
}
91+
92+
@Override
93+
public ValidationResult validateOutput(
94+
List<CommandDescription> previousCommands,
95+
CommandDescription commandDescription,
96+
List<EsqlQueryGenerator.Column> previousColumns,
97+
List<List<Object>> previousOutput,
98+
List<EsqlQueryGenerator.Column> columns,
99+
List<List<Object>> output
100+
) {
101+
// TODO validate columns
102+
return VALIDATION_OK;
103+
}
104+
}

0 commit comments

Comments
 (0)