Add generated expression to extra info in Delta Lake

trinodb · Mar 20, 2023 · 173947a · 173947a
1 parent fd47ac4
commit 173947a
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 4 deletions.
diff --git a/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakeMetadata.java b/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakeMetadata.java
@@ -452,12 +452,13 @@ public ConnectorTableMetadata getTableMetadata(ConnectorSession session, Connect
         String location = metastore.getTableLocation(tableHandle.getSchemaTableName(), session);
         Map<String, String> columnComments = getColumnComments(tableHandle.getMetadataEntry());
         Map<String, Boolean> columnsNullability = getColumnsNullability(tableHandle.getMetadataEntry());
+        Map<String, String> columnGenerations = getGeneratedColumnExpressions(tableHandle.getMetadataEntry());
         List<String> constraints = ImmutableList.<String>builder()
                 .addAll(getCheckConstraints(tableHandle.getMetadataEntry()).values())
                 .addAll(getColumnInvariants(tableHandle.getMetadataEntry()).values()) // The internal logic for column invariants in Delta Lake is same as check constraints
                 .build();
         List<ColumnMetadata> columns = getColumns(tableHandle.getMetadataEntry()).stream()
-                .map(column -> getColumnMetadata(column, columnComments.get(column.getName()), columnsNullability.getOrDefault(column.getName(), true)))
+                .map(column -> getColumnMetadata(column, columnComments.get(column.getName()), columnsNullability.getOrDefault(column.getName(), true), columnGenerations.get(column.getName())))
                 .collect(toImmutableList());
 
         ImmutableMap.Builder<String, Object> properties = ImmutableMap.<String, Object>builder()
@@ -507,7 +508,8 @@ public ColumnMetadata getColumnMetadata(ConnectorSession session, ConnectorTable
         return getColumnMetadata(
                 column,
                 getColumnComments(table.getMetadataEntry()).get(column.getName()),
-                getColumnsNullability(table.getMetadataEntry()).getOrDefault(column.getName(), true));
+                getColumnsNullability(table.getMetadataEntry()).getOrDefault(column.getName(), true),
+                getGeneratedColumnExpressions(table.getMetadataEntry()).get(column.getName()));
     }
 
     /**
@@ -566,8 +568,9 @@ public Iterator<TableColumnsMetadata> streamTableColumns(ConnectorSession sessio
                         return metastore.getMetadata(metastore.getSnapshot(table, session), session).stream().map(metadata -> {
                             Map<String, String> columnComments = getColumnComments(metadata);
                             Map<String, Boolean> columnsNullability = getColumnsNullability(metadata);
+                            Map<String, String> columnGenerations = getGeneratedColumnExpressions(metadata);
                             List<ColumnMetadata> columnMetadata = getColumns(metadata).stream()
-                                    .map(column -> getColumnMetadata(column, columnComments.get(column.getName()), columnsNullability.getOrDefault(column.getName(), true)))
+                                    .map(column -> getColumnMetadata(column, columnComments.get(column.getName()), columnsNullability.getOrDefault(column.getName(), true), columnGenerations.get(column.getName())))
                                     .collect(toImmutableList());
                             return TableColumnsMetadata.forTable(table, columnMetadata);
                         });
@@ -2556,14 +2559,15 @@ private List<CommitInfoEntry> getCommitInfoEntries(SchemaTableName table, Connec
         }
     }
 
-    private static ColumnMetadata getColumnMetadata(DeltaLakeColumnHandle column, @Nullable String comment, boolean nullability)
+    private static ColumnMetadata getColumnMetadata(DeltaLakeColumnHandle column, @Nullable String comment, boolean nullability, @Nullable String generation)
     {
         return ColumnMetadata.builder()
                 .setName(column.getName())
                 .setType(column.getType())
                 .setHidden(column.getColumnType() == SYNTHESIZED)
                 .setComment(Optional.ofNullable(comment))
                 .setNullable(nullability)
+                .setExtraInfo(generation == null ? Optional.empty() : Optional.of("generated: " + generation))
                 .build();
     }
 

diff --git a/.../src/main/java/io/trino/tests/product/deltalake/TestDeltaLakeAlterTableCompatibility.java b/.../src/main/java/io/trino/tests/product/deltalake/TestDeltaLakeAlterTableCompatibility.java
@@ -368,6 +368,13 @@ public void testTrinoAlterTablePreservesGeneratedColumn()
             onDelta().executeQuery("INSERT INTO default." + tableName + " (a, c) VALUES (1, 3)");
             assertThat(onTrino().executeQuery("SELECT * FROM delta.default." + tableName))
                     .containsOnly(row(1, 2, 3));
+
+            assertThat(onTrino().executeQuery("SELECT column_name, extra_info FROM delta.information_schema.columns WHERE table_schema = 'default' AND table_name = '" + tableName + "'"))
+                    .containsOnly(row("a", null), row("b", "generated: a * 2"), row("c", null));
+            assertThat(onTrino().executeQuery("DESCRIBE delta.default." + tableName).project(1, 3))
+                    .containsOnly(row("a", ""), row("b", "generated: a * 2"), row("c", ""));
+            assertThat(onTrino().executeQuery("SHOW COLUMNS FROM delta.default." + tableName).project(1, 3))
+                    .containsOnly(row("a", ""), row("b", "generated: a * 2"), row("c", ""));
         }
         finally {
             dropDeltaTableWithRetry("default." + tableName);