Looking join on multiple fields WIP

Author: julian-elastic

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/lookup/EnrichQuerySourceOperator.java

@@ -37,7 +37,7 @@
  */
 public final class EnrichQuerySourceOperator extends SourceOperator {
     private final BlockFactory blockFactory;
-    private final QueryList queryList;
+    private final LookupEnrichQueryGenerator queryList;
     private int queryPosition = -1;
     private final ShardContext shardContext;
     private final IndexReader indexReader;
@@ -51,7 +51,7 @@ public final class EnrichQuerySourceOperator extends SourceOperator {
     public EnrichQuerySourceOperator(
         BlockFactory blockFactory,
         int maxPageSize,
-        QueryList queryList,
+        LookupEnrichQueryGenerator queryList,
         ShardContext shardContext,
         Warnings warnings
     ) {

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/lookup/ExpressionQueryList.java

@@ -0,0 +1,54 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.compute.operator.lookup;
+
+import org.apache.lucene.search.BooleanClause;
+import org.apache.lucene.search.BooleanQuery;
+import org.apache.lucene.search.Query;
+
+import java.util.List;
+
+public class ExpressionQueryList implements LookupEnrichQueryGenerator {
+    List<QueryList> queryLists;
+
+    public ExpressionQueryList(List<QueryList> queryLists) {
+        if (queryLists.size() < 2) {
+            throw new IllegalArgumentException("ExpressionQueryList must have at least two QueryLists");
+        }
+        this.queryLists = queryLists;
+    }
+
+    @Override
+    public Query getQuery(int position) {
+        // JULIAN: Do we have a guarantee that the positions will match across all QueryLists?
+        // for now we only support AND of the queries in the lists
+        BooleanQuery.Builder builder = new BooleanQuery.Builder();
+        for (QueryList queryList : queryLists) {
+            Query q = queryList.getQuery(position);
+            builder.add(q, BooleanClause.Occur.MUST);
+
+        }
+        return builder.build();
+    }
+
+    @Override
+    public int getPositionCount() {
+        int positionCount = queryLists.get(0).getPositionCount();
+        for (QueryList queryList : queryLists) {
+            if (queryList.getPositionCount() != positionCount) {
+                throw new IllegalArgumentException(
+                    "All QueryLists must have the same position count, expected: "
+                        + positionCount
+                        + ", but got: "
+                        + queryList.getPositionCount()
+                );
+            }
+        }
+        return queryLists.get(0).getPositionCount();
+    }
+}

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/lookup/LookupEnrichQueryGenerator.java

@@ -0,0 +1,23 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.compute.operator.lookup;
+
+import org.apache.lucene.search.Query;
+import org.elasticsearch.core.Nullable;
+
+public interface LookupEnrichQueryGenerator {
+
+    /**
+     * Returns the query at the given position.
+     */
+    @Nullable
+    Query getQuery(int position);
+
+    int getPositionCount();
+
+}

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/lookup/QueryList.java

@@ -49,7 +49,7 @@
 /**
  * Generates a list of Lucene queries based on the input block.
  */
-public abstract class QueryList {
+public abstract class QueryList implements LookupEnrichQueryGenerator {
     protected final SearchExecutionContext searchExecutionContext;
     protected final AliasFilter aliasFilter;
     protected final MappedFieldType field;
@@ -74,7 +74,7 @@ protected QueryList(
     /**
      * Returns the number of positions in this query list
      */
-    int getPositionCount() {
+    public int getPositionCount() {
         return block.getPositionCount();
     }
 
@@ -87,7 +87,7 @@ int getPositionCount() {
      */
     public abstract QueryList onlySingleValues(Warnings warnings, String multiValueWarningMessage);
 
-    final Query getQuery(int position) {
+    public final Query getQuery(int position) {
         final int valueCount = block.getValueCount(position);
         if (onlySingleValueParams != null && valueCount != 1) {
             if (valueCount > 1) {

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

@@ -43,6 +43,7 @@
 import org.elasticsearch.compute.operator.ProjectOperator;
 import org.elasticsearch.compute.operator.Warnings;
 import org.elasticsearch.compute.operator.lookup.EnrichQuerySourceOperator;
+import org.elasticsearch.compute.operator.lookup.LookupEnrichQueryGenerator;
 import org.elasticsearch.compute.operator.lookup.MergePositionsOperator;
 import org.elasticsearch.compute.operator.lookup.QueryList;
 import org.elasticsearch.core.AbstractRefCounted;
@@ -191,7 +192,7 @@ public ThreadContext getThreadContext() {
     /**
      * Build a list of queries to perform inside the actual lookup.
      */
-    protected abstract QueryList queryList(
+    protected abstract LookupEnrichQueryGenerator queryList(
         T request,
         SearchExecutionContext context,
         AliasFilter aliasFilter,
@@ -271,17 +272,20 @@ protected void sendChildRequest(
     }
 
     private void doLookup(T request, CancellableTask task, ActionListener<List<Page>> listener) {
-        Block inputBlock = request.inputPage.getBlock(0);
-        if (inputBlock.areAllValuesNull()) {
-            List<Page> nullResponse = mergePages
-                ? List.of(createNullResponse(request.inputPage.getPositionCount(), request.extractFields))
-                : List.of();
-            listener.onResponse(nullResponse);
-            return;
-        }
-        final List<Releasable> releasables = new ArrayList<>(6);
         boolean started = false;
+        final List<Releasable> releasables = new ArrayList<>(6);
         try {
+            for (int j = 0; j < request.inputPage.getBlockCount(); j++) {
+                Block inputBlock = request.inputPage.getBlock(j);
+                if (inputBlock.areAllValuesNull()) {
+                    List<Page> nullResponse = mergePages
+                        ? List.of(createNullResponse(request.inputPage.getPositionCount(), request.extractFields))
+                        : List.of();
+                    listener.onResponse(nullResponse);
+                    return;
+                }
+            }
+
             var projectState = projectResolver.getProjectState(clusterService.state());
             AliasFilter aliasFilter = indicesService.buildAliasFilter(
                 projectState,
@@ -305,6 +309,8 @@ private void doLookup(T request, CancellableTask task, ActionListener<List<Page>
             final int[] mergingChannels = IntStream.range(0, request.extractFields.size()).map(i -> i + 2).toArray();
             final Operator finishPages;
             final OrdinalBytesRefBlock ordinalsBytesRefBlock;
+            // JULIAN: SHOULD WE DO THE NEXT CODE FOR EACH BLOCK IN THE PAGE?
+            Block inputBlock = request.inputPage.getBlock(0);
             if (mergePages  // TODO fix this optimization for Lookup.
                 && inputBlock instanceof BytesRefBlock bytesRefBlock
                 && (ordinalsBytesRefBlock = bytesRefBlock.asOrdinals()) != null) {
@@ -334,7 +340,7 @@ private void doLookup(T request, CancellableTask task, ActionListener<List<Page>
                 request.source.source().getColumnNumber(),
                 request.source.text()
             );
-            QueryList queryList = queryList(
+            LookupEnrichQueryGenerator queryList = queryList(
                 request,
                 shardContext.executionContext,
                 aliasFilter,

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

@@ -25,6 +25,7 @@
 import org.elasticsearch.compute.data.BlockStreamInput;
 import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.Warnings;
+import org.elasticsearch.compute.operator.lookup.LookupEnrichQueryGenerator;
 import org.elasticsearch.compute.operator.lookup.QueryList;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.Releasables;
@@ -110,7 +111,7 @@ protected TransportRequest transportRequest(EnrichLookupService.Request request,
     }
 
     @Override
-    protected QueryList queryList(
+    protected LookupEnrichQueryGenerator queryList(
         TransportRequest request,
         SearchExecutionContext context,
         AliasFilter aliasFilter,

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

@@ -146,9 +146,8 @@ public LookupFromIndexOperator(
 
     @Override
     protected void performAsync(Page inputPage, ActionListener<OngoingJoin> listener) {
-        // 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?
+        // JULIAN: should I be getting multiple blocks, and send them using the LookupFromIndexService.Request
+        // is the totalTerms supposed to be the total number of terms in all blocks combined?
         Block[] inputBlockArray = new Block[matchFields.size()];
         for (int i = 0; i < matchFields.size(); i++) {
             MatchConfig matchField = matchFields.get(i);

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

@@ -20,6 +20,8 @@
 import org.elasticsearch.compute.data.BlockStreamInput;
 import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.Warnings;
+import org.elasticsearch.compute.operator.lookup.ExpressionQueryList;
+import org.elasticsearch.compute.operator.lookup.LookupEnrichQueryGenerator;
 import org.elasticsearch.compute.operator.lookup.QueryList;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.Releasables;
@@ -38,6 +40,7 @@
 import org.elasticsearch.xpack.esql.io.stream.PlanStreamOutput;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
 import java.util.stream.Collectors;
@@ -91,22 +94,36 @@ protected TransportRequest transportRequest(LookupFromIndexService.Request reque
     }
 
     @Override
-    protected QueryList queryList(
+    protected LookupEnrichQueryGenerator queryList(
         TransportRequest request,
         SearchExecutionContext context,
         AliasFilter aliasFilter,
         Block inputBlock,
         @Nullable DataType inputDataType,
         Warnings warnings
     ) {
-        // 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");
+        if (request.matchFields.size() == 1) {
+            return termQueryList(
+                context.getFieldType(request.matchFields.get(0).fieldName().string()),
+                context,
+                aliasFilter,
+                inputBlock,
+                inputDataType
+            ).onlySingleValues(warnings, "LOOKUP JOIN encountered multi-value");
+        }
+        List<QueryList> queryLists = new ArrayList<>();
+        for (int i = 0; i < request.matchFields.size(); i++) {
+            LookupFromIndexOperator.MatchConfig matchField = request.matchFields.get(i);
+            QueryList q = termQueryList(
+                context.getFieldType(matchField.fieldName().string()),
+                context,
+                aliasFilter,
+                request.inputPage.getBlock(i),
+                matchField.type()
+            ).onlySingleValues(warnings, "LOOKUP JOIN encountered multi-value");
+            queryLists.add(q);
+        }
+        return new ExpressionQueryList(queryLists);
     }
 
     @Override

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/enrich/LookupFromIndexOperatorTests.java

@@ -142,16 +142,17 @@ protected Operator.OperatorFactory simple(SimpleOptions options) {
             new ReferenceAttribute(Source.EMPTY, "lkwd", DataType.KEYWORD),
             new ReferenceAttribute(Source.EMPTY, "lint", DataType.INTEGER)
         );
+        List<LookupFromIndexOperator.MatchConfig> matchFields = List.of(
+            new LookupFromIndexOperator.MatchConfig(matchField, inputChannel, inputDataType)
+        );
         return new LookupFromIndexOperator.Factory(
+            matchFields,
             sessionId,
             parentTask,
             maxOutstandingRequests,
-            inputChannel,
             this::lookupService,
-            inputDataType,
             lookupIndex,
             lookupIndex,
-            matchField,
             loadFields,
             Source.EMPTY
         );
@@ -165,7 +166,7 @@ protected Matcher<String> expectedDescriptionOfSimple() {
     @Override
     protected Matcher<String> expectedToStringOfSimple() {
         return matchesPattern(
-            "LookupOperator\\[index=idx input_type=LONG match_field=match load_fields=\\[lkwd\\{r}#\\d+, lint\\{r}#\\d+] inputChannel=0]"
+            "LookupOperator\\[index=idx load_fields=\\[lkwd\\{r}#\\d+, lint\\{r}#\\d+] input_type=LONG match_field=match inputChannel=0]"
         );
     }