From 0c37357b2f801747dee084cb354757f3b36eb1ea Mon Sep 17 00:00:00 2001 From: Yuya Ebihara Date: Tue, 25 Jul 2023 09:40:42 +0900 Subject: [PATCH 1/2] Support setting a field type with SET DATA TYPE in engine --- .../io/trino/execution/SetColumnTypeTask.java | 50 +++++++++++++- .../main/java/io/trino/metadata/Metadata.java | 5 ++ .../io/trino/metadata/MetadataManager.java | 8 +++ .../tracing/TracingConnectorMetadata.java | 9 +++ .../io/trino/tracing/TracingMetadata.java | 9 +++ .../io/trino/connector/MockConnector.java | 6 ++ .../execution/TestSetColumnTypeTask.java | 67 ++++++++++++++++--- .../trino/metadata/AbstractMockMetadata.java | 6 ++ .../antlr4/io/trino/sql/parser/SqlBase.g4 | 2 +- .../java/io/trino/sql/parser/AstBuilder.java | 2 +- .../java/io/trino/sql/tree/SetColumnType.java | 8 +-- .../io/trino/sql/parser/TestSqlParser.java | 4 +- .../spi/connector/ConnectorMetadata.java | 11 +++ .../ClassLoaderSafeConnectorMetadata.java | 8 +++ 14 files changed, 176 insertions(+), 19 deletions(-) diff --git a/core/trino-main/src/main/java/io/trino/execution/SetColumnTypeTask.java b/core/trino-main/src/main/java/io/trino/execution/SetColumnTypeTask.java index 8170fe7be0e80..946efa148a5a5 100644 --- a/core/trino-main/src/main/java/io/trino/execution/SetColumnTypeTask.java +++ b/core/trino-main/src/main/java/io/trino/execution/SetColumnTypeTask.java @@ -22,7 +22,10 @@ import io.trino.metadata.RedirectionAwareTableHandle; import io.trino.metadata.TableHandle; import io.trino.security.AccessControl; +import io.trino.spi.TrinoException; import io.trino.spi.connector.ColumnHandle; +import io.trino.spi.connector.ColumnMetadata; +import io.trino.spi.type.RowType; import io.trino.spi.type.Type; import io.trino.spi.type.TypeManager; import io.trino.spi.type.TypeNotFoundException; @@ -32,10 +35,15 @@ import java.util.List; import java.util.Map; +import static com.google.common.base.Preconditions.checkState; +import static com.google.common.collect.ImmutableList.toImmutableList; +import static com.google.common.collect.Iterables.getOnlyElement; import static com.google.common.util.concurrent.Futures.immediateVoidFuture; import static io.trino.metadata.MetadataUtil.createQualifiedObjectName; +import static io.trino.spi.StandardErrorCode.AMBIGUOUS_NAME; import static io.trino.spi.StandardErrorCode.COLUMN_NOT_FOUND; import static io.trino.spi.StandardErrorCode.COLUMN_TYPE_UNKNOWN; +import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; import static io.trino.spi.StandardErrorCode.TABLE_NOT_FOUND; import static io.trino.spi.StandardErrorCode.TYPE_NOT_FOUND; import static io.trino.sql.analyzer.SemanticExceptions.semanticException; @@ -94,16 +102,54 @@ else if (metadata.getView(session, qualifiedObjectName).isPresent()) { TableHandle tableHandle = redirectionAwareTableHandle.getTableHandle().get(); Map columnHandles = metadata.getColumnHandles(session, tableHandle); - ColumnHandle column = columnHandles.get(statement.getColumnName().getValue().toLowerCase(ENGLISH)); + String columnName = statement.getColumnName().getParts().get(0).toLowerCase(ENGLISH); + ColumnHandle column = columnHandles.get(columnName); if (column == null) { throw semanticException(COLUMN_NOT_FOUND, statement, "Column '%s' does not exist", statement.getColumnName()); } - metadata.setColumnType(session, tableHandle, column, getColumnType(statement)); + Type type = getColumnType(statement); + if (statement.getColumnName().getParts().size() == 1) { + metadata.setColumnType(session, tableHandle, column, type); + } + else { + ColumnMetadata columnMetadata = metadata.getColumnMetadata(session, tableHandle, column); + List fieldPath = statement.getColumnName().getParts(); + + Type currentType = columnMetadata.getType(); + for (int i = 1; i < fieldPath.size(); i++) { + String fieldName = fieldPath.get(i); + List candidates = getCandidates(currentType, fieldName); + + if (candidates.isEmpty()) { + throw semanticException(COLUMN_NOT_FOUND, statement, "Field '%s' does not exist within %s", fieldName, currentType); + } + if (candidates.size() > 1) { + throw semanticException(AMBIGUOUS_NAME, statement, "Field path %s within %s is ambiguous", fieldPath, columnMetadata.getType()); + } + currentType = getOnlyElement(candidates).getType(); + } + + checkState(fieldPath.size() >= 2, "fieldPath size must be >= 2: %s", fieldPath); + metadata.setFieldType(session, tableHandle, fieldPath, type); + } return immediateVoidFuture(); } + private static List getCandidates(Type type, String fieldName) + { + if (!(type instanceof RowType rowType)) { + throw new TrinoException(NOT_SUPPORTED, "Unsupported type: " + type); + } + List candidates = rowType.getFields().stream() + // case-insensitive match + .filter(rowField -> rowField.getName().isPresent() && rowField.getName().get().equalsIgnoreCase(fieldName)) + .collect(toImmutableList()); + + return candidates; + } + private Type getColumnType(SetColumnType statement) { Type type; diff --git a/core/trino-main/src/main/java/io/trino/metadata/Metadata.java b/core/trino-main/src/main/java/io/trino/metadata/Metadata.java index 004139f67f8f9..28da6913c2664 100644 --- a/core/trino-main/src/main/java/io/trino/metadata/Metadata.java +++ b/core/trino-main/src/main/java/io/trino/metadata/Metadata.java @@ -257,6 +257,11 @@ Optional getTableHandleForExecute( */ void setColumnType(Session session, TableHandle tableHandle, ColumnHandle column, Type type); + /** + * Set the specified type to the field. + */ + void setFieldType(Session session, TableHandle tableHandle, List fieldPath, Type type); + /** * Set the authorization (owner) of specified table's user/role */ diff --git a/core/trino-main/src/main/java/io/trino/metadata/MetadataManager.java b/core/trino-main/src/main/java/io/trino/metadata/MetadataManager.java index 584399c2d98ae..4226ba7054896 100644 --- a/core/trino-main/src/main/java/io/trino/metadata/MetadataManager.java +++ b/core/trino-main/src/main/java/io/trino/metadata/MetadataManager.java @@ -814,6 +814,14 @@ public void setColumnType(Session session, TableHandle tableHandle, ColumnHandle metadata.setColumnType(session.toConnectorSession(catalogHandle), tableHandle.getConnectorHandle(), column, type); } + @Override + public void setFieldType(Session session, TableHandle tableHandle, List fieldPath, Type type) + { + CatalogHandle catalogHandle = tableHandle.getCatalogHandle(); + ConnectorMetadata metadata = getMetadataForWrite(session, catalogHandle); + metadata.setFieldType(session.toConnectorSession(catalogHandle), tableHandle.getConnectorHandle(), fieldPath, type); + } + @Override public void setTableAuthorization(Session session, CatalogSchemaTableName table, TrinoPrincipal principal) { diff --git a/core/trino-main/src/main/java/io/trino/tracing/TracingConnectorMetadata.java b/core/trino-main/src/main/java/io/trino/tracing/TracingConnectorMetadata.java index b56338301f90e..0a62aca6b15cb 100644 --- a/core/trino-main/src/main/java/io/trino/tracing/TracingConnectorMetadata.java +++ b/core/trino-main/src/main/java/io/trino/tracing/TracingConnectorMetadata.java @@ -459,6 +459,15 @@ public void setColumnType(ConnectorSession session, ConnectorTableHandle tableHa } } + @Override + public void setFieldType(ConnectorSession session, ConnectorTableHandle tableHandle, List fieldPath, Type type) + { + Span span = startSpan("setFieldType", tableHandle); + try (var ignored = scopedSpan(span)) { + delegate.setFieldType(session, tableHandle, fieldPath, type); + } + } + @Override public void setTableAuthorization(ConnectorSession session, SchemaTableName tableName, TrinoPrincipal principal) { diff --git a/core/trino-main/src/main/java/io/trino/tracing/TracingMetadata.java b/core/trino-main/src/main/java/io/trino/tracing/TracingMetadata.java index 8a4c64cc4afd4..13a903e9ca6f7 100644 --- a/core/trino-main/src/main/java/io/trino/tracing/TracingMetadata.java +++ b/core/trino-main/src/main/java/io/trino/tracing/TracingMetadata.java @@ -476,6 +476,15 @@ public void setColumnType(Session session, TableHandle tableHandle, ColumnHandle } } + @Override + public void setFieldType(Session session, TableHandle tableHandle, List fieldPath, Type type) + { + Span span = startSpan("setFieldType", tableHandle); + try (var ignored = scopedSpan(span)) { + delegate.setFieldType(session, tableHandle, fieldPath, type); + } + } + @Override public void setTableAuthorization(Session session, CatalogSchemaTableName table, TrinoPrincipal principal) { diff --git a/core/trino-main/src/test/java/io/trino/connector/MockConnector.java b/core/trino-main/src/test/java/io/trino/connector/MockConnector.java index d37384791b713..cd84f44afce2c 100644 --- a/core/trino-main/src/test/java/io/trino/connector/MockConnector.java +++ b/core/trino-main/src/test/java/io/trino/connector/MockConnector.java @@ -596,6 +596,12 @@ public void addColumn(ConnectorSession session, ConnectorTableHandle tableHandle @Override public void setColumnType(ConnectorSession session, ConnectorTableHandle tableHandle, ColumnHandle column, Type type) {} + @Override + public void setFieldType(ConnectorSession session, ConnectorTableHandle tableHandle, List fieldPath, Type type) + { + throw new UnsupportedOperationException(); + } + @Override public void setTableAuthorization(ConnectorSession session, SchemaTableName tableName, TrinoPrincipal principal) {} diff --git a/core/trino-main/src/test/java/io/trino/execution/TestSetColumnTypeTask.java b/core/trino-main/src/test/java/io/trino/execution/TestSetColumnTypeTask.java index 338effaaf3196..7106a491418c8 100644 --- a/core/trino-main/src/test/java/io/trino/execution/TestSetColumnTypeTask.java +++ b/core/trino-main/src/test/java/io/trino/execution/TestSetColumnTypeTask.java @@ -20,18 +20,24 @@ import io.trino.metadata.TableHandle; import io.trino.security.AllowAllAccessControl; import io.trino.spi.connector.ColumnMetadata; +import io.trino.spi.connector.ConnectorTableMetadata; +import io.trino.spi.type.RowType; +import io.trino.spi.type.RowType.Field; import io.trino.sql.tree.DataType; -import io.trino.sql.tree.Identifier; import io.trino.sql.tree.NodeLocation; import io.trino.sql.tree.QualifiedName; import io.trino.sql.tree.SetColumnType; import org.testng.annotations.Test; +import java.util.Optional; + import static io.airlift.concurrent.MoreFutures.getFutureValue; +import static io.trino.spi.StandardErrorCode.AMBIGUOUS_NAME; import static io.trino.spi.StandardErrorCode.COLUMN_NOT_FOUND; import static io.trino.spi.StandardErrorCode.TABLE_NOT_FOUND; import static io.trino.spi.type.BigintType.BIGINT; import static io.trino.spi.type.IntegerType.INTEGER; +import static io.trino.spi.type.RowType.rowType; import static io.trino.sql.analyzer.TypeSignatureTranslator.toSqlType; import static io.trino.testing.TestingHandles.TEST_CATALOG_NAME; import static io.trino.testing.assertions.TrinoExceptionAssert.assertTrinoExceptionThrownBy; @@ -51,12 +57,12 @@ public void testSetDataType() .isEqualTo(ImmutableList.of(new ColumnMetadata("test", BIGINT))); // Change the column type to integer from bigint - getFutureValue(executeSetColumnType(asQualifiedName(tableName), new Identifier("test"), toSqlType(INTEGER), false)); + getFutureValue(executeSetColumnType(asQualifiedName(tableName), QualifiedName.of("test"), toSqlType(INTEGER), false)); assertThat(metadata.getTableMetadata(testSession, table).getColumns()) .isEqualTo(ImmutableList.of(new ColumnMetadata("test", INTEGER))); // Specify the same column type - getFutureValue(executeSetColumnType(asQualifiedName(tableName), new Identifier("test"), toSqlType(INTEGER), false)); + getFutureValue(executeSetColumnType(asQualifiedName(tableName), QualifiedName.of("test"), toSqlType(INTEGER), false)); assertThat(metadata.getTableMetadata(testSession, table).getColumns()) .isEqualTo(ImmutableList.of(new ColumnMetadata("test", INTEGER))); } @@ -66,7 +72,7 @@ public void testSetDataTypeNotExistingTable() { QualifiedObjectName tableName = qualifiedObjectName("not_existing_table"); - assertTrinoExceptionThrownBy(() -> getFutureValue(executeSetColumnType(asQualifiedName(tableName), new Identifier("test"), toSqlType(INTEGER), false))) + assertTrinoExceptionThrownBy(() -> getFutureValue(executeSetColumnType(asQualifiedName(tableName), QualifiedName.of("test"), toSqlType(INTEGER), false))) .hasErrorCode(TABLE_NOT_FOUND) .hasMessageContaining("Table '%s' does not exist", tableName); } @@ -76,7 +82,7 @@ public void testSetDataTypeNotExistingTableIfExists() { QualifiedName tableName = qualifiedName("not_existing_table"); - getFutureValue(executeSetColumnType(tableName, new Identifier("test"), toSqlType(INTEGER), true)); + getFutureValue(executeSetColumnType(tableName, QualifiedName.of("test"), toSqlType(INTEGER), true)); // no exception } @@ -84,7 +90,7 @@ public void testSetDataTypeNotExistingTableIfExists() public void testSetDataTypeNotExistingColumn() { QualifiedObjectName tableName = qualifiedObjectName("existing_table"); - Identifier columnName = new Identifier("not_existing_column"); + QualifiedName columnName = QualifiedName.of("not_existing_column"); metadata.createTable(testSession, TEST_CATALOG_NAME, someTable(tableName), false); assertTrinoExceptionThrownBy(() -> getFutureValue(executeSetColumnType(asQualifiedName(tableName), columnName, toSqlType(INTEGER), false))) @@ -98,7 +104,7 @@ public void testSetDataTypeOnView() QualifiedObjectName viewName = qualifiedObjectName("existing_view"); metadata.createView(testSession, viewName, someView(), false); - assertTrinoExceptionThrownBy(() -> getFutureValue(executeSetColumnType(asQualifiedName(viewName), new Identifier("test"), toSqlType(INTEGER), false))) + assertTrinoExceptionThrownBy(() -> getFutureValue(executeSetColumnType(asQualifiedName(viewName), QualifiedName.of("test"), toSqlType(INTEGER), false))) .hasErrorCode(TABLE_NOT_FOUND) .hasMessageContaining("Table '%s' does not exist, but a view with that name exists.", viewName); } @@ -109,12 +115,55 @@ public void testSetDataTypeOnMaterializedView() QualifiedObjectName materializedViewName = qualifiedObjectName("existing_materialized_view"); metadata.createMaterializedView(testSession, QualifiedObjectName.valueOf(materializedViewName.toString()), someMaterializedView(), false, false); - assertTrinoExceptionThrownBy(() -> getFutureValue(executeSetColumnType(asQualifiedName(materializedViewName), new Identifier("test"), toSqlType(INTEGER), false))) + assertTrinoExceptionThrownBy(() -> getFutureValue(executeSetColumnType(asQualifiedName(materializedViewName), QualifiedName.of("test"), toSqlType(INTEGER), false))) .hasErrorCode(TABLE_NOT_FOUND) .hasMessageContaining("Table '%s' does not exist, but a materialized view with that name exists.", materializedViewName); } - private ListenableFuture executeSetColumnType(QualifiedName table, Identifier column, DataType type, boolean exists) + @Test + public void testSetFieldDataTypeNotExistingColumn() + { + QualifiedObjectName tableName = qualifiedObjectName("existing_table"); + metadata.createTable(testSession, TEST_CATALOG_NAME, rowTable(tableName, new Field(Optional.of("a"), BIGINT)), false); + + assertTrinoExceptionThrownBy(() -> getFutureValue(executeSetColumnType(asQualifiedName(tableName), QualifiedName.of("test", "a"), toSqlType(INTEGER), false))) + .hasErrorCode(COLUMN_NOT_FOUND) + .hasMessageContaining("Column 'test.a' does not exist"); + } + + @Test + public void testSetFieldDataTypeNotExistingField() + { + QualifiedObjectName tableName = qualifiedObjectName("existing_table"); + metadata.createTable(testSession, TEST_CATALOG_NAME, rowTable(tableName, new Field(Optional.of("a"), BIGINT)), false); + + assertTrinoExceptionThrownBy(() -> getFutureValue(executeSetColumnType(asQualifiedName(tableName), QualifiedName.of("col", "b"), toSqlType(INTEGER), false))) + .hasErrorCode(COLUMN_NOT_FOUND) + .hasMessageContaining("Field 'b' does not exist within row(a bigint)"); + } + + @Test + public void testUnsupportedSetDataTypeDuplicatedField() + { + QualifiedObjectName tableName = qualifiedObjectName("existing_table"); + metadata.createTable(testSession, TEST_CATALOG_NAME, rowTable(tableName, new RowType.Field(Optional.of("a"), BIGINT), new RowType.Field(Optional.of("a"), BIGINT)), false); + TableHandle table = metadata.getTableHandle(testSession, tableName).get(); + assertThat(metadata.getTableMetadata(testSession, table).getColumns()) + .isEqualTo(ImmutableList.of(new ColumnMetadata("col", RowType.rowType( + new RowType.Field(Optional.of("a"), BIGINT), new RowType.Field(Optional.of("a"), BIGINT))))); + + assertTrinoExceptionThrownBy(() -> getFutureValue(executeSetColumnType(asQualifiedName(tableName), QualifiedName.of("col", "a"), toSqlType(INTEGER), false))) + .hasErrorCode(AMBIGUOUS_NAME) + .hasMessageContaining("Field path [col, a] within row(a bigint, a bigint) is ambiguous"); + } + + private static ConnectorTableMetadata rowTable(QualifiedObjectName tableName, Field... fields) + { + return new ConnectorTableMetadata(tableName.asSchemaTableName(), ImmutableList.of( + new ColumnMetadata("col", rowType(fields)))); + } + + private ListenableFuture executeSetColumnType(QualifiedName table, QualifiedName column, DataType type, boolean exists) { return new SetColumnTypeTask(metadata, plannerContext.getTypeManager(), new AllowAllAccessControl()) .execute(new SetColumnType(new NodeLocation(1, 1), table, column, type, exists), queryStateMachine, ImmutableList.of(), WarningCollector.NOOP); diff --git a/core/trino-main/src/test/java/io/trino/metadata/AbstractMockMetadata.java b/core/trino-main/src/test/java/io/trino/metadata/AbstractMockMetadata.java index a38dda8b79adf..1f0364537d43a 100644 --- a/core/trino-main/src/test/java/io/trino/metadata/AbstractMockMetadata.java +++ b/core/trino-main/src/test/java/io/trino/metadata/AbstractMockMetadata.java @@ -342,6 +342,12 @@ public void setColumnType(Session session, TableHandle tableHandle, ColumnHandle throw new UnsupportedOperationException(); } + @Override + public void setFieldType(Session session, TableHandle tableHandle, List fieldPath, Type type) + { + throw new UnsupportedOperationException(); + } + @Override public void setTableAuthorization(Session session, CatalogSchemaTableName table, TrinoPrincipal principal) { diff --git a/core/trino-parser/src/main/antlr4/io/trino/sql/parser/SqlBase.g4 b/core/trino-parser/src/main/antlr4/io/trino/sql/parser/SqlBase.g4 index bd4023c0a3e8b..f2a0d046c3801 100644 --- a/core/trino-parser/src/main/antlr4/io/trino/sql/parser/SqlBase.g4 +++ b/core/trino-parser/src/main/antlr4/io/trino/sql/parser/SqlBase.g4 @@ -80,7 +80,7 @@ statement | ALTER TABLE (IF EXISTS)? tableName=qualifiedName DROP COLUMN (IF EXISTS)? column=qualifiedName #dropColumn | ALTER TABLE (IF EXISTS)? tableName=qualifiedName - ALTER COLUMN columnName=identifier SET DATA TYPE type #setColumnType + ALTER COLUMN columnName=qualifiedName SET DATA TYPE type #setColumnType | ALTER TABLE tableName=qualifiedName SET AUTHORIZATION principal #setTableAuthorization | ALTER TABLE tableName=qualifiedName SET PROPERTIES propertyAssignments #setTableProperties diff --git a/core/trino-parser/src/main/java/io/trino/sql/parser/AstBuilder.java b/core/trino-parser/src/main/java/io/trino/sql/parser/AstBuilder.java index 25e491e065c9c..2887e69f7a3bf 100644 --- a/core/trino-parser/src/main/java/io/trino/sql/parser/AstBuilder.java +++ b/core/trino-parser/src/main/java/io/trino/sql/parser/AstBuilder.java @@ -801,7 +801,7 @@ public Node visitSetColumnType(SqlBaseParser.SetColumnTypeContext context) return new SetColumnType( getLocation(context), getQualifiedName(context.tableName), - (Identifier) visit(context.columnName), + getQualifiedName(context.columnName), (DataType) visit(context.type()), context.EXISTS() != null); } diff --git a/core/trino-parser/src/main/java/io/trino/sql/tree/SetColumnType.java b/core/trino-parser/src/main/java/io/trino/sql/tree/SetColumnType.java index 9f5913d0fcb3f..bcf1991c135dd 100644 --- a/core/trino-parser/src/main/java/io/trino/sql/tree/SetColumnType.java +++ b/core/trino-parser/src/main/java/io/trino/sql/tree/SetColumnType.java @@ -26,11 +26,11 @@ public class SetColumnType extends Statement { private final QualifiedName tableName; - private final Identifier columnName; + private final QualifiedName columnName; private final DataType type; private final boolean tableExists; - public SetColumnType(NodeLocation location, QualifiedName tableName, Identifier columnName, DataType type, boolean tableExists) + public SetColumnType(NodeLocation location, QualifiedName tableName, QualifiedName columnName, DataType type, boolean tableExists) { super(Optional.of(location)); this.tableName = requireNonNull(tableName, "tableName is null"); @@ -44,7 +44,7 @@ public QualifiedName getTableName() return tableName; } - public Identifier getColumnName() + public QualifiedName getColumnName() { return columnName; } @@ -68,7 +68,7 @@ public R accept(AstVisitor visitor, C context) @Override public List getChildren() { - return ImmutableList.of(columnName, type); + return ImmutableList.of(type); } @Override diff --git a/core/trino-parser/src/test/java/io/trino/sql/parser/TestSqlParser.java b/core/trino-parser/src/test/java/io/trino/sql/parser/TestSqlParser.java index a5a47bff89e91..98215d3e691fa 100644 --- a/core/trino-parser/src/test/java/io/trino/sql/parser/TestSqlParser.java +++ b/core/trino-parser/src/test/java/io/trino/sql/parser/TestSqlParser.java @@ -3235,7 +3235,7 @@ public void testAlterColumnSetDataType() QualifiedName.of(ImmutableList.of( new Identifier(location(1, 13), "foo", false), new Identifier(location(1, 17), "t", false))), - new Identifier(location(1, 32), "a", false), + QualifiedName.of(ImmutableList.of(new Identifier(location(1, 32), "a", false))), simpleType(location(1, 48), "bigint"), false)); @@ -3245,7 +3245,7 @@ public void testAlterColumnSetDataType() QualifiedName.of(ImmutableList.of( new Identifier(location(1, 23), "foo", false), new Identifier(location(1, 27), "t", false))), - new Identifier(location(1, 42), "b", false), + QualifiedName.of(ImmutableList.of(new Identifier(location(1, 42), "b", false))), simpleType(location(1, 58), "double"), true)); } diff --git a/core/trino-spi/src/main/java/io/trino/spi/connector/ConnectorMetadata.java b/core/trino-spi/src/main/java/io/trino/spi/connector/ConnectorMetadata.java index 730af0bbc234b..a3ea88e0ac781 100644 --- a/core/trino-spi/src/main/java/io/trino/spi/connector/ConnectorMetadata.java +++ b/core/trino-spi/src/main/java/io/trino/spi/connector/ConnectorMetadata.java @@ -470,6 +470,17 @@ default void setColumnType(ConnectorSession session, ConnectorTableHandle tableH throw new TrinoException(NOT_SUPPORTED, "This connector does not support setting column types"); } + /** + * Set the specified field type + * + * @param fieldPath path starting with column name. The path is always lower-cased. It cannot be an empty or a single element. + */ + @Experimental(eta = "2023-09-01") + default void setFieldType(ConnectorSession session, ConnectorTableHandle tableHandle, List fieldPath, Type type) + { + throw new TrinoException(NOT_SUPPORTED, "This connector does not support setting field types"); + } + /** * Sets the user/role on the specified table. */ diff --git a/lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/classloader/ClassLoaderSafeConnectorMetadata.java b/lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/classloader/ClassLoaderSafeConnectorMetadata.java index 19f452c808345..7725a5bbb5651 100644 --- a/lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/classloader/ClassLoaderSafeConnectorMetadata.java +++ b/lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/classloader/ClassLoaderSafeConnectorMetadata.java @@ -343,6 +343,14 @@ public void setColumnType(ConnectorSession session, ConnectorTableHandle tableHa } } + @Override + public void setFieldType(ConnectorSession session, ConnectorTableHandle tableHandle, List fieldPath, Type type) + { + try (ThreadContextClassLoader ignored = new ThreadContextClassLoader(classLoader)) { + delegate.setFieldType(session, tableHandle, fieldPath, type); + } + } + @Override public void setTableAuthorization(ConnectorSession session, SchemaTableName table, TrinoPrincipal principal) { From 4c6319eeabe671e34b41b527a24c6acfd7c3129b Mon Sep 17 00:00:00 2001 From: Yuya Ebihara Date: Tue, 25 Jul 2023 09:40:49 +0900 Subject: [PATCH 2/2] Support setting a field type with SET DATA TYPE in Iceberg --- .../trino/plugin/iceberg/IcebergMetadata.java | 30 ++++ .../iceberg/BaseIcebergConnectorTest.java | 42 +++++ .../mongodb/TestMongoConnectorTest.java | 1 + .../TestIcebergSparkCompatibility.java | 20 +++ .../io/trino/testing/BaseConnectorTest.java | 151 ++++++++++++++++++ .../testing/TestingConnectorBehavior.java | 1 + 6 files changed, 245 insertions(+) diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergMetadata.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergMetadata.java index f2e167b4d0591..20aa1795d856c 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergMetadata.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergMetadata.java @@ -1839,6 +1839,36 @@ private static boolean fieldExists(StructType structType, String fieldName) return false; } + @Override + public void setFieldType(ConnectorSession session, ConnectorTableHandle tableHandle, List fieldPath, io.trino.spi.type.Type type) + { + Table icebergTable = catalog.loadTable(session, ((IcebergTableHandle) tableHandle).getSchemaTableName()); + String parentPath = String.join(".", fieldPath.subList(0, fieldPath.size() - 1)); + NestedField parent = icebergTable.schema().caseInsensitiveFindField(parentPath); + + String caseSensitiveParentName = icebergTable.schema().findColumnName(parent.fieldId()); + NestedField field = parent.type().asStructType().caseInsensitiveField(getLast(fieldPath)); + // TODO: Add support for changing non-primitive field type + if (!field.type().isPrimitiveType()) { + throw new TrinoException(NOT_SUPPORTED, "Iceberg doesn't support changing field type from non-primitive types"); + } + + String name = caseSensitiveParentName + "." + field.name(); + // Pass dummy AtomicInteger. The field id will be discarded because the subsequent logic disallows non-primitive types. + Type icebergType = toIcebergTypeForNewColumn(type, new AtomicInteger()); + if (!icebergType.isPrimitiveType()) { + throw new TrinoException(NOT_SUPPORTED, "Iceberg doesn't support changing field type to non-primitive types"); + } + try { + icebergTable.updateSchema() + .updateColumn(name, icebergType.asPrimitiveType()) + .commit(); + } + catch (RuntimeException e) { + throw new TrinoException(ICEBERG_COMMIT_ERROR, "Failed to set field type: " + firstNonNull(e.getMessage(), e), e); + } + } + private List getColumnMetadatas(Schema schema) { ImmutableList.Builder columns = ImmutableList.builder(); diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/BaseIcebergConnectorTest.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/BaseIcebergConnectorTest.java index 171cb711510e6..3a4c87bac2140 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/BaseIcebergConnectorTest.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/BaseIcebergConnectorTest.java @@ -6951,6 +6951,48 @@ protected void verifySetColumnTypeFailurePermissible(Throwable e) "|Type not supported for Iceberg: char\\(20\\)).*"); } + @Override + protected Optional filterSetFieldTypesDataProvider(SetColumnTypeSetup setup) + { + switch ("%s -> %s".formatted(setup.sourceColumnType(), setup.newColumnType())) { + case "bigint -> integer": + case "decimal(5,3) -> decimal(5,2)": + case "varchar -> char(20)": + case "time(6) -> time(3)": + case "timestamp(6) -> timestamp(3)": + case "array(integer) -> array(bigint)": + case "row(x integer) -> row(x bigint)": + case "row(x integer) -> row(y integer)": + case "row(x integer, y integer) -> row(x integer, z integer)": + case "row(x integer) -> row(x integer, y integer)": + case "row(x integer, y integer) -> row(x integer)": + case "row(x integer, y integer) -> row(y integer, x integer)": + case "row(x integer, y integer) -> row(z integer, y integer, x integer)": + case "row(x row(nested integer)) -> row(x row(nested bigint))": + case "row(x row(a integer, b integer)) -> row(x row(b integer, a integer))": + // Iceberg allows updating column types if the update is safe. Safe updates are: + // - int to bigint + // - float to double + // - decimal(P,S) to decimal(P2,S) when P2 > P (scale cannot change) + // https://iceberg.apache.org/docs/latest/spark-ddl/#alter-table--alter-column + return Optional.of(setup.asUnsupported()); + + case "varchar(100) -> varchar(50)": + // Iceberg connector ignores the varchar length + return Optional.empty(); + } + return Optional.of(setup); + } + + @Override + protected void verifySetFieldTypeFailurePermissible(Throwable e) + { + assertThat(e).hasMessageMatching(".*(Failed to set field type: Cannot change (column type:|type from .* to )" + + "|Time(stamp)? precision \\(3\\) not supported for Iceberg. Use \"time(stamp)?\\(6\\)\" instead" + + "|Type not supported for Iceberg: char\\(20\\)" + + "|Iceberg doesn't support changing field type (from|to) non-primitive types).*"); + } + @Override protected boolean supportsPhysicalPushdown() { diff --git a/plugin/trino-mongodb/src/test/java/io/trino/plugin/mongodb/TestMongoConnectorTest.java b/plugin/trino-mongodb/src/test/java/io/trino/plugin/mongodb/TestMongoConnectorTest.java index 533cde454bdc7..dc48d0af5f071 100644 --- a/plugin/trino-mongodb/src/test/java/io/trino/plugin/mongodb/TestMongoConnectorTest.java +++ b/plugin/trino-mongodb/src/test/java/io/trino/plugin/mongodb/TestMongoConnectorTest.java @@ -113,6 +113,7 @@ protected boolean hasBehavior(TestingConnectorBehavior connectorBehavior) case SUPPORTS_ADD_FIELD: case SUPPORTS_RENAME_FIELD: case SUPPORTS_DROP_FIELD: + case SUPPORTS_SET_FIELD_TYPE: return false; case SUPPORTS_CREATE_VIEW: diff --git a/testing/trino-product-tests/src/main/java/io/trino/tests/product/iceberg/TestIcebergSparkCompatibility.java b/testing/trino-product-tests/src/main/java/io/trino/tests/product/iceberg/TestIcebergSparkCompatibility.java index 7e442175b91c0..9d51741b127e5 100644 --- a/testing/trino-product-tests/src/main/java/io/trino/tests/product/iceberg/TestIcebergSparkCompatibility.java +++ b/testing/trino-product-tests/src/main/java/io/trino/tests/product/iceberg/TestIcebergSparkCompatibility.java @@ -2748,6 +2748,26 @@ private void testTrinoSetColumnType(boolean partitioned, StorageFormat storageFo onTrino().executeQuery("DROP TABLE " + trinoTableName); } + @Test(groups = {ICEBERG, PROFILE_SPECIFIC_TESTS}, dataProvider = "testSetColumnTypeDataProvider") + public void testTrinoSetFieldType(StorageFormat storageFormat, String sourceFieldType, String sourceValueLiteral, String newFieldType, Object newValue) + { + String baseTableName = "test_set_field_type_" + randomNameSuffix(); + String trinoTableName = trinoTableName(baseTableName); + String sparkTableName = sparkTableName(baseTableName); + + onTrino().executeQuery("CREATE TABLE " + trinoTableName + " " + + "WITH (format = '" + storageFormat + "')" + + "AS SELECT CAST(row(" + sourceValueLiteral + ") AS row(field " + sourceFieldType + ")) AS col"); + + onTrino().executeQuery("ALTER TABLE " + trinoTableName + " ALTER COLUMN col.field SET DATA TYPE " + newFieldType); + + assertEquals(getColumnType(baseTableName, "col"), "row(field " + newFieldType + ")"); + assertThat(onTrino().executeQuery("SELECT col.field FROM " + trinoTableName)).containsOnly(row(newValue)); + assertThat(onSpark().executeQuery("SELECT col.field FROM " + sparkTableName)).containsOnly(row(newValue)); + + onTrino().executeQuery("DROP TABLE " + trinoTableName); + } + @DataProvider public static Object[][] testSetColumnTypeDataProvider() { diff --git a/testing/trino-testing/src/main/java/io/trino/testing/BaseConnectorTest.java b/testing/trino-testing/src/main/java/io/trino/testing/BaseConnectorTest.java index a99389e84943c..bca0c4f6d27dc 100644 --- a/testing/trino-testing/src/main/java/io/trino/testing/BaseConnectorTest.java +++ b/testing/trino-testing/src/main/java/io/trino/testing/BaseConnectorTest.java @@ -147,6 +147,7 @@ import static io.trino.testing.TestingConnectorBehavior.SUPPORTS_ROW_LEVEL_DELETE; import static io.trino.testing.TestingConnectorBehavior.SUPPORTS_ROW_TYPE; import static io.trino.testing.TestingConnectorBehavior.SUPPORTS_SET_COLUMN_TYPE; +import static io.trino.testing.TestingConnectorBehavior.SUPPORTS_SET_FIELD_TYPE; import static io.trino.testing.TestingConnectorBehavior.SUPPORTS_TOPN_PUSHDOWN; import static io.trino.testing.TestingConnectorBehavior.SUPPORTS_TRUNCATE; import static io.trino.testing.TestingConnectorBehavior.SUPPORTS_UPDATE; @@ -3041,6 +3042,156 @@ protected void verifySetColumnTypeFailurePermissible(Throwable e) throw new AssertionError("Unexpected set column type failure", e); } + @Test + public void testSetFieldType() + { + skipTestUnless(hasBehavior(SUPPORTS_CREATE_TABLE_WITH_DATA) && hasBehavior(SUPPORTS_ROW_TYPE)); + + if (!hasBehavior(SUPPORTS_SET_FIELD_TYPE)) { + try (TestTable table = new TestTable(getQueryRunner()::execute, "test_set_field_type_", "(col row(field int))")) { + assertQueryFails( + "ALTER TABLE " + table.getName() + " ALTER COLUMN col.field SET DATA TYPE bigint", + "This connector does not support setting field types"); + } + return; + } + + try (TestTable table = new TestTable(getQueryRunner()::execute, "test_set_field_type_", "AS SELECT CAST(row(123) AS row(field integer)) AS col")) { + assertEquals(getColumnType(table.getName(), "col"), "row(field integer)"); + + assertUpdate("ALTER TABLE " + table.getName() + " ALTER COLUMN col.field SET DATA TYPE bigint"); + + assertEquals(getColumnType(table.getName(), "col"), "row(field bigint)"); + assertThat(query("SELECT * FROM " + table.getName())) + .skippingTypesCheck() + .matches("SELECT row(bigint '123')"); + } + } + + @Test(dataProvider = "setFieldTypesDataProvider") + public void testSetFieldTypes(SetColumnTypeSetup setup) + { + skipTestUnless(hasBehavior(SUPPORTS_SET_FIELD_TYPE) && hasBehavior(SUPPORTS_CREATE_TABLE_WITH_DATA)); + + TestTable table; + try { + table = new TestTable( + getQueryRunner()::execute, + "test_set_field_type_", + " AS SELECT CAST(row(" + setup.sourceValueLiteral + ") AS row(field " + setup.sourceColumnType + ")) AS col"); + } + catch (Exception e) { + verifyUnsupportedTypeException(e, setup.sourceColumnType); + throw new SkipException("Unsupported column type: " + setup.sourceColumnType); + } + try (table) { + Runnable setFieldType = () -> assertUpdate("ALTER TABLE " + table.getName() + " ALTER COLUMN col.field SET DATA TYPE " + setup.newColumnType); + if (setup.unsupportedType) { + assertThatThrownBy(setFieldType::run) + .satisfies(this::verifySetFieldTypeFailurePermissible); + return; + } + setFieldType.run(); + + assertEquals(getColumnType(table.getName(), "col"), "row(field " + setup.newColumnType + ")"); + assertThat(query("SELECT * FROM " + table.getName())) + .skippingTypesCheck() + .matches("SELECT row(" + setup.newValueLiteral + ")"); + } + } + + @DataProvider + public Object[][] setFieldTypesDataProvider() + { + return setColumnTypeSetupData().stream() + .map(this::filterSetFieldTypesDataProvider) + .flatMap(Optional::stream) + .collect(toDataProvider()); + } + + protected Optional filterSetFieldTypesDataProvider(SetColumnTypeSetup setup) + { + return Optional.of(setup); + } + + @Test + public void testSetFieldTypeCaseSensitivity() + { + skipTestUnless(hasBehavior(SUPPORTS_SET_FIELD_TYPE) && hasBehavior(SUPPORTS_NOT_NULL_CONSTRAINT)); + + try (TestTable table = new TestTable(getQueryRunner()::execute, "test_set_field_type_case_", " AS SELECT CAST(row(1) AS row(\"UPPER\" integer)) col")) { + assertEquals(getColumnType(table.getName(), "col"), "row(UPPER integer)"); + + assertUpdate("ALTER TABLE " + table.getName() + " ALTER COLUMN col.upper SET DATA TYPE bigint"); + assertEquals(getColumnType(table.getName(), "col"), "row(UPPER bigint)"); + assertThat(query("SELECT * FROM " + table.getName())) + .matches("SELECT CAST(row(1) AS row(UPPER bigint))"); + } + } + + @Test + public void testSetFieldTypeWithNotNull() + { + skipTestUnless(hasBehavior(SUPPORTS_SET_FIELD_TYPE) && hasBehavior(SUPPORTS_NOT_NULL_CONSTRAINT)); + + try (TestTable table = new TestTable(getQueryRunner()::execute, "test_set_field_type_null_", "(col row(field int) NOT NULL)")) { + assertFalse(columnIsNullable(table.getName(), "col")); + + assertUpdate("ALTER TABLE " + table.getName() + " ALTER COLUMN col.field SET DATA TYPE bigint"); + assertFalse(columnIsNullable(table.getName(), "col")); + } + } + + @Test + public void testSetFieldTypeWithComment() + { + skipTestUnless(hasBehavior(SUPPORTS_SET_FIELD_TYPE) && hasBehavior(SUPPORTS_CREATE_TABLE_WITH_COLUMN_COMMENT)); + + try (TestTable table = new TestTable(getQueryRunner()::execute, "test_set_field_type_comment_", "(col row(field int) COMMENT 'test comment')")) { + assertEquals(getColumnComment(table.getName(), "col"), "test comment"); + + assertUpdate("ALTER TABLE " + table.getName() + " ALTER COLUMN col.field SET DATA TYPE bigint"); + assertEquals(getColumnComment(table.getName(), "col"), "test comment"); + } + } + + @Test + public void testSetFieldIncompatibleType() + { + skipTestUnless(hasBehavior(SUPPORTS_SET_FIELD_TYPE) && hasBehavior(SUPPORTS_CREATE_TABLE_WITH_DATA)); + + try (TestTable table = new TestTable( + getQueryRunner()::execute, + "test_set_invalid_field_type_", + "(row_col row(field varchar), nested_col row(field row(nested int)))")) { + assertThatThrownBy(() -> assertUpdate("ALTER TABLE " + table.getName() + " ALTER COLUMN row_col.field SET DATA TYPE row(nested integer)")) + .satisfies(this::verifySetFieldTypeFailurePermissible); + assertThatThrownBy(() -> assertUpdate("ALTER TABLE " + table.getName() + " ALTER COLUMN row_col.field SET DATA TYPE integer")) + .satisfies(this::verifySetFieldTypeFailurePermissible); + assertThatThrownBy(() -> assertUpdate("ALTER TABLE " + table.getName() + " ALTER COLUMN nested_col.field SET DATA TYPE integer")) + .satisfies(this::verifySetFieldTypeFailurePermissible); + } + } + + @Test + public void testSetFieldOutOfRangeType() + { + skipTestUnless(hasBehavior(SUPPORTS_SET_FIELD_TYPE) && hasBehavior(SUPPORTS_CREATE_TABLE_WITH_DATA)); + + try (TestTable table = new TestTable( + getQueryRunner()::execute, + "test_set_field_type_invalid_range_", + "AS SELECT CAST(row(9223372036854775807) AS row(field bigint)) AS col")) { + assertThatThrownBy(() -> assertUpdate("ALTER TABLE " + table.getName() + " ALTER COLUMN col.field SET DATA TYPE integer")) + .satisfies(this::verifySetFieldTypeFailurePermissible); + } + } + + protected void verifySetFieldTypeFailurePermissible(Throwable e) + { + throw new AssertionError("Unexpected set field type failure", e); + } + protected String getColumnType(String tableName, String columnName) { return (String) computeScalar(format("SELECT data_type FROM information_schema.columns WHERE table_schema = CURRENT_SCHEMA AND table_name = '%s' AND column_name = '%s'", diff --git a/testing/trino-testing/src/main/java/io/trino/testing/TestingConnectorBehavior.java b/testing/trino-testing/src/main/java/io/trino/testing/TestingConnectorBehavior.java index 13e61516672a1..558a96db61587 100644 --- a/testing/trino-testing/src/main/java/io/trino/testing/TestingConnectorBehavior.java +++ b/testing/trino-testing/src/main/java/io/trino/testing/TestingConnectorBehavior.java @@ -88,6 +88,7 @@ public enum TestingConnectorBehavior SUPPORTS_RENAME_COLUMN, SUPPORTS_RENAME_FIELD(fallback -> fallback.test(SUPPORTS_RENAME_COLUMN) && fallback.test(SUPPORTS_ROW_TYPE)), SUPPORTS_SET_COLUMN_TYPE, + SUPPORTS_SET_FIELD_TYPE(fallback -> fallback.test(SUPPORTS_SET_COLUMN_TYPE) && fallback.test(SUPPORTS_ROW_TYPE)), SUPPORTS_COMMENT_ON_TABLE, SUPPORTS_COMMENT_ON_COLUMN(SUPPORTS_COMMENT_ON_TABLE),