python - How to define a function whose input is a column name?

Question

I want to define a function who has 2 input. A dataset and a column name. Unfortunately I can't pass the column name because I get error:

def def(df,ColName):
    
    

       F.instr('ColName', s)) for s in stop_words]),
        # If no stop words found, get the whole string
        F.length('ColName') + 1)
    ).selectExpr('trim(substring(ColName, 1, idx-1)) ColName')
    
    return Transitions

When I call the function:

Transitions= Collecting_transition(Originaltitles,originaltitle)

I get error:

name 'originaltitle' is not defined

I tried to pass string originaltitle but again because of this part ).selectExpr('trim(substring(ColName, 1, idx-1)) ColName') I get another error.

question from:https://stackoverflow.com/questions/65651701/how-to-define-a-function-whose-input-is-a-column-name

Answer 1

You have to use the ColName variable inside the function, not the string 'ColName'.

def fun(df, ColName):

    Transitions = df.withColumn('idx',
    F.coalesce(
        # Get the smallest index of a stop word in the string
        F.least(*[F.when(F.instr(ColName, s) != 0, F.instr(ColName, s)) for s in stop_words]),
        # If no stop words found, get the whole string
        F.length(ColName) + 1)
    ).selectExpr(f'trim(substring({ColName}, 1, idx-1)) {ColName}')
    
    return Transitions

Transitions = fun(Originaltitles, 'originaltitle')