Python中提升Pandas性能的5个技巧

如果您曾经使用过表格数据，您可能知道这个过程：将数据导入 pandas，清理和转换它，并将其用作模型的输入。但是，当需要扩展并将代码投入生产时，您的 pandas 管道很可能开始崩溃并运行缓慢。

提示 1：矢量化矩阵化
pandas 中的矢量化操作允许您仅使用几行代码即可操作整个数据框列，而不是手动单独循环每一行。

首先，假设您想要计算给定整数在列中出现的次数。以下是 2 种可能的方法。

def count_loop(X, target: int) -> int:
    return sum(x == target for x in X["column_1"])

def count_vectorized(X, target: int) -> int:
    return (X["column_1"] == target).sum()

在测试 df1 和 df2 时，使用 count_vectorized 方法比使用 count_loop 方法的速度提高了 82 倍。

现在，假设你有一个带有日期列的 DataFrame，并希望将其偏移给定的天数。下面有两种方法。你能猜出矢量化操作的加速系数吗？

def offset_loop(X, days: int) -> pd.DataFrame:
        d = pd.Timedelta(days=days)
    X["column_const"] = [x + d for x in X["column_10"]]
    return X

def offset_vectorized(X, days: int) -> pd.DataFrame:
    X["column_const"] = X["column_10"] + pd.Timedelta(days=days)
    return X

通过使用矢量化运算而不是循环来执行这一代价高昂的操作，我们的平均速度提高了 460 倍！

关于 offset_loop 函数的一个小说明：我们在列表理解之外构建了 offset_loop。否则，它将在每次迭代中构建，这将使其运行时间加倍（这本身就值得注意）。

虽然我们强烈推荐在 pandas 中进行矢量化操作，以获得最佳性能和效率，但需要注意的是，并不是所有操作都可以矢量化，在某些情况下可能需要使用迭代等替代方法。

技巧 2：像半专业人士一样迭代
假设我们将以下函数应用于数据集：

remove_col = "column_2"
words_to_remove_col = "column_4"

def remove_words(
    remove_from: str, words_to_remove: str, min_include_word_length: int = 4
) -> str:
    words_to_exclude = set(words_to_remove.split(" "))
    no_html = re.sub("<.*?>", " ", remove_from)
    include_words = [
        x
        for x in re.findall(r"\w+", no_html)
        if (len(x) >= min_include_word_length) and (x not in words_to_exclude)
    ]
    return " ".join(include_words)

它从给定的 remove_from 短语中删除 words_too_remove 短语中出现的所有单词，并消除 HTML 标记，同时保留长度大于或等于 min_include_word_length 的单词。

可以将 remove_from 视为一个长描述，您希望从中删除一些已经出现在另一数据列中的单词。

例如，您有以下数据帧：

df = pd.DataFrame(
    {
        "column_1": [31, 41],
        "column_2": [
            "<p>The Apple iPhone 14, launched in 2022, comes in black, has metallic bezels and 2 or 3 cameras on the back.</p>",
            "<p>The Samsung Galaxy S22 Ultra, launched in 2022, is slim, comes in purple, has metallic bezels and multiple cameras on the back.</p>",
        ],
        "column_3": [59, 26],
        "column_4": ["Apple iPhone", "Samsung Galaxy"],
    }
)

如果对所有记录应用 remove_words，结果列表如下：

[
    'launched 2022 comes black metallic bezels cameras back',
    'Ultra launched 2022 slim comes purple metallic bezels multiple cameras back',
]

1、使用for 循环
第一种也是最直观的迭代方法是使用 Python for 循环。

def loop(df: pd.DataFrame, remove_col: str, words_to_remove_col: str) -> list[str]:
    res = []
    i_remove_col = df.columns.get_loc(remove_col)
    i_words_to_remove_col = df.columns.get_loc(words_to_remove_col)
    for i_row in range(df.shape[0]):
        res.append(
            remove_words(
                df.iat[i_row, i_remove_col], df.iat[i_row, i_words_to_remove_col]
            )
        )
    return result

到处都有开销：访问 df 的单个值和使用 append 动态创建列表意味着这将是最慢的方法，我们将使用它作为比较基准。但我们希望不惜一切代价避免这种情况。

2、使用apply
下一步很简单，使用 pandas 的 apply 方法，它可以在内部处理循环。你本以为这个方法会非常快，但它只快 1.9 倍。代码如下。不过，因为有更好的方法，所以你也应该避免使用这种方法。

def apply(df: pd.DataFrame, remove_col: str, words_to_remove_col: str) -> list[str]:
    return df.apply(
        func=lambda x: remove_words(x[remove_col], x[words_to_remove_col]), axis=1
    ).tolist()

在 df.apply 的每次迭代中，所提供的可调用程序都会获得一个索引为 df.columns、值为行的序列。这意味着 pandas 必须在每个循环中生成该序列，成本很高。为了降低成本，最好在你知道会用到的 df 子集上调用 apply，就像这样：

def apply_only_used_cols(df: pd.DataFrame, remove_col: str, words_to_remove_col: str) -> list[str]:
    return df[[remove_col, words_to_remove_col]].apply(
        func=lambda x: remove_words(x[remove_col], x[words_to_remove_col]), axis=1
    )

这一简单的调整使操作速度比基准快 2.1 倍，但我们仍建议避免使用这种方法。

3、使用list-comp + itertuples
将 itertuples 与心爱的列表理解结合起来使用肯定会更好。请注意，我们再次子集了 df。这比基线平均快 4.6 倍。

def itertuples_only_used_cols(df: pd.DataFrame, remove_col: str, words_to_remove_col: str) -> list[str]:
    return [
        remove_words(x[0], x[1])
        for x in df[[remove_col, words_to_remove_col]].itertuples(
            index=False, name=None
        )
    ]

4、使用list-comp + zip
另一种略有不同的方法是再次使用列表理解，但使用 zip 对所需列进行迭代，这种方法也能获得大致相同的性能（比基线快 4.6 倍）。Zip 接收迭代项并生成元组，其中第 i 个元组按顺序包含所有给定迭代项中的第 i 个元素。这种方法看起来像这样

def zip_only_used_cols(df: pd.DataFrame, remove_col: str, words_to_remove_col: str) -> list[str]:
    return [remove_words(x, y) for x, y in zip(df[remove_col], df[words_to_remove_col])]

5、使用list-comp + to_dict
这种方法的一个较慢的变体是使用 to_dict(orient="decords")遍历数据集行的 dicts。这种方法比基准方法快约 3.9 倍。具体代码如下

def to_dict_only_used_columns(df: pd.DataFrame) -> list[str]:
        return [
            remove_words(row[remove_col], row[words_to_remove_col])
            for row in df[[remove_col, words_to_remove_col]].to_dict(orient="records")
        ]

这不是一条硬性规定，但通过观察，我们发现使用 to_dict 会使程序的内存占用几乎增加两倍，因此在使用时要慎重。

额外方法 1：缓存
除了我们讨论过的迭代技术，还有两种方法可以帮助提高代码的性能：缓存和并行化。如果你的 pandas 函数被多次调用相同的参数，缓存就会特别有用。例如，如果 remove_words 应用于有许多重复值的数据集，可以使用 functools.lru_cache 来存储函数结果，避免每次都重新计算。要使用 lru_cache，只需在 remove_words 的声明中添加 @lru_cache 装饰器，然后使用您喜欢的迭代方法将函数应用到数据集。这可以大大提高代码的速度和效率。以下面的代码为例：

@lru_cache
def remove_words(...):
    ... # Same implementation as before

def zip_only_used_cols_cached(df: pd.DataFrame, remove_col: str, words_to_remove_col: str) -> list[str]:
    return [remove_words(x, y) for x, y in zip(df[remove_col], df[words_to_remove_col])]

添加此装饰器后，函数会 "记住 "之前遇到的输入的输出，从而无需重新运行所有代码。您可以在 Python 的官方文档中阅读更多关于它的信息。别忘了使用 maxsize 参数来权衡速度与内存。

您可能会注意到，缓存函数 df1 的性能与非缓存函数类似。这说明缓存并不是万能的。在我们的例子中，df1 的结构没有那么多重复，因此缓存对它的帮助不如对 df2 的帮助大，后者的速度提高了 46 倍。

额外方法 2：并行化
我们的最后一张王牌是使用 pandarallel 在多个独立的 df 中并行调用函数。该工具使用简单：只需导入并初始化，然后将所有 .applys 改为 .parallel_applys。

from pandarallel import pandarallel
pandarallel.initialize(nb_workers=min(os.cpu_count(), 12))

def parapply_only_used_cols(df: pd.DataFrame, remove_col: str, words_to_remove_col: str) -> list[str]:
    return df[[remove_col, words_to_remove_col]].parallel_apply(
        lambda x: remove_words(x[remove_col], x[words_to_remove_col]), axis=1
    )

在我们的案例中，我们测得的改进幅度为 12 倍。

尽管功能强大，parallel_apply 也不是什么灵丹妙药。它将 df 的每块数据都拾取到 /dev/shm（实际上就是 RAM）中，因此内存可能会不够用！

如果出现这种情况，可以通过减少进程数来限制 pandarallel 的内存使用量，但这也会影响性能。

为什么不把两者结合起来呢？我们试过，但没有发现实际的好处。它要么和裸 parallel_apply 一样快，要么比带缓存的裸 list comprehension 慢。这可能与并行化的实现方式以及缓存不共享有关。不过，我们需要进一步调查才能得出确切结论。

技巧 3：拥抱 NumPy
Pandas 是在 NumPy 的基础上构建的，NumPy 以其处理大型数组和矩阵数值数据的性能和速度而著称。这有助于使 Pandas 在处理大型数据集时高效快速。既然你已经是矢量化操作方面的专家，下一步自然是学习 numpy。本节内容短小精悍：如果你需要额外的性能，可以考虑直接使用 .to_numpy() 和 numpy 方法来实现 numpy。

对这种方法进行基准测试几乎是不可能的。不同的使用情况和代码的实现细节会导致结果大相径庭。在我们的测试中，我们观察到的速度提升从 2 倍到惊人的 3000 倍不等。

继续上面的例子，如果我们尝试直接在 numpy 中计数或偏移会怎样呢？请看下面的代码：

def count_numpy(X, target: int) -> int:
    return np.count_nonzero(X["column_1"].to_numpy() == target)

def offset_numpy(X, days: int) -> pd.DataFrame:
    X["column_const"] = X["column_10"].to_numpy() + np.timedelta64(days, "D")
    return X

正如你所看到的，这很容易实现：添加 .to_numpy() 并使用 numpy 对象。在计数方面，速度总共提高了 185 倍（是矢量化 pandas 方法的 2.2 倍）。在偏移量方面，与循环操作相比，速度提高了 1200 倍。

我们听到了！这些数字有些人为，因为正常人都不会这样偏移日列。但是，熟悉其他编程语言的初学者可能会遇到这种情况，这些示例值得警惕。

这些速度的提升得益于矢量化和直接使用 numpy 的复合优势。这也消除了 pandas 的杀手级功能所带来的开销。

需要注意的是，有时这并不像在操作中的每个序列后放置 .to_numpy() 那么简单。例如：一位团队成员报告说，在带有时区的日期时间数据中使用 .to_numpy() 实际上会丢失时区。如果所有数据都在同一时区，并且您正在计算差异，那么这种情况应该没有问题。不过，这还是需要注意的。

另一位团队成员在处理日期时也遇到了一些奇怪的问题。在对一系列日期时间调用 .to_numpy() 时，他们会得到一个 datetime64[ns] 数组。但是，在对单个 pd.Timestamp 调用 .to_numpy() 时，返回的却是一个 int。因此，除非使用 .to_numpy().astype(int) 从日期数组中减去单个日期，否则将返回乱码。

考虑到这些加速因素，尽管有一些注意事项，您还是希望在您的武器库中找到这个工具，而且应该毫不犹豫地使用它。

技巧 4：了解你的 dtypes
在 pandas DataFrames 中，dtype 是指定每一列数据类型的关键属性。因此，为 DataFrame 中的每一列选择合适的 dtype 是关键。

一方面，我们可以将数字降序为使用较少位的类型，以节省内存。相反，我们可以为特定数据使用专门的类型，从而降低内存成本，并以数量级的方式优化计算。

我们将讨论 pandas 中最常见的一些类型，如 int、float、bool 和字符串。但首先，我们要介绍一下可怕的对象。

1、对象类型
关于数据类型，pandas 有许多类型可以在快速 C 层有效地映射到 numpy 类型。然而，当没有简单的映射时，它就会求助于 python 对象。对于新手来说，对象是语言中所有对象的父类[源代码]。由于内存是如何处理的，这导致了效率低下。

numpy 的内核是用 C 语言编写的，这意味着数组可以直接指向位于连续内存块中的所有值。由于高速缓冲存储器可以利用数据的空间位置性，这使得执行时间大大缩短。您可以在此阅读更多相关信息。

由于 pandas 将字符串存储为对象数组，因此它不得不依赖于速度较慢的 Python 数组。与 numpy 不同，Python list 有一个指向内存连续缓冲区指针的指针，这些指针指向存储在内存中的对象，但也引用存储在其他位置的数据。因此，相比之下，在运行时访问这些数据将相当缓慢。

一言以蔽之，pandas 在一定程度上受到生态系统的制约。但是，我们如何才能充分利用我们现在所拥有的工具呢？对于数值数据和字符串数据来说，情况就不同了。

2、数字类型
关于整数和浮点数，下投是节省内存的关键。Pandas 支持 8、16、32 和 64 位有符号和无符号整数，以及 16、32 和 64 位浮点数。默认情况下，Pandas 选择使用这两种类型的 64 位变体。

诀窍在于选择能轻松容纳数据的最小类型。您应该同时考虑当前和新的有效数据。如果有符号的 int 确实存在，就不要在管道中将自己限制在无符号的 int 上。

在实践中，对于当前的数据集，你可以使用 pandas 自带的 pd.to_numeric 像这样轻松地实现这一点：

n = 100_000

df = pd.DataFrame({
        "uint8": np.random.randint(10, 20, n),
        "uint32": np.random.randint(100_000, 200_000, n),
        "int16": np.random.randint(1_000, 2_000, n) * np.random.choice((-1, 1), n),
        "float32": np.random.uniform(100_000, 200_000, n),
    }
)

df_downcasted = (
    df
    .apply(pd.to_numeric, downcast="float")
    .apply(pd.to_numeric, downcast="integer")
    .apply(pd.to_numeric, downcast="unsigned")
)

就这样，我们将 df 的大小缩小了 1/3。不过，您的数据可能会有所不同。顺便提一句：pd.to_numeric 只将浮点数下投到 float32，而不使用 float16，尽管它与 pandas 兼容。

关于数据类型，我们发现如果没有空值，应该优先使用 bool。它占用的内存与 uint8 相同，但能使列的内容更加清晰。

"如果有 NaN 值怎么办？我们听到了。原则上，你可以将 ints 和 bools 转换为 float32，然后就可以了。虽然会占用内存，但 NaN 将得到原生处理。

最近，pandas 增加了自己的 Nullable 整数和 Nullable 布尔型数据类型。它们比 float32 更节省内存，但需要注意的是，.values 和 .to_numpy() 都不会使用 np.nan 返回正确的 numpy 数组（因为它本身是浮点数，而不是 int）。在内部，这些数组包含一个_data 数组，该数组具有相应的 numpy dtype 和一个_mask，其中 True 表示 _data 中的值缺失。如果要使用 numpy，这可能是个问题。因此，我们暂时还不能完全推荐使用这些数组。

这也是我们暂时不推荐使用 .convert_dtypes 的原因，因为它使用的是可空类型。

3、字符串类型
说完了数字，让我们深入了解一下字符串。字符串随处可见，但在处理数据时，它们却是我们的祸根。如前所述，字符串是用低效的对象类型存储的，这就是为什么字符串会占用大量空间，而且比更专业的类型速度更慢。

不过，StringDType 带来了一线希望。与我们提到的其他可归零类型一样，StringDType 仍处于试验阶段，因此我们暂时不建议在生产代码中使用它。展望未来，StringDType 将正确处理缺失值，并在性能上比对象类型有所改进。

如果列的卡入度较低（重复值较多），最明智的做法是使用 pd.Category。从内部看，分类列保存了从实际列值（即转换前列的内容）到某个任意整数的映射。然后，列被保存为 int 类型，从而利用了连续内存和高效存储的优势。此外，运行长度编码（Run Length Encoding）的使用还能进一步缩小存储空间（如果一个整数重复出现，只会保存整数和实例数，而不会保存所有整数）。

技巧 5：长期存储
CSV、parquet、pickle、JSON、XML、HDF5、Feather、ORC、STATA、SAS、SPSS 或 Excel（呸！）都是 pandas 支持的格式。啊啊啊选项太多了！我们该如何选择呢？

我们可以像其他负责任的工程师一样，先回顾一下每种格式的优缺点。然后，我们会测试一些在现实世界中比较突出的格式，最后得出结论，决定在什么时候选择哪种格式。

不用担心！我们已经做足了功课--我们保证我们的方法比荷马的更科学。我们建议您好好看看镶木地板。您不必再不断地做出这些决定了。

创建 Parquet 的目的是帮助开发人员从 "压缩、高效的列式数据表示 "中获益。它支持不同的编码方式，可以按列进行压缩。这与 ML 从业人员的用例不谋而合！

听起来真棒，我也要试试！对不对？来吧，试着做一个简单的 df.to_parquet。

df = pd.DataFrame(np.random.randint(10, 20), columns=[f"col{i}" for i in range(20)])
df.to_parquet("test.parquet")

已经返回？是否出现了 ImportError？从技术上讲，Parquet 是一种格式，与其说是一种实现，不如说是一种规范。因此，要使用它，你需要安装另一个实现该格式的库。目前，pandas 支持 pyarrow 和 fastparquet。

但你承诺过不再做任何决定！我们认为 pandas 默认使用 pyarrow 是正确的选择。我们发现它比 fastparquet 更可靠。例如，当试图保存并立即读取 df2 时，fastparquet 会抱怨溢出。

归根结底：使用 parquet 并根据需要改变压缩方法。在开发过程中，应选择未压缩或快速的 parquet，而在长期存档时则采用 gzip 压缩。

CSV 可能是五大数据分发格式之一，这部分归功于它的易用性。通过简单的 df.to_csv，您就可以创建一个文件，任何人（和他们的狗）都可以用自己选择的程序打开它。但是，CSV 在容量、大小和速度方面都有很大的限制。当你遇到这些限制时，是时候开始探索 parquet 了。

最新的 pickle 版本速度很快，但存在大量安全和长期存储问题。解压过程中执行任意代码 "这几个字应该会让你不寒而栗，而这正是 pickle 官方文档中大量红色警告中的原话。

假设 pandas 更改了 DataFrame 类中的某些内容。在解压旧 DataFrame 时可能会遇到问题。为了避免这种情况，请选择（未压缩的）parquets 而不是 pickles。

我们最近发现了一款名为 Lance 的工具，它旨在简化机器学习工作流程。它支持一种名为 Lance Format 的列式数据格式，声称比 Parquet、Iceberg 和 Delta 等其他流行选项快 50 到 100 倍。虽然我们还没有对它进行彻底测试，但其潜在的速度提升还是很吸引人的，我们鼓励你试一试。

结论
以下是需要牢记的一些关键概念：

Python 的动态特性使其比编译语言更慢。这个问题在科学计算中更加严重，因为我们要运行数百万次简单的操作。
不要在数据上循环。取而代之的是将操作矢量化（记住 Hector Vector）。
你可以在 Dataframe 上调用 .to_numpy() 直接使用 numpy，这样速度会更快。
选择尽可能小的数字 dtypes，使用分类，优先使用 float32 而不是新的可空 dtypes（暂时如此）。
使用parquet存储数据。在开发过程中使用快速压缩（或根本不压缩），在长期存档时使用 gzip 压缩。