web-dev-qa-db-ja.com

spark DataFrame関数のパラメーターとして列値を使用する

次のDataFrameについて考えてみます。

#+------+---+
#|letter|rpt|
#+------+---+
#|     X|  3|
#|     Y|  1|
#|     Z|  2|
#+------+---+

これは、次のコードを使用して作成できます。

df = spark.createDataFrame([("X", 3),("Y", 1),("Z", 2)], ["letter", "rpt"])

この question のように、列rptで指定された回数だけ各行を繰り返したいとします。

1つの方法は、次のpyspark-sqlクエリを使用して、私の solution をその質問に複製することです。

query = """
SELECT *
FROM
  (SELECT DISTINCT *,
                   posexplode(split(repeat(",", rpt), ",")) AS (index, col)
   FROM df) AS a
WHERE index > 0
"""
query = query.replace("\n", " ")  # replace newlines with spaces, avoid EOF error
spark.sql(query).drop("col").sort('letter', 'index').show()
#+------+---+-----+
#|letter|rpt|index|
#+------+---+-----+
#|     X|  3|    1|
#|     X|  3|    2|
#|     X|  3|    3|
#|     Y|  1|    1|
#|     Z|  2|    1|
#|     Z|  2|    2|
#+------+---+-----+

これは機能し、正しい答えを生成します。ただし、DataFrameAPI関数を使用してこの動作を再現することはできません。

私は試した：

import pyspark.sql.functions as f
df.select(
    f.posexplode(f.split(f.repeat(",", f.col("rpt")), ",")).alias("index", "col")
).show()

しかし、これは次の結果になります。

TypeError: 'Column' object is not callable

クエリ内のrepeatへの入力として列を渡すことができるのに、APIからは渡せないのはなぜですか？ spark DataFrame関数を使用してこの動作を複製する方法はありますか？

apache-sparkpysparkapache-spark-sqlpyspark-sql
7
pault

1つのオプションは、 pyspark.sql.functions.expr 、これにより、spark-sql関数への入力として列の値を使用できます。

@ user8371915の コメント に基づいて、次のことが機能することがわかりました。

from pyspark.sql.functions import expr

df.select(
    '*',
    expr('posexplode(split(repeat(",", rpt), ","))').alias("index", "col")
).where('index > 0').drop("col").sort('letter', 'index').show()
#+------+---+-----+
#|letter|rpt|index|
#+------+---+-----+
#|     X|  3|    1|
#|     X|  3|    2|
#|     X|  3|    3|
#|     Y|  1|    1|
#|     Z|  2|    1|
#|     Z|  2|    2|
#+------+---+-----+
7
pault