Pyspark，編寫循環根據不同條件創建多個新列

Question

假設我有一個 Pyspark DataFrame 具有以下列：

用戶、分數、國家、風險/安全、payment_id

我列出了閾值：[10,20,30]

現在我想為每個閾值創建一個新列：

1. 在所有付款中得分高於閾值的風險付款的百分比（風險和安全）
2. 在所有用戶中至少有一個分數高於閾值的有風險的不同用戶的百分比（有風險的和安全的）

兩者都應按國家/地區分組。

結果應該是這樣的：

Country | % payments thresh 10 | % users thresh 10 | % payments thresh 20 ... 
A
B
C

我能夠使其與外部 for 循環一起工作，但我希望它全部在一個 dataframe 中。

thresholds = [10, 20, 30]


for thresh in thresholds:

    
df = (df
     .select('country', 'risk/safe', 'user', 'payment')
     .where(F.col('risk\safe') == 'risk')
     .groupBy('country').agg(F.sum(F.when(
         (F.col('score') >= thresh),1 
           )) / F.count('country').alias('% payments'))

Answer 1

在agg()中使用列表推導。

pay_aggs = [(func.sum((func.col('score')>=thresh).cast('int'))/func.count('country')).alias('% pay '+str(thresh)) for thresh in thresholds]
user_aggs = [(func.countDistinct(func.when(func.col('score')>=thresh, func.col('user')))/func.countDistinct('user')).alias('% user '+str(thresh)) for thresh in thresholds]

df. \
    select('country', 'risk/safe', 'user', 'payment'). \
    where(func.col('risk\safe') == 'risk'). \
    groupBy('country'). \
    agg(*pay_aggs, *user_aggs)

pay_aggs列表將生成以下聚合（您可以輕松打印列表）

# [Column<'(sum(CAST((score >= 10) AS INT)) / count(country)) AS `% pay 10`'>,
#  Column<'(sum(CAST((score >= 20) AS INT)) / count(country)) AS `% pay 20`'>,
#  Column<'(sum(CAST((score >= 30) AS INT)) / count(country)) AS `% pay 30`'>]