Support setting a field type with SET DATA TYPE statement in engine and Iceberg

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<String, ColumnHandle> 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<String> fieldPath = statement.getColumnName().getParts();
+
+            Type currentType = columnMetadata.getType();
+            for (int i = 1; i < fieldPath.size(); i++) {
+                String fieldName = fieldPath.get(i);
+                List<RowType.Field> 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<RowType.Field> getCandidates(Type type, String fieldName)
+    {
+        if (!(type instanceof RowType rowType)) {
+            throw new TrinoException(NOT_SUPPORTED, "Unsupported type: " + type);
+        }
+        List<RowType.Field> 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;

core/trino-main/src/main/java/io/trino/metadata/Metadata.java

@@ -257,6 +257,11 @@ Optional<TableExecuteHandle> 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<String> fieldPath, Type type);
+
     /**
      * Set the authorization (owner) of specified table's user/role
      */

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<String> 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)
     {

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<String> 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)
     {

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<String> 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)
     {

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<String> fieldPath, Type type)
+        {
+            throw new UnsupportedOperationException();
+        }
+
         @Override
         public void setTableAuthorization(ConnectorSession session, SchemaTableName tableName, TrinoPrincipal principal) {}
 

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,15 +82,15 @@ 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
     }
 
     @Test
     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<Void> 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<Void> 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);

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<String> fieldPath, Type type)
+    {
+        throw new UnsupportedOperationException();
+    }
+
     @Override
     public void setTableAuthorization(Session session, CatalogSchemaTableName table, TrinoPrincipal principal)
     {

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

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);
     }

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, C> R accept(AstVisitor<R, C> visitor, C context)
     @Override
     public List<Node> getChildren()
     {
-        return ImmutableList.of(columnName, type);
+        return ImmutableList.of(type);
     }
 
     @Override

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));
     }

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<String> fieldPath, Type type)
+    {
+        throw new TrinoException(NOT_SUPPORTED, "This connector does not support setting field types");
+    }
+
     /**
      * Sets the user/role on the specified table.
      */