python - Pandas 中的 for 循环真的很糟糕吗？我什么时候应该关心？

Question

是 for循环真的“坏”？如果不是，在什么情况下它们会比使用更传统的“矢量化”方法更好？1

我熟悉“矢量化”的概念，以及 Pandas 如何使用矢量化技术来加速计算。向量化函数在整个系列或 DataFrame 上广播操作，以实现比传统迭代数据大得多的加速。

然而，我很惊讶地看到很多代码(包括来自 Stack Overflow 上的答案)为涉及使用 for 遍历数据的问题提供解决方案。循环和列表推导式。文档和 API 说循环是“坏的”，并且应该“永远不要”遍历数组、系列或数据帧。那么，为什么我有时会看到用户建议基于循环的解决方案？

1 - 虽然这个问题听起来确实有些宽泛，但事实是当 for 时存在非常特殊的情况。循环通常比传统的迭代数据更好。这篇文章旨在为后代捕捉这一点。

Answer 1

TLDR；不，for 循环不是一揽子“坏”，至少，并非总是如此。 可能更准确地说某些矢量化操作比迭代 慢，而不是说迭代比某些矢量化操作快。知道何时以及为什么是从代码中获得最大性能的关键。简而言之，这些是值得考虑替代矢量化 Pandas 函数的情况:

当你的数据很小时(...取决于你在做什么)，

当处理 object/mixed dtypes

使用 str/regex 访问器函数时

让我们分别检查这些情况。

小数据上的迭代 v/s 向量化

Pandas 在其 API 设计中遵循 "Convention Over Configuration" 方法。这意味着已经安装了相同的 API 来满足广泛的数据和用例。

调用 pandas 函数时，该函数必须在内部处理以下事项(以及其他事项)，以确保正常工作

索引/轴对齐

处理混合数据类型

处理丢失的数据

几乎每个函数都必须在不同程度上处理这些问题，这会带来 开销 。数字函数(例如 Series.add )的开销较小，而字符串函数(例如 Series.str.replace )的开销更明显。

另一方面， for 循环比您想象的要快。更好的是 list comprehensions(通过 for 循环创建列表)速度更快，因为它们是列表创建的优化迭代机制。

列表推导遵循模式

[f(x) for x in seq]

其中 seq 是 pandas 系列或 DataFrame 列。或者，当操作多列时，

[f(x, y) for x, y in zip(seq1, seq2)]

其中 seq1 和 seq2 是列。

数值比较
考虑一个简单的 bool 索引操作。列表理解方法已针对 Series.ne ( != ) 和 query 计时。以下是功能:

# Boolean indexing with Numeric value comparison.
df[df.A != df.B]                            # vectorized !=
df.query('A != B')                          # query (numexpr)
df[[x != y for x, y in zip(df.A, df.B)]]    # list comp

为简单起见，我使用 perfplot 包来运行本文中的所有 timeit 测试。上述操作的时间如下:



对于中等大小的 N，列表理解优于 query，甚至在微小 N 的情况下也优于向量化不等于比较。不幸的是，列表理解是线性扩展的，因此它不会为较大的 N 提供太多的性能提升。

Note
It is worth mentioning that much of the benefit of list comprehension come from not having to worry about the index alignment, but this means that if your code is dependent on indexing alignment, this will break. In some cases, vectorised operations over the underlying NumPy arrays can be considered as bringing in the "best of both worlds", allowing for vectorisation without all the unneeded overhead of the pandas functions. This means that you can rewrite the operation above as

df[df.A.values != df.B.values]

Which outperforms both the pandas and list comprehension equivalents:

NumPy vectorization is out of the scope of this post, but it is definitely worth considering, if performance matters.

值计数
再举一个例子——这一次，使用另一个比 for 循环更快的普通 python 结构—— collections.Counter 。一个常见的要求是计算值计数并将结果作为字典返回。这是通过 value_counts 、 np.unique 和 Counter 完成的:

# Value Counts comparison.
ser.value_counts(sort=False).to_dict()           # value_counts
dict(zip(*np.unique(ser, return_counts=True)))   # np.unique
Counter(ser)                                     # Counter

结果更加明显，Counter 在更大范围的小 N (~3500) 中胜过这两种矢量化方法。

Note
More trivia (courtesy @user2357112). The Counter is implemented with a C accelerator, so while it still has to work with python objects instead of the underlying C datatypes, it is still faster than a for loop. Python power!

当然，这里的结论是性能取决于您的数据和用例。这些示例的目的是说服您不要将这些解决方案排除为合法选项。如果这些仍然不能为您提供所需的性能，那么总是有 cython 和 numba 。让我们将此测试添加到组合中。

from numba import njit, prange

@njit(parallel=True)
def get_mask(x, y):
    result = [False] * len(x)
    for i in prange(len(x)):
        result[i] = x[i] != y[i]

    return np.array(result)

df[get_mask(df.A.values, df.B.values)] # numba

Numba 提供了将循环 Python 代码 JIT 编译为非常强大的矢量化代码。了解如何使 numba 工作涉及学习曲线。

混合/object dtypes 的操作

基于字符串的比较
回顾第一节的过滤示例，如果被比较的列是字符串怎么办？考虑上面相同的 3 个函数，但输入 DataFrame 转换为字符串。

# Boolean indexing with string value comparison.
df[df.A != df.B]                            # vectorized !=
df.query('A != B')                          # query (numexpr)
df[[x != y for x, y in zip(df.A, df.B)]]    # list comp

那么，发生了什么变化？这里要注意的是 字符串操作本质上很难矢量化。 Pandas 将字符串视为对象，对对象的所有操作都退回到缓慢、循环的实现。

现在，因为这个循环实现被上面提到的所有开销所包围，所以这些解决方案之间存在恒定的幅度差异，即使它们的规模相同。

当涉及对可变/复杂对象的操作时，没有比较。列表理解优于所有涉及字典和列表的操作。

通过键访问字典值
以下是从字典列中提取值的两个操作的时间:map 和列表推导式。设置在附录中的“代码片段”标题下。

# Dictionary value extraction.
ser.map(operator.itemgetter('value'))     # map
pd.Series([x.get('value') for x in ser])  # list comprehension

位置列表索引
从列列表(处理异常)、 map 、 str.get accessor method 和列表推导中提取第 0 个元素的 3 个操作的时间:

# List positional indexing. 
def get_0th(lst):
    try:
        return lst[0]
    # Handle empty lists and NaNs gracefully.
    except (IndexError, TypeError):
        return np.nan

ser.map(get_0th)                                          # map
ser.str[0]                                                # str accessor
pd.Series([x[0] if len(x) > 0 else np.nan for x in ser])  # list comp
pd.Series([get_0th(x) for x in ser])                      # list comp safe

Note
If the index matters, you would want to do:

pd.Series([...], index=ser.index)

When reconstructing the series.

列表扁平化
最后一个例子是扁平化列表。这是另一个常见问题，它展示了纯 Python 的强大之处。

# Nested list flattening.
pd.DataFrame(ser.tolist()).stack().reset_index(drop=True)  # stack
pd.Series(list(chain.from_iterable(ser.tolist())))         # itertools.chain
pd.Series([y for x in ser for y in x])                     # nested list comp

itertools.chain.from_iterable 和嵌套列表推导式都是纯 Python 构造，并且比 stack 解决方案可扩展性好得多。

这些时间强烈表明 Pandas 不具备使用混合 dtype 的能力，您可能应该避免使用它来这样做。在可能的情况下，数据应在单独的列中以标量值(整数/浮点数/字符串)的形式存在。

最后，这些解决方案的适用性很大程度上取决于您的数据。因此，最好的办法是在决定使用什么之前对您的数据测试这些操作。请注意我没有在这些解决方案上对 apply 计时，因为它会扭曲图形(是的，它很慢)。

正则表达式操作和 .str 访问器方法

Pandas 可以对字符串列应用正则表达式操作，例如 str.contains 、 str.extract 和 str.extractall ，以及其他“向量化”字符串操作(例如 str.split 、 str.find , str.translate` 等)。这些函数比列表推导慢，并且比其他任何函数都更方便。

预编译正则表达式模式并使用 re.compile 迭代数据通常要快得多(另请参阅 Is it worth using Python's re.compile? )。等价于 str.contains 的 list comp 看起来像这样:

p = re.compile(...)
ser2 = pd.Series([x for x in ser if p.search(x)])

或者，

ser2 = ser[[bool(p.search(x)) for x in ser]]

如果你需要处理 NaN，你可以做类似的事情

ser[[bool(p.search(x)) if pd.notnull(x) else False for x in ser]]

等价于 str.extract(无组)的 list comp 将类似于:

df['col2'] = [p.search(x).group(0) for x in df['col']]

如果您需要处理不匹配和 NaN，您可以使用自定义函数(速度更快!):

def matcher(x):
    m = p.search(str(x))
    if m:
        return m.group(0)
    return np.nan

df['col2'] = [matcher(x) for x in df['col']]

matcher 函数具有很强的可扩展性。可以根据需要为每个捕获组返回一个列表。只需提取查询匹配器对象的 group 或 groups 属性即可。

对于 str.extractall ，将 p.search 更改为 p.findall 。

字符串提取
考虑一个简单的过滤操作。这个想法是如果前面有一个大写字母，则提取 4 位数字。

# Extracting strings.
p = re.compile(r'(?<=[A-Z])(\d{4})')
def matcher(x):
    m = p.search(x)
    if m:
        return m.group(0)
    return np.nan

ser.str.extract(r'(?<=[A-Z])(\d{4})', expand=False)   #  str.extract
pd.Series([matcher(x) for x in ser])                  #  list comprehension

更多例子
完全披露 - 我是下面列出的这些帖子的作者(部分或全部)。

Fast punctuation removal with pandas

String concatenation of two pandas columns

Remove unwanted parts from strings in a column

Replace all but the last occurrence of a character in a dataframe

结论

如上面的示例所示，迭代在处理小行 DataFrame、混合数据类型和正则表达式时会发光。

您获得的加速取决于您的数据和您的问题，因此您的里程可能会有所不同。最好的办法是仔细运行测试，看看付出的努力是否值得。

“矢量化”函数因其简单性和可读性而大放异彩，因此如果性能不重要，您绝对应该更喜欢那些。

另一个注意事项，某些字符串操作处理有利于使用 NumPy 的约束。下面是两个例子，其中仔细的 NumPy 矢量化优于 python:

Create new column with incremental values in a faster and efficient way - Answer by Divakar

Fast punctuation removal with pandas - Answer by Paul Panzer

此外，有时仅通过 .values 操作底层数组而不是 Series 或 DataFrames 可以为大多数常见场景提供足够健康的加速(请参阅上面 数字比较 _0x104569 部分中的 Note )。因此，例如 df[df.A.values != df.B.values] 会比 df[df.A != df.B] 显示即时性能提升。使用 .values 可能并不适用于所有情况，但它是一个有用的技巧。

如上所述，由您决定这些解决方案是否值得实现。

附录:代码片段

import perfplot  
import operator 
import pandas as pd
import numpy as np
import re

from collections import Counter
from itertools import chain

# Boolean indexing with Numeric value comparison.
perfplot.show(
    setup=lambda n: pd.DataFrame(np.random.choice(1000, (n, 2)), columns=['A','B']),
    kernels=[
        lambda df: df[df.A != df.B],
        lambda df: df.query('A != B'),
        lambda df: df[[x != y for x, y in zip(df.A, df.B)]],
        lambda df: df[get_mask(df.A.values, df.B.values)]
    ],
    labels=['vectorized !=', 'query (numexpr)', 'list comp', 'numba'],
    n_range=[2**k for k in range(0, 15)],
    xlabel='N'
)

# Value Counts comparison.
perfplot.show(
    setup=lambda n: pd.Series(np.random.choice(1000, n)),
    kernels=[
        lambda ser: ser.value_counts(sort=False).to_dict(),
        lambda ser: dict(zip(*np.unique(ser, return_counts=True))),
        lambda ser: Counter(ser),
    ],
    labels=['value_counts', 'np.unique', 'Counter'],
    n_range=[2**k for k in range(0, 15)],
    xlabel='N',
    equality_check=lambda x, y: dict(x) == dict(y)
)

# Boolean indexing with string value comparison.
perfplot.show(
    setup=lambda n: pd.DataFrame(np.random.choice(1000, (n, 2)), columns=['A','B'], dtype=str),
    kernels=[
        lambda df: df[df.A != df.B],
        lambda df: df.query('A != B'),
        lambda df: df[[x != y for x, y in zip(df.A, df.B)]],
    ],
    labels=['vectorized !=', 'query (numexpr)', 'list comp'],
    n_range=[2**k for k in range(0, 15)],
    xlabel='N',
    equality_check=None
)

# Dictionary value extraction.
ser1 = pd.Series([{'key': 'abc', 'value': 123}, {'key': 'xyz', 'value': 456}])
perfplot.show(
    setup=lambda n: pd.concat([ser1] * n, ignore_index=True),
    kernels=[
        lambda ser: ser.map(operator.itemgetter('value')),
        lambda ser: pd.Series([x.get('value') for x in ser]),
    ],
    labels=['map', 'list comprehension'],
    n_range=[2**k for k in range(0, 15)],
    xlabel='N',
    equality_check=None
)

# List positional indexing. 
ser2 = pd.Series([['a', 'b', 'c'], [1, 2], []])        
perfplot.show(
    setup=lambda n: pd.concat([ser2] * n, ignore_index=True),
    kernels=[
        lambda ser: ser.map(get_0th),
        lambda ser: ser.str[0],
        lambda ser: pd.Series([x[0] if len(x) > 0 else np.nan for x in ser]),
        lambda ser: pd.Series([get_0th(x) for x in ser]),
    ],
    labels=['map', 'str accessor', 'list comprehension', 'list comp safe'],
    n_range=[2**k for k in range(0, 15)],
    xlabel='N',
    equality_check=None
)

# Nested list flattening.
ser3 = pd.Series([['a', 'b', 'c'], ['d', 'e'], ['f', 'g']])
perfplot.show(
    setup=lambda n: pd.concat([ser2] * n, ignore_index=True),
    kernels=[
        lambda ser: pd.DataFrame(ser.tolist()).stack().reset_index(drop=True),
        lambda ser: pd.Series(list(chain.from_iterable(ser.tolist()))),
        lambda ser: pd.Series([y for x in ser for y in x]),
    ],
    labels=['stack', 'itertools.chain', 'nested list comp'],
    n_range=[2**k for k in range(0, 15)],
    xlabel='N',    
    equality_check=None

)

# Extracting strings.
ser4 = pd.Series(['foo xyz', 'test A1234', 'D3345 xtz'])
perfplot.show(
    setup=lambda n: pd.concat([ser4] * n, ignore_index=True),
    kernels=[
        lambda ser: ser.str.extract(r'(?<=[A-Z])(\d{4})', expand=False),
        lambda ser: pd.Series([matcher(x) for x in ser])
    ],
    labels=['str.extract', 'list comprehension'],
    n_range=[2**k for k in range(0, 15)],
    xlabel='N',
    equality_check=None
)