Pyarrow 在使用 pandas to_parquet() 时应用模式答案

【问题标题】：Pyarrow apply schema when using pandas to_parquet()Pyarrow 在使用 pandas to_parquet() 时应用模式
【发布时间】：2019-03-21 21:13:30
【问题描述】：

我有一个非常宽的数据框（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

【问题讨论】：

标签： python pandas pyarrow

【解决方案1】：

尝试使用箭头而不是 pandas 进行向下转换：

def convert_arrow(df):
    table = pa.Table.from_pandas(df)
    columns = [
        c.cast(pa.float32()) if c.type == pa.float64() else c
        for c in table
    ]
    return pa.Table.from_arrays(columns, table.column_names)

我做了一个简单的基准测试，它快了 20 倍。

我认为您的代码的问题是在现有 DataFrmae 中一一分配列，这效率不高。这篇博文解释得很好：https://uwekorn.com/2020/05/24/the-one-pandas-internal.html

另一个不涉及箭头的简单解决方案是转换每一列并在最后创建数据框。下面的代码比箭头版本稍慢：

def convert_pandas_by_columns(df):
    columns = [
        df[c].astype('float32') if df[c].dtype == float else df[c]
        for c in df.columns
    ]
    return pd.DataFrame(columns)

【讨论】：

【解决方案2】：

使用 pandas 1.0.x 和 pyarrow 0.15+，可以使用取自 this post 的架构定义将 schema parameter 传递到 to_parquet 中，如下所示。用于schema definition 的pyarrow 类型。

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

【讨论】：