Pyspark : UDF에서 여러 열 전달

Question

데이터 프레임의 첫 번째 열을 제외한 모든 열을 가져 와서 합계 (또는 다른 연산)하는 사용자 정의 함수를 작성했습니다. 이제 데이터 프레임에는 3 열 또는 4 열 이상이있을 수 있습니다. 그것은 다양 할 것이다.

UDF에서 4 개의 열 이름을 하드 코딩 할 수 있지만이 경우에는 다를 수 있으므로 어떻게 처리해야하는지 알고 싶습니다.

첫 번째에는 두 개의 열을 추가하고 두 번째 열에는 추가 할 세 개의 열이 있습니다. 당신이 UDF에 전달하려는 모든 열이 같은 데이터가있는 경우

Answer 1

예를 들어, 당신은 입력 매개 변수로 배열을 사용할 수 있습니다 입력 :

>>> from pyspark.sql.types import IntegerType 
>>> from pyspark.sql.functions import udf, array 
>>> sum_cols = udf(lambda arr: sum(arr), IntegerType()) 
>>> spark.createDataFrame([(101, 1, 16)], ['ID', 'A', 'B']) \ 
...  .withColumn('Result', sum_cols(array('A', 'B'))).show() 
+---+---+---+------+ 
| ID| A| B|Result| 
+---+---+---+------+ 
|101| 1| 16| 17| 
+---+---+---+------+ 

>>> spark.createDataFrame([(101, 1, 16, 8)], ['ID', 'A', 'B', 'C'])\ 
...  .withColumn('Result', sum_cols(array('A', 'B', 'C'))).show() 
+---+---+---+---+------+ 
| ID| A| B| C|Result| 
+---+---+---+---+------+ 
|101| 1| 16| 8| 25| 
+---+---+---+---+------+ 

Answer 2

사용 구조체 대신 배열의

from pyspark.sql.types import IntegerType 
from pyspark.sql.functions import udf, struct 
sum_cols = udf(lambda x: x[0]+x[1], IntegerType()) 
a=spark.createDataFrame([(101, 1, 16)], ['ID', 'A', 'B']) 
a.show() 
a.withColumn('Result', sum_cols(struct('A', 'B'))).show()