使用udfs

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来表示相同的内容，但是在两种情况下，我都将新列null value用作 chargemonth 。

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 builtin functions来操纵日期。这是一个示例：

# 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用于删除列。