基于变量和列条件创建变量-pyspark

8aqjt8rx  于 2021-07-09  发布在  Spark
关注(0)|答案(2)|浏览(270)

我试图在一个简单变量的基础上创建一个新变量 ModelType 和一个df变量 model .
目前我是这样做的

if ModelType == 'FRSG':
    df=df.withColumn(MODEL_NAME+'_veh', F.when(df["model"].isin(MDL_CD), df["ford_cd"]))
elif ModelType == 'TYSG':
    df=df.withColumn(MODEL_NAME+'_veh', F.when(df["model"].isin(MDL_CD), df["toyota_cd"]))
else:
    df=df.withColumn(MODEL_NAME+'_veh', F.when(df["model"].isin(MDL_CD), df["cm_cd"]))

我也试过这个

df=df.withColumn(MODEL_NAME+'_veh', F.when((ModelType == 'FRSG') &(df["model"].isin(MDL_CD)), df["ford_cd"]))

但是由于变量 ModelType 不是列,因此它给出错误

TypeError: condition should be a Column

有没有其他有效的方法来执行同样的操作?

6rqinv9w

6rqinv9w1#

您还可以使用保存 ModelType 像这样使用:

model_mapping = {"FRSG": "ford_cd", "TYSG": "toyota_cd"}

df = df.withColumn(
    MODEL_NAME + '_veh', 
    F.when(df["model"].isin(MDL_CD), df[model_mapping.get(ModelType, "cm_cd")])
)
whhtz7ly

whhtz7ly2#

我可能会为要在 then 零件:

if ModelType == 'FRSG':
    x = "ford_cd"
elif ModelType == 'TYSG':
    x = "toyota_cd"
else:
    x = "cm_cd"

df=df.withColumn(MODEL_NAME+'_veh', F.when(df["model"].isin(MDL_CD), df[x]))

相关问题