Skip to content

Commit 75bb40a

Browse files
mromaiosmromaios
authored
Merge branch 'main' into increase_sparse_pruning_test_coverage
2 parents c49c67e + fc8fc5a commit 75bb40a

File tree

547 files changed

+13783
-5096
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

547 files changed

+13783
-5096
lines changed

‎.gitignore‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,7 @@ testfixtures_shared/
6969
# Generated
7070
checkstyle_ide.xml
7171
x-pack/plugin/esql/src/main/generated-src/generated/
72+
server/src/main/resources/transport/defined/manifest.txt
7273

7374
# JEnv
7475
.java-version

‎benchmarks/src/main/java/org/elasticsearch/benchmark/vector/DistanceBulkBenchmark.java‎

Lines changed: 145 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,145 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the "Elastic License
4+
* 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
5+
* Public License v 1"; you may not use this file except in compliance with, at
6+
* your election, the "Elastic License 2.0", the "GNU Affero General Public
7+
* License v3.0 only", or the "Server Side Public License, v 1".
8+
*/
9+
package org.elasticsearch.benchmark.vector;
10+
11+
import org.apache.lucene.util.VectorUtil;
12+
import org.apache.lucene.util.quantization.OptimizedScalarQuantizer;
13+
import org.elasticsearch.common.logging.LogConfigurator;
14+
import org.elasticsearch.simdvec.ESVectorUtil;
15+
import org.openjdk.jmh.annotations.Benchmark;
16+
import org.openjdk.jmh.annotations.BenchmarkMode;
17+
import org.openjdk.jmh.annotations.Fork;
18+
import org.openjdk.jmh.annotations.Measurement;
19+
import org.openjdk.jmh.annotations.Mode;
20+
import org.openjdk.jmh.annotations.OutputTimeUnit;
21+
import org.openjdk.jmh.annotations.Param;
22+
import org.openjdk.jmh.annotations.Scope;
23+
import org.openjdk.jmh.annotations.Setup;
24+
import org.openjdk.jmh.annotations.State;
25+
import org.openjdk.jmh.annotations.Warmup;
26+
import org.openjdk.jmh.infra.Blackhole;
27+
28+
import java.io.IOException;
29+
import java.util.Random;
30+
import java.util.concurrent.TimeUnit;
31+
32+
@BenchmarkMode(Mode.Throughput)
33+
@OutputTimeUnit(TimeUnit.MILLISECONDS)
34+
@State(Scope.Benchmark)
35+
// first iteration is complete garbage, so make sure we really warmup
36+
@Warmup(iterations = 4, time = 1)
37+
// real iterations. not useful to spend tons of time here, better to fork more
38+
@Measurement(iterations = 5, time = 1)
39+
// engage some noise reduction
40+
@Fork(value = 1)
41+
public class DistanceBulkBenchmark {
42+
43+
static {
44+
LogConfigurator.configureESLogging(); // native access requires logging to be initialized
45+
}
46+
47+
@Param({ "384", "782", "1024" })
48+
int dims;
49+
50+
int length;
51+
52+
int numVectors = 4 * 100;
53+
int numQueries = 10;
54+
55+
float[][] vectors;
56+
float[][] queries;
57+
float[] distances = new float[4];
58+
59+
@Setup
60+
public void setup() throws IOException {
61+
Random random = new Random(123);
62+
63+
this.length = OptimizedScalarQuantizer.discretize(dims, 64) / 8;
64+
65+
vectors = new float[numVectors][dims];
66+
for (float[] vector : vectors) {
67+
for (int i = 0; i < dims; i++) {
68+
vector[i] = random.nextFloat();
69+
}
70+
}
71+
72+
queries = new float[numQueries][dims];
73+
for (float[] query : queries) {
74+
for (int i = 0; i < dims; i++) {
75+
query[i] = random.nextFloat();
76+
}
77+
}
78+
}
79+
80+
@Benchmark
81+
@Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })
82+
public void squareDistance(Blackhole bh) {
83+
for (int j = 0; j < numQueries; j++) {
84+
float[] query = queries[j];
85+
for (int i = 0; i < numVectors; i++) {
86+
float[] vector = vectors[i];
87+
float distance = VectorUtil.squareDistance(query, vector);
88+
bh.consume(distance);
89+
}
90+
}
91+
}
92+
93+
@Benchmark
94+
@Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })
95+
public void soarDistance(Blackhole bh) {
96+
for (int j = 0; j < numQueries; j++) {
97+
float[] query = queries[j];
98+
for (int i = 0; i < numVectors; i++) {
99+
float[] vector = vectors[i];
100+
float distance = ESVectorUtil.soarDistance(query, vector, vector, 1.0f, 1.0f);
101+
bh.consume(distance);
102+
}
103+
}
104+
}
105+
106+
@Benchmark
107+
@Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })
108+
public void squareDistanceBulk(Blackhole bh) {
109+
for (int j = 0; j < numQueries; j++) {
110+
float[] query = queries[j];
111+
for (int i = 0; i < numVectors; i += 4) {
112+
ESVectorUtil.squareDistanceBulk(query, vectors[i], vectors[i + 1], vectors[i + 2], vectors[i + 3], distances);
113+
for (float distance : distances) {
114+
bh.consume(distance);
115+
}
116+
117+
}
118+
}
119+
}
120+
121+
@Benchmark
122+
@Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })
123+
public void soarDistanceBulk(Blackhole bh) {
124+
for (int j = 0; j < numQueries; j++) {
125+
float[] query = queries[j];
126+
for (int i = 0; i < numVectors; i += 4) {
127+
ESVectorUtil.soarDistanceBulk(
128+
query,
129+
vectors[i],
130+
vectors[i + 1],
131+
vectors[i + 2],
132+
vectors[i + 3],
133+
vectors[i],
134+
1.0f,
135+
1.0f,
136+
distances
137+
);
138+
for (float distance : distances) {
139+
bh.consume(distance);
140+
}
141+
142+
}
143+
}
144+
}
145+
}

‎build-tools-internal/build.gradle‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -220,6 +220,14 @@ gradlePlugin {
220220
id = 'elasticsearch.internal-yaml-rest-test'
221221
implementationClass = 'org.elasticsearch.gradle.internal.test.rest.InternalYamlRestTestPlugin'
222222
}
223+
transportVersionManagementPlugin {
224+
id = 'elasticsearch.transport-version-management'
225+
implementationClass = 'org.elasticsearch.gradle.internal.transport.TransportVersionManagementPlugin'
226+
}
227+
globalTransportVersionManagementPlugin {
228+
id = 'elasticsearch.global-transport-version-management'
229+
implementationClass = 'org.elasticsearch.gradle.internal.transport.GlobalTransportVersionManagementPlugin'
230+
}
223231
}
224232
}
225233

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
import org.elasticsearch.gradle.internal.info.BuildParameterExtension;
1616
import org.elasticsearch.gradle.internal.precommit.JarHellPrecommitPlugin;
1717
import org.elasticsearch.gradle.internal.test.ClusterFeaturesMetadataPlugin;
18+
import org.elasticsearch.gradle.internal.transport.TransportVersionManagementPlugin;
1819
import org.elasticsearch.gradle.plugin.PluginBuildPlugin;
1920
import org.elasticsearch.gradle.plugin.PluginPropertiesExtension;
2021
import org.elasticsearch.gradle.util.GradleUtils;
@@ -36,6 +37,7 @@ public void apply(Project project) {
3637
project.getPluginManager().apply(JarHellPrecommitPlugin.class);
3738
project.getPluginManager().apply(ElasticsearchJavaPlugin.class);
3839
project.getPluginManager().apply(ClusterFeaturesMetadataPlugin.class);
40+
project.getPluginManager().apply(TransportVersionManagementPlugin.class);
3941
boolean isCi = project.getRootProject().getExtensions().getByType(BuildParameterExtension.class).getCi();
4042
// Clear default dependencies added by public PluginBuildPlugin as we add our
4143
// own project dependencies for internal builds

‎build-tools-internal/src/main/java/org/elasticsearch/gradle/internal/toolchain/OracleOpenJdkToolchainResolver.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -62,8 +62,8 @@ record EarlyAccessJdkBuild(JavaLanguageVersion languageVersion, String buildNumb
6262
@Override
6363
public String url(String os, String arch, String extension) {
6464
// example:
65-
// http://builds.es-jdk-archive.com/jdks/openjdk/26/openjdk-26-ea+6/openjdk-26-ea+6_linux-aarch64_bin.tar.gz
66-
return "http://builds.es-jdk-archive.com/jdks/openjdk/"
65+
// https://builds.es-jdk-archive.com/jdks/openjdk/26/openjdk-26-ea+6/openjdk-26-ea+6_linux-aarch64_bin.tar.gz
66+
return "https://builds.es-jdk-archive.com/jdks/openjdk/"
6767
+ languageVersion.asInt()
6868
+ "/"
6969
+ "openjdk-"

‎build-tools-internal/src/main/java/org/elasticsearch/gradle/internal/transport/CollectTransportVersionReferencesTask.java‎

Lines changed: 133 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,133 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the "Elastic License
4+
* 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
5+
* Public License v 1"; you may not use this file except in compliance with, at
6+
* your election, the "Elastic License 2.0", the "GNU Affero General Public
7+
* License v3.0 only", or the "Server Side Public License, v 1".
8+
*/
9+
10+
package org.elasticsearch.gradle.internal.transport;
11+
12+
import org.gradle.api.DefaultTask;
13+
import org.gradle.api.file.ConfigurableFileCollection;
14+
import org.gradle.api.file.RegularFileProperty;
15+
import org.gradle.api.tasks.CacheableTask;
16+
import org.gradle.api.tasks.Classpath;
17+
import org.gradle.api.tasks.OutputFile;
18+
import org.gradle.api.tasks.TaskAction;
19+
import org.objectweb.asm.ClassReader;
20+
import org.objectweb.asm.ClassVisitor;
21+
import org.objectweb.asm.Label;
22+
import org.objectweb.asm.MethodVisitor;
23+
import org.objectweb.asm.Opcodes;
24+
import org.objectweb.asm.tree.LdcInsnNode;
25+
import org.objectweb.asm.tree.MethodNode;
26+
27+
import java.io.IOException;
28+
import java.io.InputStream;
29+
import java.nio.file.FileVisitResult;
30+
import java.nio.file.Files;
31+
import java.nio.file.Path;
32+
import java.nio.file.SimpleFileVisitor;
33+
import java.nio.file.attribute.BasicFileAttributes;
34+
import java.util.HashSet;
35+
import java.util.Set;
36+
37+
/**
38+
* This task locates all method invocations of org.elasticsearch.TransportVersion#fromName(java.lang.String) in the
39+
* provided directory, and then records the value of string literals passed as arguments. It then records each
40+
* string on a newline along with path and line number in the provided output file.
41+
*/
42+
@CacheableTask
43+
public abstract class CollectTransportVersionReferencesTask extends DefaultTask {
44+
public static final String TRANSPORT_VERSION_SET_CLASS = "org/elasticsearch/TransportVersion";
45+
public static final String TRANSPORT_VERSION_SET_METHOD_NAME = "fromName";
46+
public static final String CLASS_EXTENSION = ".class";
47+
public static final String MODULE_INFO = "module-info.class";
48+
49+
/**
50+
* The directory to scan for method invocations.
51+
*/
52+
@Classpath
53+
public abstract ConfigurableFileCollection getClassPath();
54+
55+
/**
56+
* The output file, with each newline containing the string literal argument of each method
57+
* invocation.
58+
*/
59+
@OutputFile
60+
public abstract RegularFileProperty getOutputFile();
61+
62+
@TaskAction
63+
public void checkTransportVersion() throws IOException {
64+
var results = new HashSet<TransportVersionUtils.TransportVersionReference>();
65+
66+
for (var cpElement : getClassPath()) {
67+
Path file = cpElement.toPath();
68+
if (Files.isDirectory(file)) {
69+
addNamesFromClassesDirectory(results, file);
70+
}
71+
}
72+
73+
Path outputFile = getOutputFile().get().getAsFile().toPath();
74+
Files.writeString(outputFile, String.join("\n", results.stream().map(Object::toString).sorted().toList()));
75+
}
76+
77+
private void addNamesFromClassesDirectory(Set<TransportVersionUtils.TransportVersionReference> results, Path file) throws IOException {
78+
Files.walkFileTree(file, new SimpleFileVisitor<>() {
79+
@Override
80+
public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) throws IOException {
81+
String filename = file.getFileName().toString();
82+
if (filename.endsWith(CLASS_EXTENSION) && filename.endsWith(MODULE_INFO) == false) {
83+
try (var inputStream = Files.newInputStream(file)) {
84+
addNamesFromClass(results, inputStream, classname(file.toString()));
85+
}
86+
}
87+
return FileVisitResult.CONTINUE;
88+
}
89+
});
90+
}
91+
92+
private void addNamesFromClass(Set<TransportVersionUtils.TransportVersionReference> results, InputStream classBytes, String classname)
93+
throws IOException {
94+
ClassVisitor classVisitor = new ClassVisitor(Opcodes.ASM9) {
95+
@Override
96+
public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions) {
97+
return new MethodNode(Opcodes.ASM9, access, name, descriptor, signature, exceptions) {
98+
int lineNumber = -1;
99+
100+
@Override
101+
public void visitLineNumber(int line, Label start) {
102+
lineNumber = line;
103+
}
104+
105+
@Override
106+
public void visitMethodInsn(int opcode, String owner, String name, String descriptor, boolean isInterface) {
107+
if (owner.equals(TRANSPORT_VERSION_SET_CLASS) && name.equals(TRANSPORT_VERSION_SET_METHOD_NAME)) {
108+
var abstractInstruction = this.instructions.getLast();
109+
String location = classname + " line " + lineNumber;
110+
if (abstractInstruction instanceof LdcInsnNode ldcInsnNode
111+
&& ldcInsnNode.cst instanceof String tvName
112+
&& tvName.isEmpty() == false) {
113+
results.add(new TransportVersionUtils.TransportVersionReference(tvName, location));
114+
} else {
115+
// The instruction is not a LDC with a String constant (or an empty String), which is not allowed.
116+
throw new RuntimeException(
117+
"TransportVersion.fromName must be called with a non-empty String literal. " + "See " + location + "."
118+
);
119+
}
120+
}
121+
super.visitMethodInsn(opcode, owner, name, descriptor, isInterface);
122+
}
123+
};
124+
}
125+
};
126+
ClassReader classReader = new ClassReader(classBytes);
127+
classReader.accept(classVisitor, 0);
128+
}
129+
130+
private static String classname(String filename) {
131+
return filename.substring(0, filename.length() - CLASS_EXTENSION.length()).replaceAll("[/\\\\]", ".");
132+
}
133+
}

‎build-tools-internal/src/main/java/org/elasticsearch/gradle/internal/transport/GenerateTransportVersionManifestTask.java‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the "Elastic License
4+
* 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
5+
* Public License v 1"; you may not use this file except in compliance with, at
6+
* your election, the "Elastic License 2.0", the "GNU Affero General Public
7+
* License v3.0 only", or the "Server Side Public License, v 1".
8+
*/
9+
10+
package org.elasticsearch.gradle.internal.transport;
11+
12+
import org.gradle.api.DefaultTask;
13+
import org.gradle.api.file.DirectoryProperty;
14+
import org.gradle.api.file.RegularFileProperty;
15+
import org.gradle.api.tasks.InputDirectory;
16+
import org.gradle.api.tasks.OutputFile;
17+
import org.gradle.api.tasks.TaskAction;
18+
19+
import java.io.IOException;
20+
import java.nio.file.Files;
21+
import java.nio.file.Path;
22+
23+
public abstract class GenerateTransportVersionManifestTask extends DefaultTask {
24+
@InputDirectory
25+
public abstract DirectoryProperty getDefinitionsDirectory();
26+
27+
@OutputFile
28+
public abstract RegularFileProperty getManifestFile();
29+
30+
@TaskAction
31+
public void generateTransportVersionManifest() throws IOException {
32+
Path constantsDir = getDefinitionsDirectory().get().getAsFile().toPath();
33+
Path manifestFile = getManifestFile().get().getAsFile().toPath();
34+
try (var writer = Files.newBufferedWriter(manifestFile)) {
35+
try (var stream = Files.list(constantsDir)) {
36+
for (String filename : stream.map(p -> p.getFileName().toString()).toList()) {
37+
if (filename.equals(manifestFile.getFileName().toString())) {
38+
// don't list self
39+
continue;
40+
}
41+
writer.write(filename + "\n");
42+
}
43+
}
44+
}
45+
}
46+
}

0 commit comments

Comments
 (0)