UDF函数中的数据长度

Question

我需要编写一个复杂的用户定义函数(UDF)，它以多个列作为输入。类似于：

val uudf = udf{(val:Int, lag:Int, cumsum_p:Double) => val + lag + cum_p} // actually a more complex function but let's make it simple

第三个参数cumsum_p表示的是p的累积和，其中p是计算出的组的长度。因为这个udf随后将在groupby中使用。

，我想出了一个解决方案，，这几乎可以：

val uudf = udf{(val:Int, lag:Int, cumsum_p:Double) => val + lag + cum_p}
val w = Window.orderBy($"sale_qty")
df.withColumn("needThat", 
    uudf(col("sale_qty"),
       lead("sale_qty",1).over(w), sum(lit(1/length_group)).over(w)
    )
).show()

问题是，如果我用lit(1/count("sale_qty"))替换lit(1/length_group)，创建的列现在只包含一个元素，这会导致错误.

Answer 1

您应该首先计算count("sale_qty")：

val w = Window.orderBy($"sale_qty")
df
.withColumn("cnt",count($"sale_qty").over()) 
.withColumn("needThat", 
    uudf(col("sale_qty"),
       lead("sale_qty",1).over(w), sum(lit(1)/$"cnt").over(w)
    )
).show()