Pandas：合并数据帧并将多个连接值合并到一个数组中

Question

我是Python的新手，我正在使用Pandas将一堆MySQL表转换为JSON。 我目前的解决方案工作得很好，但（1）它不是非常pythonic，（2）我觉得必须有一些预先出炉的熊猫功能做我需要的......？ 对以下问题的任何指导都会有所帮助。

假设我有两个数据框， authors和连接表plays_authors ，代表作者与戏剧的1：多关系。

print authors
>   author_id  dates notes
> 0         1  1700s     a 
> 1         2  1800s     b 
> 2         3  1900s     c 


print plays_authors
>      author_id play_id
> 0         1      12
> 1         1      13
> 2         1      21
> 3         2      18
> 4         3       3
> 5         3       7

我想将plays_authors合并到authors plays_authors ，但不是每个作者有多行（每个play_id 1个），我想要每个作者一行，并且有一个play_id值数组，以便我可以轻松地将它们导出为json记录。

print authors
>   author_id  dates notes       play_id
> 0         1  1700s     a  [12, 13, 21]
> 1         2  1800s     b          [18]
> 2         3  1900s     c        [3, 7]

authors.to_json(orient="records")
> '[{
>    "author_id":"1",
>    "dates":"1700s",
>    "notes":"a",
>    "play_id":["12","13","21"]
>   },
>   {
>    "author_id":"2",
>    "dates":"1800s",
>    "notes":"b",
>    "play_id":["18"]
>   },
>   {
>    "author_id":"3",
>    "dates":"1900s",
>    "notes":"c",
>    "play_id":["3","7"]
>  }]'

我目前的解决方案

# main_df: main dataframe to transform
# join_df: the dataframe of the join table w values to add to df
# main_index: name of main_df index column
# multi_index: name of column w/ multiple values per main_index, added by merge with join_df
# jointype: type of merge to perform, e.g. left, right, inner, outer

def consolidate(main_df, join_df, main_index, multi_index, jointype):
    # merge
    main_df = pd.merge(main_df, join_df, on=main_index, how=jointype)
    # consolidate
    new_df = pd.DataFrame({})

    for i in main_df[main_index].unique():
        i_rows = main_df.loc[main_df[main_index] == i]
        values = []

        for column in main_df.columns:
            values.append(i_rows[:1][column].values[0])

        row_dict = dict(zip(main_df.columns, values))
        row_dict[multi_index] = list(i_rows[multi_index])
        new_df = new_df.append(row_dict, ignore_index=True)

    return new_df


authors = consolidate(authors, plays_authors, 'author_id', 'play_id', 'left')

那里有一个简单的groupby /更好的dict解决方案吗？

Answer 1

数据：

In [131]: a
Out[131]:
   author_id  dates notes
0          1  1700s     a
1          2  1800s     b
2          3  1900s     c

In [132]: pa
Out[132]:
   author_id  play_id
0          1       12
1          1       13
2          1       21
3          2       18
4          3        3
5          3        7

解：

In [133]: a.merge(pa.groupby('author_id')['play_id'].apply(list).reset_index())
Out[133]:
   author_id  dates notes       play_id
0          1  1700s     a  [12, 13, 21]
1          2  1800s     b          [18]
2          3  1900s     c        [3, 7]