在 groupby 上应用返回多个值的函数的最快方法

Question

我有以下数据框：

           pair    price          ts
0       BTCUSDT  4724.89   985446400
1       ETHUSDT   384.79   985446400
2       BTCUSDT  4724.89   985447400
3       ETHUSDT   384.79   985447400
4       BTCUSDT  4724.89   985448400
...         ...      ...         ...
172795  BTCUSDT  4834.91  1071843400
172796  ETHUSDT   388.46  1071844400
172797  BTCUSDT  4834.91  1071844400
172798  ETHUSDT   388.46  1071845400
172799  BTCUSDT  4834.91  1071845400

我想使用 pandas_ta 库对数据框中每一对的值应用一组预定义的指标。

到目前为止，我已经设法获得了这样最快的方法：

seconds[['EMA', 'SMA', 'RSI']] = seconds.groupby("pair", group_keys=False, sort=False)["price"].apply(
    lambda x: pd.DataFrame(
        {
            "EMA": ta.ema(x, length=14),
            "SMA": ta.sma(x, length=14),
            "RSI": ta.rsi(x, length=14),
        }
    )
)

但是，我觉得为每一对创建一个新的数据框，尤其是当有两个以上的数据框时，并不是真正的最佳选择。

我也试过像这样返回一个元组：

seconds.groupby("pair", group_keys=False, sort=False)["price"].apply(
    lambda x: (
        ta.ema(x, length=14),
        ta.sma(x, length=14),
        ta.rsi(x, length=14),
    )
)

但我无法弄清楚如何以有效的方式将该元组的结果附加到我的 df 中。我可以注意到元组的每个元素都包含所有需要附加的列，但不知道如何从这里开始。这似乎比为每对创建一个单独的数据框快 30% 左右。

(0             NaN
 2             NaN
 4             NaN
 6             NaN
 8             NaN
            ...   
 172791    4834.91
 172793    4834.91
 172795    4834.91
 172797    4834.91
 172799    4834.91
 Name: EMA_14, Length: 86400, dtype: float64,
 0             NaN
 2             NaN
 4             NaN
 6             NaN
 8             NaN
            ...   
 172791    4834.91
 172793    4834.91
 172795    4834.91
 172797    4834.91
 172799    4834.91
 Name: SMA_14, Length: 86400, dtype: float64,
 0               NaN
 2               NaN
 4               NaN
 6               NaN
 8               NaN
             ...    
 172791    99.993504
 172793    99.993504
 172795    99.993504
 172797    99.993504
 172799    99.993504
 Name: RSI_14, Length: 86400, dtype: float64)

预先感谢您的建议！

Answer 1

您可以使用 pd.concat 组合元组的结果。只需连接元组的所有第一个元素，然后连接所有第二个元素，依此类推。然后将它们添加到数据框，每个作为一个新列：

your_tuple = seconds.groupby("pair", group_keys=False, sort=False)["price"].apply(
    lambda x: (
        ta.ema(x, length=14),
        ta.sma(x, length=14),
        ta.rsi(x, length=14),
    )
)

columns = ["EMA", "SMA", "RSI"]
for idx, col in enumerate(columns):
    # Grab the "pieces" of each calculation from your_tuple, and concatenate them
    seconds[col] = pd.concat([group_results[idx] for group_results in your_tuple])

your_tuple 中的系列与数据帧具有相同的索引，因此当您执行 seconds[col] = pd.concat... 时，pandas 会对它们进行解读以匹配数据帧索引

话虽这么说，根据我对一些虚拟数据的测试，它的运行速度与创建数据帧的初始方法一样快（我认为这一点也不坏）。

如果您想让事情变得更快，您也可以尝试将 apply 替换为另一种操作 groupby 对象的方法。来自大熊猫docs：

虽然 apply 是一种非常灵活的方法，但它的缺点是使用它可能比使用更具体的方法（如 agg 或 transform）慢很多。 Pandas 提供了广泛的方法，比使用 apply 用于特定目的要快得多，因此在使用 apply 之前尝试使用它们。

您可以做的另一件事是并行化您的操作。根据您的任务，并行化甚至可能不值得，因此您必须针对您的具体情况进行测试。下面是一个关于如何使用 multiprocessing 并行化 groupby 计算的示例：

from multiprocessing import Pool, cpu_count
import pandas_ta as ta

metrics = [ta.ema, ta.sma, ta.rsi]
def calculate_metrics_for_group(group):
    return [metric(group, length=14) for metric in metrics]

grouped = seconds.groupby("pair", group_keys=False, sort=False)["price"]
with Pool(cpu_count()) as pool:
    your_tuple = pool.map(calculate_metrics_for_group, [group for _, group in grouped])

columns = ["EMA", "SMA", "RSI"]
for idx, col in enumerate(columns):
    seconds[col] = pd.concat([group_results[idx] for group_results in your_tuple])