如何通过Pyspark在Dataframe中使用具有多列的平面映射

from pyspark.sql import SparkSession
from pyspark.sql.functions import udf
from pyspark.sql.types import ArrayType, DoubleType

spark = SparkSession.builder.appName("Multi-column Mapping").getOrCreate()

data = [("John", 25, 180),
        ("Alice", 30, 165),
        ("Bob", 35, 175)]

df = spark.createDataFrame(data, ["name", "age", "height"])
df.show()

+-----+---+------+
| name|age|height|
+-----+---+------+
| John| 25|   180|
|Alice| 30|   165|
|  Bob| 35|   175|
+-----+---+------+

def map_columns(age, height):
    # 在这里实现你的映射逻辑
    return age * height

# 将自定义函数转换为UDF
map_columns_udf = udf(map_columns, DoubleType())

# 使用withColumn函数将新列添加到Dataframe中
df = df.withColumn("score", map_columns_udf(df["age"], df["height"]))
df.show()

+-----+---+------+------+
| name|age|height| score|
+-----+---+------+------+
| John| 25|   180| 4500.0|
|Alice| 30|   165| 4950.0|
|  Bob| 35|   175| 6125.0|
+-----+---+------+------+