使用 udfs 提取 pyspark 数据帧中的列

Question

我正在使用下面的代码片段来提取数据帧列的一部分。

df.withColumn("chargemonth",getBookedMonth1(df['chargedate']))
def getBookedMonth1(chargedate):
booked_year=chargedate[0:3]
booked_month=chargedate[5:7]
return booked_year+"-"+booked_month

我也使用了相同的getBookedMonth，但是在这两种情况下，我都对新列chargemonthnull value。

from pyspark.sql.functions import substring
def getBookedMonth(chargedate):
booked_year=substring(chargedate, 1,4)
booked_month=substring(chargedate,5, 6)
return booked_year+"-"+booked_month

这是在 pyspark 中提取/列子字符串的正确方法吗？

Answer 1

请不要为此使用 udf！UDF 以性能不佳而闻名。

我建议你使用 Spark 内置函数来操作日期。下面是一个示例：

# DF sample
data = [(1, "2019-12-05"), (2, "2019-12-06"), (3, "2019-12-07")]
df = spark.createDataFrame(data, ["id", "chargedate"])
# format dates as 'yyyy-MM'
df.withColumn("chargemonth", date_format(to_date(col("chargedate")), "yyyy-MM")).show()
+---+----------+-----------+
| id|chargedate|chargemonth|
+---+----------+-----------+
|  1|2019-12-05|    2019-12|
|  2|2019-12-06|    2019-12|
|  3|2019-12-07|    2019-12|
+---+----------+-----------+

Answer 2

你需要创建一个新函数作为 Pyspark UDF。

>>> from pyspark.sql.functions import udf
>>> data = [
...     {"chargedate":"2019-01-01"},
...     {"chargedate":"2019-02-01"},
...     {"chargedate":"2019-03-01"},
...     {"chargedate":"2019-04-01"}
... ]
>>>
>>> booked_month = udf(lambda a:"{0}-{1}".format(a[0:4], a[5:7]))
>>>
>>> df = spark.createDataFrame(data)
>>> df = df.withColumn("chargemonth",booked_month(df['chargedate'])).drop('chargedate')
>>> df.show()
+-----------+
|chargemonth|
+-----------+
|    2019-01|
|    2019-02|
|    2019-03|
|    2019-04|
+-----------+
>>>

withColumn是添加列的正确方法，drop用于删除列。