避免重复列条件

Question

让我们假设我有以下df -

students = spark.createDataFrame(
[
("amit",),
("amit",),
("itay",),
],
["student"],
)

我想根据student列中的值创建很多列。我知道我可以在这个数据帧上只有两个值。

的例子:

students = students.withColumn(
"address", f.when(f.col("student") == "amit", f.lit("berlin")).otherwise(f.lit("paris"))
).withColumn(
"studies", f.when(f.col("student") == "amit", f.lit("CS")).otherwise(f.lit("physics"))
).withColumn(
"age", f.when(f.col("student") == "amit", f.lit("25")).otherwise(f.lit("27"))
)

我可以做得更干净，而不重复所有的f.when(f.col("student") == "amit"或创建这些列在一起吗?任何建议都可以。

Answer 1

您可以创建一个包含创建列所需的所有信息的三元组列表:

values = [
("address", "berlin", "paris"),
("studies", "CS", "physics"),
("age", "25", "27")
]

然后，您可以通过迭代values来创建您的火花列:

cols = [
f.when(f.col('student') == "amit", f.lit(val1))
.otherwise(f.lit(val2)).alias(col_name) 
for (col_name, val1, val2) in values
]
students.select("*", *cols).show()

+-------+-------+-------+---+
|student|address|studies|age|
+-------+-------+-------+---+
|   amit| berlin|     CS| 25|
|   amit| berlin|     CS| 25|
|   itay|  paris|physics| 27|
+-------+-------+-------+---+