From 5270a3fc031448f6f807fc398e7579596cb221a0 Mon Sep 17 00:00:00 2001 From: Yiwen Yuan Date: Tue, 23 Jul 2024 21:08:03 +0000 Subject: [PATCH] drop dangling fkey in user_friends table --- relbench/datasets/event.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/relbench/datasets/event.py b/relbench/datasets/event.py index 003fbc62..09cc95dd 100644 --- a/relbench/datasets/event.py +++ b/relbench/datasets/event.py @@ -81,6 +81,13 @@ def make_db(self) -> Database: user=lambda df: df["user"].astype(int), friend=lambda df: df["friend"].astype(int), ) + + # Some friends are not present in the user table, so we drop those friends + # in the user_friends table + user_friends_flattened_df = user_friends_flattened_df.merge( + users_df, how="inner", left_on="friend", right_on="user_id" + ) + user_friends_flattened_df = user_friends_flattened_df[["user", "friend"]] user_friends_flattened_df.to_csv( os.path.join(path, "user_friends_flattened.csv") )