pandas 保存到parquet时,Int/Float类型的分类变量丢失

jdzmm42g  于 2023-05-27  发布在  其他
关注(0)|答案(2)|浏览(101)

我在pandas中有以下 Dataframe ,它被保存为一个parquet

import pandas as pd
df = pd.DataFrame({"a":['1','2','3']}).astype("category")

在检查唯一的字段时,我们看到

>>> df['a'].cat.categories
Index(['1', '2', '3'], dtype='object')

现在保存并加载parquet文件,并检查列:

>>> df.to_parquet("1.parquet")
>>> pd.read_parquet("1.parquet")['a']
0    1
1    2
2    3
Name: a, dtype: category
Categories (3, object): ['1', '2', '3']

到目前为止一切都很好,根据pyarrowpandascategorical数据类型文档,一切都如预期的那样,其中两个框架都声称支持分类数据类型。
但是,如果我们执行以下操作(将源列类型从object更改为int),则分类变量在加载或保存期间丢失:

>>> df = pd.DataFrame({"a":[1,2,3]}).astype("category")
>>> df['a']
0    1
1    2
2    3
Name: a, dtype: category
Categories (3, int64): [1, 2, 3]
>>> df.to_parquet("1.parquet")
>>> pd.read_parquet("1.parquet")['a']
0    1
1    2
2    3
Name: a, dtype: int64

两个问题:
1.为什么基于数值的分类列在此保存/加载例程期间丢失?
1.有什么办法可以防止这种情况发生吗?
谢谢!

svdrlsy4

svdrlsy41#

this tutorial上,似乎使用fastparquet引擎将保留分类信息:

import pandas as pd
df = pd.DataFrame({"a":[1,2,3]}).astype("category")

df.to_parquet("1.parquet", engine='fastparquet')
pd.read_parquet("1.parquet", engine='fastparquet')['a']

给了我

0    1
1    2
2    3
Name: a, dtype: category
Categories (3, int64): [1, 2, 3]
sauutmhj

sauutmhj2#

如果显式地遍历pyarrow.Table对象,则会维护类型:

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

df = pd.DataFrame({"a":['1','2','3']}).astype("category")
tab = pa.Table.from_pandas(df)

path = "/tmp/foo.parquet"
pq.write_table(tab, path)
df2 = pq.read_table(path).to_pandas()
df2['a']

返回

0    1
1    2
2    3
Name: a, dtype: category
Categories (3, object): ['1', '2', '3']

您也可以使用pa.dictionary()类型,尽管pandas对它的支持是非常新的。

相关问题