Pyarrow：如何在分区镶木地板数据集中指定分区键的数据类型？

Question

我想创建一个以字符串作为分区键的分区 pyarrow 数据集：

import pandas as pd
import pyarrow as pa
from pyarrow import parquet as pq

data = {'key': ['001', '001', '002', '002'],
        'value_1': [10, 20, 100, 200],
        'value_2': ['a', 'b', 'a', 'b']}

df = pd.DataFrame(data)
tbl = pa.Table.from_pandas(df)

# write data to partitioned dataset
pq.write_to_dataset(tbl, 'partitioned_data', partition_cols=['key'])

print(pa.__version__)
print(df)
print(tbl)

4.0.0

   key  value_1 value_2
0  001       10       a
1  001       20       b
2  002      100       a
3  002      200       b

pyarrow.Table
keys: string
values_1: int64
values_2: string

分区数据集在文件系统中显示如下：

partitioned_data
├── key=001
│   └── 9acb170b99d14f1eba72af3697c71b8c.parquet
└── key=002
    └── 836f365800f0449b956eb35de67bbc8c.parquet

键和值按预期显示为不同分区的文件夹名称。现在我再次导入分区数据：

# read the partitioned data and convert to DataFrame
imported_tbl = pq.read_table('partitioned_data')
imported_df = imported_tbl.to_pandas()

print(imported_tbl)
print(imported_df)

pyarrow.Table
value_1: int64
value_2: string
key: dictionary<values=int32, indices=int32, ordered=0>


   value_1 value_2 key
0       10       a   1
1       20       b   1
2      100       a   2
3      200       b   2

导入的数据中，'key'的dtype由string变为dictionary，结果在不正确的值。特别是，尾随零丢失（“001”变为 1）。

是否有任何方法可以在导出或导入时指定键的数据类型以保留值？

Answer 1

使用这种文件结构，没有关于存储在任何地方的分区键的明确元数据（或架构信息）。所以 pq.read_table 尝试猜测类型。在您的情况下（即使尾随零）它也无法猜测它是一个字符串并且认为 key 是一个整数。

您可以使用 dataset api 来提供有关分区的一些信息：

my_partitioning = pa.dataset.partitioning(pa.schema([pa.field("key", pa.string())]), flavor='hive')
my_data_set = pa.dataset.dataset("partitioned_data", partitioning=my_partitioning)
table = my_data_set.to_table()
table.to_pandas()

|    |   value_1 | value_2   |   key |
|---:|----------:|:----------|------:|
|  0 |        10 | a         |   001 |
|  1 |        20 | b         |   001 |
|  2 |       100 | a         |   002 |
|  3 |       200 | b         |   002 |

Answer 2

如果您按照 0x26res 的答案所示指定分区，您实际上可以使用 pq.read_table：

my_partitioning = pa.dataset.partitioning(pa.schema([pa.field("key", pa.string())]), flavor='hive')

table = pq.read_table('partitioned_data', partitioning=my_partitioning)

Pyarrow：如何在分区镶木地板数据集中指定分区键的数据类型？

Pyarrow: How to specify the dtype of partition keys in partitioned parquet datasets?

python

pandas

parquet

pyarrow