Lookup Join on Multiple Columns POC WIP

Author: julian-elastic

server/src/main/java/org/elasticsearch/TransportVersions.java

@@ -342,6 +342,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion NODE_USAGE_STATS_FOR_THREAD_POOLS_IN_CLUSTER_INFO = def(9_121_0_00);
     public static final TransportVersion ESQL_CATEGORIZE_OPTIONS = def(9_122_0_00);
     public static final TransportVersion ML_INFERENCE_AZURE_AI_STUDIO_RERANK_ADDED = def(9_123_0_00);
+    public static final TransportVersion ESQL_LOOKUP_JOIN_ON_MANY_FIELDS = def(9_124_0_00);
 
     /*
      * STOP! READ THIS FIRST! No, really,

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/enrich/LookupFromIndexOperator.java

@@ -12,6 +12,7 @@
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.AsyncOperator;
@@ -27,6 +28,7 @@
 import org.elasticsearch.xpack.esql.core.expression.NamedExpression;
 import org.elasticsearch.xpack.esql.core.tree.Source;
 import org.elasticsearch.xpack.esql.core.type.DataType;
+import org.elasticsearch.xpack.esql.planner.Layout;
 
 import java.io.IOException;
 import java.util.Iterator;
@@ -37,47 +39,64 @@
 
 // TODO rename package
 public final class LookupFromIndexOperator extends AsyncOperator<LookupFromIndexOperator.OngoingJoin> {
+    public record MatchConfig(FieldAttribute.FieldName fieldName, int channel, DataType type) implements Writeable {
+        public MatchConfig(FieldAttribute match, Layout.ChannelAndType input) {
+            // TODO: Using exactAttribute was supposed to handle TEXT fields with KEYWORD subfields - but we don't allow these in lookup
+            // indices, so the call to exactAttribute looks redundant now.
+            this(match.exactAttribute().fieldName(), input.channel(), input.type());
+        }
+
+        public MatchConfig(StreamInput in) throws IOException {
+            this(new FieldAttribute.FieldName(in.readString()), in.readInt(), DataType.readFrom(in));
+        }
+
+        @Override
+        public void writeTo(StreamOutput out) throws IOException {
+            out.writeString(fieldName.string());
+            out.writeInt(channel);
+            type.writeTo(out);
+        }
+
+    }
+
     public record Factory(
+        List<MatchConfig> matchFields,
         String sessionId,
         CancellableTask parentTask,
         int maxOutstandingRequests,
-        int inputChannel,
         Function<DriverContext, LookupFromIndexService> lookupService,
-        DataType inputDataType,
         String lookupIndexPattern,
         String lookupIndex,
-        FieldAttribute.FieldName matchField,
         List<NamedExpression> loadFields,
         Source source
     ) implements OperatorFactory {
         @Override
         public String describe() {
-            return "LookupOperator[index="
-                + lookupIndex
-                + " input_type="
-                + inputDataType
-                + " match_field="
-                + matchField.string()
-                + " load_fields="
-                + loadFields
-                + " inputChannel="
-                + inputChannel
-                + "]";
+            StringBuilder stringBuilder = new StringBuilder();
+            stringBuilder.append("LookupOperator[index=").append(lookupIndex).append(" load_fields=").append(loadFields);
+            for (MatchConfig matchField : matchFields) {
+                stringBuilder.append(" input_type=")
+                    .append(matchField.type)
+                    .append(" match_field=")
+                    .append(matchField.fieldName.string())
+                    .append(" inputChannel=")
+                    .append(matchField.channel);
+            }
+            stringBuilder.append("]");
+            return stringBuilder.toString();
         }
 
         @Override
         public Operator get(DriverContext driverContext) {
             return new LookupFromIndexOperator(
+                matchFields,
                 sessionId,
                 driverContext,
                 parentTask,
                 maxOutstandingRequests,
-                inputChannel,
                 lookupService.apply(driverContext),
-                inputDataType,
                 lookupIndexPattern,
                 lookupIndex,
-                matchField.string(),
                 loadFields,
                 source
             );
@@ -87,14 +106,12 @@ public Operator get(DriverContext driverContext) {
     private final LookupFromIndexService lookupService;
     private final String sessionId;
     private final CancellableTask parentTask;
-    private final int inputChannel;
-    private final DataType inputDataType;
     private final String lookupIndexPattern;
     private final String lookupIndex;
-    private final String matchField;
     private final List<NamedExpression> loadFields;
     private final Source source;
     private long totalTerms = 0L;
+    private List<MatchConfig> matchFields;
     /**
      * Total number of pages emitted by this {@link Operator}.
      */
@@ -105,43 +122,47 @@ public Operator get(DriverContext driverContext) {
     private OngoingJoin ongoing = null;
 
     public LookupFromIndexOperator(
+        List<MatchConfig> matchFields,
         String sessionId,
         DriverContext driverContext,
         CancellableTask parentTask,
         int maxOutstandingRequests,
-        int inputChannel,
         LookupFromIndexService lookupService,
-        DataType inputDataType,
         String lookupIndexPattern,
         String lookupIndex,
-        String matchField,
         List<NamedExpression> loadFields,
         Source source
     ) {
         super(driverContext, lookupService.getThreadContext(), maxOutstandingRequests);
+        this.matchFields = matchFields;
         this.sessionId = sessionId;
         this.parentTask = parentTask;
-        this.inputChannel = inputChannel;
         this.lookupService = lookupService;
-        this.inputDataType = inputDataType;
         this.lookupIndexPattern = lookupIndexPattern;
         this.lookupIndex = lookupIndex;
-        this.matchField = matchField;
         this.loadFields = loadFields;
         this.source = source;
     }
 
     @Override
     protected void performAsync(Page inputPage, ActionListener<OngoingJoin> listener) {
-        final Block inputBlock = inputPage.getBlock(inputChannel);
-        totalTerms += inputBlock.getTotalValueCount();
+        // what is happening here?
+        // should I be getting multiple bloks, and send them using the LookupFromIndexService.Request
+        // is the totalTerms supposed to be the total number of terms in all blocks?
+        Block[] inputBlockArray = new Block[matchFields.size()];
+        for (int i = 0; i < matchFields.size(); i++) {
+            MatchConfig matchField = matchFields.get(i);
+            int inputChannel = matchField.channel;
+            final Block inputBlock = inputPage.getBlock(inputChannel);
+            totalTerms += inputBlock.getTotalValueCount();
+            inputBlockArray[i] = inputBlock;
+        }
         LookupFromIndexService.Request request = new LookupFromIndexService.Request(
             sessionId,
             lookupIndex,
             lookupIndexPattern,
-            inputDataType,
-            matchField,
-            new Page(inputBlock),
+            matchFields,
+            new Page(inputBlockArray),
             loadFields,
             source
         );
@@ -190,17 +211,18 @@ protected void releaseFetchedOnAnyThread(OngoingJoin ongoingJoin) {
 
     @Override
     public String toString() {
-        return "LookupOperator[index="
-            + lookupIndex
-            + " input_type="
-            + inputDataType
-            + " match_field="
-            + matchField
-            + " load_fields="
-            + loadFields
-            + " inputChannel="
-            + inputChannel
-            + "]";
+        StringBuilder stringBuilder = new StringBuilder();
+        stringBuilder.append("LookupOperator[index=").append(lookupIndex).append(" load_fields=").append(loadFields);
+        for (MatchConfig matchField : matchFields) {
+            stringBuilder.append(" input_type=")
+                .append(matchField.type)
+                .append(" match_field=")
+                .append(matchField.fieldName.string())
+                .append(" inputChannel=")
+                .append(matchField.channel);
+        }
+        stringBuilder.append("]");
+        return stringBuilder.toString();
     }
 
     @Override

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/enrich/LookupFromIndexService.java

@@ -40,6 +40,7 @@
 import java.io.IOException;
 import java.util.List;
 import java.util.Objects;
+import java.util.stream.Collectors;
 
 /**
  * {@link LookupFromIndexService} performs lookup against a Lookup index for
@@ -84,7 +85,7 @@ protected TransportRequest transportRequest(LookupFromIndexService.Request reque
             request.inputPage,
             null,
             request.extractFields,
-            request.matchField,
+            request.matchFields,
             request.source
         );
     }
@@ -98,10 +99,14 @@ protected QueryList queryList(
         @Nullable DataType inputDataType,
         Warnings warnings
     ) {
-        return termQueryList(context.getFieldType(request.matchField), context, aliasFilter, inputBlock, inputDataType).onlySingleValues(
-            warnings,
-            "LOOKUP JOIN encountered multi-value"
-        );
+        // TODO: THIS NEEDS IMPLEMENTATION FOR MULTI-FIELD MATCHING
+        return termQueryList(
+            context.getFieldType(request.matchFields.get(0).fieldName().string()),
+            context,
+            aliasFilter,
+            inputBlock,
+            inputDataType
+        ).onlySingleValues(warnings, "LOOKUP JOIN encountered multi-value");
     }
 
     @Override
@@ -115,25 +120,24 @@ protected AbstractLookupService.LookupResponse readLookupResponse(StreamInput in
     }
 
     public static class Request extends AbstractLookupService.Request {
-        private final String matchField;
+        private final List<LookupFromIndexOperator.MatchConfig> matchFields;
 
         Request(
             String sessionId,
             String index,
             String indexPattern,
-            DataType inputDataType,
-            String matchField,
+            List<LookupFromIndexOperator.MatchConfig> matchFields,
             Page inputPage,
             List<NamedExpression> extractFields,
             Source source
         ) {
-            super(sessionId, index, indexPattern, inputDataType, inputPage, extractFields, source);
-            this.matchField = matchField;
+            super(sessionId, index, indexPattern, null, inputPage, extractFields, source);
+            this.matchFields = matchFields;
         }
     }
 
     protected static class TransportRequest extends AbstractLookupService.TransportRequest {
-        private final String matchField;
+        private final List<LookupFromIndexOperator.MatchConfig> matchFields;
 
         TransportRequest(
             String sessionId,
@@ -143,11 +147,11 @@ protected static class TransportRequest extends AbstractLookupService.TransportR
             Page inputPage,
             Page toRelease,
             List<NamedExpression> extractFields,
-            String matchField,
+            List<LookupFromIndexOperator.MatchConfig> matchFields,
             Source source
         ) {
             super(sessionId, shardId, indexPattern, inputDataType, inputPage, toRelease, extractFields, source);
-            this.matchField = matchField;
+            this.matchFields = matchFields;
         }
 
         static TransportRequest readFrom(StreamInput in, BlockFactory blockFactory) throws IOException {
@@ -181,6 +185,10 @@ static TransportRequest readFrom(StreamInput in, BlockFactory blockFactory) thro
                 String sourceText = in.readString();
                 source = new Source(source.source(), sourceText);
             }
+            List<LookupFromIndexOperator.MatchConfig> matchFields = null;
+            if (in.getTransportVersion().onOrAfter(TransportVersions.ESQL_LOOKUP_JOIN_ON_MANY_FIELDS)) {
+                matchFields = in.readCollectionAsList(LookupFromIndexOperator.MatchConfig::new);
+            }
             TransportRequest result = new TransportRequest(
                 sessionId,
                 shardId,
@@ -189,7 +197,7 @@ static TransportRequest readFrom(StreamInput in, BlockFactory blockFactory) thro
                 inputPage,
                 inputPage,
                 extractFields,
-                matchField,
+                matchFields,
                 source
             );
             result.setParentTask(parentTaskId);
@@ -213,18 +221,21 @@ public void writeTo(StreamOutput out) throws IOException {
             out.writeWriteable(inputPage);
             PlanStreamOutput planOut = new PlanStreamOutput(out, null);
             planOut.writeNamedWriteableCollection(extractFields);
-            out.writeString(matchField);
+            out.writeString(matchFields.get(0).fieldName().string());
             if (out.getTransportVersion().onOrAfter(TransportVersions.V_8_17_0)) {
                 source.writeTo(planOut);
             }
             if (out.getTransportVersion().onOrAfter(TransportVersions.ESQL_LOOKUP_JOIN_SOURCE_TEXT)) {
                 out.writeString(source.text());
             }
+            if (out.getTransportVersion().onOrAfter(TransportVersions.ESQL_LOOKUP_JOIN_ON_MANY_FIELDS)) {
+                out.writeCollection(matchFields, (o, matchConfig) -> matchConfig.writeTo(o));
+            }
         }
 
         @Override
         protected String extraDescription() {
-            return " ,match_field=" + matchField;
+            return " ,match_fields=" + matchFields.stream().map(x -> x.fieldName().string()).collect(Collectors.joining(", "));
         }
     }
 

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/parser/LogicalPlanBuilder.java

@@ -634,11 +634,6 @@ public PlanFactory visitJoinCommand(EsqlBaseParser.JoinCommandContext ctx) {
             }
         }
 
-        var matchFieldsCount = joinFields.size();
-        if (matchFieldsCount > 1) {
-            throw new ParsingException(source, "JOIN ON clause only supports one field at the moment, found [{}]", matchFieldsCount);
-        }
-
         return p -> {
             boolean hasRemotes = p.anyMatch(node -> {
                 if (node instanceof UnresolvedRelation r) {

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/LocalExecutionPlanner.java

@@ -769,48 +769,33 @@ private PhysicalOperation planLookupJoin(LookupJoinExec join, LocalExecutionPlan
         if (join.leftFields().size() != join.rightFields().size()) {
             throw new IllegalArgumentException("can't plan [" + join + "]: mismatching left and right field count");
         }
-        List<MatchConfig> matchFields = new ArrayList<>(join.leftFields().size());
+        List<LookupFromIndexOperator.MatchConfig> matchFields = new ArrayList<>(join.leftFields().size());
         for (int i = 0; i < join.leftFields().size(); i++) {
             TypedAttribute left = (TypedAttribute) join.leftFields().get(i);
             FieldAttribute right = (FieldAttribute) join.rightFields().get(i);
             Layout.ChannelAndType input = source.layout.get(left.id());
             if (input == null) {
                 throw new IllegalArgumentException("can't plan [" + join + "][" + left + "]");
             }
-            matchFields.add(new MatchConfig(right, input));
+            matchFields.add(new LookupFromIndexOperator.MatchConfig(right, input));
         }
-        if (matchFields.size() != 1) {
-            throw new IllegalArgumentException("can't plan [" + join + "]: multiple join predicates are not supported");
-        }
-        // TODO support multiple match fields, and support more than equality predicates
-        MatchConfig matchConfig = matchFields.getFirst();
 
         return source.with(
             new LookupFromIndexOperator.Factory(
+                matchFields,
                 sessionId,
                 parentTask,
                 context.queryPragmas().enrichMaxWorkers(),
-                matchConfig.channel(),
                 ctx -> lookupFromIndexService,
-                matchConfig.type(),
                 localSourceExec.indexPattern(),
                 indexName,
-                matchConfig.fieldName(),
                 join.addedFields().stream().map(f -> (NamedExpression) f).toList(),
                 join.source()
             ),
             layout
         );
     }
 
-    private record MatchConfig(FieldAttribute.FieldName fieldName, int channel, DataType type) {
-        private MatchConfig(FieldAttribute match, Layout.ChannelAndType input) {
-            // TODO: Using exactAttribute was supposed to handle TEXT fields with KEYWORD subfields - but we don't allow these in lookup
-            // indices, so the call to exactAttribute looks redundant now.
-            this(match.exactAttribute().fieldName(), input.channel(), input.type());
-        }
-    }
-
     private PhysicalOperation planLocal(LocalSourceExec localSourceExec, LocalExecutionPlannerContext context) {
         Layout.Builder layout = new Layout.Builder();
         layout.append(localSourceExec.output());