python - 如何查找 Pyspark 中列中值最大的行名称

Question

我有一个这样的表:

+-------+-----+------+------+
|user_id|apple|good banana|carrot|
+-------+-----+------+------+
| user_0|    0|     3|     1|
| user_1|    1|     0|     2|
| user_2|    5|     1|     2|
+-------+-----+------+------+

在这里，对于每种水果，我想获取购买最多商品的顾客列表。所需的输出如下:

                max_user max_count
apple           [user_2]         5
banana          [user_0]         3
carrot  [user_1, user_2]         2

MWE

import numpy as np
import pandas as pd
import pyspark
from pyspark.sql import functions as F

spark = pyspark.sql.SparkSession.builder.getOrCreate()
sc = spark.sparkContext
sqlContext = pyspark.SQLContext(sc)


# pandas dataframe
pdf = pd.DataFrame({'user_id': ['user_0','user_1','user_2'],
                   'apple': [0,1,5],
                   'good banana': [3,0,1],
                   'carrot': [1,2,2]})


# spark dataframe
df = sqlContext.createDataFrame(pdf)
# df.show()


df.createOrReplaceTempView("grocery")
spark.sql('select * from grocery').show()

问题 1

如何使用Pyspark获得所需的输出？

问题 2

如何使用Pyspark sql获取所需的输出？

引用文献

我已经做了一些研究并搜索了多个页面。到目前为止，我已经想出了一个接近的答案，但它需要转置表，而这里我的表是正常的。另外，我正在学习多种方法，例如Spark方法和SQL方法。

• how to get the name of column with maximum value in pyspark dataframe

Answer 1

Pyspark 解决方案。与 pandas 解决方案类似，首先使用 stack 融化数据帧，然后使用rank 过滤最大计数的行，按 Fruit 分组，并使用 collect_list 获取用户列表。

from pyspark.sql import functions as F, Window

df2 = df.selectExpr(
    'user_id',
    'stack(3, ' + ', '.join(["'%s', %s" % (c, c) for c in df.columns[1:]]) + ') as (fruit, items)'
).withColumn(
    'rn',
    F.rank().over(Window.partitionBy('fruit').orderBy(F.desc('items')))
).filter('rn = 1').groupBy('fruit').agg(
    F.collect_list('user_id').alias('max_user'),
    F.max('items').alias('max_count')
)

df2.show()
+------+----------------+---------+
| fruit|        max_user|max_count|
+------+----------------+---------+
| apple|        [user_2]|        5|
|banana|        [user_0]|        3|
|carrot|[user_1, user_2]|        2|
+------+----------------+---------+

对于 Spark SQL:

df.createOrReplaceTempView("grocery")

df2 = spark.sql("""
    select
        fruit,
        collect_list(user_id) as max_user,
        max(items) as max_count
    from (
        select *,
            rank() over (partition by fruit order by items desc) as rn
        from (
            select
                user_id,
                stack(3, 'apple', apple, 'banana', banana, 'carrot', carrot) as (fruit, items)
            from grocery
        )
    )
    where rn = 1 group by fruit
""")

df2.show()
+------+----------------+---------+
| fruit|        max_user|max_count|
+------+----------------+---------+
| apple|        [user_2]|        5|
|banana|        [user_0]|        3|
|carrot|[user_1, user_2]|        2|
+------+----------------+---------+