Derive a FLATTEN's KEY and PATH from its input

crates/flowscope-core/src/analyzer/visitor.rs

@@ -472,9 +472,10 @@ impl<'a, 'b> LineageVisitor<'a, 'b> {
     /// Registers Snowflake's `LATERAL FLATTEN(INPUT => expr) AS f` as a
     /// relation with the columns FLATTEN returns, so that `f.value` resolves.
     ///
-    /// VALUE and THIS carry the input, so they derive from the columns it
-    /// reads, with the whole call as their expression. SEQ, KEY, PATH and
-    /// INDEX say where an element sits rather than carry it, and read nothing.
+    /// VALUE and THIS carry the input, and KEY and PATH name where an element
+    /// sits in it, which for an object is part of its content: all four derive
+    /// from the columns it reads, with the whole call as their expression. SEQ
+    /// and INDEX number the rows and the elements, and read nothing.
     fn visit_flatten(
         &mut self,
         name: &ast::ObjectName,
@@ -530,8 +531,9 @@ impl<'a, 'b> LineageVisitor<'a, 'b> {
         );
 
         let checkpoint = self.ctx.projection_checkpoint();
+        let mut first_carrier = true;
         for column in ["SEQ", "KEY", "PATH", "INDEX", "VALUE", "THIS"] {
-            let carries_input = matches!(column, "VALUE" | "THIS");
+            let carries_input = matches!(column, "KEY" | "PATH" | "VALUE" | "THIS");
             if carries_input && !sources.is_empty() {
                 let reported = self.analyzer.issues.len();
                 self.analyzer.add_output_column_with_aggregation(
@@ -546,11 +548,12 @@ impl<'a, 'b> LineageVisitor<'a, 'b> {
                         aggregation: None,
                     },
                 );
-                // THIS reads the input VALUE has just read, so whatever resolving
-                // it reports has been reported once already.
-                if column == "THIS" {
+                // Each reads the input the first one has just read, so whatever
+                // resolving it reports has been reported once already.
+                if !first_carrier {
                     self.analyzer.issues.truncate(reported);
                 }
+                first_carrier = false;
             } else {
                 self.analyzer
                     .add_unsourced_column(self.ctx, column, &node_id);

crates/flowscope-core/tests/lineage_engine.rs

@@ -9939,6 +9939,8 @@ fn flatten_value_comes_from_its_input_not_from_a_namesake() {
     assert_eq!(
         expressions,
         HashSet::from([
+            ("key", Some("FLATTEN(input => o.tags)")),
+            ("path", Some("FLATTEN(input => o.tags)")),
             ("value", Some("FLATTEN(input => o.tags)")),
             ("this", Some("FLATTEN(input => o.tags)")),
         ])
@@ -10029,7 +10031,32 @@ fn flatten_position_columns_read_nothing() {
         .into_iter()
         .map(|(column, _)| column)
         .collect();
-    assert_eq!(fed, expected_sources(&["value", "this"]));
+    assert_eq!(fed, expected_sources(&["key", "path", "value", "this"]));
+}
+
+#[test]
+fn flatten_key_and_path_come_from_the_input() {
+    // An object's keys are its content: flattening parsed JSON, the key and
+    // the path are read from the column the JSON is in. The index numbers the
+    // elements and reads nothing.
+    let sql = r#"
+        SELECT h.id, f.key AS element_name, f.path AS element_path,
+            f.value::STRING AS element_value, f.index AS position
+        FROM hx h,
+        LATERAL FLATTEN(input => TRY_PARSE_JSON(h.data)) f
+    "#;
+    let schema = schema_of(vec![schema_table(None, None, "hx", &["id", "data"])]);
+
+    let result = run_analysis(sql, Dialect::Snowflake, Some(schema));
+    let stmt = first_statement(&result);
+    for output in ["element_name", "element_path", "element_value"] {
+        assert_eq!(
+            base_sources_of(&stmt, output),
+            expected_sources(&["hx.data"]),
+            "{output}"
+        );
+    }
+    assert!(base_sources_of(&stmt, "position").is_empty());
 }
 
 #[test]
@@ -10104,14 +10131,14 @@ fn star_over_flatten_returns_its_columns() {
         columns,
         expected_sources(&["id", "arr", "seq", "key", "path", "index", "value", "this"])
     );
-    for carried in ["value", "this"] {
+    for carried in ["key", "path", "value", "this"] {
         assert_eq!(
             base_sources_of(&stmt, carried),
             expected_sources(&["items.arr"]),
             "sources of {carried}"
         );
     }
-    for position in ["seq", "key", "path", "index"] {
+    for position in ["seq", "index"] {
         assert!(
             base_sources_of(&stmt, position).is_empty(),
             "{position} should read nothing"

crates/flowscope-core/tests/snapshots/snapshots__run_snowflake_snapshot_test@snowflake_lateral_flatten.snap

@@ -21,7 +21,8 @@ expression: clean_result
       },
       "statementIds": [
         0
-      ]
+      ],
+      "expression": "FLATTEN(input => b.cool_ids)"
     },
     {
       "id": "column_268ea64f9e076ed3",
@@ -80,7 +81,8 @@ expression: clean_result
       },
       "statementIds": [
         0
-      ]
+      ],
+      "expression": "FLATTEN(input => b.cool_ids)"
     },
     {
       "id": "column_9312b887682fcd76",
@@ -308,6 +310,30 @@ expression: clean_result
         0
       ]
     },
+    {
+      "id": "edge_b92c790c00e1761f",
+      "from": "column_268ea64f9e076ed3",
+      "to": "column_8b8a696c1d07ba23",
+      "type": "derivation",
+      "expression": "FLATTEN(input => b.cool_ids)",
+      "joinType": "INNER",
+      "joinCondition": "b.c_id = a.c_id",
+      "statementIds": [
+        0
+      ]
+    },
+    {
+      "id": "edge_bc9ac2bfeded5339",
+      "from": "column_268ea64f9e076ed3",
+      "to": "column_1f95f87efeb20682",
+      "type": "derivation",
+      "expression": "FLATTEN(input => b.cool_ids)",
+      "joinType": "INNER",
+      "joinCondition": "b.c_id = a.c_id",
+      "statementIds": [
+        0
+      ]
+    },
     {
       "id": "edge_cb4d48c39a8c8c26",
       "from": "flatten_c1bebf2d8b83f913",