Python pandas: merge loses categorical columns

Question

I am working with large DataFrames of categorical data and I found that when I use a pandas.merge on two dataframes any columns of categorical data are automatically upcast to a larger datatype. (This can dramatically increase RAM consumption.) A simple example to illustrate:

EDIT: made a more appropriate example

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

I would like value1 and value2 to remain categorical in the result DataFrame. The conversion to object type of string labels can be particularly costly.

From https://github.com/pydata/pandas/issues/8938 this may be as intended? Is there anyway to avoid this?

Answer 1

As a workaround, you could convert the categorical columns to integer-valued codes, and store the mapping of columns to categories in a dict. For example,

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')}

Now you can merge as usual since there is no trouble merging integer valued columns.

Later, you can re-constitute the categorical columns using the data in 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()

yields

<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

Answer 2

You could split the column categories into indices (pandas.Series.cat.categories) and codes (pandas.Series.cat.codes), merge the dataframes and then recreate the categorical Series using from_codes function. It's ugly but it seems to be fast and memory-efficient.

# categorical indices
indices = [x.cat.categories for x in [df1.value1, df2.value2]]
# in-place setting columns with their categorical codes
for df, col in zip([df1, df2], ['value1', 'value2']):
    df[col] = df[col].cat.codes
# merging updated dataframes
result = pandas.merge(df1, df2, on='ID')
# recreating categorical series
for col, index in zip(['value1', 'value2'], indices):
    result[col] = pandas.Categorical.from_codes(result[col], index)