Improve Expanding Lookup Join performance by pushing a filter to the right side of the lookup join

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

@@ -369,6 +369,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion SCRIPT_RESCORER = def(9_143_0_00);
     public static final TransportVersion ESQL_LOOKUP_OPERATOR_EMITTED_ROWS = def(9_144_0_00);
     public static final TransportVersion ALLOCATION_DECISION_NOT_PREFERRED = def(9_145_0_00);
+    public static final TransportVersion ESQL_LOOKUP_JOIN_PRE_JOIN_FILTER = def(9_146_0_00);
 
     /*
      * STOP! READ THIS FIRST! No, really,

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 Lucine 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/EnrichQuerySourceOperator.java

@@ -159,7 +159,7 @@ Page buildPage(int positions, IntVector.Builder positionsBuilder, IntVector.Buil
         return page;
     }
 
-    private Query nextQuery() {
+    private Query nextQuery() throws IOException {
         ++queryPosition;
         while (isFinished() == false) {
             Query query = queryList.getQuery(queryPosition);

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());