如何使用 ExtensionType 制作 pyarrow.DictionaryArray？使用 from_buffers？使用演员表？

Question

最终，我的目标是创建一个pyarrow.DictionaryArray 和一个ExtensionType。

对于所有其他类型的 Arrow 数组，我可以使用 Array.from_buffers 静态方法来构造它并将 ExtensionType 作为它的第一个参数传递。但是，我看不到在 DictionaryArray 上使用 from_buffers 的方法，因为我需要传递它的 dictionary，但它需要零个子项。

给定一个 DicationaryArray a 并天真地使用它的 from_buffers（假设字典在它的 type 中，我很确定它不是）会导致段错误。

>>> import pyarrow as pa
>>> a = pa.array(["one", "two", "three", "two", "one"]).dictionary_encode()
>>> b = pa.DictionaryArray.from_buffers(a.type, len(a), a.indices.buffers())
/arrow/cpp/src/arrow/array/array_dict.cc:83:  Check failed: (data->dictionary) != (nullptr) 
/home/jpivarski/miniconda3/lib/python3.9/site-packages/pyarrow/libarrow.so.600(+0xf17fe8)[0x7f9dc4008fe8]
/home/jpivarski/miniconda3/lib/python3.9/site-packages/pyarrow/libarrow.so.600(_ZN5arrow4util8ArrowLogD1Ev+0xed)[0x7f9dc400978d]
/home/jpivarski/miniconda3/lib/python3.9/site-packages/pyarrow/libarrow.so.600(_ZN5arrow15DictionaryArrayC2ERKSt10shared_ptrINS_9ArrayDataEE+0x11f)[0x7f9dc427d00f]
/home/jpivarski/miniconda3/lib/python3.9/site-packages/pyarrow/libarrow.so.600(_ZN5arrow9MakeArrayERKSt10shared_ptrINS_9ArrayDataEE+0x2c6)[0x7f9dc41381b6]
/home/jpivarski/miniconda3/lib/python3.9/site-packages/pyarrow/lib.cpython-39-x86_64-linux-gnu.so(+0x213a77)[0x7f9dc5504a77]
python(+0x15f995)[0x55b45da91995]
python(_PyObject_MakeTpCall+0x316)[0x55b45da783d6]
python(_PyEval_EvalFrameDefault+0x52de)[0x55b45db162ce]
python(+0x138e20)[0x55b45da6ae20]
python(_PyEval_EvalCodeWithName+0x47)[0x55b45db4f977]
python(PyEval_EvalCodeEx+0x39)[0x55b45db4f9b9]
python(PyEval_EvalCode+0x1b)[0x55b45db4f9db]
python(+0x2506c9)[0x55b45db826c9]
python(+0x28b994)[0x55b45dbbd994]
python(+0x1142bf)[0x55b45da462bf]
python(PyRun_InteractiveLoopFlags+0xeb)[0x55b45da4646a]
python(+0x11487a)[0x55b45da4687a]
python(+0x114e14)[0x55b45da46e14]
python(Py_BytesMain+0x39)[0x55b45dbc4329]
/lib/x86_64-linux-gnu/libc.so.6(__libc_start_main+0xf3)[0x7f9dc6abd0b3]
python(+0x20aa51)[0x55b45db3ca51]
Aborted (core dumped)

无论如何，我认为这行不通。

为了完整起见，我还认为可能存在从每种存储类型到具有该存储类型的 ExtensionType 的强制转换规则，但不，cast 也不起作用。

>>> import json
>>> import pyarrow as pa
>>> a = pa.array(["one", "two", "three", "two", "one"]).dictionary_encode()
>>> 
>>> 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
... 
>>> b = a.cast(AnnotatedType(a.type, {"some": "data"}))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "pyarrow/array.pxi", line 825, in pyarrow.lib.Array.cast
  File "/home/jpivarski/miniconda3/lib/python3.9/site-packages/pyarrow/compute.py", line 309, in cast
    return call_function("cast", [arr], options)
  File "pyarrow/_compute.pyx", line 528, in pyarrow._compute.call_function
  File "pyarrow/_compute.pyx", line 327, in pyarrow._compute.Function.call
  File "pyarrow/error.pxi", line 143, in pyarrow.lib.pyarrow_internal_check_status
  File "pyarrow/error.pxi", line 120, in pyarrow.lib.check_status
pyarrow.lib.ArrowNotImplementedError: Unsupported cast from dictionary<values=string, indices=int32, ordered=0> to extension<my:app<AnnotatedType>> (no available cast function for target type)

DictionaryArrays 可以有 ExtensionType 吗？（在具有 ExtensionType 的数组上调用 dictionary_encode 不是一种选择。我的 DictionaryArrays 将完全构建；我希望用新类型重建它们而不扩展他们出去。）

Answer 1

DictionaryArray.from_buffers的crash好像是个bug（我打开https://issues.apache.org/jira/browse/ARROW-14495，我觉得其实可以修复工作）。

但特别是对于 DictionaryArray，还有一个替代构造函数：DictionaryArray.from_arrays，可以在这里使用。

使用您的示例 AnnotatedType 扩展类型，让我们首先为 DictionaryArray 的字典创建一个这种类型的小数组：

>>> dictionary = pa.array(["one", "two", "three"], pa.string())
>>> dictionary_ext = pa.ExtensionArray.from_storage(AnnotatedType(pa.string(), "annotation"), dictionary)

现在我们可以使用这个字典和索引一起创建 DictionaryArray：

>>> arr = pa.DictionaryArray.from_arrays(pa.array([0,1,2,0,1]), dictionary_ext)
>>> arr
<pyarrow.lib.DictionaryArray object at 0x7f7b0e630c80>

-- dictionary:
  [
    "one",
    "two",
    "three"
  ]
-- indices:
  [
    0,
    1,
    2,
    0,
    1
  ]

>>> arr.type
DictionaryType(dictionary<values=extension<my:app<AnnotatedType>>, indices=int64, ordered=0>)