如何使用 pyarrow 更改列数据类型

Question

我正在读取一组箭头文件并将它们写入镶木地板文件：

import pathlib
from pyarrow import parquet as pq
from pyarrow import feather
import pyarrow as pa

base_path = pathlib.Path('../mydata')

fields = [
    pa.field('value', pa.int64()),
    pa.field('code', pa.dictionary(pa.int32(), pa.uint64(), ordered=False)),
]
schema = pa.schema(fields)

with pq.ParquetWriter('sample.parquet', schema) as pqwriter:
    for file_path in base_path.glob('*.arrow'):
        table = feather.read_table(file_path)
        pqwriter.write_table(table)

我的问题是code箭头文件中的字段是用索引int8而不是int32. 然而范围int8还不够。因此，我定义了一个模式，其中包含parquet 文件中int32字段的索引。code

但是，将箭头表写入 parquet 现在会抱怨架构不匹配。

如何更改箭头列的数据类型？我检查了 pyarrow API，没有找到更改架构的方法。这可以在不往返熊猫的情况下完成吗？

Answer 1

Arrow ChunkedArray 有一个强制转换函数，但不幸的是它不适用于您想要执行的操作：

>>> table['code'].cast(pa.dictionary(pa.int32(), pa.uint64(), ordered=False))
Unsupported cast from dictionary<values=uint64, indices=int8, ordered=0> to dictionary<values=uint64, indices=int32, ordered=0> (no available cast function for target type)

相反，您可以将其转换为pa.uint64()并将其编码为字典：

>>> table['code'].cast(pa.uint64()).dictionary_encode().type
DictionaryType(dictionary<values=uint64, indices=int32, ordered=0>)

这是一个独立的示例：

import pyarrow as pa

source_schema = pa.schema([
    pa.field('value', pa.int64()),
    pa.field('code', pa.dictionary(pa.int8(), pa.uint64(), ordered=False)),
])

source_table = pa.Table.from_arrays([
    pa.array([1, 2, 3], pa.int64()),
    pa.array([1, 2, 1000], pa.dictionary(pa.int8(), pa.uint64(), ordered=False)),

], schema=source_schema)

destination_schema = pa.schema([
    pa.field('value', pa.int64()),
    pa.field('code', pa.dictionary(pa.int32(), pa.uint64(), ordered=False)),
])

destination_data = pa.Table.from_arrays([
    source_table['value'],
    source_table['code'].cast(pa.uint64()).dictionary_encode(),
], schema=destination_schema)