apache spark - structtype - Pyspark: pasa múltiples columnas en UDF

spark sql tutorial

Question

Estoy escribiendo una Función definida por el usuario que tomará todas las columnas, excepto la primera en un marco de datos y la suma (o cualquier otra operación). Ahora el marco de datos a veces puede tener 3 columnas o 4 columnas o más. Va a variar

Sé que puedo codificar 4 nombres de columna como pass en el UDF, pero en este caso variarán, así que me gustaría saber cómo hacerlo.

Aquí hay dos ejemplos en el primero, tenemos dos columnas para agregar y en el segundo, tenemos tres columnas para agregar.

Answer 1

Si todas las columnas que desea pasar a UDF tienen el mismo tipo de datos, puede usar una matriz como parámetro de entrada, por ejemplo:

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