ESQL - dense vector support cosine normalization #132721

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

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

ESQL - dense vector support cosine normalization #132721

ESQL - dense vector support cosine normalization #132721

Uh oh!

carlosdelest commented Aug 12, 2025

carlosdelest Aug 12, 2025

carlosdelest Aug 12, 2025

carlosdelest Aug 12, 2025

carlosdelest Aug 12, 2025

carlosdelest Aug 12, 2025

elasticsearchmachine commented Aug 12, 2025

elasticsearchmachine commented Aug 12, 2025

carlosdelest commented Aug 12, 2025

kderusso left a comment

kderusso Aug 13, 2025

carlosdelest Aug 13, 2025

Uh oh!

kderusso Aug 13, 2025

carlosdelest Aug 13, 2025

kderusso Aug 13, 2025

carlosdelest Aug 13, 2025

ChrisHegarty left a comment

Uh oh!

Labels

4 participants

ESQL - dense vector support cosine normalization #132721

ESQL - dense vector support cosine normalization #132721

Uh oh!

Conversation

carlosdelest commented Aug 12, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

elasticsearchmachine commented Aug 12, 2025

elasticsearchmachine commented Aug 12, 2025

carlosdelest commented Aug 12, 2025

kderusso left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ChrisHegarty left a comment

Choose a reason for hiding this comment

Uh oh!

Labels

4 participants