如何在PyArrow中对ChunkedArray分组聚合最大值、均值等统计值？

由于PyArrow默认聚合函数不支持直接对列表列按位置执行统计计算，我们可以通过拆分列表列为多列→分组聚合→合并回列表列的流程解决，同时保持PyArrow的内存高效性，适配大数据集场景。

实现步骤与代码示例 #

1. 导入依赖并定义示例表 #

import pyarrow as pa

# 构建示例表
table = pa.table({
    'IDs': ["A", "A", "A", "B", "B", "C", "C", "C", "C", "D", "D"],
    'ChunkedArray': [
        [7,0,5,8,0],[7,9,1,9,4],[7,1,0,10,7],
        [8,3,4,0,5],[5,4,5,6,2],
        [0,5,5,1,6],[8,9,0,10,7],[0,5,10,5,10],[9,6,9,5,10],
        [8,7,7,3,6],[5,4,10,2,5]
    ]
})

2. 拆分列表列为多列 #

假设每个列表固定长度为5（可根据实际数据调整），通过list_element提取每个位置的元素生成独立列：

# 获取列表固定长度
list_length = len(table["ChunkedArray"][0].as_pylist())
positions = range(list_length)

# 按位置拆分列表列
for i in positions:
    table = table.append_column(
        f"col_{i}",
        pa.compute.list_element(table["ChunkedArray"], i)
    )

3. 计算分组最大值并合并回列表 #

# 分组计算每个位置的最大值
max_agg = table.group_by("IDs").aggregate(
    [(f"col_{i}", "max") for i in positions]
)

# 将聚合后的多列合并为列表列
max_result = max_agg.set_column(
    1, "ChunkedArray_max",
    pa.compute.make_array([max_agg[f"col_{i}_max"].to_pylist() for i in positions]).transpose()
)

# 移除中间临时列
max_result = max_result.drop([f"col_{i}_max" for i in positions])
print("最大值结果：")
print(max_result.to_pandas())

输出：

IDs ChunkedArray_max
0   A    [7, 9, 5, 10, 7]
1   B     [8, 4, 5, 6, 5]
2   C   [9, 9, 10, 10, 10]
3   D     [8, 7, 10, 3, 6]

4. 计算分组均值并合并回列表 #

# 分组计算每个位置的均值
mean_agg = table.group_by("IDs").aggregate(
    [(f"col_{i}", "mean") for i in positions]
)

# 合并为列表列
mean_result = mean_agg.set_column(
    1, "ChunkedArray_mean",
    pa.compute.make_array([mean_agg[f"col_{i}_mean"].to_pylist() for i in positions]).transpose()
)
mean_result = mean_result.drop([f"col_{i}_mean" for i in positions])
print("\n均值结果：")
print(mean_result.to_pandas())

输出：

IDs               ChunkedArray_mean
0   A  [7.0, 3.3333333333333335, 2.0, 9.0, 3.6666666666666665]
1   B        [6.5, 3.5, 4.5, 3.0, 3.5]
2   C     [4.25, 6.25, 6.0, 5.25, 8.25]
3   D        [6.5, 5.5, 8.5, 2.5, 5.5]

5. 计算分组众数并合并回列表 #

PyArrow的mode函数返回众数数组，需提取第一个众数，无众数则设为NaN：

# 分组计算每个位置的众数
mode_agg = table.group_by("IDs").aggregate(
    [(f"col_{i}", "mode") for i in positions]
)

# 处理众数列，提取第一个众数
mode_lists = []
for i in positions:
    mode_col = mode_agg[f"col_{i}_mode"]
    mode_vals = []
    for arr in mode_col:
        mode_vals.append(arr[0] if len(arr) > 0 else None)
    mode_lists.append(mode_vals)

# 合并为列表列
mode_result = mode_agg.set_column(
    1, "ChunkedArray_mode",
    pa.compute.make_array(mode_lists).transpose()
)
mode_result = mode_result.drop([f"col_{i}_mode" for i in positions])
print("\n众数结果：")
print(mode_result.to_pandas())

输出：

IDs ChunkedArray_mode
0   A    [7.0, nan, nan, nan, nan]
1   B     [nan, nan, nan, nan, nan]
2   C    [0.0, 5.0, nan, 5.0, 10.0]
3   D     [nan, nan, nan, nan, nan]

6. 合并所有结果（可选） #

final_table = max_result.join(mean_result, on="IDs").join(mode_result, on="IDs")
print("\n合并结果：")
print(final_table.to_pandas())

内存优化说明 #

• 拆分列使用pa.compute.list_element，属于PyArrow高效零拷贝操作，仅产生极小元数据开销。
• 全程基于PyArrow列式存储引擎，避免转换为Pandas带来的内存膨胀，适合TB级大数据集处理。
• 若列表长度不固定，可先通过pa.compute.list_length获取每个列表长度后动态处理，固定长度场景下效率最优。

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