обработка данных
March 25, 2023

Скрытые ошибки выборки колонок по типам, которые сделают отладку кода мучительной

В большинстве случаев аналитики данных используют способ выборки колонок по типам через атрибут dtypes Pandas датафрейма. В то же время данный способ имеет подводные камни:

import pandas as pd
import numpy as np

df = pd.DataFrame([[1, 1, 1.2, 1., 'fio1', True, 'cat1',
                        pd.Period('2022-01', freq='M'), pd.Timestamp('2022-01-01')],
                  [2, 2, 2.3, 0.2, 'fio2', True, 'cat1', 
                        pd.Period('2022-02', freq='M'), pd.Timestamp('2022-02-01')],
                  [3, 3, 1, 5.3, 'fio3', False, 'cat2', 
                        pd.Period('2022-03', freq='M'), pd.Timestamp('2022-03-01')]],
                  columns=['id1','id2', 'val1', 'val2', 'fio', 'is_mail', 'cat', 'mon', 'day'])\
                  .astype({'id1':'int8', 'val2':'float16', 'cat':'category'})

df

А вот перечень типов данных:

types_ser = df.dtypes
types_ser

Результат выборки колонок по types_ser плохо предсказуем, примечательно, что сравнение осуществляется как со строками, так и типами данных:

types_ser[types_ser==np.number]
types_ser[types_ser=='int']

При этом в первом случае возвращается только одна численная колонка, а во втором только int64. Причина такого странного поведения в том, что dtypes возвращает объекты типов, а не строки с их названиями:

type(types_ser.loc['id1']), type(types_ser.loc['id2']), type(types_ser.loc['mon'])

Чтобы сделать результат работы с dtypes более предсказуемым, явно преобразуйте тип серии в строчный:

types_s = df.dtypes.astype(str)
types_s
types_s[(types_s.str.contains('int'))|(types_s.str.contains('float'))]
types_s[(types_s.str.contains('period'))|(types_s.str.contains('datetime'))]

Таким образом, ключом к селекции типов с dtypes является преобразование серии к строчному типу.