如何在pyspark中应用函数？

from pyspark import SparkContext

# 初始化 SparkContext
sc = SparkContext("local", "RDD Example")

# 创建一个 RDD
data = [1, 2, 3, 4, 5]
rdd = sc.parallelize(data)

# 定义一个函数
def square(x):
    return x * x

# 使用 map() 应用函数
squared_rdd = rdd.map(square)

# 收集结果
result = squared_rdd.collect()
print(result)  # 输出: [1, 4, 9, 16, 25]

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, lit

# 初始化 SparkSession
spark = SparkSession.builder.appName("DataFrame Example").getOrCreate()

# 创建一个 DataFrame
data = [("Alice", 29), ("Bob", 31), ("Catherine", 25)]
columns = ["Name", "Age"]
df = spark.createDataFrame(data, columns)

# 使用 withColumn() 应用函数
df_with_age_plus_one = df.withColumn("AgePlusOne", col("Age") + lit(1))

# 使用 filter() 应用函数
filtered_df = df.filter(col("Age") > 25)

# 显示结果
df_with_age_plus_one.show()
filtered_df.show()

# 示例：类型错误
def add(x, y):
    return x + y

rdd = sc.parallelize([(1, "2"), (3, 4)])
result = rdd.map(add).collect()  # 会报错

# 解决方法：确保数据类型一致
rdd = sc.parallelize([(1, 2), (3, 4)])
result = rdd.map(add).collect()  # 正常运行

# 示例：性能问题
def complex_function(x):
    # 复杂的计算逻辑
    return x * x + x * x * x

rdd = sc.parallelize(range(1000000))
result = rdd.map(complex_function).collect()  # 可能会很慢

# 解决方法：优化函数逻辑
def optimized_function(x):
    return x * (x + x * x)

result = rdd.map(optimized_function).collect()  # 提高性能