web-dev-qa-db-ja.com

Pyspark：複数の配列列を行に分割します

1つの行と複数の列を持つデータフレームがあります。一部の列は単一の値であり、他の列はリストです。すべてのリスト列は同じ長さです。リスト以外の列はそのままにして、各リスト列を個別の行に分割します。

サンプル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|
# +---+---------+---------+---+

私が欲しいもの：

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

リスト列が1つしかない場合は、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|
# +---+---+---------+---+

ただし、explodec列も試してみると、必要な長さの2乗の長さのデータフレームになります。

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|
# +---+---+---+---+

私が欲しいのは、各列について、その列の配列のn番目の要素を取得し、それを新しい行に追加します。データフレーム内のすべての列にわたって爆発をマッピングしようとしましたが、それもうまくいかないようです：

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

スパーク> = 2.4

Zip_udfをarrays_Zip関数に置き換えることができます

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"))

スパーク<2.4

DataFramesおよび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"))

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"]))

両方のソリューションは、Python通信オーバーヘッドのために非効率的です。データサイズが固定されている場合、次のようなことができます。

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")

あるいは：

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"))

これは、UDFまたはRDDと比較して大幅に高速です。任意の数の列をサポートするように一般化：

# 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

各入力行から複数​​の出力行を作成するため、flatMapではなくmapを使用する必要があります。

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