PySpark:获取数据帧中每行的前k列

Question

我有一个数据框,每个联系人的每个优惠都有分数.我想创建一个新的数据帧,其中包含每个联系人的前3个优惠.

输入数据框是这样的:

=======================================================================
| 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    |
-----------------------------------------------------------------------

我想将它转换为数据帧,如下所示:

===============================================================
| 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

你需要一些进口:

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

使用问题中显示的数据:

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)])

你可以组装和排序一个数组structs:

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

并且select:

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

这应该给出以下结果:

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

根据您的需要重新命名列,您就可以开始了.