[Spark] Spark에서 사용자 정의 함수를 적용 & When 사용

👉  Spark에서 사용자 정의 함수를 작성하고 적용하기
      (파이썬에서의 함수 정의 및 적용과 비슷한 듯 다른 부분이 있음)

👉  SQL의 case when 과 유사한 When 을 사용해서 함수와 유사한 결과를 만들기

👉 SQL의 case when 문을 그대로 적용해보기

---

1. 파이썬 함수 작성하기

• 먼저 일반 python용 함수를 작성한다. 
• 이때, 반드시 입력 값과 반환 값을 설정 & 값이 None일 경우 처리 방식도 설정

# 일반 python용 UDF를 작성
def get_category(age):
    cat = ''
    
    # age 값이 None일 경우는 NA를 Return
    if age == None:
        return 'NA'
    
    if age <= 5: cat = 'Baby'
    elif age <= 12: cat = 'Child'
    elif age <= 18: cat = 'Teenager'
    elif age <= 25: cat = 'Student'
    elif age <= 35: cat = 'Young Adult'
    elif age <= 60: cat = 'Adult'
    else : cat = 'Elderly'
    
    return cat

 

 

2. python용 함수를 Pyspark용 함수로 변환

함수 객체 = udf(lambda x : 함수명(x), 해당 함수의 반환형)
               = udf(lambda 입력변수: 일반 UDF, 해당 일반 UDF의 반환형)

from pyspark.sql.functions import udf,col
from pyspark.sql.types import StringType

udf_get_category = udf(lambda x:get_category(x), StringType() )

 

3. pySpark용 함수 적용

2에서 만든 함수 객체를 적용하는데, 이때 객체명(함수를 적용할 컬럼명) 의 방식으로 적용

titanic_sdf_filled_01 = titanic_sdf_filled.withColumn("Age_Category",udf_get_category(col("Age")))
titanic_sdf_filled_01.show()

---

👉 SQL의 case when 과 유사한 When 을 사용해보기

위에서 적용한 case 별 반환 값이 다른 함수를 when 으로 구현할 수 있다.

- 유의할 점 1 : when 으로 케이스를 나눌 때 .when으로 결과객체에 대해서 조건을 써줘야 함

- 유의할 점 2 : 구조는 sql과 유사하지만, 마지막에 else가 아닌 otherwise로 마무리한다.

from pyspark.sql.functions import when
                   
titanic_sdf_filled_02 = titanic_sdf_filled.withColumn('Age_category', when(F.col('Age') <= 5, 'Baby')
                                                                      .when(F.col('Age') <= 12, 'Child')
                                                                      .when(F.col('Age') <= 18, 'Teenage')
                                                                      .when(F.col('Age') <= 25, 'Student')
                                                                      .when(F.col('Age') <= 35, 'Young Adult')
                                                                      .when(F.col('Age') <= 60, 'Adult')
                                                                      .when(F.col('Age').isNull(), 'NA')
                                                                      .otherwise('Elderly'))

titanic_sdf_filled_02.limit(10).show()

 

👉 SQL의 case when 문을 그대로 적용해보기

아래와 같이 expr을 호출하여, SQL 문을 큰 따옴표 안에 넣어서, case when 을 적용해줄 수도 있음!

from pyspark.sql.functions import expr, col

titanic_sdf_filled_03 = titanic_sdf.withColumn('Age_category', expr("CASE WHEN age = 12 THEN 'Child' " + 
                                               " WHEN Age <= 18 THEN 'Teenage' " +
                                               " WHEN Age <= 25 THEN 'Student' " +
                                               " WHEN Age <= 35 THEN 'Young Adult' " + 
                                               " WHEN Age <= 60 THEN 'Adult' " + 
                                               " WHEN Age IS NULL THEN 'NA' " +
                                               " ELSE 'Elderly' "))
titanic_sdf_filled_03.limit(10).show()