Spark Dataframe Scala: добавление новых столбцов по некоторым условиям

Question

Я изменил свой вопрос, чтобы его было легче понять.

исходный df выглядит так:

+---+----------+-------+----+------+
| id|tim       |price  | qty|qtyChg|
+---+----------+-------+----+------+
|  1| 31951.509|  0.370|   1|     1|
|  2| 31951.515|145.380| 100|   100|
|  3| 31951.519|149.370| 100|   100|
|  4| 31951.520|144.370| 100|   100|
|  5| 31951.520|119.370|   5|     5|
|  6| 31951.520|149.370| 300|   200|
|  7| 31951.521|149.370| 400|   100|
|  8| 31951.522|149.370| 410|    10|
|  9| 31951.522|149.870|  50|    50|
| 10| 31951.522|109.370|  50|    50|
| 11| 31951.522|144.370| 400|   300|
| 12| 31951.524|149.370| 610|   200|
| 13| 31951.526|135.130|  22|    22|
| 14| 31951.527|149.370| 750|   140|
| 15| 31951.528| 89.370| 100|   100|
| 16| 31951.528|145.870|  50|    50|
| 17| 31951.528|139.370| 100|   100|
| 18| 31951.531|144.370| 410|    10|
| 19| 31951.531|149.370| 769|    19|
| 20| 31951.538|149.370| 869|   100|
| 21| 31951.538|144.880| 200|   200|
| 22| 31951.541|139.370| 221|   121|
| 23| 31951.542|149.370|1199|   330|
| 24| 31951.542|139.370| 236|    15|
| 25| 31951.542|144.370| 510|   100|
| 26| 31951.543|146.250|  50|    50|
| 27| 31951.543|143.820| 100|   100|
| 28| 31951.543|139.370| 381|   145|
| 29| 31951.544|149.370|1266|    67|
| 30| 31951.544|150.000|  50|    50|
| 31| 31951.544|137.870| 300|   300|
| 32| 31951.544|140.470|  10|    10|
| 33| 31951.545|150.000|  53|     3|
| 34| 31951.545|140.000|  25|    25|
| 35| 31951.545|148.310|   8|     8|
| 36| 31951.547|149.000|  20|    20|
| 37| 31951.549|143.820| 102|     2|
| 38| 31951.549|150.110|  75|    75|
+---+----------+-------+----+------+

затем я запускаю код

val ww = Window.partitionBy().orderBy($"tim") 

val step1 = df.withColumn("sequence",sort_array(collect_set(col("price")).over(ww),asc=false))
.withColumn("top1price",col("sequence").getItem(0))
.withColumn("top2price",col("sequence").getItem(1))
.drop("sequence")

Новый фрейм данных выглядит так:

+---+---------+-------+----+------+---------+---------+
| id|      tim|  price| qty|qtyChg|top1price|top2price|
+---+---------+-------+----+------+---------+---------+
|  1|31951.509|  0.370|   1|     1|    0.370|     null|
|  2|31951.515|145.380| 100|   100|  145.380|    0.370|
|  3|31951.519|149.370| 100|   100|  149.370|  145.380|
|  4|31951.520|149.370| 300|   200|  149.370|  145.380|
|  5|31951.520|144.370| 100|   100|  149.370|  145.380|
|  6|31951.520|119.370|   5|     5|  149.370|  145.380|
|  7|31951.521|149.370| 400|   100|  149.370|  145.380|
|  8|31951.522|109.370|  50|    50|  149.870|  149.370|
|  9|31951.522|144.370| 400|   300|  149.870|  149.370|
| 10|31951.522|149.870|  50|    50|  149.870|  149.370|
| 11|31951.522|149.370| 410|    10|  149.870|  149.370|
| 12|31951.524|149.370| 610|   200|  149.870|  149.370|
| 13|31951.526|135.130|  22|    22|  149.870|  149.370|
| 14|31951.527|149.370| 750|   140|  149.870|  149.370|
| 15|31951.528| 89.370| 100|   100|  149.870|  149.370|
| 16|31951.528|139.370| 100|   100|  149.870|  149.370|
| 17|31951.528|145.870|  50|    50|  149.870|  149.370|
| 18|31951.531|144.370| 410|    10|  149.870|  149.370|
| 19|31951.531|149.370| 769|    19|  149.870|  149.370|
| 20|31951.538|144.880| 200|   200|  149.870|  149.370|
| 21|31951.538|149.370| 869|   100|  149.870|  149.370|
| 22|31951.541|139.370| 221|   121|  149.870|  149.370|
| 23|31951.542|144.370| 510|   100|  149.870|  149.370|
| 24|31951.542|139.370| 236|    15|  149.870|  149.370|
| 25|31951.542|149.370|1199|   330|  149.870|  149.370|
| 26|31951.543|139.370| 381|   145|  149.870|  149.370|
| 27|31951.543|143.820| 100|   100|  149.870|  149.370|
| 28|31951.543|146.250|  50|    50|  149.870|  149.370|
| 29|31951.544|140.470|  10|    10|  150.000|  149.870|
| 30|31951.544|137.870| 300|   300|  150.000|  149.870|
| 31|31951.544|150.000|  50|    50|  150.000|  149.870|
| 32|31951.544|149.370|1266|    67|  150.000|  149.870|
| 33|31951.545|140.000|  25|    25|  150.000|  149.870|
| 34|31951.545|150.000|  53|     3|  150.000|  149.870|
| 35|31951.545|148.310|   8|     8|  150.000|  149.870|
| 36|31951.547|149.000|  20|    20|  150.000|  149.870|
| 37|31951.549|150.110|  75|    75|  150.110|  150.000|
| 38|31951.549|143.820| 102|     2|  150.110|  150.000|
+---+---------+-------+----+------+---------+---------+

Я надеюсь получить два новых столбца top1priceQty, top2priceQty, в которых будут храниться наиболее обновленные соответствующие значения top1price и top2price.

Например, в строке 6 top1price = 149,370, на основе этого значения я хочу получить соответствующее количество, равное 400 (а не 100 или 300). в строке 33, когда top1price = 150.00000000, я хочу получить соответствующее количество 53 из строки 32, а не 50 из строки 28. То же правило применяется к top2price

Заранее всем спасибо!

Answer 1

Вы сами были очень близки к ответу. Вместо того, чтобы собирать набор из одного столбца, соберите массив «LMTPRICE» и соответствующий ему «qty». Затем используйте getItem (0) .getItem (0) для top1price и getItem (0) .getItem (1) для top1priceQty. Чтобы сохранить порядок до INTEREST_TIME для получения правильного количества, используйте INTEREST_TIME также после «LMTPRICE» и перед «qty».

df.withColumn("sequence",sort_array(collect_set(array("LMTPRICE","INTEREST_TIME","qty")).over(ww),asc=false)).withColumn("top1price",col("sequence").getItem(0).getItem(0)).withColumn("top1priceQty",col("sequence").getItem(0).getItem(2).cast("int")).drop("sequence").show(false)

+-----+-------------+--------+---+------+---------+------------+
|index|INTEREST_TIME|LMTPRICE|qty|qtyChg|top1price|top1priceQty|
+-----+-------------+--------+---+------+---------+------------+
|0    |31951.509    |0.37    |1  |1     |0.37     |1           |
|1    |31951.515    |145.38  |100|100   |145.38   |100         |
|2    |31951.519    |149.37  |100|100   |149.37   |100         |
|3    |31951.52     |119.37  |5  |5     |149.37   |300         |
|4    |31951.52     |144.37  |100|100   |149.37   |300         |
|5    |31951.52     |149.37  |300|200   |149.37   |300         |
|6    |31951.521    |149.37  |400|100   |149.37   |400         |
|7    |31951.522    |149.87  |50 |50    |149.87   |50          |
|8    |31951.522    |149.37  |410|10    |149.87   |50          |
|9    |31951.522    |109.37  |50 |50    |149.87   |50          |
|10   |31951.522    |144.37  |400|300   |149.87   |50          |
|11   |31951.524    |149.87  |610|200   |149.87   |610         |
|12   |31951.526    |135.13  |22 |22    |149.87   |610         |
|13   |31951.527    |149.37  |750|140   |149.87   |610         |
|14   |31951.528    |139.37  |100|100   |149.87   |610         |
|15   |31951.528    |145.87  |50 |50    |149.87   |610         |
|16   |31951.528    |89.37   |100|100   |149.87   |610         |
|17   |31951.531    |144.37  |410|10    |149.87   |610         |
|18   |31951.531    |149.37  |769|19    |149.87   |610         |
|19   |31951.538    |149.37  |869|100   |149.87   |610         |
+-----+-------------+--------+---+------+---------+------------+