Pyspark: Vượt qua nhiều cột trong UDF

Question

Tôi đang viết một hàm do người dùng định nghĩa sẽ lấy tất cả các cột ngoại trừ cột đầu tiên trong một khung dữ liệu và tổng (hoặc bất kỳ thao tác nào khác). Giờ đây, khung dữ liệu đôi khi có thể có 3 cột hoặc 4 cột trở lên. Nó sẽ thay đổi.

Tôi biết tôi có thể mã cứng 4 tên cột như vượt qua trong UDF nhưng trong trường hợp này nó sẽ khác nhau vì vậy tôi muốn biết làm thế nào để hoàn thành nó?

Dưới đây là hai ví dụ trong ví dụ đầu tiên, chúng tôi có hai cột để thêm vào và cột thứ hai sẽ có thêm ba cột.

Answer 1

Nếu tất cả các cột bạn muốn vượt qua để UDF có cùng dữ liệu gõ bạn có thể sử dụng mảng như tham số đầu vào, ví dụ:

>>> 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

Sử dụng struct thay vì mảng

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()