Fix some tests.

Author: afoucret

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

@@ -554,7 +554,7 @@ public static LogicalOptimizerContext unboundLogicalOptimizerContext() {
         mock(ProjectResolver.class),
         mock(IndexNameExpressionResolver.class),
         null,
-        new InferenceService(mock(Client.class)),
+        new InferenceService(createMockClient()),
         new BlockFactoryProvider(PlannerUtils.NON_BREAKING_BLOCK_FACTORY),
         TEST_PLANNER_SETTINGS,
         new CrossProjectModeDecider(Settings.EMPTY)
@@ -563,6 +563,7 @@ public static LogicalOptimizerContext unboundLogicalOptimizerContext() {
     private static ClusterService createMockClusterService() {
         var service = mock(ClusterService.class);
         doReturn(new ClusterName("test-cluster")).when(service).getClusterName();
+        doReturn(Settings.EMPTY).when(service).getSettings();
         return service;
     }
 
@@ -578,6 +579,13 @@ private static ThreadPool createMockThreadPool() {
         return threadPool;
     }
 
+    private static Client createMockClient() {
+        var client = mock(Client.class);
+        doReturn(Settings.EMPTY).when(client).settings();
+        doReturn(createMockThreadPool()).when(client).threadPool();
+        return client;
+    }
+
     private EsqlTestUtils() {}
 
     public static Configuration configuration(QueryPragmas pragmas, String query, EsqlStatement statement) {
@@ -732,6 +740,8 @@ private static <T> List<T> toList(Iterator<T> iterator) {
     public static List<String> withDefaultLimitWarning(List<String> warnings) {
         List<String> result = warnings == null ? new ArrayList<>() : new ArrayList<>(warnings);
         result.add("No limit defined, adding default limit of [1000]");
+
+        result.add("No limit defined, adding default limit of [100]");
         return result;
     }
 

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

@@ -8943,13 +8943,54 @@ public void testKnnWithStats() {
     }
 
     public void testKnnWithRerankAmdTopN() {
-        assertThat(typesError("""
+        var query = """
             from types metadata _score
             | where knn(dense_vector, [0, 1, 2])
+            | LIMIT 100
             | rerank "some text" on text with { "inference_id" : "reranking-inference-id" }
             | sort _score desc
             | limit 10
-            """), containsString("Knn function must be used with a LIMIT clause"));
+            """;
+
+        var optimized = planTypes(query);
+
+        var topN = as(optimized, TopN.class);
+        assertThat(topN.limit().fold(FoldContext.small()), equalTo(10));
+
+        var rerank = as(topN.child(), Rerank.class);
+        var rerankLimit = as(rerank.child(), Limit.class);
+        assertThat(rerankLimit.limit().fold(FoldContext.small()), equalTo(100));
+        var filter = as(rerankLimit.child(), Filter.class);
+
+        // KNN is using the first limit (100)
+        var knn = as(filter.condition(), Knn.class);
+        assertThat(knn.implicitK(), equalTo(100));
+    }
+
+    public void testKnnWithRerankImplicitLimitAmdTopN() {
+        var query = """
+            from types metadata _score
+            | where knn(dense_vector, [0, 1, 2])
+            | rerank "some text" on text with { "inference_id" : "reranking-inference-id" }
+            | sort _score desc
+            | limit 10
+            """;
+
+        var optimized = planTypes(query);
+
+        var topN = as(optimized, TopN.class);
+        assertThat(topN.limit().fold(FoldContext.small()), equalTo(10));
+
+        var rerank = as(topN.child(), Rerank.class);
+
+        // RERANK implicit limit is set by setting (1000)
+        var rerankLimit = as(rerank.child(), Limit.class);
+        assertThat(rerankLimit.limit().fold(FoldContext.small()), equalTo(1000));
+        var filter = as(rerankLimit.child(), Filter.class);
+
+        // KNN is using the implicit limit of RERANK (1000)
+        var knn = as(filter.condition(), Knn.class);
+        assertThat(knn.implicitK(), equalTo(1000));
     }
 
     public void testKnnWithRerankAmdLimit() {
@@ -8964,7 +9005,8 @@ public void testKnnWithRerankAmdLimit() {
 
         var rerank = as(optimized, Rerank.class);
         var limit = as(rerank.child(), Limit.class);
-        assertThat(limit.limit().fold(FoldContext.small()), equalTo(100));
+        assertThat(limit.limit().fold(FoldContext.small()), equalTo(1_000));
+
         var filter = as(limit.child(), Filter.class);
         var knn = as(filter.condition(), Knn.class);
         assertThat(knn.implicitK(), equalTo(100));

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/EnforceRowLimitsTests.java

@@ -60,13 +60,13 @@ public String toString() {
 
     public static class CompletionTestCase extends TestCase {
         public String name() {
-            return "Rerank";
+            return "Completion";
         }
 
         public Completion createPlan(LogicalPlan child, Expression rowLimit) {
             Source source = EMPTY;
             Expression prompt = Literal.keyword(source, "test prompt");
-            Attribute targetField = new ReferenceAttribute(source, null, "completion", DataType.KEYWORD);
+            Attribute targetField = new ReferenceAttribute(source, "completion", DataType.KEYWORD);
             return new Completion(source, child, Literal.keyword(source, "test-inference-id"), prompt, targetField, rowLimit);
         }
     }
@@ -80,7 +80,7 @@ public String name() {
         public Rerank createPlan(LogicalPlan child, Expression rowLimit) {
             Source source = EMPTY;
             Expression queryText = Literal.keyword(source, "test query");
-            Attribute scoreAttribute = new ReferenceAttribute(source, null, "score", DataType.DOUBLE);
+            Attribute scoreAttribute = new ReferenceAttribute(source, "score", DataType.DOUBLE);
             List<Alias> rerankFields = List.of();
             return new Rerank(
                 source,