如何让pandas中的保留计算更加高效？

Question

我正在尝试根据用户 ID 计算 7 天保留率（用户是否在 7 天内回来过？）。目前，我正在使用这段代码：

df_retention['seven_day_retention']=df_retention.groupby('user_id')['date'].transform(lambda x: ((x.shift(-1) - x).dt.days< 8).astype(int))

这个跨越 10M 行的过程需要几个小时，而且不可行。在 Databricks 中是否有更好的工作方式？

Answer 1

我测试了这个，它似乎比你的方法快得多。您的方法随着用户数量的增加而扩展。我想 groupby + lambda 在这里是一个特别糟糕的组合。

就像 @Confused Learner 所说，你需要使用内置的

pandas

方法，因为它们是用 C 编写的，并避免 lambda，这显然是用 Python 编写的。

import datetime
import random

import pandas as pd


# some synthetic data
k = int(1e3)
user_ids = random.choices(population=range(k), k=k)
months = random.choices(population=range(1, 12), k=k)
days = random.choices(population=range(1, 28), k=k)

# our synthetic dataframe
df_retention = pd.DataFrame(
    [   
        [user_id, datetime.datetime(2022, month, day)]
        for user_id, month, day in zip(user_ids, months, days)
    ],
    columns=["user_id", "date"]
)

df_retention.sort_values(by=["user_id", "date"], inplace=True)  # sort by user, then date

df_diff = df_retention[["user_id", "date"]].diff()  # take the difference of all the rows
retained = (df_diff["date"] <= datetime.timedelta(days=7)) & (df_diff["user_id"] == 0)  # True if diff is <= 7 days & it is the same user

retained.iloc[:-1] = retained.iloc[1:]  # shift the results
retained.iloc[-1] = False  # pad with False, since it's the last entry and we don't know if they ever returned
df_retention['seven_day_retention'] = retained

以下是强制使用

user_id=0

和

k=10

时的输出示例：

   user_id       date  seven_day_retention
4        0 2022-01-02                 True
2        0 2022-01-08                False
9        0 2022-02-14                False
0        0 2022-03-06                False
1        0 2022-04-21                False
6        0 2022-05-23                 True
3        0 2022-05-25                False
5        0 2022-07-21                False
7        0 2022-08-06                False
8        0 2022-10-12                False

Answer 2

你的代码非常慢。我认为你必须改变你的方法。您可以首先根据人员 ID 和日期对数据框进行排序。然后您可以使用 for 循环来比较每一行和下一行。这段代码的复杂度为 O(n)。如果您愿意，可以使用更快的方式。例如，在第二部分中，您可以使用示例代码，无需 groupby 和转换，只需计算每行和下一行之间的差异

Answer 3

df['时间戳'] = pd.to_datetime(df['时间戳']) df_sorted = df.sort_values(by=['user_id', 'timestamp'])

first_visit_dates = df_sorted.groupby('user_id')['timestamp'].nth(1)

visit_after_first_visit = df_sorted[df_sorted['timestamp'].isin(first_visit_dates)] Visit_after_first_visit['days_after_first_visit'] = (visit_after_first_visit['timestamp'] - Visit_after_first_visit.groupby('user_id')['timestamp'].transform('first')).dt.days

保留 = Visit_after_first_visit.groupby('days_after_first_visit').size() / df_sorted['user_id'].nunique()

打印（保留）