Pyspark Pipeline Custom Transformer

Question

I'm having some trouble understanding the creation of custom transformers for Pyspark pipelines.

I am writing a custom transformer that will take the dataframe column Company and remove stray commas:

from pyspark.sql.functions import *

class DFCommaDropper(Transformer):

    def__init__(self, *args, **kwargs):
        self.name = CommaDropper

    def transform(self,df):
        df = df.withColumn('Company', regexp_replace('Company',',','')
        return df

The above code is obviously wrong. I'm unsure what/how to initialize this and then how to use the initialized class instance in the transform function.

Thanks in advance for your help.

Answer 1

class StrayCommaRemover(TransformerMixin):
    def __init__(self): //Initialize self by setting some variables here which can be passed as a input to transformer


def fit(self, X, y=None):
    self.columns = X.columns //Setting context based on input Data
    return self

def transform(self, X, y=None): // Actual transformation logic
    X= X.withColumn('Company', regexp_replace('Company',',','')
    return X

You can add above transformer as a step in your pipeline and can call init() and fit() on it.