Spark Dataframe: пользовательская функция в окне

Question

У меня есть DataFrame df со столбцом, который содержит значения Double.

Мне удалось вычислить скользящее среднее следующим образом:

val myWindow =  Window.partitionBy("id")
  .orderBy(asc("timestamp"))
  .rowsBetween(-2,2)

val newDf = df.withColumn("slidingAvg",avg($"value").over(myWindow))

Но теперь я хочу улучшитьнемного мои вычисления.Вместо этого базового скользящего среднего я хочу вычислить «квадрат среднего скользящего корня» (значения ближе к 0 должны влиять на мой результат больше, чем на более высокие или более низкие значения).

Вот моя попытка на данный момент:

val myWindow =  Window.partitionBy("id")
  .orderBy(asc("timestamp"))
  .rowsBetween(-2,2)

val newDf = df.withColumn("squaredSlidingRootAverage",avg(sqrt($"value")).over(myWindow) * avg(sqrt($"value")).over(myWindow))

Это работает, когда у меня есть только положительные значения, но с отрицательными значениями это не дает мне ожидаемый результат (из-за sqrt).Я хотел бы иметь «что-то», которое делает этот расчет, когда есть только положительные значения:

avg(sqrt($"value")).over(myWindow) * avg(sqrt($"value")).over(myWindow)

Это, когда есть только отрицательные значения:

- avg( - sqrt(abs($"value"))).over(myWindow) * avg( - sqrt(abs($"value"))).over(myWindow))

И смесь обоихкогда положительные и отрицательные значения смешиваются в окне.

Поэтому я пытаюсь реализовать следующее:

val myWindow =  Window.partitionBy("id")
  .orderBy(asc("timestamp"))
  .rowsBetween(-2,2)

val newDf = df.withColumn("squaredSlidingRootAverage",
           [AVG_SIGN] * avg([VALUE_SIGN] * sqrt(abs($"value"))).over(myWindow) * 
                        avg([VALUE_SIGN] * sqrt(abs($"value"))).over(myWindow))

Как я могу это реализовать?

Образец ввода:

+---+----------+------+
| id| timestamp| value|
+---+----------+------+
|  0|         0|     0|
|  0|         1|     0|
|  0|         2|     0|
|  0|         3|     0|
|  0|         4|    25|
|  0|         5|  -100|
|  0|         6|   -25|
|  0|         7|     9|
|  0|         8|     0|
|  0|         9|     0|
|  0|        10|   -25|
|  0|        11|   100|
|  0|        12|     0|
|  0|        13|     0|
|  0|        14|     0|
|  0|        15|     0|

Образец вывода:

+---+----------+------+
| id| timestamp| value|
+---+----------+------+
|  0|         0|     0|
|  0|         1|     0|
|  0|         2|     1|
|  0|         3|    -1|
|  0|         4|    -4|
|  0|         5| -1.96|
|  0|         6| -1.96|
|  0|         7| -5.76|
|  0|         8| -1.96|
|  0|         9|  2.56|
|  0|        10|     1|
|  0|        11|     1|
|  0|        12|     1|
|  0|        13|     4|
|  0|        14|     0|
|  0|        15|     0|

Answer 1

Вы можете использовать функции when и otherwise для обработки двух разных случаев:

df
    .withColumn("sqrt", when('value <0, -sqrt(- 'value)).otherwise(sqrt('value)))
    .withColumn("avg", avg('sqrt) over myWindow)
    .withColumn("value", when('avg <0, -'avg * 'avg).otherwise('avg * 'avg))
    .select("id", "timestamp", "value")