KNN k is set via optimizer and limit

Author: carlosdelest

docs/reference/query-languages/esql/images/functions/knn.svg

@@ -490,7 +490,7 @@ private static FunctionDefinition[][] snapshotFunctions() {
                 def(FirstOverTime.class, uni(FirstOverTime::new), "first_over_time"),
                 def(Score.class, uni(Score::new), Score.NAME),
                 def(Term.class, bi(Term::new), "term"),
-                def(Knn.class, quad(Knn::new), "knn"),
+                def(Knn.class, tri(Knn::new), "knn"),
                 def(StGeohash.class, StGeohash::new, "st_geohash"),
                 def(StGeohashToLong.class, StGeohashToLong::new, "st_geohash_to_long"),
                 def(StGeohashToString.class, StGeohashToString::new, "st_geohash_to_string"),

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/EsqlFunctionRegistry.java

@@ -7,8 +7,6 @@
 
 package org.elasticsearch.xpack.esql.expression.function.vector;
 
-import org.apache.logging.log4j.LogManager;
-import org.apache.logging.log4j.Logger;
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
@@ -56,14 +54,12 @@
 import static java.util.Map.entry;
 import static org.elasticsearch.common.logging.LoggerMessageFormat.format;
 import static org.elasticsearch.index.query.AbstractQueryBuilder.BOOST_FIELD;
-import static org.elasticsearch.search.vectors.KnnVectorQueryBuilder.K_FIELD;
 import static org.elasticsearch.search.vectors.KnnVectorQueryBuilder.NUM_CANDS_FIELD;
 import static org.elasticsearch.search.vectors.KnnVectorQueryBuilder.VECTOR_SIMILARITY_FIELD;
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.ParamOrdinal.FIRST;
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.ParamOrdinal.FOURTH;
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.ParamOrdinal.SECOND;
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.ParamOrdinal.THIRD;
-import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.isFoldable;
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.isNotNull;
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.isType;
 import static org.elasticsearch.xpack.esql.core.type.DataType.DENSE_VECTOR;
@@ -73,13 +69,12 @@
 import static org.elasticsearch.xpack.esql.expression.function.FunctionUtils.resolveTypeQuery;
 
 public class Knn extends FullTextFunction implements OptionalArgument, VectorFunction, PostAnalysisPlanVerificationAware {
-    private final Logger log = LogManager.getLogger(getClass());
 
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(Expression.class, "Knn", Knn::readFrom);
 
     private final Expression field;
     // k is not serialized as it's already included in the query builder on the rewrite step before being sent to data nodes
-    private final transient Expression k;
+    private final transient Integer k;
     private final Expression options;
     // Expressions to be used as prefilters in knn query
     private final List<Expression> filterExpressions;
@@ -107,13 +102,6 @@ public Knn(
             type = { "dense_vector" },
             description = "Vector value to find top nearest neighbours for."
         ) Expression query,
-        @Param(
-            name = "k",
-            type = { "integer" },
-            description = "The number of nearest neighbors to return from each shard. "
-                + "Elasticsearch collects k results from each shard, then merges them to find the global top results. "
-                + "This value must be less than or equal to num_candidates."
-        ) Expression k,
         @MapParam(
             name = "options",
             params = {
@@ -125,12 +113,13 @@ public Knn(
                         + "Defaults to 1.0."
                 ),
                 @MapParam.MapParamEntry(
-                    name = "num_candidates",
+                    name = "min_candidates",
                     type = "integer",
                     valueHint = { "10" },
-                    description = "The number of nearest neighbor candidates to consider per shard while doing knn search. "
-                        + "Cannot exceed 10,000. Increasing num_candidates tends to improve the accuracy of the final results. "
-                        + "Defaults to 1.5 * k"
+                    description = "The minimum number of nearest neighbor candidates to consider per shard while doing knn search. " +
+                        " KNN may use a higher number of candidates in case the query can't use a approximate results. "
+                        + "Cannot exceed 10,000. Increasing min_candidates tends to improve the accuracy of the final results. "
+                        + "Defaults to 1.5 * LIMIT used for the query."
                 ),
                 @MapParam.MapParamEntry(
                     name = "similarity",
@@ -152,32 +141,29 @@ public Knn(
             optional = true
         ) Expression options
     ) {
-        this(source, field, query, k, options, null, List.of());
+        this(source, field, query, options, null, null, List.of());
     }
 
     public Knn(
         Source source,
         Expression field,
         Expression query,
-        Expression k,
         Expression options,
+        Integer k,
         QueryBuilder queryBuilder,
         List<Expression> filterExpressions
     ) {
-        super(source, query, expressionList(field, query, k, options), queryBuilder);
+        super(source, query, expressionList(field, query, options), queryBuilder);
         this.field = field;
         this.k = k;
         this.options = options;
         this.filterExpressions = filterExpressions;
     }
 
-    private static List<Expression> expressionList(Expression field, Expression query, Expression k, Expression options) {
+    private static List<Expression> expressionList(Expression field, Expression query, Expression options) {
         List<Expression> result = new ArrayList<>();
         result.add(field);
         result.add(query);
-        if (k != null) {
-            result.add(k);
-        }
         if (options != null) {
             result.add(options);
         }
@@ -188,7 +174,7 @@ public Expression field() {
         return field;
     }
 
-    public Expression k() {
+    public Integer k() {
         return k;
     }
 
@@ -207,7 +193,7 @@ public DataType dataType() {
 
     @Override
     protected TypeResolution resolveParams() {
-        return resolveField().and(resolveQuery()).and(resolveK()).and(Options.resolve(options(), source(), FOURTH, ALLOWED_OPTIONS));
+        return resolveField().and(resolveQuery()).and(Options.resolve(options(), source(), THIRD, ALLOWED_OPTIONS));
     }
 
     private TypeResolution resolveField() {
@@ -227,14 +213,9 @@ private TypeResolution resolveQuery() {
         return TypeResolution.TYPE_RESOLVED;
     }
 
-    private TypeResolution resolveK() {
-        if (k == null) {
-            // Function has already been rewritten and included in QueryBuilder - otherwise parsing would have failed
-            return TypeResolution.TYPE_RESOLVED;
-        }
-
-        return isType(k(), dt -> dt == INTEGER, sourceText(), THIRD, "integer").and(isFoldable(k(), sourceText(), THIRD))
-            .and(isNotNull(k(), sourceText(), THIRD));
+    public Knn replaceK(Integer k) {
+        Check.notNull(k, "k must not be null");
+        return new Knn(source(), field(), query(), options(), k, queryBuilder(), filterExpressions());
     }
 
     public List<Number> queryAsObject() {
@@ -248,16 +229,9 @@ public List<Number> queryAsObject() {
         throw new EsqlIllegalArgumentException(format(null, "Query value must be a list of numbers in [{}], found [{}]", source(), query));
     }
 
-    int getKIntValue() {
-        if (k() instanceof Literal literal) {
-            return (int) (Number) literal.value();
-        }
-        throw new EsqlIllegalArgumentException(format(null, "K value must be a constant integer in [{}], found [{}]", source(), k()));
-    }
-
     @Override
     public Expression replaceQueryBuilder(QueryBuilder queryBuilder) {
-        return new Knn(source(), field(), query(), k(), options(), queryBuilder, filterExpressions());
+        return new Knn(source(), field(), query(), options(), k(), queryBuilder, filterExpressions());
     }
 
     @Override
@@ -273,15 +247,12 @@ public Translatable translatable(LucenePushdownPredicates pushdownPredicates) {
 
     @Override
     protected Query translate(LucenePushdownPredicates pushdownPredicates, TranslatorHandler handler) {
+        assert k() != null : "Knn function must have a k value set before translation";
         var fieldAttribute = Match.fieldAsFieldAttribute(field());
 
         Check.notNull(fieldAttribute, "Knn must have a field attribute as the first argument");
         String fieldName = getNameFromFieldAttribute(fieldAttribute);
         float[] queryAsFloats = queryAsFloats();
-        int kValue = getKIntValue();
-
-        Map<String, Object> opts = queryOptions();
-        opts.put(K_FIELD.getPreferredName(), kValue);
 
         List<QueryBuilder> filterQueries = new ArrayList<>();
         for (Expression filterExpression : filterExpressions()) {
@@ -295,7 +266,7 @@ protected Query translate(LucenePushdownPredicates pushdownPredicates, Translato
             }
         }
 
-        return new KnnQuery(source(), fieldName, queryAsFloats, opts, filterQueries);
+        return new KnnQuery(source(), fieldName, queryAsFloats, k(), queryOptions(), filterQueries);
     }
 
     private float[] queryAsFloats() {
@@ -308,7 +279,7 @@ private float[] queryAsFloats() {
     }
 
     public Expression withFilters(List<Expression> filterExpressions) {
-        return new Knn(source(), field(), query(), k(), options(), queryBuilder(), filterExpressions);
+        return new Knn(source(), field(), query(), options(), k(), queryBuilder(), filterExpressions);
     }
 
     private Map<String, Object> queryOptions() throws InvalidArgumentException {
@@ -343,16 +314,16 @@ public Expression replaceChildren(List<Expression> newChildren) {
             source(),
             newChildren.get(0),
             newChildren.get(1),
-            newChildren.get(2),
-            newChildren.size() > 3 ? newChildren.get(3) : null,
+            newChildren.size() > 2 ? newChildren.get(2) : null,
+            k(),
             queryBuilder(),
             filterExpressions()
         );
     }
 
     @Override
     protected NodeInfo<? extends Expression> info() {
-        return NodeInfo.create(this, Knn::new, field(), query(), k(), options(), queryBuilder(), filterExpressions());
+        return NodeInfo.create(this, Knn::new, field(), query(), options(), k(), queryBuilder(), filterExpressions());
     }
 
     @Override

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/vector/Knn.java

@@ -44,6 +44,7 @@
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PushDownInferencePlan;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PushDownJoinPastProject;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PushDownRegexExtract;
+import org.elasticsearch.xpack.esql.optimizer.rules.logical.PushLimitToKnn;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.RemoveStatsOverride;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.ReplaceAggregateAggExpressionWithEval;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.ReplaceAggregateNestedExpressionWithEval;
@@ -192,6 +193,7 @@ protected static Batch<LogicalPlan> operators(boolean local) {
             new PruneColumns(),
             new PruneLiteralsInOrderBy(),
             new PushDownAndCombineLimits(),
+            new PushLimitToKnn(),
             new PushDownAndCombineFilters(),
             new PushDownConjunctionsToKnnPrefilters(),
             new PushDownAndCombineSample(),

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

@@ -0,0 +1,70 @@
+/*
+ * 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.optimizer.rules.logical;
+
+import org.elasticsearch.xpack.esql.core.expression.Expression;
+import org.elasticsearch.xpack.esql.core.util.Holder;
+import org.elasticsearch.xpack.esql.expression.function.vector.Knn;
+import org.elasticsearch.xpack.esql.optimizer.LogicalOptimizerContext;
+import org.elasticsearch.xpack.esql.plan.logical.Aggregate;
+import org.elasticsearch.xpack.esql.plan.logical.Filter;
+import org.elasticsearch.xpack.esql.plan.logical.Limit;
+import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
+import org.elasticsearch.xpack.esql.plan.logical.TopN;
+import org.elasticsearch.xpack.esql.plan.logical.UnaryPlan;
+import org.elasticsearch.xpack.esql.plan.logical.inference.Rerank;
+
+/**
+ * Traverses the logical plan and pushes down the limit to the KNN function(s) in filter expressions, so KNN can use
+ * it to set k if not specified.
+ */
+public class PushLimitToKnn extends OptimizerRules.ParameterizedOptimizerRule<Limit, LogicalOptimizerContext> {
+
+    public PushLimitToKnn() {
+        super(OptimizerRules.TransformDirection.DOWN);
+    }
+
+    @Override
+    public LogicalPlan rule(Limit limit, LogicalOptimizerContext ctx) {
+        Holder<Boolean> breakerReached = new Holder<>(false);
+        Holder<Boolean> firstLimit = new Holder<>(false);
+        return limit.transformDown(plan -> {
+            if (breakerReached.get()) {
+                // We reached a breaker and don't want to continue processing
+                return plan;
+            }
+            if (plan instanceof Filter filter) {
+                Expression limitAppliedExpression = limitFilterExpressions(filter.condition(), limit, ctx);
+                if (limitAppliedExpression.equals(filter.condition()) == false) {
+                    return filter.with(limitAppliedExpression);
+                }
+            } else if (plan instanceof Limit) {
+                // Break if it's not the initial limit
+                breakerReached.set(firstLimit.get());
+                firstLimit.set(true);
+            } else if (plan instanceof TopN || plan instanceof Rerank || plan instanceof Aggregate) {
+                breakerReached.set(true);
+            }
+
+            return plan;
+        });
+    }
+
+    /**
+     * Applies a limit to the filter expressions of a condition. Some filter expressions, such as KNN function,
+     * can be optimized by applying the limit directly to them.
+     */
+    private Expression limitFilterExpressions(Expression condition, Limit limit, LogicalOptimizerContext ctx) {
+        return condition.transformDown(exp -> {
+            if (exp instanceof Knn knn) {
+                return knn.replaceK((Integer) limit.limit().fold(ctx.foldCtx()));
+            }
+            return exp;
+        });
+    }
+}

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

@@ -32,9 +32,11 @@ public class KnnQuery extends Query {
     private final List<QueryBuilder> filterQueries;
 
     public static final String RESCORE_OVERSAMPLE_FIELD = "rescore_oversample";
+    private final Integer k;
 
-    public KnnQuery(Source source, String field, float[] query, Map<String, Object> options, List<QueryBuilder> filterQueries) {
+    public KnnQuery(Source source, String field, float[] query, Integer k, Map<String, Object> options, List<QueryBuilder> filterQueries) {
         super(source);
+        this.k = k;
         assert options != null;
         this.field = field;
         this.query = query;
@@ -44,7 +46,6 @@ public KnnQuery(Source source, String field, float[] query, Map<String, Object>
 
     @Override
     protected QueryBuilder asBuilder() {
-        Integer k = (Integer) options.get(K_FIELD.getPreferredName());
         Integer numCands = (Integer) options.get(NUM_CANDS_FIELD.getPreferredName());
         RescoreVectorBuilder rescoreVectorBuilder = null;
         Float oversample = (Float) options.get(RESCORE_OVERSAMPLE_FIELD);

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/querydsl/query/KnnQuery.java

@@ -121,7 +121,7 @@ private static List<TestCaseSupplier> addFunctionNamedParams(List<TestCaseSuppli
 
     @Override
     protected Expression build(Source source, List<Expression> args) {
-        Knn knn = new Knn(source, args.get(0), args.get(1), args.get(2), args.size() > 3 ? args.get(3) : null);
+        Knn knn = new Knn(source, args.get(0), args.get(1), args.size() > 2 ? args.get(2) : null);
         // We need to add the QueryBuilder to the match expression, as it is used to implement equals() and hashCode() and
         // thus test the serialization methods. But we can only do this if the parameters make sense .
         if (args.get(0) instanceof FieldAttribute && args.get(1).foldable()) {