Pandas性能优化:矢量化操作

近日在进行数据分析项目的时候,涉及到对DataFrame的遍历操作。由于数据量较大,每次的遍历都会消耗较长的时间,遂开始进行性能优化方面的尝试,在偶然间了解到Pandas的矢量化操作(Vectorization)。本文将对其具体应用进行探究,并将其与其他的遍历方式进行性能对比。

何谓矢量化?

矢量化,有别于对每一个单独的值(标量)进行操作,是Pandas底层支持的对于整个Array进行的操作。这种操作适用于对某一列的全体数据进行普适的操作。举一个最简单的例子,比如在某次考试中,教授给了所有学生5分的curve,这种情况下就可以采用这种矢量化的思想。与其采取“行”的思想,遍历每一个学生的成绩并修改,我们采用列的思想,将这一列的全体数据直接应用加5分的操作,Pandas会自动将操作应用到每一个单元格。并且在实际操作中,Pandas的矢量化操作由底层的C语言实现,会带来显著的效率提升。

矢量化操作适用于Pandas的Dataframe,Series对象,同时也Numpy的Series对象。

矢量化操作:小试牛刀

在这部分,我们以上面我提到的考试加分的例子进行实际测试。首先,我们生成一个由学生ID和对应成绩两列构成的DataFrame,共1000条数据,如图:

如何进行矢量化操作呢?其实Pandas对于常用的函数,例如求和,平均值,方差等常用统计函数(所有支持矢量化的内置函数请查看Pandas官方文档)做了非常好的矢量化支持,在大多数简单的场景下,我们需要做的只是把一整列的元素,当成一个元素去处理,Pandas会自动把函数应用到每一个单元格上。例如:

grade_df["grade"] += 5

对比两次的输出,可以看到grade列已经改变。我们使用Jupyter Notebook对以上操作进行时间分析:

性能比较

接下来,我们使用传统的遍历方式进行操作,进行性能比较。我们分别采用.iterrows()方法和.apply()方法。

.iterrows()方法:

.apply()方法:

由此可见,矢量化操作的效率比最快的.apply()方法还快了37%,比iterrows()方法更是快了一个数量级。

矢量化操作:自定义矢量化函数

请注意:并非所有函数都支持矢量化,Pandas官方同样未给出明确的规范,只是简单的提到了一句话“accept NumPy arrays and return another array or value”(接收Numpy array对象并返回另一个array或值)。本文列出的情况之一仅供参考,具体应用场景有待进一步研究。

尽管Pandas官方内置了许多常用的矢量化函数,但是显然不足以覆盖稍复杂的需求。我们继续对上面的例子进行功能完善,当教授完成了加分后,需要将每个学生的百分制成绩转为字母等第。熟悉.apply()方法的读者会想到写一个条件判断函数,然后使用lambda调用。如图:

对于矢量化操作,我们是否可以使用同样的思路去调用呢?我们前面曾经讨论过,我们只需要将一列数据当成一个数据传递进去,Pandas会自动进行处理。由此尝试如下:

我们将整个grade列当做参数传入了letter_convert函数,但返回了错误:

ValueError: The truth value of a Series is ambiguous. Use a.empty, a.bool(), a.item(), a.any() or a.all().”

由返回的错误信息可知,Pandas并未执行矢量化操作,反而是将整个Series对象当做参数传入了letter_convert函数,if语句无法判断整个Series对象的布尔值。报错信息建议使用的a.any()、a.all()等函数只有在需要对所有元素生成一个统一输出时适用。例如:是否至少有一个人拿了A,是否所有人都通过考试等。显然,上述函数并不适用于本情况,因为我们期望对每一个元素生成对应的字母成绩。

Pandas的矢量化函数不支持显式的if条件语句。

但对于以上类型的,使用if判断单元格值,并且返回新值的函数,我们可以找到另一种方法进行代替:Numpy.where()函数。

Numpy.where()函数

根据Numpy官方文档的描述,Numpy.where函数的格式为:numpy.where(condition[,x,y])。其接收一个包含ndarray对象的布尔表达式,并遍历其元素,如果该元素为真,则yield x,否则yield y,最终返回修改后的ndarray对象。

在本例中,我们可以将letter_convert函数改写为如下图所示的形式:

def vectorized_letter_convert(x):
    r = np.where(x>=60, "D", "F")
    r = np.where(x>=70, "C", r)
    r = np.where(x>=80, "B", r)
    r = np.where(x>=90, "A", r)
    return r

注意:在本例和其他涉及到数值比较的情况中,如上面的代码所示,可能需要将显式的if-elif-else条件排列顺序倒转。在多个where语句并列执行时,其相当于多条if语句并列执行,也就是说每一条where语句都会在每个单元格上执行。

最终如愿以偿:

我们将数据量提升至10,000条,进行性能分析:

.iterrows()方法:

.apply()方法:

矢量化操作:

由此足见矢量化方法所带来的可观性能提升。在涉及到对于DataFrame的遍历操作时,尽量避免使用iterrows方法,转向apply方法及矢量化方法是更好的选择。对于更复杂函数的矢量化仍在探索之中,希望可以在日后得到广泛应用。

参考文档:

  1. Pandas官方文档
  2. Numpy官方文档
  3. A Beginner’s Guide to Optimizing Pandas Code for Speed by Sofia Heisler