Merge branch 'main' into esql_named_writeable_exchange

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

@@ -0,0 +1,295 @@
+/*
+ * 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.benchmark.compute.operator;
+
+import org.apache.lucene.util.BytesRef;
+import org.elasticsearch.compute.data.Block;
+import org.elasticsearch.compute.data.BooleanBlock;
+import org.elasticsearch.compute.data.BooleanVector;
+import org.elasticsearch.compute.data.BytesRefBlock;
+import org.elasticsearch.compute.data.DoubleBlock;
+import org.elasticsearch.compute.data.IntBlock;
+import org.elasticsearch.compute.data.LongBlock;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Level;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OperationsPerInvocation;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Param;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.TearDown;
+import org.openjdk.jmh.annotations.Warmup;
+
+import java.util.concurrent.TimeUnit;
+
+@Warmup(iterations = 5)
+@Measurement(iterations = 7)
+@BenchmarkMode(Mode.AverageTime)
+@OutputTimeUnit(TimeUnit.NANOSECONDS)
+@State(Scope.Thread)
+@Fork(1)
+public class BlockKeepMaskBenchmark extends BlockBenchmark {
+ static {
+ // Smoke test all the expected values and force loading subclasses more like prod
+ int totalPositions = 10;
+ for (String paramString : RELEVANT_TYPE_BLOCK_COMBINATIONS) {
+ String[] params = paramString.split("/");
+ String dataType = params[0];
+ String blockKind = params[1];
+ BooleanVector mask = buildMask(totalPositions);
+
+ BenchmarkBlocks data = buildBenchmarkBlocks(dataType, blockKind, mask, totalPositions);
+ Block[] results = new Block[NUM_BLOCKS_PER_ITERATION];
+ run(data, mask, results);
+ assertCheckSums(dataType, blockKind, data, results, totalPositions);
+ }
+ }
+
+ record BenchmarkBlocks(Block[] blocks, long[] checkSums) {};
+
+ static BenchmarkBlocks buildBenchmarkBlocks(String dataType, String blockKind, BooleanVector mask, int totalPositions) {
+ Block[] blocks = BlockBenchmark.buildBlocks(dataType, blockKind, totalPositions);
+ return new BenchmarkBlocks(blocks, checksumsFor(dataType, blocks, mask));
+ }
+
+ static long[] checksumsFor(String dataType, Block[] blocks, BooleanVector mask) {
+ long[] checkSums = new long[NUM_BLOCKS_PER_ITERATION];
+ switch (dataType) {
+ case "boolean" -> {
+ for (int blockIndex = 0; blockIndex < NUM_BLOCKS_PER_ITERATION; blockIndex++) {
+ BooleanBlock block = (BooleanBlock) blocks[blockIndex];
+ checkSums[blockIndex] = computeBooleanCheckSum(block, mask);
+ }
+ }
+ case "BytesRef" -> {
+ for (int blockIndex = 0; blockIndex < NUM_BLOCKS_PER_ITERATION; blockIndex++) {
+ BytesRefBlock block = (BytesRefBlock) blocks[blockIndex];
+ checkSums[blockIndex] = computeBytesRefCheckSum(block, mask);
+ }
+ }
+ case "double" -> {
+ for (int blockIndex = 0; blockIndex < NUM_BLOCKS_PER_ITERATION; blockIndex++) {
+ DoubleBlock block = (DoubleBlock) blocks[blockIndex];
+ checkSums[blockIndex] = computeDoubleCheckSum(block, mask);
+ }
+ }
+ case "int" -> {
+ for (int blockIndex = 0; blockIndex < NUM_BLOCKS_PER_ITERATION; blockIndex++) {
+ IntBlock block = (IntBlock) blocks[blockIndex];
+ checkSums[blockIndex] = computeIntCheckSum(block, mask);
+ }
+ }
+ case "long" -> {
+ for (int blockIndex = 0; blockIndex < NUM_BLOCKS_PER_ITERATION; blockIndex++) {
+ LongBlock block = (LongBlock) blocks[blockIndex];
+ checkSums[blockIndex] = computeLongCheckSum(block, mask);
+ }
+ }
+ // TODO float
+ default -> throw new IllegalStateException("illegal data type [" + dataType + "]");
+ }
+ return checkSums;
+ }
+
+ static BooleanVector buildMask(int totalPositions) {
+ try (BooleanVector.FixedBuilder builder = blockFactory.newBooleanVectorFixedBuilder(totalPositions)) {
+ for (int p = 0; p < totalPositions; p++) {
+ builder.appendBoolean(p % 2 == 0);
+ }
+ return builder.build();
+ }
+ }
+
+ private static void run(BenchmarkBlocks data, BooleanVector mask, Block[] results) {
+ for (int blockIndex = 0; blockIndex < NUM_BLOCKS_PER_ITERATION; blockIndex++) {
+ results[blockIndex] = data.blocks[blockIndex].keepMask(mask);
+ }
+ }
+
+ private static void assertCheckSums(String dataType, String blockKind, BenchmarkBlocks data, Block[] results, int positionCount) {
+ long[] checkSums = checksumsFor(dataType, results, blockFactory.newConstantBooleanVector(true, positionCount));
+ for (int blockIndex = 0; blockIndex < NUM_BLOCKS_PER_ITERATION; blockIndex++) {
+ if (checkSums[blockIndex] != data.checkSums[blockIndex]) {
+ throw new AssertionError(
+ "checksums do not match for block ["
+ + blockIndex
+ + "]["
+ + dataType
+ + "]["
+ + blockKind
+ + "]: "
+ + checkSums[blockIndex]
+ + " vs "
+ + data.checkSums[blockIndex]
+ );
+ }
+ }
+ }
+
+ private static long computeBooleanCheckSum(BooleanBlock block, BooleanVector mask) {
+ long sum = 0;
+
+ for (int p = 0; p < block.getPositionCount(); p++) {
+ if (block.isNull(p) || mask.getBoolean(p) == false) {
+ continue;
+ }
+ int start = block.getFirstValueIndex(p);
+ int end = start + block.getValueCount(p);
+ for (int i = start; i < end; i++) {
+ sum += block.getBoolean(i) ? 1 : 0;
+ }
+ }
+
+ return sum;
+ }
+
+ private static long computeBytesRefCheckSum(BytesRefBlock block, BooleanVector mask) {
+ long sum = 0;
+ BytesRef scratch = new BytesRef();
+
+ for (int p = 0; p < block.getPositionCount(); p++) {
+ if (block.isNull(p) || mask.getBoolean(p) == false) {
+ continue;
+ }
+ int start = block.getFirstValueIndex(p);
+ int end = start + block.getValueCount(p);
+ for (int i = start; i < end; i++) {
+ BytesRef v = block.getBytesRef(i, scratch);
+ sum += v.length > 0 ? v.bytes[v.offset] : 0;
+ }
+ }
+
+ return sum;
+ }
+
+ private static long computeDoubleCheckSum(DoubleBlock block, BooleanVector mask) {
+ long sum = 0;
+
+ for (int p = 0; p < block.getPositionCount(); p++) {
+ if (block.isNull(p) || mask.getBoolean(p) == false) {
+ continue;
+ }
+ int start = block.getFirstValueIndex(p);
+ int end = start + block.getValueCount(p);
+ for (int i = start; i < end; i++) {
+ sum += (long) block.getDouble(i);
+ }
+ }
+
+ return sum;
+ }
+
+ private static long computeIntCheckSum(IntBlock block, BooleanVector mask) {
+ int sum = 0;
+
+ for (int p = 0; p < block.getPositionCount(); p++) {
+ if (block.isNull(p) || mask.getBoolean(p) == false) {
+ continue;
+ }
+ int start = block.getFirstValueIndex(p);
+ int end = start + block.getValueCount(p);
+ for (int i = start; i < end; i++) {
+ sum += block.getInt(i);
+ }
+ }
+
+ return sum;
+ }
+
+ private static long computeLongCheckSum(LongBlock block, BooleanVector mask) {
+ long sum = 0;
+
+ for (int p = 0; p < block.getPositionCount(); p++) {
+ if (block.isNull(p) || mask.getBoolean(p) == false) {
+ continue;
+ }
+ int start = block.getFirstValueIndex(p);
+ int end = start + block.getValueCount(p);
+ for (int i = start; i < end; i++) {
+ sum += block.getLong(i);
+ }
+ }
+
+ return sum;
+ }
+
+ /**
+ * Must be a subset of {@link BlockBenchmark#RELEVANT_TYPE_BLOCK_COMBINATIONS}
+ */
+ @Param(
+ {
+ "boolean/array",
+ "boolean/array-multivalue-null",
+ "boolean/big-array",
+ "boolean/big-array-multivalue-null",
+ "boolean/vector",
+ "boolean/vector-big-array",
+ "boolean/vector-const",
+ "BytesRef/array",
+ "BytesRef/array-multivalue-null",
+ "BytesRef/vector",
+ "BytesRef/vector-const",
+ "double/array",
+ "double/array-multivalue-null",
+ "double/big-array",
+ "double/big-array-multivalue-null",
+ "double/vector",
+ "double/vector-big-array",
+ "double/vector-const",
+ "int/array",
+ "int/array-multivalue-null",
+ "int/big-array",
+ "int/big-array-multivalue-null",
+ "int/vector",
+ "int/vector-big-array",
+ "int/vector-const",
+ "long/array",
+ "long/array-multivalue-null",
+ "long/big-array",
+ "long/big-array-multivalue-null",
+ "long/vector",
+ "long/vector-big-array",
+ "long/vector-const" }
+ )
+ public String dataTypeAndBlockKind;
+
+ private BenchmarkBlocks data;
+
+ private final BooleanVector mask = buildMask(BLOCK_TOTAL_POSITIONS);
+
+ private final Block[] results = new Block[NUM_BLOCKS_PER_ITERATION];
+
+ @Setup
+ public void setup() {
+ String[] params = dataTypeAndBlockKind.split("/");
+ String dataType = params[0];
+ String blockKind = params[1];
+
+ data = buildBenchmarkBlocks(dataType, blockKind, mask, BLOCK_TOTAL_POSITIONS);
+ }
+
+ @Benchmark
+ @OperationsPerInvocation(NUM_BLOCKS_PER_ITERATION * BLOCK_TOTAL_POSITIONS)
+ public void run() {
+ run(data, mask, results);
+ }
+
+ @TearDown(Level.Iteration)
+ public void assertCheckSums() {
+ String[] params = dataTypeAndBlockKind.split("/");
+ String dataType = params[0];
+ String blockKind = params[1];
+ assertCheckSums(dataType, blockKind, data, results, BLOCK_TOTAL_POSITIONS);
+ }
+}