Utilize optimized dot_product where possible when calculating vector …

…magnitude (elastic#99448) Lucene provides an optimized `dot_product` calculation for vectors. We should use that when calculating a vector's magnitude.
nik9000 · Sep 12, 2023 · 2072be9 · 2072be9
1 parent e26dca4
commit 2072be9
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 24 deletions.
diff --git a/server/src/main/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldMapper.java b/server/src/main/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldMapper.java
@@ -26,6 +26,7 @@
 import org.apache.lucene.search.KnnFloatVectorQuery;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.util.BytesRef;
+import org.apache.lucene.util.VectorUtil;
 import org.elasticsearch.common.xcontent.support.XContentMapValues;
 import org.elasticsearch.index.IndexVersion;
 import org.elasticsearch.index.fielddata.FieldDataContext;
@@ -859,10 +860,7 @@ public Query createKnnQuery(byte[] queryVector, int numCands, Query filter, Floa
             }
 
             if (similarity == VectorSimilarity.DOT_PRODUCT || similarity == VectorSimilarity.COSINE) {
-                float squaredMagnitude = 0.0f;
-                for (byte b : queryVector) {
-                    squaredMagnitude += b * b;
-                }
+                float squaredMagnitude = VectorUtil.dotProduct(queryVector, queryVector);
                 elementType.checkVectorMagnitude(similarity, elementType.errorByteElementsAppender(queryVector), squaredMagnitude);
             }
             Query knnQuery = new KnnByteVectorQuery(name(), queryVector, numCands, filter);
@@ -891,10 +889,7 @@ public Query createKnnQuery(float[] queryVector, int numCands, Query filter, Flo
             elementType.checkVectorBounds(queryVector);
 
             if (similarity == VectorSimilarity.DOT_PRODUCT || similarity == VectorSimilarity.COSINE) {
-                float squaredMagnitude = 0.0f;
-                for (float e : queryVector) {
-                    squaredMagnitude += e * e;
-                }
+                float squaredMagnitude = VectorUtil.dotProduct(queryVector, queryVector);
                 elementType.checkVectorMagnitude(similarity, elementType.errorFloatElementsAppender(queryVector), squaredMagnitude);
             }
             Query knnQuery = switch (elementType) {

diff --git a/server/src/main/java/org/elasticsearch/index/mapper/vectors/VectorEncoderDecoder.java b/server/src/main/java/org/elasticsearch/index/mapper/vectors/VectorEncoderDecoder.java
@@ -9,6 +9,7 @@
 package org.elasticsearch.index.mapper.vectors;
 
 import org.apache.lucene.util.BytesRef;
+import org.apache.lucene.util.VectorUtil;
 import org.elasticsearch.index.IndexVersion;
 
 import java.nio.ByteBuffer;
@@ -46,12 +47,7 @@ public static float decodeMagnitude(IndexVersion indexVersion, BytesRef vectorBR
      * Calculates vector magnitude
      */
     private static float calculateMagnitude(float[] decodedVector) {
-        double magnitude = 0.0f;
-        for (int i = 0; i < decodedVector.length; i++) {
-            magnitude += decodedVector[i] * decodedVector[i];
-        }
-        magnitude = Math.sqrt(magnitude);
-        return (float) magnitude;
+        return (float) Math.sqrt(VectorUtil.dotProduct(decodedVector, decodedVector));
     }
 
     public static float getMagnitude(IndexVersion indexVersion, BytesRef vectorBR, float[] decodedVector) {

diff --git a/server/src/main/java/org/elasticsearch/script/field/vectors/DenseVector.java b/server/src/main/java/org/elasticsearch/script/field/vectors/DenseVector.java
@@ -8,6 +8,8 @@
 
 package org.elasticsearch.script.field.vectors;
 
+import org.apache.lucene.util.VectorUtil;
+
 import java.util.List;
 
 /**
@@ -151,11 +153,7 @@ default double cosineSimilarity(Object queryVector) {
     int size();
 
     static float getMagnitude(byte[] vector) {
-        int mag = 0;
-        for (int elem : vector) {
-            mag += elem * elem;
-        }
-        return (float) Math.sqrt(mag);
+        return (float) Math.sqrt(VectorUtil.dotProduct(vector, vector));
     }
 
     static float getMagnitude(byte[] vector, int dims) {
@@ -170,11 +168,7 @@ static float getMagnitude(byte[] vector, int dims) {
     }
 
     static float getMagnitude(float[] vector) {
-        double mag = 0.0f;
-        for (float elem : vector) {
-            mag += elem * elem;
-        }
-        return (float) Math.sqrt(mag);
+        return (float) Math.sqrt(VectorUtil.dotProduct(vector, vector));
     }
 
     static float getMagnitude(List<Number> vector) {