web-dev-qa-db-fra.com

Pyspark: Passer plusieurs colonnes dans UDF

J'écris une fonction définie par l'utilisateur qui prendra toutes les colonnes sauf la première dans un cadre de données et fera la somme (ou toute autre opération). Maintenant, le cadre de données peut parfois avoir 3 colonnes ou 4 colonnes ou plus. Ça va varier.

Je sais que je peux coder durement 4 noms de colonne comme transmis dans le fichier UDF, mais dans ce cas, cela variera et j'aimerais donc savoir comment le faire.

Voici deux exemples dans la première, nous avons deux colonnes à ajouter et dans la seconde, nous avons trois colonnes à ajouter.

enter image description here

apache-sparkpysparkspark-dataframe
28
sjishan

Si toutes les colonnes que vous souhaitez transmettre à UDF ont le même type de données, vous pouvez utiliser un tableau comme paramètre d'entrée, par exemple:

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

Utilisez struct au lieu de tableau

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

Une autre manière simple sans Array et Struct.

from pyspark.sql.types import IntegerType
from pyspark.sql.functions import udf, struct

def sum(x, y):
    return x + y

sum_cols = udf(sum, IntegerType())

a=spark.createDataFrame([(101, 1, 16)], ['ID', 'A', 'B'])
a.show()
a.withColumn('Result', sum_cols('A', 'B')).show()
10
neeraj bhadani

Voici comment j'ai essayé et semblé fonctionner:

colsToSum = df.columns[1:]
df_sum = df.withColumn("rowSum", sum([df[col] for col in colsToSum]))
0
Travis Hensrud