Use byte or float vectors as appropriate

thecoop · thecoop · commit 2540fddbb171 · 2025-02-18T17:30:17.000Z
diff --git a/modules/lang-painless/src/yamlRestTest/resources/rest-api-spec/test/painless/145_dense_vector_byte_basic.yml b/modules/lang-painless/src/yamlRestTest/resources/rest-api-spec/test/painless/145_dense_vector_byte_basic.yml
@@ -126,18 +126,18 @@ setup:
               script:
                 source: "dotProduct(params.query_vector, 'vector')"
                 params:
-                  query_vector: [0.0, 111.0, -13.0, 14.0, -124.0]
+                  query_vector: [0.5, 111.3, -13.0, 14.8, -156.0]
 
   - match: {hits.total: 3}
 
   - match: {hits.hits.0._id: "2"}
-  - match: {hits.hits.0._score: 28732.0}
+  - match: {hits.hits.0._score: 32865.2}
 
   - match: {hits.hits.1._id: "3"}
-  - match: {hits.hits.1._score: 17439.0}
+  - match: {hits.hits.1._score: 21413.4}
 
   - match: {hits.hits.2._id: "1"}
-  - match: {hits.hits.2._score: 1632.0}
+  - match: {hits.hits.2._score: 1862.3}
 ---
 "Cosine Similarity":
   - do:
@@ -251,18 +251,18 @@ setup:
               script:
                 source: "cosineSimilarity(params.query_vector, 'vector')"
                 params:
-                  query_vector: [0.0, 111.0, -13.0, 14.0, -124.0]
+                  query_vector: [0.5, 111.3, -13.0, 14.8, -156.0]
 
   - match: {hits.total: 3}
 
   - match: {hits.hits.0._id: "2"}
-  - gte: {hits.hits.0._score: 0.995}
-  - lte: {hits.hits.0._score: 0.998}
+  - gte: {hits.hits.0._score: 0.989}
+  - lte: {hits.hits.0._score: 0.992}
 
   - match: {hits.hits.1._id: "3"}
-  - gte: {hits.hits.1._score: 0.829}
-  - lte: {hits.hits.1._score: 0.832}
+  - gte: {hits.hits.1._score: 0.885}
+  - lte: {hits.hits.1._score: 0.888}
 
   - match: {hits.hits.2._id: "1"}
-  - gte: {hits.hits.2._score: 0.509}
-  - lte: {hits.hits.2._score: 0.512}
+  - gte: {hits.hits.2._score: 0.505}
+  - lte: {hits.hits.2._score: 0.508}
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
@@ -346,16 +346,17 @@ IndexFieldData.Builder fielddataBuilder(DenseVectorFieldType denseVectorFieldTyp
             }
 
             @Override
-            public void checkVectorBounds(float[] vector) {
-                checkNanAndInfinite(vector);
-
-                StringBuilder errorBuilder = null;
+            StringBuilder checkVectorErrors(float[] vector) {
+                StringBuilder errors = checkNanAndInfinite(vector);
+                if (errors != null) {
+                    return errors;
+                }
 
                 for (int index = 0; index < vector.length; ++index) {
                     float value = vector[index];
 
                     if (value % 1.0f != 0.0f) {
-                        errorBuilder = new StringBuilder(
+                        errors = new StringBuilder(
                             "element_type ["
                                 + this
                                 + "] vectors only support non-decimal values but found decimal value ["
@@ -368,7 +369,7 @@ public void checkVectorBounds(float[] vector) {
                     }
 
                     if (value < Byte.MIN_VALUE || value > Byte.MAX_VALUE) {
-                        errorBuilder = new StringBuilder(
+                        errors = new StringBuilder(
                             "element_type ["
                                 + this
                                 + "] vectors only support integers between ["
@@ -385,9 +386,7 @@ public void checkVectorBounds(float[] vector) {
                     }
                 }
 
-                if (errorBuilder != null) {
-                    throw new IllegalArgumentException(appendErrorElements(errorBuilder, vector).toString());
-                }
+                return errors;
             }
 
             @Override
@@ -614,8 +613,8 @@ public FloatVectorValues getFloatVectorValues(String fieldName) throws IOExcepti
             }
 
             @Override
-            public void checkVectorBounds(float[] vector) {
-                checkNanAndInfinite(vector);
+            StringBuilder checkVectorErrors(float[] vector) {
+                return checkNanAndInfinite(vector);
             }
 
             @Override
@@ -768,16 +767,17 @@ IndexFieldData.Builder fielddataBuilder(DenseVectorFieldType denseVectorFieldTyp
             }
 
             @Override
-            public void checkVectorBounds(float[] vector) {
-                checkNanAndInfinite(vector);
-
-                StringBuilder errorBuilder = null;
+            StringBuilder checkVectorErrors(float[] vector) {
+                StringBuilder errors = checkNanAndInfinite(vector);
+                if (errors != null) {
+                    return errors;
+                }
 
                 for (int index = 0; index < vector.length; ++index) {
                     float value = vector[index];
 
                     if (value % 1.0f != 0.0f) {
-                        errorBuilder = new StringBuilder(
+                        errors = new StringBuilder(
                             "element_type ["
                                 + this
                                 + "] vectors only support non-decimal values but found decimal value ["
@@ -790,7 +790,7 @@ public void checkVectorBounds(float[] vector) {
                     }
 
                     if (value < Byte.MIN_VALUE || value > Byte.MAX_VALUE) {
-                        errorBuilder = new StringBuilder(
+                        errors = new StringBuilder(
                             "element_type ["
                                 + this
                                 + "] vectors only support integers between ["
@@ -807,9 +807,7 @@ public void checkVectorBounds(float[] vector) {
                     }
                 }
 
-                if (errorBuilder != null) {
-                    throw new IllegalArgumentException(appendErrorElements(errorBuilder, vector).toString());
-                }
+                return errors;
             }
 
             @Override
@@ -993,7 +991,37 @@ public abstract VectorData parseKnnVector(
 
         public abstract ByteBuffer createByteBuffer(IndexVersion indexVersion, int numBytes);
 
-        public abstract void checkVectorBounds(float[] vector);
+        public static ElementType checkValidVector(float[] vector, ElementType... possibleTypes) {
+            assert possibleTypes.length != 0;
+            // we're looking for one valid allowed type
+            StringBuilder[] errors = new StringBuilder[possibleTypes.length];
+            for (int i = 0; i < possibleTypes.length; i++) {
+                StringBuilder error = possibleTypes[i].checkVectorErrors(vector);
+                if (error == null) {
+                    // this one works - all ok
+                    return possibleTypes[i];
+                } else {
+                    errors[i] = error;
+                }
+            }
+
+            // oh dear, none of the types work with this vector. Generate the error message and throw.
+            StringBuilder message = new StringBuilder();
+            for (int i = 0; i < possibleTypes.length; i++) {
+                if (i > 0) message.append(" ");
+                message.append("Vector is not a ").append(possibleTypes[i]).append(" vector: ").append(errors[i]);
+            }
+            throw new IllegalArgumentException(appendErrorElements(message, vector).toString());
+        }
+
+        public void checkVectorBounds(float[] vector) {
+            StringBuilder errors = checkVectorErrors(vector);
+            if (errors != null) {
+                throw new IllegalArgumentException(appendErrorElements(errors, vector).toString());
+            }
+        }
+
+        abstract StringBuilder checkVectorErrors(float[] vector);
 
         abstract void checkVectorMagnitude(
             VectorSimilarity similarity,
@@ -1017,7 +1045,7 @@ public int parseDimensionCount(DocumentParserContext context) throws IOException
             return index;
         }
 
-        void checkNanAndInfinite(float[] vector) {
+        StringBuilder checkNanAndInfinite(float[] vector) {
             StringBuilder errorBuilder = null;
 
             for (int index = 0; index < vector.length; ++index) {
@@ -1044,9 +1072,7 @@ void checkNanAndInfinite(float[] vector) {
                 }
             }
 
-            if (errorBuilder != null) {
-                throw new IllegalArgumentException(appendErrorElements(errorBuilder, vector).toString());
-            }
+            return errorBuilder;
         }
 
         static StringBuilder appendErrorElements(StringBuilder errorBuilder, float[] vector) {
diff --git a/server/src/main/java/org/elasticsearch/script/VectorScoreScriptUtils.java b/server/src/main/java/org/elasticsearch/script/VectorScoreScriptUtils.java
@@ -11,6 +11,7 @@
 
 import org.elasticsearch.ExceptionsHelper;
 import org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper;
+import org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.ElementType;
 import org.elasticsearch.script.field.vectors.DenseVector;
 import org.elasticsearch.script.field.vectors.DenseVectorDocValuesField;
 
@@ -42,7 +43,10 @@ void setNextVector() {
     }
 
     public static class ByteDenseVectorFunction extends DenseVectorFunction {
-        protected final byte[] queryVector;
+        // either byteQueryVector or floatQueryVector will be non-null
+        protected final byte[] byteQueryVector;
+        protected final float[] floatQueryVector;
+        // only valid if byteQueryVector is used
         protected final float qvMagnitude;
 
         /**
@@ -52,21 +56,39 @@ public static class ByteDenseVectorFunction extends DenseVectorFunction {
          * @param field The vector field.
          * @param queryVector The query vector.
          */
-        public ByteDenseVectorFunction(ScoreScript scoreScript, DenseVectorDocValuesField field, List<Number> queryVector) {
+        public ByteDenseVectorFunction(
+            ScoreScript scoreScript,
+            DenseVectorDocValuesField field,
+            List<Number> queryVector,
+            ElementType... allowedTypes
+        ) {
             super(scoreScript, field);
             field.getElementType().checkDimensions(field.get().getDims(), queryVector.size());
-            this.queryVector = new byte[queryVector.size()];
-            float[] validateValues = new float[queryVector.size()];
-            int queryMagnitude = 0;
+            float[] floatValues = new float[queryVector.size()];
             for (int i = 0; i < queryVector.size(); i++) {
-                final Number number = queryVector.get(i);
-                byte value = number.byteValue();
-                this.queryVector[i] = value;
-                queryMagnitude += value * value;
-                validateValues[i] = number.floatValue();
+                floatValues[i] = queryVector.get(i).floatValue();
             }
-            this.qvMagnitude = (float) Math.sqrt(queryMagnitude);
-            field.getElementType().checkVectorBounds(validateValues);
+
+            switch (ElementType.checkValidVector(floatValues, allowedTypes)) {
+                case FLOAT:
+                    byteQueryVector = null;
+                    floatQueryVector = floatValues;
+                    qvMagnitude = -1;   // invalid valid, not used for float vectors
+                    break;
+                case BYTE:
+                    floatQueryVector = null;
+                    byteQueryVector = new byte[floatValues.length];
+                    float queryMagnitude = 0;
+                    for (int i = 0; i < floatValues.length; i++) {
+                        byteQueryVector[i] = (byte) floatValues[i];
+                        queryMagnitude += floatValues[i] * floatValues[i];
+                    }
+                    this.qvMagnitude = (float) Math.sqrt(queryMagnitude);
+                    break;
+                default:
+                    throw new AssertionError("Unexpected element type");
+            }
+
         }
 
         /**
@@ -78,7 +100,8 @@ public ByteDenseVectorFunction(ScoreScript scoreScript, DenseVectorDocValuesFiel
          */
         public ByteDenseVectorFunction(ScoreScript scoreScript, DenseVectorDocValuesField field, byte[] queryVector) {
             super(scoreScript, field);
-            this.queryVector = queryVector;
+            byteQueryVector = queryVector;
+            floatQueryVector = null;
             float queryMagnitude = 0.0f;
             for (byte value : queryVector) {
                 queryMagnitude += value * value;
@@ -115,7 +138,7 @@ public FloatDenseVectorFunction(
                 queryMagnitude += value * value;
             }
             queryMagnitude = Math.sqrt(queryMagnitude);
-            field.getElementType().checkVectorBounds(this.queryVector);
+            DenseVectorFieldMapper.ElementType.FLOAT.checkVectorBounds(this.queryVector);
 
             if (normalizeQuery) {
                 for (int dim = 0; dim < this.queryVector.length; dim++) {
@@ -133,7 +156,7 @@ public interface L1NormInterface {
     public static class ByteL1Norm extends ByteDenseVectorFunction implements L1NormInterface {
 
         public ByteL1Norm(ScoreScript scoreScript, DenseVectorDocValuesField field, List<Number> queryVector) {
-            super(scoreScript, field, queryVector);
+            super(scoreScript, field, queryVector, ElementType.BYTE);
         }
 
         public ByteL1Norm(ScoreScript scoreScript, DenseVectorDocValuesField field, byte[] queryVector) {
@@ -142,7 +165,7 @@ public ByteL1Norm(ScoreScript scoreScript, DenseVectorDocValuesField field, byte
 
         public double l1norm() {
             setNextVector();
-            return field.get().l1Norm(queryVector);
+            return field.get().l1Norm(byteQueryVector);
         }
     }
 
@@ -197,7 +220,7 @@ public interface HammingDistanceInterface {
     public static class ByteHammingDistance extends ByteDenseVectorFunction implements HammingDistanceInterface {
 
         public ByteHammingDistance(ScoreScript scoreScript, DenseVectorDocValuesField field, List<Number> queryVector) {
-            super(scoreScript, field, queryVector);
+            super(scoreScript, field, queryVector, ElementType.BYTE);
         }
 
         public ByteHammingDistance(ScoreScript scoreScript, DenseVectorDocValuesField field, byte[] queryVector) {
@@ -206,7 +229,7 @@ public ByteHammingDistance(ScoreScript scoreScript, DenseVectorDocValuesField fi
 
         public int hamming() {
             setNextVector();
-            return field.get().hamming(queryVector);
+            return field.get().hamming(byteQueryVector);
         }
     }
 
@@ -243,7 +266,7 @@ public interface L2NormInterface {
     public static class ByteL2Norm extends ByteDenseVectorFunction implements L2NormInterface {
 
         public ByteL2Norm(ScoreScript scoreScript, DenseVectorDocValuesField field, List<Number> queryVector) {
-            super(scoreScript, field, queryVector);
+            super(scoreScript, field, queryVector, ElementType.BYTE);
         }
 
         public ByteL2Norm(ScoreScript scoreScript, DenseVectorDocValuesField field, byte[] queryVector) {
@@ -252,7 +275,7 @@ public ByteL2Norm(ScoreScript scoreScript, DenseVectorDocValuesField field, byte
 
         public double l2norm() {
             setNextVector();
-            return field.get().l2Norm(queryVector);
+            return field.get().l2Norm(byteQueryVector);
         }
     }
 
@@ -388,7 +411,7 @@ public double dotProduct() {
     public static class ByteDotProduct extends ByteDenseVectorFunction implements DotProductInterface {
 
         public ByteDotProduct(ScoreScript scoreScript, DenseVectorDocValuesField field, List<Number> queryVector) {
-            super(scoreScript, field, queryVector);
+            super(scoreScript, field, queryVector, ElementType.BYTE, ElementType.FLOAT);
         }
 
         public ByteDotProduct(ScoreScript scoreScript, DenseVectorDocValuesField field, byte[] queryVector) {
@@ -397,7 +420,11 @@ public ByteDotProduct(ScoreScript scoreScript, DenseVectorDocValuesField field,
 
         public double dotProduct() {
             setNextVector();
-            return field.get().dotProduct(queryVector);
+            if (floatQueryVector != null) {
+                return field.get().dotProduct(floatQueryVector);
+            } else {
+                return field.get().dotProduct(byteQueryVector);
+            }
         }
     }
 
@@ -461,7 +488,7 @@ public interface CosineSimilarityInterface {
     public static class ByteCosineSimilarity extends ByteDenseVectorFunction implements CosineSimilarityInterface {
 
         public ByteCosineSimilarity(ScoreScript scoreScript, DenseVectorDocValuesField field, List<Number> queryVector) {
-            super(scoreScript, field, queryVector);
+            super(scoreScript, field, queryVector, ElementType.BYTE, ElementType.FLOAT);
         }
 
         public ByteCosineSimilarity(ScoreScript scoreScript, DenseVectorDocValuesField field, byte[] queryVector) {
@@ -470,7 +497,11 @@ public ByteCosineSimilarity(ScoreScript scoreScript, DenseVectorDocValuesField f
 
         public double cosineSimilarity() {
             setNextVector();
-            return field.get().cosineSimilarity(queryVector, qvMagnitude);
+            if (floatQueryVector != null) {
+                return field.get().cosineSimilarity(floatQueryVector);
+            } else {
+                return field.get().cosineSimilarity(byteQueryVector, qvMagnitude);
+            }
         }
     }
 
