将UDF应用于Spark DF中的列，并根据列的不同而改变函数

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()

data = [("Alice", 25), ("Bob", 30), ("Charlie", 35)]
df = spark.createDataFrame(data, ["Name", "Age"])
df.show()

+-------+---+
|   Name|Age|
+-------+---+
|  Alice| 25|
|    Bob| 30|
|Charlie| 35|
+-------+---+

def uppercase(name):
    return name.upper()

def double_age(age):
    return age * 2

uppercase_udf = udf(uppercase, StringType())
double_age_udf = udf(double_age, IntegerType())

df = df.withColumn("UpperName", uppercase_udf(df["Name"]))
df = df.withColumn("DoubleAge", double_age_udf(df["Age"]))
df.show()

+-------+---+---------+---------+
|   Name|Age|UpperName|DoubleAge|
+-------+---+---------+---------+
|  Alice| 25|    ALICE|       50|
|    Bob| 30|      BOB|       60|
|Charlie| 35| CHARLIE|       70|
+-------+---+---------+---------+