优化 Pandas 数据处理性能：从 apply 到矢量化操作

引言

在处理大规模数据时，性能往往是一个关键问题。本文将通过一个具体示例，展示如何优化 Pandas 数据处理的性能，从而大幅提高代码运行效率。我们将从一个需求引导开始，分析为什么 apply 方法会慢，并给出多种优化方式，最后总结这些优化方法的应用场景和效果。

假设我们有一个日志文件数据集，包含数百万条记录。我们需要将这些日志记录中的信息格式化成特定的字符串，以便后续处理。数据集的字段包括时间戳、进程ID、日志级别、文件类型、文件位置和日志信息。

原始数据示例如下：

n = 4000000
data = {
'sample_ts_str': np.random.choice(pd.date_range('2021-01-01', periods=1000).strftime('%Y-%m-%d %H:%M:%S'), n),
'pid': np.random.randint(1000, 9999, n),
'level': np.random.choice(['INFO', 'DEBUG', 'WARN', 'ERROR'], n),
'file_type': np.random.choice(['log', 'txt', 'csv'], n),
'file_pos': np.random.randint(1, 1000, n),
'log': np.random.choice(['log message 1', 'log message 2', 'log message 3'], n)
}
df = pd.DataFrame(data)

我们这里使用代码模拟生成4000000条数据，需求是将这些字段组合成格式为 [timestamp] [pid] [level] [file_type] [file_pos]log 的字符串，并存储在新的列中。

原始 apply 方法

最直接的方法是使用 apply 方法逐行处理数据：

df['origin_log'] = df.apply(
    lambda x: '[%s] [%s] [%s] [%s] [%s]%s' % (
        x['sample_ts_str'], x['pid'], x['level'], x['file_type'], x['file_pos'], x['log']), axis=1)

虽然这种方法简洁明了，但在处理大规模数据时性能较差。通过时间测试，我们发现处理400万条记录需要约53.45秒。

Original apply method time: 53.45 seconds

为什么 apply 方法会慢？

apply 方法的主要问题在于它逐行调用 Python 函数，即使底层数据是由高效的 C 语言实现的 Pandas 数组结构支持。这种逐行操作的开销非常大，尤其是在处理大型数据集时，性能瓶颈更加明显。

矢量化操作

为了解决 apply 方法的性能问题，我们可以采用 Pandas 提供的矢量化操作。这种方法通过一次性对整个数组进行操作，极大地提高了处理速度。矢量化操作的实现如下：

df['origin_log'] = (
    '[' + df['sample_ts_str'].astype(str) + '] ' +
    '[' + df['pid'].astype(str) + '] ' +
    '[' + df['level'].astype(str) + '] ' +
    '[' + df['file_type'].astype(str) + '] ' +
    '[' + df['file_pos'].astype(str) + ']' +
    df['log'].astype(str)
)

通过时间测试，处理同样的400万条记录，矢量化方法只需要约6.61秒。

Vectorized method time: 6.61 seconds

为什么矢量化操作更快？

矢量化操作之所以更快，是因为它利用了底层高效的数组操作，避免了逐行调用 Python 函数的开销。Pandas 的底层实现基于高度优化的 C 语言库，如 NumPy，使得矢量化操作能够直接对内存中的数据进行批量处理，从而显著提高性能。

这种优化方法适用于任何需要对大规模数据进行逐行处理的场景，例如日志分析、数据清洗和格式转换等。在实际应用中，采用矢量化操作可以显著减少处理时间，提高数据处理效率。