Allow changing field type in records wrapped in a map

trinodb · Nov 27, 2024 · 9d98c66 · 9d98c66
1 parent 332a658
commit 9d98c66
Show file tree

Hide file tree

Showing 4 changed files with 116 additions and 0 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
@@ -27,6 +27,7 @@
 import io.trino.spi.connector.ColumnHandle;
 import io.trino.spi.connector.ColumnMetadata;
 import io.trino.spi.type.ArrayType;
+import io.trino.spi.type.MapType;
 import io.trino.spi.type.RowType;
 import io.trino.spi.type.Type;
 import io.trino.spi.type.TypeManager;
@@ -153,6 +154,15 @@ private static List<RowType.Field> getCandidates(Type type, String fieldName)
             // return nameless Field to denote unwrapping of container
             return ImmutableList.of(RowType.field(arrayType.getElementType()));
         }
+        if (type instanceof MapType mapType) {
+            if (fieldName.equals("key")) {
+                return ImmutableList.of(RowType.field(mapType.getKeyType()));
+            }
+            if (fieldName.equals("value")) {
+                return ImmutableList.of(RowType.field(mapType.getValueType()));
+            }
+            throw new TrinoException(NOT_SUPPORTED, "MAP type should be denoted by 'key' or 'value' in the path; found '%s'".formatted(fieldName));
+        }
         if (!(type instanceof RowType rowType)) {
             throw new TrinoException(NOT_SUPPORTED, "Unsupported type: " + type);
         }

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
@@ -8567,6 +8567,24 @@ public void testCreateTableWithDataLocationButObjectStoreDisabled()
                 "Data location can only be set when object store is enabled");
     }
 
+    @Test
+    @Override
+    public void testSetFieldMapKeyType()
+    {
+        // Iceberg doesn't support change a map 'key' column. Only map values can be changed.
+        assertThatThrownBy(super::testSetFieldMapKeyType)
+                .hasMessageContaining("Failed to set field type: Cannot alter map keys");
+    }
+
+    @Test
+    @Override
+    public void testSetNestedFieldMapKeyType()
+    {
+        // Iceberg doesn't support change a map 'key' column. Only map values can be changed.
+        assertThatThrownBy(super::testSetNestedFieldMapKeyType)
+                .hasMessageContaining("Failed to set field type: Cannot alter map keys");
+    }
+
     @Override
     protected Optional<SetColumnTypeSetup> filterSetColumnTypesDataProvider(SetColumnTypeSetup setup)
     {

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
@@ -151,6 +151,7 @@
 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_SET_FIELD_TYPE_IN_MAP;
 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;
@@ -3295,6 +3296,92 @@ public void testSetFieldTypeInNestedArray()
         }
     }
 
+    @Test
+    public void testSetFieldMapKeyType()
+    {
+        skipTestUnless(hasBehavior(SUPPORTS_CREATE_TABLE_WITH_DATA) && hasBehavior(SUPPORTS_ROW_TYPE));
+
+        String tableDefinition = "AS SELECT CAST(map(array[row(1)], array[2]) AS map(row(field integer), integer)) AS col";
+        if (!hasBehavior(SUPPORTS_SET_FIELD_TYPE_IN_MAP)) {
+            try (TestTable table = new TestTable(getQueryRunner()::execute, "test_set_field_type_in_map", tableDefinition)) {
+                assertQueryFails("ALTER TABLE " + table.getName() + " ALTER COLUMN col.key.field SET DATA TYPE bigint", ".*does not support.*");
+            }
+            return;
+        }
+
+        try (TestTable table = new TestTable(getQueryRunner()::execute, "test_set_field_type_in_map", tableDefinition)) {
+            assertThat(getColumnType(table.getName(), "col")).isEqualTo("map(row(field integer), integer)");
+
+            assertUpdate("ALTER TABLE " + table.getName() + " ALTER COLUMN col.key.field SET DATA TYPE bigint");
+
+            assertThat(getColumnType(table.getName(), "col")).isEqualTo("map(row(field bigint), integer)");
+            assertThat(query("SELECT * FROM " + table.getName()))
+                    .matches("SELECT CAST(map(array[row(1)], array[2]) AS map(row(field bigint), integer))");
+        }
+    }
+
+    @Test
+    public void testSetFieldMapValueType()
+    {
+        skipTestUnless(hasBehavior(SUPPORTS_CREATE_TABLE_WITH_DATA) && hasBehavior(SUPPORTS_ROW_TYPE));
+
+        String tableDefinition = "AS SELECT CAST(map(array[1], array[row(2)]) AS map(integer, row(field integer))) AS col";
+        if (!hasBehavior(SUPPORTS_SET_FIELD_TYPE_IN_MAP)) {
+            try (TestTable table = new TestTable(getQueryRunner()::execute, "test_set_field_type_in_map", tableDefinition)) {
+                assertQueryFails("ALTER TABLE " + table.getName() + " ALTER COLUMN col.value.field SET DATA TYPE bigint", ".*does not support.*");
+            }
+            return;
+        }
+
+        try (TestTable table = new TestTable(getQueryRunner()::execute, "test_set_field_type_in_map", tableDefinition)) {
+            assertThat(getColumnType(table.getName(), "col")).isEqualTo("map(integer, row(field integer))");
+
+            assertUpdate("ALTER TABLE " + table.getName() + " ALTER COLUMN col.value.field SET DATA TYPE bigint");
+
+            assertThat(getColumnType(table.getName(), "col")).isEqualTo("map(integer, row(field bigint))");
+            assertThat(query("SELECT * FROM " + table.getName()))
+                    .matches("SELECT CAST(map(array[1], array[row(2)]) AS map(integer, row(field bigint)))");
+        }
+    }
+
+    @Test
+    public void testSetNestedFieldMapKeyType()
+    {
+        skipTestUnless(hasBehavior(SUPPORTS_SET_FIELD_TYPE_IN_ARRAY) && hasBehavior(SUPPORTS_SET_FIELD_TYPE_IN_MAP) && hasBehavior(SUPPORTS_CREATE_TABLE_WITH_DATA) && hasBehavior(SUPPORTS_ARRAY) && hasBehavior(SUPPORTS_ROW_TYPE));
+
+        try (TestTable table = new TestTable(
+                getQueryRunner()::execute,
+                "test_set_nested_field_type_in_map",
+                "AS SELECT CAST(array[map(array[row(1)], array[2])] AS array(map(row(field integer), integer))) AS col")) {
+            assertThat(getColumnType(table.getName(), "col")).isEqualTo("array(map(row(field integer), integer))");
+
+            assertUpdate("ALTER TABLE " + table.getName() + " ALTER COLUMN col.element.key.field SET DATA TYPE bigint");
+
+            assertThat(getColumnType(table.getName(), "col")).isEqualTo("array(map(row(field bigint), integer))");
+            assertThat(query("SELECT * FROM " + table.getName()))
+                    .matches("SELECT CAST(array[map(array[row(1)], array[2])] AS array(map(row(field bigint), integer)))");
+        }
+    }
+
+    @Test
+    public void testSetNestedFieldMapValueType()
+    {
+        skipTestUnless(hasBehavior(SUPPORTS_SET_FIELD_TYPE_IN_ARRAY) && hasBehavior(SUPPORTS_SET_FIELD_TYPE_IN_MAP) && hasBehavior(SUPPORTS_CREATE_TABLE_WITH_DATA) && hasBehavior(SUPPORTS_ARRAY) && hasBehavior(SUPPORTS_ROW_TYPE));
+
+        try (TestTable table = new TestTable(
+                getQueryRunner()::execute,
+                "test_set_nested_field_type_in_map",
+                "AS SELECT CAST(array[map(array[1], array[row(2)])] AS array(map(integer, row(field integer)))) AS col")) {
+            assertThat(getColumnType(table.getName(), "col")).isEqualTo("array(map(integer, row(field integer)))");
+
+            assertUpdate("ALTER TABLE " + table.getName() + " ALTER COLUMN col.element.value.field SET DATA TYPE bigint");
+
+            assertThat(getColumnType(table.getName(), "col")).isEqualTo("array(map(integer, row(field bigint)))");
+            assertThat(query("SELECT * FROM " + table.getName()))
+                    .matches("SELECT CAST(array[map(array[1], array[row(2)])] AS array(map(integer, row(field bigint))))");
+        }
+    }
+
     protected void verifySetFieldTypeFailurePermissible(Throwable e)
     {
         throw new AssertionError("Unexpected set field type failure", e);

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
@@ -94,6 +94,7 @@ public enum TestingConnectorBehavior
     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_SET_FIELD_TYPE_IN_MAP(SUPPORTS_SET_FIELD_TYPE),
 
     SUPPORTS_COMMENT_ON_TABLE,
     SUPPORTS_COMMENT_ON_COLUMN(SUPPORTS_COMMENT_ON_TABLE),