Allow changing field type in records wrapped in an array

trinodb · May 25, 2024 · ef5489f · ef5489f
1 parent 447fee4
commit ef5489f
Show file tree

Hide file tree

Showing 4 changed files with 64 additions and 2 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
@@ -25,6 +25,7 @@
 import io.trino.spi.TrinoException;
 import io.trino.spi.connector.ColumnHandle;
 import io.trino.spi.connector.ColumnMetadata;
+import io.trino.spi.type.ArrayType;
 import io.trino.spi.type.RowType;
 import io.trino.spi.type.Type;
 import io.trino.spi.type.TypeManager;
@@ -139,7 +140,11 @@ else if (metadata.isView(session, qualifiedObjectName)) {
 
     private static List<RowType.Field> getCandidates(Type type, String fieldName)
     {
-        if (!(type instanceof RowType rowType)) {
+        Type analyzedType = type;
+        if (type instanceof ArrayType arrayType) {
+            analyzedType = arrayType.getElementType();
+        }
+        if (!(analyzedType instanceof RowType rowType)) {
             throw new TrinoException(NOT_SUPPORTED, "Unsupported type: " + type);
         }
         List<RowType.Field> candidates = rowType.getFields().stream()

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
@@ -2132,7 +2132,19 @@ public void setFieldType(ConnectorSession session, ConnectorTableHandle tableHan
         NestedField parent = icebergTable.schema().caseInsensitiveFindField(parentPath);
 
         String caseSensitiveParentName = icebergTable.schema().findColumnName(parent.fieldId());
-        NestedField field = parent.type().asStructType().caseInsensitiveField(getLast(fieldPath));
+
+        Types.StructType structType;
+        if (parent.type().isListType()) {
+            // list(struct...)
+            structType = parent.type().asListType().elementType().asStructType();
+            caseSensitiveParentName += ".element";
+        }
+        else {
+            // just struct
+            structType = parent.type().asStructType();
+        }
+        NestedField field = structType.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");

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
@@ -149,6 +149,7 @@
 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_SET_FIELD_TYPE_IN_ARRAY;
 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;
@@ -3159,6 +3160,49 @@ public void testSetFieldOutOfRangeType()
         }
     }
 
+    @Test
+    public void testSetFieldTypeInArray()
+    {
+        skipTestUnless(hasBehavior(SUPPORTS_CREATE_TABLE_WITH_DATA) && hasBehavior(SUPPORTS_ARRAY) && hasBehavior(SUPPORTS_ROW_TYPE));
+
+        if (!hasBehavior(SUPPORTS_SET_FIELD_TYPE_IN_ARRAY)) {
+            try (TestTable table = new TestTable(getQueryRunner()::execute, "test_set_field_type_in_array_", "(col array(row(field int)))")) {
+                assertQueryFails(
+                        "ALTER TABLE " + table.getName() + " ALTER COLUMN col.field SET DATA TYPE bigint",
+                        ".*does not support.*");
+            }
+            return;
+        }
+
+        try (TestTable table = new TestTable(getQueryRunner()::execute, "test_set_field_type_in_array_", "AS SELECT CAST(array[row(123)] AS array(row(field integer))) AS col")) {
+            assertThat(getColumnType(table.getName(), "col")).isEqualTo("array(row(field integer))");
+
+            assertUpdate("ALTER TABLE " + table.getName() + " ALTER COLUMN col.field SET DATA TYPE bigint");
+
+            assertThat(getColumnType(table.getName(), "col")).isEqualTo("array(row(field bigint))");
+            assertThat(query("SELECT * FROM " + table.getName()))
+                    .skippingTypesCheck()
+                    .matches("SELECT array[row(bigint '123')]");
+        }
+    }
+
+    @Test
+    public void testSetFieldTypeInNestedArray()
+    {
+        skipTestUnless(hasBehavior(SUPPORTS_SET_FIELD_TYPE_IN_ARRAY) && hasBehavior(SUPPORTS_CREATE_TABLE_WITH_DATA) && hasBehavior(SUPPORTS_ARRAY) && hasBehavior(SUPPORTS_ROW_TYPE));
+
+        try (TestTable table = new TestTable(getQueryRunner()::execute, "test_set_field_type_in_nested_array_", "AS SELECT CAST(array[row(array[row(123)])] AS array(row(field array(row(a integer))))) AS col")) {
+            assertThat(getColumnType(table.getName(), "col")).isEqualTo("array(row(field array(row(a integer))))");
+
+            assertUpdate("ALTER TABLE " + table.getName() + " ALTER COLUMN col.field.a SET DATA TYPE bigint");
+
+            assertThat(getColumnType(table.getName(), "col")).isEqualTo("array(row(field array(row(a bigint))))");
+            assertThat(query("SELECT * FROM " + table.getName()))
+                    .skippingTypesCheck()
+                    .matches("SELECT array[row(array[row(bigint '123')])]");
+        }
+    }
+
     protected void verifySetFieldTypeFailurePermissible(Throwable e)
     {
         throw new AssertionError("Unexpe" +

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
@@ -93,6 +93,7 @@ public enum TestingConnectorBehavior
     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_SET_FIELD_TYPE_IN_ARRAY(fallback -> fallback.test(SUPPORTS_SET_FIELD_TYPE) && fallback.test(SUPPORTS_ADD_FIELD_IN_ARRAY)),
 
     SUPPORTS_COMMENT_ON_TABLE,
     SUPPORTS_COMMENT_ON_COLUMN(SUPPORTS_COMMENT_ON_TABLE),