web-dev-qa-db-fra.com

Ajouter la somme des colonnes en tant que nouvelle colonne dans la base de données PySpark

J'utilise PySpark et j'ai une base de données Spark avec un tas de colonnes numériques. Je veux ajouter une colonne qui est la somme de toutes les autres colonnes.

Supposons que mon cadre de données comporte des colonnes "a", "b" et "c". Je sais que je peux le faire:

df.withColumn('total_col', df.a + df.b + df.c)

Le problème est que je ne veux pas taper chaque colonne individuellement et les ajouter, surtout si j'ai beaucoup de colonnes. Je veux pouvoir le faire automatiquement ou en spécifiant une liste de noms de colonnes que je veux ajouter. Y a-t-il une autre façon de faire cela?

pythonapache-sparkpysparkspark-dataframe
17
plam

Ce n'était pas évident. Je ne vois aucune somme basée sur les lignes des colonnes définies dans l'API spark Dataframes.

Version 2

Cela peut être fait de manière assez simple:

newdf = df.withColumn('total', sum(df[col] for col in df.columns))

df.columns est fourni par pyspark sous forme de liste de chaînes donnant tous les noms de colonnes du cadre de données Spark. Pour une somme différente, vous pouvez également fournir une autre liste de noms de colonnes.

Je n’ai pas essayé cela comme première solution car je n’étais pas certain de la manière dont cela se comporterait. Mais ça marche.

Version 1

C'est trop compliqué, mais ça marche aussi.

Tu peux le faire:

  1. utilisez df.columns pour obtenir une liste des noms des colonnes
  2. utiliser cette liste de noms pour faire une liste des colonnes
  3. transmettez cette liste à quelque chose qui invoquera la fonction add surchargée de la colonne de manière de manière fonctionnelle, type de pli)

Avec réduire de python, quelques informations sur le fonctionnement de la surcharge d'opérateur et le code pyspark des colonnes ici qui devient:

def column_add(a,b):
     return  a.__add__(b)

newdf = df.withColumn('total_col', 
         reduce(column_add, ( df[col] for col in df.columns ) ))

Notez que ceci est une réduction python, pas une réduction spark RDD, et le terme de parenthèse dans le deuxième paramètre à réduire requiert la parenthèse car il s'agit d'une expression de générateur de liste. 

Testé, ça marche!

$ pyspark
>>> df = sc.parallelize([{'a': 1, 'b':2, 'c':3}, {'a':8, 'b':5, 'c':6}, {'a':3, 'b':1, 'c':0}]).toDF().cache()
>>> df
DataFrame[a: bigint, b: bigint, c: bigint]
>>> df.columns
['a', 'b', 'c']
>>> def column_add(a,b):
...     return a.__add__(b)
...
>>> df.withColumn('total', reduce(column_add, ( df[col] for col in df.columns ) )).collect()
[Row(a=1, b=2, c=3, total=6), Row(a=8, b=5, c=6, total=19), Row(a=3, b=1, c=0, total=4)]
32
Paul

Mon problème était semblable au précédent (un peu plus complexe) car je devais ajouter des sommes de colonnes consécutives en tant que nouvelles colonnes dans le cadre de données PySpark. Cette approche utilise le code de la version 1 de Paul ci-dessus: 

import pyspark
from pyspark.sql import SparkSession
import pandas as pd

spark = SparkSession.builder.appName('addColAsCumulativeSUM').getOrCreate()
df=spark.createDataFrame(data=[(1,2,3),(4,5,6),(3,2,1)\
                              ,(6,1,-4),(0,2,-2),(6,4,1)\
                              ,(4,5,2),(5,-3,-5),(6,4,-1)]\
                              ,schema=['x1','x2','x3'])
df.show()

+---+---+---+
| x1| x2| x3|
+---+---+---+
|  1|  2|  3|
|  4|  5|  6|
|  3|  2|  1|
|  6|  1| -4|
|  0|  2| -2|
|  6|  4|  1|
|  4|  5|  2|
|  5| -3| -5|
|  6|  4| -1|
+---+---+---+

colnames=df.columns

ajouter de nouvelles colonnes qui sont des sommes cumulatives (consécutives):

for i in range(0,len(colnames)):
    colnameLst= colnames[0:i+1]
    colname = 'cm'+ str(i+1)
    df = df.withColumn(colname, sum(df[col] for col in colnameLst))

df.show ()

+---+---+---+---+---+---+
| x1| x2| x3|cm1|cm2|cm3|
+---+---+---+---+---+---+
|  1|  2|  3|  1|  3|  6|
|  4|  5|  6|  4|  9| 15|
|  3|  2|  1|  3|  5|  6|
|  6|  1| -4|  6|  7|  3|
|  0|  2| -2|  0|  2|  0|
|  6|  4|  1|  6| 10| 11|
|  4|  5|  2|  4|  9| 11|
|  5| -3| -5|  5|  2| -3|
|  6|  4| -1|  6| 10|  9|
+---+---+---+---+---+---+

Les colonnes 'somme cumulative' ajoutées sont les suivantes:

cm1 = x1
cm2 = x1 + x2
cm3 = x1 + x2 + x3
0
Grant Shannon

La solution 

newdf = df.withColumn('total', sum(df[col] for col in df.columns))

posté par @Paul travaille. Néanmoins, j'ai eu l'erreur, comme beaucoup d'autres que j'ai vu,

TypeError: 'Column' object is not callable

Après un certain temps, j'ai trouvé le problème (du moins dans mon cas). Le problème est que j’avais déjà importé certaines fonctions de pyspark avec la ligne

from pyspark.sql.functions import udf, col, count, sum, when, avg, mean, min

la ligne a donc importé la commande sum pyspark alors que df.withColumn('total', sum(df[col] for col in df.columns)) est censé utiliser la fonction python sum normale.

Vous pouvez supprimer la référence de la fonction pyspark avec del sum.

Sinon, dans mon cas, j'ai modifié l'importation en 

import pyspark.sql.functions as F

et ensuite référencé les fonctions en tant que F.sum.

0
Francesco Boi

La façon la plus simple de le faire est d’utiliser la fonction expr

from pyspark.sql.functions import *
data = data.withColumn('total', expr("col1 + col2 + col3 + col4"))
0
Jonathan