[Python] How to optimally join a DataFrame on a column with comma separated IDs

I have a dataframe of users something like this:

Name Sex AgeGroup SuitablePlaces
User1 Male 30 1,4
User2 Female 40 5
User3 Male 50 1,3
User4 Male 40 1,3,4
User5 Female 30 [null]
User6 Female 50
User7 Male 30 3

It has columns [Name, Sex, AgeGroup and SuitablePlaces] where

• AgeGroup is an id essentially like 20,30,40,50
• SuitablePlaces is also an Id, but it can be:
  • single comma separated id (1)
  • multiple comma separated ids (1,4,5)
  • null
  • empty

I have another Dataframe with the AgeGroup to PlaceId mapping and Place name:

AgeGroup PlaceID PlaceName
30 1 Place1
30 5 Place2
40 3 Place3
40 5 Place4
50 4 Place5
50 2 Place6

It has columns [AgeGroup, PlaceId, PlaceName] where essentially AgeGroup and PlaceId together are a primary key.

I want to join the two tables on SuitablePlaces such that:

• if SuitablePlaces is [null] or empty, remove user from Dataframe
• if SuitablePlaces has single value:
  • find the AgeGroup and PlaceId in the Places Dataframe and add that to the Users DataFrame
  • if not found, remove user from Dataframe
• if SuitablePlaces has comma separated values:
  • split the values
  • find the first AgeGroup and PlaceId combo that matches in the Places Dataframe
  • add that to the Users DataFrame
  • if not found, remove user from Dataframe

I want to end up with a Dataframe like such in the end:

Name Sex AgeGroup SuitablePlaces PlaceName
User1 Male 30 1,4 Place1
User2 Female 40 5 Place4
User4 Male 40 1,3,4 Place3

Where Users what did not meet the criteria are not present. I am very new to DataFrame and not sure if this is possible without writing the logic myself. I have written code that does this. But it seems very inefficient and it may be possible to make it better. Also in my code i am replacing the suitablePlaces column instead of adding a new column called places.

Please let me know if i can improve the code below. Thanks!

def map_user_places(user_df, places_df):
suitable_places_field = "SuitablePlaces"
age_group_field = "AgeGroup"

user_df = user_df.copy()
places_df = places_df.copy()

def map_single_value(suitable_places, age_group):
"""
Map a denial code or comma-separated codes to their
corresponding category.
"""
if suitable_places is None or suitable_places == "":
return None

suitable_places_str = str(suitable_places).strip()

# Handle comma-separated values
if "," in suitable_places_str:
places = [
int(code.strip())
for code in suitable_places_str.split(",")]
filtered_df = places_df.query(
f"AgeGroup=={age_group} & PlaceId in {places}"
)
if len(filtered_df) > 0:
return filtered_df.iloc[0]["PlaceName"]
else:
return None

# Handle single value
filtered_df = places_df.query(
f"AgeGroup=={age_group} & PlaceId=={suitable_places_str}"
)
if len(filtered_df) > 0:
return filtered_df.iloc[0]["PlaceName"]
else:
return None

user_df[suitable_places_field] = user_df.apply(
lambda row: map_single_value(
row[suitable_places_field],
row[age_group_field]),
axis=1
)

df = df.dropna(subset=[suitable_places_field])

return df

Continue reading...