[Pyspark] UDF 함수에 parameter값 추가

데이터 전처리를 하는중, udf에 파라미터를 넘겨줘서 파라미터에 따라서 함수를 실행하는걸 원했다

다음과 같이 진행

 

파라미터값을 넘겨줄 땐, F.lit()함수를 이용해서 파라미터 값을 넘겨주면된다.

즉, F.lit('play')일 경우에는 reason 함수에서 logtype == 'play' 쪽 알고리즘이 실행되고 아니면 그 밑에 알고리즘이 실행이된다.

import pyspark.sql.functions as F

def reason(string, logtype):
    if logtype == 'play':
        try:
            return string.split(",")[0] + "," + string.split(",")[1]
        except:
            return string.split(",")[0]
    # unplay
    else:
        try:
            return string.split(",")[0]
        except:
            return string

change_reason = F.udf(lambda x, logtype: reason(x, logtype))

# udf 사용시엔
temp_dataframe\
	.withColumn('reason', change_reason(F.col('reason'), F.lit('play'))\
    .show(10,False)