将列添加到Pyspark，该列将多个组分配给相应的行

from pyspark.sql import SparkSession
from pyspark.sql.functions import when

# 创建SparkSession
spark = SparkSession.builder.getOrCreate()

# 创建示例DataFrame
data = [("Alice", 25), ("Bob", 30), ("Charlie", 35)]
df = spark.createDataFrame(data, ["Name", "Age"])

# 定义多个组和对应的值
groups = {
    "A": "Group A",
    "B": "Group B",
    "C": "Group C"
}

# 使用withColumn方法添加新列，并使用when和otherwise函数进行条件判断和赋值
df_with_group = df.withColumn("Group", when(df.Age < 30, groups["A"])
                                .when(df.Age < 35, groups["B"])
                                .otherwise(groups["C"]))

# 显示结果
df_with_group.show()

+-------+---+-------+
|   Name|Age|  Group|
+-------+---+-------+
|  Alice| 25|Group A|
|    Bob| 30|Group B|
|Charlie| 35|Group C|
+-------+---+-------+