Add first version of KnnTests and generated docs

Author: carlosdelest

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

@@ -152,7 +152,8 @@ public boolean equals(Object obj) {
             return false;
         }
 
-        return Objects.equals(queryBuilder, ((FullTextFunction) obj).queryBuilder);
+        return Objects.equals(queryBuilder, ((FullTextFunction) obj).queryBuilder)
+            && Objects.equals(query, ((FullTextFunction) obj).query);
     }
 
     @Override

docs/reference/query-languages/esql/kibana/definition/functions/knn.json

@@ -72,10 +72,8 @@ public class Knn extends FullTextFunction implements OptionalArgument, VectorFun
     @FunctionInfo(
         returnType = "boolean",
         preview = true,
-        description = """
-            Finds the k nearest vectors to a query vector, as measured by a similarity metric.
-            knn function finds nearest vectors through approximate search on indexed dense_vectors
-            """,
+        description = "Finds the k nearest vectors to a query vector, as measured by a similarity metric. " +
+                "knn function finds nearest vectors through approximate search on indexed dense_vectors.",
         examples = {
             @Example(file = "knn-function", tag = "knn-function"),
             @Example(file = "knn-function", tag = "knn-function-options"), },
@@ -96,7 +94,7 @@ public Knn(
                     name = "boost",
                     type = "float",
                     valueHint = { "2.5" },
-                    description = "Floating point number used to decrease or increase the relevance scores of the query. "
+                    description = "Floating point number used to decrease or increase the relevance scores of the query."
                         + "Defaults to 1.0."
                 ),
                 @MapParam.MapParamEntry(
@@ -120,7 +118,7 @@ public Knn(
                     type = "double",
                     valueHint = { "0.01" },
                     description = "The minimum similarity required for a document to be considered a match. "
-                        + "The similarity value calculated relates to the raw similarity used, not the document score"
+                        + "The similarity value calculated relates to the raw similarity used, not the document score."
                 ),
                 @MapParam.MapParamEntry(
                     name = "rescore_oversample",
@@ -237,12 +235,13 @@ public void writeTo(StreamOutput out) throws IOException {
 
     @Override
     public boolean equals(Object o) {
+        // Knn does not serialize options, as they get included in the query builder. We need to override equals and hashcode to
+        // ignore options when comparing two Knn functions
         if (o == null || getClass() != o.getClass()) return false;
-        if (super.equals(o) == false) return false;
         Knn knn = (Knn) o;
-        return Objects.equals(field, knn.field)
-            && Objects.equals(query(), knn.query())
-            && Objects.equals(queryBuilder(), knn.queryBuilder());
+        return Objects.equals(field(), knn.field())
+                && Objects.equals(query(), knn.query())
+                && Objects.equals(queryBuilder(), knn.queryBuilder());
     }
 
     @Override

docs/reference/query-languages/esql/kibana/docs/functions/knn.md

@@ -24,6 +24,7 @@
 import org.elasticsearch.index.query.TermQueryBuilder;
 import org.elasticsearch.index.query.TermsQueryBuilder;
 import org.elasticsearch.index.query.WildcardQueryBuilder;
+import org.elasticsearch.search.vectors.KnnVectorQueryBuilder;
 import org.elasticsearch.test.EqualsHashCodeTestUtils;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.expression.ExpressionWritables;
@@ -111,6 +112,7 @@ public static NamedWriteableRegistry writableRegistry() {
         entries.add(new NamedWriteableRegistry.Entry(QueryBuilder.class, WildcardQueryBuilder.NAME, WildcardQueryBuilder::new));
         entries.add(new NamedWriteableRegistry.Entry(QueryBuilder.class, RegexpQueryBuilder.NAME, RegexpQueryBuilder::new));
         entries.add(new NamedWriteableRegistry.Entry(QueryBuilder.class, ExistsQueryBuilder.NAME, ExistsQueryBuilder::new));
+        entries.add(new NamedWriteableRegistry.Entry(QueryBuilder.class, KnnVectorQueryBuilder.NAME, KnnVectorQueryBuilder::new));
         entries.add(SingleValueQuery.ENTRY);
         entries.addAll(ExpressionWritables.getNamedWriteables());
         entries.addAll(PlanWritables.getNamedWriteables());

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

@@ -10,25 +10,121 @@
 import com.carrotsearch.randomizedtesting.annotations.Name;
 import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;
 
+import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
+import org.elasticsearch.xpack.esql.core.expression.FieldAttribute;
+import org.elasticsearch.xpack.esql.core.expression.Literal;
+import org.elasticsearch.xpack.esql.core.expression.MapExpression;
 import org.elasticsearch.xpack.esql.core.tree.Source;
+import org.elasticsearch.xpack.esql.core.type.DataType;
+import org.elasticsearch.xpack.esql.expression.function.AbstractFunctionTestCase;
 import org.elasticsearch.xpack.esql.expression.function.TestCaseSupplier;
+import org.elasticsearch.xpack.esql.expression.function.vector.Knn;
+import org.elasticsearch.xpack.esql.io.stream.PlanStreamOutput;
+import org.elasticsearch.xpack.esql.optimizer.rules.physical.local.LucenePushdownPredicates;
 
+import java.util.ArrayList;
 import java.util.List;
 import java.util.function.Supplier;
 
-public class KnnTests extends NoneFieldFullTextFunctionTestCase {
+import static org.elasticsearch.xpack.esql.SerializationTestUtils.serializeDeserialize;
+import static org.elasticsearch.xpack.esql.core.type.DataType.BOOLEAN;
+import static org.elasticsearch.xpack.esql.core.type.DataType.DENSE_VECTOR;
+import static org.elasticsearch.xpack.esql.core.type.DataType.KEYWORD;
+import static org.elasticsearch.xpack.esql.core.type.DataType.UNSUPPORTED;
+import static org.elasticsearch.xpack.esql.planner.TranslatorHandler.TRANSLATOR_HANDLER;
+import static org.hamcrest.Matchers.equalTo;
+
+public class KnnTests extends AbstractFunctionTestCase {
+
     public KnnTests(@Name("TestCase") Supplier<TestCaseSupplier.TestCase> testCaseSupplier) {
-        super(testCaseSupplier);
+        this.testCase = testCaseSupplier.get();
     }
 
     @ParametersFactory
     public static Iterable<Object[]> parameters() {
-        return generateParameters();
+        return parameterSuppliersFromTypedData(addFunctionNamedParams(testCaseSuppliers()));
+    }
+
+    private static List<TestCaseSupplier> testCaseSuppliers() {
+        List<TestCaseSupplier> suppliers = new ArrayList<>();
+
+        suppliers.add(
+            TestCaseSupplier.testCaseSupplier(
+                new TestCaseSupplier.TypedDataSupplier("dense_vector field", KnnTests::randomDenseVector, DENSE_VECTOR),
+                new TestCaseSupplier.TypedDataSupplier("query", KnnTests::randomDenseVector, DENSE_VECTOR, true),
+                (d1, d2) -> equalTo("string"),
+                BOOLEAN,
+                (o1, o2) -> true
+            )
+        );
+
+        return suppliers;
+    }
+
+    private static List<Float> randomDenseVector() {
+        int dimensions = randomIntBetween(64, 128);
+        List<Float> vector = new ArrayList<>();
+        for (int i = 0; i < dimensions; i++) {
+            vector.add(randomFloat());
+        }
+        return vector;
+    }
+
+    /**
+     * Adds function named parameters to all the test case suppliers provided
+     */
+    private static List<TestCaseSupplier> addFunctionNamedParams(List<TestCaseSupplier> suppliers) {
+        // TODO get to a common class with MatchTests
+        List<TestCaseSupplier> result = new ArrayList<>();
+        for (TestCaseSupplier supplier : suppliers) {
+            List<DataType> dataTypes = new ArrayList<>(supplier.types());
+            dataTypes.add(UNSUPPORTED);
+            result.add(new TestCaseSupplier(supplier.name() + ", options", dataTypes, () -> {
+                List<TestCaseSupplier.TypedData> values = new ArrayList<>(supplier.get().getData());
+                values.add(
+                    new TestCaseSupplier.TypedData(
+                        new MapExpression(
+                            Source.EMPTY,
+                            List.of(
+                                new Literal(Source.EMPTY, randomAlphaOfLength(10), KEYWORD)
+                            )
+                        ),
+                        UNSUPPORTED,
+                        "options"
+                    ).forceLiteral()
+                );
+
+                return new TestCaseSupplier.TestCase(values, equalTo("KnnEvaluator"), BOOLEAN, equalTo(true));
+            }));
+        }
+        return result;
     }
 
     @Override
     protected Expression build(Source source, List<Expression> args) {
-        return new Kql(source, args.get(0));
+        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()) {
+            QueryBuilder queryBuilder = TRANSLATOR_HANDLER.asQuery(LucenePushdownPredicates.DEFAULT, knn).toQueryBuilder();
+            knn = (Knn) knn.replaceQueryBuilder(queryBuilder);
+        }
+        return knn;
+    }
+
+    /**
+     * Copy of the overridden method that doesn't check for children size, as the {@code options} child isn't serialized in Match.
+     */
+    @Override
+    protected Expression serializeDeserializeExpression(Expression expression) {
+        Expression newExpression = serializeDeserialize(
+            expression,
+            PlanStreamOutput::writeNamedWriteable,
+            in -> in.readNamedWriteable(Expression.class),
+            testCase.getConfiguration() // The configuration query should be == to the source text of the function for this to work
+        );
+        // Fields use synthetic sources, which can't be serialized. So we use the originals instead.
+        return newExpression.replaceChildren(expression.children());
     }
 }

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

@@ -82,7 +82,7 @@ protected Expression build(Source source, List<Expression> args) {
         // 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()) {
             QueryBuilder queryBuilder = TRANSLATOR_HANDLER.asQuery(LucenePushdownPredicates.DEFAULT, match).toQueryBuilder();
-            match.replaceQueryBuilder(queryBuilder);
+            match = (Match) match.replaceQueryBuilder(queryBuilder);
         }
         return match;
     }