web-dev-qa-db-fra.com

Pyspark: diviser plusieurs colonnes de tableau en lignes

J'ai un dataframe qui a une ligne et plusieurs colonnes. Certaines des colonnes sont des valeurs uniques et d'autres des listes. Toutes les colonnes de la liste ont la même longueur. Je souhaite scinder chaque colonne de la liste en une ligne distincte, tout en conservant telle quelle la colonne non-liste.

Échantillon DF:

from pyspark import Row
from pyspark.sql import SQLContext
from pyspark.sql.functions import explode

sqlc = SQLContext(sc)

df = sqlc.createDataFrame([Row(a=1, b=[1,2,3],c=[7,8,9], d='foo')])
# +---+---------+---------+---+
# |  a|        b|        c|  d|
# +---+---------+---------+---+
# |  1|[1, 2, 3]|[7, 8, 9]|foo|
# +---+---------+---------+---+

Ce que je veux:

+---+---+----+------+
|  a|  b|  c |    d |
+---+---+----+------+
|  1|  1|  7 |  foo |
|  1|  2|  8 |  foo |
|  1|  3|  9 |  foo |
+---+---+----+------+

Si je n'avais qu'une colonne de liste, ce serait facile en faisant simplement un explode:

df_exploded = df.withColumn('b', explode('b'))
# >>> df_exploded.show()
# +---+---+---------+---+
# |  a|  b|        c|  d|
# +---+---+---------+---+
# |  1|  1|[7, 8, 9]|foo|
# |  1|  2|[7, 8, 9]|foo|
# |  1|  3|[7, 8, 9]|foo|
# +---+---+---------+---+

Cependant, si j'essaie également de explode la colonne c, je me retrouve avec un cadre de données de longueur égale à ce que je veux:

df_exploded_again = df_exploded.withColumn('c', explode('c'))
# >>> df_exploded_again.show()
# +---+---+---+---+
# |  a|  b|  c|  d|
# +---+---+---+---+
# |  1|  1|  7|foo|
# |  1|  1|  8|foo|
# |  1|  1|  9|foo|
# |  1|  2|  7|foo|
# |  1|  2|  8|foo|
# |  1|  2|  9|foo|
# |  1|  3|  7|foo|
# |  1|  3|  8|foo|
# |  1|  3|  9|foo|
# +---+---+---+---+

Ce que je veux, c'est - pour chaque colonne, prenons le nième élément du tableau dans cette colonne et l'ajoutons à une nouvelle ligne. J'ai essayé de mapper un éclat sur toutes les colonnes du cadre de données, mais cela ne semble pas fonctionner non plus:

df_split = df.rdd.map(lambda col: df.withColumn(col, explode(col))).toDF()
pythonapache-sparkdataframepysparkapache-spark-sql
45
Steve

étincelle> = 2.4

Vous pouvez remplacer Zip_udf avec arrays_Zip une fonction

from pyspark.sql.functions import arrays_Zip, col

(df
    .withColumn("tmp", arrays_Zip("b", "c"))
    .withColumn("tmp", explode("tmp"))
    .select("a", col("tmp.b"), col("tmp.c"), "d"))

étincelle <2.4

Avec DataFrames et UDF:

from pyspark.sql.types import ArrayType, StructType, StructField, IntegerType
from pyspark.sql.functions import col, udf, explode

Zip_ = udf(
  lambda x, y: list(Zip(x, y)),
  ArrayType(StructType([
      # Adjust types to reflect data types
      StructField("first", IntegerType()),
      StructField("second", IntegerType())
  ]))
)

(df
    .withColumn("tmp", Zip_("b", "c"))
    # UDF output cannot be directly passed to explode
    .withColumn("tmp", explode("tmp"))
    .select("a", col("tmp.first").alias("b"), col("tmp.second").alias("c"), "d"))

Avec RDDs:

(df
    .rdd
    .flatMap(lambda row: [(row.a, b, c, row.d) for b, c in Zip(row.b, row.c)])
    .toDF(["a", "b", "c", "d"]))

Les deux solutions sont inefficaces en raison de la surcharge de communication de Python). Si la taille des données est fixe, vous pouvez procéder comme suit:

from functools import reduce
from pyspark.sql import DataFrame

# Length of array
n = 3

# For legacy Python you'll need a separate function
# in place of method accessor 
reduce(
    DataFrame.unionAll, 
    (df.select("a", col("b").getItem(i), col("c").getItem(i), "d")
        for i in range(n))
).toDF("a", "b", "c", "d")

ou même:

from pyspark.sql.functions import array, struct

# SQL level Zip of arrays of known size
# followed by explode
tmp = explode(array(*[
    struct(col("b").getItem(i).alias("b"), col("c").getItem(i).alias("c"))
    for i in range(n)
]))

(df
    .withColumn("tmp", tmp)
    .select("a", col("tmp").getItem("b"), col("tmp").getItem("c"), "d"))

Cela devrait être nettement plus rapide par rapport à UDF ou RDD. Généralisé pour supporter un nombre arbitraire de colonnes:

# This uses keyword only arguments
# If you use legacy Python you'll have to change signature
# Body of the function can stay the same
def Zip_and_explode(*colnames, n):
    return explode(array(*[
        struct(*[col(c).getItem(i).alias(c) for c in colnames])
        for i in range(n)
    ]))

df.withColumn("tmp", Zip_and_explode("b", "c", n=3))
53
user6910411

Vous devez utiliser flatMap et non pas map pour créer plusieurs lignes en sortie à partir de chaque ligne en entrée.

from pyspark.sql import Row
def dualExplode(r):
    rowDict = r.asDict()
    bList = rowDict.pop('b')
    cList = rowDict.pop('c')
    for b,c in Zip(bList, cList):
        newDict = dict(rowDict)
        newDict['b'] = b
        newDict['c'] = c
        yield Row(**newDict)

df_split = sqlContext.createDataFrame(df.rdd.flatMap(dualExplode))
9
David