Блог

У меня есть фрейм данных PySpark-

 df1 = spark.createDataFrame([
    ("u1", 1),
    ("u1", 2),
    ("u2", 1),
    ("u2", 1),
    ("u2", 1),
    ("u3", 3),
    ],
    ['user_id', 'var1'])

print(df1.printSchema())
df1.show(truncate=False)
 

Вывод-

 root
 |-- user_id: string (nullable = true)
 |-- var1: long (nullable = true)

None
 ------- ---- 
|user_id|var1|
 ------- ---- 
|u1     |1   |
|u1     |2   |
|u2     |1   |
|u2     |1   |
|u2     |1   |
|u3     |3   |
 ------- ---- 
 

Теперь я хочу сгруппировать всех уникальных пользователей и показать количество уникальных переменных для них в новом столбце. Желаемый результат будет выглядеть так-

  ------- --------------- 
|user_id|num_unique_var1|
 ------- --------------- 
|u1     |2              |
|u2     |1              |
|u3     |1              |
 ------- --------------- 
 

Я могу использовать collect_set и создать udf, чтобы найти длину набора. Но я думаю, что должен быть лучший способ сделать это.
Как мне добиться этого в одной строке кода?

 df1.groupBy('user_id').agg(F.countDistinct('var1').alias('num')).show()
 

countDistinct — это именно то, что мне было нужно.

Вывод-

  ------- --- 
|user_id|num|
 ------- --- 
|     u3|  1|
|     u2|  1|
|     u1|  2|
 ------- --- 
 

countDistinct это, безусловно, лучший способ сделать это, но для полноты картины то, что вы сказали в своем вопросе, также возможно без использования UDF. Вы можете использовать size , чтобы получить длину collect_set :

 df1.groupBy('user_id').agg(F.size(F.collect_set('var1')).alias('num'))
 

это полезно, если вы хотите использовать его в оконной функции, поскольку countDistinct он не поддерживается в оконной функции.

например ,

 df1.withColumn('num', F.countDistinct('var1').over(Window.partitionBy('user_id')))
 

потерпел бы неудачу, но

 df1.withColumn('num', F.size(F.collect_set('var1')).over(Window.partitionBy('user_id')))
 

это сработало бы.