Pyspark - DataFrame 在循环中应用函数时未更新

Question

I'm trying to apply different functions to various columns of a DataFrame depending on a condition. When I do this in a loop, fn1 is applied successfully on the first iteration. But the df turns None on the second iteration. I guess the problem is the way I'm initializing the df in the scope of a loop.

df = spark.createDataFrame([(10,4,2,3),(20,1,3,4),(30,7,4,5),(40,2,1,9)], schema=['id','metric_1','metric_2', 'metric_3'])

cols_info = [{'name':'metric_1','apply_func':'True','method':'fn1'},{'name':'metric_2','apply_func':'True','method':'fn2'}, {'name':'metric_3','apply_func':'True','method':'fn3'}]

def fn1(df, col):
    return df.withColumn(col, F.pow(df[col], 2))

def fn2(df, col):
    return df.withColumn(col, F.hash(df[col]))

def fn3(df, col):
    return df.withColumn(col, F.log2(df[col]))

def process_data(df, columns):
    for col in columns:
        if col["apply_func"] == "True":
            if column["method"] == "fn1":
                df = fn1(df, col["name"])
            if column["method"] == "fn2":
                df = fn2(df, col["name"])
            if column["method"] == "fn3":
                df = fn3(df, col["name"])

    return df

What is the correct way to apply such transformations with Pyspark DataFrame API?

Answer 1

Can you try to write the functions in this way. This way worked for me:

def fn1(df, col):
    df = df.withColumn(col, F.pow(df[col], 2))
    return df


def fn2(df, col):
    df = df.withColumn(col, F.hash(df[col]))
    return df

def fn3(df, col):
    df = df.withColumn(col, F.log2(df[col]))
    return df

def process_data(df, columns):
    for col in columns:
        if col["apply_func"] == "True":
            if col["method"] == "fn1":
                df = fn1(df, col["name"])
            if col["method"] == "fn2":
                df = fn2(df, col["name"])
            if col["method"] == "fn3":
                df = fn3(df, col["name"])
    return df

I think the assignment is necessary but not very sure. Someone could improve on my answer