Длина кадра данных внутри функции 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 будет затем использоваться в групповом режиме.

Я придумаю это решение , которое почти нормально:

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/length_group) на lit(1/count("sale_qty")), то созданный столбец теперь содержит только 1 элемент, который приводит к ошибке ...

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()