Pyspark 如何修改多个案例的列数据。（适用于Pypark）

假设列“company_num”在数据帧中具有5,5,5,1,1,2,2。 我想为每个公司添加不同的编号

如果谁有公司编号5,2，那么+2

如果谁有公司编号1，那么+7

我预计最终公司数量将更改为7,7,7,8,8,4,4

为了解决这个问题，我尝试如下：

 new = (
    df
    .where("company_num in (5,2)")
    .withColumn('company_num', (col('company_num')+2).cast('integer'))
    .where("company_num in (1)")
    .withColumn('company_num', (col('company_num')+7).cast('integer'))
    )
    

---

但是，这不能正常工作

使用

when

，而不是

where

，来执行条件更新：

import pyspark.sql.functions as F

df2 = df.withColumn(
    'company_num', 
    F.when(F.col('company_num').isin(5,2), F.col('company_num')+2)
     .when(F.col('company_num') == 1, F.col('company_num')+7)
     .otherwise(F.col('company_num'))
)

df2.show()
+-----------+
|company_num|
+-----------+
|          7|
|          7|
|          7|
|          8|
|          8|
|          4|
|          4|
+-----------+

---

这段代码可能很有用

>>> from pyspark.sql.types import IntegerType
>>> from pyspark.sql.functions import when
>>> data = [5,5,5,1,1,2,2]
>>> input_df = spark.createDataFrame(data, IntegerType())
>>> input_df.show()
+-----+
|value|
+-----+
|    5|
|    5|
|    5|
|    1|
|    1|
|    2|
|    2|
+-----+

>>> input_df.withColumn("new_value", when(input_df.value.isin(5, 2), input_df.value+2).otherwise(when(input_df.value == 1, input_df.value+7))).show()
+-----+---------+
|value|new_value|
+-----+---------+
|    5|        7|
|    5|        7|
|    5|        7|
|    1|        8|
|    1|        8|
|    2|        4|
|    2|        4|
+-----+---------+

---

您是否尝试过使用udf？能否提供一个输入和输出的清晰示例？