本文深入探讨了在pandas中对大型数据集进行分组抽样的优化策略，特别是当每个组需要不同的样本数量，并根据组内元素总数动态调整是否允许重复抽样（`replace`参数）时。通过介绍一种结合字典映射和`groupby().apply()`方法的解决方案，我们展示了如何避免低效的循环操作，从而显著提升处理效率，实现灵活且高性能的数据抽样。
在数据分析和机器学习任务中，我们经常需要从大型数据集中进行抽样。当数据需要根据某个或多个列进行分组，并且每个组的抽样需求（例如，抽样数量 n 和是否允许重复 replace）都不同时，传统的 df.groupby().sample() 方法可能无法满足所有要求。尤其对于包含数千万甚至上亿条记录的数据集，以及数十万个唯一分组的情况，低效的抽样方法会导致严重的性能瓶颈。
问题背景与传统方法的局限性
假设我们有一个大型DataFrame df，其中包含一个分组列 "a"，以及一个记录每个组所需抽样数量的DataFrame df_counts。我们的目标是根据 df_counts 中为每个 "a" 组指定的 count 值进行抽样。
1. 基础的 groupby().sample()： Pandas提供了 df.groupby("a").sample(n=k)，这可以对每个组抽取固定数量 k 的样本。但这种方法无法为每个组指定不同的 n 值。
# 示例：每个组抽取1个样本
# df.groupby("a").sample(n=1, random_state=1)这无法满足不同组不同 n 的需求。
2. 循环迭代的低效性： 一种直观但效率低下的方法是遍历每个唯一的组，筛选出该组的数据，然后进行抽样，最后将所有结果拼接起来。此外，为了更灵活地控制抽样行为，我们可能还需要根据组内元素的总数与所需样本量 n 的关系，动态地决定 replace 参数（即当组内元素不足 n 时允许重复抽样，否则不允许）。
# 伪代码：循环迭代的低效方法 # sampled_dfs = [] # for group_val in df['a'].unique(): # filter_df = df.loc[df['a'] == group_val] # n_samples = get_n_from_df_counts(group_val) # 从df_counts获取该组的n值 # # if len(filter_df) >= n_samples: # sampled_group = filter_df.sample(n=n_samples, random_state=6, replace=False) # else: # sampled_group = filter_df.sample(n=n_samples, random_state=6, replace=True) # sampled_dfs.append(sampled_group) # # final_sampled_df = pd.concat(sampled_dfs)
对于拥有100k个唯一值的分组列和90M行的数据集，这种基于Python循环的逐组筛选和抽样会带来巨大的性能开销，因为每次迭代都会产生新的DataFrame对象和额外的内存操作。
优化方案：结合 groupby().apply() 和字典映射
为了解决上述问题，我们可以利用Pandas的 groupby().apply() 方法，结合一个预先构建的字典来高效地传递每个组的抽样参数。apply() 方法虽然在某些情况下不如完全向量化的操作快，但它将对每个组的操作封装在一个函数中，并由Pandas在内部进行管理，通常比显式的Python循环效率更高。
核心思路是：
- 构建样本量字典： 将包含每个组所需样本量的数据转换成一个字典，其中键是分组列的值，值是对应的样本量。
- 定义自定义抽样函数： 创建一个函数，该函数接收一个组的DataFrame，并根据字典查找该组的样本量 n。同时，根据组的实际大小与 n 的关系，动态设置 sample() 方法的 replace 参数。
- 应用自定义函数： 使用 df.groupby().apply() 将自定义抽样函数应用到每个组。
步骤一：准备样本计数数据
首先，我们需要一个DataFrame来定义每个分组的样本数量。
import pandas as pd
import numpy as np
# 示例输入数据
# df_counts 定义了每个 'a' 组需要抽样的数量
df_counts = pd.DataFrame({
'a': [1, 2, 3],
'count': [1, 3, 2]
})
# 原始数据 df_original，我们将从中抽样
df_original = pd.DataFrame({
'a': [1, 1, 1, 2, 2, 3, 3],
'x': ['a', 'b', 'c', 'd', 'e', 'f', 'g']
})
print("df_counts:")
print(df_counts)
print("\ndf_original:")
print(df_original)输出:
df_counts: a count 0 1 1 1 2 3 2 3 2 df_original: a x 0 1 a 1 1 b 2 1 c 3 2 d 4 2 e 5 3 f 6 3 g
接下来，将 df_counts 转换为一个字典，以便在抽样函数中快速查找。
# 构建样本量字典
sample_counts_dict = df_counts.set_index("a")["count"].to_dict()
print("\nsample_counts_dict:")
print(sample_counts_dict)输出:
sample_counts_dict:
{1: 1, 2: 3, 3: 2}步骤二：定义自定义抽样函数
这个函数将是 groupby().apply() 的核心。它接收一个组的DataFrame，一个包含所有组样本量的字典，以及一个随机种子。
def get_sample(group_df, sample_dict, random_state=None):
"""
对给定的组DataFrame进行抽样。
参数:
group_df (pd.DataFrame): 当前分组的DataFrame。
sample_dict (dict): 包含每个组所需样本量的字典。
random_state (int, optional): 随机种子，用于结果可复现性。
返回:
pd.DataFrame: 抽样后的DataFrame，如果该组没有对应的样本量则返回None。
"""
# 获取当前组的键（例如 'a' 列的值）
group_key = group_df["a"].iat[0]
# 从字典中获取该组的样本量n
n_samples = sample_dict.get(group_key)
# 如果字典中没有该组的样本量，则返回None（表示不抽样或跳过）
if n_samples is None:
return None
# 动态设置 replace 参数
# 如果组内元素数量小于所需样本量n，则必须允许重复抽样 (replace=True)
# 否则，如果组内元素数量足够，则默认不允许重复抽样 (replace=False)
# 这样可以最大化地获取唯一行
replace_flag = len(group_df) <= n_samples
return group_df.sample(n=n_samples, random_state=random_state, replace=replace_flag)
步骤三：应用自定义函数进行抽样
最后，我们将 get_sample 函数应用到 df_original 的每个组。
# 应用自定义函数进行分组抽样
# group_keys=False 可以避免将分组键作为额外的索引添加到结果中，保持输出整洁
sampled_output_df = df_original.groupby("a", group_keys=False).apply(
get_sample,
sample_dict=sample_counts_dict,
random_state=6
)
print("\nDesired sampled output:")
print(sampled_output_df)输出:
Desired sampled output: a x 0 1 a 3 2 d 4 2 e 4 2 e 5 3 f 6 3 g
可以看到，对于 a=1，原始数据有3条，需要抽样1条，结果是 a。 对于 a=2，原始数据有2条（'d', 'e'），需要抽样3条。由于 len(group_df)
性能考量与注意事项
- groupby().apply() 的优势： 相较于显式的Python循环，apply() 在内部进行了优化，尤其是在处理大量组时，可以减少Python解释器和Pandas对象之间的切换开销。它允许我们将自定义逻辑高效地应用于每个组。
- group_keys=False： 在 groupby().apply() 中设置 group_keys=False 可以防止分组键作为额外的索引添加到结果DataFrame中，这在结果拼接时可以避免不必要的索引重置操作，并保持输出结构的简洁。
- replace 参数的动态控制： 这种方法允许我们根据每个组的实际情况，灵活地设置 replace 参数，确保在组内元素不足以满足抽样数量时，能够通过重复抽样来达到目标数量，而在元素充足时，则优先进行无重复抽样。
- 随机种子 random_state： 使用 random_state 参数可以确保抽样结果的可复现性，这在调试和结果验证时非常重要。
- 大数据集的挑战： 尽管 apply() 比循环更优，但对于拥有极其庞大数量的组（例如数百万个组），且每个组的数据量都非常小的情况，apply() 仍然可能面临性能挑战。在这种极端情况下，可能需要考虑使用更底层的库（如Numba）或分布式计算框架（如Spark）来进一步优化。然而，对于大多数常见的大数据集场景，groupby().apply() 配合自定义函数是一个非常高效且灵活的解决方案。
总结
本文介绍了一种在Pandas中对大型数据集进行高效分组抽样的专业方法。通过将每个组的样本数量预先存储在一个字典中，并结合 groupby().apply() 方法和一个动态控制 replace 参数的自定义函数，我们能够克服传统方法的局限性，实现灵活且高性能的抽样。这种模式在处理复杂分组抽样需求时，是提升代码效率和可维护性的关键策略。
