python Pandas : merge loses categorical columns

Question

我正在处理分类数据的大型数据帧，我发现当我在两个数据帧上使用 pandas.merge 时，任何分类数据列都会自动向上转换为更大的数据类型。 (这会显着增加 RAM 消耗。)一个简单的例子来说明:

编辑:做了一个更合适的例子

import pandas
import numpy

df1 = pandas.DataFrame(
    {'ID': [5, 3, 6, 7, 0, 4, 8, 2, 9, 1, 6, 5, 4, 9, 7, 2, 1, 8, 3, 0], 
     'value1': pandas.Categorical(numpy.random.randint(0, 2, 20))})

df2 = pandas.DataFrame(
    {'ID': [5, 3, 6, 7, 0, 4, 8, 2, 9, 1],  
     'value2': pandas.Categorical(['c', 'a', 'c', 'a', 'c', 'b', 'b', 'a', 'a', 'b'])})

result = pandas.merge(df1, df2, on="ID")
result.dtypes


Out []:
ID         int32
value1     int64
value2    object
dtype: object

我希望 value1 和 value2 在结果 DataFrame 中保持分类。字符串标签转换为对象类型的成本可能特别高。

来自 https://github.com/pydata/pandas/issues/8938这可能是预期的？有什么办法可以避免这种情况吗？

Answer 1

我可能遗漏了您的目标，但目的是让用户在需要时转换为(或不转换)类别。我认为在这种特殊情况下，这可以自动完成。老实说，分类转换无论如何都会在最后完成，所以这实际上不会为您节省任何东西(通过在 merge 中进行)。

In [57]: result = pandas.merge(df1, df2, on="ID")

In [58]: result['value1'] = result['value1'].astype('category')

In [59]: result['value2'] = result['value2'].astype('category')

In [60]: result
Out[60]: 
    ID value1 value2
0    5      0      c
1    5      1      c
2    3      0      a
3    3      1      a
4    6      0      c
5    6      0      c
6    7      0      a
7    7      1      a
8    0      1      c
9    0      1      c
10   4      1      b
11   4      1      b
12   8      0      b
13   8      1      b
14   2      1      a
15   2      1      a
16   9      0      a
17   9      1      a
18   1      0      b
19   1      1      b

In [61]: result.dtypes
Out[61]: 
ID           int64
value1    category
value2    category
dtype: object

In [62]: result.info()
<class 'pandas.core.frame.DataFrame'>
Int64Index: 20 entries, 0 to 19
Data columns (total 3 columns):
ID        20 non-null int64
value1    20 non-null category
value2    20 non-null category
dtypes: category(2), int64(1)
memory usage: 400.0 byte