Python识别3000万行订单数据年内连续购买留存年数的高效方法

开发者阿泽

2026-10-4

问题描述

我现有约3000万行数据，共包含6个字段：

DISTINCT_IRECIPIENTID | ORDERNUMBER | ORDERDATE | ORDERDATE_OF_NEXT_ORDER | RETAINED_OR_NOT

其中RETAINED_OR_NOT字段有三类取值：

"Retained for one year"：ORDERDATE与ORDERDATE_OF_NEXT_ORDER的差值≤365天
Next_purchase_but_not_retained：ORDERDATE与ORDERDATE_OF_NEXT_ORDER的差值>365天
"Only one lifetime purchase"：该用户终身仅产生1笔订单
我需要计算每个消费者连续留存的年数：例如某消费者下单4次，前3笔订单相邻间隔均≤1年，第3笔与第4笔间隔>1年，则前3笔对应的计数为3，最后1笔计数为0。
当前数据已按DISTINCT_IRECIPIENTID、ORDERDATE降序排序，我编写了如下代码，但执行速度极慢，请问有什么优化方案可以提升运行效率？

def find_consecutive_purchases_in_a_year(input):
    count = 0
    exit_loop = 0
    sub_data = prepared_main_data_backup[prepared_main_data_backup['DISTINCT_IRECIPIENTID']
                                         == input]

    for index, row in sub_data.iterrows():

        if exit_loop == 1:
            return count
        if exit_loop == 0:
            if row['RETAINED_OR_NOT'] == 'retained_for_one_year':
                count += 1
            else:
                exit_loop = 1
    return count


data_test = prepared_main_data_backup

data_test['retain_counter'] = data_test['DISTINCT_IRECIPIENTID'].apply(
    find_consecutive_purchases_in_a_year)

参考样例数据如下：

DISTINCT_IRECIPIENTID   TSORDERDATETIME FIRST_TRANS_DATE    ORDER_DATE_AFTER    DIFFERENCE_BETWEEN_ORDERS   RETAINED_OR_NOT Output
    1   2017-04-24-09.33.21.000000  2017-04-24-09.33.21.000000          only one lifetime purchase  0
    2   2017-04-24-09.35.16.000000  2017-04-24-09.35.16.000000          only one lifetime purchase  0
    3   2017-04-27-14.45.48.000000  2017-04-27-14.45.48.000000  2017-04-29-14.53.46.000000  2   retained_for_one_year   2
    3   2017-04-29-14.53.46.000000  2017-04-27-14.45.48.000000  2017-05-10-09.06.25.000000  11  retained_for_one_year   2
    3   2017-05-10-09.06.25.000000  2017-04-27-14.45.48.000000  2018-09-22-05.54.07.000000  500 next_purchase_but_not_retained  0
    3   2018-09-22-05.54.07.000000  2017-04-27-14.45.48.000000  2020-09-12-19.12.59.000000  721 next_purchase_but_not_retained  0
    3   2020-09-12-19.12.59.000000  2017-04-27-14.45.48.000000  2020-09-14-11.49.33.000000  2   retained_for_one_year   2
    3   2020-09-14-11.49.33.000000  2017-04-27-14.45.48.000000  2021-06-08-07.18.42.000000  267 retained_for_one_year   2
    4   2017-04-24-09.35.27.000000  2017-04-24-09.35.27.000000  2017-04-30-12.00.14.000000  6   retained_for_one_year   1
    4   2017-04-30-12.00.14.000000  2017-04-24-09.35.27.000000  2018-06-18-09.15.23.000000  414 next_purchase_but_not_retained  0

优化方案

原始代码性能差有两个核心原因：

每次apply调用都全表筛选对应用户的数据，时间复杂度达到O(N²)，3000万行量级下完全不可行
iterrows迭代本身速度极慢，进一步放大了耗时

直接用pandas向量化分组操作即可，时间复杂度为O(N)，3000万行数据常规配置下几分钟即可跑完：

import pandas as pd

# 按用户分组，给非留存行打分割标记，区分不同的连续留存序列
prepared_main_data_backup['break_point'] = (prepared_main_data_backup['RETAINED_OR_NOT'] != 'retained_for_one_year').groupby(prepared_main_data_backup['DISTINCT_IRECIPIENTID']).cumsum()

# 按用户+分割点分组，计算每个连续留存序列的长度
grouped = prepared_main_data_backup.groupby(['DISTINCT_IRECIPIENTID', 'break_point'])
prepared_main_data_backup['retain_counter'] = grouped['RETAINED_OR_NOT'].transform('count')

# 非留存行的计数直接置0
prepared_main_data_backup.loc[prepared_main_data_backup['RETAINED_OR_NOT'] != 'retained_for_one_year', 'retain_counter'] = 0

# 清理临时辅助列
prepared_main_data_backup.drop('break_point', axis=1, inplace=True)

如果内存不足以一次性加载全量数据，可以按DISTINCT_IRECIPIENTID做哈希分块，逐块处理后再合并结果即可。

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