Python pandas:merge失去了分类列

Question

我正在处理分类数据的大型DataFrame,我发现当我在两个数据帧上使用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

作为解决方法，您可以将分类列转换为整数值代码，并将列到类别的映射存储在字典中。例如，

def decat(df):
    """
    Convert categorical columns to (integer) codes; return the categories in catmap
    """
    catmap = dict()
    for col, dtype in df.dtypes.iteritems():
        if com.is_categorical_dtype(dtype):
            c = df[col].cat
            catmap[col] = c.categories
            df[col] = c.codes
    return df, catmap

In [304]: df
Out[304]: 
   ID value2
0   5      c
1   3      a
2   6      c
3   7      a
4   0      c
5   4      b
6   8      b
7   2      a
8   9      a
9   1      b

In [305]: df, catmap = decat(df)

In [306]: df
Out[306]: 
   ID  value2
0   5       2
1   3       0
2   6       2
3   7       0
4   0       2
5   4       1
6   8       1
7   2       0
8   9       0
9   1       1

In [307]: catmap
Out[307]: {'value2': Index([u'a', u'b', u'c'], dtype='object')}

现在您可以照常合并，因为合并整数值列没有问题。

稍后，您可以使用以下数据重新构建分类列catmap：

def recat(df, catmap):
    """
    Use catmap to reconstitute columns in df to categorical dtype
    """
    for col, categories in catmap.iteritems():
        df[col] = pd.Categorical(categories[df[col]])
        df[col].cat.categories = categories
    return df

import numpy as np
import pandas as pd
import pandas.core.common as com

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

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

def decat(df):
    """
    Convert categorical columns to (integer) codes; return the categories in catmap
    """
    catmap = dict()
    for col, dtype in df.dtypes.iteritems():
        if com.is_categorical_dtype(dtype):
            c = df[col].cat
            catmap[col] = c.categories
            df[col] = c.codes
    return df, catmap

def recat(df, catmap):
    """
    Use catmap to reconstitute columns in df to categorical dtype
    """
    for col, categories in catmap.iteritems():
        df[col] = pd.Categorical(categories[df[col]])
        df[col].cat.categories = categories
    return df

def mergecat(left, right, *args, **kwargs):
    left, left_catmap = decat(left)
    right, right_catmap = decat(right)
    left_catmap.update(right_catmap)
    result = pd.merge(left, right, *args, **kwargs)
    return recat(result, left_catmap)

result = mergecat(df1, df2, on='ID')
result.info()

产量

<class 'pandas.core.frame.DataFrame'>
Int64Index: 20 entries, 0 to 19
Data columns (total 3 columns):
ID        20 non-null int32
value1    20 non-null category
value2    20 non-null category
dtypes: category(2), int32(1)
memory usage: 320.0 bytes