Skip to content

Complete example

We're going to write a dataframe-agnostic "Standard Scaler". This class will have fit and transform methods (like scikit-learn transformers), and will work agnostically for pandas and Polars.

We'll need to write two methods:

  • fit: find the mean and standard deviation for each column from a given training set;
  • transform: scale a given dataset with the mean and standard deviations calculated during fit.

The fit method is a bit complicated, so let's start with transform. Suppose we've already calculated the mean and standard deviation of each column, and have stored them in attributes self.means and self.std_devs.

Transform method

We're going to take in a dataframe, and return a dataframe of the same type. Therefore, we use @nw.narwhalify:

import narwhals as nw

class StandardScaler:
    @nw.narwhalify
    def transform(self, df):
        return df.with_columns(
            (nw.col(col) - self._means[col]) / self._std_devs[col]
            for col in df.columns
        )

Note that all the calculations here can stay lazy if the underlying library permits it, so we don't pass in any extra keyword-arguments such as eager_only, we just use the default eager_only=False.

Fit method

Unlike the transform method, fit cannot stay lazy, as we need to compute concrete values for the means and standard deviations.

To be able to get Series out of our DataFrame, we'll pass eager_only=True to nw.from_native. This is because Polars doesn't have a concept of lazy Series, and so Narwhals doesn't either.

We can specify that in the @nw.narwhalify decorator by setting eager_only=True, and the argument will be propagated to nw.from_native.

import narwhals as nw

class StandardScaler:
    @nw.narwhalify(eager_only=True)
    def fit(self, df_any):
        self._means = {col: df[col].mean() for col in df.columns}
        self._std_devs = {col: df[col].std() for col in df.columns}

Putting it all together

Here is our dataframe-agnostic standard scaler:

import narwhals as nw

class StandardScaler:
    @nw.narwhalify(eager_only=True)
    def fit(self, df):
        self._means = {col: df[col].mean() for col in df.columns}
        self._std_devs = {col: df[col].std() for col in df.columns}

    @nw.narwhalify
    def transform(self, df):
        return df.with_columns(
            (nw.col(col) - self._means[col]) / self._std_devs[col]
            for col in df.columns
        )

Next, let's try running it. Notice how, as transform doesn't use any eager-only features, so we can pass a Polars LazyFrame to it and have it stay lazy!

import pandas as pd

df_train = pd.DataFrame({'a': [1, 2, 3], 'b': [4, 5, 7]})
df_test = pd.DataFrame({'a': [1, 2, 3], 'b': [4, 5, 7]})
scaler = StandardScaler()
scaler.fit(df_train)
print(scaler.transform(df_test))
     a         b
0 -1.0 -0.872872
1  0.0 -0.218218
2  1.0  1.091089
import polars as pl

df_train = pl.DataFrame({'a': [1, 2, 3], 'b': [4, 5, 7]})
df_test = pl.LazyFrame({'a': [1, 2, 3], 'b': [4, 5, 7]})
scaler = StandardScaler()
scaler.fit(df_train)
print(scaler.transform(df_test).collect())
shape: (3, 2)
┌──────┬───────────┐
 a     b         
 ---   ---       
 f64   f64       
╞══════╪═══════════╡
 -1.0  -0.872872 
 0.0   -0.218218 
 1.0   1.091089  
└──────┴───────────┘