Repository: reata/sqllineage (python)

Issue:
Failure to Handle Multiple Lineage Path for Same Column
Why holders.get_column_lineage return simple_paths which equal to one only? Are there any reasons?
In my case, i have a target column (i.e. with alias X) which is the result of a CASE construction, and all of those columns (A,B,C) are from same column and same table (i.e. BASE_TABLE.Z), so result of nx.all_simple_paths returns the next paths:
1) BASE_TABLE.Z->A->X
2) BASE_TABLE.Z->B->X
3) BASE_TABLE.Z->C->X

But get_column_lineage not return any one, the reason is the following lines:
if len(simple_paths) == 1:
    columns.add(tuple(simple_paths[0]))

(simplified example)
CASE
when JOIN_TABLE_1.A not null JOIN_TABLE_1.A
when JOIN_TABLE_2.B not null JOIN_TABLE_2.B
else JOIN_TABLE_3.C
END as X
JOIN SELECT Z as A FROM BASE_TABLE as JOIN_TABLE_1
JOIN SELECT Z as B FROM BASE_TABLE as JOIN_TABLE_2
JOIN SELECT Z as C FROM BASE_TABLE as JOIN_TABLE_3

Thank you in advance!


Pull request description:
fix: handle multiple lineage paths for same column
Closes #228 

Patch:
diff --git a/sqllineage/core/holders.py b/sqllineage/core/holders.py
index d82a013..7bd0b8c 100644
--- a/sqllineage/core/holders.py
+++ b/sqllineage/core/holders.py
@@ -30,9 +30,8 @@ class ColumnLineageMixin:
         columns = set()
         for (source, target) in itertools.product(source_columns, target_columns):
             simple_paths = list(nx.all_simple_paths(self.graph, source, target))
-            if len(simple_paths) == 1:
-                columns.add(tuple(simple_paths[0]))
-            # we can ignore when simple path doesn't exist, but could there be more than one simple path?
+            for path in simple_paths:
+                columns.add(tuple(path))
         return columns
 
 


Test patch:
diff --git a/tests/test_columns.py b/tests/test_columns.py
index 99ec782..1885d6f 100644
--- a/tests/test_columns.py
+++ b/tests/test_columns.py
@@ -737,3 +737,29 @@ FROM tab2"""
             ),
         ],
     )
+
+
+def test_column_lineage_multiple_paths_for_same_column():
+    sql = """INSERT OVERWRITE TABLE tab2
+SELECT tab1.id,
+       coalesce(join_table_1.col1, join_table_2.col1, join_table_3.col1) AS col1
+FROM tab1
+         LEFT JOIN (SELECT id, col1 FROM tab1 WHERE flag = 1) AS join_table_1
+                   ON tab1.id = join_table_1.id
+         LEFT JOIN (SELECT id, col1 FROM tab1 WHERE flag = 2) AS join_table_2
+                   ON tab1.id = join_table_2.id
+         LEFT JOIN (SELECT id, col1 FROM tab1 WHERE flag = 3) AS join_table_3
+                   ON tab1.id = join_table_3.id"""
+    assert_column_lineage_equal(
+        sql,
+        [
+            (
+                ColumnQualifierTuple("id", "tab1"),
+                ColumnQualifierTuple("id", "tab2"),
+            ),
+            (
+                ColumnQualifierTuple("col1", "tab1"),
+                ColumnQualifierTuple("col1", "tab2"),
+            ),
+        ],
+    )