从 PyArrow 写入 Parquet 文件时如何指定逻辑类型?

How to specify logical types when writing Parquet files from PyArrow?

我在 Python 中使用 PyArrow to write Parquet files from some Pandas 数据帧。

有没有一种方法可以指定写入 parquet 文件的逻辑类型?

例如,在 PyArrow 中写入一个 np.uint32 列会在 parquet 文件中产生一个 INT64 列,而使用 fastparquet 模块写入相同的列会产生一个具有逻辑类型的 INT32 列UINT_32(这是我从 PyArrow 获得的行为)。

例如:

import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq
import fastparquet as fp
import numpy as np

df = pd.DataFrame.from_records(data=[(1, 'foo'), (2, 'bar')], columns=['id', 'name'])
df['id'] = df['id'].astype(np.uint32)

# write parquet file using PyArrow
pq.write_table(pa.Table.from_pandas(df, preserve_index=False), 'pyarrow.parquet')

# write parquet file using fastparquet
fp.write('fastparquet.parquet', df)

# print schemas of both written files
print('PyArrow:', pq.ParquetFile('pyarrow.parquet').schema)
print('fastparquet:', pq.ParquetFile('fastparquet.parquet').schema)

这个输出:

PyArrow: <pyarrow._parquet.ParquetSchema object at 0x10ecf9048>
id: INT64
name: BYTE_ARRAY UTF8

fastparquet: <pyarrow._parquet.ParquetSchema object at 0x10f322848>
id: INT32 UINT_32
name: BYTE_ARRAY UTF8

我在使用其他列类型时遇到了类似的问题,因此我真的在寻找一种通用方法来指定在使用 PyArrow 进行写入时使用的逻辑类型。

PyArrow默认写入parquet 1.0版本文件,需要2.0版本才能使用UINT_32逻辑类型

解决方法是写table时指定版本,即

pq.write_table(pa.Table.from_pandas(df, preserve_index=False), 'pyarrow.parquet', version='2.0')

这将导致写入预期的镶木地板架构。