Make apply return two series

Question

Say I have the following dataframe

id | dict_col
---+---------
1    {"age":[1,2],"name":["john","doe"]}
2    {"age":[3,4],"name":["foo","bar"]}

and the following

def unnest_dict(row):
    age = row["age"]
    name = row["name"]
    return age,name

age,name = df["dict_col"].apply(unnest_dict)

that throws a ValueError:

ValueError: too many values to unpack (expected 2)

I can wrap the return into a tuple (age,name) but then I need to loop over each of them afterwards like so

def unnest_dict(row):
    .
    .
    return (age,name)
data = df["dict_col"].apply(unnest_dict)
age = [p[0] for p in data]
name = [p[1] for p in data]

print(age)
# [[1,2],[3,4]]

print(name)
# [["john","doe"],["foo","bar"]]

Isnt there a way to extract the series directly from the function?

NOTE: This is not the same question as this SO since that it is specific of how to explode a dict - my question is regarding how to extract two (or more) series directly from the function-return. The example provided is just an example, and could've been any operation.

Not quite the same. I just want a function to return a series - and not elements in the series — CutePoison, Apr 21 '21 at 07:40

score 0 · Answer 1 · answered Apr 21 '21 at 07:56

import pandas as pd

df = pd.DataFrame({'id': [1,2], 'dict_col': [ {"age":[1,2],"name":["john","doe"]}, {"age":[3,4],"name":["foo","bar"]}]})

def unnest_dict(row):
    age = row['dict_col']["age"]
    name = row['dict_col']["name"]
    return {'age': age,'name': name}


result = df.apply(unnest_dict, axis=1, result_type='expand')
age = result['age'].tolist()
name = result['name'].tolist()
print(age)
print(name)

Output:

[[1, 2], [3, 4]]
[['john', 'doe'], ['foo', 'bar']]

Make apply return two series

1 Answers1