Pandas groupby聚合不同形状数组列求和报广播错误如何解决

使用pandas的aggregate逻辑对分组后列的所有值求和，测试构造DataFrame代码如下：

import pandas as pd
import numpy as np

df2 = pd.DataFrame([[1,  np.array([2, 5, 3])],
                   [1,  np.array([2, 5, 3])],
                   [1, np.array([2, 5, 3, 5, 3])]],
                 columns=['doc_id', 'topic_dist'])

初始编写的聚合代码如下，目标是对不同形状的数组完成聚合计算：

def getsumcolumns(dfsource):
    grouped = dfsource.groupby('doc_id')
    aggregate = list((k, v["topic_dist"].sum()) for k, v in grouped) 
    df_results = pd.DataFrame(aggregate, columns=['doc_id','topic_dist'])
    print(df_results)
    return df_results 

运行代码触发如下错误：

operands could not be broadcast together with shapes (3,) (5,) 

预期得到的输出结果为：

doc_id  topic_dist
0       1  [6, 15, 9, 5, 3]

直接对存储numpy数组的Series调用.sum()方法时，pandas会沿用numpy的数组运算逻辑，要求所有参与运算的数组形状一致，按位置做逐元素相加。当前同组内的数组长度分别为3、3、5，不满足numpy广播规则，因此触发形状不匹配报错。

分组后对每组内的数组，先按组内最长数组的长度做对齐，长度不足的位置默认补0，再逐位求和即可。修改后的聚合代码如下：

from itertools import zip_longest

def getsumcolumns(dfsource):
    agg_result = []
    for doc_id, group_data in dfsource.groupby('doc_id'):
        # 按位置对齐所有数组，缺位补0后逐位求和
        sum_array = np.array([
            sum(position_values) 
            for position_values in zip_longest(*group_data['topic_dist'], fillvalue=0)
        ])
        agg_result.append((doc_id, sum_array))
    df_results = pd.DataFrame(agg_result, columns=['doc_id','topic_dist'])
    print(df_results)
    return df_results

传入测试DataFrame调用该函数，即可得到符合预期的输出结果，该方案可以自动适配组内任意不同长度的数组，不需要手动指定数组维度。

内容的提问来源于stack exchange，提问作者Wiseman2022