Add Generative tests for Lookup Join On Multiple Columns

Author: julian-elastic

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

@@ -215,13 +215,24 @@ private List<String> availableIndices() throws IOException {
             .toList();
     }
 
-    public record LookupIdx(String idxName, String key, String keyType) {}
+    public record LookupIdxColumn(String name, String type) {}
+
+    public record LookupIdx(String idxName, List<LookupIdxColumn> keys) {}
 
     private List<LookupIdx> lookupIndices() {
         List<LookupIdx> result = new ArrayList<>();
         // we don't have key info from the dataset loader, let's hardcode it for now
-        result.add(new LookupIdx("languages_lookup", "language_code", "integer"));
-        result.add(new LookupIdx("message_types_lookup", "message", "keyword"));
+        result.add(new LookupIdx("languages_lookup", List.of(new LookupIdxColumn("language_code", "integer"))));
+        result.add(new LookupIdx("message_types_lookup", List.of(new LookupIdxColumn("message", "keyword"))));
+        List<LookupIdxColumn> multiColumnJoinableLookupKeys = List.of(
+            new LookupIdxColumn("id_int", "integer"),
+            new LookupIdxColumn("name_str", "keyword"),
+            new LookupIdxColumn("is_active_bool", "boolean"),
+            new LookupIdxColumn("ip_addr", "ip"),
+            new LookupIdxColumn("other1", "keyword"),
+            new LookupIdxColumn("other2", "integer")
+        );
+        result.add(new LookupIdx("multi_column_joinable_lookup", multiColumnJoinableLookupKeys));
         return result;
     }
 

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

@@ -11,10 +11,15 @@
 import org.elasticsearch.xpack.esql.qa.rest.generative.GenerativeRestTest;
 import org.elasticsearch.xpack.esql.qa.rest.generative.command.CommandGenerator;
 
+import java.util.ArrayList;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 import static org.elasticsearch.test.ESTestCase.randomFrom;
+import static org.elasticsearch.test.ESTestCase.randomInt;
+import static org.elasticsearch.test.ESTestCase.randomSubsetOf;
 
 public class LookupJoinGenerator implements CommandGenerator {
 
@@ -29,15 +34,50 @@ public CommandDescription generate(
     ) {
         GenerativeRestTest.LookupIdx lookupIdx = randomFrom(schema.lookupIndices());
         String lookupIdxName = lookupIdx.idxName();
-        String idxKey = lookupIdx.key();
-        String keyType = lookupIdx.keyType();
+        int joinColumnsCount = randomInt(lookupIdx.keys().size() - 1) + 1; // at least one column must be used for the join
+        List<GenerativeRestTest.LookupIdxColumn> joinColumns = randomSubsetOf(joinColumnsCount, lookupIdx.keys());
+        List<String> keyNames = new ArrayList<>();
+        List<String> joinOn = new ArrayList<>();
+        Set<String> usedColumns = new HashSet<>();
+        for (GenerativeRestTest.LookupIdxColumn joinColumn : joinColumns) {
+            String idxKey = joinColumn.name();
+            String keyType = joinColumn.type();
 
-        var candidateKeys = previousOutput.stream().filter(x -> x.type().equals(keyType)).toList();
-        if (candidateKeys.isEmpty()) {
+            var candidateKeys = previousOutput.stream().filter(x -> x.type().equals(keyType)).toList();
+            if (candidateKeys.isEmpty()) {
+                continue; // no candidate keys of the right type, skip this column
+            }
+            EsqlQueryGenerator.Column key = randomFrom(candidateKeys);
+            if (usedColumns.contains(key.name())) {
+                continue; // already used this column, skip
+            } else {
+                usedColumns.add(key.name());
+            }
+            keyNames.add(key.name());
+            joinOn.add(idxKey);
+        }
+        if (keyNames.isEmpty()) {
             return EMPTY_DESCRIPTION;
         }
-        EsqlQueryGenerator.Column key = randomFrom(candidateKeys);
-        String cmdString = "| rename " + key.name() + " as " + idxKey + " | lookup join " + lookupIdxName + " on " + idxKey;
+        StringBuilder stringBuilder = new StringBuilder();
+        for (int i = 0; i < keyNames.size(); i++) {
+            stringBuilder.append("| rename ");
+            stringBuilder.append(keyNames.get(i));
+            stringBuilder.append(" as ");
+            stringBuilder.append(joinOn.get(i));
+        }
+        stringBuilder.append(" | lookup join ").append(lookupIdxName).append(" on ");
+        for (int i = 0; i < keyNames.size(); i++) {
+            stringBuilder.append(joinOn.get(i));
+            if (i < keyNames.size() - 1) {
+                stringBuilder.append(", ");
+            }
+        }
+        stringBuilder.append("| LIMIT 10");
+        if (keyNames.size() > 1) {
+            System.out.println("GENERATED JOIN ON MORE THAN ONE COLUMN: " + stringBuilder);
+        }
+        String cmdString = stringBuilder.toString();
         return new CommandDescription(LOOKUP_JOIN, this, cmdString, Map.of());
     }