如何获取PyArrow ChunkedArray的所有类别？求更优实现方案

问题场景 #

我有一个PyArrow ChunkedArray（字典类型），想要提取它包含的所有类别。目前的实现是遍历每个子数组，收集各自的字典类别后求并集，代码如下：

import pyarrow as pa
import functools

# 构造测试数据
one = pa.array(["a", "b", None, "d"], pa.dictionary(pa.int64(), pa.utf8()))
two = pa.array(["a", "b", None, "d", 'e'], pa.dictionary(pa.int64(), pa.utf8()))
ca = pa.chunked_array([one, two])

# 当前实现：遍历chunk求类别并集
cats = functools.reduce(lambda x, y: x | y, (set(x.dictionary) for x in ca.chunks))
out = pa.chunked_array([list(cats)])

这个方法可行但步骤繁琐，有没有更高效简洁的实现方式？

更优实现方案 #

方案1：合并字典数组后取唯一值 #

直接收集所有chunk的字典数组，合并后提取唯一值，这种方式避免了处理整个大数组，效率更高：

import pyarrow as pa

# 构造测试数据（同上）
one = pa.array(["a", "b", None, "d"], pa.dictionary(pa.int64(), pa.utf8()))
two = pa.array(["a", "b", None, "d", 'e'], pa.dictionary(pa.int64(), pa.utf8()))
ca = pa.chunked_array([one, two])

# 合并所有chunk的字典数组
all_dicts = pa.concat_arrays([chunk.dictionary for chunk in ca.chunks])
# 提取唯一类别
unique_cats = pa.unique(all_dicts)
# 转为ChunkedArray（如果需要）
out = pa.chunked_array([unique_cats])

方案2：直接调用ChunkedArray的unique()方法 #

PyArrow的ChunkedArray原生支持unique()方法，内部会高效遍历所有chunk并提取唯一值，代码最简洁：

import pyarrow as pa

# 构造测试数据（同上）
one = pa.array(["a", "b", None, "d"], pa.dictionary(pa.int64(), pa.utf8()))
two = pa.array(["a", "b", None, "d", 'e'], pa.dictionary(pa.int64(), pa.utf8()))
ca = pa.chunked_array([one, two])

# 提取唯一类别，drop_nulls=True可排除缺失值None
unique_cats = ca.unique(drop_nulls=True)
# 转为ChunkedArray（如果需要）
out = pa.chunked_array([unique_cats])

方案优势 #

• 方案1更适合数据量较大的场景：仅操作体积更小的字典数组，避免扁平化整个ChunkedArray，性能更优。
• 方案2代码最简洁：利用PyArrow内置方法，无需手动处理遍历和集合合并，可读性更强。

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