create variable based on variable and column condition - pyspark

Question

I'm trying to create a new variable based on a simple variable ModelType and a df variable model.

Currently I'm doing it in this way

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"]))

I have tried this as well

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

but since the variable ModelType is not a column so it gives an error

TypeError: condition should be a Column

Is there any other efficient method also to perform the same?

blackbishop · Accepted Answer · 2021-03-25 12:30:59Z

2

You can also use a dict that holds the possible mappings for ModelType and use it like this:

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")])
)

answered Mar 25, 2021 at 12:30

blackbishop

32.8k11 gold badges61 silver badges86 bronze badges

Sign up to request clarification or add additional context in comments.

Comments

mck · Accepted Answer · 2021-03-25 10:45:02Z

1

I would probably use a variable for the column to be chosen in the then part:

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]))

answered Mar 25, 2021 at 10:45

mck

42.7k13 gold badges44 silver badges62 bronze badges

Collectives™ on Stack Overflow

create variable based on variable and column condition - pyspark

2 Answers 2

Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related