Skip to content

Commit a5d2807

Browse files
ioanatiaioanatia
committed
Add FORK generative tests
1 parent 53f3ab2 commit a5d2807

File tree

5 files changed

+180
-14
lines changed

5 files changed

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

‎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: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
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+
public abstract class GenerativeForkRestTest extends EsqlSpecTestCase {
19+
public GenerativeForkRestTest(String fileName, String groupName, String testName, Integer lineNumber, CsvSpecReader.CsvTestCase testCase, String instructions, Mode mode) {
20+
super(fileName, groupName, testName, lineNumber, testCase, instructions, mode);
21+
}
22+
23+
@Override
24+
protected void doTest() throws Throwable {
25+
String query = testCase.query + " | FORK (WHERE true) (WHERE true) | WHERE _fork == \"fork1\" | DROP _fork";
26+
doTest(query);
27+
}
28+
29+
@Override
30+
protected void shouldSkipTest(String testName) throws IOException {
31+
super.shouldSkipTest(testName);
32+
33+
assumeFalse(
34+
"Tests using FORK or RRF already are skipped since we don't support multiple FORKs",
35+
testCase.requiredCapabilities.contains(FORK_V7.capabilityName()) || testCase.requiredCapabilities.contains(RRF.capabilityName())
36+
);
37+
38+
assumeFalse(
39+
"Tests using INSIST are not supported for now",
40+
testCase.requiredCapabilities.contains(UNMAPPED_FIELDS.capabilityName())
41+
);
42+
43+
assumeFalse(
44+
"Tests using implicit_casting_date_and_date_nanos are not supported for now",
45+
testCase.requiredCapabilities.contains(IMPLICIT_CASTING_DATE_AND_DATE_NANOS.capabilityName())
46+
);
47+
48+
assumeTrue(
49+
"Cluster needs to support FORK",
50+
hasCapabilities(client(), List.of(FORK_V7.capabilityName()))
51+
);
52+
}
53+
}

‎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
@@ -1615,20 +1615,21 @@ record TypeResolutionKey(String fieldName, DataType fieldType) {}
16151615
@Override
16161616
public LogicalPlan apply(LogicalPlan plan) {
16171617
unionFieldAttributes = new ArrayList<>();
1618+
return plan.transformUp(LogicalPlan.class, p -> p.childrenResolved() == false ? p : doRule(p));
1619+
}
1620+
1621+
private LogicalPlan doRule(LogicalPlan plan) {
1622+
Holder<Integer> alreadyAddedUnionFieldAttributes = new Holder<>(unionFieldAttributes.size());
16181623
// Collect field attributes from previous runs
1619-
plan.forEachUp(EsRelation.class, rel -> {
1624+
if (plan instanceof EsRelation rel) {
1625+
unionFieldAttributes.clear();
16201626
for (Attribute attr : rel.output()) {
16211627
if (attr instanceof FieldAttribute fa && fa.field() instanceof MultiTypeEsField && fa.synthetic()) {
16221628
unionFieldAttributes.add(fa);
16231629
}
16241630
}
1625-
});
1626-
1627-
return plan.transformUp(LogicalPlan.class, p -> p.childrenResolved() == false ? p : doRule(p));
1628-
}
1631+
}
16291632

1630-
private LogicalPlan doRule(LogicalPlan plan) {
1631-
int alreadyAddedUnionFieldAttributes = unionFieldAttributes.size();
16321633
// See if the eval function has an unresolved MultiTypeEsField field
16331634
// Replace the entire convert function with a new FieldAttribute (containing type conversion knowledge)
16341635
plan = plan.transformExpressionsOnly(e -> {
@@ -1637,8 +1638,9 @@ private LogicalPlan doRule(LogicalPlan plan) {
16371638
}
16381639
return e;
16391640
});
1641+
16401642
// If no union fields were generated, return the plan as is
1641-
if (unionFieldAttributes.size() == alreadyAddedUnionFieldAttributes) {
1643+
if (unionFieldAttributes.size() == alreadyAddedUnionFieldAttributes.get()) {
16421644
return plan;
16431645
}
16441646

0 commit comments

Comments
 (0)