Merge branch 'main' into unmute-testSecurityNativeRealm

Author: ldematte

distribution/docker/src/docker/dockerfiles/cloud_ess_fips/Dockerfile

@@ -25,7 +25,7 @@
 # Extract Elasticsearch artifact
 ################################################################################
 
-FROM docker.elastic.co/wolfi/chainguard-base-fips:latest@sha256:cbdc968f357706186e2b83ea1fea92382f5d51842e867a662541b67c28f4e809 AS builder
+FROM docker.elastic.co/wolfi/chainguard-base-fips:latest@sha256:7b18967a1329291d113228149b97382406f039171c6537b1afa4067701a6ba9e AS builder
 
 # Install required packages to extract the Elasticsearch distribution
 RUN <%= retry.loop(package_manager, "export DEBIAN_FRONTEND=noninteractive && ${package_manager} update && ${package_manager} update && ${package_manager} add --no-cache curl") %>
@@ -104,7 +104,7 @@ WORKDIR /usr/share/elasticsearch/config
 # Add entrypoint
 ################################################################################
 
-FROM docker.elastic.co/wolfi/chainguard-base-fips:latest@sha256:cbdc968f357706186e2b83ea1fea92382f5d51842e867a662541b67c28f4e809
+FROM docker.elastic.co/wolfi/chainguard-base-fips:latest@sha256:7b18967a1329291d113228149b97382406f039171c6537b1afa4067701a6ba9e
 
 RUN <%= retry.loop(package_manager,
           "export DEBIAN_FRONTEND=noninteractive && \n" +

distribution/docker/src/docker/dockerfiles/wolfi/Dockerfile

@@ -25,7 +25,7 @@
 # Extract Elasticsearch artifact
 ################################################################################
 
-FROM docker.elastic.co/wolfi/chainguard-base:latest@sha256:b85d54c9019ff80e88aab7d357ede3341d1442ce190173eed42ae5a116753e4e AS builder
+FROM docker.elastic.co/wolfi/chainguard-base:latest@sha256:46fd486e4716d592aa1ad7dad15bb3e6d0aa5b08b2070e9daec9f6b8c8c990a6 AS builder
 
 # Install required packages to extract the Elasticsearch distribution
 RUN <%= retry.loop(package_manager, "export DEBIAN_FRONTEND=noninteractive && ${package_manager} update && ${package_manager} update && ${package_manager} add --no-cache curl") %>
@@ -80,7 +80,7 @@ RUN sed -i -e 's/ES_DISTRIBUTION_TYPE=tar/ES_DISTRIBUTION_TYPE=docker/' bin/elas
 # Add entrypoint
 ################################################################################
 
-FROM docker.elastic.co/wolfi/chainguard-base:latest@sha256:b85d54c9019ff80e88aab7d357ede3341d1442ce190173eed42ae5a116753e4e
+FROM docker.elastic.co/wolfi/chainguard-base:latest@sha256:46fd486e4716d592aa1ad7dad15bb3e6d0aa5b08b2070e9daec9f6b8c8c990a6
 
 RUN <%= retry.loop(package_manager,
           "export DEBIAN_FRONTEND=noninteractive && \n" +

docs/changelog/136988.yaml

@@ -0,0 +1,5 @@
+pr: 136988
+summary: Enable score function in release builds
+area: ES|QL
+type: feature
+issues: []

docs/reference/query-languages/esql/_snippets/functions/layout/score.md

@@ -53,6 +53,9 @@ for information on the limitations of full text search.
 :::{include} ../_snippets/functions/layout/qstr.md
 :::
 
+:::{include} ../_snippets/functions/layout/score.md
+:::
+
 % TERM is currently a hidden feature
 % To make it visible again, uncomment this and the line in
 lists/search-functions.md

docs/reference/query-languages/esql/functions-operators/search-functions.md

@@ -48,8 +48,8 @@ public class ESNextOSQVectorsScorer {
 
     /** Sole constructor, called by sub-classes. */
     public ESNextOSQVectorsScorer(IndexInput in, byte queryBits, byte indexBits, int dimensions, int dataLength) {
-        if (queryBits != 4 || indexBits != 1) {
-            throw new IllegalArgumentException("Only asymmetric 4-bit query and 1-bit index supported");
+        if (queryBits != 4 || (indexBits != 1 && indexBits != 2)) {
+            throw new IllegalArgumentException("Only asymmetric 4-bit query and 1 or 2-bit index supported");
         }
         this.in = in;
         this.queryBits = queryBits;
@@ -65,15 +65,27 @@ public ESNextOSQVectorsScorer(IndexInput in, byte queryBits, byte indexBits, int
     public long quantizeScore(byte[] q) throws IOException {
         if (indexBits == 1) {
             if (queryBits == 4) {
-                return quantized4BitScore(q);
+                return quantized4BitScore(q, length);
             }
             throw new IllegalArgumentException("Only asymmetric 4-bit query supported");
         }
+        if (indexBits == 2) {
+            if (queryBits == 4) {
+                return quantized4BitScore2BitIndex(q);
+            }
+        }
         throw new IllegalArgumentException("Only 1-bit index supported");
+    }
 
+    private long quantized4BitScore2BitIndex(byte[] q) throws IOException {
+        assert q.length == length * 2;
+        assert length % 2 == 0 : "length must be even for 2-bit index length: " + length + " dimensions: " + dimensions;
+        int lower = (int) quantized4BitScore(q, length / 2);
+        int upper = (int) quantized4BitScore(q, length / 2);
+        return lower + ((long) upper << 1);
     }
 
-    private long quantized4BitScore(byte[] q) throws IOException {
+    private long quantized4BitScore(byte[] q, int length) throws IOException {
         assert q.length == length * 4;
         final int size = length;
         long subRet0 = 0;
@@ -120,6 +132,15 @@ public void quantizeScoreBulk(byte[] q, int count, float[] scores) throws IOExce
             }
             throw new IllegalArgumentException("Only asymmetric 4-bit query supported");
         }
+        if (indexBits == 2) {
+            if (queryBits == 4) {
+                for (int i = 0; i < count; i++) {
+                    scores[i] = quantizeScore(q);
+                }
+                return;
+            }
+            throw new IllegalArgumentException("Only asymmetric 4-bit query supported");
+        }
     }
 
     /**
@@ -140,9 +161,9 @@ public float score(
     ) {
         float ax = lowerInterval;
         // Here we assume `lx` is simply bit vectors, so the scaling isn't necessary
-        float lx = (upperInterval - ax) * BIT_SCALES[indexBits];
+        float lx = (upperInterval - ax) * BIT_SCALES[indexBits - 1];
         float ay = queryLowerInterval;
-        float ly = (queryUpperInterval - ay) * BIT_SCALES[queryBits];
+        float ly = (queryUpperInterval - ay) * BIT_SCALES[queryBits - 1];
         float y1 = queryComponentSum;
         float score = ax * ay * dimensions + ay * lx * (float) targetComponentSum + ax * ly * y1 + lx * ly * qcDist;
         // For euclidean, we need to invert the score and apply the additional correction, which is

docs/reference/query-languages/esql/kibana/definition/functions/score.json

@@ -393,6 +393,13 @@ public static void packAsBinary(int[] vector, byte[] packed) {
         IMPL.packAsBinary(vector, packed);
     }
 
+    public static void packDibit(int[] vector, byte[] packed) {
+        if (packed.length * Byte.SIZE / 2 < vector.length) {
+            throw new IllegalArgumentException("packed array is too small: " + packed.length * Byte.SIZE / 2 + " < " + vector.length);
+        }
+        IMPL.packDibit(vector, packed);
+    }
+
     /**
      * The idea here is to organize the query vector bits such that the first bit
      * of every dimension is in the first set dimensions bits, or (dimensions/8) bytes. The second,

libs/simdvec/src/main/java/org/elasticsearch/simdvec/ESNextOSQVectorsScorer.java

@@ -321,11 +321,58 @@ public void soarDistanceBulk(
         distances[3] = soarDistance(v1, c3, originalResidual, soarLambda, rnorm);
     }
 
+    @Override
+    public void packDibit(int[] vector, byte[] packed) {
+        packDibitImpl(vector, packed);
+    }
+
     @Override
     public void packAsBinary(int[] vector, byte[] packed) {
         packAsBinaryImpl(vector, packed);
     }
 
+    /**
+     * Packs two bit vector (values 0-3) into a byte array with lower bits first.
+     * The striding is similar to transposeHalfByte
+     *
+     * @param vector the input vector with values 0-3
+     * @param packed the output packed byte array
+     */
+    public static void packDibitImpl(int[] vector, byte[] packed) {
+        int limit = vector.length - 7;
+        int i = 0;
+        int index = 0;
+        for (; i < limit; i += 8, index++) {
+            assert vector[i] >= 0 && vector[i] <= 3;
+            assert vector[i + 1] >= 0 && vector[i + 1] <= 3;
+            assert vector[i + 2] >= 0 && vector[i + 2] <= 3;
+            assert vector[i + 3] >= 0 && vector[i + 3] <= 3;
+            assert vector[i + 4] >= 0 && vector[i + 4] <= 3;
+            assert vector[i + 5] >= 0 && vector[i + 5] <= 3;
+            assert vector[i + 6] >= 0 && vector[i + 6] <= 3;
+            assert vector[i + 7] >= 0 && vector[i + 7] <= 3;
+            int lowerByte = (vector[i] & 1) << 7 | (vector[i + 1] & 1) << 6 | (vector[i + 2] & 1) << 5 | (vector[i + 3] & 1) << 4
+                | (vector[i + 4] & 1) << 3 | (vector[i + 5] & 1) << 2 | (vector[i + 6] & 1) << 1 | (vector[i + 7] & 1);
+            int upperByte = ((vector[i] >> 1) & 1) << 7 | ((vector[i + 1] >> 1) & 1) << 6 | ((vector[i + 2] >> 1) & 1) << 5 | ((vector[i
+                + 3] >> 1) & 1) << 4 | ((vector[i + 4] >> 1) & 1) << 3 | ((vector[i + 5] >> 1) & 1) << 2 | ((vector[i + 6] >> 1) & 1) << 1
+                | ((vector[i + 7] >> 1) & 1);
+            packed[index] = (byte) lowerByte;
+            packed[index + packed.length / 2] = (byte) upperByte;
+        }
+        if (i == vector.length) {
+            return;
+        }
+        int lowerByte = 0;
+        int upperByte = 0;
+        for (int j = 7; i < vector.length; j--, i++) {
+            assert vector[i] >= 0 && vector[i] <= 3;
+            lowerByte |= (vector[i] & 1) << j;
+            upperByte |= ((vector[i] >> 1) & 1) << j;
+        }
+        packed[index] = (byte) lowerByte;
+        packed[index + packed.length / 2] = (byte) upperByte;
+    }
+
     public static void packAsBinaryImpl(int[] vector, byte[] packed) {
         int limit = vector.length - 7;
         int i = 0;

libs/simdvec/src/main/java/org/elasticsearch/simdvec/ESVectorUtil.java

@@ -66,6 +66,8 @@ void soarDistanceBulk(
 
     void packAsBinary(int[] vector, byte[] packed);
 
+    void packDibit(int[] vector, byte[] packed);
+
     void transposeHalfByte(int[] q, byte[] quantQueryByte);
 
     int indexOf(byte[] bytes, int offset, int length, byte marker);