Skip to content

Commit 77c5835

Browse files
alex-spiesalex-spies
committed
Add union type test cases to LookupJoinTypesIT
1 parent 213f5df commit 77c5835

File tree

1 file changed

+101
-30
lines changed

1 file changed

+101
-30
lines changed

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

Lines changed: 101 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,6 @@
5959
import static org.hamcrest.Matchers.containsString;
6060
import static org.hamcrest.Matchers.equalTo;
6161
import static org.hamcrest.Matchers.is;
62-
import static org.hamcrest.Matchers.nullValue;
6362

6463
/**
6564
* This test suite tests the lookup join functionality in ESQL with various data types.
@@ -146,6 +145,17 @@ protected Collection<Class<? extends Plugin>> nodePlugins() {
146145
}
147146
}
148147

148+
// Union types; non-exhaustive and can be extended
149+
{
150+
TestConfigs configs = testConfigurations.computeIfAbsent("union-types", TestConfigs::new);
151+
configs.addUnionTypePasses(SHORT, INTEGER, INTEGER);
152+
configs.addUnionTypePasses(BYTE, DOUBLE, LONG);
153+
configs.addUnionTypePasses(DATETIME, DATE_NANOS, DATE_NANOS);
154+
configs.addUnionTypePasses(DATE_NANOS, DATETIME, DATETIME);
155+
configs.addUnionTypePasses(SCALED_FLOAT, HALF_FLOAT, DOUBLE);
156+
configs.addUnionTypePasses(TEXT, KEYWORD, KEYWORD);
157+
}
158+
149159
// Tests for all unsupported types
150160
DataType[] unsupported = Join.UNSUPPORTED_TYPES;
151161
{
@@ -276,6 +286,10 @@ public void testLookupJoinOthers() {
276286
testLookupJoinTypes("others");
277287
}
278288

289+
public void testLookupJoinUnionTypes() {
290+
testLookupJoinTypes("union-types");
291+
}
292+
279293
private void testLookupJoinTypes(String group) {
280294
TestConfigs configs = testConfigurations.get(group);
281295
initIndexes(configs);
@@ -298,6 +312,7 @@ private void initIndexes(TestConfigs configs) {
298312
if (mapping.settings != null) {
299313
builder = builder.setSettings(mapping.settings);
300314
}
315+
301316
assertAcked(builder);
302317
}
303318
}
@@ -309,12 +324,8 @@ private void initData(TestConfigs configs) {
309324
}
310325
}
311326

312-
private static String lookupPropertyFor(TestConfig config) {
313-
return String.format(Locale.ROOT, "\"%s\": %s", config.lookupFieldName(), sampleDataTextFor(config.lookupType()));
314-
}
315-
316-
private static String mainPropertyFor(TestConfig config) {
317-
return String.format(Locale.ROOT, "\"%s\": %s", config.mainFieldName(), sampleDataTextFor(config.mainType()));
327+
private static String propertyFor(String fieldName, DataType type) {
328+
return String.format(Locale.ROOT, "\"%s\": %s", fieldName, sampleDataTextFor(type));
318329
}
319330

320331
private static String sampleDataTextFor(DataType type) {
@@ -425,12 +436,13 @@ public List<TestDocument> docs() {
425436
%s,
426437
"other": "value"
427438
}
428-
""", lookupPropertyFor(config));
439+
""", propertyFor(config.lookupFieldName(), config.lookupType()));
429440
results.add(new TestDocument(config.lookupIndexName(), "" + (++docId), doc));
430441
}
442+
431443
List<String> mainProperties = configs.values()
432444
.stream()
433-
.map(LookupJoinTypesIT::mainPropertyFor)
445+
.map(c -> propertyFor(c.mainFieldName(), c.mainType()))
434446
.distinct()
435447
.collect(Collectors.toList());
436448
results.add(new TestDocument(MAIN_INDEX, "1", String.format(Locale.ROOT, """
@@ -439,6 +451,20 @@ public List<TestDocument> docs() {
439451
}
440452
""", String.join(",\n ", mainProperties))));
441453

454+
for (TestConfig config : configs.values()) {
455+
for (TestMapping addtionalIndex : config.additionalIndexes()) {
456+
String doc = String.format(Locale.ROOT, """
457+
{
458+
%s,
459+
"other": "value"
460+
}
461+
""", propertyFor(config.mainFieldName(), ((TestConfigPassesUnionType) config).otherMainType()));
462+
// Casting to TestConfigPassesUnionType is an ugly hack; better to derive the test data from the TestMapping or from the
463+
// TestConfig.
464+
results.add(new TestDocument(addtionalIndex.indexName, "1", doc));
465+
}
466+
}
467+
442468
return results;
443469
}
444470

@@ -461,6 +487,10 @@ private void addPasses(DataType mainType, DataType lookupType) {
461487
add(new TestConfigPasses(mainType, lookupType));
462488
}
463489

490+
private void addUnionTypePasses(DataType mainType, DataType otherMainType, DataType lookupType) {
491+
add(new TestConfigPassesUnionType(mainType, otherMainType, lookupType));
492+
}
493+
464494
private void addFails(DataType mainType, DataType lookupType) {
465495
String fieldName = LOOKUP_INDEX_PREFIX + lookupType.esType();
466496
String errorMessage = String.format(
@@ -500,22 +530,6 @@ private void addFailsUnsupported(DataType mainType, DataType lookupType) {
500530
}
501531
}
502532

503-
private static class UnionTypeTestConfigs extends TestConfigs {
504-
UnionTypeTestConfigs(String group) {
505-
super(group);
506-
}
507-
508-
@Override
509-
public List<TestMapping> indices() {
510-
return super.indices();
511-
}
512-
513-
@Override
514-
public List<TestDocument> docs() {
515-
return super.docs();
516-
}
517-
}
518-
519533
interface TestConfig {
520534
DataType mainType();
521535

@@ -529,7 +543,7 @@ default String mainFieldName() {
529543
}
530544

531545
default TestMapping mainIndex() {
532-
return new TestMapping(MAIN_INDEX, List.of(propertySpecFor(mainFieldName(), mainType(), "")), null);
546+
return new TestMapping(MAIN_INDEX, List.of(propertySpecFor(mainFieldName(), mainType())), null);
533547
}
534548

535549
/**
@@ -560,7 +574,7 @@ default String lookupFieldName() {
560574
default TestMapping lookupIndex() {
561575
return new TestMapping(
562576
lookupIndexName(),
563-
List.of(propertySpecFor(lookupFieldName(), lookupType(), ", \"other\": { \"type\" : \"keyword\" }")),
577+
List.of(propertySpecFor(lookupFieldName(), lookupType()), "\"other\": { \"type\" : \"keyword\" }"),
564578
Settings.builder().put("index.number_of_shards", 1).put("index.number_of_replicas", 0).put("index.mode", "lookup").build()
565579
);
566580
}
@@ -589,17 +603,17 @@ default String testQuery() {
589603
void doTest();
590604
}
591605

592-
private static String propertySpecFor(String fieldName, DataType type, String extra) {
606+
private static String propertySpecFor(String fieldName, DataType type) {
593607
if (type == SCALED_FLOAT) {
594608
return String.format(
595609
Locale.ROOT,
596610
"\"%s\": { \"type\" : \"%s\", \"scaling_factor\": %f }",
597611
fieldName,
598612
type.esType(),
599613
SCALING_FACTOR
600-
) + extra;
614+
);
601615
}
602-
return String.format(Locale.ROOT, "\"%s\": { \"type\" : \"%s\" }", fieldName, type.esType().replaceAll("cartesian_", "")) + extra;
616+
return String.format(Locale.ROOT, "\"%s\": { \"type\" : \"%s\" }", fieldName, type.esType().replaceAll("cartesian_", ""));
603617
}
604618

605619
private static void validateIndex(String indexName, String fieldName, Object expectedValue) {
@@ -627,6 +641,63 @@ public void doTest() {
627641
}
628642
}
629643

644+
private record TestConfigPassesUnionType(DataType mainType, DataType otherMainType, DataType lookupType) implements TestConfig {
645+
@Override
646+
public String lookupIndexName() {
647+
// Override so it doesn't clash with other lookup indices from non-union type tests.
648+
return LOOKUP_INDEX_PREFIX + mainType().esType() + "_union_" + otherMainType().esType() + "_" + lookupType().esType();
649+
}
650+
651+
private String additionalIndexName() {
652+
return MAIN_INDEX + "_" + mainFieldName() + "_as_" + otherMainType().typeName();
653+
}
654+
655+
@Override
656+
public List<TestMapping> additionalIndexes() {
657+
return List.of(new TestMapping(additionalIndexName(), List.of(propertySpecFor(mainFieldName(), otherMainType)), null));
658+
}
659+
660+
@Override
661+
public void validateAdditionalIndexes() {
662+
validateIndex(additionalIndexName(), mainFieldName(), sampleDataFor(otherMainType));
663+
}
664+
665+
@Override
666+
public String testQuery() {
667+
String mainField = mainFieldName();
668+
String lookupField = lookupFieldName();
669+
String lookupIndex = lookupIndexName();
670+
671+
return String.format(
672+
Locale.ROOT,
673+
"FROM %s, %s | EVAL %s = %s::%s | LOOKUP JOIN %s ON %s | KEEP other",
674+
MAIN_INDEX,
675+
additionalIndexName(),
676+
lookupField,
677+
mainField,
678+
lookupType.typeName(),
679+
lookupIndex,
680+
lookupField
681+
);
682+
}
683+
684+
@Override
685+
public void doTest() {
686+
String query = testQuery();
687+
try (var response = EsqlQueryRequestBuilder.newRequestBuilder(client()).query(query).get()) {
688+
Iterator<Object> results = response.response().column(0).iterator();
689+
690+
assertTrue("Expected at least two results for query, but result was empty: " + query, results.hasNext());
691+
Object indexedResult = results.next();
692+
assertThat("Expected valid result: " + query, indexedResult, equalTo("value"));
693+
694+
assertTrue("Expected at least two results for query: " + query, results.hasNext());
695+
indexedResult = results.next();
696+
assertThat("Expected valid result: " + query, indexedResult, equalTo("value"));
697+
}
698+
}
699+
}
700+
630701
private record TestConfigFails<E extends Exception>(DataType mainType, DataType lookupType, Class<E> exception, Consumer<E> assertion)
631702
implements
632703
TestConfig {

0 commit comments

Comments
 (0)