Append aggregate stats from loop in Pyspark table - pyspark

Question

I have the following Pyspark code. In each iteration in the loop, I filter out all rows with a particular string in column H. Then I calculate some aggregate stats on column G (results in 3 values). I want to save all the aggregate counts in a single table (rows: CM, NCM, FP; and columns: POP, POP N, POP SN, POP QP).
from pyspark.sql import SparkSession
import pandas as pd
import numpy as np
import pyspark.sql.functions as F
spark = SparkSession.builder.master("local").appName("Word Count").config("spark.some.config.option", "some-value").getOrCreate()
np.random.seed(0)
rows = 1000
df_pandas = pd.DataFrame({ 'A' : 1.,
'B' : pd.Timestamp('20130102'),
'C' : pd.Series(1,index=list(range(rows)),dtype='float32'),
'D' : np.array([3] * rows,dtype='int32'),
'E' : pd.Categorical(np.random.choice(["test","train","frog", "chicken"], rows)),
'F' : 'foo' ,
'G' : np.random.choice(['CM', 'NCM', 'FP'], rows),
'H' : np.random.choice(['POP', 'POP N', 'POP SN', 'POP QP'], rows)})
df_spark = spark.createDataFrame(df_pandas)
blocks = ['POP', 'POP N', 'POP SN', 'POP QP']
for block in blocks:
df_spark_trimmed = df_spark.filter(~F.col('H').isin(block))
counts = df_spark_trimmed.groupby('G').count()
counts.show()

Answer 1

Use join:
import pyspark.sql.functions as F
count_by_g = df_spark.groupBy('G').agg(F.count('*').alias('CountByG'))
count_by_gh = df_spark.groupBy(['G', 'H']).agg(F.count('*').alias('CountByGH'))
count_by_g.join(count_by_gh, ['G']).selectExpr(
'G', 'H', 'CountByG - CountByGH as count'
).groupBy('G').pivot('H').agg(F.max('count').alias('count')).show()
+---+---+-----+------+------+
| G|POP|POP N|POP QP|POP SN|
+---+---+-----+------+------+
| CM|256| 260| 245| 250|
|NCM|265| 254| 248| 262|
| FP|246| 236| 239| 239|
+---+---+-----+------+------+
Or another solution with window functions:
df_spark.groupBy(['G', 'H']).count().selectExpr(
'G', 'H', 'sum(count) over (partition by G) - count as count'
).groupBy('G').pivot('H').agg(F.max('count').alias('count')).show()
+---+---+-----+------+------+
| G|POP|POP N|POP QP|POP SN|
+---+---+-----+------+------+
| CM|256| 260| 245| 250|
|NCM|265| 254| 248| 262|
| FP|246| 236| 239| 239|
+---+---+-----+------+------+