python - 加速 pandas 上的复杂功能

Question

我使用以下代码在数据帧的一列中填充 NaN 值:

for i in tqdm(range(nadf.shape[0])):
    a = nadf["primary"][i]
    nadf["count"][i] = np.ceil(d[a]*a)

此代码通过将字典 d 中“primary”的相应值与“primary”的值相乘来替换“count”中的 NaN 值。 nadf 有 1600 万行。我知道执行速度会很慢，但是有没有办法加快速度？

Answer 1

如果我以正确的方式理解您的问题和数据框值，则可以使用 pandas 内部功能通过以下方式解决问题:

请关注代码中的注释，随时提问。

import pandas as pd
import numpy as np
import math


def fill_nan(row, _d):
    """fill nan values in "count" column based on "primary" column value and dictionary _d"""

    if math.isnan(row["count"]):
        return np.ceil(_d[row["primary"]]) * row["primary"]
    return row["count"]  # else not nan


if __name__ == "__main__":
    d = {1: 10, 2: 20, 3: 30}

    df = pd.DataFrame({
        "primary": [1, 2, 3, 1, 2, 1, 2],
        "count": [10.1, 4, 5, np.nan, np.nan, 4, np.nan]
    })

    df["count"] = df.apply(lambda row: fill_nan(row, d), axis=1)  # changes nan here

    print(df)