根据分组数据框的大小过滤熊猫udf

Question

我想将pandas udf函数应用于pyspark.sql.DataFrame，因此代码的蓝图如下所示。

udf_schema = df.withColumn('output', F.lit(0)).schema

@F.pandas_udf(udf_schema, F.PandasUDFType.GROUPED_MAP)
def my_udf(pdf):
    # DO SOMETHING
    df['output'] = 1.2345
    return df

df.groupby(['col1', 'col2']).apply(my_udf)

但是有一些特定的col1＆col2对具有非常低的观察值。因此，我的主要动机是这样的，如果每个col1＆col2对的观察数少于10，则不要将其转换为pandas数据框（因为它很昂贵）并计算实际结果，而只需将1放入输出列即可。有办法解决吗？