首页 > 解决方案 > 根据其他列中的值查询熊猫数据框的值的更快方法

问题描述

我有这个来自 Kaggle 的数据集,经过一些过滤,它包含大约 460 万行数据 在此处输入图像描述

我想查询这个数据框以获取close_adjustedand 等于某个值的symboldate。这是我到目前为止的实现

def price_diff(date_one, date_two, symbol):

    data = df.loc[df['symbol'] == symbol]
    date_one_data = data.loc[data['date'] == date_one]
    date_two_data = data.loc[data['date'] == date_two]

    # Only if data from both dates exist
    if (len(date_one_data) > 0 and len(date_two_data) > 0):
        close_price_one = date_one_data.squeeze()['close_adjusted']
        close_price_two = date_two_data.squeeze()['close_adjusted']
        return close_price_one - close_price_two

    return dummy

这段代码运行了大约 523 毫秒:523 ms ± 1.54 ms per loop

有没有更快的方法可以实现我想要的?现在,在 Kaggle 笔记本上将上述函数应用于包含 25000 行日期和符号的数据集需要一个多小时。

谢谢!

标签: pythonpandas

解决方案


在进行了一些分析之后,我发现你的大部分计算时间都花在了过滤掉symbol. 您可以通过访问为数据帧提供支持的底层 numpy 数组来提高速度,而不是使用数据帧进行过滤,因为它要快得多。可以使用该ravel()函数访问该数组。

def price_diff_ME(date_one, date_two, symbol):

    #data = df.loc[df['symbol'] == symbol]
    data = df[df["symbol"].ravel() == symbol]

    date_one_data = data.loc[data['date'] == date_one]
    date_two_data = data.loc[data['date'] == date_two]

    # Only if data from both dates exist
    if (len(date_one_data) > 0 and len(date_two_data) > 0):
        close_price_one = date_one_data.squeeze()['close_adjusted']
        close_price_two = date_two_data.squeeze()['close_adjusted']
        return close_price_one - close_price_two

    return False

在使用ravel()获取 numpy 数组后,我只需运行与以前相同的过滤器即可生成二进制掩码。然后,我将掩码应用于数据框以重现您之前使用的相同过滤器。

仅仅通过改变那一行,我就能看到速度提高了 70%。如果您愿意,也可以使用相同的方法过滤掉其他日期。


推荐阅读