Skip to content

ESQL: Speed up spec tests #131949

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 12 commits into from
Jul 30, 2025
Merged

ESQL: Speed up spec tests #131949

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
79ebb5e
ESQL: Check speed up spec tests
nik9000 Jul 25, 2025
75e4f98
[CI] Auto commit changes from spotless
Jul 25, 2025
c09517b
Fix
nik9000 Jul 27, 2025
8a5fae6
Merge remote-tracking branch 'nik9000/tet' into tet
nik9000 Jul 27, 2025
2d6d9cc
Update x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/x…
nik9000 Jul 28, 2025
dbdc721
Fixup
nik9000 Jul 28, 2025
4a1db88
Merge branch 'main' into tet
nik9000 Jul 28, 2025
c4e8750
Merge branch 'main' into tet
nik9000 Jul 30, 2025
2b86b14
Switch to admin client
nik9000 Jul 30, 2025
7f41613
Fix
nik9000 Jul 30, 2025
e81956a
Cache better
nik9000 Jul 30, 2025
e8d54fa
[CI] Auto commit changes from spotless
Jul 30, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 2 additions & 9 deletions ...r/src/javaRestTest/java/org/elasticsearch/xpack/esql/qa/mixed/MixedClusterEsqlSpecIT.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,6 @@
import org.elasticsearch.test.rest.TestFeatureService;
import org.elasticsearch.xpack.esql.CsvSpecReader.CsvTestCase;
import org.elasticsearch.xpack.esql.qa.rest.EsqlSpecTestCase;
import org.elasticsearch.xpack.esql.qa.rest.RestEsqlTestCase.Mode;
import org.junit.AfterClass;
import org.junit.Before;
import org.junit.ClassRule;
Expand Down Expand Up @@ -43,11 +42,6 @@ public void extractOldClusterFeatures() {
}
}

protected static boolean oldClusterHasFeature(String featureId) {
assert oldClusterTestFeatureService != null;
return oldClusterTestFeatureService.clusterHasFeature(featureId);
}
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Unused.


@AfterClass
public static void cleanUp() {
oldClusterTestFeatureService = null;
Expand All @@ -59,10 +53,9 @@ public MixedClusterEsqlSpecIT(
String testName,
Integer lineNumber,
CsvTestCase testCase,
String instructions,
Mode mode
String instructions
) {
super(fileName, groupName, testName, lineNumber, testCase, instructions, mode);
super(fileName, groupName, testName, lineNumber, testCase, instructions);
}

@Override
Expand Down
22 changes: 3 additions & 19 deletions ...i-clusters/src/javaRestTest/java/org/elasticsearch/xpack/esql/ccq/MultiClusterSpecIT.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,6 @@
import org.elasticsearch.xpack.esql.CsvTestsDataLoader;
import org.elasticsearch.xpack.esql.SpecReader;
import org.elasticsearch.xpack.esql.qa.rest.EsqlSpecTestCase;
import org.elasticsearch.xpack.esql.qa.rest.RestEsqlTestCase.Mode;
import org.junit.AfterClass;
import org.junit.ClassRule;
import org.junit.rules.RuleChain;
Expand All @@ -36,7 +35,6 @@
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.net.URL;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Locale;
Expand All @@ -59,7 +57,6 @@
import static org.elasticsearch.xpack.esql.action.EsqlCapabilities.Cap.JOIN_PLANNING_V1;
import static org.elasticsearch.xpack.esql.action.EsqlCapabilities.Cap.METADATA_FIELDS_REMOTE_TEST;
import static org.elasticsearch.xpack.esql.action.EsqlCapabilities.Cap.UNMAPPED_FIELDS;
import static org.elasticsearch.xpack.esql.qa.rest.RestEsqlTestCase.Mode.SYNC;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.mock;
Expand Down Expand Up @@ -87,19 +84,7 @@ public class MultiClusterSpecIT extends EsqlSpecTestCase {
public static List<Object[]> readScriptSpec() throws Exception {
List<URL> urls = classpathResources("/*.csv-spec");
assertTrue("Not enough specs found " + urls, urls.size() > 0);
List<Object[]> specs = SpecReader.readScriptSpec(urls, specParser());

int len = specs.get(0).length;
List<Object[]> testcases = new ArrayList<>();
for (var spec : specs) {
for (Mode mode : List.of(SYNC)) { // No async, for now
Object[] obj = new Object[len + 1];
System.arraycopy(spec, 0, obj, 0, len);
obj[len] = mode;
testcases.add(obj);
}
}
return testcases;
return SpecReader.readScriptSpec(urls, specParser());
}

public MultiClusterSpecIT(
Expand All @@ -108,10 +93,9 @@ public MultiClusterSpecIT(
String testName,
Integer lineNumber,
CsvTestCase testCase,
String instructions,
Mode mode
String instructions
) {
super(fileName, groupName, testName, lineNumber, convertToRemoteIndices(testCase), instructions, mode);
super(fileName, groupName, testName, lineNumber, convertToRemoteIndices(testCase), instructions);
}

// TODO: think how to handle this better
Expand Down
14 changes: 3 additions & 11 deletions ...lti-node/src/javaRestTest/java/org/elasticsearch/xpack/esql/qa/multi_node/EsqlSpecIT.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,9 +8,9 @@
package org.elasticsearch.xpack.esql.qa.multi_node;

import org.elasticsearch.test.cluster.ElasticsearchCluster;
import org.elasticsearch.test.junit.annotations.TestLogging;
import org.elasticsearch.xpack.esql.CsvSpecReader.CsvTestCase;
import org.elasticsearch.xpack.esql.qa.rest.EsqlSpecTestCase;
import org.elasticsearch.xpack.esql.qa.rest.RestEsqlTestCase.Mode;
import org.junit.ClassRule;

import java.io.IOException;
Expand All @@ -24,16 +24,8 @@ protected String getTestRestCluster() {
return cluster.getHttpAddresses();
}

public EsqlSpecIT(
String fileName,
String groupName,
String testName,
Integer lineNumber,
CsvTestCase testCase,
String instructions,
Mode mode
) {
super(fileName, groupName, testName, lineNumber, testCase, instructions, mode);
public EsqlSpecIT(String fileName, String groupName, String testName, Integer lineNumber, CsvTestCase testCase, String instructions) {
super(fileName, groupName, testName, lineNumber, testCase, instructions);
}

@Override
Expand Down
13 changes: 2 additions & 11 deletions ...le-node/src/javaRestTest/java/org/elasticsearch/xpack/esql/qa/single_node/EsqlSpecIT.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,6 @@
import org.elasticsearch.xpack.esql.planner.PhysicalSettings;
import org.elasticsearch.xpack.esql.plugin.ComputeService;
import org.elasticsearch.xpack.esql.qa.rest.EsqlSpecTestCase;
import org.elasticsearch.xpack.esql.qa.rest.RestEsqlTestCase.Mode;
import org.junit.Before;
import org.junit.ClassRule;

Expand All @@ -37,16 +36,8 @@ protected String getTestRestCluster() {
return cluster.getHttpAddresses();
}

public EsqlSpecIT(
String fileName,
String groupName,
String testName,
Integer lineNumber,
CsvTestCase testCase,
String instructions,
Mode mode
) {
super(fileName, groupName, testName, lineNumber, testCase, instructions, mode);
public EsqlSpecIT(String fileName, String groupName, String testName, Integer lineNumber, CsvTestCase testCase, String instructions) {
super(fileName, groupName, testName, lineNumber, testCase, instructions);
}

@Override
Expand Down
6 changes: 2 additions & 4 deletions ...e/src/javaRestTest/java/org/elasticsearch/xpack/esql/qa/single_node/GenerativeForkIT.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,6 @@
import org.elasticsearch.test.TestClustersThreadFilter;
import org.elasticsearch.test.cluster.ElasticsearchCluster;
import org.elasticsearch.xpack.esql.CsvSpecReader;
import org.elasticsearch.xpack.esql.qa.rest.RestEsqlTestCase.Mode;
import org.elasticsearch.xpack.esql.qa.rest.generative.GenerativeForkRestTest;
import org.junit.ClassRule;

Expand All @@ -32,10 +31,9 @@ public GenerativeForkIT(
String testName,
Integer lineNumber,
CsvSpecReader.CsvTestCase testCase,
String instructions,
Mode mode
String instructions
) {
super(fileName, groupName, testName, lineNumber, testCase, instructions, mode);
super(fileName, groupName, testName, lineNumber, testCase, instructions);
}

@Override
Expand Down
80 changes: 20 additions & 60 deletions ...n/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/EsqlSpecTestCase.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,6 @@
import static org.elasticsearch.xpack.esql.CsvTestUtils.ExpectedResults;
import static org.elasticsearch.xpack.esql.CsvTestUtils.isEnabled;
import static org.elasticsearch.xpack.esql.CsvTestUtils.loadCsvSpecValues;
import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.availableDatasetsForEs;
import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.createInferenceEndpoints;
import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.deleteInferenceEndpoints;
import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.loadDataSetIntoEs;
Expand All @@ -79,6 +78,7 @@
import static org.elasticsearch.xpack.esql.action.EsqlCapabilities.Cap.SEMANTIC_TEXT_FIELD_CAPS;
import static org.elasticsearch.xpack.esql.action.EsqlCapabilities.Cap.SOURCE_FIELD_MAPPING;
import static org.elasticsearch.xpack.esql.qa.rest.RestEsqlTestCase.assertNotPartial;
import static org.elasticsearch.xpack.esql.qa.rest.RestEsqlTestCase.hasCapabilities;

// This test can run very long in serverless configurations
@TimeoutSuite(millis = 30 * TimeUnits.MINUTE)
Expand All @@ -101,19 +101,7 @@ public abstract class EsqlSpecTestCase extends ESRestTestCase {
public static List<Object[]> readScriptSpec() throws Exception {
List<URL> urls = classpathResources("/*.csv-spec");
assertTrue("Not enough specs found " + urls, urls.size() > 0);
List<Object[]> specs = SpecReader.readScriptSpec(urls, specParser());

int len = specs.get(0).length;
List<Object[]> testcases = new ArrayList<>();
for (var spec : specs) {
for (Mode mode : Mode.values()) {
Object[] obj = new Object[len + 1];
System.arraycopy(spec, 0, obj, 0, len);
obj[len] = mode;
testcases.add(obj);
}
}
return testcases;
return SpecReader.readScriptSpec(urls, specParser());
}

protected EsqlSpecTestCase(
Expand All @@ -122,34 +110,38 @@ protected EsqlSpecTestCase(
String testName,
Integer lineNumber,
CsvTestCase testCase,
String instructions,
Mode mode
String instructions
) {
this.fileName = fileName;
this.groupName = groupName;
this.testName = testName;
this.lineNumber = lineNumber;
this.testCase = testCase;
this.instructions = instructions;
this.mode = mode;
this.mode = randomFrom(Mode.values());
}

private static boolean dataLoaded = false;

@Before
public void setup() throws IOException {
if (supportsInferenceTestService()) {
createInferenceEndpoints(adminClient());
}

boolean supportsLookup = supportsIndexModeLookup();
boolean supportsSourceMapping = supportsSourceFieldMapping();
if (indexExists(availableDatasetsForEs(client(), supportsLookup, supportsSourceMapping).iterator().next().indexName()) == false) {
loadDataSetIntoEs(client(), supportsLookup, supportsSourceMapping);
boolean supportsInferenceTestService = supportsInferenceTestService();
if (dataLoaded == false) {
if (supportsInferenceTestService) {
createInferenceEndpoints(adminClient());
}

loadDataSetIntoEs(client(), supportsLookup, supportsSourceMapping, supportsInferenceTestService);
dataLoaded = true;
}
}

@AfterClass
public static void wipeTestData() throws IOException {
try {
dataLoaded = false;
adminClient().performRequest(new Request("DELETE", "/*"));
} catch (ResponseException e) {
// 404 here just means we had no indexes
Expand All @@ -159,7 +151,6 @@ public static void wipeTestData() throws IOException {
}

deleteInferenceEndpoints(adminClient());

}

public boolean logResults() {
Expand Down Expand Up @@ -211,38 +202,6 @@ protected static void checkCapabilities(RestClient client, TestFeatureService te
}
}

protected static boolean hasCapabilities(List<String> requiredCapabilities) throws IOException {
return hasCapabilities(adminClient(), requiredCapabilities);
}

public static boolean hasCapabilities(RestClient client, List<String> requiredCapabilities) throws IOException {
if (requiredCapabilities.isEmpty()) {
return true;
}
try {
if (clusterHasCapability(client, "POST", "/_query", List.of(), requiredCapabilities).orElse(false)) {
return true;
}
LOGGER.info("capabilities API returned false, we might be in a mixed version cluster so falling back to cluster features");
} catch (ResponseException e) {
if (e.getResponse().getStatusLine().getStatusCode() / 100 == 4) {
/*
* The node we're testing against is too old for the capabilities
* API which means it has to be pretty old. Very old capabilities
* are ALSO present in the features API, so we can check them instead.
*
* It's kind of weird that we check for *any* 400, but that's required
* because old versions of Elasticsearch return 400, not the expected
* 404.
*/
LOGGER.info("capabilities API failed, falling back to cluster features");
} else {
throw e;
}
}
return false;
}
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Moved.

Also, we don't need this big comment. It was require when this test supported versions of elasticsearch without the capability api.


protected boolean supportsInferenceTestService() {
return true;
}
Expand Down Expand Up @@ -271,7 +230,9 @@ protected final void doTest(String query) throws Throwable {
builder.tables(tables());
}

Map<?, ?> prevTooks = supportsTook() ? tooks() : null;
boolean checkTook = supportsTook() && rarely();
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

About 20% of the time we check took. It's actually pretty expensive to build. Not terribly so, but it shaved some load off of the test.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's not our numbers that are expensive - it's the fact that we have to fetch the entire x-pack usage api.


Map<?, ?> prevTooks = checkTook ? tooks() : null;
Map<String, Object> answer = RestEsqlTestCase.runEsql(
builder.query(query),
testCase.assertWarnings(deduplicateExactWarnings()),
Expand All @@ -296,7 +257,7 @@ protected final void doTest(String query) throws Throwable {

assertResults(expectedColumnsWithValues, actualColumns, actualValues, testCase.ignoreOrder, logger);

if (supportsTook()) {
if (checkTook) {
LOGGER.info("checking took incremented from {}", prevTooks);
long took = ((Number) answer.get("took")).longValue();
int prevTookHisto = ((Number) prevTooks.remove(tookKey(took))).intValue();
Expand Down Expand Up @@ -413,15 +374,14 @@ protected boolean preserveClusterUponCompletion() {
return true;
}

@Before
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Saves some time at the cost of some paranoia.

@After
public void assertRequestBreakerEmptyAfterTests() throws Exception {
assertRequestBreakerEmpty();
}

public static void assertRequestBreakerEmpty() throws Exception {
assertBusy(() -> {
HttpEntity entity = adminClient().performRequest(new Request("GET", "/_nodes/stats")).getEntity();
HttpEntity entity = adminClient().performRequest(new Request("GET", "/_nodes/stats?metric=breaker")).getEntity();
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Saves a bunch of bytes across the wire. We only really want the breaker stats.

Map<?, ?> stats = XContentHelper.convertToMap(XContentType.JSON.xContent(), entity.getContent(), false);
Map<?, ?> nodes = (Map<?, ?>) stats.get("nodes");

Expand Down
22 changes: 19 additions & 3 deletions ...n/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/RestEsqlTestCase.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.Response;
import org.elasticsearch.client.ResponseException;
import org.elasticsearch.client.RestClient;
import org.elasticsearch.client.WarningsHandler;
import org.elasticsearch.common.bytes.BytesArray;
import org.elasticsearch.common.io.Streams;
Expand Down Expand Up @@ -52,6 +53,7 @@
import java.util.Locale;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.IntFunction;

import static java.util.Collections.emptySet;
Expand Down Expand Up @@ -1334,8 +1336,8 @@ public static Map<String, Object> runEsqlAsync(
checkKeepOnCompletion(requestObject, json, keepOnCompletion);
String id = (String) json.get("id");

var supportsAsyncHeaders = clusterHasCapability("POST", "/_query", List.of(), List.of("async_query_status_headers")).orElse(false);
var supportsSuggestedCast = clusterHasCapability("POST", "/_query", List.of(), List.of("suggested_cast")).orElse(false);
var supportsAsyncHeaders = hasCapabilities(client(), List.of("async_query_status_headers"));
var supportsSuggestedCast = hasCapabilities(client(), List.of("suggested_cast"));

if (id == null) {
// no id returned from an async call, must have completed immediately and without keep_on_completion
Expand Down Expand Up @@ -1409,13 +1411,27 @@ public static Map<String, Object> runEsqlAsync(
private static void prepareProfileLogger(RequestObjectBuilder requestObject, @Nullable ProfileLogger profileLogger) throws IOException {
if (profileLogger != null) {
profileLogger.clearProfile();
var isProfileSafe = clusterHasCapability("POST", "/_query", List.of(), List.of("fixed_profile_serialization")).orElse(false);
var isProfileSafe = hasCapabilities(client(), List.of("fixed_profile_serialization"));
if (isProfileSafe) {
requestObject.profile(true);
}
}
}

private static final Map<List<String>, Boolean> capabilities = new ConcurrentHashMap<>();

public static boolean hasCapabilities(RestClient client, List<String> requiredCapabilities) throws IOException {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It looks like List<String> requiredCapabilities always have 1 item in it. Is it worth simplifying? May be I missed some usage?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

csv-spec files make a list actually.

if (requiredCapabilities.isEmpty()) {
return true;
}
Boolean cap = capabilities.get(requiredCapabilities);
if (cap == null) {
cap = clusterHasCapability(client, "POST", "/_query", List.of(), requiredCapabilities).orElse(false);
capabilities.put(requiredCapabilities, cap);
}
return cap;
}
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Moved this up to this class so others could share in the cache.


private static Object removeOriginalTypesAndSuggestedCast(Object response) {
if (response instanceof ArrayList<?> columns) {
var newColumns = new ArrayList<>();
Expand Down
Loading