跳至主要内容

高性能 pandas 替代方案的基准测试

了解对 Python 强大 pandas 替代方案 Polars、Vaex 和 Datatable 的最新基准测试。比较它们在数据加载、分组、排序等方面的表现。
· 13分钟 阅读

使用 AI 探索

简介

Pandas 是 Python 数据科学生态中功能极其强大的工具,提供了多种数据处理与清洗能力。然而,它在处理中等规模数据集时表现出色,但面对更大规模数据集时可能出现性能瓶颈,因此需要考虑高性能替代方案。

本文将系统介绍这些替代方案,并通过基准测试,从数据加载时间、执行时间、内存占用、可扩展性和易用性等方面进行对比。

注意:基准测试结果会因您的硬件环境而有较大差异。

理解基准测试

基准测试是用于对软件或硬件进行性能评估的参照点。在软件性能优化中,它允许我们衡量不同方法、算法或工具的效率。在本文场景下,数据处理类库的关键基准指标包括执行时间、内存占用、可扩展性和易用性。

高性能替代方案简介

主要的 Python 替代工具包括 Polars、Vaex 和 Datatable。在深入比较之前,先快速概览一下各工具。

Polars

Polars 可由三个关键特性来概括。其一,它提供了完整的 Python API,包含大量用于处理数据框的函数。其二,它既可作为数据框库使用,也可作为数据模型的底层查询引擎。其三,通过采用 Apache Arrow 规范中安全的 Arrow2 实现,Polars 在处理海量数据时具备极高效率。

您可以查看我们的Polars 入门教程,以及pandas 2.0 与 Polars 的比较。

Vaex

Vaex 面向惰性、外存(out-of-core)的 DataFrame(类似 Pandas),用于可视化和探索大型表格数据集。其高效之处在于采用惰性求值,仅在必要时执行操作,从而降低内存占用和时间成本。

Datatable

Datatable 可在单机环境下以尽可能高的速度处理大型数据(可至 100GB)。其一大优势是与 Pandas/NumPy/纯 Python 的互操作性,便于用户在不同数据处理框架间转换。

若想进一步了解 pandas,我们的Pandas 教程:Python 中的 DataFrame是不错的起点。

此外,Pandas 速查表:Python 数据科学提供了该数据分析库的基础速览,并附带代码示例。

搭建基准测试环境

本节介绍基准测试数据的生成方法,以及参与基准分析的全部库的安装流程。

基准测试数据

本次基准测试的数据集大小为 5.7 GB,足以进行较为充分的对比。每一行被重复 100000 次,总计 76,800,000 行。原始数据使用的是糖尿病数据,可从 Kaggle 免费获取。

最终的基准测试笔记本可在GitHub上查看。

import pandas as pd

data_URL = "https://raw.githubusercontent.com/keitazoumana/Experimentation-Data/main/diabetes.csv"

original_data = pd.read_csv(data_URL)

# Duplicate each row 100000 times
benchmarking_df = original_data.loc[original_data.index.repeat(100000)]

# Save the result as a CSV file for future use.
file_name = "benchmarking_data.csv"

benchmarking_df.to_csv(file_name, index=False)

安装相关库

现在,可以开始安装参与基准分析的所有库。

在 Jupyter Notebook 环境中,可通过如下 pip install 命令安装:

%%bash
pip3 -q install -U polars
pip3 -q install vaex
pip -q install datatable

成功运行上述代码块后,即可通过 import 语句导入:

import pandas as pd
from time import time
import polars as pl
import vaex as vx
import datatable as dt

此外,我们使用 plotly 库进行可视化。

import plotly.express as px

至此,便可开始进行基准分析。

执行时间基准测试

执行时间将针对以下操作进行评估:数据加载、数据卸载、数据聚合和数据筛选。我们为每个任务实现了一个辅助函数,返回包含两个主要键的字典:库名与执行时间。

随后使用绘图函数以 Plotly Express 将结果可视化。该函数接收两个参数:来自各辅助函数的字典列表,以及图表标题。

def plot_metrics(list_exec_time, graph_title):

	df = pd.DataFrame(list_exec_time)

	fig = px.bar(df, x='library', y= 'execution_time', title=graph_title)
    
	fig.show()

数据加载

辅助函数 read_csv_with_time 有两个主要参数:要读取的文件名和库名。

def read_csv_with_time(library_name, file_name):

	final_time = 0

	start_time = time.time()

	if library_name.lower() == 'polars':
    	df = pl.read_csv(file_name)

	elif library_name.lower() == 'pandas':
    	df = pd.read_csv(file_name)

	elif library_name.lower() == 'vaex':
    	df = vx.read_csv(file_name)

	elif library_name.lower() == 'datatable':
    	df = dt.fread(file_name)

	else:
    	raise ValueError("Invalid library name. Must be 'polars', 'pandas', 'vaex', or 'datatable'")

	end_time = time.time()

	final_time = end_time - start_time

	return {"library": library_name, "execution_time": final_time}

可对每个库分别调用该函数,从 pandas 开始。

pandas_time, pandas_df = read_csv_with_time('pandas', file_name)
polars_time, polars_df = read_csv_with_time('polars', file_name)
vaex_time, vaex_df = read_csv_with_time('vaex', file_name)
datatable_time, dt_df = read_csv_with_time('datatable', file_name)

将得到的字典合并为列表,并用以下表达式进行绘图。

exec_times = [pandas_time, polars_time,
           vaex_time, datatable_time]

def plot_metrics(list_exec_time, graph_title):

 [print(exec_time) for exec_time in list_exec_time]
 df = pd.DataFrame(exec_times)

 # Plot bar plot using Plotly Express
 fig = px.bar(df, x='library', y='execution_time', title=graph_title)
 fig.show()

plot_metrics(exec_times, "Read Execution Time Comparison")
{'library': 'pandas', 'execution_time': 35.07908916473389}
{'library': 'polars', 'execution_time': 6.041005373001099}
{'library': 'vaex', 'execution_time': 36.431522607803345}
{'library': 'datatable', 'execution_time': 4.4584901332855225}

数据加载执行时间的字典输出

数据加载执行时间的图形展示

根据图表,我们可以得出以下结论:

  • Polars 比 Datatable 慢 1.35 倍
  • Pandas 比 Datatable 慢 7.87 倍。
  • Vaex 比 Datatable 慢约 8.17 倍。

数据分组

同样,group_data_with_time 用于记录对指定列进行分组的执行时间。在本例中,我们使用 Pregnancies 列。

from time import time

def group_data_with_time(library_name, df, column_name='Pregnancies'):

 start_time = time()

 if library_name.lower() == 'polars':
     df_grouped = df.group_by(column_name).first()

 elif library_name.lower() == 'vaex':
     df_grouped = df.groupby(column_name, agg='first')

 elif library_name.lower() == 'pandas':
     df_grouped = df.groupby(column_name).first()

 elif library_name.lower() == 'datatable':
     df_grouped = df[:, dt.first(dt.f[:]), dt.by(column_name)]

 else:
     raise ValueError("Invalid library name. Must be 'polars', 'vaex', or 'datatable'")

 end_time = time()

 final_time = end_time - start_time

 return {"library": library_name, "execution_time": final_time}
pandas_time = group_data_with_time('pandas', pandas_df)
polars_time = group_data_with_time('polars', polars_df)
vaex_time = group_data_with_time('vaex', vaex_df)
datatable_time = group_data_with_time('datatable', dt_df)
exec_times = [pandas_time, polars_time,
           vaex_time, datatable_time]
plot_metrics(exec_times, "Grouping Execution Time Comparison")
{'library': 'pandas', 'execution_time': 2.382519245147705}
{'library': 'pandas', 'execution_time'': 1.0898876190185547}
{'library': 'pandas', 'execution_time': 0.8506548404693604}
{'library': 'pandas', 'execution_time': 0.34698963165283203}

数据分组执行时间的图形展示

针对更新后的数据分组任务,将最快的 DataTable 与其他库相比:

可以看到,DataTable 是本次评测中最快的。

  • Vaex 比 DataTable 慢 2.45 倍。
  • Polars 比 DataTable 慢 3.14 倍。
  • Pandas 比 DataTable 慢 6.87 倍。

上述分析显示,DataTable 在数据分组任务上效率最佳,Vaex 的表现相对接近。Polars 和 Pandas 的执行时间相近,但都明显慢于 DataTable。

列排序

采用同样的方法,列排序函数会使用各库对指定列进行排序。

def sort_data_with_time(library_name, df, column_name='Pregnancies'):
    
	start_time = time()

	if library_name.lower() == 'polars':
    	df_sorted = df.sort(column_name)
	elif library_name.lower() == 'vaex':
    	df_sorted = df.sort(column_name)

	elif library_name.lower() == 'datatable':
    	df_sorted = df.sort(column_name)
   	 
	elif library_name.lower() == 'pandas':
    	df_sorted = pd.DataFrame(df).sort_values(column_name)
	else:
    	raise ValueError("Invalid library name. Must be 'polars', 'vaex', 'datatable', or 'pandas'")

	end_time = time()

	final_time = end_time - start_time

	return {"library": library_name, "execution_time": final_time}
pandas_time = sort_data_with_time('pandas', pandas_df)
polars_time = sort_data_with_time('polars', polars_df)
vaex_time = sort_data_with_time('vaex', vaex_df)
datatable_time = sort_data_with_time('datatable', dt_df)
{'library': 'pandas', 'execution_time': 6.735127687454224}
{'library': 'polars', 'execution_time': 4.119458436965942}
{'library': 'vaex', 'execution_time': 1.0371737480163574}
{'library': 'datatable', 'execution_time': 0.3971593379974365}

数据排序执行时间的字典输出

数据排序执行时间的图形展示

在数据排序任务中,DataTable 是本次评测中最快的库。

  • Vaex 比 DataTable 慢 2.61 倍。
  • Polars 比 DataTable 慢 10.37 倍。
  • Pandas 比 DataTable 慢 16.96 倍。

该结果表明 DataTable 在数据排序任务上具备显著优势,明显快于 Vaex、Polars 和 Pandas;Vaex 次之,但仍然慢于 DataTable 两倍以上。

数据卸载

数据卸载的思路是将原始数据转换为另一种格式,在这里为 numpy 数组。

def offload_data_with_time(library_name, df):
    
	start_time = time()

	if library_name.lower() == 'polars':
    	array = df.to_numpy()
   	 
	elif library_name.lower() == 'vaex':
    	array = df.to_pandas_df().values

	elif library_name.lower() == 'datatable':
    	array = df.to_numpy()
   	 
	elif library_name.lower() == 'pandas':
    	array = pd.DataFrame(df).values
	else:
    	raise ValueError("Invalid library name. Must be 'polars', 'vaex', 'datatable', or 'pandas'")

	end_time = time()

	final_time = end_time - start_time

	return {"library": library_name, "execution_time": final_time}
exec_times = [pandas_time, polars_time,
           vaex_time, datatable_time]

plot_metrics(exec_times, "Data offloading Execution Time Comparison")
{'library': 'pandas', 'execution_time': 1.8406445980072021}
{'library': 'polars', 'execution_time': 3.238591432571411}
{'library': 'vaex', 'execution_time': 6.945307970046997}
{'library': 'datatable', 'execution_time': 2.634136438369751}

数据卸载执行时间的字典输出

数据卸载执行时间的图形展示

数据卸载是执行时间基准测试的最后一轮。

本轮中,Vaex 的执行时间最高。将 pandas 与其他库对比可见,其执行时间大致为:

  • DataTable 比 Pandas 慢 1.43 倍。
  • Polars 比 Pandas 慢 1.76 倍。
  • Vaex 比 Pandas 慢 3.77 倍。

内存占用基准测试

在比较不同库时,执行时间并非唯一标准。了解它们处理内存的效率同样重要。本节使用 tracemalloc 库展示两类内存占用:

  • 当前内存占用。指使用特定库执行期间的总内存占用,对应当下被占用的 RAM 空间。
  • 峰值内存占用。指程序使用内存的最大值。需要注意,峰值内存总是高于当前内存。

通过这些可视化可以了解哪些程序内存占用更高,可能导致内存溢出,尤其在处理大型数据集时尤为重要。

除 tracemalloc 外,还需要 os 库。

import tracemalloc as tm
import os

开始之前,我们创建一个空列表,用于保存每次内存占用的结果。

list_memory_usage = []

对于每个库,内存占用的估计遵循以下通用写法(此处以 Vaex 为例)。此外,这里仅关注数据卸载任务的内存占用。

tm.start()
vaex_time = offload_data_with_time('vaex', vaex_df)
memory_usage = tm.get_traced_memory()
tm.stop()

list_memory_usage.append({
	'library': 'vaex',
	'memory_usage': memory_usage
})

其余库(Polars、pandas、Datatable)的写法相同,分别如下:

tm.start()
polars_time = offload_data_with_time('polars', polars_df)
memory_usage = tm.get_traced_memory()

tm.stop()

list_memory_usage.append({
	'library': 'polars',
	'memory_usage': memory_usage
})
tm.start()
offload_data_with_time('pandas', pandas_df)

# Get the memory usage
memory_usage = tm.get_traced_memory()

tm.stop()

list_memory_usage.append({
	'library': 'pandas',
	'memory_usage': memory_usage
})
tm.start()
datatable_time = offload_data_with_time('datatable', datatable_df)
memory_usage = tm.get_traced_memory()

tm.stop()

list_memory_usage.append({
	'library': 'datatable',
	'memory_usage': memory_usage
})

运行完上述代码后,可使用以下辅助函数绘制图表:

def plot_memory_usage(list_memory_usage, graph_title='Memory Usage by Library'):

	df = pd.DataFrame(list_memory_usage)

	# separate the memory usage tuple into two columns: current_memory and peak_memory
	df[['current_memory', 'peak_memory']] = pd.DataFrame(df['memory_usage'].tolist(), index=df.index)

	# now we no longer need the memory_usage column
	df = df.drop(columns='memory_usage')

	# melt the DataFrame to make it suitable for grouped bar chart
	df_melted = df.melt(id_vars='library', var_name='memory_type', value_name='memory')

	# create the grouped bar chart
	fig = px.bar(df_melted, x='library', y='memory', color='memory_type', barmode='group',
             	labels={'memory':'Memory Usage (bytes)', 'library':'Library', 'memory_type':'Memory Type'},
             	title=graph_title)
    
	fig.update_layout(yaxis_type="log")
	fig.show()

可以注意到:

  • Pandas 的当前内存占用大约是 Polars 的 1100 倍、Vaex 的 7.4 倍、Datatable 的 29.4 倍。
  • 在峰值内存方面,pandas 和 Vaex 约比 Polars 多用 96.3 万倍内存,比 Datatable 多用 13.1 万倍。

进一步的分析(例如 Dask 与 pandas 的速度对比)有助于获得更全面的视角。

Pandas 替代方案对比表

下方为对比表,汇总了 Polars、Vaex 与 Datatables 之间的差异:

基准维度

Polars

Vaex

Datatable

数据加载时间

快于 Vaex,慢于 Datatable

慢于 Polars 与 Datatable

三者中最快

数据分组时间

三者中最慢

快于 Polars,慢于 Datatable

三者中最快

数据排序时间

快于 Vaex,慢于 Datatable

慢于 Polars 与 Datatable

三者中最快

数据卸载时间

快于 Vaex,慢于 Datatable

三者中最慢

快于 Polars,慢于 Vaex

当前内存占用

占用最少

多于 Polars,少于 Datatable

三者中占用最多

峰值内存占用

占用最少

多于 Polars,少于 Datatable

三者中占用最多

可扩展性

可随数据规模扩展

可随数据规模扩展,但数据增大时可能使用更多内存

对数据规模高度可扩展

易用性

具备完整 Python API,兼容 Apache Arrow

采用惰性求值,兼容 Pandas

与 Pandas/NumPy/纯 Python 互操作

结论

在评估可扩展性与集成便利性时,需要综合考虑多种因素。我们的基准测试表明,尽管 pandas 成熟易用,但学习曲线更为平缓,它在处理超大规模数据集时效率并不理想。

不过,仍有多种方法可以加速 pandas,我们的高性能数据处理:Pandas 2.0 vs Polars将对此有所帮助。

Vaex 的性能因任务而异。因此,选择 Pandas 的替代方案,应根据用户的具体需求、数据集规模,以及其对学习成本与集成工作的承受能力来决定。

Zoumana Keita

作者

Zoumana Keita

一位多才多艺、乐于分享知识并回馈他人的数据科学家。Zoumana 是 YouTube 内容创作者,也是 Medium 上的顶尖科技写作者。他热爱演讲、编程与教学。Zoumana 拥有两个硕士学位:其一是在法国巴黎获得的计算机科学硕士,专注于机器学习;其二是在美国德州理工大学获得的数据科学硕士。他的职业生涯始于法国 Groupe OPEN 的软件开发工程师,随后加入 IBM 担任机器学习顾问,为保险公司构建端到端的 AI 解决方案。之后,Zoumana 加入了 Axionable——一家总部位于巴黎与蒙特利尔的首家“可持续 AI”初创公司。在那里,他担任数据科学家,为来自法国、蒙特利尔、新加坡和瑞士的客户交付 AI 产品,主要聚焦 NLP 场景。此外,他将 5% 的时间投入到研发工作。目前,他在 IFC(世界银行集团)担任高级数据科学家。

主题