PySpark: get top column k for each row in dataframe

Question

I have a dataframe with ratings for each sentence for each contact. I want to create a new dataframe from this, which has 3 best suggestions for each contact.

The input data frame looks something like this:

=======================================================================
| contact | offer 1 | offer 2 | offer 3 | offer 4 | offer 5 | offer 6 |
=======================================================================
| name 1  | 0       | 3       | 1       |   2     |    1    |    6    |
-----------------------------------------------------------------------
| name 2  | 1       | 7       | 2       |   9     |    5    |    3    |
-----------------------------------------------------------------------

I want to convert it to dataframe as follows:

===============================================================
| contact | best offer | second best offer | third best offer |
===============================================================
| name 1  | offer 6    | offer 2           | offer 4          |
---------------------------------------------------------------
| name 1  | offer 4    | offer 2           | offer 5          |
---------------------------------------------------------------

Answer 1

You will need some imports:

from pyspark.sql.functions import array, col, lit, sort_array, struct

With data as shown in the question:

df = sc.parallelize([
    ("name 1", 0, 3, 1, 2, 1, 6),
    ("name 2", 1, 7, 2, 9, 5, 3),
]).toDF(["contact"] + ["offer_{}".format(i) for i in range(1, 7)])

you can collect and sort the array structs:

offers = sort_array(array(*[
    struct(col(c).alias("v"), lit(c).alias("k")) for c in df.columns[1:]
]), asc=False)

and select:

df.select(
    ["contact"] + [offers[i]["k"].alias("_{}".format(i)) for i in [0, 1, 2]])

which should give the following result:

+-------+-------+-------+-------+
|contact|     _0|     _1|     _2|
+-------+-------+-------+-------+
| name 1|offer_6|offer_2|offer_4|
| name 2|offer_4|offer_2|offer_5|
+-------+-------+-------+-------+

, .