Pandas DataFrame aggregate a list column to a set

Code description

This code snippet shows you how to group a pandas DataFrame and then aggregate a column with list or array type to a set (with duplicates removed) or a list. To implement it, we can first expode the column (list type) and then use groupby to create a grouped DataFrame and then aggregate using set or list or a combination of both.

Input

      category      users
    0        A     [1, 2]
    1        B     [3, 4]
    2        C  [5, 6, 7]
    3        A  [1, 8, 1]
    4        B  [1, 6, 9]

Output

      category        users_set       users_list
    0        A        {8, 1, 2}        [8, 1, 2]
    1        B  {1, 3, 4, 6, 9}  [1, 3, 4, 6, 9]
    2        C        {5, 6, 7}        [5, 6, 7]

Code snippet

    import pandas as pd
    
    data = {
        "category": ['A', 'B', 'C', 'A', 'B'],
        "users": [
            [1,2],
            [3,4],
            [5,6,7],
            [1,8,1],
            [1,6,9]
        ]
    }
    
    df = pd.DataFrame(data)
    print(df)
    
    # Expode the array/list column
    df = df.explode('users')
    print(df)
        
    # Aggregate and consolidate users of each category into a set
    df = df.groupby('category')['users'].agg(
        users_set = lambda x: set(x),
        users_list = lambda x: list(set(x))
    ).reset_index()
    print(df)