Skip to content

Reduce generated code size for Row Constructors #21299

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Apr 10, 2024
Merged

Reduce generated code size for Row Constructors #21299

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
75 changes: 75 additions & 0 deletions core/trino-main/src/main/java/io/trino/sql/gen/RowConstructorCodeGenerator.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,30 +18,38 @@
import io.airlift.bytecode.Scope;
import io.airlift.bytecode.Variable;
import io.airlift.bytecode.control.IfStatement;
import io.trino.annotation.UsedByGeneratedCode;
import io.trino.spi.block.Block;
import io.trino.spi.block.BlockBuilder;
import io.trino.spi.block.BlockBuilderStatus;
import io.trino.spi.block.SqlRow;
import io.trino.spi.type.RowType;
import io.trino.spi.type.Type;
import io.trino.sql.relational.RowExpression;
import io.trino.sql.relational.SpecialForm;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

import static io.airlift.bytecode.ParameterizedType.type;
import static io.airlift.bytecode.expression.BytecodeExpressions.constantFalse;
import static io.airlift.bytecode.expression.BytecodeExpressions.constantInt;
import static io.airlift.bytecode.expression.BytecodeExpressions.constantNull;
import static io.airlift.bytecode.expression.BytecodeExpressions.invokeStatic;
import static io.airlift.bytecode.expression.BytecodeExpressions.newArray;
import static io.airlift.bytecode.expression.BytecodeExpressions.newInstance;
import static io.trino.sql.gen.SqlTypeBytecodeExpression.constantType;
import static java.lang.String.format;
import static java.util.Objects.requireNonNull;

public class RowConstructorCodeGenerator
implements BytecodeGenerator
{
private final Type rowType;
private final List<RowExpression> arguments;
// Arbitrary value chosen to balance the code size vs performance trade off. Not perf tested.
private static final int MEGAMORPHIC_FIELD_COUNT = 64;
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would assume this value is arbitrary and not perf tested. If so, we should comment here, so people don't assume 64 is the best value.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, the value is taken from a constant by the same name in RowType, but unlike that variation was not specifically arrived at by testing for performance but rather was chosen based on seeming reasonable to balance the code size vs performance trade off (ie: mostly arbitrary). We can add a comment to that effect.


public RowConstructorCodeGenerator(SpecialForm specialForm)
{
Expand All @@ -53,6 +61,10 @@ public RowConstructorCodeGenerator(SpecialForm specialForm)
@Override
public BytecodeNode generateExpression(BytecodeGeneratorContext context)
{
if (arguments.size() > MEGAMORPHIC_FIELD_COUNT) {
return generateExpressionForLargeRows(context);
}

BytecodeBlock block = new BytecodeBlock().setDescription("Constructor for " + rowType);
CallSiteBinder binder = context.getCallSiteBinder();
Scope scope = context.getScope();
Expand Down Expand Up @@ -88,4 +100,67 @@ public BytecodeNode generateExpression(BytecodeGeneratorContext context)
block.append(context.wasNull().set(constantFalse()));
return block;
}

// Avoids inline BlockBuilder and Block creation logic for each field which reduces the generated code size
// for RowTypes with many fields significantly, but does so at the cost of virtual method dispatch
private BytecodeNode generateExpressionForLargeRows(BytecodeGeneratorContext context)
{
BytecodeBlock block = new BytecodeBlock().setDescription("Constructor for " + rowType);
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you add a comment on how this is different from the above code. Specifically touch on where we get less byte code and why the performance is not as good as above (otherwise we would use the same code for both).

CallSiteBinder binder = context.getCallSiteBinder();
Scope scope = context.getScope();
List<Type> types = rowType.getTypeParameters();

Variable fieldBuilders = scope.createTempVariable(BlockBuilder[].class);
block.append(fieldBuilders.set(invokeStatic(RowConstructorCodeGenerator.class, "createFieldBlockBuildersForSingleRow", BlockBuilder[].class, constantType(binder, rowType))));

// Cache local variable declarations per java type on stack for reuse
Map<Class<?>, Variable> javaTypeTempVariables = new HashMap<>();
Variable blockBuilder = scope.createTempVariable(BlockBuilder.class);
for (int i = 0; i < arguments.size(); ++i) {
Type fieldType = types.get(i);
Variable field = javaTypeTempVariables.computeIfAbsent(fieldType.getJavaType(), scope::createTempVariable);

block.append(blockBuilder.set(fieldBuilders.getElement(constantInt(i))));

block.comment("Clean wasNull and Generate + " + i + "-th field of row");
block.append(context.wasNull().set(constantFalse()));
block.append(context.generate(arguments.get(i)));
block.putVariable(field);
block.append(new IfStatement()
.condition(context.wasNull())
.ifTrue(blockBuilder.invoke("appendNull", BlockBuilder.class).pop())
.ifFalse(constantType(binder, fieldType).writeValue(blockBuilder, field).pop()));
}

block.append(invokeStatic(RowConstructorCodeGenerator.class, "createSqlRowFromFieldBuildersForSingleRow", SqlRow.class, fieldBuilders));
block.append(context.wasNull().set(constantFalse()));
return block;
}

@UsedByGeneratedCode
public static BlockBuilder[] createFieldBlockBuildersForSingleRow(Type rowType)
{
if (!(rowType instanceof RowType)) {
throw new IllegalArgumentException("Not a row type: " + rowType);
}
List<Type> fieldTypes = rowType.getTypeParameters();
BlockBuilder[] fieldBlockBuilders = new BlockBuilder[fieldTypes.size()];
for (int i = 0; i < fieldTypes.size(); i++) {
fieldBlockBuilders[i] = fieldTypes.get(i).createBlockBuilder(null, 1);
}
return fieldBlockBuilders;
}

@UsedByGeneratedCode
public static SqlRow createSqlRowFromFieldBuildersForSingleRow(BlockBuilder[] fieldBuilders)
{
Block[] fieldBlocks = new Block[fieldBuilders.length];
for (int i = 0; i < fieldBuilders.length; i++) {
fieldBlocks[i] = fieldBuilders[i].build();
if (fieldBlocks[i].getPositionCount() != 1) {
throw new IllegalArgumentException(format("builder must only contain a single position, found: %s positions", fieldBlocks[i].getPositionCount()));
}
}
return new SqlRow(0, fieldBlocks);
}
}
48 changes: 47 additions & 1 deletion ...in/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergConnectorSmokeTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,8 @@
import io.trino.plugin.hive.metastore.HiveMetastore;
import io.trino.plugin.hive.metastore.HiveMetastoreFactory;
import io.trino.testing.QueryRunner;
import io.trino.testing.sql.TestTable;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.TestInstance;

import java.io.IOException;
Expand All @@ -26,6 +28,9 @@

import static io.trino.plugin.iceberg.IcebergQueryRunner.ICEBERG_CATALOG;
import static io.trino.plugin.iceberg.IcebergTestUtils.checkOrcFileSorting;
import static io.trino.tpch.TpchTable.NATION;
import static io.trino.tpch.TpchTable.ORDERS;
import static io.trino.tpch.TpchTable.REGION;
import static org.apache.iceberg.FileFormat.ORC;
import static org.assertj.core.api.Assertions.assertThat;
import static org.junit.jupiter.api.TestInstance.Lifecycle.PER_CLASS;
Expand All @@ -48,7 +53,7 @@ protected QueryRunner createQueryRunner()
throws Exception
{
QueryRunner queryRunner = IcebergQueryRunner.builder()
.setInitialTables(REQUIRED_TPCH_TABLES)
.setInitialTables(NATION, ORDERS, REGION)
.setIcebergProperties(ImmutableMap.of(
"iceberg.file-format", format.name(),
"iceberg.register-table-procedure.enabled", "true",
Expand Down Expand Up @@ -108,4 +113,45 @@ protected boolean isFileSorted(Location path, String sortColumnName)
{
return checkOrcFileSorting(fileSystem, path, sortColumnName);
}

@Test
public void testRowConstructorColumnLimitForMergeQuery()
{
String[] colNames = {"orderkey", "custkey", "orderstatus", "totalprice", "orderpriority", "clerk", "shippriority", "comment", "orderdate"};
String[] colTypes = {"bigint", "bigint", "varchar", "decimal(12,2)", "varchar", "varchar", "int", "varchar", "date"};
String tableDefinition = "(";
String columns = "(";
String selectQuery = "select ";
String notMatchedClause = "";
String matchedClause = "";
// Creating merge query with 325 columns
for (int i = 0; i < 36; i++) {
for (int j = 0; j < 9; j++) {
String columnName = colNames[j];
String columnType = colTypes[j];
tableDefinition += columnName + "_" + i + " " + columnType + ",";
selectQuery += columnName + " " + columnName + "_" + i + ",";
columns += columnName + "_" + i + ",";
notMatchedClause += "s." + columnName + "_" + i + ",";
matchedClause += columnName + "_" + i + " = s." + columnName + "_" + i + ",";
}
}
tableDefinition += "orderkey bigint, custkey bigint, orderstatus varchar, totalprice decimal(12,2), orderpriority varchar) ";
selectQuery += "orderkey, custkey, orderstatus, totalprice, orderpriority from orders limit 1 ";
columns += "orderkey, custkey, orderstatus, totalprice, orderpriority) ";
notMatchedClause += "s.orderkey, s.custkey, s.orderstatus, s.totalprice, s.orderpriority ";
matchedClause += "orderkey = s.orderkey, custkey = s.custkey, orderstatus = s.orderstatus, totalprice = t.totalprice, orderpriority = s.orderpriority ";
TestTable table = new TestTable(getQueryRunner()::execute, "test_merge_", tableDefinition);
assertUpdate("INSERT INTO " + table.getName() + " " + columns + " " + selectQuery, 1);
TestTable mergeTable = new TestTable(getQueryRunner()::execute, "test_table_", tableDefinition);
assertUpdate("INSERT INTO " + mergeTable.getName() + " " + columns + " " + selectQuery, 1);
assertUpdate("""
MERGE INTO %s t
USING (select * from %s ) s
ON (t.orderkey = s.orderkey)
WHEN MATCHED THEN UPDATE SET %s
WHEN NOT MATCHED THEN INSERT VALUES (%s)
""".formatted(mergeTable.getName(), table.getName(), matchedClause, notMatchedClause),
1);
}
}
10 changes: 10 additions & 0 deletions ...ng/trino-testing/src/main/java/io/trino/testing/AbstractDistributedEngineOnlyQueries.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -393,4 +393,14 @@ public void testSelectiveLimit()
"LIMIT 1",
"VALUES -1");
}

@Test
public void testRowConstructorColumnLimit()
{
// Generate a query with 859 columns: SELECT row(col1, col2, ....col859) from t
String colNames = "orderkey, custkey, orderstatus, totalprice, orderpriority, clerk, shippriority, comment, orderdate";
String rowFields = colNames + (", " + colNames).repeat(94) + ", orderkey, custkey, orderstatus, totalprice";
@Language("SQL") String query = "SELECT row(" + rowFields + ") FROM (select * from tpch.tiny.orders limit 1) t(" + colNames + ")";
assertThat(getQueryRunner().execute(query).getOnlyValue()).isNotNull();
}
}