[PySpark] 문법 예제 : when

다음 문서 내용을 정리함

sparkbyexamples.com/pyspark/pyspark-when-otherwise/

 

pyspark 의 dataframe 에서,

if 문처럼 사용할 수 있는 SQL function 인 case-when 과 비슷한

when (+ otherwise ) 을 사용하는 방법을 예제를 들어 설명한다.

 

 

샘플 dataframe 을 만들어본다.

 

data = [("James","M",60000),("Michael","M",70000), ("Robert",None,400000),("Maria","F",500000), ("Jen","",None)]

columns = ["name","gender","salary"]

df = spark.createDataFrame(data = data, schema = columns)

 

 

다음과 같은 조건문이 들어가는 위치에 when 과 otherwise 를 사용할 수 있다.

예를 들어 withColumn 함수를 사용할 때, 다음과 같이 두번째 파라미터 값으로 when 을 줄 수 있다.

 

from pyspark.sql.functions import when

df2 = df.withColumn("new_gender", when(df.gender == "M","Male")

    .when(df.gender == "F","Female")

    .when(df.gender.isNull() ,"")

    .otherwise(df.gender))

 

혹은

 

df2.withColumn("new_column", when(col("code") == "a" | col("code") == "d", "A")

    .when(col("code") == "b" & col("amt") == "4", "B")

    .otherwise("A1"))

 

혹은 select 문에 쓸 수 있다.

 

df2=df.select(col("*"), when(df.gender == "M","Male")

    .when(df.gender == "F","Female")

    .when(df.gender.isNull() ,"")

    .otherwise(df.gender)

    .alias("new_gender"))

 

 

물론 SQL 을 지원하는 spark 특성상,

sql 문법을 실행하는 `expr` 을 사용하면

case-when 문도 pyspark 에서 사용할 수 있다.

 

df4 = df.select(col("*"), expr("CASE WHEN gender = 'M' THEN 'Male' " +

    "WHEN gender = 'F' THEN 'Female' WHEN gender IS NULL THEN ''"

    + "ELSE gender END")

    .alias("new_gender"))