web-dev-qa-db-fra.com

Diviser la colonne de chaîne Spark Dataframe en plusieurs colonnes

J'ai vu différentes personnes suggérer que Dataframe.explode est un moyen utile de le faire, mais il en résulte plus de lignes que le cadre de données d'origine, ce qui n'est pas du tout ce que je veux. Je veux simplement faire l'équivalent Dataframe du très simple:

rdd.map(lambda row: row + [row.my_str_col.split('-')])

qui prend quelque chose qui ressemble à:

col1 | my_str_col
-----+-----------
  18 |  856-yygrm
 201 |  777-psgdg

et le convertit en ceci:

col1 | my_str_col | _col3 | _col4
-----+------------+-------+------
  18 |  856-yygrm |   856 | yygrm
 201 |  777-psgdg |   777 | psgdg

Je connais pyspark.sql.functions.split(), mais il en résulte une colonne de tableau imbriquée au lieu de deux colonnes de niveau supérieur comme je le souhaite. 

Idéalement, je souhaite que ces nouvelles colonnes soient également nommées.

apache-sparkpysparkapache-spark-sqlspark-dataframepyspark-sql
32
Peter Gaultney

pyspark.sql.functions.split() est la bonne approche: vous devez simplement aplatir la colonne ArrayType imbriquée en plusieurs colonnes de niveau supérieur. Dans ce cas, où chaque tableau ne contient que 2 éléments, c'est très simple. Vous utilisez simplement Column.getItem() pour récupérer chaque partie du tableau sous forme de colonne:

split_col = pyspark.sql.functions.split(df['my_str_col'], '-')
df = df.withColumn('NAME1', split_col.getItem(0))
df = df.withColumn('NAME2', split_col.getItem(1))

Le résultat sera:

col1 | my_str_col | NAME1 | NAME2
-----+------------+-------+------
  18 |  856-yygrm |   856 | yygrm
 201 |  777-psgdg |   777 | psgdg

Je ne sais pas comment résoudre ce problème dans un cas général où les tableaux imbriqués ne sont pas de la même taille, rangée à rangée.

54
Peter Gaultney

Voici une solution au cas général qui ne nécessite pas de connaître la longueur du tableau à l’avance, en utilisant collect ou en utilisant udfs. Malheureusement, cela ne fonctionne que pour spark version 2.1 et ultérieure, car il nécessite la fonction posexplode .

Supposons que vous ayez le DataFrame suivant:

df = spark.createDataFrame(
    [
        [1, 'A, B, C, D'], 
        [2, 'E, F, G'], 
        [3, 'H, I'], 
        [4, 'J']
    ]
    , ["num", "letters"]
)
df.show()
#+---+----------+
#|num|   letters|
#+---+----------+
#|  1|A, B, C, D|
#|  2|   E, F, G|
#|  3|      H, I|
#|  4|         J|
#+---+----------+

Fractionnez la colonne letters, puis utilisez posexplode pour faire éclater le tableau résultant, ainsi que la position dans le tableau. Utilisez ensuite pyspark.sql.functions.expr pour récupérer l’élément à l’index pos dans ce tableau.

import pyspark.sql.functions as f

df.select(
        "num",
        f.split("letters", ", ").alias("letters"),
        f.posexplode(f.split("letters", ", ")).alias("pos", "val")
    )\
    .show()
#+---+------------+---+---+
#|num|     letters|pos|val|
#+---+------------+---+---+
#|  1|[A, B, C, D]|  0|  A|
#|  1|[A, B, C, D]|  1|  B|
#|  1|[A, B, C, D]|  2|  C|
#|  1|[A, B, C, D]|  3|  D|
#|  2|   [E, F, G]|  0|  E|
#|  2|   [E, F, G]|  1|  F|
#|  2|   [E, F, G]|  2|  G|
#|  3|      [H, I]|  0|  H|
#|  3|      [H, I]|  1|  I|
#|  4|         [J]|  0|  J|
#+---+------------+---+---+

Nous créons maintenant deux nouvelles colonnes à partir de ce résultat. La première est le nom de notre nouvelle colonne, qui sera une concaténation de letter et de l'index dans le tableau. La deuxième colonne sera la valeur à l'index correspondant dans le tableau. Nous obtenons ce dernier en exploitant la fonctionnalité de pyspark.sql.functions.expr qui nous permet utilise les valeurs de colonne en tant que paramètres .

df.select(
        "num",
        f.split("letters", ", ").alias("letters"),
        f.posexplode(f.split("letters", ", ")).alias("pos", "val")
    )\
    .drop("val")\
    .select(
        "num",
        f.concat(f.lit("letter"),f.col("pos").cast("string")).alias("name"),
        f.expr("letters[pos]").alias("val")
    )\
    .show()
#+---+-------+---+
#|num|   name|val|
#+---+-------+---+
#|  1|letter0|  A|
#|  1|letter1|  B|
#|  1|letter2|  C|
#|  1|letter3|  D|
#|  2|letter0|  E|
#|  2|letter1|  F|
#|  2|letter2|  G|
#|  3|letter0|  H|
#|  3|letter1|  I|
#|  4|letter0|  J|
#+---+-------+---+

Maintenant, nous pouvons simplement groupBy la num et pivot le DataFrame. En réunissant tout cela, nous obtenons:

df.select(
        "num",
        f.split("letters", ", ").alias("letters"),
        f.posexplode(f.split("letters", ", ")).alias("pos", "val")
    )\
    .drop("val")\
    .select(
        "num",
        f.concat(f.lit("letter"),f.col("pos").cast("string")).alias("name"),
        f.expr("letters[pos]").alias("val")
    )\
    .groupBy("num").pivot("name").agg(f.first("val"))\
    .show()
#+---+-------+-------+-------+-------+
#|num|letter0|letter1|letter2|letter3|
#+---+-------+-------+-------+-------+
#|  1|      A|      B|      C|      D|
#|  3|      H|      I|   null|   null|
#|  2|      E|      F|      G|   null|
#|  4|      J|   null|   null|   null|
#+---+-------+-------+-------+-------+
12
pault

J'ai trouvé une solution pour le cas général inégal (ou lorsque vous obtenez les colonnes imbriquées, obtenues avec la fonction .split ()):

import pyspark.sql.functions as f

@f.udf(StructType([StructField(col_3, StringType(), True),
                   StructField(col_4, StringType(), True)]))

 def splitCols(array):
    return array[0],  ''.join(array[1:len(array)])

 df = df.withColumn("name", splitCols(f.split(f.col("my_str_col"), '-')))\
        .select(df.columns+['name.*'])

Fondamentalement, il vous suffit de sélectionner toutes les colonnes précédentes + les colonnes imbriquées 'nom_colonne. *' Et vous les obtiendrez sous forme de deux colonnes de niveau supérieur dans ce cas.

0
Jasminyas