Pyspark: Pass mehrere Spalten in UDF

Question

Ich schreibe eine benutzerdefinierte Funktion, die alle Spalten mit Ausnahme der ersten in einem Datenrahmen und Summe (oder eine andere Operation) nehmen wird. Jetzt kann der Datenrahmen manchmal 3 Spalten oder 4 Spalten oder mehr haben. Es wird variieren.

Ich weiß, ich kann 4 Spaltennamen als Pass in der UDF fest codieren, aber in diesem Fall wird es variieren, also würde ich gerne wissen, wie man es schafft?

Hier sind zwei Beispiele in der ersten haben wir zwei Spalten hinzufügen und in der zweiten haben wir drei Spalten hinzuzufügen.

Answer 1

Wenn alle Spalten, die Sie UDF die gleichen Daten müssen geben möchten geben Sie Array als Eingabeparameter, zum Beispiel verwenden können:

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

Verwenden Struktur anstelle von Array

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