12

What is the equivalent of this operation in Pyspark?

import pandas as pd
import numpy as np

df = pd.DataFrame({'Type':list('ABBC'), 'Set':list('ZZXY')})
df['color'] = np.where(df['Set']=='Z', 'green', 'red')
print(df)

output

   Set Type  color
0   Z    A  green
1   Z    B  green
2   X    B    red
3   Y    C    red
adil blanco
  • 335
  • 1
  • 4
  • 14

1 Answers1

22

You're looking for pyspark.sql.functions.when():

from pyspark.sql.functions import when, col

df = df.withColumn('color', when(col('Set') == 'Z', 'green').otherwise('red'))
df.show()
#+---+----+-----+
#|Set|Type|color|
#+---+----+-----+
#|  Z|   A|green|
#|  Z|   B|green|
#|  X|   B|  red|
#|  Y|   C|  red|
#+---+----+-----+

If you have multiple conditions to check, you can chain together calls to when() as shown in this answer.

pault
  • 41,343
  • 15
  • 107
  • 149