在for循环中使用udf在Pyspark中创建多个列

from pyspark.sql import SparkSession
from pyspark.sql.functions import udf
from pyspark.sql.types import StringType, IntegerType

spark = SparkSession.builder.appName("UDF Example").getOrCreate()

def calculate_length(string):
    return len(string)

udf_calculate_length = udf(calculate_length, IntegerType())

def calculate_square_length(string):
    return len(string) ** 2

udf_calculate_square_length = udf(calculate_square_length, IntegerType())

data = [("John",), ("Alice",), ("Bob",)]
df = spark.createDataFrame(data, ["name"])

for column_name in ["length", "square_length"]:
    df = df.withColumn(column_name, udf_calculate_length(df["name"]))

df.show()