How to pass all dataframe cloumns to UDF?

在本範例你會學到：

簡易UDF (User Define function) 使用方式
傳遞多參數給UDF

1.定義UDF

用一個簡單加總所有欄位的 function 作為實現範例(當然有更好的方式做加總這個動作)。
熟悉 python 的人應該已經熟悉 *columns 的用法，代表所有的參數
有興趣可以參考以下網址，講的蠻詳盡的。

https://skylinelimit.blogspot.com/2018/04/python-args-kwargs.html

# function for summing
def get_total(*columns):
    total = 0
    for col_value in columns:
        total += col_value
    return total

# Define udf type and func
udf_get_total = F.udf(get_total,IntegerType())

2.將所有column傳給UDF

沒什麼特別的用法，就是用全部的欄位作為一個 list 傳入 *[col for col in df.columns]
df.columns 為全部的欄位

df.withColumn('sum',udf_get_total(*[col for col in df.columns])).show()

3.完整的範例

from pyspark import SparkConf, SparkContext
from pyspark.sql import SparkSession
import pyspark.sql.functions as F
from pyspark.sql.types import *

spark = SparkSession.builder.appName('pass_columns').getOrCreate()

# function for summing
def get_total(*columns):
    total = 0
    for col_value in columns:
        total += col_value
    return total

# Create dataframe
df = spark.createDataFrame([(1, 2, 3, 4, 5, 6),(12, 24, 33, 44, 54, 66)], ["A", "B", "C", "D", "E", "F"])

# Define udf type and func
udf_get_total = F.udf(get_total,IntegerType())

# pass all column to UDF and create a new column for summing result.
df.withColumn('sum',udf_get_total(*[col for col in df.columns])).show()