在pyspark

Question

我的pyspark数据框中有两个时间戳列。我想创建第三列，其中包含两个时间戳之间的时间戳小时数组。

这是我为此编写的代码。

# Creating udf function
def getBetweenStamps(st_date, dc_date):
    import numpy as np
    hr = 0
    date_list = []
    runnig_date = st_date
    while (dc_date>runnig_date):
        runnig_date = st_date+timedelta(hours=hr)
        date_list.append(runnig_date)
        hr+=1
    dates = np.array(date_list)
    return dates
udf_betweens = F.udf(getBetweenStamps, ArrayType(DateType()))

# Using udf function
orders.withColumn('date_array', udf_betweens(F.col('start_date'), F.col('ICUDischargeDate'))).show()

但是这显示了错误

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

我认为对函数的输入是以两个数组而不是两个引起错误的日期时间输入的。有没有办法解决？还有其他解决方法吗？

非常感谢。

Answer 1

从udf返回numpy数组时出现错误。您只需返回date_list即可。

def getBetweenStamps(st_date, dc_date):
    import numpy as np
    hr = 0
    date_list = []
    runnig_date = st_date
    while (dc_date>runnig_date):
        runnig_date = st_date+timedelta(hours=hr)
        date_list.append(runnig_date)
        hr+=1
    return date_list
udf_betweens = F.udf(getBetweenStamps, ArrayType(DateType()))

要测试以上功能：

df = spark.sql("select current_timestamp() as t1").withColumn("t2", col("t1") + expr("INTERVAL 1 DAYS"))
df.withColumn('date_array', udf_betweens(F.col('t1'), F.col('t2'))).show()