Add basic implementations of float-byte script comparisons

thecoop · thecoop · commit 9f42da9a6b85 · 2025-02-12T13:54:36.000Z
diff --git a/libs/simdvec/src/main/java/org/elasticsearch/simdvec/ESVectorUtil.java b/libs/simdvec/src/main/java/org/elasticsearch/simdvec/ESVectorUtil.java
@@ -80,6 +80,19 @@ public static float ipFloatBit(float[] q, byte[] d) {
         return IMPL.ipFloatBit(q, d);
     }
 
+    /**
+     * Compute the inner product of two vectors, where the query vector is a float vector and the document vector is a byte vector.
+     * @param q the query vector
+     * @param d the document vector
+     * @return the inner product of the two vectors
+     */
+    public static float ipFloatByte(float[] q, byte[] d) {
+        if (q.length != d.length) {
+            throw new IllegalArgumentException("vector dimensions incompatible: " + q.length + "!= " + d.length);
+        }
+        return IMPL.ipFloatByte(q, d);
+    }
+
     /**
      * AND bit count computed over signed bytes.
      * Copied from Lucene's XOR implementation
diff --git a/libs/simdvec/src/main/java/org/elasticsearch/simdvec/internal/vectorization/DefaultESVectorUtilSupport.java b/libs/simdvec/src/main/java/org/elasticsearch/simdvec/internal/vectorization/DefaultESVectorUtilSupport.java
@@ -39,6 +39,11 @@ public float ipFloatBit(float[] q, byte[] d) {
         return ipFloatBitImpl(q, d);
     }
 
+    @Override
+    public float ipFloatByte(float[] q, byte[] d) {
+        return ipFloatByteImpl(q, d);
+    }
+
     public static int ipByteBitImpl(byte[] q, byte[] d) {
         assert q.length == d.length * Byte.SIZE;
         int acc0 = 0;
@@ -101,4 +106,12 @@ public static long ipByteBinByteImpl(byte[] q, byte[] d) {
         }
         return ret;
     }
+
+    public static float ipFloatByteImpl(float[] q, byte[] d) {
+        float ret = 0;
+        for (int i = 0; i < q.length; i++) {
+            ret += q[i] * d[i];
+        }
+        return ret;
+    }
 }
diff --git a/libs/simdvec/src/main/java/org/elasticsearch/simdvec/internal/vectorization/ESVectorUtilSupport.java b/libs/simdvec/src/main/java/org/elasticsearch/simdvec/internal/vectorization/ESVectorUtilSupport.java
@@ -18,4 +18,6 @@ public interface ESVectorUtilSupport {
     int ipByteBit(byte[] q, byte[] d);
 
     float ipFloatBit(float[] q, byte[] d);
+
+    float ipFloatByte(float[] q, byte[] d);
 }
diff --git a/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/PanamaESVectorUtilSupport.java b/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/PanamaESVectorUtilSupport.java
@@ -58,6 +58,11 @@ public float ipFloatBit(float[] q, byte[] d) {
         return DefaultESVectorUtilSupport.ipFloatBitImpl(q, d);
     }
 
+    @Override
+    public float ipFloatByte(float[] q, byte[] d) {
+        return DefaultESVectorUtilSupport.ipFloatByteImpl(q, d);
+    }
+
     private static final VectorSpecies<Byte> BYTE_SPECIES_128 = ByteVector.SPECIES_128;
     private static final VectorSpecies<Byte> BYTE_SPECIES_256 = ByteVector.SPECIES_256;
 
diff --git a/libs/simdvec/src/test/java/org/elasticsearch/simdvec/ESVectorUtilTests.java b/libs/simdvec/src/test/java/org/elasticsearch/simdvec/ESVectorUtilTests.java
@@ -32,11 +32,28 @@ public void testIpByteBit() {
     public void testIpFloatBit() {
         float[] q = new float[16];
         byte[] d = new byte[] { (byte) Integer.parseInt("01100010", 2), (byte) Integer.parseInt("10100111", 2) };
-        random().nextFloat();
+        for (int i = 0; i < q.length; i++) {
+            q[i] = random().nextFloat();
+        }
         float expected = q[1] + q[2] + q[6] + q[8] + q[10] + q[13] + q[14] + q[15];
         assertEquals(expected, ESVectorUtil.ipFloatBit(q, d), 1e-6);
     }
 
+    public void testIpFloatByte() {
+        float[] q = new float[16];
+        byte[] d = new byte[16];
+        for (int i = 0; i < q.length; i++) {
+            q[i] = random().nextFloat();
+        }
+        random().nextBytes(d);
+
+        float expected = 0;
+        for (int i = 0; i < q.length; i++) {
+            expected += q[i] * d[i];
+        }
+        assertEquals(expected, ESVectorUtil.ipFloatByte(q, d), 1e-6);
+    }
+
     public void testBitAndCount() {
         testBasicBitAndImpl(ESVectorUtil::andBitCountLong);
     }
diff --git a/server/src/main/java/org/elasticsearch/script/field/vectors/ByteBinaryDenseVector.java b/server/src/main/java/org/elasticsearch/script/field/vectors/ByteBinaryDenseVector.java
@@ -12,6 +12,7 @@
 import org.apache.lucene.util.BytesRef;
 import org.apache.lucene.util.VectorUtil;
 import org.elasticsearch.core.SuppressForbidden;
+import org.elasticsearch.simdvec.ESVectorUtil;
 
 import java.nio.ByteBuffer;
 import java.util.List;
@@ -61,7 +62,7 @@ public int dotProduct(byte[] queryVector) {
 
     @Override
     public double dotProduct(float[] queryVector) {
-        throw new UnsupportedOperationException("use [int dotProduct(byte[] queryVector)] instead");
+        return ESVectorUtil.ipFloatByte(queryVector, vectorValue);
     }
 
     @Override
@@ -142,7 +143,11 @@ public double cosineSimilarity(byte[] queryVector, float qvMagnitude) {
 
     @Override
     public double cosineSimilarity(float[] queryVector, boolean normalizeQueryVector) {
-        throw new UnsupportedOperationException("use [double cosineSimilarity(byte[] queryVector, float qvMagnitude)] instead");
+        if (normalizeQueryVector) {
+            return dotProduct(queryVector) / (DenseVector.getMagnitude(queryVector) * getMagnitude());
+        }
+
+        return dotProduct(queryVector) / getMagnitude();
     }
 
     @Override
diff --git a/server/src/main/java/org/elasticsearch/script/field/vectors/ByteKnnDenseVector.java b/server/src/main/java/org/elasticsearch/script/field/vectors/ByteKnnDenseVector.java
@@ -11,6 +11,7 @@
 
 import org.apache.lucene.util.VectorUtil;
 import org.elasticsearch.core.SuppressForbidden;
+import org.elasticsearch.simdvec.ESVectorUtil;
 
 import java.util.List;
 
@@ -51,12 +52,12 @@ public float getMagnitude() {
 
     @Override
     public int dotProduct(byte[] queryVector) {
-        return VectorUtil.dotProduct(docVector, queryVector);
+        return VectorUtil.dotProduct(queryVector, docVector);
     }
 
     @Override
     public double dotProduct(float[] queryVector) {
-        throw new UnsupportedOperationException("use [int dotProduct(byte[] queryVector)] instead");
+        return ESVectorUtil.ipFloatByte(queryVector, docVector);
     }
 
     @Override
@@ -145,7 +146,11 @@ public double cosineSimilarity(byte[] queryVector, float qvMagnitude) {
 
     @Override
     public double cosineSimilarity(float[] queryVector, boolean normalizeQueryVector) {
-        throw new UnsupportedOperationException("use [double cosineSimilarity(byte[] queryVector, float qvMagnitude)] instead");
+        if (normalizeQueryVector) {
+            return dotProduct(queryVector) / (DenseVector.getMagnitude(queryVector) * getMagnitude());
+        }
+
+        return dotProduct(queryVector) / getMagnitude();
     }
 
     @Override
diff --git a/server/src/test/java/org/elasticsearch/script/VectorScoreScriptUtilsTests.java b/server/src/test/java/org/elasticsearch/script/VectorScoreScriptUtilsTests.java
@@ -183,11 +183,12 @@ public void testByteVectorClassBindings() throws IOException {
             for (int i = 0; i < queryVectorArray.length; i++) {
                 queryVectorArray[i] = queryVector.get(i).floatValue();
             }
-            UnsupportedOperationException uoe = expectThrows(
-                UnsupportedOperationException.class,
-                () -> field.getInternal().cosineSimilarity(queryVectorArray, true)
+            assertEquals(
+                "cosineSimilarity result is not equal to the expected value!",
+                cosineSimilarityExpected,
+                field.getInternal().cosineSimilarity(queryVectorArray, true),
+                0.001
             );
-            assertThat(uoe.getMessage(), containsString("use [double cosineSimilarity(byte[] queryVector, float qvMagnitude)] instead"));
 
             // Check each function rejects query vectors with the wrong dimension
             IllegalArgumentException e = expectThrows(
@@ -342,11 +343,7 @@ public void testByteVsFloatSimilarity() throws IOException {
             switch (field.getElementType()) {
                 case BYTE -> {
                     assertEquals(field.getName(), dotProductExpected, field.get().dotProduct(byteVector));
-                    UnsupportedOperationException e = expectThrows(
-                        UnsupportedOperationException.class,
-                        () -> field.get().dotProduct(floatVector)
-                    );
-                    assertThat(e.getMessage(), containsString("use [int dotProduct(byte[] queryVector)] instead"));
+                    assertEquals(field.getName(), dotProductExpected, field.get().dotProduct(floatVector), 0.001);
                 }
                 case FLOAT -> {
                     assertEquals(field.getName(), dotProductExpected, field.get().dotProduct(floatVector), 0.001);
@@ -423,14 +420,7 @@ public void testByteVsFloatSimilarity() throws IOException {
             switch (field.getElementType()) {
                 case BYTE -> {
                     assertEquals(field.getName(), cosineSimilarityExpected, field.get().cosineSimilarity(byteVector), 0.001);
-                    UnsupportedOperationException e = expectThrows(
-                        UnsupportedOperationException.class,
-                        () -> field.get().cosineSimilarity(floatVector)
-                    );
-                    assertThat(
-                        e.getMessage(),
-                        containsString("use [double cosineSimilarity(byte[] queryVector, float qvMagnitude)] instead")
-                    );
+                    assertEquals(field.getName(), cosineSimilarityExpected, field.get().cosineSimilarity(floatVector), 0.001);
                 }
                 case FLOAT -> {
                     assertEquals(field.getName(), cosineSimilarityExpected, field.get().cosineSimilarity(floatVector), 0.001);

Original file line number	Diff line number	Diff line change
`@@ -39,6 +39,11 @@ public float ipFloatBit(float[] q, byte[] d) {`
`39`	`39`	`return ipFloatBitImpl(q, d);`
`40`	`40`	`}`
`41`	`41`
	`42`	`+ @Override`
	`43`	`+ public float ipFloatByte(float[] q, byte[] d) {`
	`44`	`+ return ipFloatByteImpl(q, d);`
	`45`	`+ }`
	`46`	`+`
`42`	`47`	`public static int ipByteBitImpl(byte[] q, byte[] d) {`
`43`	`48`	`assert q.length == d.length * Byte.SIZE;`
`44`	`49`	`int acc0 = 0;`
`@@ -101,4 +106,12 @@ public static long ipByteBinByteImpl(byte[] q, byte[] d) {`
`101`	`106`	`}`
`102`	`107`	`return ret;`
`103`	`108`	`}`
	`109`	`+`
	`110`	`+ public static float ipFloatByteImpl(float[] q, byte[] d) {`
	`111`	`+ float ret = 0;`
	`112`	`+ for (int i = 0; i < q.length; i++) {`
	`113`	`+ ret += q[i] * d[i];`
	`114`	`+ }`
	`115`	`+ return ret;`
	`116`	`+ }`
`104`	`117`	`}`
Original file line number	Diff line number	Diff line change
`@@ -18,4 +18,6 @@ public interface ESVectorUtilSupport {`
`18`	`18`	`int ipByteBit(byte[] q, byte[] d);`
`19`	`19`
`20`	`20`	`float ipFloatBit(float[] q, byte[] d);`
	`21`	`+`
	`22`	`+ float ipFloatByte(float[] q, byte[] d);`
`21`	`23`	`}`