如何在单次遍历中按值划分Arrow Table？

完全可以做到仅单次遍历原表就完成按指定列取值的划分，不用针对每个分组值循环扫描全表。核心思路是先通过矢量化操作一次性计算出每行的分组归属，再基于归属索引拆分表，全程只扫一次原表。

具体步骤与代码示例（以PyArrow为例） #

• 先为已知的分组值建立值-索引映射，方便后续快速定位分组
• 用Arrow的计算函数一次性生成每行的分组索引（仅遍历原表一次）
• 基于分组索引提取对应行生成子表（此步骤无需再次扫描原表数据）

import pyarrow as pa
import pyarrow.compute as pc

# 示例Arrow Table
sample_data = {
    'group_col': ['X', 'Y', 'X', 'Z', 'Y', 'X'],
    'data_col': [10, 20, 30, 40, 50, 60]
}
source_table = pa.Table.from_pydict(sample_data)

# 已知的分组取值集合
target_groups = ['X', 'Y', 'Z']
# 建立分组值到索引的映射
group_mapping = {val: idx for idx, val in enumerate(target_groups)}

# 单次遍历计算所有行的分组索引（矢量化操作，效率极高）
row_group_indices = pc.index_in(source_table['group_col'], pa.array(target_groups))

# 根据索引拆分表（循环仅遍历分组数，而非原表）
split_tables = {}
for group_val, idx in group_mapping.items():
    # 筛选出当前分组的行索引
    group_mask = pc.equal(row_group_indices, idx)
    selected_rows = pc.select_indices(group_mask)
    # 提取子表
    split_tables[group_val] = source_table.take(selected_rows)

# 输出验证
for g, tbl in split_tables.items():
    print(f"分组 {g}:")
    print(tbl.to_pandas())

关键说明 #

1. 矢量化计算的优势：pc.index_in是Arrow底层优化的矢量化操作，仅需遍历原表一次就能完成所有行的分组匹配，比循环扫描n次的效率提升显著，尤其适合大表场景。
2. 关于分组扫描的误区：Arrow的scan_groups主要用于流式处理分组数据（比如逐组聚合），并非直接拆分出独立的分组表。如果目标是拆分表，上面的方法更直接高效。
3. Dataset场景扩展：如果是处理磁盘上的Arrow Dataset，可以直接用partition_by参数写入分区文件，实现物理层面的分组，但内存中拆分仍推荐上述方法。

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