python - Pyarrow 在使用 Pandas to_parquet() 时应用模式

Question

我有一个非常宽的数据框(20,000 列)，主要由 Pandas 中的 float64 列组成。我想将这些列转换为 float32 并写入 Parquet 格式。我这样做是因为这些文件的下游用户是内存有限的小容器。

我目前在 Pandas 中转换，但这在广泛的数据集上非常慢，然后写出 Parquet 。是否可以在写入 to_parquet 过程本身时转换类型？下面显示了一个虚拟示例。

import pandas as pd
import numpy as np
import pyarrow
df = pd.DataFrame(np.random.randn(3000, 15000)) # make dummy data set
df.columns = [str(x) for x in list(df)] # make column names string for parquet
df[list(df.loc[:, df.dtypes == float])] = df[list(df.loc[:, df.dtypes == float])].astype('float32') # cast the data
df.to_parquet("myfile.parquet") # write out the df

Answer 1

使用 pandas 1.0.x 和 pyarrow 0.15+ 可以通过 schema parameter在 to_parquet 中，如下所示，使用取自 this post 的模式定义. pyarrow 中用于 schema definition 的类型.

import pandas as pd
import pyarrow as pa

FILE_PATH = "/tmp/df.parquet"
df = pd.DataFrame({'a': [None, None]})
df.to_parquet(FILE_PATH)
pd.read_parquet(FILE_PATH).dtypes

这给出了以下类型:

a    object
dtype: object

定义了架构:

SCHEMA = pa.schema([('a', pa.int32())])
df.to_parquet(FILE_PATH, schema=SCHEMA)

pd.read_parquet(FILE_PATH).dtypes

它现在提供以下类型:

a    float64
dtype: object