Как удалить все категориальные столбцы из pandas DataFrame?

Question

Категориальные столбцы - отличный способ сэкономить оперативную память в pandas, однако бывают случаи, когда они просто замедляют работу. Особенно после того, как вы прошли этап, на котором у вас есть большой фрейм данных, и теперь вы работаете в подмножестве. Например, похоже, что это не очень хорошо работает с печатью в Jupyter или с использованием таких библиотек, как qgrid.

Я в основном хотел бы удалить все категориальные столбцы из фрейма данных, чтобы ускорить простые вещи:

Вот пример:

df = pd.DataFrame({"A": ["a", "b", "c", "a"], 
                   "B": ["a", "b", "c", "a"],
                   "C": [0,3,0,3],
                   "D": [0.2,0.2,0.3,0.3],
                   "F": [0,1,2,3]
                  }
                 )
df["B"] = df["B"].astype('category')
df["C"] = df["C"].astype('category')
df["D"] = df["D"].astype('category')

В результате некоторые столбцы становятся Категориальными (с разными типами: int, float, str).

df.dtypes
A      object
B    category
C    category
D    category
F       int64
dtype: object

В идеале что-то вроде:

df = df.remove_all_categorical_columns();

Это вернет исходные базовые c типы:

df.dtypes
A     object
B     object
C      int64
D    float64
F      int64
dtype: object

Answer 1

Вы можете восстановить исходный тип данных, используя df['column'].cat.categories.dtype. Остальное - это перебрать все столбцы, используя df['column']astype(df['column'].cat.categories.dtype).

В вашем примере будет работать следующее (и, надеюсь, generi c достаточно для других случаев):

def uncategorize(col):
    if col.dtype.name == 'category':
        try:
            return col.astype(col.cat.categories.dtype)
        except:
            # In case there is pd.NA (pandas >= 1.0), Int64 should be used instead of int64
            return col.astype(col.cat.categories.dtype.name.title())           
    else:
        return col

df = df.apply(lambda x: uncategorize(x), axis=0)

Затем вы восстанавливаете исходные типы данных.

df.dtypes
A     object
B     object
C      int64
D    float64
F      int64
dtype: object