How to drop columns in ColumnTransformer?

Viewed 19

I created a custom Pipeline that adds one column 'Message Length', encodes categorical & boolean columns, and drops selected columns.

def custom_pipeline(to_drop: list = [], features_out: bool = False) -> Pipeline:
    # Add 'Message Length' attribute based on the 'Raw Message' column
    attrib_adder = AttributeAdder(attribs_in=['Raw Message'], attribs_out=['Message Length'], func=get_message_length)

    # Define the column transformer
    preprocessor = ColumnTransformer(transformers=[
        ('virus_scanned', enumerate_virus_scanned, ['X-Virus-Scanned']),
        ('priority', enumerate_priority, ['X-Priority']),
        ('encoding', enumerate_encoding, ['Encoding']),
        ('flags', enumerate_bool, ['Is HTML', 'Is JavaScript', 'Is CSS']),
        ('select', 'passthrough', ['Attachments', 'URLs', 'IPs', 'Images', 'Message Length']),
        ('drop_out', 'drop', to_drop) # --> This does not work
    ])

    # Define pipeline
    pipe = Pipeline(steps=[
        ('attrib_adder', attrib_adder),
        ('preprocessor', preprocessor),
        ('scaler', MinMaxScaler())
    ])

    # Get features out
    if features_out:
        features = [col for col in chain(*[cols for _,_,cols in preprocessor.transformers[:-1]]) if col not in to_drop]
        
        # Return pipeline and features
        return pipe, features
    
    # Return pipeline
    return pipe

Unfortunately, the last 'drop_out' transformer does not drop columns.

For example, even if I pass to_drop = ['Attachments', 'Message Length'] it still preserves them in the output.

What might be the possible solution?

0 Answers
Related