ESQL: Ignore multivalued key columns in lookup index on JOIN (#120726) (#121037)

Fixes #118780

Second part of #120519

In the first PR, we avoid matching multivalue keys in lookup when they come from the query.
Now, we avoid matching multivalues when the lookup index has multivalues in the key column.

Author: ivancea

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 renamed to 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;
     }