Expanding Array Columns in Pandas DataFrames and wrap up in scikit-learn transformer

Sometimes, you may find yourself working with a Pandas DataFrame that contains a column of arrays with the same length. In some cases, it may be more useful to “explode” this column of arrays into multiple columns, with one column for each value in the arrays. This can make it easier to perform analysis or modeling on the data.

In this blog post, we’ll explore how to use Pandas to expand an array column into multiple columns, and how to encapsulate this functionality into a scikit-learn transformer for use in machine learning pipelines.

Expanding an Array Column with Pandas

To illustrate how to expand an array column in a Pandas DataFrame, let’s start with an example DataFrame that contains a column of arrays:

import pandas as pd

df = pd.DataFrame({'array_col': [[1, 2, 3], [4, 5, 6], [7, 8, 9]]})

This DataFrame has a single column named array_col with three rows, each containing an array of three integers. To expand this column into multiple columns, we can use the apply method to apply a function to each row of the DataFrame. This function will return a new DataFrame with the values of the array column, which will be automatically assigned to new columns in the resulting DataFrame.

Here’s an example of how to do this:

def explode_array_column(row):
    return pd.Series(row['array_col'])

expanded_cols = df.apply(explode_array_column, axis=1)
expanded_cols.columns = ['col_{}'.format(i) for i in range(expanded_cols.shape[1])]

df = pd.concat([df, expanded_cols], axis=1)
df = df.drop('array_col', axis=1)

print(df)

This will output a new DataFrame with three columns (col_0, col_1, and col_2) that contain the values from the original array column:

   col_0  col_1  col_2
0      1      2      3
1      4      5      6
2      7      8      9

Creating a Custom Transformer with scikit-learn

While the above approach works well for a single DataFrame, it can be cumbersome to repeat the same steps for multiple DataFrames. One way to simplify this process is to encapsulate the functionality into a custom transformer that can be used in scikit-learn pipelines.

Here’s an example implementation of a custom transformer that expands an array column into multiple columns:

import pandas as pd
from sklearn.base import BaseEstimator, TransformerMixin

class ArrayExpander(BaseEstimator, TransformerMixin):
    def __init__(self):
        pass
        
    def fit(self, X, y=None):
        return self
    
    def transform(self, X):
        def explode_array_column(row):
            return pd.Series(row['array_col'])
        
        df = X.copy()
        expanded_cols = df.apply(explode_array_column, axis=1)
        expanded_cols.columns = ['embed_col_{}'.format(i) for i in range(expanded_cols.shape[1])]
        df = pd.concat([df, expanded_cols], axis=1)
        df = df.drop('array_col', axis=1)
        return df

      
array_expander = ArrayExpander()
df_transformed = array_expander.fit_transform(df)