ESQL: LOOKUP JOIN with expressions

docs/changelog/134098.yaml

@@ -0,0 +1,5 @@
+pr: 134098
+summary: LOOKUP JOIN with expressions
+area: ES|QL
+type: enhancement
+issues: []

docs/reference/query-languages/esql/_snippets/commands/types/lookup-join.md

@@ -5,17 +5,17 @@
 | field from the left index | field from the lookup index | 
 | --- | --- | 
 | boolean | boolean |
-| byte | half_float, float, double, scaled_float, byte, short, integer, long |
+| byte | byte, short, integer, long, half_float, float, double, scaled_float |
 | date | date |
 | date_nanos | date_nanos |
 | double | half_float, float, double, scaled_float, byte, short, integer, long |
 | float | half_float, float, double, scaled_float, byte, short, integer, long |
 | half_float | half_float, float, double, scaled_float, byte, short, integer, long |
-| integer | half_float, float, double, scaled_float, byte, short, integer, long |
+| integer | byte, short, integer, long, half_float, float, double, scaled_float |
 | ip | ip |
 | keyword | keyword |
-| long | half_float, float, double, scaled_float, byte, short, integer, long |
+| long | byte, short, integer, long, half_float, float, double, scaled_float |
 | scaled_float | half_float, float, double, scaled_float, byte, short, integer, long |
-| short | half_float, float, double, scaled_float, byte, short, integer, long |
+| short | byte, short, integer, long, half_float, float, double, scaled_float |
 | text | keyword |
 

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

@@ -327,6 +327,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion TIMESERIES_DEFAULT_LIMIT = def(9_160_0_00);
     public static final TransportVersion INFERENCE_API_OPENAI_HEADERS = def(9_161_0_00);
     public static final TransportVersion NEW_SEMANTIC_QUERY_INTERCEPTORS = def(9_162_0_00);
+    public static final TransportVersion ESQL_LOOKUP_JOIN_ON_EXPRESSION = def(9_163_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

@@ -566,6 +566,7 @@ private Response query(String query, String filterPath) throws IOException {
                 .setRequestConfig(RequestConfig.custom().setSocketTimeout(Math.toIntExact(TimeValue.timeValueMinutes(6).millis())).build())
                 .setWarningsHandler(WarningsHandler.PERMISSIVE)
         );
+        logger.info("Running query:" + query);
         return runQuery(() -> client().performRequest(request));
     }
 
@@ -735,28 +736,48 @@ private Map<String, Object> fetchMvLongs() throws IOException {
     public void testLookupExplosion() throws IOException {
         int sensorDataCount = 400;
         int lookupEntries = 10000;
-        Map<?, ?> map = lookupExplosion(sensorDataCount, lookupEntries, 1, lookupEntries);
+        Map<?, ?> map = lookupExplosion(sensorDataCount, lookupEntries, 1, lookupEntries, false);
         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, lookupEntries);
+        Map<?, ?> map = lookupExplosion(sensorDataCount, lookupEntries, joinFieldsCount, lookupEntries, false);
+        assertMap(map, matchesMap().extraOk().entry("values", List.of(List.of(sensorDataCount * lookupEntries))));
+    }
+
+    public void testLookupExplosionExpression() throws IOException {
+        int sensorDataCount = 400;
+        int lookupEntries = 10000;
+        Map<?, ?> map = lookupExplosion(sensorDataCount, lookupEntries, 1, lookupEntries, true);
+        assertMap(map, matchesMap().extraOk().entry("values", List.of(List.of(sensorDataCount * lookupEntries))));
+    }
+
+    public void testLookupExplosionManyFieldsExpression() throws IOException {
+        int sensorDataCount = 400;
+        int lookupEntries = 1000;
+        int joinFieldsCount = 399;// only join on 399 columns due to max expression size of 400
+        Map<?, ?> map = lookupExplosion(sensorDataCount, lookupEntries, joinFieldsCount, lookupEntries, true);
         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.
         int lookupEntries = 10000;
-        assertCircuitBreaks(attempt -> lookupExplosion(attempt * 1500, lookupEntries, 30, lookupEntries));
+        assertCircuitBreaks(attempt -> lookupExplosion(attempt * 1500, lookupEntries, 30, lookupEntries, false));
     }
 
     public void testLookupExplosionManyMatches() throws IOException {
         // 1500, 10000 is enough locally, but some CI machines need more.
         int lookupEntries = 10000;
-        assertCircuitBreaks(attempt -> lookupExplosion(attempt * 1500, lookupEntries, 1, lookupEntries));
+        assertCircuitBreaks(attempt -> lookupExplosion(attempt * 1500, lookupEntries, 1, lookupEntries, false));
+    }
+
+    public void testLookupExplosionManyMatchesExpression() throws IOException {
+        int lookupEntries = 10000;
+        assertCircuitBreaks(attempt -> lookupExplosion(attempt * 1500, lookupEntries, 1, lookupEntries, true));
     }
 
     public void testLookupExplosionManyMatchesFiltered() throws IOException {
@@ -768,9 +789,21 @@ public void testLookupExplosionManyMatchesFiltered() throws IOException {
         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);
+        Map<?, ?> map = lookupExplosion(sensorDataCount, lookupEntries, 1, lookupEntries / reductionFactor, false);
         assertMap(map, matchesMap().extraOk().entry("values", List.of(List.of(sensorDataCount * lookupEntries / reductionFactor))));
+    }
 
+    public void testLookupExplosionManyMatchesFilteredExpression() 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, true);
+        assertMap(map, matchesMap().extraOk().entry("values", List.of(List.of(sensorDataCount * lookupEntries / reductionFactor))));
     }
 
     public void testLookupExplosionNoFetch() throws IOException {
@@ -797,17 +830,33 @@ public void testLookupExplosionBigStringManyMatches() throws IOException {
         assertCircuitBreaks(attempt -> lookupExplosionBigString(attempt * 500, 1));
     }
 
-    private Map<String, Object> lookupExplosion(int sensorDataCount, int lookupEntries, int joinFieldsCount, int lookupEntriesToKeep)
-        throws IOException {
+    private Map<String, Object> lookupExplosion(
+        int sensorDataCount,
+        int lookupEntries,
+        int joinFieldsCount,
+        int lookupEntriesToKeep,
+        boolean expressionBasedJoin
+    ) throws IOException {
         try {
-            lookupExplosionData(sensorDataCount, lookupEntries, joinFieldsCount);
+            lookupExplosionData(sensorDataCount, lookupEntries, joinFieldsCount, expressionBasedJoin);
             StringBuilder query = startQuery();
             query.append("FROM sensor_data | LOOKUP JOIN sensor_lookup ON ");
-            for (int i = 0; i < joinFieldsCount; i++) {
-                if (i != 0) {
-                    query.append(",");
+            if (expressionBasedJoin) {
+                for (int i = 0; i < joinFieldsCount; i++) {
+                    if (i != 0) {
+                        query.append(" AND ");
+                    }
+                    query.append("id_left").append(i);
+                    query.append("==");
+                    query.append("id_right").append(i);
+                }
+            } else {
+                for (int i = 0; i < joinFieldsCount; i++) {
+                    if (i != 0) {
+                        query.append(",");
+                    }
+                    query.append("id").append(i);
                 }
-                query.append("id").append(i);
             }
             if (lookupEntries != lookupEntriesToKeep) {
                 // add a filter to reduce the number of matches
@@ -826,7 +875,7 @@ private Map<String, Object> lookupExplosion(int sensorDataCount, int lookupEntri
 
     private Map<String, Object> lookupExplosionNoFetch(int sensorDataCount, int lookupEntries) throws IOException {
         try {
-            lookupExplosionData(sensorDataCount, lookupEntries, 1);
+            lookupExplosionData(sensorDataCount, lookupEntries, 1, false);
             StringBuilder query = startQuery();
             query.append("FROM sensor_data | LOOKUP JOIN sensor_lookup ON id0 | STATS COUNT(*)\"}");
             return responseAsMap(query(query.toString(), null));
@@ -836,14 +885,15 @@ private Map<String, Object> lookupExplosionNoFetch(int sensorDataCount, int look
         }
     }
 
-    private void lookupExplosionData(int sensorDataCount, int lookupEntries, int joinFieldCount) throws IOException {
-        initSensorData(sensorDataCount, 1, joinFieldCount);
-        initSensorLookup(lookupEntries, 1, i -> "73.9857 40.7484", joinFieldCount);
+    private void lookupExplosionData(int sensorDataCount, int lookupEntries, int joinFieldCount, boolean expressionBasedJoin)
+        throws IOException {
+        initSensorData(sensorDataCount, 1, joinFieldCount, expressionBasedJoin);
+        initSensorLookup(lookupEntries, 1, i -> "73.9857 40.7484", joinFieldCount, expressionBasedJoin);
     }
 
     private Map<String, Object> lookupExplosionBigString(int sensorDataCount, int lookupEntries) throws IOException {
         try {
-            initSensorData(sensorDataCount, 1, 1);
+            initSensorData(sensorDataCount, 1, 1, false);
             initSensorLookupString(lookupEntries, 1, i -> {
                 int target = Math.toIntExact(ByteSizeValue.ofMb(1).getBytes());
                 StringBuilder str = new StringBuilder(Math.toIntExact(ByteSizeValue.ofMb(2).getBytes()));
@@ -876,7 +926,7 @@ public void testEnrichExplosionManyMatches() throws IOException {
 
     private Map<String, Object> enrichExplosion(int sensorDataCount, int lookupEntries) throws IOException {
         try {
-            initSensorData(sensorDataCount, 1, 1);
+            initSensorData(sensorDataCount, 1, 1, false);
             initSensorEnrich(lookupEntries, 1, i -> "73.9857 40.7484");
             try {
                 StringBuilder query = startQuery();
@@ -1050,7 +1100,7 @@ private void initMvLongsIndex(int docs, int fields, int fieldValues) throws IOEx
         initIndex("mv_longs", bulk.toString());
     }
 
-    private void initSensorData(int docCount, int sensorCount, int joinFieldCount) throws IOException {
+    private void initSensorData(int docCount, int sensorCount, int joinFieldCount, boolean expressionBasedJoin) throws IOException {
         logger.info("loading sensor data");
         // We cannot go over 1000 fields, due to failed on parsing mappings on index creation
         // [sensor_data] java.lang.IllegalArgumentException: Limit of total fields [1000] has been exceeded
@@ -1061,8 +1111,9 @@ private void initSensorData(int docCount, int sensorCount, int joinFieldCount) t
                  "properties": {
                      "@timestamp": { "type": "date" },
             """);
+        String suffix = expressionBasedJoin ? "_left" : "";
         for (int i = 0; i < joinFieldCount; i++) {
-            createIndexBuilder.append("\"id").append(i).append("\": { \"type\": \"long\" },");
+            createIndexBuilder.append("\"id").append(suffix).append(i).append("\": { \"type\": \"long\" },");
         }
         createIndexBuilder.append("""
                     "value": { "type": "double" }
@@ -1083,7 +1134,7 @@ private void initSensorData(int docCount, int sensorCount, int joinFieldCount) t
                 {"create":{}}
                 {"timestamp":"%s",""", DateFieldMapper.DEFAULT_DATE_TIME_FORMATTER.formatMillis(i * 10L + firstDate)));
             for (int j = 0; j < joinFieldCount; j++) {
-                data.append(String.format(Locale.ROOT, "\"id%d\":%d, ", j, i % sensorCount));
+                data.append(String.format(Locale.ROOT, "\"id%s%d\":%d, ", suffix, j, i % sensorCount));
             }
             data.append(String.format(Locale.ROOT, "\"value\": %f}\n", i * 1.1));
             if (i % docsPerBulk == docsPerBulk - 1) {
@@ -1094,8 +1145,13 @@ private void initSensorData(int docCount, int sensorCount, int joinFieldCount) t
         initIndex("sensor_data", data.toString());
     }
 
-    private void initSensorLookup(int lookupEntries, int sensorCount, IntFunction<String> location, int joinFieldsCount)
-        throws IOException {
+    private void initSensorLookup(
+        int lookupEntries,
+        int sensorCount,
+        IntFunction<String> location,
+        int joinFieldsCount,
+        boolean expressionBasedJoin
+    ) throws IOException {
         logger.info("loading sensor lookup");
         // cannot go over 1000 fields, due to failed on parsing mappings on index creation
         // [sensor_data] java.lang.IllegalArgumentException: Limit of total fields [1000] has been exceeded
@@ -1105,8 +1161,9 @@ private void initSensorLookup(int lookupEntries, int sensorCount, IntFunction<St
             {
                 "properties": {
             """);
+        String suffix = expressionBasedJoin ? "_right" : "";
         for (int i = 0; i < joinFieldsCount; i++) {
-            createIndexBuilder.append("\"id").append(i).append("\": { \"type\": \"long\" },");
+            createIndexBuilder.append("\"id").append(suffix).append(i).append("\": { \"type\": \"long\" },");
         }
         createIndexBuilder.append("""
                     "location": { "type": "geo_point" },
@@ -1127,7 +1184,7 @@ private void initSensorLookup(int lookupEntries, int sensorCount, IntFunction<St
                 {"create":{}}
                 {"""));
             for (int j = 0; j < joinFieldsCount; j++) {
-                data.append(String.format(Locale.ROOT, "\"id%d\":%d, ", j, sensor));
+                data.append(String.format(Locale.ROOT, "\"id%s%d\":%d, ", suffix, j, sensor));
             }
             data.append(String.format(Locale.ROOT, """
                 "location": "POINT(%s)", "filter_key": %d}\n""", location.apply(sensor), i));
@@ -1165,7 +1222,7 @@ private void initSensorLookupString(int lookupEntries, int sensorCount, IntFunct
     }
 
     private void initSensorEnrich(int lookupEntries, int sensorCount, IntFunction<String> location) throws IOException {
-        initSensorLookup(lookupEntries, sensorCount, location, 1);
+        initSensorLookup(lookupEntries, sensorCount, location, 1, false);
         logger.info("loading sensor enrich");
 
         Request create = new Request("PUT", "/_enrich/policy/sensor");

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";
@@ -155,17 +155,11 @@ private Query wrapSingleValueQuery(Query query) {
     }
 
     /**
-     * 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.
+     * Returns a function that reads values from the given block. The function
+     * takes the offset of the value to read and returns the value as an {@link Object}.
      */
-    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 +190,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 +304,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 +367,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 +419,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 +460,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/OperatorTestCase.java

@@ -97,7 +97,7 @@ protected ByteSizeValue enoughMemoryForSimple() {
      * asserting both that this throws a {@link CircuitBreakingException} and releases
      * all pages.
      */
-    public final void testSimpleCircuitBreaking() {
+    public void testSimpleCircuitBreaking() {
         /*
          * Build the input before building `simple` to handle the rare
          * cases where `simple` need some state from the input - mostly