Adds an advanced binary quantization format copied from Lucene

server/src/main/java/org/elasticsearch/index/codec/vectors/BQSpaceUtils.java

@@ -0,0 +1,78 @@
+/*
+ * @notice
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ * Modifications copyright (C) 2024 Elasticsearch B.V.
+ */
+package org.elasticsearch.index.codec.vectors;
+
+/** Utility class for quantization calculations */
+public class BQSpaceUtils {
+
+    public static final short B_QUERY = 4;
+    // the first four bits masked
+    private static final int B_QUERY_MASK = 15;
+
+    /**
+     * Copied from Lucene, replace with Lucene's implementation sometime after Lucene 10
+     * @param q the query vector, assumed to be half-byte quantized with values between 0 and 15
+     * @param dimensions the number of dimensions in the query vector
+     * @param quantQueryByte the byte array to store the transposed query vector
+     */
+    public static void transposeBin(byte[] q, int dimensions, byte[] quantQueryByte) {
+        // TODO: rewrite this in Panama Vector API
+        int qOffset = 0;
+        final byte[] v1 = new byte[4];
+        final byte[] v = new byte[32];
+        for (int i = 0; i < dimensions; i += 32) {
+            // for every four bytes we shift left (with remainder across those bytes)
+            for (int j = 0; j < v.length; j += 4) {
+                v[j] = (byte) (q[qOffset + j] << B_QUERY | ((q[qOffset + j] >>> B_QUERY) & B_QUERY_MASK));
+                v[j + 1] = (byte) (q[qOffset + j + 1] << B_QUERY | ((q[qOffset + j + 1] >>> B_QUERY) & B_QUERY_MASK));
+                v[j + 2] = (byte) (q[qOffset + j + 2] << B_QUERY | ((q[qOffset + j + 2] >>> B_QUERY) & B_QUERY_MASK));
+                v[j + 3] = (byte) (q[qOffset + j + 3] << B_QUERY | ((q[qOffset + j + 3] >>> B_QUERY) & B_QUERY_MASK));
+            }
+            for (int j = 0; j < B_QUERY; j++) {
+                moveMaskEpi8Byte(v, v1);
+                for (int k = 0; k < 4; k++) {
+                    quantQueryByte[(B_QUERY - j - 1) * (dimensions / 8) + i / 8 + k] = v1[k];
+                    v1[k] = 0;
+                }
+                for (int k = 0; k < v.length; k += 4) {
+                    v[k] = (byte) (v[k] + v[k]);
+                    v[k + 1] = (byte) (v[k + 1] + v[k + 1]);
+                    v[k + 2] = (byte) (v[k + 2] + v[k + 2]);
+                    v[k + 3] = (byte) (v[k + 3] + v[k + 3]);
+                }
+            }
+            qOffset += 32;
+        }
+    }
+
+    private static void moveMaskEpi8Byte(byte[] v, byte[] v1b) {
+        int m = 0;
+        for (int k = 0; k < v.length; k++) {
+            if ((v[k] & 0b10000000) == 0b10000000) {
+                v1b[m] |= 0b00000001;
+            }
+            if (k % 8 == 7) {
+                m++;
+            } else {
+                v1b[m] <<= 1;
+            }
+        }
+    }
+}

server/src/main/java/org/elasticsearch/index/codec/vectors/BQVectorUtils.java

@@ -0,0 +1,97 @@
+/*
+ * @notice
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ * Modifications copyright (C) 2024 Elasticsearch B.V.
+ */
+package org.elasticsearch.index.codec.vectors;
+
+import org.apache.lucene.util.ArrayUtil;
+import org.apache.lucene.util.BitUtil;
+import org.apache.lucene.util.VectorUtil;
+
+/** Utility class for vector quantization calculations */
+public class BQVectorUtils {
+    private static final float EPSILON = 1e-4f;
+
+    public static boolean isUnitVector(float[] v) {
+        double l1norm = VectorUtil.dotProduct(v, v);
+        return Math.abs(l1norm - 1.0d) <= EPSILON;
+    }
+
+    public static int discretize(int value, int bucket) {
+        return ((value + (bucket - 1)) / bucket) * bucket;
+    }
+
+    public static float[] pad(float[] vector, int dimensions) {
+        if (vector.length >= dimensions) {
+            return vector;
+        }
+        return ArrayUtil.growExact(vector, dimensions);
+    }
+
+    public static byte[] pad(byte[] vector, int dimensions) {
+        if (vector.length >= dimensions) {
+            return vector;
+        }
+        return ArrayUtil.growExact(vector, dimensions);
+    }
+
+    /**
+     * Copied from Lucene, replace with Lucene's implementation sometime after Lucene 10
+     * @param d the byte array to count the number of set bits in
+     * @return count of flipped bits in the byte array
+     */
+    public static int popcount(byte[] d) {
+        int r = 0;
+        int cnt = 0;
+        for (final int upperBound = d.length & -Integer.BYTES; r < upperBound; r += Integer.BYTES) {
+            cnt += Integer.bitCount((int) BitUtil.VH_NATIVE_INT.get(d, r));
+        }
+        for (; r < d.length; r++) {
+            cnt += Integer.bitCount(d[r] & 0xFF);
+        }
+        return cnt;
+    }
+
+    // TODO: move to VectorUtil & vectorize?
+    public static void divideInPlace(float[] a, float b) {
+        for (int j = 0; j < a.length; j++) {
+            a[j] /= b;
+        }
+    }
+
+    public static float[] subtract(float[] a, float[] b) {
+        float[] result = new float[a.length];
+        subtract(a, b, result);
+        return result;
+    }
+
+    public static void subtractInPlace(float[] target, float[] other) {
+        subtract(target, other, target);
+    }
+
+    private static void subtract(float[] a, float[] b, float[] result) {
+        for (int j = 0; j < a.length; j++) {
+            result[j] = a[j] - b[j];
+        }
+    }
+
+    public static float norm(float[] vector) {
+        float magnitude = VectorUtil.dotProduct(vector, vector);
+        return (float) Math.sqrt(magnitude);
+    }
+}

server/src/main/java/org/elasticsearch/index/codec/vectors/BinarizedByteVectorValues.java

@@ -0,0 +1,58 @@
+/*
+ * @notice
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ * Modifications copyright (C) 2024 Elasticsearch B.V.
+ */
+package org.elasticsearch.index.codec.vectors;
+
+import org.apache.lucene.search.DocIdSetIterator;
+import org.apache.lucene.search.VectorScorer;
+
+import java.io.IOException;
+
+/**
+ * Copied from Lucene, replace with Lucene's implementation sometime after Lucene 10
+ */
+public abstract class BinarizedByteVectorValues extends DocIdSetIterator {
+
+    public abstract float[] getCorrectiveTerms();
+
+    public abstract byte[] vectorValue() throws IOException;
+
+    /** Return the dimension of the vectors */
+    public abstract int dimension();
+
+    /**
+     * Return the number of vectors for this field.
+     *
+     * @return the number of vectors returned by this iterator
+     */
+    public abstract int size();
+
+    @Override
+    public final long cost() {
+        return size();
+    }
+
+    /**
+     * Return a {@link VectorScorer} for the given query vector.
+     *
+     * @param query the query vector
+     * @return a {@link VectorScorer} instance or null
+     */
+    public abstract VectorScorer scorer(float[] query) throws IOException;
+}