本文介绍在不加载全部数据到内存的前提下，使用 pyarrow 原生 api 对百gb级 arrow 表（如 hugging face dataset 导出的 `.arrow` 文件）进行高效、可复现、带/不带放回的随机采样，并保持原始索引完整性。
当处理像 181 GB、3000 万行这样的超大规模 PyArrow 数据集时，常规思路（如 to_pandas() 后调用 df.sample()）极易因内存溢出导致 Python 进程崩溃。根本原因在于：Pandas 需将整个 Arrow 表解码为列式 NumPy 数组并载入内存——这不仅浪费资源，更违背了 Arrow “零拷贝”与“延迟计算”的设计哲学。
幸运的是，PyArrow 提供了完全内存友好的原生采样能力：Table.take(indices)。该方法仅根据整数索引序列，以 O(1) 时间复杂度按需提取指定行，底层不触发全量解码，也不复制原始数据块，真正实现“采样即取用”。
以下是一个健壮、可复用的采样函数，支持带放回（replace=True）与不带放回（replace=False）两种模式，并保留原始行索引（便于后续溯源或去重分析）：
import pyarrow as pa
import random
import numpy as np
def sample_table(
table: pa.Table,
n_sample_rows: int,
replace: bool = False,
seed: int = None
) -> pa.Table:
"""
从 PyArrow Table 中随机采样指定行数，支持带/不带放回。
Args:
table: 输入的 PyArrow Table（可直接来自 Dataset.to_table()）
n_sample_rows: 采样行数
replace: 是否允许重复采样同一行（True=带放回）
seed: 随机种子，确保结果可复现
Returns:
新的 PyArrow Table，包含采样后的行，原始列结构与元数据完全保留
"""
if seed is not None:
random.seed(seed)
np.random.seed(seed) # 兼容 np.random.choice 的行为
if n_sample_rows <= 0:
raise ValueError("n_sample_rows must be positive")
if n_sample_rows >= table.num_rows and not replace:
return table # 全量返回，无需采样
if replace:
# 带放回：使用 np.random.choice 更高效（支持重复索引）
indices = np.random.choice(table.num_rows, size=n_sample_rows, replace=True)
else:
# 不带放回：使用 random.sample（更高效且无 numpy 依赖）
indices = random.sample(range(table.num_rows), k=n_sample_rows)
return table.take(indices)
# ✅ 使用示例：从 Hugging Face Dataset 加载并采样（无需转 Pandas！）
from datasets import Dataset
# 直接加载 .arrow 文件为 Dataset，再转为 Table（轻量，不加载数据）
ds = Dataset.from_file("embeddings_job/combined_embeddings_small/data-00000-of-00001.arrow")
table = ds.to_table() # 此步仅读取元数据，毫秒级完成
# 采样 100 行（带放回，用于 20 次独立训练）
for i in range(20):
sampled_table = sample_table(table, n_sample_rows=100, replace=True, seed=42 + i)
# → sampled_table 是标准 pa.Table，可直接送入 sklearn / XGBoost 等库
# （例如：scikit-learn 支持 pa.Array / pa.ChunkedArray 作为输入）⚠️ 关键注意事项：
- 索引保留性：table.take(indices) 返回的新表中，每行仍携带其原始全局索引（可通过 sampled_table.schema.metadata 或自定义字段记录），不会被重置为 0,1,2,...；若需显式保留原始行号，可在采样前添加索引列：table = table.add_column(0, "original_idx", pa.array(range(table.num_rows)))。
- 性能优势：相比 shuffle().select()（需全表重排，O(n log n)），take() 是纯索引查找，时间复杂度 O(k)，k 为采样数，对 100 行采样几乎瞬时完成。
- 避免终端 shuf：shuf 等命令行工具操作的是文本文件，而 Arrow 是二进制列式格式，无法直接使用；强行转换将丢失类型、压缩和零拷贝优势，得不偿失。
- 与下游模型兼容：现代机器学习库（如 scikit-learn ≥ 1.3、XGBoost ≥ 2.0）已原生支持 PyArrow 数组作为特征输入，无需中间转 Pandas，进一步规避内存瓶颈。
总结而言，table.take() + random.sample() 或 np.random.choice() 是处理超大 Arrow 数据集随机采样的黄金组合。它轻量、可靠、可复现，且完全契合 Arrow 的设计范式——让数据留在磁盘或内存映射区，只把真正需要的子集“拉”出来计算。对于您计划运行的 20 次随机森林训练，只需循环调用该函数生成 20 个独立 pa.Table 实例，即可安全、高效地完成全流程。
