如何为 pyarrow Table 列设置 'category' 数据类型？

Question

我知道在 parquet 文件中写入 pandas DataFrame 时可以保留 category 类型，使用 to_parquet.

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

感谢您的帮助！最佳，

Answer 1

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

>>> 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
    ]
]

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

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

如何为 pyarrow Table 列设置 'category' 数据类型？

How to set the 'category' data type for a pyarrow Table column?

python

parquet

pyarrow