[WIP]Add support for LOOKUP JOIN on an expression

docs/changelog/133166.yaml

@@ -0,0 +1,6 @@
+pr: 133166
+summary: Improve Expanding Lookup Join performance by pushing a filter to the right
+  side of the lookup join
+area: ES|QL
+type: enhancement
+issues: []

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

@@ -357,6 +357,8 @@ static TransportVersion def(int id) {
     public static final TransportVersion ESQL_SAMPLE_OPERATOR_STATUS = def(9_127_0_00);
     public static final TransportVersion ALLOCATION_DECISION_NOT_PREFERRED = def(9_145_0_00);
     public static final TransportVersion ESQL_QUALIFIERS_IN_ATTRIBUTES = def(9_146_0_00);
+    public static final TransportVersion ESQL_LOOKUP_JOIN_PRE_JOIN_FILTER = def(9_147_0_00);
+    public static final TransportVersion ESQL_LOOKUP_JOIN_ON_EXPRESSION = def(9_148_0_00);
 
     /*
      * STOP! READ THIS FIRST! No, really,

server/src/main/java/org/elasticsearch/index/mapper/MappedFieldType.java

@@ -387,7 +387,7 @@ public Query regexpQuery(
     }
 
     /**
-     * Returns a Lucine pushable Query for the current field
+     * Returns a Lucene pushable Query for the current field
      * For now can only be AutomatonQuery or MatchAllDocsQuery() or MatchNoDocsQuery()
      */
     public Query automatonQuery(

test/external-modules/esql-heap-attack/src/javaRestTest/java/org/elasticsearch/xpack/esql/heap_attack/HeapAttackIT.java

@@ -698,26 +698,42 @@ private Map<String, Object> fetchMvLongs() throws IOException {
     public void testLookupExplosion() throws IOException {
         int sensorDataCount = 400;
         int lookupEntries = 10000;
-        Map<?, ?> map = lookupExplosion(sensorDataCount, lookupEntries, 1);
+        Map<?, ?> map = lookupExplosion(sensorDataCount, lookupEntries, 1, lookupEntries);
         assertMap(map, matchesMap().extraOk().entry("values", List.of(List.of(sensorDataCount * lookupEntries))));
     }
 
     public void testLookupExplosionManyFields() throws IOException {
         int sensorDataCount = 400;
         int lookupEntries = 1000;
         int joinFieldsCount = 990;
-        Map<?, ?> map = lookupExplosion(sensorDataCount, lookupEntries, joinFieldsCount);
+        Map<?, ?> map = lookupExplosion(sensorDataCount, lookupEntries, joinFieldsCount, lookupEntries);
         assertMap(map, matchesMap().extraOk().entry("values", List.of(List.of(sensorDataCount * lookupEntries))));
     }
 
     public void testLookupExplosionManyMatchesManyFields() throws IOException {
         // 1500, 10000 is enough locally, but some CI machines need more.
-        assertCircuitBreaks(attempt -> lookupExplosion(attempt * 1500, 10000, 30));
+        int lookupEntries = 10000;
+        assertCircuitBreaks(attempt -> lookupExplosion(attempt * 1500, lookupEntries, 30, lookupEntries));
     }
 
     public void testLookupExplosionManyMatches() throws IOException {
         // 1500, 10000 is enough locally, but some CI machines need more.
-        assertCircuitBreaks(attempt -> lookupExplosion(attempt * 1500, 10000, 1));
+        int lookupEntries = 10000;
+        assertCircuitBreaks(attempt -> lookupExplosion(attempt * 1500, lookupEntries, 1, lookupEntries));
+    }
+
+    public void testLookupExplosionManyMatchesFiltered() throws IOException {
+        // This test will only work with the expanding join optimization
+        // that pushes the filter to the right side of the lookup.
+        // Without the optimization, it will fail with circuit_breaking_exception
+        int sensorDataCount = 10000;
+        int lookupEntries = 10000;
+        int reductionFactor = 1000; // reduce the number of matches by this factor
+        // lookupEntries % reductionFactor must be 0 to ensure the number of rows returned matches the expected value
+        assertTrue(0 == lookupEntries % reductionFactor);
+        Map<?, ?> map = lookupExplosion(sensorDataCount, lookupEntries, 1, lookupEntries / reductionFactor);
+        assertMap(map, matchesMap().extraOk().entry("values", List.of(List.of(sensorDataCount * lookupEntries / reductionFactor))));
+
     }
 
     public void testLookupExplosionNoFetch() throws IOException {
@@ -744,7 +760,8 @@ public void testLookupExplosionBigStringManyMatches() throws IOException {
         assertCircuitBreaks(attempt -> lookupExplosionBigString(attempt * 500, 1));
     }
 
-    private Map<String, Object> lookupExplosion(int sensorDataCount, int lookupEntries, int joinFieldsCount) throws IOException {
+    private Map<String, Object> lookupExplosion(int sensorDataCount, int lookupEntries, int joinFieldsCount, int lookupEntriesToKeep)
+        throws IOException {
         try {
             lookupExplosionData(sensorDataCount, lookupEntries, joinFieldsCount);
             StringBuilder query = startQuery();
@@ -755,7 +772,14 @@ private Map<String, Object> lookupExplosion(int sensorDataCount, int lookupEntri
                 }
                 query.append("id").append(i);
             }
-            query.append(" | STATS COUNT(location)\"}");
+            if (lookupEntries != lookupEntriesToKeep) {
+                // add a filter to reduce the number of matches
+                // we add both a Lucene pushable filter and a non-pushable filter
+                // this is to make sure that even if there are non-pushable filters the pushable filters is still applied
+                query.append(" | WHERE ABS(filter_key) > -1 AND filter_key < ").append(lookupEntriesToKeep);
+
+            }
+            query.append(" | STATS COUNT(location) | LIMIT 100\"}");
             return responseAsMap(query(query.toString(), null));
         } finally {
             deleteIndex("sensor_data");
@@ -1038,7 +1062,8 @@ private void initSensorLookup(int lookupEntries, int sensorCount, IntFunction<St
             createIndexBuilder.append("\"id").append(i).append("\": { \"type\": \"long\" },");
         }
         createIndexBuilder.append("""
-                    "location": { "type": "geo_point" }
+                    "location": { "type": "geo_point" },
+                    "filter_key": { "type": "integer" }
                 }
             }""");
         CreateIndexResponse response = createIndex(
@@ -1058,7 +1083,7 @@ private void initSensorLookup(int lookupEntries, int sensorCount, IntFunction<St
                 data.append(String.format(Locale.ROOT, "\"id%d\":%d, ", j, sensor));
             }
             data.append(String.format(Locale.ROOT, """
-                "location": "POINT(%s)"}\n""", location.apply(sensor)));
+                "location": "POINT(%s)", "filter_key": %d}\n""", location.apply(sensor), i));
             if (i % docsPerBulk == docsPerBulk - 1) {
                 bulk("sensor_lookup", data.toString());
                 data.setLength(0);

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

@@ -125,7 +125,7 @@ public final Query getQuery(int position) {
      * Returns the query at the given position.
      */
     @Nullable
-    abstract Query doGetQuery(int position, int firstValueIndex, int valueCount);
+    public abstract Query doGetQuery(int position, int firstValueIndex, int valueCount);
 
     private Query wrapSingleValueQuery(Query query) {
         assert onlySingleValueParams != null : "Requested to wrap single value query without single value params";
@@ -159,13 +159,8 @@ private Query wrapSingleValueQuery(Query query) {
      * using only the {@link ElementType} of the {@link Block} to determine the
      * query.
      */
-    public static QueryList rawTermQueryList(
-        MappedFieldType field,
-        SearchExecutionContext searchExecutionContext,
-        AliasFilter aliasFilter,
-        Block block
-    ) {
-        IntFunction<Object> blockToJavaObject = switch (block.elementType()) {
+    public static IntFunction<Object> createBlockValueReader(Block block) {
+        return switch (block.elementType()) {
             case BOOLEAN -> {
                 BooleanBlock booleanBlock = (BooleanBlock) block;
                 yield booleanBlock::getBoolean;
@@ -196,7 +191,20 @@ public static QueryList rawTermQueryList(
             case AGGREGATE_METRIC_DOUBLE -> throw new IllegalArgumentException("can't read values from [aggregate metric double] block");
             case UNKNOWN -> throw new IllegalArgumentException("can't read values from [" + block + "]");
         };
-        return new TermQueryList(field, searchExecutionContext, aliasFilter, block, null, blockToJavaObject);
+    }
+
+    /**
+     * Returns a list of term queries for the given field and the input block
+     * using only the {@link ElementType} of the {@link Block} to determine the
+     * query.
+     */
+    public static QueryList rawTermQueryList(
+        MappedFieldType field,
+        SearchExecutionContext searchExecutionContext,
+        AliasFilter aliasFilter,
+        Block block
+    ) {
+        return new TermQueryList(field, searchExecutionContext, aliasFilter, block, null, createBlockValueReader(block));
     }
 
     /**
@@ -297,7 +305,7 @@ public TermQueryList onlySingleValues(Warnings warnings, String multiValueWarnin
         }
 
         @Override
-        Query doGetQuery(int position, int firstValueIndex, int valueCount) {
+        public Query doGetQuery(int position, int firstValueIndex, int valueCount) {
             return switch (valueCount) {
                 case 0 -> null;
                 case 1 -> field.termQuery(blockValueReader.apply(firstValueIndex), searchExecutionContext);
@@ -360,7 +368,7 @@ public DateNanosQueryList onlySingleValues(Warnings warnings, String multiValueW
         }
 
         @Override
-        Query doGetQuery(int position, int firstValueIndex, int valueCount) {
+        public Query doGetQuery(int position, int firstValueIndex, int valueCount) {
             return switch (valueCount) {
                 case 0 -> null;
                 case 1 -> dateFieldType.equalityQuery(blockValueReader.apply(firstValueIndex), searchExecutionContext);
@@ -412,7 +420,7 @@ public GeoShapeQueryList onlySingleValues(Warnings warnings, String multiValueWa
         }
 
         @Override
-        Query doGetQuery(int position, int firstValueIndex, int valueCount) {
+        public Query doGetQuery(int position, int firstValueIndex, int valueCount) {
             return switch (valueCount) {
                 case 0 -> null;
                 case 1 -> shapeQuery.apply(firstValueIndex);
@@ -453,5 +461,5 @@ private IntFunction<Query> shapeQuery() {
         }
     }
 
-    protected record OnlySingleValueParams(Warnings warnings, String multiValueWarningMessage) {}
+    public record OnlySingleValueParams(Warnings warnings, String multiValueWarningMessage) {}
 }

x-pack/plugin/esql/compute/test/src/main/java/org/elasticsearch/compute/test/AnyOperatorTestCase.java

@@ -87,7 +87,7 @@ protected final Operator.OperatorFactory simple() {
     /**
      * Makes sure the description of {@link #simple} matches the {@link #expectedDescriptionOfSimple}.
      */
-    public final void testSimpleDescription() {
+    public void testSimpleDescription() {
         Operator.OperatorFactory factory = simple();
         String description = factory.describe();
         assertThat(description, expectedDescriptionOfSimple());