如何设置 pyarrow 表列的“类别”数据类型？

Question

我知道在 parquet 文件中使用category.pandas 编写 pandas 时可以保留类型。DataFrameto_parquet

一开始，就我而言，我已经有一个 pyarrow Table。我可以将其一列设置为具有该category类型吗？如果是，怎么办？（我无法在 Google 和 pyarrow 文档中找到提示）

谢谢你的帮助！最好的，

Answer 1

在pyarrow中，分类类型被称为“字典类型”。可以使用以下方法将 pyarrow 数组转换为此类类型dictionary_encode()：

>>> import pyarrow as pa
>>> table = pa.table({'a': ['A', 'B', 'A']})
>>> table.schema
a: string

>>> table.column('a')
<pyarrow.lib.ChunkedArray object at 0x7f1f94fb9938>
[
  [
    "A",
    "B",
    "A"
  ]
]

>>> table.column('a').dictionary_encode()
<pyarrow.lib.ChunkedArray object at 0x7f1f94fb9b48>
[

  -- dictionary:
    [
      "A",
      "B"
    ]
  -- indices:
    [
      0,
      1,
      0
    ]
]

然后用这个新编码的列更改表有点复杂，但可以通过以下方式完成：

>>> table2 = table.set_column(0, "a", table.column('a').dictionary_encode())
>>> table2.schema
a: dictionary<values=string, indices=int32, ordered=0>