Merge branch 'main' into outputChecker_v2

Author: julian-elastic

benchmarks/src/main/java/org/elasticsearch/benchmark/_nightly/esql/ValuesSourceReaderBenchmark.java

@@ -92,7 +92,8 @@ public class ValuesSourceReaderBenchmark {
         "double",
         "keyword",
         "stored_keyword",
-        "3_stored_keywords" };
+        "3_stored_keywords",
+        "keyword_mv" };
 
     private static final int BLOCK_LENGTH = 16 * 1024;
     private static final int INDEX_SIZE = 10 * BLOCK_LENGTH;
@@ -332,7 +333,7 @@ public FieldNamesFieldMapper.FieldNamesFieldType fieldNames() {
     @Param({ "in_order", "shuffled" })
     public String layout;
 
-    @Param({ "long", "keyword", "stored_keyword" })
+    @Param({ "long", "keyword", "stored_keyword", "keyword_mv" })
     public String name;
 
     private Directory directory;
@@ -398,6 +399,22 @@ public void benchmark() {
                         }
                     }
                 }
+                case "keyword_mv" -> {
+                    BytesRef scratch = new BytesRef();
+                    BytesRefBlock values = op.getOutput().<BytesRefBlock>getBlock(1);
+                    for (int p = 0; p < values.getPositionCount(); p++) {
+                        int count = values.getValueCount(p);
+                        if (count > 0) {
+                            int first = values.getFirstValueIndex(p);
+                            for (int i = 0; i < count; i++) {
+                                BytesRef r = values.getBytesRef(first + i, scratch);
+                                r.offset++;
+                                r.length--;
+                                sum += Integer.parseInt(r.utf8ToString());
+                            }
+                        }
+                    }
+                }
             }
         }
         long expected = 0;
@@ -407,6 +424,16 @@ public void benchmark() {
                     expected += i % 1000;
                 }
                 break;
+            case "keyword_mv":
+                for (int i = 0; i < INDEX_SIZE; i++) {
+                    int v1 = i % 1000;
+                    expected += v1;
+                    int v2 = i % 500;
+                    if (v1 != v2) {
+                        expected += v2;
+                    }
+                }
+                break;
             case "3_stored_keywords":
                 for (int i = 0; i < INDEX_SIZE; i++) {
                     expected += 3 * (i % 1000);
@@ -461,7 +488,9 @@ private void setupIndex() throws IOException {
                         new StoredField("double", (double) i),
                         new KeywordFieldMapper.KeywordField("keyword_1", new BytesRef(c + i % 1000), keywordFieldType),
                         new KeywordFieldMapper.KeywordField("keyword_2", new BytesRef(c + i % 1000), keywordFieldType),
-                        new KeywordFieldMapper.KeywordField("keyword_3", new BytesRef(c + i % 1000), keywordFieldType)
+                        new KeywordFieldMapper.KeywordField("keyword_3", new BytesRef(c + i % 1000), keywordFieldType),
+                        new KeywordFieldMapper.KeywordField("keyword_mv", new BytesRef(c + i % 1000), keywordFieldType),
+                        new KeywordFieldMapper.KeywordField("keyword_mv", new BytesRef(c + i % 500), keywordFieldType)
                     )
                 );
                 if (i % COMMIT_INTERVAL == 0) {

benchmarks/src/main/java/org/elasticsearch/benchmark/compute/operator/AggregatorBenchmark.java

@@ -191,7 +191,7 @@ private static Operator operator(DriverContext driverContext, String grouping, S
                 new BlockHash.GroupSpec(2, ElementType.BYTES_REF)
             );
             case TOP_N_LONGS -> List.of(
-                new BlockHash.GroupSpec(0, ElementType.LONG, false, new BlockHash.TopNDef(0, true, true, TOP_N_LIMIT))
+                new BlockHash.GroupSpec(0, ElementType.LONG, null, new BlockHash.TopNDef(0, true, true, TOP_N_LIMIT))
             );
             default -> throw new IllegalArgumentException("unsupported grouping [" + grouping + "]");
         };

benchmarks/src/main/java/org/elasticsearch/benchmark/compute/operator/ValuesAggregatorBenchmark.java

@@ -95,8 +95,7 @@ static void selfTest() {
         try {
             for (String groups : ValuesAggregatorBenchmark.class.getField("groups").getAnnotationsByType(Param.class)[0].value()) {
                 for (String dataType : ValuesAggregatorBenchmark.class.getField("dataType").getAnnotationsByType(Param.class)[0].value()) {
-                    run(Integer.parseInt(groups), dataType, 10, 0);
-                    run(Integer.parseInt(groups), dataType, 10, 1);
+                    run(Integer.parseInt(groups), dataType, 10);
                 }
             }
         } catch (NoSuchFieldException e) {
@@ -114,10 +113,7 @@ static void selfTest() {
     @Param({ BYTES_REF, INT, LONG })
     public String dataType;
 
-    @Param({ "0", "1" })
-    public int numOrdinalMerges;
-
-    private static Operator operator(DriverContext driverContext, int groups, String dataType, int numOrdinalMerges) {
+    private static Operator operator(DriverContext driverContext, int groups, String dataType) {
         if (groups == 1) {
             return new AggregationOperator(
                 List.of(supplier(dataType).aggregatorFactory(AggregatorMode.SINGLE, List.of(0)).apply(driverContext)),
@@ -132,20 +128,8 @@ private static Operator operator(DriverContext driverContext, int groups, String
         ) {
             @Override
             public Page getOutput() {
-                mergeOrdinal();
                 return super.getOutput();
             }
-
-            // simulate OrdinalsGroupingOperator
-            void mergeOrdinal() {
-                var merged = supplier(dataType).groupingAggregatorFactory(AggregatorMode.SINGLE, List.of(1)).apply(driverContext);
-                for (int i = 0; i < numOrdinalMerges; i++) {
-                    for (int p = 0; p < groups; p++) {
-                        merged.addIntermediateRow(p, aggregators.getFirst(), p);
-                    }
-                }
-                aggregators.set(0, merged);
-            }
         };
     }
 
@@ -352,12 +336,12 @@ private static Block groupingBlock(int groups) {
 
     @Benchmark
     public void run() {
-        run(groups, dataType, OP_COUNT, numOrdinalMerges);
+        run(groups, dataType, OP_COUNT);
     }
 
-    private static void run(int groups, String dataType, int opCount, int numOrdinalMerges) {
+    private static void run(int groups, String dataType, int opCount) {
         DriverContext driverContext = driverContext();
-        try (Operator operator = operator(driverContext, groups, dataType, numOrdinalMerges)) {
+        try (Operator operator = operator(driverContext, groups, dataType)) {
             Page page = page(groups, dataType);
             for (int i = 0; i < opCount; i++) {
                 operator.addInput(page.shallowCopy());

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

@@ -8,12 +8,14 @@
  */
 package org.elasticsearch.benchmark.vector;
 
+import org.apache.lucene.index.VectorSimilarityFunction;
 import org.apache.lucene.store.Directory;
 import org.apache.lucene.store.IOContext;
 import org.apache.lucene.store.IndexInput;
 import org.apache.lucene.store.IndexOutput;
 import org.apache.lucene.store.MMapDirectory;
 import org.apache.lucene.util.VectorUtil;
+import org.apache.lucene.util.quantization.OptimizedScalarQuantizer;
 import org.elasticsearch.common.logging.LogConfigurator;
 import org.elasticsearch.core.IOUtils;
 import org.elasticsearch.simdvec.ES91Int4VectorsScorer;
@@ -52,20 +54,26 @@ public class Int4ScorerBenchmark {
         LogConfigurator.configureESLogging(); // native access requires logging to be initialized
     }
 
-    @Param({ "384", "702", "1024" })
+    @Param({ "384", "782", "1024" })
     int dims;
 
-    int numVectors = 200;
-    int numQueries = 10;
+    int numVectors = 20 * ES91Int4VectorsScorer.BULK_SIZE;
+    int numQueries = 5;
 
     byte[] scratch;
     byte[][] binaryVectors;
     byte[][] binaryQueries;
+    float[] scores = new float[ES91Int4VectorsScorer.BULK_SIZE];
+
+    float[] scratchFloats = new float[3];
 
     ES91Int4VectorsScorer scorer;
     Directory dir;
     IndexInput in;
 
+    OptimizedScalarQuantizer.QuantizationResult queryCorrections;
+    float centroidDp;
+
     @Setup
     public void setup() throws IOException {
         binaryVectors = new byte[numVectors][dims];
@@ -77,9 +85,19 @@ public void setup() throws IOException {
                     binaryVector[i] = (byte) ThreadLocalRandom.current().nextInt(16);
                 }
                 out.writeBytes(binaryVector, 0, binaryVector.length);
+                ThreadLocalRandom.current().nextBytes(binaryVector);
+                out.writeBytes(binaryVector, 0, 14); // corrections
             }
         }
 
+        queryCorrections = new OptimizedScalarQuantizer.QuantizationResult(
+            ThreadLocalRandom.current().nextFloat(),
+            ThreadLocalRandom.current().nextFloat(),
+            ThreadLocalRandom.current().nextFloat(),
+            Short.toUnsignedInt((short) ThreadLocalRandom.current().nextInt())
+        );
+        centroidDp = ThreadLocalRandom.current().nextFloat();
+
         in = dir.openInput("vectors", IOContext.DEFAULT);
         binaryQueries = new byte[numVectors][dims];
         for (byte[] binaryVector : binaryVectors) {
@@ -105,18 +123,66 @@ public void scoreFromArray(Blackhole bh) throws IOException {
             in.seek(0);
             for (int i = 0; i < numVectors; i++) {
                 in.readBytes(scratch, 0, dims);
-                bh.consume(VectorUtil.int4DotProduct(binaryQueries[j], scratch));
+                int dp = VectorUtil.int4DotProduct(binaryQueries[j], scratch);
+                in.readFloats(scratchFloats, 0, 3);
+                float score = scorer.applyCorrections(
+                    queryCorrections.lowerInterval(),
+                    queryCorrections.upperInterval(),
+                    queryCorrections.quantizedComponentSum(),
+                    queryCorrections.additionalCorrection(),
+                    VectorSimilarityFunction.EUCLIDEAN,
+                    centroidDp, // assuming no centroid dot product for this benchmark
+                    scratchFloats[0],
+                    scratchFloats[1],
+                    Short.toUnsignedInt(in.readShort()),
+                    scratchFloats[2],
+                    dp
+                );
+                bh.consume(score);
             }
         }
     }
 
     @Benchmark
     @Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })
-    public void scoreFromMemorySegmentOnlyVector(Blackhole bh) throws IOException {
+    public void scoreFromMemorySegment(Blackhole bh) throws IOException {
         for (int j = 0; j < numQueries; j++) {
             in.seek(0);
             for (int i = 0; i < numVectors; i++) {
-                bh.consume(scorer.int4DotProduct(binaryQueries[j]));
+                bh.consume(
+                    scorer.score(
+                        binaryQueries[j],
+                        queryCorrections.lowerInterval(),
+                        queryCorrections.upperInterval(),
+                        queryCorrections.quantizedComponentSum(),
+                        queryCorrections.additionalCorrection(),
+                        VectorSimilarityFunction.EUCLIDEAN,
+                        centroidDp
+                    )
+                );
+            }
+        }
+    }
+
+    @Benchmark
+    @Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })
+    public void scoreFromMemorySegmentBulk(Blackhole bh) throws IOException {
+        for (int j = 0; j < numQueries; j++) {
+            in.seek(0);
+            for (int i = 0; i < numVectors; i += ES91Int4VectorsScorer.BULK_SIZE) {
+                scorer.scoreBulk(
+                    binaryQueries[j],
+                    queryCorrections.lowerInterval(),
+                    queryCorrections.upperInterval(),
+                    queryCorrections.quantizedComponentSum(),
+                    queryCorrections.additionalCorrection(),
+                    VectorSimilarityFunction.EUCLIDEAN,
+                    centroidDp,
+                    scores
+                );
+                for (float score : scores) {
+                    bh.consume(score);
+                }
             }
         }
     }

build-tools-internal/src/main/java/org/elasticsearch/gradle/internal/ElasticsearchTestBasePlugin.java

@@ -34,6 +34,7 @@
 import java.io.File;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.stream.Stream;
 
 import javax.inject.Inject;
@@ -50,6 +51,8 @@ public abstract class ElasticsearchTestBasePlugin implements Plugin<Project> {
 
     public static final String DUMP_OUTPUT_ON_FAILURE_PROP_NAME = "dumpOutputOnFailure";
 
+    public static final Set<String> TEST_TASKS_WITH_ENTITLEMENTS = Set.of("test", "internalClusterTest");
+
     @Inject
     protected abstract ProviderFactory getProviderFactory();
 
@@ -174,14 +177,23 @@ public void execute(Task t) {
             nonInputProperties.systemProperty("workspace.dir", Util.locateElasticsearchWorkspace(project.getGradle()));
             // we use 'temp' relative to CWD since this is per JVM and tests are forbidden from writing to CWD
             nonInputProperties.systemProperty("java.io.tmpdir", test.getWorkingDir().toPath().resolve("temp"));
+            if (test.getName().equals("internalClusterTest")) {
+                // configure a node home directory independent of the Java temp dir so that entitlements can be properly enforced
+                nonInputProperties.systemProperty("tempDir", test.getWorkingDir().toPath().resolve("nodesTemp"));
+            }
 
             SourceSetContainer sourceSets = project.getExtensions().getByType(SourceSetContainer.class);
             SourceSet mainSourceSet = sourceSets.findByName(SourceSet.MAIN_SOURCE_SET_NAME);
             SourceSet testSourceSet = sourceSets.findByName(SourceSet.TEST_SOURCE_SET_NAME);
-            if ("test".equals(test.getName()) && mainSourceSet != null && testSourceSet != null) {
+            SourceSet internalClusterTestSourceSet = sourceSets.findByName("internalClusterTest");
+
+            if (TEST_TASKS_WITH_ENTITLEMENTS.contains(test.getName()) && mainSourceSet != null && testSourceSet != null) {
                 FileCollection mainRuntime = mainSourceSet.getRuntimeClasspath();
                 FileCollection testRuntime = testSourceSet.getRuntimeClasspath();
-                FileCollection testOnlyFiles = testRuntime.minus(mainRuntime);
+                FileCollection internalClusterTestRuntime = ("internalClusterTest".equals(test.getName())
+                    && internalClusterTestSourceSet != null) ? internalClusterTestSourceSet.getRuntimeClasspath() : project.files();
+                FileCollection testOnlyFiles = testRuntime.plus(internalClusterTestRuntime).minus(mainRuntime);
+
                 test.doFirst(task -> test.environment("es.entitlement.testOnlyPath", testOnlyFiles.getAsPath()));
             }
 
@@ -241,14 +253,15 @@ public void execute(Task t) {
      * Computes and sets the {@code --patch-module=java.base} and {@code --add-opens=java.base} JVM command line options.
      */
     private void configureJavaBaseModuleOptions(Project project) {
-        project.getTasks().withType(Test.class).matching(task -> task.getName().equals("test")).configureEach(test -> {
-            FileCollection patchedImmutableCollections = patchedImmutableCollections(project);
+        project.getTasks().withType(Test.class).configureEach(test -> {
+            // patch immutable collections only for "test" task
+            FileCollection patchedImmutableCollections = test.getName().equals("test") ? patchedImmutableCollections(project) : null;
             if (patchedImmutableCollections != null) {
                 test.getInputs().files(patchedImmutableCollections);
                 test.systemProperty("tests.hackImmutableCollections", "true");
             }
 
-            FileCollection entitlementBridge = entitlementBridge(project);
+            FileCollection entitlementBridge = TEST_TASKS_WITH_ENTITLEMENTS.contains(test.getName()) ? entitlementBridge(project) : null;
             if (entitlementBridge != null) {
                 test.getInputs().files(entitlementBridge);
             }
@@ -312,27 +325,30 @@ private static void configureEntitlements(Project project) {
         }
         FileCollection bridgeFiles = bridgeConfig;
 
-        project.getTasks().withType(Test.class).configureEach(test -> {
-            // See also SystemJvmOptions.maybeAttachEntitlementAgent.
-
-            // Agent
-            if (agentFiles.isEmpty() == false) {
-                test.getInputs().files(agentFiles);
-                test.systemProperty("es.entitlement.agentJar", agentFiles.getAsPath());
-                test.systemProperty("jdk.attach.allowAttachSelf", true);
-            }
+        project.getTasks()
+            .withType(Test.class)
+            .matching(test -> TEST_TASKS_WITH_ENTITLEMENTS.contains(test.getName()))
+            .configureEach(test -> {
+                // See also SystemJvmOptions.maybeAttachEntitlementAgent.
+
+                // Agent
+                if (agentFiles.isEmpty() == false) {
+                    test.getInputs().files(agentFiles);
+                    test.systemProperty("es.entitlement.agentJar", agentFiles.getAsPath());
+                    test.systemProperty("jdk.attach.allowAttachSelf", true);
+                }
 
-            // Bridge
-            if (bridgeFiles.isEmpty() == false) {
-                String modulesContainingEntitlementInstrumentation = "java.logging,java.net.http,java.naming,jdk.net";
-                test.getInputs().files(bridgeFiles);
-                // Tests may not be modular, but the JDK still is
-                test.jvmArgs(
-                    "--add-exports=java.base/org.elasticsearch.entitlement.bridge=ALL-UNNAMED,"
-                        + modulesContainingEntitlementInstrumentation
-                );
-            }
-        });
+                // Bridge
+                if (bridgeFiles.isEmpty() == false) {
+                    String modulesContainingEntitlementInstrumentation = "java.logging,java.net.http,java.naming,jdk.net";
+                    test.getInputs().files(bridgeFiles);
+                    // Tests may not be modular, but the JDK still is
+                    test.jvmArgs(
+                        "--add-exports=java.base/org.elasticsearch.entitlement.bridge=ALL-UNNAMED,"
+                            + modulesContainingEntitlementInstrumentation
+                    );
+                }
+            });
     }
 
 }

build-tools-internal/src/main/java/org/elasticsearch/gradle/internal/RestrictedBuildApiService.java

@@ -56,8 +56,6 @@ private static ListMultimap<Class<?>, String> createLegacyRestTestBasePluginUsag
         map.put(LegacyRestTestBasePlugin.class, ":x-pack:qa:third-party:jira");
         map.put(LegacyRestTestBasePlugin.class, ":x-pack:qa:third-party:pagerduty");
         map.put(LegacyRestTestBasePlugin.class, ":x-pack:qa:third-party:slack");
-        map.put(LegacyRestTestBasePlugin.class, ":x-pack:plugin:async-search:qa:rest");
-        map.put(LegacyRestTestBasePlugin.class, ":x-pack:plugin:autoscaling:qa:rest");
         map.put(LegacyRestTestBasePlugin.class, ":x-pack:plugin:deprecation:qa:early-deprecation-rest");
         map.put(LegacyRestTestBasePlugin.class, ":x-pack:plugin:deprecation:qa:rest");
         map.put(LegacyRestTestBasePlugin.class, ":x-pack:plugin:downsample:qa:with-security");