加速大型数组和数据集的操作（Pandas 慢，Numpy 更好，进一步改进？）

Question

我有一个包含数百万行和大约 6 列的大型数据集。数据当前位于 Pandas 数据框中，我正在寻找对其进行操作的最快方法。例如，假设我想删除一列中值为“1”的所有行。这是我的最小工作示例：# Create dummy data arrays and pandas dataframearray_size = int(5e6)array1 = np.random.rand(array_size)array2 = np.random.rand(array_size)array3 = np.random.rand(array_size)array_condition = np.random.randint(0, 3, size=array_size)df = pd.DataFrame({'array_condition': array_condition, 'array1': array1, 'array2': array2, 'array3': array3})def method1():    df_new = df.drop(df[df.array_condition == 1].index)编辑：正如 Henry Yik 在评论中指出的，更快的 Pandas 方法是这样的：def method1b():    df_new = df[df.array_condition != 1]我相信 Pandas 在这种事情上可能会很慢，所以我还使用 numpy 实现了一个方法，将每一列作为一个单独的数组处理：def method2():    masking = array_condition != 1    array1_new = array1[masking]    array2_new = array2[masking]    array3_new = array3[masking]    array_condition_new = array_condition[masking]    结果：%timeit method1()625 ms ± 7.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)%timeit methodb()158 ms ± 7.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)%timeit method2()138 ms ± 3.8 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)因此，我们确实看到使用 numpy 的性能略有显着提升。然而，这是以可读性低得多的代码为代价的（即必须创建一个掩码并将其应用于每个数组）。这种方法似乎不像我有 30 列数据那样可扩展，我需要很多代码行来将掩码应用于每个数组！此外，允许可选列会很有用，因此此方法可能会在尝试对空数组进行操作时失败。因此，我有两个问题：1）在numpy中是否有更清洁/更灵活的方法来实现这一点？2）或者更好，我可以在这里使用任何更高性能的方法吗？例如 JIT（numba？）、Cython 还是其他？PS，在实践中，可以使用就地操作，一旦数据被丢弃，就用新数组替换旧数组

Answer 1

第 1 部分：Pandas 和（也许）Numpy

比较您的method1b和method2：

• method1b生成一个DataFrame，这可能是你想要的，

• method2生成一个Numpy 数组，因此要获得完全可比较的结果，您应该随后从中生成一个DataFrame。

所以我将您的方法2更改为：

def method2():

    masking = array_condition != 1

    array1_new = array1[masking]

    array2_new = array2[masking]

    array3_new = array3[masking]

    array_condition_new = array_condition[masking]

    df_new = pd.DataFrame({ 'array_condition': array_condition[masking],

        'array1': array1_new, 'array2': array2_new, 'array3': array3_new})

然后比较执行时间（使用%timeit）。

结果是我的method2 （扩展）版本的执行时间 比method1b长约5％（请自行检查）。

所以我的观点是，只要是单一的操作，可能还是和Pandas在一起比较好。

但是，如果您想在源 DataFrame 上按顺序执行几个操作和/或您对Numpy数组的结果感到满意，那么值得：

• 调用arr = df.values以获取底层Numpy数组。

• 使用Numpy方法对其执行所有必需的操作。

• （可选）从最终结果创建一个 DataFrame。

我尝试了method1b的Numpy版本：

def method3():
    a = df.values
    arr = a[a[:,0] != 1]

但执行时间要长约40%。

原因可能是Numpy数组具有相同类型的所有元素，因此array_condition列被强制浮动，然后创建整个Numpy数组，这需要一些时间。

第 2 部分：Numpy 和 Numba

要考虑的替代方法是使用Numba包 - 一种即时 Python 编译器。

我做了这样的测试：

创建了一个Numpy数组（作为初步步骤）：

a = df.values

原因是 JIT 编译的方法能够使用Numpy方法和类型，但不能使用Pandas的方法和类型。

为了执行测试，我使用了与上面几乎相同的方法，但使用了@njit注释（需要来自 numba import njit）：

@njit
def method4():
    arr = a[a[:,0] != 1]

这次：

• 执行时间约为method1b时间的 45% 。

• 但由于a = df.values已经在测试循环之前执行过，因此这个结果是否与之前的测试有可比性存在疑问。

无论如何，自己尝试Numba，也许这对您来说是一个有趣的选择。

Answer 2

您可能会发现在这里使用numpy.where很有用。它将布尔掩码转换为数组索引，使生活更便宜。将其与 numpy.vstack 结合可以实现一些内存便宜的操作：

def method3():

    wh = np.where(array_condition == 1)

    return np.vstack(tuple(col[wh] for col in (array1, array2, array3)))

这给出了以下时间：

>>> %timeit method2()

180 ms ± 6.66 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

>>> %timeit method3()

96.9 ms ± 2.5 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

元组解包允许该操作在内存上相当轻，因为当对象被 vstack-ed 重新组合在一起时，它会更小。如果您需要直接从 DataFrame 中获取列，则以下代码段可能有用：

def method3b():

    wh = np.where(array_condition == 1)

    col_names = ['array1','array2','array3']

    return np.vstack(tuple(col[wh] for col in tuple(df[col_name].to_numpy() 

        for col_name in col_names)))

这允许人们从 DataFrame 中按名称获取列，然后在运行中对这些列进行元组解包。速度差不多：

>>> %timeit method3b()

96.6 ms ± 3.09 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)