ESQL: Fix LOOKUP JOIN with limit (elastic#120411)

For queries like 
... | LOOKUP JOIN lookup_index ON key | LIMIT 10
the limit cannot be simply pushed past the join - but it can be duplicated past the join.

In such cases, leave an explicit Limit plan node downstream from the Join (in addition to pushing down the limit), but mark it in a way that prevents being duplicated multiple times (which would cause infinite loops).

Align the logic for MV_EXPAND, which used to, instead, internalize a limit into the MvExpand node.

Author: alex-spies

x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/EsqlTestUtils.java

@@ -66,6 +66,7 @@
 import org.elasticsearch.xpack.esql.parser.QueryParam;
 import org.elasticsearch.xpack.esql.plan.logical.Enrich;
 import org.elasticsearch.xpack.esql.plan.logical.EsRelation;
+import org.elasticsearch.xpack.esql.plan.logical.Limit;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
 import org.elasticsearch.xpack.esql.plan.logical.local.LocalRelation;
 import org.elasticsearch.xpack.esql.plan.logical.local.LocalSupplier;
@@ -111,6 +112,7 @@
 import static java.util.Collections.emptyList;
 import static java.util.Collections.emptyMap;
 import static java.util.Collections.unmodifiableMap;
+import static org.elasticsearch.test.ESTestCase.assertEquals;
 import static org.elasticsearch.test.ESTestCase.between;
 import static org.elasticsearch.test.ESTestCase.randomAlphaOfLength;
 import static org.elasticsearch.test.ESTestCase.randomBoolean;
@@ -403,6 +405,21 @@ public static <T> T as(Object node, Class<T> type) {
         return type.cast(node);
     }
 
+    public static Limit asLimit(Object node, Integer limitLiteral) {
+        return asLimit(node, limitLiteral, null);
+    }
+
+    public static Limit asLimit(Object node, Integer limitLiteral, Boolean duplicated) {
+        Limit limit = as(node, Limit.class);
+        if (limitLiteral != null) {
+            assertEquals(as(limit.limit(), Literal.class).value(), limitLiteral);
+        }
+        if (duplicated != null) {
+            assertEquals(limit.duplicated(), duplicated);
+        }
+        return limit;
+    }
+
     public static Map<String, EsField> loadMapping(String name) {
         return LoadMapping.loadMapping(name);
     }

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

@@ -273,6 +273,58 @@ emp_no:integer
 10001
 ;
 
+
+lookupIndexInFromRepeatedRowBug
+// Test for https://github.com/elastic/elasticsearch/issues/118852
+required_capability: join_lookup_v12
+FROM languages_lookup_non_unique_key
+| WHERE language_code == 1
+| LOOKUP JOIN languages_lookup ON language_code
+| KEEP language_code, language_name, country
+| SORT language_code, language_name, country
+;
+
+language_code:integer | language_name:keyword       | country:text
+1                     | English                     | Canada
+1                     | English                     | United Kingdom
+1                     | English                     | United States of America
+1                     | English                     | null
+;
+
+nonUniqueRightKeyOnTheCoordinatorLateLimit
+required_capability: join_lookup_v12
+required_capability: join_lookup_fix_limit_pushdown
+
+FROM employees
+| SORT emp_no
+| EVAL language_code = emp_no % 10
+| LOOKUP JOIN languages_lookup_non_unique_key ON language_code
+| KEEP emp_no, language_code, language_name, country
+| LIMIT 4
+| SORT country
+;
+
+emp_no:integer | language_code:integer | language_name:keyword | country:text
+10001          | 1                     | English               | Canada
+10001          | 1                     | null                  | United Kingdom
+10001          | 1                     | English               | United States of America
+10001          | 1                     | English               | null
+;
+
+nonUniqueRightKeyLateLimitWithEmptyRelation
+required_capability: join_lookup_v12
+required_capability: join_lookup_fix_limit_pushdown
+
+ROW language_code = 1
+| WHERE language_code != 1
+| LOOKUP JOIN languages_lookup_non_unique_key ON language_code
+| LIMIT 1
+| KEEP language_code, language_name
+;
+
+language_code:integer | language_name:keyword
+;
+
 ###########################################################################
 # null and multi-value behavior with languages_lookup_non_unique_key index
 ###########################################################################
@@ -1278,23 +1330,6 @@ ignoreOrder:true
 2023-10-23T12:15:03.360Z | 172.21.2.162      | 3450233             | QA              | null
 ;
 
-lookupIndexInFromRepeatedRowBug
-// Test for https://github.com/elastic/elasticsearch/issues/118852
-required_capability: join_lookup_v12
-FROM languages_lookup_non_unique_key
-| WHERE language_code == 1
-| LOOKUP JOIN languages_lookup ON language_code
-| KEEP language_code, language_name, country
-| SORT language_code, language_name, country
-;
-
-language_code:integer | language_name:keyword       | country:text
-1                     | English                     | Canada
-1                     | English                     | United Kingdom
-1                     | English                     | United States of America
-1                     | English                     | null
-;
-
 lookupIndexQuoting
 required_capability: join_lookup_v12
 FROM languages_lookup_non_unique_key

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java

@@ -594,6 +594,11 @@ public enum Cap {
          */
         JOIN_LOOKUP_SKIP_MV_ON_LOOKUP_KEY(JOIN_LOOKUP_V12.isEnabled()),
 
+        /**
+         * Fix pushing down LIMIT past LOOKUP JOIN in case of multiple matching join keys.
+         */
+        JOIN_LOOKUP_FIX_LIMIT_PUSHDOWN(JOIN_LOOKUP_V12.isEnabled()),
+
         /**
          * Fix for https://github.com/elastic/elasticsearch/issues/117054
          */

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

@@ -549,8 +549,7 @@ private LogicalPlan resolveMvExpand(MvExpand p, List<Attribute> childrenOutput)
                     resolved,
                     resolved.resolved()
                         ? new ReferenceAttribute(resolved.source(), resolved.name(), resolved.dataType(), resolved.nullable(), null, false)
-                        : resolved,
-                    p.limit()
+                        : resolved
                 );
             }
             return p;

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/AddDefaultTopN.java

@@ -34,7 +34,7 @@
  * | sort first_name
  * | limit 15
  * <p>
- * PushDownAndCombineLimits rule will copy the "limit 15" after "sort emp_no" if there is no filter on the expanded values
+ * {@link PushDownAndCombineLimits} will copy the "limit 15" after "sort emp_no" if there is no filter on the expanded values
  * OR if there is no sort between "limit" and "mv_expand".
  * But, since this type of query has such a filter, the "sort emp_no" will have no limit when it reaches the current rule.
  */

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/PushDownAndCombineLimits.java

@@ -7,7 +7,6 @@
 
 package org.elasticsearch.xpack.esql.optimizer.rules.logical;
 
-import org.elasticsearch.xpack.esql.core.expression.Literal;
 import org.elasticsearch.xpack.esql.optimizer.LogicalOptimizerContext;
 import org.elasticsearch.xpack.esql.plan.logical.Aggregate;
 import org.elasticsearch.xpack.esql.plan.logical.Enrich;
@@ -21,6 +20,9 @@
 import org.elasticsearch.xpack.esql.plan.logical.join.Join;
 import org.elasticsearch.xpack.esql.plan.logical.join.JoinTypes;
 
+import java.util.ArrayList;
+import java.util.List;
+
 public final class PushDownAndCombineLimits extends OptimizerRules.ParameterizedOptimizerRule<Limit, LogicalOptimizerContext> {
 
     public PushDownAndCombineLimits() {
@@ -31,27 +33,18 @@ public PushDownAndCombineLimits() {
     public LogicalPlan rule(Limit limit, LogicalOptimizerContext ctx) {
         if (limit.child() instanceof Limit childLimit) {
             var limitSource = limit.limit();
-            var l1 = (int) limitSource.fold(ctx.foldCtx());
-            var l2 = (int) childLimit.limit().fold(ctx.foldCtx());
-            return new Limit(limit.source(), Literal.of(limitSource, Math.min(l1, l2)), childLimit.child());
+            var parentLimitValue = (int) limitSource.fold(ctx.foldCtx());
+            var childLimitValue = (int) childLimit.limit().fold(ctx.foldCtx());
+            // We want to preserve the duplicated() value of the smaller limit, so we'll use replaceChild.
+            return parentLimitValue < childLimitValue ? limit.replaceChild(childLimit.child()) : childLimit;
         } else if (limit.child() instanceof UnaryPlan unary) {
             if (unary instanceof Eval || unary instanceof Project || unary instanceof RegexExtract || unary instanceof Enrich) {
                 return unary.replaceChild(limit.replaceChild(unary.child()));
-            } else if (unary instanceof MvExpand mvx) {
+            } else if (unary instanceof MvExpand) {
                 // MV_EXPAND can increase the number of rows, so we cannot just push the limit down
                 // (we also have to preserve the LIMIT afterwards)
-                //
-                // To avoid infinite loops, ie.
-                // | MV_EXPAND | LIMIT -> | LIMIT | MV_EXPAND | LIMIT -> ... | MV_EXPAND | LIMIT
-                // we add an inner limit to MvExpand and just push down the existing limit, ie.
-                // | MV_EXPAND | LIMIT N -> | LIMIT N | MV_EXPAND (with limit N)
-                var limitSource = limit.limit();
-                var limitVal = (int) limitSource.fold(ctx.foldCtx());
-                Integer mvxLimit = mvx.limit();
-                if (mvxLimit == null || mvxLimit > limitVal) {
-                    mvx = new MvExpand(mvx.source(), mvx.child(), mvx.target(), mvx.expanded(), limitVal);
-                }
-                return mvx.replaceChild(limit.replaceChild(mvx.child()));
+                // To avoid repeating this infinitely, we have to set duplicated = true.
+                return duplicateLimitAsFirstGrandchild(limit);
             }
             // check if there's a 'visible' descendant limit lower than the current one
             // and if so, align the current limit since it adds no value
@@ -62,17 +55,15 @@ public LogicalPlan rule(Limit limit, LogicalOptimizerContext ctx) {
                     var l1 = (int) limit.limit().fold(ctx.foldCtx());
                     var l2 = (int) descendantLimit.limit().fold(ctx.foldCtx());
                     if (l2 <= l1) {
-                        return new Limit(limit.source(), Literal.of(limit.limit(), l2), limit.child());
+                        return limit.withLimit(descendantLimit.limit());
                     }
                 }
             }
-        } else if (limit.child() instanceof Join join) {
-            if (join.config().type() == JoinTypes.LEFT) {
-                // NOTE! This is only correct because our LEFT JOINs preserve the number of rows from the left hand side.
-                // This deviates from SQL semantics. In SQL, multiple matches on the right hand side lead to multiple rows in the output.
-                // For us, multiple matches on the right hand side are collected into multi-values.
-                return join.replaceChildren(limit.replaceChild(join.left()), join.right());
-            }
+        } else if (limit.child() instanceof Join join && join.config().type() == JoinTypes.LEFT) {
+            // Left joins increase the number of rows if any join key has multiple matches from the right hand side.
+            // Therefore, we cannot simply push down the limit - but we can add another limit before the join.
+            // To avoid repeating this infinitely, we have to set duplicated = true.
+            return duplicateLimitAsFirstGrandchild(limit);
         }
         return limit;
     }
@@ -100,4 +91,27 @@ private static Limit descendantLimit(UnaryPlan unary) {
         }
         return null;
     }
+
+    /**
+     * Duplicate the limit past its child if it wasn't duplicated yet. The duplicate is placed on top of its leftmost grandchild.
+     * Idempotent. (Sets {@link Limit#duplicated()} to {@code true} on the limit that remains at the top.)
+     */
+    private static Limit duplicateLimitAsFirstGrandchild(Limit limit) {
+        if (limit.duplicated()) {
+            return limit;
+        }
+
+        List<LogicalPlan> grandChildren = limit.child().children();
+        LogicalPlan firstGrandChild = grandChildren.getFirst();
+        LogicalPlan newFirstGrandChild = limit.replaceChild(firstGrandChild);
+
+        List<LogicalPlan> newGrandChildren = new ArrayList<>();
+        newGrandChildren.add(newFirstGrandChild);
+        for (int i = 1; i < grandChildren.size(); i++) {
+            newGrandChildren.add(grandChildren.get(i));
+        }
+
+        LogicalPlan newChild = limit.child().replaceChildren(newGrandChildren);
+        return limit.replaceChild(newChild).withDuplicated(true);
+    }
 }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/Limit.java

@@ -21,21 +21,52 @@ public class Limit extends UnaryPlan {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(LogicalPlan.class, "Limit", Limit::new);
 
     private final Expression limit;
-
+    /**
+     * Important for optimizations. This should be {@code false} in most cases, which allows this instance to be duplicated past a child
+     * plan node that increases the number of rows, like for LOOKUP JOIN and MV_EXPAND.
+     * Needs to be set to {@code true} in {@link org.elasticsearch.xpack.esql.optimizer.rules.logical.PushDownAndCombineLimits} to avoid
+     * infinite loops from adding a duplicate of the limit past the child over and over again.
+     */
+    private final transient boolean duplicated;
+
+    /**
+     * Default way to create a new instance. Do not use this to copy an existing instance, as this sets {@link Limit#duplicated} to
+     * {@code false}.
+     */
     public Limit(Source source, Expression limit, LogicalPlan child) {
+        this(source, limit, child, false);
+    }
+
+    public Limit(Source source, Expression limit, LogicalPlan child, boolean duplicated) {
         super(source, child);
         this.limit = limit;
+        this.duplicated = duplicated;
     }
 
+    /**
+     * Omits reading {@link Limit#duplicated}, c.f. {@link Limit#writeTo}.
+     */
     private Limit(StreamInput in) throws IOException {
-        this(Source.readFrom((PlanStreamInput) in), in.readNamedWriteable(Expression.class), in.readNamedWriteable(LogicalPlan.class));
+        this(
+            Source.readFrom((PlanStreamInput) in),
+            in.readNamedWriteable(Expression.class),
+            in.readNamedWriteable(LogicalPlan.class),
+            false
+        );
     }
 
+    /**
+     * Omits serializing {@link Limit#duplicated} because when sent to a data node, this should always be {@code false}.
+     * That's because if it's true, this means a copy of this limit was pushed down below an MvExpand or Join, and thus there's
+     * another pipeline breaker further upstream - we're already on the coordinator node.
+     */
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         Source.EMPTY.writeTo(out);
         out.writeNamedWriteable(limit());
         out.writeNamedWriteable(child());
+        // Let's make sure we notice during tests if we ever serialize a duplicated Limit.
+        assert duplicated == false;
     }
 
     @Override
@@ -45,18 +76,30 @@ public String getWriteableName() {
 
     @Override
     protected NodeInfo<Limit> info() {
-        return NodeInfo.create(this, Limit::new, limit, child());
+        return NodeInfo.create(this, Limit::new, limit, child(), duplicated);
     }
 
     @Override
     public Limit replaceChild(LogicalPlan newChild) {
-        return new Limit(source(), limit, newChild);
+        return new Limit(source(), limit, newChild, duplicated);
     }
 
     public Expression limit() {
         return limit;
     }
 
+    public Limit withLimit(Expression limit) {
+        return new Limit(source(), limit, child(), duplicated);
+    }
+
+    public boolean duplicated() {
+        return duplicated;
+    }
+
+    public Limit withDuplicated(boolean duplicated) {
+        return new Limit(source(), limit, child(), duplicated);
+    }
+
     @Override
     public String commandName() {
         return "LIMIT";
@@ -69,7 +112,7 @@ public boolean expressionsResolved() {
 
     @Override
     public int hashCode() {
-        return Objects.hash(limit, child());
+        return Objects.hash(limit, child(), duplicated);
     }
 
     @Override
@@ -83,6 +126,6 @@ public boolean equals(Object obj) {
 
         Limit other = (Limit) obj;
 
-        return Objects.equals(limit, other.limit) && Objects.equals(child(), other.child());
+        return Objects.equals(limit, other.limit) && Objects.equals(child(), other.child()) && (duplicated == other.duplicated);
     }
 }