[8.x] ESQL: Ignore multivalued key columns in lookup index on JOIN (#120726)

x-pack/plugin/esql/compute/src/main/java/module-info.java

@@ -35,4 +35,5 @@
     exports org.elasticsearch.compute.operator.mvdedupe;
     exports org.elasticsearch.compute.aggregation.table;
     exports org.elasticsearch.compute.data.sort;
+    exports org.elasticsearch.compute.querydsl.query;
 }

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

@@ -9,6 +9,9 @@
 
 import org.apache.lucene.document.InetAddressPoint;
 import org.apache.lucene.geo.GeoEncodingUtils;
+import org.apache.lucene.search.BooleanClause;
+import org.apache.lucene.search.BooleanQuery;
+import org.apache.lucene.search.MatchAllDocsQuery;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.common.geo.ShapeRelation;
@@ -20,6 +23,8 @@
 import org.elasticsearch.compute.data.FloatBlock;
 import org.elasticsearch.compute.data.IntBlock;
 import org.elasticsearch.compute.data.LongBlock;
+import org.elasticsearch.compute.operator.Warnings;
+import org.elasticsearch.compute.querydsl.query.SingleValueMatchQuery;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.geometry.Geometry;
 import org.elasticsearch.geometry.Point;
@@ -30,6 +35,7 @@
 import org.elasticsearch.index.mapper.RangeFieldMapper;
 import org.elasticsearch.index.query.SearchExecutionContext;
 
+import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.function.IntFunction;
@@ -38,10 +44,14 @@
  * Generates a list of Lucene queries based on the input block.
  */
 public abstract class QueryList {
+    protected final SearchExecutionContext searchExecutionContext;
+    protected final MappedFieldType field;
     protected final Block block;
     protected final boolean onlySingleValues;
 
-    protected QueryList(Block block, boolean onlySingleValues) {
+    protected QueryList(MappedFieldType field, SearchExecutionContext searchExecutionContext, Block block, boolean onlySingleValues) {
+        this.searchExecutionContext = searchExecutionContext;
+        this.field = field;
         this.block = block;
         this.onlySingleValues = onlySingleValues;
     }
@@ -59,11 +69,52 @@ int getPositionCount() {
      */
     public abstract QueryList onlySingleValues();
 
+    final Query getQuery(int position) {
+        final int valueCount = block.getValueCount(position);
+        if (onlySingleValues && valueCount != 1) {
+            return null;
+        }
+        final int firstValueIndex = block.getFirstValueIndex(position);
+
+        Query query = doGetQuery(position, firstValueIndex, valueCount);
+
+        if (onlySingleValues) {
+            query = wrapSingleValueQuery(query);
+        }
+
+        return query;
+    }
+
     /**
      * Returns the query at the given position.
      */
     @Nullable
-    abstract Query getQuery(int position);
+    abstract Query doGetQuery(int position, int firstValueIndex, int valueCount);
+
+    private Query wrapSingleValueQuery(Query query) {
+        SingleValueMatchQuery singleValueQuery = new SingleValueMatchQuery(
+            searchExecutionContext.getForField(field, MappedFieldType.FielddataOperation.SEARCH),
+            // Not emitting warnings for multivalued fields not matching
+            Warnings.NOOP_WARNINGS
+        );
+
+        Query rewrite = singleValueQuery;
+        try {
+            rewrite = singleValueQuery.rewrite(searchExecutionContext.searcher());
+            if (rewrite instanceof MatchAllDocsQuery) {
+                // nothing to filter
+                return query;
+            }
+        } catch (IOException e) {
+            // ignore
+            // TODO: Should we do something with the exception?
+        }
+
+        BooleanQuery.Builder builder = new BooleanQuery.Builder();
+        builder.add(query, BooleanClause.Occur.FILTER);
+        builder.add(rewrite, BooleanClause.Occur.FILTER);
+        return builder.build();
+    }
 
     /**
      * Returns a list of term queries for the given field and the input block
@@ -146,8 +197,6 @@ public static QueryList geoShapeQueryList(MappedFieldType field, SearchExecution
     }
 
     private static class TermQueryList extends QueryList {
-        private final MappedFieldType field;
-        private final SearchExecutionContext searchExecutionContext;
         private final IntFunction<Object> blockValueReader;
 
         private TermQueryList(
@@ -157,9 +206,7 @@ private TermQueryList(
             boolean onlySingleValues,
             IntFunction<Object> blockValueReader
         ) {
-            super(block, onlySingleValues);
-            this.field = field;
-            this.searchExecutionContext = searchExecutionContext;
+            super(field, searchExecutionContext, block, onlySingleValues);
             this.blockValueReader = blockValueReader;
         }
 
@@ -169,19 +216,14 @@ public TermQueryList onlySingleValues() {
         }
 
         @Override
-        Query getQuery(int position) {
-            final int count = block.getValueCount(position);
-            if (onlySingleValues && count != 1) {
-                return null;
-            }
-            final int first = block.getFirstValueIndex(position);
-            return switch (count) {
+        Query doGetQuery(int position, int firstValueIndex, int valueCount) {
+            return switch (valueCount) {
                 case 0 -> null;
-                case 1 -> field.termQuery(blockValueReader.apply(first), searchExecutionContext);
+                case 1 -> field.termQuery(blockValueReader.apply(firstValueIndex), searchExecutionContext);
                 default -> {
-                    final List<Object> terms = new ArrayList<>(count);
-                    for (int i = 0; i < count; i++) {
-                        final Object value = blockValueReader.apply(first + i);
+                    final List<Object> terms = new ArrayList<>(valueCount);
+                    for (int i = 0; i < valueCount; i++) {
+                        final Object value = blockValueReader.apply(firstValueIndex + i);
                         terms.add(value);
                     }
                     yield field.termsQuery(terms, searchExecutionContext);
@@ -192,8 +234,6 @@ Query getQuery(int position) {
 
     private static class GeoShapeQueryList extends QueryList {
         private final BytesRef scratch = new BytesRef();
-        private final MappedFieldType field;
-        private final SearchExecutionContext searchExecutionContext;
         private final IntFunction<Geometry> blockValueReader;
         private final IntFunction<Query> shapeQuery;
 
@@ -203,10 +243,8 @@ private GeoShapeQueryList(
             Block block,
             boolean onlySingleValues
         ) {
-            super(block, onlySingleValues);
+            super(field, searchExecutionContext, block, onlySingleValues);
 
-            this.field = field;
-            this.searchExecutionContext = searchExecutionContext;
             this.blockValueReader = blockToGeometry(block);
             this.shapeQuery = shapeQuery();
         }
@@ -217,15 +255,10 @@ public GeoShapeQueryList onlySingleValues() {
         }
 
         @Override
-        Query getQuery(int position) {
-            final int count = block.getValueCount(position);
-            if (onlySingleValues && count != 1) {
-                return null;
-            }
-            final int first = block.getFirstValueIndex(position);
-            return switch (count) {
+        Query doGetQuery(int position, int firstValueIndex, int valueCount) {
+            return switch (valueCount) {
                 case 0 -> null;
-                case 1 -> shapeQuery.apply(first);
+                case 1 -> shapeQuery.apply(firstValueIndex);
                 // TODO: support multiple values
                 default -> throw new IllegalArgumentException("can't read multiple Geometry values from a single position");
             };

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/querydsl/query/SingleValueMatchQuery.java → x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/querydsl/query/SingleValueMatchQuery.java

@@ -5,7 +5,7 @@
  * 2.0.
  */
 
-package org.elasticsearch.xpack.esql.querydsl.query;
+package org.elasticsearch.compute.querydsl.query;
 
 import org.apache.lucene.index.DocValues;
 import org.apache.lucene.index.LeafReaderContext;
@@ -39,7 +39,7 @@
 /**
  * Finds all fields with a single-value. If a field has a multi-value, it emits a {@link Warnings}.
  */
-final class SingleValueMatchQuery extends Query {
+public final class SingleValueMatchQuery extends Query {
 
     /**
      * Choose a big enough value so this approximation never drives the iteration.
@@ -52,7 +52,7 @@ final class SingleValueMatchQuery extends Query {
     private final IndexFieldData<?> fieldData;
     private final Warnings warnings;
 
-    SingleValueMatchQuery(IndexFieldData<?> fieldData, Warnings warnings) {
+    public SingleValueMatchQuery(IndexFieldData<?> fieldData, Warnings warnings) {
         this.fieldData = fieldData;
         this.warnings = warnings;
     }