Skip to content

Commit 01de61e

Browse files
ioanatiamridula-s109
ioanatia
authored and
mridula-s109
committed
ES|QL: Add FORK generative tests (elastic#129135)
1 parent ac213d5 commit 01de61e

File tree

5 files changed

+196
-14
lines changed

5 files changed

+196
-14
lines changed

‎x-pack/plugin/esql/qa/server/single-node/src/javaRestTest/java/org/elasticsearch/xpack/esql/qa/single_node/GenerativeForkIT.java

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
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.CsvSpecReader;
15+
import org.elasticsearch.xpack.esql.qa.rest.generative.GenerativeForkRestTest;
16+
import org.junit.ClassRule;
17+
18+
@ThreadLeakFilters(filters = TestClustersThreadFilter.class)
19+
public class GenerativeForkIT extends GenerativeForkRestTest {
20+
@ClassRule
21+
public static ElasticsearchCluster cluster = Clusters.testCluster(spec -> spec.plugin("inference-service-test"));
22+
23+
@Override
24+
protected String getTestRestCluster() {
25+
return cluster.getHttpAddresses();
26+
}
27+
28+
public GenerativeForkIT(
29+
String fileName,
30+
String groupName,
31+
String testName,
32+
Integer lineNumber,
33+
CsvSpecReader.CsvTestCase testCase,
34+
String instructions,
35+
Mode mode
36+
) {
37+
super(fileName, groupName, testName, lineNumber, testCase, instructions, mode);
38+
}
39+
40+
@Override
41+
protected boolean enableRoundingDoubleValuesOnAsserting() {
42+
// This suite runs with more than one node and three shards in serverless
43+
return cluster.getNumNodes() > 1;
44+
}
45+
46+
@Override
47+
protected boolean supportsSourceFieldMapping() {
48+
return cluster.getNumNodes() == 1;
49+
}
50+
}

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

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -259,15 +259,19 @@ protected boolean supportsSourceFieldMapping() throws IOException {
259259
return true;
260260
}
261261

262-
protected final void doTest() throws Throwable {
262+
protected void doTest() throws Throwable {
263+
doTest(testCase.query);
264+
}
265+
266+
protected final void doTest(String query) throws Throwable {
263267
RequestObjectBuilder builder = new RequestObjectBuilder(randomFrom(XContentType.values()));
264268

265-
if (testCase.query.toUpperCase(Locale.ROOT).contains("LOOKUP_\uD83D\uDC14")) {
269+
if (query.toUpperCase(Locale.ROOT).contains("LOOKUP_\uD83D\uDC14")) {
266270
builder.tables(tables());
267271
}
268272

269273
Map<?, ?> prevTooks = supportsTook() ? tooks() : null;
270-
Map<String, Object> answer = runEsql(builder.query(testCase.query), testCase.assertWarnings(deduplicateExactWarnings()));
274+
Map<String, Object> answer = runEsql(builder.query(query), testCase.assertWarnings(deduplicateExactWarnings()));
271275

272276
var expectedColumnsWithValues = loadCsvSpecValues(testCase.expectedResults);
273277

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

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
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;
9+
10+
import org.elasticsearch.xpack.esql.CsvSpecReader;
11+
import org.elasticsearch.xpack.esql.qa.rest.EsqlSpecTestCase;
12+
13+
import java.io.IOException;
14+
import java.util.List;
15+
16+
import static org.elasticsearch.xpack.esql.action.EsqlCapabilities.Cap.*;
17+
18+
/**
19+
* Tests for FORK. We generate tests for FORK from existing CSV tests.
20+
* We append a `| FORK (WHERE true) (WHERE true) | WHERE _fork == "fork1" | DROP _fork` suffix to existing
21+
* CSV test cases. This will produce a query that executes multiple FORK branches but expects the same results
22+
* as the initial CSV test case.
23+
* For now, we skip tests that already require FORK, since multiple FORK commands are not allowed.
24+
*/
25+
public abstract class GenerativeForkRestTest extends EsqlSpecTestCase {
26+
public GenerativeForkRestTest(
27+
String fileName,
28+
String groupName,
29+
String testName,
30+
Integer lineNumber,
31+
CsvSpecReader.CsvTestCase testCase,
32+
String instructions,
33+
Mode mode
34+
) {
35+
super(fileName, groupName, testName, lineNumber, testCase, instructions, mode);
36+
}
37+
38+
@Override
39+
protected void doTest() throws Throwable {
40+
String query = testCase.query + " | FORK (WHERE true) (WHERE true) | WHERE _fork == \"fork1\" | DROP _fork";
41+
doTest(query);
42+
}
43+
44+
@Override
45+
protected void shouldSkipTest(String testName) throws IOException {
46+
super.shouldSkipTest(testName);
47+
48+
assumeFalse(
49+
"Tests using FORK or RRF already are skipped since we don't support multiple FORKs",
50+
testCase.requiredCapabilities.contains(FORK_V7.capabilityName()) || testCase.requiredCapabilities.contains(RRF.capabilityName())
51+
);
52+
53+
assumeFalse(
54+
"Tests using INSIST are not supported for now",
55+
testCase.requiredCapabilities.contains(UNMAPPED_FIELDS.capabilityName())
56+
);
57+
58+
assumeFalse(
59+
"Tests using implicit_casting_date_and_date_nanos are not supported for now",
60+
testCase.requiredCapabilities.contains(IMPLICIT_CASTING_DATE_AND_DATE_NANOS.capabilityName())
61+
);
62+
63+
assumeTrue("Cluster needs to support FORK", hasCapabilities(client(), List.of(FORK_V7.capabilityName())));
64+
}
65+
}

‎x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/ForkIT.java

Lines changed: 64 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
import org.elasticsearch.action.support.WriteRequest;
1212
import org.elasticsearch.common.settings.Settings;
1313
import org.elasticsearch.compute.operator.DriverProfile;
14+
import org.elasticsearch.test.junit.annotations.TestLogging;
1415
import org.elasticsearch.xpack.esql.VerificationException;
1516
import org.elasticsearch.xpack.esql.parser.ParsingException;
1617
import org.junit.Before;
@@ -26,13 +27,13 @@
2627
import static org.elasticsearch.xpack.esql.EsqlTestUtils.getValuesList;
2728
import static org.hamcrest.Matchers.equalTo;
2829

29-
// @TestLogging(value = "org.elasticsearch.xpack.esql:TRACE,org.elasticsearch.compute:TRACE", reason = "debug")
30+
@TestLogging(value = "org.elasticsearch.xpack.esql:TRACE,org.elasticsearch.compute:TRACE", reason = "debug")
3031
public class ForkIT extends AbstractEsqlIntegTestCase {
3132

3233
@Before
3334
public void setupIndex() {
3435
assumeTrue("requires FORK capability", EsqlCapabilities.Cap.FORK.isEnabled());
35-
createAndPopulateIndex();
36+
createAndPopulateIndices();
3637
}
3738

3839
public void testSimple() {
@@ -706,6 +707,52 @@ public void testWithLookUpAfterFork() {
706707
}
707708
}
708709

710+
public void testWithUnionTypesBeforeFork() {
711+
var query = """
712+
FROM test,test-other
713+
| EVAL x = id::keyword
714+
| EVAL id = id::keyword
715+
| EVAL content = content::keyword
716+
| FORK (WHERE x == "2")
717+
(WHERE x == "1")
718+
| SORT _fork, x, content
719+
| KEEP content, id, x, _fork
720+
""";
721+
722+
try (var resp = run(query)) {
723+
assertColumnNames(resp.columns(), List.of("content", "id", "x", "_fork"));
724+
Iterable<Iterable<Object>> expectedValues = List.of(
725+
List.of("This is a brown dog", "2", "2", "fork1"),
726+
List.of("This is a brown dog", "2", "2", "fork1"),
727+
List.of("This is a brown fox", "1", "1", "fork2"),
728+
List.of("This is a brown fox", "1", "1", "fork2")
729+
);
730+
assertValues(resp.values(), expectedValues);
731+
}
732+
}
733+
734+
public void testWithUnionTypesInBranches() {
735+
var query = """
736+
FROM test,test-other
737+
| EVAL content = content::keyword
738+
| FORK (EVAL x = id::keyword | WHERE x == "2" | EVAL id = x::integer)
739+
(EVAL x = "a" | WHERE id::keyword == "1" | EVAL id = id::integer)
740+
| SORT _fork, x
741+
| KEEP content, id, x, _fork
742+
""";
743+
744+
try (var resp = run(query)) {
745+
assertColumnNames(resp.columns(), List.of("content", "id", "x", "_fork"));
746+
Iterable<Iterable<Object>> expectedValues = List.of(
747+
List.of("This is a brown dog", 2, "2", "fork1"),
748+
List.of("This is a brown dog", 2, "2", "fork1"),
749+
List.of("This is a brown fox", 1, "a", "fork2"),
750+
List.of("This is a brown fox", 1, "a", "fork2")
751+
);
752+
assertValues(resp.values(), expectedValues);
753+
}
754+
}
755+
709756
public void testWithEvalWithConflictingTypes() {
710757
var query = """
711758
FROM test
@@ -833,7 +880,7 @@ public void testProfile() {
833880
}
834881
}
835882

836-
private void createAndPopulateIndex() {
883+
private void createAndPopulateIndices() {
837884
var indexName = "test";
838885
var client = client().admin().indices();
839886
var createRequest = client.prepareCreate(indexName)
@@ -867,6 +914,20 @@ private void createAndPopulateIndex() {
867914
.setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE)
868915
.get();
869916
ensureYellow(lookupIndex);
917+
918+
var otherTestIndex = "test-other";
919+
920+
createRequest = client.prepareCreate(otherTestIndex)
921+
.setSettings(Settings.builder().put("index.number_of_shards", 1))
922+
.setMapping("id", "type=keyword", "content", "type=keyword");
923+
assertAcked(createRequest);
924+
client().prepareBulk()
925+
.add(new IndexRequest(otherTestIndex).id("1").source("id", "1", "content", "This is a brown fox"))
926+
.add(new IndexRequest(otherTestIndex).id("2").source("id", "2", "content", "This is a brown dog"))
927+
.add(new IndexRequest(otherTestIndex).id("3").source("id", "3", "content", "This dog is really brown"))
928+
.setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE)
929+
.get();
930+
ensureYellow(indexName);
870931
}
871932

872933
static Iterator<Iterator<Object>> valuesFilter(Iterator<Iterator<Object>> values, Predicate<Iterator<Object>> filter) {

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/Analyzer.java

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1638,20 +1638,21 @@ record TypeResolutionKey(String fieldName, DataType fieldType) {}
16381638
@Override
16391639
public LogicalPlan apply(LogicalPlan plan) {
16401640
unionFieldAttributes = new ArrayList<>();
1641+
return plan.transformUp(LogicalPlan.class, p -> p.childrenResolved() == false ? p : doRule(p));
1642+
}
1643+
1644+
private LogicalPlan doRule(LogicalPlan plan) {
1645+
Holder<Integer> alreadyAddedUnionFieldAttributes = new Holder<>(unionFieldAttributes.size());
16411646
// Collect field attributes from previous runs
1642-
plan.forEachUp(EsRelation.class, rel -> {
1647+
if (plan instanceof EsRelation rel) {
1648+
unionFieldAttributes.clear();
16431649
for (Attribute attr : rel.output()) {
16441650
if (attr instanceof FieldAttribute fa && fa.field() instanceof MultiTypeEsField && fa.synthetic()) {
16451651
unionFieldAttributes.add(fa);
16461652
}
16471653
}
1648-
});
1649-
1650-
return plan.transformUp(LogicalPlan.class, p -> p.childrenResolved() == false ? p : doRule(p));
1651-
}
1654+
}
16521655

1653-
private LogicalPlan doRule(LogicalPlan plan) {
1654-
int alreadyAddedUnionFieldAttributes = unionFieldAttributes.size();
16551656
// See if the eval function has an unresolved MultiTypeEsField field
16561657
// Replace the entire convert function with a new FieldAttribute (containing type conversion knowledge)
16571658
plan = plan.transformExpressionsOnly(e -> {
@@ -1660,8 +1661,9 @@ private LogicalPlan doRule(LogicalPlan plan) {
16601661
}
16611662
return e;
16621663
});
1664+
16631665
// If no union fields were generated, return the plan as is
1664-
if (unionFieldAttributes.size() == alreadyAddedUnionFieldAttributes) {
1666+
if (unionFieldAttributes.size() == alreadyAddedUnionFieldAttributes.get()) {
16651667
return plan;
16661668
}
16671669

0 commit comments

Comments
 (0)