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

Posted 2023-03-23

【中文标题】使用列值作为 spark DataFrame 函数的参数

【英文标题】：Using a column value as a parameter to a spark DataFrame function

【发布时间】：2018-12-10 22:26:45

【问题描述】：

考虑以下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 函数来复制这种行为？

【问题讨论】：

f.expr("""repeat(",", rpt)""") 而不是 f.repeat(",", f.col("rpt"))?

@user8371915 df.select('*', f.expr('posexplode(split(repeat(",", rpt), ","))').alias("index", "col")).where('index > 0').drop("col").sort('letter', 'index').show() 有效。你知道这是否是使用 Column 作为参数的唯一方法吗？为什么它在 sql 语法中起作用？

@user8371915 请考虑将您的建议作为答案发布（并且可以在我的问题之外对其进行编辑）。我认为这将对其他人有益。

【参考方案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|
#+------+---+-----+