熊猫根据数据类型选择列

我认为构建自定义转换器更容易/更清晰。此外,它可以很容易地应用于管道中

class SelectDtypeColumnsTransfomer(TransformerMixin):

    def __init__(self, dtype=object):
        self.dtype = dtype

    def transform(self, X, **transform_params):
        """ X : pandas DataFrame """

        columns = X.columns[X.dtypes == self.dtype]
        trans = X[columns].copy()
        return trans

    def fit(self, X, y=None, **fit_params):
        return self

例如:

df = pd.DataFrame({'A':[1, 2], 'B': ['s', 'd'], 'c':['test', 'r']})
print(SelectDtypeColumnsTransfomer(np.int64).transform(df))
   A
0  1
1  2
print(SelectDtypeColumnsTransfomer(object).transform(df))
   B     c
0  s  test
1  d     r

关于管道中的使用:

您应该确保训练集和测试集中的列具有相同的数据类型。根据预处理数据的方式,例如在训练集中,列的类型可能是浮点型(包括nan),而在测试集中,列的类型可能是int型(无nan),反之亦然。在这种情况下,您需要调整fit函数,该函数在拟合期间固定列,并进一步考虑在管道的以下步骤中确保数据类型一致

你可以用这样的

df_dtypes = df.dtypes.values.tolist()
select_dtype = np.dtype('int64')
select_cols = [True if x ==  select_dtype else False for x in df_dtypes]