[ES|QL] pushing down eval expression when it requires data access (elastic#136610)

Author: pmpailis

docs/changelog/136610.yaml

@@ -0,0 +1,6 @@
+pr: 136610
+summary: Pushing down eval expression when it requires data access
+area: ES|QL
+type: bug
+issues:
+ - 133462

x-pack/plugin/esql/qa/testFixtures/src/main/resources/score-function.csv-spec

@@ -125,3 +125,78 @@ book_no:keyword | title:text
 2714            | Return of the King Being the Third Part of The Lord of the Rings | 1.9245924949645996
 7350            | Return of the Shadow                                             | 0.0
 ;
+
+evalUsingScoreAndLimit
+required_capability: score_function
+required_capability: match_function
+required_capability: pushing_down_eval_with_score
+
+FROM books
+| EVAL title_match = SCORE(MATCH(title, "aardvark"))
+| LIMIT 2
+| KEEP title_match
+;
+
+title_match:double
+0
+0
+;
+
+evalUsingScoreAndLimit2
+required_capability: score_function
+required_capability: match_function
+required_capability: pushing_down_eval_with_score
+
+FROM books
+| WHERE author == "J R R Tolkien"
+| EVAL title_match_ring = SCORE(MATCH(title, "ring")), title_match_hobbit = SCORE(MATCH(title, "hobbit"))
+| keep title_match_ring, title_match_hobbit
+;
+
+title_match_ring:double   | title_match_hobbit:double
+0                         | 3.45605206489563
+;
+
+evalUsingScoreAndTopN
+required_capability: score_function
+required_capability: match_function
+required_capability: pushing_down_eval_with_score
+
+FROM books
+| SORT book_no desc
+| EVAL title_match_towers = SCORE(MATCH(title, "towers"))
+| limit 10
+| keep book_no, title_match_towers
+;
+
+book_no:keyword | title_match_towers:double
+9896            | 0
+9801            | 0
+9607            | 0
+9032            | 0
+8956            | 0
+8875            | 3.5007452964782715
+8873            | 0
+8678            | 0
+8615            | 0
+8605            | 0
+;
+
+evalUsingScoreAndFilter
+required_capability: score_function
+required_capability: match_function
+required_capability: pushing_down_eval_with_score
+
+FROM books
+| sort book_no desc
+| EVAL title_match_world = SCORE(MATCH(title, "world"))
+| WHERE title_match_world > 0 and to_integer(book_no) < 8500
+| LIMIT 10
+| KEEP book_no, title_match_world
+| SORT book_no asc
+| LIMIT 1
+;
+
+book_no:keyword  |  title_match_world:double
+7912             |  2.7460334300994873
+;

x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/plugin/ScoreFunctionIT.java

@@ -21,6 +21,7 @@
 import java.util.Collection;
 import java.util.List;
 
+import static org.elasticsearch.common.collect.Iterators.toList;
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
 import static org.hamcrest.CoreMatchers.containsString;
 
@@ -33,6 +34,58 @@ public void setupIndex() {
         createAndPopulateIndex();
     }
 
+    public void testPushingDownEvalWithScore() {
+        var query = """
+            FROM test
+            | EVAL first_score = score(match(content, "dog")), s = [1]
+            | LIMIT 2
+            """;
+
+        try (var resp = run(query)) {
+            assertColumnTypes(resp.columns(), List.of("text", "integer", "double", "integer"));
+            assertColumnNames(resp.columns(), List.of("content", "id", "first_score", "s"));
+            var results = resp.values();
+            final int expectedNumRows = 2;
+            final int expectedNumColumns = 4;
+            for (int i = 0; i < expectedNumRows; i++) {
+                var row = toList(results.next());
+                assertEquals(expectedNumColumns, row.size());
+                for (var column : row) {
+                    assertNotNull(column);
+                }
+            }
+            if (results.hasNext()) {
+                fail("Expected only 2 rows but got more");
+            }
+        }
+    }
+
+    public void testPushingDownEvalWithNestedScore() {
+        var query = """
+            FROM test
+            | EVAL first_score = TO_INTEGER(0.2 + SCORE(MATCH(content, "dog")))
+            | LIMIT 3
+            """;
+
+        try (var resp = run(query)) {
+            assertColumnTypes(resp.columns(), List.of("text", "integer", "integer"));
+            assertColumnNames(resp.columns(), List.of("content", "id", "first_score"));
+            var results = resp.values();
+            final int expectedRows = 3;
+            final int expectedColumns = 3;
+            for (int i = 0; i < expectedRows; i++) {
+                var row = toList(results.next());
+                assertEquals(expectedColumns, row.size());
+                for (var column : row) {
+                    assertNotNull(column);
+                }
+            }
+            if (results.hasNext()) {
+                fail("Expected only 2 rows but got more");
+            }
+        }
+    }
+
     public void testScoreSingleNoMetadata() {
         var query = """
             FROM test

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

@@ -1508,7 +1508,7 @@ public enum Cap {
 
         /**
          * Support for the literal {@code m} suffix as an alias for {@code minute} in temporal amounts.
-         */
+        */
         TEMPORAL_AMOUNT_M,
 
         /**
@@ -1529,7 +1529,13 @@ public enum Cap {
         /**
          * Fix double release in inline stats when LocalRelation is reused
          */
-        INLINE_STATS_DOUBLE_RELEASE_FIX(INLINESTATS_V11.enabled)
+        INLINE_STATS_DOUBLE_RELEASE_FIX(INLINESTATS_V11.enabled),
+
+        /**
+         * Support for pushing down EVAL with SCORE
+         * https://github.com/elastic/elasticsearch/issues/133462
+         */
+        PUSHING_DOWN_EVAL_WITH_SCORE
 
         ;
 

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizer.java

@@ -192,7 +192,7 @@ protected static Batch<LogicalPlan> operators(boolean local) {
             new PruneFilters(),
             new PruneColumns(),
             new PruneLiteralsInOrderBy(),
-            new PushDownAndCombineLimits(),
+            new PushDownAndCombineLimits(local),
             new PushLimitToKnn(),
             new PushDownAndCombineFilters(),
             new PushDownConjunctionsToKnnPrefilters(),

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

@@ -7,6 +7,9 @@
 
 package org.elasticsearch.xpack.esql.optimizer.rules.logical;
 
+import org.elasticsearch.xpack.esql.core.expression.Alias;
+import org.elasticsearch.xpack.esql.core.util.Holder;
+import org.elasticsearch.xpack.esql.expression.function.fulltext.Score;
 import org.elasticsearch.xpack.esql.optimizer.LogicalOptimizerContext;
 import org.elasticsearch.xpack.esql.plan.logical.Aggregate;
 import org.elasticsearch.xpack.esql.plan.logical.Enrich;
@@ -27,8 +30,11 @@
 
 public final class PushDownAndCombineLimits extends OptimizerRules.ParameterizedOptimizerRule<Limit, LogicalOptimizerContext> {
 
-    public PushDownAndCombineLimits() {
+    private final boolean local;
+
+    public PushDownAndCombineLimits(boolean local) {
         super(OptimizerRules.TransformDirection.DOWN);
+        this.local = local;
     }
 
     @Override
@@ -45,7 +51,12 @@ public LogicalPlan rule(Limit limit, LogicalOptimizerContext ctx) {
                 || unary instanceof RegexExtract
                 || unary instanceof Enrich
                 || unary instanceof InferencePlan<?>) {
-                return unary.replaceChild(limit.replaceChild(unary.child()));
+                if (false == local && unary instanceof Eval && evalAliasNeedsData((Eval) unary)) {
+                    // do not push down the limit through an eval that needs data (e.g. a score function) during initial planning
+                    return limit;
+                } else {
+                    return unary.replaceChild(limit.replaceChild(unary.child()));
+                }
             } 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)
@@ -76,6 +87,21 @@ public LogicalPlan rule(Limit limit, LogicalOptimizerContext ctx) {
         return limit;
     }
 
+    /**
+     * Determines if the provided {@link Alias} expression depends on document data by traversing its expression tree.
+     * Returns {@code true} if any child expression requires access to document-specific values, such as the {@link Score} function.
+     * This is used to prevent pushing down limits past operations that need to evaluate expressions using document data.
+     */
+    private boolean evalAliasNeedsData(Eval eval) {
+        Holder<Boolean> hasScore = new Holder<>(false);
+        eval.forEachExpression(expr -> {
+            if (expr instanceof Score) {
+                hasScore.set(true);
+            }
+        });
+        return hasScore.get();
+    }
+
     /**
      * Checks the existence of another 'visible' Limit, that exists behind an operation that doesn't produce output more data than
      * its input (that is not a relation/source nor aggregation).

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizerTests.java

@@ -58,6 +58,7 @@
 import org.elasticsearch.xpack.esql.expression.function.aggregate.SummationMode;
 import org.elasticsearch.xpack.esql.expression.function.fulltext.Match;
 import org.elasticsearch.xpack.esql.expression.function.fulltext.MultiMatch;
+import org.elasticsearch.xpack.esql.expression.function.fulltext.Score;
 import org.elasticsearch.xpack.esql.expression.function.grouping.Bucket;
 import org.elasticsearch.xpack.esql.expression.function.grouping.Categorize;
 import org.elasticsearch.xpack.esql.expression.function.scalar.convert.ToDouble;
@@ -209,6 +210,33 @@
 public class LogicalPlanOptimizerTests extends AbstractLogicalPlanOptimizerTests {
     private static final LiteralsOnTheRight LITERALS_ON_THE_RIGHT = new LiteralsOnTheRight();
 
+    public void testEvalWithScoreImplicitLimit() {
+        var plan = plan("""
+            FROM test
+            | EVAL s = SCORE(MATCH(last_name, "high"))
+            """);
+        var limit = as(plan, Limit.class);
+        assertThat(limit.child(), instanceOf(Eval.class));
+        assertThat(((Literal) limit.limit()).value(), equalTo(1000));
+        var eval = as(limit.child(), Eval.class);
+        assertThat(eval.fields().size(), equalTo(1));
+        assertThat(eval.fields().get(0).child(), instanceOf(Score.class));
+    }
+
+    public void testEvalWithScoreExplicitLimit() {
+        var plan = plan("""
+            FROM test
+            | EVAL s = SCORE(MATCH(last_name, "high"))
+            | LIMIT 42
+            """);
+        var limit = as(plan, Limit.class);
+        assertThat(limit.child(), instanceOf(Eval.class));
+        assertThat(((Literal) limit.limit()).value(), equalTo(42));
+        var eval = as(limit.child(), Eval.class);
+        assertThat(eval.fields().size(), equalTo(1));
+        assertThat(eval.fields().get(0).child(), instanceOf(Score.class));
+    }
+
     public void testEmptyProjections() {
         var plan = plan("""
             from test
@@ -1032,12 +1060,12 @@ public void testCombineLimits() {
         var anotherLimit = new Limit(EMPTY, L(limitValues[secondLimit]), oneLimit);
         assertEquals(
             new Limit(EMPTY, L(Math.min(limitValues[0], limitValues[1])), emptySource()),
-            new PushDownAndCombineLimits().rule(anotherLimit, logicalOptimizerCtx)
+            new PushDownAndCombineLimits(false).rule(anotherLimit, logicalOptimizerCtx)
         );
     }
 
     public void testPushdownLimitsPastLeftJoin() {
-        var rule = new PushDownAndCombineLimits();
+        var rule = new PushDownAndCombineLimits(false);
 
         var leftChild = emptySource();
         var rightChild = new LocalRelation(Source.EMPTY, List.of(fieldAttribute()), EmptyLocalSupplier.EMPTY);