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.

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.

| improve this answer | |
  • $\begingroup$ What do you mean by 'set context based on input d' $\endgroup$ – Windstorm1981 Jul 10 at 12:09
  • $\begingroup$ What do you mean by 'set context based on input data'? $\endgroup$ – Windstorm1981 Jul 10 at 12:11
  • $\begingroup$ Building Transformer context by setting the variables in it to use it further. $\endgroup$ – Akshay Tilekar Jul 23 at 8:27

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.