如何将 Pandas 数据框转换为模式中具有联合类型的 PyArrow 表？

Question

我有一个 Pandas 数据框，其中有一列包含 dict/structs 列表。其中一个键（在下面的示例中为thing）可以具有一个 int 或 string 的值。有没有办法定义一个 PyArrow 类型，允许将此数据帧转换为 PyArrow 表，以便最终输出到 Parquet 文件？

我尝试为此使用pa.union，但我似乎在做一些不受支持/未实施的事情。

import pandas as pd
import pyarrow as pa


df = pd.DataFrame(data={"id": [1, 2], "dict": [{"thing": 1}, {"thing": "two"}]})

schema = pa.schema([
    pa.field("id", pa.int64()),
    pa.field("dict", pa.struct([
        ("thing", pa.union([
            pa.field("int64", pa.int64()),
            pa.field("string", pa.string()),
        ], "sparse"))
    ]))
])

t = pa.Table.from_pandas(df, schema=schema)

结果：

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "pyarrow/table.pxi", line 1394, in pyarrow.lib.Table.from_pandas
  File "/usr/local/lib/python3.8/site-packages/pyarrow/pandas_compat.py", line 587, in dataframe_to_arrays
    arrays = [convert_column(c, f)
  File "/usr/local/lib/python3.8/site-packages/pyarrow/pandas_compat.py", line 587, in <listcomp>
    arrays = [convert_column(c, f)
  File "/usr/local/lib/python3.8/site-packages/pyarrow/pandas_compat.py", line 574, in convert_column
    raise e
  File "/usr/local/lib/python3.8/site-packages/pyarrow/pandas_compat.py", line 568, in convert_column
    result = pa.array(col, type=type_, from_pandas=True, safe=safe)
  File "pyarrow/array.pxi", line 292, in pyarrow.lib.array
  File "pyarrow/array.pxi", line 83, in pyarrow.lib._ndarray_to_array
  File "pyarrow/error.pxi", line 105, in pyarrow.lib.check_status
pyarrow.lib.ArrowNotImplementedError: ('sparse_union', 'Conversion failed for column dict with type object')

pa.union 的帮助文本没有给出如何使用它的示例。

>>> help(pa.union)
Help on built-in function union in module pyarrow.lib:

union(...)
    union(children_fields, mode, type_codes=None)

    Create UnionType from children fields.

    A union is defined by an ordered sequence of types; each slot in the union
    can have a value chosen from these types.

    Parameters
    ----------
    fields : sequence of Field values
        Each field must have a UTF8-encoded name, and these field names are
        part of the type metadata.
    mode : str
        Either 'dense' or 'sparse'.
    type_codes : list of integers, default None

    Returns
    -------
    type : DataType

Answer 1

在 pyarrow 2.0.0 中好像还没有实现：

import pandas as pd
import pyarrow as pa

union  = pa.union([
            pa.field("int64", pa.int64()),
            pa.field("string", pa.string()),
        ], 'sparse')

pa.array([1, 'two'], union)

---------------------------------------------------------------------------
ArrowNotImplementedError                  Traceback (most recent call last)
<ipython-input-72-f7ec6792b124> in <module>
     10         ], 'sparse')
     11 
---> 12 pa.array([1, 'two'], union)

/nix/store/aagq4nyc9m4ikjda1mykgv125v792zk7-python3-3.7.7-env/lib/python3.7/site-packages/pyarrow/array.pxi in pyarrow.lib.array()

/nix/store/aagq4nyc9m4ikjda1mykgv125v792zk7-python3-3.7.7-env/lib/python3.7/site-packages/pyarrow/array.pxi in pyarrow.lib._sequence_to_array()

/nix/store/aagq4nyc9m4ikjda1mykgv125v792zk7-python3-3.7.7-env/lib/python3.7/site-packages/pyarrow/error.pxi in pyarrow.lib.pyarrow_internal_check_status()

/nix/store/aagq4nyc9m4ikjda1mykgv125v792zk7-python3-3.7.7-env/lib/python3.7/site-packages/pyarrow/error.pxi in pyarrow.lib.check_status()

ArrowNotImplementedError: sparse_union

Answer 2

PyArrow 有一个内置方法 .from_pandas()

https://arrow.apache.org/docs/python/generated/pyarrow.Table.html#pyarrow.Table.from_pandas

import pandas as pd
import pyarrow as pa
df = pd.DataFrame({
    ...     'int': [1, 2],
    ...     'str': ['a', 'b']
    ... })
pa.Table.from_pandas(df)
<pyarrow.lib.Table object at 0x7f05d1fb1b40>