将列值用作参数传递给 Spark DataFrame 函数

Question

考虑以下DataFrame：

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

可以使用以下代码来创建：

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

假设我想按列rpt中指定的次数重复每一行，就像这个问题中所示。

一种方法是使用以下pyspark-sql查询来复制我对该问题的解决方案：

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'对象不可调用

为什么我能够在查询中将列作为输入传递给repeat，但无法从API中传递？是否有一种方法可以使用Spark DataFrame函数复制此行为？

- pault

---

f.expr("""repeat(",", rpt)""") 而不是 f.repeat(",", f.col("rpt"))？ - Alper t. Turker

@user8371915 df.select('*', f.expr('posexplode(split(repeat(",", rpt), ","))').alias("index", "col")).where('index > 0').drop("col").sort('letter', 'index').show() 是有效的。你知道是否只有这种方式可以将列作为参数使用吗？为什么它在SQL语法中有效？ - pault

@user8371915 请考虑将您的建议发布为答案（并且可以将其从我的问题中编辑出来）。我认为这对未来的其他人会有好处。 - pault

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

- pault