web-dev-qa-db-ja.com

Apache Spark-UDFの結果を複数のデータフレーム列に割り当てます

私はpysparkを使用して、大きなcsvファイルをspark-csvでデータフレームにロードし、前処理ステップとして、列の1つ（json文字列を含む）で利用可能なデータにさまざまな操作を適用する必要があります。これはX値を返します。それぞれの値は独自の個別の列に格納する必要があります。

その機能はUDFで実装されます。ただし、そのUDFから値のリストを返し、それらを個々の列にフィードする方法はわかりません。以下に簡単な例を示します。

(...)
from pyspark.sql.functions import udf
def udf_test(n):
    return [n/2, n%2]

test_udf=udf(udf_test)


df.select('amount','trans_date').withColumn("test", test_udf("amount")).show(4)

これにより、以下が生成されます。

+------+----------+--------------------+
|amount|trans_date|                test|
+------+----------+--------------------+
|  28.0|2016-02-07|         [14.0, 0.0]|
| 31.01|2016-02-07|[15.5050001144409...|
| 13.41|2016-02-04|[6.70499992370605...|
| 307.7|2015-02-17|[153.850006103515...|
| 22.09|2016-02-05|[11.0450000762939...|
+------+----------+--------------------+
only showing top 5 rows

Udfによって返される2つの（この例の）値を別々の列に保存する最良の方法は何でしょうか？現在、それらは文字列として入力されています：

df.select('amount','trans_date').withColumn("test", test_udf("amount")).printSchema()

root
 |-- amount: float (nullable = true)
 |-- trans_date: string (nullable = true)
 |-- test: string (nullable = true)
pythonapache-sparkpysparkapache-spark-sqluser-defined-functions
38
Everaldo Aguiar

単一のUDF呼び出しから複数のトップレベル列を作成することはできませんが、新しいstructを作成できます。 returnTypeを指定したUDFが必要です。

from pyspark.sql.functions import udf
from pyspark.sql.types import *

schema = StructType([
    StructField("foo", FloatType(), False),
    StructField("bar", FloatType(), False)
])

def udf_test(n):
    return (n / 2, n % 2) if n and n != 0.0 else (float('nan'), float('nan'))

test_udf = udf(udf_test, schema)
df = sc.parallelize([(1, 2.0), (2, 3.0)]).toDF(["x", "y"])

foobars = df.select(test_udf("y").alias("foobar"))
foobars.printSchema()
## root
##  |-- foobar: struct (nullable = true)
##  |    |-- foo: float (nullable = false)
##  |    |-- bar: float (nullable = false)

単純なselectでスキーマをさらにフラット化します。

foobars.select("foobar.foo", "foobar.bar").show()
## +---+---+
## |foo|bar|
## +---+---+
## |1.0|0.0|
## |1.5|1.0|
## +---+---+

Spark DataFrame の単一の列から複数の列を取得する]も参照してください。

58
zero323