Merge branch 'elasticsearch/main' into fork_streaming

Author: ioanatia

benchmarks/src/main/java/org/elasticsearch/benchmark/vector/VectorScorerBenchmark.java

@@ -19,6 +19,7 @@
 import org.apache.lucene.store.MMapDirectory;
 import org.apache.lucene.util.hnsw.RandomVectorScorer;
 import org.apache.lucene.util.hnsw.RandomVectorScorerSupplier;
+import org.apache.lucene.util.hnsw.UpdateableRandomVectorScorer;
 import org.apache.lucene.util.quantization.QuantizedByteVectorValues;
 import org.apache.lucene.util.quantization.ScalarQuantizer;
 import org.elasticsearch.common.logging.LogConfigurator;
@@ -76,10 +77,10 @@ public class VectorScorerBenchmark {
     float vec2Offset;
     float scoreCorrectionConstant;
 
-    RandomVectorScorer luceneDotScorer;
-    RandomVectorScorer luceneSqrScorer;
-    RandomVectorScorer nativeDotScorer;
-    RandomVectorScorer nativeSqrScorer;
+    UpdateableRandomVectorScorer luceneDotScorer;
+    UpdateableRandomVectorScorer luceneSqrScorer;
+    UpdateableRandomVectorScorer nativeDotScorer;
+    UpdateableRandomVectorScorer nativeSqrScorer;
 
     RandomVectorScorer luceneDotScorerQuery;
     RandomVectorScorer nativeDotScorerQuery;
@@ -118,12 +119,16 @@ public void setup() throws IOException {
         in = dir.openInput("vector.data", IOContext.DEFAULT);
         var values = vectorValues(dims, 2, in, VectorSimilarityFunction.DOT_PRODUCT);
         scoreCorrectionConstant = values.getScalarQuantizer().getConstantMultiplier();
-        luceneDotScorer = luceneScoreSupplier(values, VectorSimilarityFunction.DOT_PRODUCT).scorer(0);
+        luceneDotScorer = luceneScoreSupplier(values, VectorSimilarityFunction.DOT_PRODUCT).scorer();
+        luceneDotScorer.setScoringOrdinal(0);
         values = vectorValues(dims, 2, in, VectorSimilarityFunction.EUCLIDEAN);
-        luceneSqrScorer = luceneScoreSupplier(values, VectorSimilarityFunction.EUCLIDEAN).scorer(0);
+        luceneSqrScorer = luceneScoreSupplier(values, VectorSimilarityFunction.EUCLIDEAN).scorer();
+        luceneSqrScorer.setScoringOrdinal(0);
 
-        nativeDotScorer = factory.getInt7SQVectorScorerSupplier(DOT_PRODUCT, in, values, scoreCorrectionConstant).get().scorer(0);
-        nativeSqrScorer = factory.getInt7SQVectorScorerSupplier(EUCLIDEAN, in, values, scoreCorrectionConstant).get().scorer(0);
+        nativeDotScorer = factory.getInt7SQVectorScorerSupplier(DOT_PRODUCT, in, values, scoreCorrectionConstant).get().scorer();
+        nativeDotScorer.setScoringOrdinal(0);
+        nativeSqrScorer = factory.getInt7SQVectorScorerSupplier(EUCLIDEAN, in, values, scoreCorrectionConstant).get().scorer();
+        nativeSqrScorer.setScoringOrdinal(0);
 
         // setup for getInt7SQVectorScorer / query vector scoring
         float[] queryVec = new float[dims];

build-tools-internal/src/main/java/org/elasticsearch/gradle/internal/dependencies/patches/PatcherInfo.java

@@ -0,0 +1,57 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.gradle.internal.dependencies.patches;
+
+import org.objectweb.asm.ClassVisitor;
+import org.objectweb.asm.ClassWriter;
+
+import java.util.Arrays;
+import java.util.HexFormat;
+import java.util.function.Function;
+
+public final class PatcherInfo {
+    private final String jarEntryName;
+    private final byte[] classSha256;
+    private final Function<ClassWriter, ClassVisitor> visitorFactory;
+
+    private PatcherInfo(String jarEntryName, byte[] classSha256, Function<ClassWriter, ClassVisitor> visitorFactory) {
+        this.jarEntryName = jarEntryName;
+        this.classSha256 = classSha256;
+        this.visitorFactory = visitorFactory;
+    }
+
+    /**
+     * Creates a patcher info entry, linking a jar entry path name and its SHA256 digest to a patcher factory (a factory to create an ASM
+     * visitor)
+     *
+     * @param jarEntryName   the jar entry path, as a string
+     * @param classSha256    the SHA256 digest of the class bytes, as a HEX string
+     * @param visitorFactory the factory to create an ASM visitor from a ASM writer
+     */
+    public static PatcherInfo classPatcher(String jarEntryName, String classSha256, Function<ClassWriter, ClassVisitor> visitorFactory) {
+        return new PatcherInfo(jarEntryName, HexFormat.of().parseHex(classSha256), visitorFactory);
+    }
+
+    boolean matches(byte[] otherClassSha256) {
+        return Arrays.equals(this.classSha256, otherClassSha256);
+    }
+
+    public String jarEntryName() {
+        return jarEntryName;
+    }
+
+    public byte[] classSha256() {
+        return classSha256;
+    }
+
+    public ClassVisitor createVisitor(ClassWriter classWriter) {
+        return visitorFactory.apply(classWriter);
+    }
+}

build-tools-internal/src/main/java/org/elasticsearch/gradle/internal/dependencies/patches/Utils.java

@@ -0,0 +1,140 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.gradle.internal.dependencies.patches;
+
+import org.objectweb.asm.ClassReader;
+import org.objectweb.asm.ClassWriter;
+
+import java.io.File;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.security.MessageDigest;
+import java.security.NoSuchAlgorithmException;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Enumeration;
+import java.util.HexFormat;
+import java.util.Locale;
+import java.util.function.Function;
+import java.util.jar.JarEntry;
+import java.util.jar.JarFile;
+import java.util.jar.JarOutputStream;
+import java.util.stream.Collectors;
+
+import static org.objectweb.asm.ClassWriter.COMPUTE_FRAMES;
+import static org.objectweb.asm.ClassWriter.COMPUTE_MAXS;
+
+public class Utils {
+
+    private static final MessageDigest SHA_256;
+
+    static {
+        try {
+            SHA_256 = MessageDigest.getInstance("SHA-256");
+        } catch (NoSuchAlgorithmException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    private record MismatchInfo(String jarEntryName, String expectedClassSha256, String foundClassSha256) {
+        @Override
+        public String toString() {
+            return "[class='"
+                + jarEntryName
+                + '\''
+                + ", expected='"
+                + expectedClassSha256
+                + '\''
+                + ", found='"
+                + foundClassSha256
+                + '\''
+                + ']';
+        }
+    }
+
+    /**
+     * Patches the classes in the input JAR file, using the collection of patchers. Each patcher specifies a target class (its jar entry
+     * name) and the SHA256 digest on the class bytes.
+     * This digest is checked against the class bytes in the JAR, and if it does not match, an IllegalArgumentException is thrown.
+     * If the input file does not contain all the classes to patch specified in the patcher info collection, an IllegalArgumentException
+     * is also thrown.
+     * @param inputFile the JAR file to patch
+     * @param outputFile the output (patched) JAR file
+     * @param patchers list of patcher info (classes to patch (jar entry name + optional SHA256 digest) and ASM visitor to transform them)
+     */
+    public static void patchJar(File inputFile, File outputFile, Collection<PatcherInfo> patchers) {
+        var classPatchers = patchers.stream().collect(Collectors.toMap(PatcherInfo::jarEntryName, Function.identity()));
+        var mismatchedClasses = new ArrayList<MismatchInfo>();
+        try (JarFile jarFile = new JarFile(inputFile); JarOutputStream jos = new JarOutputStream(new FileOutputStream(outputFile))) {
+            Enumeration<JarEntry> entries = jarFile.entries();
+            while (entries.hasMoreElements()) {
+                JarEntry entry = entries.nextElement();
+                String entryName = entry.getName();
+                // Add the entry to the new JAR file
+                jos.putNextEntry(new JarEntry(entryName));
+
+                var classPatcher = classPatchers.remove(entryName);
+                if (classPatcher != null) {
+                    byte[] classToPatch = jarFile.getInputStream(entry).readAllBytes();
+                    var classSha256 = SHA_256.digest(classToPatch);
+
+                    if (classPatcher.matches(classSha256)) {
+                        ClassReader classReader = new ClassReader(classToPatch);
+                        ClassWriter classWriter = new ClassWriter(classReader, COMPUTE_MAXS | COMPUTE_FRAMES);
+                        classReader.accept(classPatcher.createVisitor(classWriter), 0);
+                        jos.write(classWriter.toByteArray());
+                    } else {
+                        mismatchedClasses.add(
+                            new MismatchInfo(
+                                classPatcher.jarEntryName(),
+                                HexFormat.of().formatHex(classPatcher.classSha256()),
+                                HexFormat.of().formatHex(classSha256)
+                            )
+                        );
+                    }
+                } else {
+                    // Read the entry's data and write it to the new JAR
+                    try (InputStream is = jarFile.getInputStream(entry)) {
+                        is.transferTo(jos);
+                    }
+                }
+                jos.closeEntry();
+            }
+        } catch (IOException ex) {
+            throw new RuntimeException(ex);
+        }
+
+        if (mismatchedClasses.isEmpty() == false) {
+            throw new IllegalArgumentException(
+                String.format(
+                    Locale.ROOT,
+                    """
+                        Error patching JAR [%s]: SHA256 digest mismatch (%s). This JAR was updated to a version that contains different \
+                        classes, for which this patcher was not designed. Please check if the patcher still \
+                        applies correctly, and update the SHA256 digest(s).""",
+                    inputFile.getName(),
+                    mismatchedClasses.stream().map(MismatchInfo::toString).collect(Collectors.joining())
+                )
+            );
+        }
+
+        if (classPatchers.isEmpty() == false) {
+            throw new IllegalArgumentException(
+                String.format(
+                    Locale.ROOT,
+                    "error patching [%s]: the jar does not contain [%s]",
+                    inputFile.getName(),
+                    String.join(", ", classPatchers.keySet())
+                )
+            );
+        }
+    }
+}