Skip to content

Commit

Permalink
Add support for RowType to BenchmarkDataGenerator and benchmarks
Browse files Browse the repository at this point in the history
  • Loading branch information
@sopel39
rkondziolka authored and sopel39 committed Feb 23, 2022
1 parent 083035d commit e94383d
Show file tree
Hide file tree
Showing 2 changed files with 97 additions and 29 deletions.
102 changes: 73 additions & 29 deletions core/trino-main/src/test/java/io/trino/execution/buffer/BenchmarkBlockSerde.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
import io.trino.spi.type.DecimalType;
import io.trino.spi.type.Decimals;
import io.trino.spi.type.Int128;
import io.trino.spi.type.RowType;
import io.trino.spi.type.SqlDecimal;
import io.trino.spi.type.Type;
import io.trino.spi.type.VarcharType;
Expand All @@ -42,7 +43,9 @@
import org.openjdk.jmh.annotations.State;
import org.openjdk.jmh.annotations.Warmup;
import org.testng.annotations.Test;
import oshi.util.tuples.Pair;

import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.Random;
Expand Down Expand Up @@ -176,6 +179,18 @@ public Object deserializeLineitem(LineitemBenchmarkData data)
return ImmutableList.copyOf(readPages(data.getPagesSerde(), new BasicSliceInput(data.getDataSource())));
}

@Benchmark
public Object serializeRow(RowTypeBenchmarkData data)
{
return serializePages(data);
}

@Benchmark
public Object deserializeRow(RowTypeBenchmarkData data)
{
return ImmutableList.copyOf(readPages(data.getPagesSerde(), new BasicSliceInput(data.getDataSource())));
}

private static List<Slice> serializePages(BenchmarkData data)
{
PagesSerdeContext context = new PagesSerdeContext();
Expand All @@ -199,35 +214,7 @@ public void setup(Type type, Function<Random, ?> valueGenerator)

Iterator<?> values = createValues(ROWS, valueGenerator, nullChance);
while (values.hasNext()) {
Object value = values.next();
if (value == null) {
blockBuilder.appendNull();
}
else if (BIGINT.equals(type)) {
BIGINT.writeLong(blockBuilder, ((Number) value).longValue());
}
else if (Decimals.isLongDecimal(type)) {
type.writeObject(blockBuilder, Int128.valueOf(((SqlDecimal) value).toBigDecimal().unscaledValue()));
}
else if (type instanceof VarcharType) {
Slice slice = truncateToLength(utf8Slice((String) value), type);
type.writeSlice(blockBuilder, slice);
}
else if (TIMESTAMP_PICOS.equals(type)) {
TIMESTAMP_PICOS.writeObject(blockBuilder, value);
}
else if (INTEGER.equals(type)) {
blockBuilder.writeInt((int) value);
}
else if (SMALLINT.equals(type)) {
blockBuilder.writeShort((short) value);
}
else if (TINYINT.equals(type)) {
blockBuilder.writeByte((byte) value);
}
else {
throw new IllegalArgumentException("Unsupported type " + type);
}
writeValue(type, values.next(), blockBuilder);
pageBuilder.declarePosition();
if (pageBuilder.isFull()) {
pagesBuilder.add(pageBuilder.build());
Expand All @@ -245,6 +232,51 @@ else if (TINYINT.equals(type)) {

setup(sliceOutput.slice(), pagesSerde, pages);
}

private void writeValue(Type type, Object value, BlockBuilder blockBuilder)
{
if (value == null) {
blockBuilder.appendNull();
}
else if (BIGINT.equals(type)) {
BIGINT.writeLong(blockBuilder, ((Number) value).longValue());
}
else if (Decimals.isLongDecimal(type)) {
type.writeObject(blockBuilder, Int128.valueOf(((SqlDecimal) value).toBigDecimal().unscaledValue()));
}
else if (type instanceof VarcharType) {
Slice slice = truncateToLength(utf8Slice((String) value), type);
type.writeSlice(blockBuilder, slice);
}
else if (TIMESTAMP_PICOS.equals(type)) {
TIMESTAMP_PICOS.writeObject(blockBuilder, value);
}
else if (INTEGER.equals(type)) {
blockBuilder.writeInt((int) value);
}
else if (SMALLINT.equals(type)) {
blockBuilder.writeShort((short) value);
}
else if (TINYINT.equals(type)) {
blockBuilder.writeByte((byte) value);
}
else if (type instanceof RowType) {
BlockBuilder row = blockBuilder.beginBlockEntry();
List<?> values = (List<?>) value;
if (values.size() != type.getTypeParameters().size()) {
throw new IllegalArgumentException("Size of types and values must have the same size");
}
List<Pair<Type, Object>> pairs = new ArrayList<>();
for (int i = 0; i < type.getTypeParameters().size(); i++) {
pairs.add(new Pair<Type, Object>(type.getTypeParameters().get(i), ((List<?>) value).get(i)));
}
pairs.forEach(p -> writeValue(p.getA(), p.getB(), row));
blockBuilder.closeEntry();
}
else {
throw new IllegalArgumentException("Unsupported type " + type);
}
}
}

public abstract static class BenchmarkData
Expand Down Expand Up @@ -369,6 +401,18 @@ public void setup()
}
}

@State(Thread)
public static class RowTypeBenchmarkData
extends TypeBenchmarkData
{
@Setup
public void setup()
{
RowType type = RowType.anonymous(ImmutableList.of(BIGINT));
super.setup(type, (random -> BenchmarkDataGenerator.randomRow(type.getTypeParameters(), random)));
}
}

@Test
public void test()
{
Expand Down
24 changes: 24 additions & 0 deletions core/trino-main/src/test/java/io/trino/execution/buffer/BenchmarkDataGenerator.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,8 +13,10 @@
*/
package io.trino.execution.buffer;

import io.trino.spi.type.DecimalType;
import io.trino.spi.type.LongTimestamp;
import io.trino.spi.type.SqlDecimal;
import io.trino.spi.type.Type;

import java.math.BigInteger;
import java.util.ArrayList;
Expand All @@ -23,7 +25,9 @@
import java.util.Random;
import java.util.function.Function;

import static io.trino.spi.type.BigintType.BIGINT;
import static io.trino.spi.type.Timestamps.PICOSECONDS_PER_MICROSECOND;
import static io.trino.spi.type.VarcharType.VARCHAR;

public class BenchmarkDataGenerator
{
Expand Down Expand Up @@ -93,4 +97,24 @@ public static byte randomByte(Random random)
{
return (byte) random.nextInt();
}

public static List<Object> randomRow(List<Type> fieldTypes, Random random)
{
List<Object> row = new ArrayList<>(fieldTypes.size());
for (Type type : fieldTypes) {
if (type == VARCHAR) {
row.add(randomAsciiString(random));
}
else if (type == BIGINT) {
row.add(random.nextLong());
}
else if (type instanceof DecimalType) {
row.add(randomLongDecimal(random));
}
else {
throw new UnsupportedOperationException(String.format("The %s is not supported", type));
}
}
return row;
}
}

0 comments on commit e94383d

Please sign in to comment.