WIP POC for expression join

Author: julian-elastic

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

@@ -10,6 +10,8 @@
 import org.apache.lucene.search.Query;
 import org.elasticsearch.core.Nullable;
 
+import java.io.IOException;
+
 /**
  * An interface to generates queries for the lookup and enrich operators.
  * This interface is used to retrieve queries based on a position index.
@@ -20,7 +22,7 @@ public interface LookupEnrichQueryGenerator {
      * Returns the query at the given position.
      */
     @Nullable
-    Query getQuery(int position);
+    Query getQuery(int position) throws IOException;
 
     /**
      * Returns the number of queries in this generator

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

@@ -89,7 +89,7 @@ public int getPositionCount() {
     public abstract QueryList onlySingleValues(Warnings warnings, String multiValueWarningMessage);
 
     @Override
-    public final Query getQuery(int position) {
+    public final Query getQuery(int position) throws IOException {
         final int valueCount = block.getValueCount(position);
         if (onlySingleValueParams != null && valueCount != 1) {
             if (valueCount > 1) {
@@ -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) throws IOException;
 
     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);

x-pack/plugin/esql/qa/testFixtures/src/main/resources/lookup-join.csv-spec

@@ -2006,6 +2006,100 @@ language_code_float:integer | language_name:keyword
 127                         | max_byte
 ;
 
+lookupWithExpressionEquals
+required_capability: join_lookup_v12
+required_capability: lookup_join_on_expression
+
+FROM employees
+| WHERE emp_no == 10001
+| LOOKUP JOIN languages_lookup ON languages == language_code
+| KEEP emp_no, languages, language_code, language_name
+;
+
+emp_no:integer | languages:integer | language_code:integer | language_name:keyword
+10001          | 2                 | 2                     | French
+;
+
+lookupWithExpressionNotEquals
+required_capability: join_lookup_v12
+required_capability: lookup_join_on_expression
+
+FROM employees
+| WHERE emp_no == 10001
+| LOOKUP JOIN languages_lookup ON languages != language_code
+| KEEP emp_no, languages, language_code, language_name
+| SORT language_code
+;
+
+emp_no:integer | languages:integer | language_code:integer | language_name:keyword
+10001          | 2                 | 1                     | English
+10001          | 2                 | 3                     | Spanish
+10001          | 2                 | 4                     | German
+;
+
+lookupWithExpressionGreater
+required_capability: join_lookup_v12
+required_capability: lookup_join_on_expression
+
+FROM employees
+| WHERE emp_no == 10001
+| LOOKUP JOIN languages_lookup ON languages > language_code
+| KEEP emp_no, languages, language_code, language_name
+;
+
+emp_no:integer | languages:integer | language_code:integer | language_name:keyword
+10001          | 2                 | 1                     | English
+;
+
+lookupWithExpressionGreaterOrEquals
+required_capability: join_lookup_v12
+required_capability: lookup_join_on_expression
+
+FROM employees
+| WHERE emp_no == 10001
+| LOOKUP JOIN languages_lookup ON languages >= language_code
+| KEEP emp_no, languages, language_code, language_name
+| SORT language_code
+;
+
+emp_no:integer | languages:integer | language_code:integer | language_name:keyword
+10001          | 2                 | 1                     | English
+10001          | 2                 | 2                     | French
+;
+
+lookupWithExpressionLess
+required_capability: join_lookup_v12
+required_capability: lookup_join_on_expression
+
+FROM employees
+| WHERE emp_no == 10001
+| LOOKUP JOIN languages_lookup ON languages < language_code
+| KEEP emp_no, languages, language_code, language_name
+| SORT language_code
+;
+
+emp_no:integer | languages:integer | language_code:integer | language_name:keyword
+10001          | 2                 | 3                     | Spanish
+10001          | 2                 | 4                     | German
+;
+
+lookupWithExpressionLessOrEquals
+required_capability: join_lookup_v12
+required_capability: lookup_join_on_expression
+
+FROM employees
+| WHERE emp_no == 10001
+| LOOKUP JOIN languages_lookup ON languages <= language_code
+| KEEP emp_no, languages, language_code, language_name
+| SORT language_code
+;
+
+emp_no:integer | languages:integer | language_code:integer | language_name:keyword
+10001          | 2                 | 2                     | French
+10001          | 2                 | 3                     | Spanish
+10001          | 2                 | 4                     | German
+;
+
 byteJoinDouble
 required_capability: join_lookup_v12
 required_capability: lookup_join_on_mixed_numeric_fields

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/Analyzer.java

@@ -85,6 +85,7 @@
 import org.elasticsearch.xpack.esql.expression.function.vector.VectorFunction;
 import org.elasticsearch.xpack.esql.expression.predicate.operator.arithmetic.DateTimeArithmeticOperation;
 import org.elasticsearch.xpack.esql.expression.predicate.operator.arithmetic.EsqlArithmeticOperation;
+import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.EsqlBinaryComparison;
 import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.In;
 import org.elasticsearch.xpack.esql.index.EsIndex;
 import org.elasticsearch.xpack.esql.index.IndexResolution;
@@ -704,10 +705,25 @@ private LogicalPlan resolveLookup(Lookup l, List<Attribute> childrenOutput) {
             return l;
         }
 
+        private List<Expression> resolveJoinFilters(List<Expression> filters, List<Attribute> leftOutput, List<Attribute> rightOutput) {
+            if (filters.isEmpty()) {
+                return emptyList();
+            }
+            List<Attribute> childrenOutput = new ArrayList<>(leftOutput);
+            childrenOutput.addAll(rightOutput);
+
+            List<Expression> resolvedFilters = new ArrayList<>(filters.size());
+            for (Expression filter : filters) {
+                resolvedFilters.add(filter.transformUp(UnresolvedAttribute.class, ua -> maybeResolveAttribute(ua, childrenOutput)));
+            }
+            return resolvedFilters;
+        }
+
         private Join resolveLookupJoin(LookupJoin join) {
             JoinConfig config = join.config();
             // for now, support only (LEFT) USING clauses
             JoinType type = config.type();
+
             // rewrite the join into an equi-join between the field with the same name between left and right
             if (type instanceof UsingJoinType using) {
                 List<Attribute> cols = using.columns();
@@ -727,12 +743,33 @@ private Join resolveLookupJoin(LookupJoin join) {
                     );
                     return join.withConfig(new JoinConfig(type, singletonList(errorAttribute), emptyList(), emptyList()));
                 }
-                // resolve the using columns against the left and the right side then assemble the new join config
-                List<Attribute> leftKeys = resolveUsingColumns(cols, join.left().output(), "left");
-                List<Attribute> rightKeys = resolveUsingColumns(cols, join.right().output(), "right");
+                List<Attribute> leftKeys = new ArrayList<>();
+                List<Attribute> rightKeys = new ArrayList<>();
+                List<Expression> resolvedFilters = new ArrayList<>();
+                List<Attribute> matchKeys;
+                if (join.candidateRightHandFilters().isEmpty() == false) {
+                    resolvedFilters = resolveJoinFilters(join.candidateRightHandFilters(), join.left().output(), join.right().output());
+                    // build leftKeys and rightKeys using the left side of the resolvedFilters.
+                    for (Expression expression : resolvedFilters) {
+                        if (expression instanceof EsqlBinaryComparison binaryComparison) {
+                            leftKeys.add((Attribute) binaryComparison.left());
+                            rightKeys.add((Attribute) binaryComparison.right());
+                        } else {
+                            throw new EsqlIllegalArgumentException("Unsupported join filter expression: " + expression);
+                        }
+                    }
+                    Set<Attribute> matchKeysSet = new HashSet<>(leftKeys);
+                    matchKeysSet.addAll(rightKeys);
+                    matchKeys = new ArrayList<>(matchKeysSet);
+                } else {
+                    // resolve the using columns against the left and the right side then assemble the new join config
+                    leftKeys = resolveUsingColumns(cols, join.left().output(), "left");
+                    rightKeys = resolveUsingColumns(cols, join.right().output(), "right");
+                    matchKeys = leftKeys;
+                }
 
-                config = new JoinConfig(coreJoin, leftKeys, leftKeys, rightKeys);
-                join = new LookupJoin(join.source(), join.left(), join.right(), config, join.isRemote());
+                config = new JoinConfig(coreJoin, matchKeys, leftKeys, rightKeys);
+                return new LookupJoin(join.source(), join.left(), join.right(), config, join.isRemote(), resolvedFilters);
             } else if (type != JoinTypes.LEFT) {
                 // everything else is unsupported for now
                 // LEFT can only happen by being mapped from a USING above. So we need to exclude this as well because this rule can be run
@@ -741,6 +778,7 @@ private Join resolveLookupJoin(LookupJoin join) {
                 // add error message
                 return join.withConfig(new JoinConfig(type, singletonList(errorAttribute), emptyList(), emptyList()));
             }
+
             return join;
         }
 

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/enrich/AbstractLookupService.java

@@ -558,6 +558,10 @@ abstract static class TransportRequest extends AbstractTransportRequest implemen
             this.source = source;
         }
 
+        public Page getInputPage() {
+            return inputPage;
+        }
+
         @Override
         public final String[] indices() {
             return new String[] { indexPattern };

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/enrich/BinaryComparisonQueryList.java

@@ -0,0 +1,71 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.esql.enrich;
+
+import org.apache.lucene.search.Query;
+import org.elasticsearch.cluster.service.ClusterService;
+import org.elasticsearch.compute.data.Block;
+import org.elasticsearch.compute.operator.Warnings;
+import org.elasticsearch.compute.operator.lookup.QueryList;
+import org.elasticsearch.index.mapper.MappedFieldType;
+import org.elasticsearch.index.query.SearchExecutionContext;
+import org.elasticsearch.xpack.esql.core.expression.Literal;
+import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.EsqlBinaryComparison;
+import org.elasticsearch.xpack.esql.optimizer.rules.physical.local.LucenePushdownPredicates;
+import org.elasticsearch.xpack.esql.planner.TranslatorHandler;
+import org.elasticsearch.xpack.esql.plugin.EsqlFlags;
+import org.elasticsearch.xpack.esql.stats.SearchContextStats;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.function.IntFunction;
+
+public class BinaryComparisonQueryList extends QueryList {
+    private final EsqlBinaryComparison binaryComparison;
+    private final IntFunction<Object> blockValueReader;
+    private final ClusterService clusterService;
+    private final SearchExecutionContext searchExecutionContext;
+
+    public BinaryComparisonQueryList(
+        MappedFieldType field,
+        SearchExecutionContext searchExecutionContext,
+        Block block,
+        EsqlBinaryComparison binaryComparison,
+        ClusterService clusterService
+    ) {
+        super(field, searchExecutionContext, null, block, null);
+        this.binaryComparison = binaryComparison;
+        this.blockValueReader = QueryList.createBlockValueReader(block);
+        this.clusterService = clusterService;
+        this.searchExecutionContext = searchExecutionContext;
+    }
+
+    @Override
+    public QueryList onlySingleValues(Warnings warnings, String multiValueWarningMessage) {
+        throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public Query doGetQuery(int position, int firstValueIndex, int valueCount) throws IOException {
+        if (valueCount == 0) {
+            return null;
+        }
+        Object value = blockValueReader.apply(firstValueIndex);
+        // create a new comparison with the value from the block as a literal
+        EsqlBinaryComparison swapped = (EsqlBinaryComparison) binaryComparison.swapLeftAndRight();
+        EsqlBinaryComparison comparison = swapped.getFunctionType()
+            .buildNewInstance(swapped.source(), swapped.left(), new Literal(swapped.right().source(), value, swapped.right().dataType()));
+        LucenePushdownPredicates lucenePushdownPredicates = LucenePushdownPredicates.from(
+            SearchContextStats.from(List.of(searchExecutionContext)),
+            new EsqlFlags(clusterService.getClusterSettings())
+        );
+        return comparison.asQuery(lucenePushdownPredicates, TranslatorHandler.TRANSLATOR_HANDLER)
+            .toQueryBuilder()
+            .toQuery(searchExecutionContext);
+    }
+}