Skip to content

Commit

Permalink
Merge branch 'main' of github.com:elastic/elasticsearch into ml-simil…
Browse files Browse the repository at this point in the history 
…arity
  • Loading branch information
@jonathan-buttner
jonathan-buttner committed Mar 21, 2024
2 parents 0975793 + e5b6040 commit f58e7ef
Show file tree
Hide file tree
Showing 809 changed files with 12,532 additions and 2,738 deletions.
34 changes: 32 additions & 2 deletions benchmarks/src/main/java/org/elasticsearch/benchmark/compute/operator/EvalBenchmark.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@

package org.elasticsearch.benchmark.compute.operator;

import org.apache.lucene.util.BytesRef;
import org.elasticsearch.common.breaker.NoopCircuitBreaker;
import org.elasticsearch.common.util.BigArrays;
import org.elasticsearch.compute.data.Block;
Expand All @@ -26,11 +27,13 @@
import org.elasticsearch.xpack.esql.expression.function.scalar.date.DateTrunc;
import org.elasticsearch.xpack.esql.expression.function.scalar.math.Abs;
import org.elasticsearch.xpack.esql.expression.function.scalar.multivalue.MvMin;
import org.elasticsearch.xpack.esql.expression.function.scalar.string.RLike;
import org.elasticsearch.xpack.esql.expression.predicate.operator.arithmetic.Add;
import org.elasticsearch.xpack.esql.planner.Layout;
import org.elasticsearch.xpack.esql.type.EsqlDataTypes;
import org.elasticsearch.xpack.ql.expression.FieldAttribute;
import org.elasticsearch.xpack.ql.expression.Literal;
import org.elasticsearch.xpack.ql.expression.predicate.regex.RLikePattern;
import org.elasticsearch.xpack.ql.tree.Source;
import org.elasticsearch.xpack.ql.type.DataTypes;
import org.elasticsearch.xpack.ql.type.EsField;
Expand Down Expand Up @@ -58,7 +61,6 @@
@State(Scope.Thread)
@Fork(1)
public class EvalBenchmark {
private static final BigArrays BIG_ARRAYS = BigArrays.NON_RECYCLING_INSTANCE; // TODO real big arrays?
private static final BlockFactory blockFactory = BlockFactory.getInstance(
new NoopCircuitBreaker("noop"),
BigArrays.NON_RECYCLING_INSTANCE
Expand All @@ -82,7 +84,9 @@ public class EvalBenchmark {
}
}

@Param({ "abs", "add", "date_trunc", "equal_to_const", "long_equal_to_long", "long_equal_to_int", "mv_min", "mv_min_ascending" })
@Param(
{ "abs", "add", "date_trunc", "equal_to_const", "long_equal_to_long", "long_equal_to_int", "mv_min", "mv_min_ascending", "rlike" }
)
public String operation;

private static Operator operator(String operation) {
Expand Down Expand Up @@ -134,6 +138,11 @@ private static EvalOperator.ExpressionEvaluator evaluator(String operation) {
FieldAttribute longField = longField();
yield EvalMapper.toEvaluator(new MvMin(Source.EMPTY, longField), layout(longField)).get(driverContext);
}
case "rlike" -> {
FieldAttribute keywordField = keywordField();
RLike rlike = new RLike(Source.EMPTY, keywordField, new RLikePattern(".ar"));
yield EvalMapper.toEvaluator(rlike, layout(keywordField)).get(driverContext);
}
default -> throw new UnsupportedOperationException();
};
}
Expand All @@ -146,6 +155,10 @@ private static FieldAttribute intField() {
return new FieldAttribute(Source.EMPTY, "int", new EsField("int", DataTypes.INTEGER, Map.of(), true));
}

private static FieldAttribute keywordField() {
return new FieldAttribute(Source.EMPTY, "keyword", new EsField("keyword", DataTypes.KEYWORD, Map.of(), true));
}

private static Layout layout(FieldAttribute... fields) {
Layout.Builder layout = new Layout.Builder();
layout.append(Arrays.asList(fields));
Expand Down Expand Up @@ -205,6 +218,15 @@ private static void checkExpected(String operation, Page actual) {
}
}
}
case "rlike" -> {
BooleanVector v = actual.<BooleanBlock>getBlock(1).asVector();
for (int i = 0; i < BLOCK_LENGTH; i++) {
boolean expected = i % 2 == 1;
if (v.getBoolean(i) != expected) {
throw new AssertionError("[" + operation + "] expected [" + expected + "] but was [" + v.getBoolean(i) + "]");
}
}
}
default -> throw new UnsupportedOperationException();
}
}
Expand Down Expand Up @@ -250,6 +272,14 @@ private static Page page(String operation) {
}
yield new Page(builder.build());
}
case "rlike" -> {
var builder = blockFactory.newBytesRefVectorBuilder(BLOCK_LENGTH);
BytesRef[] values = new BytesRef[] { new BytesRef("foo"), new BytesRef("bar") };
for (int i = 0; i < BLOCK_LENGTH; i++) {
builder.appendBytesRef(values[i % 2]);
}
yield new Page(builder.build().asBlock());
}
default -> throw new UnsupportedOperationException();
};
}
Expand Down
4 changes: 4 additions & 0 deletions build-tools-internal/build.gradle
View file
Original file line number Diff line number Diff line change
Expand Up @@ -179,6 +179,10 @@ gradlePlugin {
id = 'elasticsearch.test.fixtures'
implementationClass = 'org.elasticsearch.gradle.internal.testfixtures.TestFixturesPlugin'
}
deployTestFixtures {
id = 'elasticsearch.deploy-test-fixtures'
implementationClass = 'org.elasticsearch.gradle.internal.testfixtures.TestFixturesDeployPlugin'
}
testBase {
id = 'elasticsearch.test-base'
implementationClass = 'org.elasticsearch.gradle.internal.ElasticsearchTestBasePlugin'
Expand Down
14 changes: 11 additions & 3 deletions ...nal/src/main/java/org/elasticsearch/gradle/internal/ElasticsearchBuildCompletePlugin.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,12 @@
import org.gradle.api.tasks.Input;
import org.jetbrains.annotations.NotNull;

import java.io.*;
import java.io.BufferedInputStream;
import java.io.BufferedOutputStream;
import java.io.File;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.io.OutputStream;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.ArrayList;
Expand Down Expand Up @@ -210,12 +215,15 @@ private static void createBuildArchiveTar(List<File> files, File projectDir, Fil
throw new IOException("Support only file!");
}

long entrySize = Files.size(path);
TarArchiveEntry tarEntry = new TarArchiveEntry(path.toFile(), calculateArchivePath(path, projectPath));
tarEntry.setSize(Files.size(path));
tarEntry.setSize(entrySize);
tOut.putArchiveEntry(tarEntry);

// copy file to TarArchiveOutputStream
Files.copy(path, tOut);
try (BufferedInputStream bin = new BufferedInputStream(Files.newInputStream(path))) {
IOUtils.copyLarge(bin, tOut, 0, entrySize);
}
tOut.closeArchiveEntry();

}
Expand Down
129 changes: 92 additions & 37 deletions build-tools-internal/src/main/java/org/elasticsearch/gradle/internal/MrjarPlugin.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@

package org.elasticsearch.gradle.internal;

import org.elasticsearch.gradle.internal.precommit.CheckForbiddenApisTask;
import org.elasticsearch.gradle.util.GradleUtils;
import org.gradle.api.Plugin;
import org.gradle.api.Project;
Expand Down Expand Up @@ -40,6 +41,7 @@

import javax.inject.Inject;

import static de.thetaphi.forbiddenapis.gradle.ForbiddenApisPlugin.FORBIDDEN_APIS_TASK_NAME;
import static org.objectweb.asm.Opcodes.V_PREVIEW;

public class MrjarPlugin implements Plugin<Project> {
Expand All @@ -58,48 +60,30 @@ public void apply(Project project) {
project.getPluginManager().apply(ElasticsearchJavaBasePlugin.class);
var javaExtension = project.getExtensions().getByType(JavaPluginExtension.class);

var srcDir = project.getProjectDir().toPath().resolve("src");
List<Integer> mainVersions = new ArrayList<>();
try (var subdirStream = Files.list(srcDir)) {
for (Path sourceset : subdirStream.toList()) {
assert Files.isDirectory(sourceset);
String sourcesetName = sourceset.getFileName().toString();
Matcher sourcesetMatcher = MRJAR_SOURCESET_PATTERN.matcher(sourcesetName);
if (sourcesetMatcher.matches()) {
mainVersions.add(Integer.parseInt(sourcesetMatcher.group(1)));
}
}
} catch (IOException e) {
throw new UncheckedIOException(e);
}

Collections.sort(mainVersions);
List<String> parentSourceSets = new ArrayList<>();
parentSourceSets.add(SourceSet.MAIN_SOURCE_SET_NAME);
List<Integer> mainVersions = findSourceVersions(project);
List<String> mainSourceSets = new ArrayList<>();
mainSourceSets.add(SourceSet.MAIN_SOURCE_SET_NAME);
List<String> testSourceSets = new ArrayList<>(mainSourceSets);
testSourceSets.add(SourceSet.TEST_SOURCE_SET_NAME);
for (int javaVersion : mainVersions) {
String sourcesetName = "main" + javaVersion;
addMrjarSourceset(project, javaExtension, sourcesetName, parentSourceSets, javaVersion);
parentSourceSets.add(sourcesetName);
String mainSourceSetName = SourceSet.MAIN_SOURCE_SET_NAME + javaVersion;
SourceSet mainSourceSet = addSourceSet(project, javaExtension, mainSourceSetName, mainSourceSets, javaVersion);
configureSourceSetInJar(project, mainSourceSet, javaVersion);
mainSourceSets.add(mainSourceSetName);
testSourceSets.add(mainSourceSetName);

String testSourceSetName = SourceSet.TEST_SOURCE_SET_NAME + javaVersion;
SourceSet testSourceSet = addSourceSet(project, javaExtension, testSourceSetName, testSourceSets, javaVersion);
testSourceSets.add(testSourceSetName);
createTestTask(project, testSourceSet, javaVersion, mainSourceSets);
}
}

private void addMrjarSourceset(
Project project,
JavaPluginExtension javaExtension,
String sourcesetName,
List<String> parentSourceSets,
int javaVersion
) {
SourceSet sourceSet = javaExtension.getSourceSets().maybeCreate(sourcesetName);
for (String parentSourceSetName : parentSourceSets) {
GradleUtils.extendSourceSet(project, parentSourceSetName, sourcesetName);
}
configureMrjar(project);
}

private void configureMrjar(Project project) {
var jarTask = project.getTasks().withType(Jar.class).named(JavaPlugin.JAR_TASK_NAME);
jarTask.configure(task -> {
task.into("META-INF/versions/" + javaVersion, copySpec -> copySpec.from(sourceSet.getOutput()));
task.manifest(manifest -> { manifest.attributes(Map.of("Multi-Release", "true")); });
});
jarTask.configure(task -> { task.manifest(manifest -> { manifest.attributes(Map.of("Multi-Release", "true")); }); });

project.getTasks().withType(Test.class).named(JavaPlugin.TEST_TASK_NAME).configure(testTask -> {
testTask.dependsOn(jarTask);
Expand All @@ -109,6 +93,19 @@ private void addMrjarSourceset(
FileCollection testRuntime = sourceSets.getByName(SourceSet.TEST_SOURCE_SET_NAME).getRuntimeClasspath();
testTask.setClasspath(testRuntime.minus(mainRuntime).plus(project.files(jarTask)));
});
}

private SourceSet addSourceSet(
Project project,
JavaPluginExtension javaExtension,
String sourceSetName,
List<String> parentSourceSets,
int javaVersion
) {
SourceSet sourceSet = javaExtension.getSourceSets().maybeCreate(sourceSetName);
for (String parentSourceSetName : parentSourceSets) {
GradleUtils.extendSourceSet(project, parentSourceSetName, sourceSetName);
}

project.getTasks().withType(JavaCompile.class).named(sourceSet.getCompileJavaTaskName()).configure(compileTask -> {
compileTask.getJavaCompiler()
Expand All @@ -121,6 +118,64 @@ private void addMrjarSourceset(

compileTask.doLast(t -> { stripPreviewFromFiles(compileTask.getDestinationDirectory().getAsFile().get().toPath()); });
});

// Since we configure MRJAR sourcesets to allow preview apis, class signatures for those
// apis are not known by forbidden apis, so we must ignore all missing classes. We could, in theory,
// run forbidden apis in a separate jvm matching the sourceset jvm, but it's not worth
// the complexity (according to forbidden apis author!)
String forbiddenApisTaskName = sourceSet.getTaskName(FORBIDDEN_APIS_TASK_NAME, null);
project.getTasks().withType(CheckForbiddenApisTask.class).named(forbiddenApisTaskName).configure(forbiddenApisTask -> {
forbiddenApisTask.setIgnoreMissingClasses(true);
});

return sourceSet;
}

private void configureSourceSetInJar(Project project, SourceSet sourceSet, int javaVersion) {
var jarTask = project.getTasks().withType(Jar.class).named(JavaPlugin.JAR_TASK_NAME);
jarTask.configure(task -> task.into("META-INF/versions/" + javaVersion, copySpec -> copySpec.from(sourceSet.getOutput())));
}

private void createTestTask(Project project, SourceSet sourceSet, int javaVersion, List<String> mainSourceSets) {
var jarTask = project.getTasks().withType(Jar.class).named(JavaPlugin.JAR_TASK_NAME);
var testTaskProvider = project.getTasks().register(JavaPlugin.TEST_TASK_NAME + javaVersion, Test.class);
testTaskProvider.configure(testTask -> {
testTask.dependsOn(jarTask);

SourceSetContainer sourceSets = GradleUtils.getJavaSourceSets(project);
FileCollection testRuntime = sourceSet.getRuntimeClasspath();
for (String mainSourceSetName : mainSourceSets) {
FileCollection mainRuntime = sourceSets.getByName(mainSourceSetName).getOutput();
testRuntime = testRuntime.minus(mainRuntime);
}
testTask.setClasspath(testRuntime.plus(project.files(jarTask)));
testTask.setTestClassesDirs(sourceSet.getOutput().getClassesDirs());

testTask.getJavaLauncher()
.set(javaToolchains.launcherFor(spec -> spec.getLanguageVersion().set(JavaLanguageVersion.of(javaVersion))));
});

project.getTasks().named("check").configure(checkTask -> checkTask.dependsOn(testTaskProvider));
}

private static List<Integer> findSourceVersions(Project project) {
var srcDir = project.getProjectDir().toPath().resolve("src");
List<Integer> versions = new ArrayList<>();
try (var subdirStream = Files.list(srcDir)) {
for (Path sourceSetPath : subdirStream.toList()) {
assert Files.isDirectory(sourceSetPath);
String sourcesetName = sourceSetPath.getFileName().toString();
Matcher sourcesetMatcher = MRJAR_SOURCESET_PATTERN.matcher(sourcesetName);
if (sourcesetMatcher.matches()) {
versions.add(Integer.parseInt(sourcesetMatcher.group(1)));
}
}
} catch (IOException e) {
throw new UncheckedIOException(e);
}

Collections.sort(versions);
return versions;
}

private static void stripPreviewFromFiles(Path compileDir) {
Expand Down
37 changes: 37 additions & 0 deletions ...l/src/main/java/org/elasticsearch/gradle/internal/testfixtures/TestFixtureDeployment.java
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,37 @@
/*
* 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 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 or the Server
* Side Public License, v 1.
*/

package org.elasticsearch.gradle.internal.testfixtures;

import org.gradle.api.Named;
import org.gradle.api.provider.ListProperty;
import org.gradle.api.provider.Property;

import java.io.File;

public abstract class TestFixtureDeployment implements Named {

private final String name;

public TestFixtureDeployment(String name) {
this.name = name;
}

@Override
public String getName() {
return name;
}

public abstract Property<String> getDockerRegistry();

public abstract Property<File> getDockerContext();

public abstract Property<String> getVersion();

public abstract ListProperty<String> getBaseImages();
}
Loading

0 comments on commit f58e7ef

Please sign in to comment.