Question

我有一个Databricks数据框，其中有多列，而UDF则根据其他列的值生成新列的内容。

原始数据集的示例是：

interval_group_id     control_value     pulse_value     device_timestamp

2797895314            5                 5               2020-09-12 09:08:44
0                     5                 5               2020-09-12 09:08:45
0                     6                 5               2020-09-12 09:08:46
0                     0                 5               2020-09-12 09:08:47

现在，我根据上述各列的某些逻辑，尝试添加一个名为group_id的新列。我的UDF代码是：

@udf('integer')
def udf_calculate_group_id_new (interval_group_id, prev_interval_group_id, control_val, pulse_val):
  
  if interval_group_id != 0:
    return interval_group_id
  elif control_val >= pulse_val and prev_interval_group_id != 0:
    return prev_interval_group_id
  else:
    return -1

添加到我的数据框中的新列的操作如下：

df = df.withColumn('group_id'
                   , udf_calculate_group_id_new(
                          df.interval_group_id                                                   
                          , lag(col('interval_group_id')).over(Window.orderBy('device_timestamp'))            
                          , df.control_value
                          , df.pulse_value)
                  )

我的预期结果是：

interval_group_id    control_value    pulse_value   device_timestamp       group_id

2797895314           5                5             2020-09-12 09:08:44    2797895314
0                    5                5             2020-09-12 09:08:45    2797895314
0                    6                5             2020-09-12 09:08:46    2797895314
0                    0                5             2020-09-12 09:08:47    -1

但是，添加新的group_id列的结果是：

interval_group_id    control_value    pulse_value   device_timestamp       group_id

2797895314           5                5             2020-09-12 09:08:44    null
0                    5                5             2020-09-12 09:08:45    null
0                    6                5             2020-09-12 09:08:46    -1
0                    0                5             2020-09-12 09:08:47    -1

我的目标是根据上述条件，将值2797895314沿group_id列传播，但这种情况不会发生，并且会填充结果null和-1的错误。

这是UDF的错误还是我对UDF的工作方式的期望不正确？还是我只是不擅长编码？

PySpark UDF不返回预期结果

0 个答案: