apache-spark - 使用列值作为spark DataFrame函数的参数

Question

考虑以下DataFrame:

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

可以使用以下代码创建:

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

假设我想重复每行 rpt列中指定的次数，就像此 question一样。

一种方法是使用以下 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|
#+------+---+-----+

这可以工作并产生正确的答案。但是，我无法使用DataFrame API函数来复制此行为。

我试过了:

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函数复制此行为？

Answer 1

一种选择是使用 pyspark.sql.functions.expr ，它允许您将列值用作spark-sql函数的输入。

基于@ user8371915的comment，我发现以下作品有效:

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