Put customized functions in Sklearn pipeline

sklearn

machine learning

python

data science

custom functions

Put customized functions in Sklearn pipeline

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Machine Learning Data Science Python

Introduction

Scikit-learn pipelines chain preprocessing steps and models into a single object. When you need custom data transformations that are not available in sklearn (like domain-specific feature engineering, text cleaning, or custom scaling), you can integrate your own functions using FunctionTransformer, custom transformer classes, or TransformerMixin. This keeps your entire workflow in one pipeline, ensuring consistent training and inference.

Method 1: FunctionTransformer (Simple Functions)

For stateless transformations that do not need to learn from the training data:

python

1from sklearn.preprocessing import FunctionTransformer
2from sklearn.pipeline import Pipeline
3from sklearn.linear_model import LogisticRegression
4import numpy as np
5
6# Custom function: log transform positive features
7def log_transform(X):
8    return np.log1p(X)
9
10# Custom function: clip outliers
11def clip_outliers(X, lower=1, upper=99):
12    low = np.percentile(X, lower, axis=0)
13    high = np.percentile(X, upper, axis=0)
14    return np.clip(X, low, high)
15
16pipeline = Pipeline([
17    ('clip', FunctionTransformer(clip_outliers)),
18    ('log', FunctionTransformer(log_transform)),
19    ('model', LogisticRegression())
20])
21
22pipeline.fit(X_train, y_train)
23predictions = pipeline.predict(X_test)

With Parameters

Pass additional arguments via kw_args:

python

1clip_transformer = FunctionTransformer(
2    clip_outliers,
3    kw_args={'lower': 5, 'upper': 95}
4)

With Inverse Transform

If your function is invertible:

python

1log_transformer = FunctionTransformer(
2    func=np.log1p,
3    inverse_func=np.expm1
4)

Method 2: Custom Transformer Class (Stateful)

When your transformation needs to learn parameters from the training data (like computing means, thresholds, or vocabularies), create a class that inherits from BaseEstimator and TransformerMixin:

python

1from sklearn.base import BaseEstimator, TransformerMixin
2import numpy as np
3
4class OutlierClipper(BaseEstimator, TransformerMixin):
5    """Clip values to percentile-based bounds learned from training data."""
6
7    def __init__(self, lower_percentile=1, upper_percentile=99):
8        self.lower_percentile = lower_percentile
9        self.upper_percentile = upper_percentile
10
11    def fit(self, X, y=None):
12        self.lower_ = np.percentile(X, self.lower_percentile, axis=0)
13        self.upper_ = np.percentile(X, self.upper_percentile, axis=0)
14        return self
15
16    def transform(self, X):
17        return np.clip(X, self.lower_, self.upper_)
18
19
20pipeline = Pipeline([
21    ('clip', OutlierClipper(lower_percentile=5, upper_percentile=95)),
22    ('model', LogisticRegression())
23])
24
25pipeline.fit(X_train, y_train)  # OutlierClipper learns bounds from X_train
26pipeline.predict(X_test)         # Uses training bounds to clip X_test

Key rules:

fit() learns parameters from training data and returns self
transform() applies the transformation using learned parameters
__init__() stores hyperparameters only — never modify them in fit()
TransformerMixin provides fit_transform() automatically

Method 3: Column-Specific Transformations with ColumnTransformer

Apply different custom functions to different columns:

python

1from sklearn.compose import ColumnTransformer
2from sklearn.preprocessing import FunctionTransformer, StandardScaler
3
4def extract_hour(X):
5    """Extract hour from a datetime column."""
6    return X.apply(lambda x: x.hour).values.reshape(-1, 1)
7
8def text_length(X):
9    """Compute text length."""
10    return X.apply(len).values.reshape(-1, 1)
11
12preprocessor = ColumnTransformer([
13    ('hour', FunctionTransformer(extract_hour), ['timestamp']),
14    ('text_len', FunctionTransformer(text_length), ['description']),
15    ('numeric', StandardScaler(), ['price', 'quantity']),
16])
17
18pipeline = Pipeline([
19    ('preprocess', preprocessor),
20    ('model', LogisticRegression())
21])

Method 4: Custom Feature Selector

A transformer that selects features based on criteria:

python

1class VarianceThresholdCustom(BaseEstimator, TransformerMixin):
2    """Select features with variance above a threshold."""
3
4    def __init__(self, threshold=0.01):
5        self.threshold = threshold
6
7    def fit(self, X, y=None):
8        self.variances_ = np.var(X, axis=0)
9        self.mask_ = self.variances_ > self.threshold
10        return self
11
12    def transform(self, X):
13        return X[:, self.mask_]
14
15    def get_feature_names_out(self, input_features=None):
16        if input_features is not None:
17            return input_features[self.mask_]
18        return np.arange(self.mask_.sum())

Complete Example: Feature Engineering Pipeline

python

1from sklearn.pipeline import Pipeline, make_pipeline
2from sklearn.preprocessing import StandardScaler, FunctionTransformer
3from sklearn.ensemble import RandomForestClassifier
4from sklearn.model_selection import cross_val_score
5import numpy as np
6import pandas as pd
7
8class DateFeatureExtractor(BaseEstimator, TransformerMixin):
9    """Extract useful features from datetime columns."""
10
11    def __init__(self, date_column='date'):
12        self.date_column = date_column
13
14    def fit(self, X, y=None):
15        return self
16
17    def transform(self, X):
18        df = X.copy()
19        dt = pd.to_datetime(df[self.date_column])
20        df['hour'] = dt.dt.hour
21        df['day_of_week'] = dt.dt.dayofweek
22        df['month'] = dt.dt.month
23        df['is_weekend'] = (dt.dt.dayofweek >= 5).astype(int)
24        df = df.drop(columns=[self.date_column])
25        return df
26
27pipeline = Pipeline([
28    ('dates', DateFeatureExtractor(date_column='timestamp')),
29    ('scale', StandardScaler()),
30    ('model', RandomForestClassifier(n_estimators=100))
31])
32
33scores = cross_val_score(pipeline, X, y, cv=5)
34print(f"Accuracy: {scores.mean():.3f} (+/- {scores.std():.3f})")

Using with GridSearchCV

Custom transformer parameters are accessible for hyperparameter tuning:

python

1from sklearn.model_selection import GridSearchCV
2
3pipeline = Pipeline([
4    ('clip', OutlierClipper()),
5    ('scale', StandardScaler()),
6    ('model', LogisticRegression())
7])
8
9param_grid = {
10    'clip__lower_percentile': [1, 5, 10],
11    'clip__upper_percentile': [90, 95, 99],
12    'model__C': [0.1, 1.0, 10.0]
13}
14
15grid_search = GridSearchCV(pipeline, param_grid, cv=5)
16grid_search.fit(X_train, y_train)
17print(f"Best params: {grid_search.best_params_}")

Common Pitfalls

Data Integrity: Always validate input to custom functions. A transformer that receives unexpected NaN values or wrong dtypes can produce silent errors.
Feature Alignment: Ensure custom transformations output consistent shapes between fit and transform. If fit sees 10 features, transform must also process 10 features.
Modifying __init__ params in fit: Sklearn's get_params()/set_params() (used by GridSearchCV) relies on __init__ parameters being stored unmodified. Use trailing underscores (like self.mean_) for learned attributes.
Pandas vs NumPy: Some sklearn steps return NumPy arrays, losing column names. If your custom transformer expects a DataFrame, use set_output(transform="pandas") (sklearn 1.2+) or convert explicitly.
Pickle serialization: Custom transformers must be picklable for joblib.dump(). Avoid lambda functions or closures in FunctionTransformer if you need to save the pipeline.

Summary

Approach	When to Use
`FunctionTransformer`	Stateless transforms (log, clip, encode)
Custom class with `TransformerMixin`	Stateful transforms (learn from training data)
`ColumnTransformer`	Different transforms per column
`make_pipeline`	Quick pipeline without naming steps

Use FunctionTransformer for simple, stateless functions
Create a class inheriting BaseEstimator + TransformerMixin for stateful transformations
Custom transformers work with GridSearchCV, cross_val_score, and Pipeline seamlessly
Always implement fit() returning self and transform() returning the transformed data

Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course

Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Put customized functions in Sklearn pipeline

ML System Design practice on Codemia

Introduction

Method 1: FunctionTransformer (Simple Functions)

With Parameters

With Inverse Transform

Method 2: Custom Transformer Class (Stateful)

Method 3: Column-Specific Transformations with ColumnTransformer

Method 4: Custom Feature Selector

Complete Example: Feature Engineering Pipeline

Using with GridSearchCV

Common Pitfalls

Summary

Related reading

Tackling System Design Interview Problems

Track what you have practised

ML System Design practice on Codemia