使用窗口函数：

import pyspark.sql.functions as F
from pyspark.sql import Window

df = spark.createDataFrame(
    [
    ('33000120304','629138','1','20210507')
    ,('33000009746','629138','0','20210129')
    ,('50000441489','629138','1','20210511')
    ,('22000020680','187367','0','20220210')
    ,('41419714737','187367','1','20210628')
    ,('50000577980','176129','0','20220225')
    ,('50000001607','633782','1','20210419')
    ,('50000001607','633782','1','20210419')
    ,('50000001608','633782','1','20210419')
    ],
    ['SKU','Store','Flag','Date']
).withColumn('Date', F.to_date(F.col('Date'), 'yyyyMMdd'))

w = Window.partitionBy("Store").orderBy(F.desc('Date'))

df_res = df\
    .withColumn('max_date', F.max(F.col('Date')).over(w))\
    .filter((F.col('Date') == F.col('max_date')))\
    .filter(F.col('Flag')==1)\
    .groupBy('Store').agg(F.countDistinct("SKU").alias('count_distinct_sku'))

df\
.select('Store')\
.dropDuplicates()\
.join(df_res, on='Store', how='left')\
.fillna(0)\
.show()

# +------+------------------+
# | Store|count_distinct_sku|
# +------+------------------+
# |629138|                 1|
# |187367|                 0|
# |176129|                 0|
# |633782|                 2|
# +------+------------------+