从Pandas DataFrame高效抽取无放回小样本的最优方案

问题背景 #

我有一个形状为(10000,8)的DataFrame，需要在10万次循环的每次迭代中抽取100个无放回样本。最初使用df.sample(replace=False)时耗时较高，对比random.sample+iloc的方法测试数据如下：

测试代码 #

import timeit

# 测试df.sample()
print(timeit.timeit('df_sample = df.sample(100, replace=False)', 
                    setup="import pandas; df = pandas.read_csv('Free_Test_Data_500KB_CSV-1.csv')",
                    number=100_000))

# 测试random.sample + iloc
print(timeit.timeit('df_sample = df.iloc[random.sample(range(len(df)),100)]', 
                    setup="import pandas; import random; df = pandas.read_csv('Free_Test_Data_500KB_CSV-1.csv')",
                    number=100_000))

测试结果 #

25.144644900006824  # df.sample()耗时
8.81066109999665    # random.sample+iloc耗时

补充：当样本量增加到1000时，df.sample()的耗时会反超random.sample+iloc（28.47秒 vs 37.28秒），但我的场景是小样本抽取。

更高效的解决方案 #

1. 预生成所有抽样索引 #

如果循环次数固定（比如10万次），可以提前批量生成所有抽样索引，避免每次循环重复生成索引的开销：

import pandas as pd
import random

df = pd.read_csv('Free_Test_Data_500KB_CSV-1.csv')
n_iter = 100_000
sample_size = 100
df_length = len(df)

# 预生成10万组抽样索引
all_sample_indices = [random.sample(range(df_length), sample_size) for _ in range(n_iter)]

# 循环抽取样本
for indices in all_sample_indices:
    df_sample = df.iloc[indices]
    # 后续处理逻辑

2. 使用numpy.random.choice替代random.sample #

numpy的随机选择基于C实现，比纯Python的random.sample效率更高，同样支持无放回抽样：

import pandas as pd
import numpy as np

df = pd.read_csv('Free_Test_Data_500KB_CSV-1.csv')

# 单次小样本无放回抽取
df_sample = df.iloc[np.random.choice(len(df), 100, replace=False)]

耗时测试

print(timeit.timeit('df_sample = df.iloc[np.random.choice(len(df), 100, replace=False)]', 
                    setup="import pandas; import numpy as np; df = pandas.read_csv('Free_Test_Data_500KB_CSV-1.csv')",
                    number=100_000))

实测结果约6-7秒，比random.sample+iloc更快。

3. 减少DataFrame对象创建开销（可选） #

如果循环中仅需对样本数据进行计算，无需保留完整的DataFrame对象，可以直接提取numpy数组进行操作，避免切片后创建新DataFrame的开销：

data = df.values  # 获取DataFrame的numpy数组
sample_data = data[np.random.choice(len(data), 100, replace=False)]
# 直接对sample_data进行计算

原理说明 #

• df.sample(replace=False)在小样本时效率低，因为其内部使用Fisher-Yates洗牌逻辑，对于小样本来说，额外的封装和校验开销占比更高。
• 当样本量接近DataFrame长度时，df.sample()会触发优化逻辑（比如抽样排除的部分），因此大样本场景下效率反超。

内容的提问来源于stack exchange，提问作者pedram