在 Pyspark 的 groupby 上创建一个新的计算列

时间:2021-02-25 15:06:12

标签: python apache-spark pyspark apache-spark-sql

我在 Pyspark 中有以下数据框,它已经在“帐户名”列的 groupby 中。

accountname |   namespace   |   cost    |   cost_to_pay
account001  |   ns1         |   93      |   9
account001  |   Transversal |   93      |   25
account002  |   ns2         |   50      |   27
account002  |   Transversal |   50      |   12

我需要一个新列,它是 "cost" - "cost_to_pay" where "namespace" == "Transversal",我需要这个结果在新列的所有字段中,如下所示:

accountname |   namespace   |   cost    |   cost_to_pay |   new_column1                                         
account001  |   ns1         |   93      |   9           |   68                    
account001  |   Transversal |   93      |   25          |   68
account002  |   ns2         |   50      |   27          |   38
account002  |   Transversal |   50      |   12          |   38

68 是从 account001 中减去 groupby 的 93 - 25 的结果。 而38就是account002减去50-12的结果。

知道如何实现这一目标吗?

2 个答案:

答案 0 :(得分:2)

您可以使用最大掩码差异来获取每个帐户名的差异:

from pyspark.sql import functions as F, Window

df2 = df.withColumn(
    'new_column1',
    F.max(
        F.when(
            F.col('namespace') == 'Transversal',
            F.col('cost') - F.col('cost_to_pay')
        )
    ).over(Window.partitionBy('accountname'))
)

df2.show()
+-----------+-----------+----+-----------+-----------+
|accountname|  namespace|cost|cost_to_pay|new_column1|
+-----------+-----------+----+-----------+-----------+
| account001|        ns1|  93|          9|         68|
| account001|Transversal|  93|         25|         68|
| account002|        ns2|  50|         27|         38|
| account002|Transversal|  50|         12|         38|
+-----------+-----------+----+-----------+-----------+

答案 1 :(得分:1)

如果 dfgroupby 之后的数据框,您可以使用以下方法找到 df_temp

df_temp = df.filter(F.col('namespace')=='Transversal')
df_temp = df_temp.withcolumn('new_column1', F.col('cost') - F.col('cost_to_pay'))
df_temp = df_temp.select('accountname', 'new_column1') ## keep only relevant columns
## you might want to have some extra checks, like droping duplicates, etc

## and finally join df_temp with you main dataframe df
df = df.join(df_temp, on='accountname', how='left')
df = df.na.fill({'accountname':'some predefined value, like 0}) ## if you wish to fill nulls
相关问题