ESQL: dense_vector cosine similarity function by carlosdelest · Pull Request #130641 · elastic/elasticsearch

carlosdelest · 2025-07-04T14:50:12Z

tracked in #130828
Implements CosineSimilarityFunction for ES|QL, and adds basic infrastructure for other vector similarity functions.

FROM colors
 | EVAL similarity = V_COSINE(rgb_vector, [0, 255, 255]) 
 | SOR similarity DESC

Adds a base class, VectorSimilarityFunction, that provides the building block for vector similarity functions.

There are pending validations that should be done for the function parameters:

Same number of dimensions
Same underlying dense_vector field type

We can work on these validations as follow ups, as they may depend on field_caps API returning that information.

carlosdelest · 2025-07-04T14:53:49Z

...n/java/org/elasticsearch/xpack/esql/expression/function/vector/VectorSimilarityFunction.java

+                        FloatBlock leftBlock = (FloatBlock) left.get(context).eval(page);
+                        FloatBlock rightBlock = (FloatBlock) right.get(context).eval(page)
+                    ) {
+                        int positionCount = page.getPositionCount();


@ChrisHegarty I'm wondering if this is the right way to provide an evaluation for dense_vector based operations. Besides vector similarity functions, we will create vector operations (add, substract, dot product, etc).

Do you think we should create the necessary infrastructure for template based evaluators, or should having this ad-hoc evaluation work?

Is there anything we should be careful about when doing ad-hoc evaluation for vectorization purposes?

I think that this is fine. The vectorization that we are looking for here is in the comparison operation itself, so when comparing float[]'s. Ultimately tho, we would want to be able to compare against mmap'ed off-heap data, but that is completely separate and can come later - since it would require a block backed by a memory segment. We had similar(ish), though different, with big array blocks. Would need to re-check the details.

… FloatBlocks

carlosdelest · 2025-07-07T09:16:10Z

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

            if (f instanceof In in) {
                return processIn(in);
            }
+            if (f instanceof VectorFunction) {


Needed to change the order to ensure VectorFunction are processed first, as similarity functions are scalar functions as well

carlosdelest · 2025-07-07T09:33:32Z

x-pack/plugin/esql/qa/testFixtures/src/main/resources/vector-cosine-similarity.csv-spec

+required_capability: cosine_vector_similarity_function
+
+row vector = [1, 2, 3] 
+| eval similarity = round(v_cosine(vector, [0, 1, 2]), 3) 


For this to work properly, we need to implement a conversion function so we can convert non-foldable values to dense_vector.

elasticsearchmachine · 2025-07-07T11:58:50Z

Pinging @elastic/es-analytical-engine (Team:Analytics)

kderusso

Nice work @carlosdelest !

kderusso · 2025-07-07T13:54:55Z

...l/src/main/java/org/elasticsearch/xpack/esql/expression/function/vector/VectorWritables.java

+/**
+ * Defines the named writables for vector functions in ESQL.
+ */
+public final class VectorWritables {


Not sure if we need this utility class just yet, but I'll assume you have plans to add more :)

Haha yeah, it's a bit premature yet - but we will be adding a number of vector similarity functions soon enough, and I wanted to provide places where it would be easy to look for them.

kderusso · 2025-07-07T13:55:55Z

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/evaluator/EvalMapper.java

                }
-                var wrapper = BlockUtils.wrapperFor(blockFactory, ElementType.fromJava(multiValue.get(0).getClass()), positions);
+                // dense_vector create internally float values, even if they are specified as doubles
+                ElementType elementType = lit.dataType() == DataType.DENSE_VECTOR


Should this logic be in its own method?

I'd say no as this is a one-liner for getting the correct ElementType - there's no more logic than doing a specific check for dense_vector. I'd say, ff more special cases come into play then let's add it as it will become confusing.

svilen-mihaylov-elastic · 2025-07-07T15:02:17Z

.../src/main/java/org/elasticsearch/xpack/esql/expression/function/vector/CosineSimilarity.java

+
+import static org.apache.lucene.index.VectorSimilarityFunction.COSINE;
+
+public class CosineSimilarity extends VectorSimilarityFunction {


Do you need to subclass different types of functions here? Why not just have an enum which specifies the type in VectorSimilarityFunction?

Good point - I think this aligns better with the current way ESQL functions work. I'm not sure that docs generation work with enums as of now as well.

Happy to review this when adding more functions though!

…ch-functions-basics' into non-issue/esql-vector-search-functions-basics

…-search-functions-basics

ioanatia

are we missing the docs that will be generated for the v_cosine function?
otherwise LGTM

…milarityFunction to extend BinaryScalarFunction

…and folding done

github-actions · 2025-07-08T11:49:48Z

🔍 Preview links for changed docs

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

…nts when casting

carlosdelest · 2025-07-08T12:01:19Z

...n/java/org/elasticsearch/xpack/esql/expression/function/vector/VectorSimilarityFunction.java

+/**
+ * Base class for vector similarity functions, which compute a similarity score between two dense vectors
+ */
+public abstract class VectorSimilarityFunction extends BinaryScalarFunction implements EvaluatorMapper, VectorFunction {


Now VectorSimilarityFunction extends BinaryScalarFunction. That brings some simplifications to the code as we already have two params.

carlosdelest · 2025-07-08T12:02:10Z

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/analysis/AnalyzerTests.java

+    public void testDenseVectorImplicitCastingSimilarityFunctions() {
+        if (EsqlCapabilities.Cap.COSINE_VECTOR_SIMILARITY_FUNCTION.isEnabled()) {
+            checkDenseVectorImplicitCastingSimilarityFunction("v_cosine(vector, [0.342, 0.164, 0.234])", List.of(0.342f, 0.164f, 0.234f));
+            checkDenseVectorImplicitCastingSimilarityFunction("v_cosine(vector, [1, 2, 3])", List.of(1f, 2f, 3f));


Checks casting is done for non-float values, and creates a float Literal

carlosdelest · 2025-07-08T12:03:03Z

...icsearch/xpack/esql/expression/function/vector/AbstractVectorSimilarityFunctionTestCase.java

+import static org.elasticsearch.xpack.esql.core.type.DataType.DOUBLE;
+import static org.hamcrest.Matchers.equalTo;
+
+public abstract class AbstractVectorSimilarityFunctionTestCase extends AbstractScalarFunctionTestCase {


New test case added that extends AbstractScalarFunctionTestCase. This brings quite a few tests like checking what happens with null values, evaluator type checks, etc.

carlosdelest · 2025-07-08T12:03:44Z

...test/java/org/elasticsearch/xpack/esql/expression/function/vector/CosineSimilarityTests.java

+import java.util.function.Supplier;
+
+@FunctionName("v_cosine")
+public class CosineSimilarityTests extends AbstractVectorSimilarityFunctionTestCase {


New functions test cases should be simple, all the heavy lifting is done in the abstract class

carlosdelest · 2025-07-08T12:05:01Z

are we missing the docs that will be generated for the v_cosine function?

@ioanatia 🤦 yes we were. There were no AbstractFunctionTestCase test. I just added that along with some changes to implicit casting. Thanks!

…-search-functions-basics

…ch-functions-basics' into non-issue/esql-vector-search-functions-basics

carlosdelest added 6 commits July 4, 2025 13:12

Add CosineSimilarity function

0c18eba

Add IT

31faca1

Refactor

16542b5

Extract superclass, use overriden method

5098373

Use lambda instead of overriden methods

b441d60

Refactoring

4508154

carlosdelest added >non-issue :Analytics/ES|QL AKA ESQL Team:Search Relevance Meta label for the Search Relevance team in Elasticsearch v9.2.0 labels Jul 4, 2025

Refactoring

58bd1c0

carlosdelest commented Jul 4, 2025

View reviewed changes

carlosdelest added 7 commits July 4, 2025 18:42

Allow using literals for dense_vectors, converting them internally to…

fc88621

… FloatBlocks

Add test

597741d

Test non-null args

85e2426

Rename to v_cosine, add analyzer tests for implicit casting

6cc2115

First CSV tests

1a9b44c

Add tests

a53eab3

Spotless

8c80b72

carlosdelest changed the title ~~ESQL: Basic infrastructure for dense vector similarity functions~~ ESQL: dense_vector cosine similarity function Jul 7, 2025

carlosdelest commented Jul 7, 2025

View reviewed changes

carlosdelest and others added 3 commits July 7, 2025 11:18

Add comment

7517557

Merge branch 'main' into non-issue/esql-vector-search-functions-basics

58e6ac7

[CI] Auto commit changes from spotless

ebafdf8

carlosdelest commented Jul 7, 2025

View reviewed changes

carlosdelest marked this pull request as ready for review July 7, 2025 11:58

elasticsearchmachine added Team:Analytics Meta label for analytical engine team (ESQL/Aggs/Geo) and removed Team:Search Relevance Meta label for the Search Relevance team in Elasticsearch labels Jul 7, 2025

carlosdelest requested a review from tteofili July 7, 2025 11:58

carlosdelest requested a review from svilen-mihaylov-elastic July 7, 2025 11:58

kderusso reviewed Jul 7, 2025

View reviewed changes

svilen-mihaylov-elastic reviewed Jul 7, 2025

View reviewed changes

carlosdelest added 3 commits July 7, 2025 19:25

Add checks for different number of dimensions

3ab9a72

Merge remote-tracking branch 'carlosdelest/non-issue/esql-vector-sear…

b1d6f85

…ch-functions-basics' into non-issue/esql-vector-search-functions-basics

Merge remote-tracking branch 'origin/main' into non-issue/esql-vector…

4b0b772

…-search-functions-basics

ioanatia reviewed Jul 8, 2025

View reviewed changes

carlosdelest added 4 commits July 8, 2025 12:49

Add test infrastructure for VectorSimilarityFunction. Change VectorSi…

53e96f9

…milarityFunction to extend BinaryScalarFunction

Ensure casting is done using floats so we get the appropriate blocks …

290dbe1

…and folding done

Generate docs for function

312a727

Generate docs for function

08364ee

carlosdelest added 3 commits July 8, 2025 13:51

Fix tests

fb3bec7

Remove unnecessary change to BlockUtis now that we create float eleme…

3d82d86

…nts when casting

Remove unnecessary change to BlockUtis now that we create float eleme…

7601522

…nts when casting

carlosdelest commented Jul 8, 2025

View reviewed changes

carlosdelest requested a review from ioanatia July 8, 2025 12:05

ioanatia approved these changes Jul 8, 2025

View reviewed changes

carlosdelest mentioned this pull request Jul 8, 2025

ES|QL - dense_vector exact (brute force) search support #130828

Closed

11 tasks

Merge branch 'main' into non-issue/esql-vector-search-functions-basics

2a9e322

carlosdelest enabled auto-merge (squash) July 8, 2025 14:58

carlosdelest and others added 7 commits July 8, 2025 18:36

Merge remote-tracking branch 'origin/main' into non-issue/esql-vector…

dd09bf8

…-search-functions-basics

Fix telemetry test

2400b7a

Merge remote-tracking branch 'carlosdelest/non-issue/esql-vector-sear…

2bdb35d

…ch-functions-basics' into non-issue/esql-vector-search-functions-basics

Merge branch 'main' into non-issue/esql-vector-search-functions-basics

81f26ec

Merge branch 'main' into non-issue/esql-vector-search-functions-basics

707b3c2

Merge branch 'main' into non-issue/esql-vector-search-functions-basics

7fcdb36

[CI] Auto commit changes from spotless

5e9bcd2

carlosdelest merged commit f1ddd4c into elastic:main Jul 15, 2025
32 of 33 checks passed


		import static org.apache.lucene.index.VectorSimilarityFunction.COSINE;

		public class CosineSimilarity extends VectorSimilarityFunction {

Conversation

carlosdelest commented Jul 4, 2025 • edited by ioanatia Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

elasticsearchmachine commented Jul 7, 2025

Uh oh!

kderusso left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ioanatia left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions bot commented Jul 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔍 Preview links for changed docs

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

carlosdelest commented Jul 8, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

carlosdelest commented Jul 4, 2025 •

edited by ioanatia

Loading

github-actions bot commented Jul 8, 2025 •

edited

Loading