Fix column level lineage miss when use unnest

Author: hqbhoho

core/trino-main/src/main/java/io/trino/sql/analyzer/StatementAnalyzer.java

@@ -1656,6 +1656,7 @@ else if (expressionType instanceof MapType mapType) {
 
                 outputFields.addAll(expressionOutputs);
                 mappings.put(NodeRef.of(expression), expressionOutputs);
+                expressionOutputs.forEach(field -> analysis.addSourceColumns(field, analysis.getExpressionSourceColumns(expression)));
             }
 
             Optional<Field> ordinalityField = Optional.empty();

testing/trino-tests/src/test/java/io/trino/execution/TestEventListenerBasic.java

@@ -51,6 +51,7 @@
 import io.trino.spi.eventlistener.TableInfo;
 import io.trino.spi.metrics.Metrics;
 import io.trino.spi.security.ViewExpression;
+import io.trino.spi.type.ArrayType;
 import io.trino.spi.type.Type;
 import io.trino.spi.type.TypeManager;
 import io.trino.spi.type.TypeSignature;
@@ -156,7 +157,9 @@ public Iterable<ConnectorFactory> getConnectorFactories()
                             }
                             return ImmutableList.of(
                                     new ColumnMetadata("test_varchar", createVarcharType(15)),
-                                    new ColumnMetadata("test_bigint", BIGINT));
+                                    new ColumnMetadata("test_bigint", BIGINT),
+                                    new ColumnMetadata("test_varchar_array", new ArrayType(createVarcharType(15))),
+                                    new ColumnMetadata("test_bigint_array", new ArrayType(BIGINT)));
                         })
                         .withGetTableHandle((session, schemaTableName) -> {
                             if (!schemaTableName.getTableName().startsWith("create")) {
@@ -761,7 +764,7 @@ public void testReferencedTablesWithColumnMask()
             throws Exception
     {
         QueryEvents queryEvents = runQueryAndWaitForEvents(
-                "CREATE TABLE mock.default.create_table_with_referring_mask AS SELECT * FROM mock.default.test_table_with_column_mask"
+                "CREATE TABLE mock.default.create_table_with_referring_mask AS SELECT test_varchar, test_bigint FROM mock.default.test_table_with_column_mask"
         ).getQueryEvents();
 
         QueryCompletedEvent event = queryEvents.getQueryCompletedEvent();
@@ -1372,7 +1375,9 @@ public void testCreateTableLike()
                 .containsExactly(
                         new OutputColumnMetadata("test_column", BIGINT_TYPE, ImmutableSet.of()),
                         new OutputColumnMetadata("test_varchar", VARCHAR_TYPE, ImmutableSet.of()),
-                        new OutputColumnMetadata("test_bigint", BIGINT_TYPE, ImmutableSet.of()));
+                        new OutputColumnMetadata("test_bigint", BIGINT_TYPE, ImmutableSet.of()),
+                        new OutputColumnMetadata("test_varchar_array", "array(varchar(15))", ImmutableSet.of()),
+                        new OutputColumnMetadata("test_bigint_array", "array(bigint)", ImmutableSet.of()));
     }
 
     @Test
@@ -1419,6 +1424,41 @@ private void testOutputColumnsForSetOperations(String setOperator)
                                 new ColumnDetail("tpch", "sf1", "orders", "custkey"))));
     }
 
+    @Test
+    public void testOutputColumnsWithUnnestQueries()
+            throws Exception
+    {
+        // test unnest with one array column
+        assertLineage(
+                "SELECT test_varchar_unnest AS test_varchar, test_bigint AS test_bigint FROM mock.default.tests_table_unnest CROSS JOIN UNNEST(test_varchar_array) AS t(test_varchar_unnest)",
+                ImmutableSet.of("mock.default.tests_table_unnest"),
+                new OutputColumnMetadata("test_varchar", VARCHAR_TYPE, ImmutableSet.of(new ColumnDetail("mock", "default", "tests_table_unnest", "test_varchar_array"))),
+                new OutputColumnMetadata("test_bigint", BIGINT_TYPE, ImmutableSet.of(new ColumnDetail("mock", "default", "tests_table_unnest", "test_bigint"))));
+        // test unnest with one array column and with ordinality
+        assertLineage(
+                "SELECT test_varchar_unnest AS test_varchar, row_number_unnest AS test_bigint FROM mock.default.tests_table_unnest CROSS JOIN UNNEST(test_varchar_array) WITH ORDINALITY AS t(test_varchar_unnest, row_number_unnest)",
+                ImmutableSet.of("mock.default.tests_table_unnest"),
+                new OutputColumnMetadata("test_varchar", VARCHAR_TYPE, ImmutableSet.of(new ColumnDetail("mock", "default", "tests_table_unnest", "test_varchar_array"))),
+                new OutputColumnMetadata("test_bigint", BIGINT_TYPE, ImmutableSet.of()));
+        // test unnest with two array column
+        assertLineage(
+                "SELECT test_varchar_unnest AS test_varchar, test_bigint_unnest AS test_bigint FROM mock.default.tests_table_unnest CROSS JOIN UNNEST(test_varchar_array, test_bigint_array) AS t(test_varchar_unnest, test_bigint_unnest)",
+                ImmutableSet.of("mock.default.tests_table_unnest"),
+                new OutputColumnMetadata("test_varchar", VARCHAR_TYPE, ImmutableSet.of(new ColumnDetail("mock", "default", "tests_table_unnest", "test_varchar_array"))),
+                new OutputColumnMetadata("test_bigint", BIGINT_TYPE, ImmutableSet.of(new ColumnDetail("mock", "default", "tests_table_unnest", "test_bigint_array"))));
+        // test unnest with two array column and with ordinality
+        assertLineage(
+                "SELECT test_varchar_unnest AS test_varchar, row_number_unnest AS test_bigint FROM mock.default.tests_table_unnest CROSS JOIN UNNEST(test_varchar_array, test_bigint_array)  WITH ORDINALITY AS t(test_varchar_unnest, test_bigint_unnest, row_number_unnest)",
+                ImmutableSet.of("mock.default.tests_table_unnest"),
+                new OutputColumnMetadata("test_varchar", VARCHAR_TYPE, ImmutableSet.of(new ColumnDetail("mock", "default", "tests_table_unnest", "test_varchar_array"))),
+                new OutputColumnMetadata("test_bigint", BIGINT_TYPE, ImmutableSet.of()));
+        assertLineage(
+                "SELECT CAST(test_bigint_unnest AS varchar(15)) AS test_varchar, row_number_unnest AS test_bigint FROM mock.default.tests_table_unnest CROSS JOIN UNNEST(test_varchar_array, test_bigint_array)  WITH ORDINALITY AS t(test_varchar_unnest, test_bigint_unnest, row_number_unnest)",
+                ImmutableSet.of("mock.default.tests_table_unnest"),
+                new OutputColumnMetadata("test_varchar", VARCHAR_TYPE, ImmutableSet.of(new ColumnDetail("mock", "default", "tests_table_unnest", "test_bigint_array"))),
+                new OutputColumnMetadata("test_bigint", BIGINT_TYPE, ImmutableSet.of()));
+    }
+
     @Test
     public void testTableStats()
             throws Exception