如何构造带ExtensionType的pyarrow.DictionaryArray？可用from_buffers或cast吗？

DictionaryArray完全支持搭配ExtensionType使用，你尝试的两种方案均不符合构造逻辑，可通过pa.ExtensionArray.from_storage方法实现需求，全程不需要展开底层字典数据，零额外开销。

错误原因说明 #

• 调用DictionaryArray.from_buffers触发段错误，是因为该方法构造字典数组时必须额外传入字典参数，仅传入索引缓冲区会导致空指针异常，且该方法本身不适合直接构造Extension类型的字典数组。
• 直接调用cast方法报错，是因为pyarrow默认没有实现普通字典类型到自定义ExtensionType的转换规则，未手动注册转换规则的情况下该操作不被支持。

正确实现方案 #

你只需要将已有的DictionaryArray作为存储载体，直接套入ExtensionType的包装即可：

import json
import pyarrow as pa

# 预构建的DictionaryArray
a = pa.array(["one", "two", "three", "two", "one"]).dictionary_encode()

# 自定义ExtensionType
class AnnotatedType(pa.ExtensionType):
    def __init__(self, storage_type, annotation):
        self.annotation = annotation
        super().__init__(storage_type, "my:app")
    def __arrow_ext_serialize__(self):
        return json.dumps(self.annotation).encode()
    @classmethod
    def __arrow_ext_deserialize__(cls, storage_type, serialized):
        annotation = json.loads(serialized.decode())
        return cls(storage_type, annotation)
    @property
    def num_buffers(self):
        return self.storage_type.num_buffers
    @property
    def num_fields(self):
        return self.storage_type.num_fields

# 核心构造逻辑，零拷贝无数据展开
ext_type = AnnotatedType(a.type, {"some": "data"})
ext_array = pa.ExtensionArray.from_storage(ext_type, a)

# 验证：原始字典数组完全保留
print(ext_array.storage.dictionary)
# 输出：["one", "two", "three"]

方案优势 #

• 完全不需要修改、拷贝或展开原有DictionaryArray的索引和字典数据，仅做类型包装，性能无损耗
• 生成的ext_array是标准的ExtensionArray实例，所有ExtensionType相关特性均可正常使用
• 后续可随时通过.storage属性还原回原始的DictionaryArray

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