Comment faire correspondre une colonne à plusieurs colonnes dans pyspark ?

Question

J'ai un cadre de données pyspark qui ressemble à ceci :

df.show()
+---+
|dim|
+---+
|1x1|
|0x0|
|1x0|
+---+

Le type de données dans dim es str . Maintenant je veux séparer dim en 2 colonnes, et avoir quelque chose comme ça :

df.show()
+---+----+----+
|dim|dim1|dim2|
+---+----+----+
|1x1|   1|   1|
|0x0|   0|   0|
|1x0|   1|   0|
+---+----+----+

Je sais que si je devais opérer sur une seule chaîne de caractères, j'utiliserais simplement la fonction split() en python : "1x1".split("x") Mais comment créer simultanément plusieurs colonnes à partir d'une colonne mappée par une fonction de fractionnement ?

Answer 1

Vous pouvez essayer

from pyspark.sql.functions import split
df_temp = df.withColumn("dim1", split("dim", "x")[0])
df_new = df_temp.withColumn("dim2", split("dim", "x")[1])
df_new.show()

Ou vous pouvez les enchaîner en un seul appel :

df_new = df.withColumn("dim1", split("dim", "x")[0])\
    .withColumn("dim2", split("dim", "x")[1])

Answer 2

Une autre option consiste à utiliser rdd y map() :

df = df.rdd.map(lambda row: (row['dim'],) + tuple(row['dim'].split('x')))\
    .toDF(["dim", "dim1", "dim2"])
df.show()
#+---+----+----+
#|dim|dim1|dim2|
#+---+----+----+
#|1x1|   1|   1|
#|0x0|   0|   0|
#|1x0|   1|   0|
#+---+----+----+

Et voici une adaptation de @Grace O'Halloran 's réponse en utilisant pyspark.sql.functions.split() dans une boucle au lieu d'énumérer les possibilités (utile si vous avez un grand nombre de dimensions) :

ndim = 2
import pyspark.sql.functions as f
for d in range(ndim):
    df = df.withColumn('dim{}'.format(d+1), f.split('dim', 'x')[d])