fix(parse_edge_data): error when edge type has _

tigergraph · Jan 2, 2024 · 4d334a1 · 4d334a1
1 parent 7e0de75
commit 4d334a1
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/pyTigerGraph/gds/dataloaders.py b/pyTigerGraph/gds/dataloaders.py
@@ -1044,10 +1044,9 @@ def _parse_edge_data(
             # If seeds are given, create the is_seed column
             if seeds:
                 seed_df = pd.DataFrame.from_records(
-                    [i.split("_") for i in seeds], 
-                    columns=["source", "etype", "target"])
+                    [(i.split("_", 1)[0], i.rsplit("_", 1)[-1]) for i in seeds], 
+                    columns=["source", "target"])
                 seed_df["is_seed"] = True
-                del seed_df["etype"]
                 data = data.merge(seed_df, on=["source", "target"], how="left")
                 data.fillna({"is_seed": False}, inplace=True)
         else:
@@ -1058,7 +1057,7 @@ def _parse_edge_data(
             # If seeds are given, create the is_seed column
             if seeds:
                 seed_df = pd.DataFrame.from_records(
-                    [i.split("_") for i in seeds], 
+                    [(i.split("_", 1)[0], i.split("_", 1)[1].rsplit("_", 1)[0], i.rsplit("_", 1)[-1]) for i in seeds], 
                     columns=["source", "etype", "target"])
                 seed_df["is_seed"] = True
             for etype in e_file_dict: