Spark 中的 Python 函数

Question

我正在尝试根据activity_prioritization_rounding数据帧上的标志向前或向后舍入deadline_date列。 -1 表示向后，0 表示无，1 表示向前。

当我使用单个日期作为变量时，该函数可以工作，但我很难将其应用于整个数据集。我收到的错误是“ValueError:无法将列转换为 bool 值:请使用 '&' 表示 'and'、'|'构建 DataFrame bool 表达式时，代表“或”，“~”代表“非”。”当尝试传递函数的列部分时。 python 中相对较新的构建函数。

from pyspark.sql.functions import next_day, date_sub
from pyspark.sql.functions import to_date


def next_date(column,date,dayOfWEek):
  if column == -1:
    return date_sub(next_day(date,dayOfWEek),0)
  elif column == 1:
    return date_sub(next_day(date,dayOfWEek),7)
  else:
    return date


activity_prioritization_rounding= sql("""select * from spa.activity_master""")
activity_prioritization_rounding.withColumn(
   "New_Date",
    next_date(col("deadline_rounding"),col("deadline_date"),"Friday"))
)

Answer 1

您需要从 python 函数中创建一个 udf 并将 Frid​​ay 作为列发送，因为它不会在数据帧中广播。您可以使用 lit 来做到这一点。

from pyspark.sql.functions import udf, next_day, date_sub, to_date, lit
from pyspark.sql.types import DateType

activity_prioritization_rounding.withColumn("New_Date",udf(next_date(col("deadline_rounding"),col("deadline_date"),lit("Friday"), DateType())))

编辑:正如 @jxc 正确提到的，您不能在 UDF 中使用 Spark 函数。

将其简化为 when().when().otherwise()

from pyspark.sql.functions import udf, next_day, date_sub, to_date, lit, when, col

day_of_week = "Friday"
activity_prioritization_rounding.withColumn("New_Date", when(
    col("deadline_rounding") == -1, date_sub(next_day(col("deadline_date"), day_of_week), 0)).when(
    col("deadline_rounding") == 1, date_sub(next_day(col("deadline_date"), day_of_week), 7)).otherwise(
    col("deadline_date")))