web-dev-qa-db-ja.com

Pyspark：UDFで複数の列を渡す

データフレームの最初の列を除くすべての列を取得し、合計（またはその他の操作）を実行するユーザー定義関数を作成しています。現在、データフレームには3列または4列以上の列が含まれることがあります。それは異なります。

4つの列名をUDFで渡すようにハードコーディングできることは知っていますが、この場合は異なるため、どのようにそれを実現するのか知りたいですか？

最初の2つの例には2つの列を追加し、2番目の例には3つの列を追加します。

enter image description here

apache-sparkpysparkspark-dataframe
28
sjishan

UDFに渡すすべての列のデータ型が同じである場合、入力パラメーターとして配列を使用できます。次に例を示します。

>>> from pyspark.sql.types import IntegerType
>>> from pyspark.sql.functions import udf, array
>>> sum_cols = udf(lambda arr: sum(arr), IntegerType())
>>> spark.createDataFrame([(101, 1, 16)], ['ID', 'A', 'B']) \
...     .withColumn('Result', sum_cols(array('A', 'B'))).show()
+---+---+---+------+
| ID|  A|  B|Result|
+---+---+---+------+
|101|  1| 16|    17|
+---+---+---+------+

>>> spark.createDataFrame([(101, 1, 16, 8)], ['ID', 'A', 'B', 'C'])\
...     .withColumn('Result', sum_cols(array('A', 'B', 'C'))).show()
+---+---+---+---+------+
| ID|  A|  B|  C|Result|
+---+---+---+---+------+
|101|  1| 16|  8|    25|
+---+---+---+---+------+
29
Mariusz

配列の代わりに構造体を使用

from pyspark.sql.types import IntegerType
from pyspark.sql.functions import udf, struct
sum_cols = udf(lambda x: x[0]+x[1], IntegerType())
a=spark.createDataFrame([(101, 1, 16)], ['ID', 'A', 'B'])
a.show()
a.withColumn('Result', sum_cols(struct('A', 'B'))).show()
16
kavi

配列と構造体を使用しない別の簡単な方法。

from pyspark.sql.types import IntegerType
from pyspark.sql.functions import udf, struct

def sum(x, y):
    return x + y

sum_cols = udf(sum, IntegerType())

a=spark.createDataFrame([(101, 1, 16)], ['ID', 'A', 'B'])
a.show()
a.withColumn('Result', sum_cols('A', 'B')).show()
10
neeraj bhadani

これは私が試した方法であり、動作しているように見えました：

colsToSum = df.columns[1:]
df_sum = df.withColumn("rowSum", sum([df[col] for col in colsToSum]))
0
Travis Hensrud