PySpark выбрать топ N из нескольких столбцов

Question

Фрейм данных Pyspark

Tick | Long Position  | Short Position
ADM  | $25.5          | $12.2
INT  | $300           | $120
CITI | $-12           | $9.1

Мне нужно выбрать топ-3 из «Длинной позиции» и «Короткой позиции». Из данных выше, что я хочу, будет что-то вроде ниже. Как это сделать?

Tick | Position | PnL
INT  | Long     | $300
INT  | Short    | $120
ADM  | Long     | $25.5

Answer 1

Моя первоначальная идея состояла в том, чтобы объединить два столбца и создать один фрейм данных, который находится в структуре вашего вывода, а затем взять верхний 3. Я только что попробовал это в Scala, он должен работать и с pyspark.

val y = someDF.select("Tick", "Long").withColumn("type", lit("Long")).withColumnRenamed("Long", "PnL")
val z = someDF.select("Tick", "Short").withColumn("type", lit("Short")).withColumnRenamed("Short", "PnL")

println(y.union(z).sort(desc("PnL")).take(3))

Поскольку вы все делаете с помощью API API, я думаю, что это также должно дать вам хорошую производительность. Дайте мне знать, как это происходит!