时间戳

首页 > 解决方案 > 加速大型数组和数据集的操作（Pandas 慢，Numpy 更好，进一步改进？）

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

问题描述

我有一个包含数百万行和大约 6 列的大型数据集。数据当前位于 Pandas 数据框中，我正在寻找对其进行操作的最快方法。例如，假设我想删除一列中值为“1”的所有行。

这是我的最小工作示例：

# Create dummy data arrays and pandas dataframe
array_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，在实践中，可以使用就地操作，一旦数据被丢弃，就用新数组替换旧数组

标签： pythonpandasperformancenumpyoptimization

解决方案

第 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，也许这对您来说是一个有趣的选择。

推荐阅读