apache / apache/arrow

[C++][Parquet] Writing DictionaryArrays with ExtensionType to Parquet

Open
#30,080 1 comment 0 reactions 0 assignees View on GitHub
Component: Python Type: bug
Dominant language
C++
Stars
17.1k
Forks
4.3k
Avg merge
3d 13h
Merged PRs (30d)
88

Description

Thanks to some help I got from @jorisvandenbossche, I can create DictionaryArrays with ExtensionType (on just the dictionary, the dictionary array itself, or both). However, these extended-DictionaryArrays can't be written to Parquet files.

To start, let's set up my minimal reproducer ExtensionType, this time with an explicit ExtensionArray:
```python

>>> import json
>>> import numpy as np
>>> import pyarrow as pa
>>> import pyarrow.parquet
>>>
>>> class AnnotatedArray(pa.ExtensionArray):
... pass
...
>>> 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)
... def __arrow_ext_class__(self):
... return AnnotatedArray
...
>>> pa.register_extension_type(AnnotatedType(pa.null(), None))
```
A non-extended DictionaryArray could be built like this:
```python

>>> dictarray = pa.DictionaryArray.from_arrays(
... np.array([3, 2, 2, 2, 0, 1, 3], np.int32),
... pa.Array.from_buffers(
... pa.float64(),
... 4,
... [
... None,
... pa.py_buffer(np.array([0.0, 1.1, 2.2, 3.3])),
... ],
... ),
... )
>>> dictarray

-- dictionary:
[
0,
1.1,
2.2,
3.3
]
-- indices:
[
3,
2,
2,
2,
0,
1,
3
]
```
I can write it to a file and read it back, though the fact that it comes back as a non-DictionaryArray might be part of the problem. Is some decision being made about the array of indices being too short to warrant dictionary encoding?
```python

>>> pa.parquet.write_table(pa.table({"": dictarray}), "tmp.parquet")
>>> pa.parquet.read_table("tmp.parquet")
pyarrow.Table
: double
----
: [[3.3,2.2,2.2,2.2,0,1.1,3.3]]
```
Anyway, the next step is to make a DictionaryArray with ExtensionTypes. In this example, I'm making both the dictionary and the outer DictionaryArray itself be extended:
```python

>>> dictionary_type = AnnotatedType(pa.float64(), "inner annotation")
>>> dictarray_type = AnnotatedType(
... pa.dictionary(pa.int32(), dictionary_type), "outer annotation"
... )
>>> dictarray_ext = AnnotatedArray.from_storage(
... dictarray_type,
... pa.DictionaryArray.from_arrays(
... np.array([3, 2, 2, 2, 0, 1, 3], np.int32),
... pa.Array.from_buffers(
... dictionary_type,
... 4,
... [
... None,
... pa.py_buffer(np.array([0.0, 1.1, 2.2, 3.3])),
... ],
... ),
... )
... )
>>> dictarray_ext
<__main__.AnnotatedArray object at 0x7f8c71ec7ee0>

-- dictionary:
[
0,
1.1,
2.2,
3.3
]
-- indices:
[
3,
2,
2,
2,
0,
1,
3
]
```
This can't be written to a Parquet file:
```python

>>> pa.parquet.write_table(pa.table({"": dictarray_ext}), "tmp2.parquet")
Traceback (most recent call last):
File "", line 1, in
File "/home/jpivarski/miniconda3/lib/python3.9/site-packages/pyarrow/parquet.py", line 2034, in write_table
writer.write_table(table, row_group_size=row_group_size)
File "/home/jpivarski/miniconda3/lib/python3.9/site-packages/pyarrow/parquet.py", line 701, in write_table
self.writer.write_table(table, row_group_size=row_group_size)
File "pyarrow/_parquet.pyx", line 1451, in pyarrow._parquet.ParquetWriter.write_table
File "pyarrow/error.pxi", line 120, in pyarrow.lib.check_status
pyarrow.lib.ArrowNotImplementedError: Unsupported cast from dictionary>, indices=int32, ordered=0> to extension> (no available cast function for target type)
```
My first thought was maybe the data used in the dictionary must be simple (it's usually strings). So how about making the outer DictionaryArray extended, but the inner dictionary not extended? The type definitions are now inline.
```python

>>> dictarray_partial = AnnotatedArray.from_storage(
... AnnotatedType( # extended, but the content is not
... pa.dictionary(pa.int32(), pa.float64()), "only annotation"
... ),
... pa.DictionaryArray.from_arrays(
... np.array([3, 2, 2, 2, 0, 1, 3], np.int32),
... pa.Array.from_buffers(
... pa.float64(), # not extended
... 4,
... [
... None,
... pa.py_buffer(np.array([0.0, 1.1, 2.2, 3.3])),
... ],
... ),
... )
... )
>>> dictarray_partial
<__main__.AnnotatedArray object at 0x7f8c71ee5100>

-- dictionary:
[
0,
1.1,
2.2,
3.3
]
-- indices:
[
3,
2,
2,
2,
0,
1,
3
]
```
I can write this, but it comes back as a non-extended type, maybe because it's a non-DictionaryArray with the type of the original's dictionary (non-extended).
```python

>>> pa.parquet.write_table(pa.table({"": dictarray_partial}), "tmp3.parquet")
>>> pa.parquet.read_table("tmp3.parquet")
pyarrow.Table
: double
----
: [[3.3,2.2,2.2,2.2,0,1.1,3.3]]
```
Okay, since there's four possibilities here, what about making the dictionary an ExtensionType, but the outer DictionaryArray is not?
```python

>>> dictarray_other = pa.DictionaryArray.from_arrays(
... np.array([3, 2, 2, 2, 0, 1, 3], np.int32),
... pa.Array.from_buffers(
... AnnotatedType(pa.float64(), "only annotation"),
... 4,
... [
... None,
... pa.py_buffer(np.array([0.0, 1.1, 2.2, 3.3])),
... ],
... )
... )
>>> dictarray_other

-- dictionary:
[
0,
1.1,
2.2,
3.3
]
-- indices:
[
3,
2,
2,
2,
0,
1,
3
]
```
Nope, can't write this, either:
```python

>>> pa.parquet.write_table(pa.table({"": dictarray_other}), "tmp4.parquet")
Traceback (most recent call last):
File "", line 1, in
File "/home/jpivarski/miniconda3/lib/python3.9/site-packages/pyarrow/parquet.py", line 2034, in write_table
writer.write_table(table, row_group_size=row_group_size)
File "/home/jpivarski/miniconda3/lib/python3.9/site-packages/pyarrow/parquet.py", line 701, in write_table
self.writer.write_table(table, row_group_size=row_group_size)
File "pyarrow/_parquet.pyx", line 1451, in pyarrow._parquet.ParquetWriter.write_table
File "pyarrow/error.pxi", line 120, in pyarrow.lib.check_status
pyarrow.lib.ArrowNotImplementedError: Unsupported cast from dictionary>, indices=int32, ordered=0> to extension> (no available cast function for target type)
```
I'm pretty sure I aligned all the types right. Perhaps only one of these cases should be supported as the way it ought to work, but there ought to be some way to get the annotations into a Parquet file and read them back. (Other than un-dictencoding the array.)

**Reporter**: [Jim Pivarski](https://issues.apache.org/jira/browse/ARROW-14525) / @jpivarski
#### Related issues:
- [[C++][Python] Support for pandas Categoricals with Intervals](https://github.com/apache/arrow/issues/30119) (is related to)
- [[C++] Cast dictionary of extension type to extension type](https://github.com/apache/arrow/issues/31015) (depends upon)

**Note**: *This issue was originally created as [ARROW-14525](https://issues.apache.org/jira/browse/ARROW-14525). Please see the [migration documentation](https://github.com/apache/arrow/issues/14542) for further details.*

Contributor guide

Open the contributing guide

Research direction

Start by reproducing the failing pyarrow.parquet.write_table call and trace the cast path through pyarrow/parquet.py and pyarrow/_parquet.pyx. Compare the four extension/dictionary cases and the related cast issue; done means supported extension-typed dictionary arrays can be written to Parquet and read back with their annotations preserved.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, python
Domain
data-engineering
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.