PySpark сворачивает информационный фрейм

Question

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

>>> my_df.show(3)
+------------+---------+-------+--------------+
|     user_id|  address|   type|count| country|
+------------+---------+-------+-----+--------+
|      ABC123|  yyy,USA| animal|    2|     USA|
|      ABC123|  xxx,USA| animal|    3|     USA|
|      qwerty|  55A,AUS|  human|    3|     AUS|
|      ABC123|  zzz,RSA| animal|    4|     RSA|
+------------+---------+-------+--------------+

Как мне свернуть этот фрейм данных, чтобы получить следующий результат -

>>> new_df.show(3)
+------------+---------+-------+--------------+
|     user_id|  address|   type|count| country|
+------------+---------+-------+-----+--------+
|      qwerty|  55A,AUS|  human|    3|     AUS|
|      ABC123|  xxx,USA| animal|    5|     USA|
+------------+---------+-------+--------------+ 

Для данного user_id:

1. Получите country с наибольшей суммой отсчетов
2. Для country, полученного на шаге 1, получите address с наибольшим подсчетом

Полагаю, мне придется разделить my_df на 2 разных кадра данных и получить country и address отдельно.Но я точно не знаю синтаксис для этого.Ваша помощь приветствуется.Спасибо.

Answer 1

Я имел в виду что-то вроде этого:

>>> import pandas as pd
>>> from pyspark.sql.functions import *
>>> from pyspark.sql.window import *
>>> from pyspark.sql import SparkSession

>>> spark = SparkSession.builder.appName('abc').getOrCreate()

>>> data = {"user_id": ["ABC123", "ABC123", "qwerty", "ABC123"], "address": ["yyy,USA", "xxx,USA", "55A,AUS", "zzz,RSA"], "type": ["animal", "animal", "human", "animal"], "count": [2,3,3,4], "country": ["USA", "USA", "AUS", "RSA"]}

>>> df = pd.DataFrame(data=data)

>>> df_pyspark = spark.createDataFrame(df)

>>> w = Window().partitionBy("user_id", "country").orderBy((col("count").desc()))

>>> w2 = Window().partitionBy("user_id").orderBy(col("sum_country").desc())

>>> df_pyspark.select("user_id", "address", "type", "count", "country", sum("count").over(w).alias("sum_country")).select("user_id", first("country").over(w2).alias("top_country"), first("address").over(w).alias("top_address"), "country").where(col("top_country")==col("country")).distinct().show()
+-------+-----------+-----------+-------+
|user_id|top_country|top_address|country|
+-------+-----------+-----------+-------+
| qwerty|        AUS|    55A,AUS|    AUS|
| ABC123|        USA|    xxx,USA|    USA|
+-------+-----------+-----------+-------+

Вы можете добавить тип, число и т. Д. В зависимости от того, какую логику вы хотите использовать для этого - вы можете сделать то же самое, что и для top_address (т.е. first функция), или вы можете groupBy и agg