как получить k-самый большой элемент и индекс в массиве данных pyspark

Question

У меня есть следующий фрейм данных в pyspark:

+------------------------------------------------------------+
|probability                                                 |
+------------------------------------------------------------+
|[0.27047928569511825,0.5312608102025099,0.19825990410237174]|
|[0.06711381377029987,0.8775456658890036,0.05534052034069637]|
|[0.10847074295048188,0.04602848157663474,0.8455007754728833]|
+------------------------------------------------------------+

, и я хочу получить наибольшее, 2-самое большое значение и их индекс:

+-------------------------------------------------------------------------------------------------------------- -----+
|probability                                                 |     largest_1    |index_1|largest_2          |index_2 |
+------------------------------------------------------------|------------------|-------|-------------------|--------+
|[0.27047928569511825,0.5312608102025099,0.19825990410237174]|0.5312608102025099|   1   |0.27047928569511825| 0      |
|[0.06711381377029987,0.8775456658890036,0.05534052034069637]|0.8775456658890036|   1   |0.06711381377029987| 0      |
|[0.10847074295048188,0.04602848157663474,0.8455007754728833]|0.8455007754728833|   2   |0.10847074295048188| 0      |
+--------------------------------------------------------------------------------------------------------------------+

Answer 1

Вот еще один способ использования transform (требуется искра 2.4+) для преобразования array of doubles в array of structs, содержащее значение и индекс каждого элемента в исходном массиве, sort_array ( по убыванию), а затем возьмите первое N:

from pyspark.sql.functions import expr

df.withColumn('d', expr('sort_array(transform(probability, (x,i) -> (x as val, i as idx)), False)')) \
  .selectExpr(
    'probability',
    'd[0].val as largest_1',
    'd[0].idx as index_1',
    'd[1].val as largest_2',
    'd[1].idx as index_2'
).show(truncate=False)
+--------------------------------------------------------------+------------------+-------+-------------------+-------+
|probability                                                   |largest_1         |index_1|largest_2          |index_2|
+--------------------------------------------------------------+------------------+-------+-------------------+-------+
|[0.27047928569511825, 0.5312608102025099, 0.19825990410237174]|0.5312608102025099|1      |0.27047928569511825|0      |
|[0.06711381377029987, 0.8775456658890036, 0.05534052034069637]|0.8775456658890036|1      |0.06711381377029987|0      |
|[0.10847074295048188, 0.04602848157663474, 0.8455007754728833]|0.8455007754728833|2      |0.10847074295048188|0      |
+--------------------------------------------------------------+------------------+-------+-------------------+-------+

Answer 2

С Spark-2.4 +

Вы можете использовать array_sort and array_position встроенные функции для этого случая.

Example:

df=spark.sql("select array(0.27047928569511825,0.5312608102025099,0.19825990410237174) probability union select array(0.06711381377029987,0.8775456658890036,0.05534052034069637) prbability union select array(0.10847074295048188,0.04602848157663474,0.8455007754728833) probability")

#DataFrame[probability: array<decimal(17,17)>]

#sample data
df.show(10,False)
#+---------------------------------------------------------------+
#|probability                                                    |
#+---------------------------------------------------------------+
#|[0.06711381377029987, 0.87754566588900360, 0.05534052034069637]|
#|[0.27047928569511825, 0.53126081020250990, 0.19825990410237174]|
#|[0.10847074295048188, 0.04602848157663474, 0.84550077547288330]|
#+---------------------------------------------------------------+      


df.withColumn("sort_arr",array_sort(col("probability"))).\
withColumn("largest_1",element_at(col("sort_arr"),-1)).\
withColumn("largest_2",element_at(col("sort_arr"),-2)).\
selectExpr("*","array_position(probability,largest_1) -1 index_1","array_position(probability,largest_2) -1 index_2").\
drop("sort_arr").\
show(10,False)


#+---------------------------------------------------------------+-------------------+-------------------+-------+-------+
#|probability                                                    |largest_1          |largest_2          |index_1|index_2|
#+---------------------------------------------------------------+-------------------+-------------------+-------+-------+
#|[0.06711381377029987, 0.87754566588900360, 0.05534052034069637]|0.87754566588900360|0.06711381377029987|1      |0      |
#|[0.27047928569511825, 0.53126081020250990, 0.19825990410237174]|0.53126081020250990|0.27047928569511825|1      |0      |
#|[0.10847074295048188, 0.04602848157663474, 0.84550077547288330]|0.84550077547288330|0.10847074295048188|2      |0      |
#+---------------------------------------------------------------+-------------------+-------------------+-------+-------+