优化Pandas多DataFrame按区域编码匹配累计体积的日期映射方法

示例数据 #

import pandas as pd
import numpy as np

# 订单数据
orders = pd.DataFrame({'area':['a','b','c'], 'code': [177, 23,10], 'requires':[10.0,3.0,6.0]})
# 每日入库数据
graphic = pd.DataFrame({'area': ['a','a','a','c','c'], 'code': [177, 177,177,10,10], 
                        'date':['2022-06-11', '2022-06-12','2022-06-13','2022-07-01','2022-07-03'], 
                        'volume':[7.0,2.0,1.0,6.0,3.0]})

# orders输出：
#   area  code  requires
# 0    a   177      10.0
# 1    b    23       3.0
# 2    c    10       6.0

# graphic输出：
#   area  code        date  volume
# 0    a   177  2022-06-11     7.0
# 1    a   177  2022-06-12     2.0
# 2    a   177  2022-06-13     1.0
# 3    c    10  2022-07-01     6.0
# 4    c    10  2022-07-03     3.0

需求说明 #

以area和code为匹配键，找到graphic中累计volume首次满足orders的requires的最早日期；若累计体积不足则返回NaN。具体规则：

• area='a'、code='177'需求10.0：累计到2022-06-13时体积(7+2+1=10)达标，对应日期为2022-06-13
• area='b'无匹配数据，返回NaN
• area='c'、code='10'需求6.0：2022-07-01的体积已满足，对应日期为2022-07-01

期望结果：

area  code  requires        date
0    a   177      10.0  2022-06-13
1    b    23       3.0         NaN
2    c    10       6.0  2022-07-01

现有实现的性能问题 #

当前实现通过逐行遍历处理，在2万行数据场景下速度极慢，代码如下：

def data_shipment(city, code, required_to_make):
    if graphic[(graphic['area']==city) & (graphic['code']==code)].empty:
        return np.nan
    elif required_to_make == 0:
        return np.nan   
    elif graphic.loc[(graphic['area']==city) & (graphic['code']==code), 'volume'].iloc[0] >= required_to_make:   
        data_of_making = graphic.loc[(graphic['area']==city) & (graphic['code']==code) & (graphic['date']==graphic.loc[(graphic['area']==city) 
                                    & (graphic['code']==code), 'date'].iloc[0]),'date'].iloc[0]
        graphic.loc[(graphic['area']==city) & (graphic['code']==code) & (graphic['date']==data_of_making),
                    'volume'] -= required_to_make
        data = graphic[(graphic['area']==city) & (graphic['code']==code)]['date'].iloc[0]
        return data
    else:
        if len(graphic.loc[(graphic['area']==city) & (graphic['code']==code)]) == 1:
            return np.nan
        elif len(graphic.loc[(graphic['area']==city) & (graphic['code']==code)]) > 1:
            len_of_dates = len(graphic.loc[(graphic['area']==city) & (graphic['code']==code)])
            i = 0
            storage_var = graphic.loc[(graphic['area']==city) & (graphic['code']==code), 'volume'].iloc[0]
            list_of_dates_to_be_deleted = []
            while storage_var<required_to_make and len_of_dates>1:
                len_of_dates-=1
                i+=1
                storage_var+= graphic.loc[(graphic['area']==city) & (graphic['code']==code), 'volume'].iloc[i]
                list_of_dates_to_be_deleted.append(graphic.loc[(graphic['area']==city) & (graphic['code']==code) & 
                                             (graphic['date']==graphic.loc[(graphic['area']==city) 
                                            & (graphic['code']==code), 'date'].iloc[i-1]),
                                             'date'].iloc[0])                    
            if storage_var<required_to_make:
                return np.nan
            else:
                for j in (list_of_dates_to_be_deleted):
                    graphic.drop(graphic[(graphic['area']==city) & (graphic['code']==code) & 
                           (graphic['date']==j)].index, inplace=True)
                    i-=1
                data_of_making = graphic.loc[(graphic['area']==city) & (graphic['code']==code) & 
                                             (graphic['date']==graphic.loc[(graphic['area']==city) 
                                            & (graphic['code']==code), 'date'].iloc[i]),
                                             'date'].iloc[0]
            graphic.loc[(graphic['area']==city) & (graphic['code']==code) & 
                        (graphic['date']==data_of_making), 'volume'] = storage_var-required_to_make
 
            data = graphic[(graphic['area']==city) & (graphic['code']==code)]['date'].iloc[i] 
            return data

高效实现方案 #

利用Pandas的分组、窗口函数和合并操作，实现矢量化计算，避免逐行遍历：

步骤1：预处理graphic数据 #

按area、code分组，计算累计体积，并保留日期顺序：

# 先确保date是日期类型
graphic['date'] = pd.to_datetime(graphic['date'])

# 分组计算累计volume，同时保留日期
graphic['cum_volume'] = graphic.groupby(['area', 'code'])['volume'].cumsum()

步骤2：筛选每个分组中首次满足条件的日期 #

对每个area+code分组，找到累计体积>=对应requires的最早日期：

# 先将orders和graphic合并，保留所有orders的行
merged = orders.merge(graphic, on=['area', 'code'], how='left')

# 筛选出累计体积满足要求的行，然后按area、code分组取最早的date
result = merged[merged['cum_volume'] >= merged['requires']].groupby(['area', 'code', 'requires'])['date'].min().reset_index()

# 补全没有满足条件的行（返回NaN）
final_result = orders.merge(result, on=['area', 'code', 'requires'], how='left')

步骤3：查看结果 #

print(final_result)
# 输出：
#   area  code  requires       date
# 0    a   177      10.0 2022-06-13
# 1    b    23       3.0        NaT
# 2    c    10       6.0 2022-07-01

如果需要把NaT转为NaN，可以执行：

final_result['date'] = final_result['date'].astype('object').where(final_result['date'].notna(), np.nan)

补充：模拟库存消耗（扣减已使用的volume） #

如果需要像原代码那样扣减graphic中的库存，可以在找到达标日期后，对graphic进行批量更新：

# 先获取每个分组需要扣减的总量
required_map = orders.set_index(['area', 'code'])['requires'].to_dict()

# 遍历每个分组
for (area, code), group in graphic.groupby(['area', 'code']):
    if (area, code) not in required_map:
        continue
    req = required_map[(area, code)]
    cum = group['cum_volume'].values
    # 找到首次达标位置
    idx = np.argmax(cum >= req)
    if cum[idx] < req:
        continue
    # 扣减前面的全部volume，最后一行扣减剩余量
    graphic.loc[group.index[:idx], 'volume'] = 0
    graphic.loc[group.index[idx], 'volume'] = cum[idx] - req

这样处理后，graphic中的库存会更新为扣除需求后的剩余量，同时整体计算为矢量化操作，处理2万行数据的速度会大幅提升。

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