ESQL - dense vector support cosine normalization by carlosdelest · Pull Request #132721 · elastic/elasticsearch

carlosdelest · 2025-08-12T12:22:47Z

ESQL did not support cosine normalization. Vectors were retrieved directly from doc values without applying denormalization.

This PR fixes that, including test cases for the different vectors similarities.

carlosdelest · 2025-08-12T12:24:09Z

server/src/main/java/org/elasticsearch/index/mapper/BlockDocValuesReader.java

+        private final DenseVectorFieldMapper.DenseVectorFieldType fieldType;

-        public DenseVectorBlockLoader(String fieldName, int dimensions) {
+        public DenseVectorBlockLoader(String fieldName, int dimensions, DenseVectorFieldMapper.DenseVectorFieldType fieldType) {


Passing the DenseVectorFieldType looked better than including other specific attributes needed for normalization (fieldType.name(), IndexVersion). I opted for creating a isNormalized() method in the DenseVectorFieldType instead.

carlosdelest · 2025-08-12T12:24:30Z

server/src/main/java/org/elasticsearch/index/mapper/BlockDocValuesReader.java

-        private final FloatVectorValues floatVectorValues;
-        private final KnnVectorValues.DocIndexIterator iterator;
-        private final int dimensions;
+    private abstract static class DenseVectorValuesBlockReader<T extends KnnVectorValues> extends BlockDocValuesReader {


Using an abstract class is previous work for supporting byte element type.

carlosdelest · 2025-08-12T12:26:30Z

server/src/main/java/org/elasticsearch/index/mapper/BlockDocValuesReader.java

+            assert floats.length == dimensions
+                : "unexpected dimensions for vector value; expected " + dimensions + " but got " + floats.length;
+
+            assert magnitudeDocValues.advanceExact(iterator.docID());


I tried to use DenormalizedCosineFloatVectorValues instead, but that class depends on the ordinal and not the docId() for retrieving the magnitude.

Using the docId() to retrieve the magnitude works.

carlosdelest · 2025-08-12T12:27:07Z

server/src/main/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldMapper.java

-                            public CacheHelper getCoreCacheHelper() {
-                                return r.getCoreCacheHelper();
-                            }
+                    denseVectorFieldType.isNormalized() && denseVectorFieldType.indexed ? r -> new FilterLeafReader(r) {


Used the isNormalized() method to improve readability - that messed up some indentation

carlosdelest · 2025-08-12T12:27:42Z

...n/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/DenseVectorFieldTypeIT.java

-        for (String indexType : DENSE_VECTOR_INDEX_TYPES) {
-            params.add(new Object[] { indexType, true, false });
-        }
+        Supplier<ElementType> elementTypeProvider = () -> ElementType.FLOAT;


Testing for all combinations took > 1 min. I decided to rely on randomization for testing the different combinations.

…r-support-normalization' into non-issue/esql-dense-vector-support-normalization # Conflicts: # x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/DenseVectorFieldTypeIT.java

…ation

elasticsearchmachine · 2025-08-12T17:09:13Z

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

elasticsearchmachine · 2025-08-12T17:09:13Z

Pinging @elastic/es-search-relevance (Team:Search Relevance)

carlosdelest · 2025-08-12T17:11:24Z

@nik9000 I'd appreciate your review here, as the doc values reading was tricky. I'd like to ensure I am doing the right thing!

…ation

…r-support-normalization' into non-issue/esql-dense-vector-support-normalization

kderusso

Overall changes look good, but I'll defer to someone more expert in ES|QL and KNN.

kderusso · 2025-08-13T12:22:40Z

server/src/main/java/org/elasticsearch/index/mapper/BlockDocValuesReader.java

+
+        @Override
+        protected void appendDoc(BlockLoader.FloatBuilder builder) throws IOException {
+            float[] floats = vectorValues.vectorValue(iterator.index());


Could we refactor checking the dimension count out, since it's shared between both block readers?

Yes, there's a dimension() common method for KnnVectorValues that I wasn't aware of. Done in e98bbcb

...n/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/DenseVectorFieldTypeIT.java

kderusso · 2025-08-13T12:31:41Z

...n/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/DenseVectorFieldTypeIT.java

 public class DenseVectorFieldTypeIT extends AbstractEsqlIntegTestCase {

-    private static final Set<String> DENSE_VECTOR_INDEX_TYPES = Set.of(
+    public static final Set<String> ALL_DENSE_VECTOR_INDEX_TYPES = Set.of(


Would it make sense to use VectorIndexType.values() here to create the set, rather than hard coding the specific index types?

Yes, done in 0cb587a

kderusso · 2025-08-13T12:36:43Z

...gin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/plugin/KnnFunctionIT.java

+    private final String indexType;
+
+    @ParametersFactory
+    public static Iterable<Object[]> parameters() throws Exception {


You could simplify this a bit, with something like this (untested):

public static Iterable<Object[]> parameters() throws Exception { return ALL_DENSE_VECTOR_INDEX_TYPES.stream() .filter(indexType -> indexType.contains("flat") == false) .map(indexType -> new Object[] { DenseVectorFieldMapper.ElementType.FLOAT, indexType }) .collect(Collectors.toList()); }

I find good ol' loops better in terms of clarity here. Also, soon enough we'll need to add the bytes as well.

…vector-support-normalization

ChrisHegarty

LGTM

…ation

…vector-support-normalization # Conflicts: # server/src/main/java/org/elasticsearch/index/mapper/BlockDocValuesReader.java # x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/DenseVectorFieldTypeIT.java # x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/plugin/KnnFunctionIT.java

* Take into account normalization for dense vector support * Fix cherry pick * [CI] Auto commit changes from spotless * Remove debugging code * Check that we may not have magnitudes at all, or for normalized vectors * Better parameterized test * Refactor dimension check * Refactor index names * Remove comment * Fix test --------- Co-authored-by: elasticsearchmachine <infra-root+elasticsearchmachine@elastic.co>

* upstream/main: (278 commits) ESQL - dense vector support cosine normalization (elastic#132721) [ML] Add support for dimensions in google vertex ai request (elastic#132689) ESQL - Add byte element support for dense_vector data type (elastic#131863) ESQL: Fix async operator warnings not always sent when blocking (elastic#132744) Method not needed anymore (elastic#132912) [Test] Excercise shutdown more reliably in snapshot stress IT (elastic#132909) Update Gradle shadow plugin to 9.0.1 (elastic#132637) Mute org.elasticsearch.test.rest.yaml.CcsCommonYamlTestSuiteIT test {p0=search/410_named_queries/named_queries_with_score} elastic#132906 Update docker.elastic.co/wolfi/chainguard-base-fips:latest Docker digest to fa6cb69 (elastic#132735) Remove unnecessary calls to fold() (elastic#131870) Use consistent terminology for transport version resources/references (elastic#132882) Mute org.elasticsearch.test.rest.yaml.CcsCommonYamlTestSuiteIT test {p0=search.vectors/40_knn_search_cosine/kNN search only regular query} elastic#132890 Finalize release notes for v9.1.2 release (elastic#132745) Finalize release notes for v9.0.5 release (elastic#132718) Move inner records out of TransportVersionUtils (elastic#132872) Add support for Lookup Join on Multiple Fields (elastic#131559) Bootstrap PR-based benchmarks (elastic#132717) Refactor MetadataIndexTemplateService to use template maps instead of project metadata (elastic#132662) [Gradle] Update nebula ospackage plugin to 12.1.0 (elastic#132640) Mute org.elasticsearch.xpack.esql.CsvTests test {csv-spec:ip.CdirMatchEqualsInsOrs} elastic#132860 ...

carlosdelest added 2 commits August 12, 2025 14:11

Take into account normalization for dense vector support

763fe63

Fix cherry pick

80b48cf

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

carlosdelest commented Aug 12, 2025

View reviewed changes

elasticsearchmachine and others added 5 commits August 12, 2025 12:35

[CI] Auto commit changes from spotless

40edca3

Remove debugging code

8bd7f79

Merge remote-tracking branch 'carlosdelest/non-issue/esql-dense-vecto…

f9447f7

…r-support-normalization' into non-issue/esql-dense-vector-support-normalization # Conflicts: # x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/DenseVectorFieldTypeIT.java

Check that we may not have magnitudes at all, or for normalized vectors

7d2625c

Merge branch 'main' into non-issue/esql-dense-vector-support-normaliz…

968c73f

…ation

carlosdelest marked this pull request as ready for review August 12, 2025 17:08

elasticsearchmachine added the Team:Analytics Meta label for analytical engine team (ESQL/Aggs/Geo) label Aug 12, 2025

carlosdelest requested review from ChrisHegarty, ioanatia and nik9000 August 12, 2025 17:10

carlosdelest requested review from afoucret and svilen-mihaylov-elastic August 12, 2025 17:11

Merge branch 'main' into non-issue/esql-dense-vector-support-normaliz…

84f56ac

…ation

carlosdelest requested review from kderusso and tteofili August 13, 2025 07:31

carlosdelest added 2 commits August 13, 2025 11:39

Better parameterized test

4371b69

Merge remote-tracking branch 'carlosdelest/non-issue/esql-dense-vecto…

c19db57

…r-support-normalization' into non-issue/esql-dense-vector-support-normalization

kderusso reviewed Aug 13, 2025

View reviewed changes

Refactor dimension check

e98bbcb

carlosdelest added 3 commits August 13, 2025 16:57

Refactor index names

0cb587a

Remove comment

e16c7cf

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

c9d7974

…vector-support-normalization

ChrisHegarty approved these changes Aug 13, 2025

View reviewed changes

carlosdelest and others added 3 commits August 14, 2025 10:17

Merge branch 'main' into non-issue/esql-dense-vector-support-normaliz…

276f98c

…ation

Fix test

8936d33

carlosdelest merged commit 7a34391 into elastic:main Aug 14, 2025
33 checks passed

Conversation

carlosdelest commented Aug 12, 2025

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!

elasticsearchmachine commented Aug 12, 2025

Uh oh!

elasticsearchmachine commented Aug 12, 2025

Uh oh!

carlosdelest commented Aug 12, 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!

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!

ChrisHegarty left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants