PySpark分组行减法最优方案对比：两种实现哪个更优？

按CORP和BRANCH分组，将MARKET为AMERICA行的AMOUNT值减去同组内MARKET为NO_AMERICA与SO_AMERICA行的AMOUNT之和（AMERICA行原值已包含后两者）。

原始DataFrame #

计算示例 #

C0 B1 AMERICA行的计算结果为：4595 - (221 + 545) = 3829

已实现的两种PySpark解决方案 #

通过explain(mode='cost')查看两者成本相同，不确定哪种方案更优。

方案1 #

df_res1 = df_res1.withColumn('america', F.when(F.col('MARKET') == 'AMERICA', F.col('AMOUNT')).otherwise(F.lit(0)))
                 .withColumn('no_so_america', F.when(F.col('MARKET').isin(['NO_AMERICA', 'SO_AMERICA']), F.col('AMOUNT'))
                                               .otherwise(F.lit(0)))

win_perc = W.partitionBy('CORP', 'BRANCH')

df_res1 = df_res1.withColumn('tmp', F.col('america')-F.sum('no_so_america').over(win_perc))
                 .withColumn('AMOUNT',
                            F.when(F.col('tmp') > F.lit(0), F.col('tmp'))
                             .otherwise(F.col('AMOUNT')))
                 .drop('america', 'no_so_america', 'tmp') 

df_res1.explain(mode='cost')

方案2 #

df_sum_no_so = df_res2.where(F.col('MARKET').isin(['NO_AMERICA', 'SO_AMERICA']))
                      .groupBy('CORP', 'BRANCH')
                      .agg(F.sum('AMOUNT').alias('sum_no_so_america'))

df_res2 = df_res2.join(df_sum_no_so, ['CORP', 'BRANCH'], 'left')

df_res2 = df_res2.withColumn('AMOUNT',
                            F.when(F.col('MARKET') == 'AMERICA', F.col('AMOUNT') - F.col('sum_no_so_america'))
                            .otherwise(F.col('AMOUNT')))

df_res2.explain(mode='cost')

内容的提问来源于stack exchange，提问作者Flammy