使用pandas rolling+apply处理数组列遇DataError，如何实现自定义函数？

问题场景 #

现有如下Pandas DataFrame：

import pandas as pd

my_dict = {
    'id': [1, 2, 3, 4, 5],
    'salary': [100, 200, 300, 400, 500],
    'location_vector' : [[1, 2, 3], [2, 3, 4], [2, 3, 5], [4, 5, 4], [7, 5, 4]]
}

df = pd.DataFrame(my_dict)

需要对location_vector列逐行应用自定义函数my_func（接收两个数组，返回新数组），实现类似rolling(2)的效果——即每一行用当前行数组和上一行数组作为参数传入函数。原尝试代码：

df['result'] = df['location_vector'].rolling(2).apply(lambda x, y: my_func(x, y))

运行报错：DataError: No numeric types to aggregate。

错误原因 #

1. rolling.apply的回调函数仅接收一个参数：窗口内的所有元素集合（而非拆分好的两个数组），原lambda写法参数数量错误。
2. Pandas的rolling默认针对数值类型列做聚合计算，但location_vector是存储数组的object类型列，不符合默认聚合要求，因此抛出类型错误。

解决方案 #

方法一：使用shift()实现（推荐，更高效直观） #

利用shift()获取上一行的数组，再逐行调用自定义函数：

# 定义示例自定义函数（可替换为你的业务逻辑）
def my_func(arr1, arr2):
    # 示例：两个数组对应元素相减
    return [a - b for a, b in zip(arr1, arr2)]

# 获取上一行的location_vector，第一行自动为NaN
df['prev_vector'] = df['location_vector'].shift(1)

# 逐行应用函数，跳过第一行（无前置数组的情况）
df['result'] = df.apply(
    lambda row: my_func(row['location_vector'], row['prev_vector']) if pd.notna(row['prev_vector']) else None,
    axis=1
)

# 清理临时列
df.drop('prev_vector', axis=1, inplace=True)

执行后得到的结果符合预期：

id  salary location_vector     result
0   1     100      [1, 2, 3]       None
1   2     200      [2, 3, 4]  [1, 1, 1]
2   3     300      [2, 3, 5]  [0, 0, 1]
3   4     400      [4, 5, 4]  [2, 2, -1]
4   5     500      [7, 5, 4]  [3, 0, 0]

方法二：修改rolling.apply的用法 #

如果坚持使用rolling，需调整回调函数逻辑，直接处理窗口内的元素集合，并设置raw=False以保留原数据结构：

def rolling_func(window):
    # window是窗口内的Series，包含当前行和上一行的数组
    if len(window) < 2:
        return None
    # 取出当前行（index=1）和上一行（index=0）的数组
    current_arr, prev_arr = window.iloc[1], window.iloc[0]
    return my_func(current_arr, prev_arr)

df['result'] = df['location_vector'].rolling(2).apply(rolling_func, raw=False)

此方法同样能实现需求，但效率略低于shift()方案，适合必须使用rollingAPI的场景。

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