PySpark：基于列值对数据框进行分区并将结果数据框存储在列表中

Question

我有一个带有4列的pyspark数据框：城市，季节，天气变量，变量值。我必须将框架划分为不同的城市，季节和天气变量组合。接下来，我将在这些分区上应用k均值。

我正在使用以下代码来创建分区：

a = df_in.select('city', 'season', 'variable').distinct().toPandas().as_matrix()
dfArray = [df_in.filter("city = '{}' and season = '{}' and variable = '{}'".format(x[0], x[1], x[2])) for x in a]

但是问题在于，该过程非常缓慢，因为它反复使用过滤器。有人可以提出一种使流程更有效的方法的建议。诸如groupby之类的方法可能会有所帮助。

请在下面找到要转换为pyspark的代码的python版本：

from sklearn.cluster import KMeans

def k_means_on_partition(group):

    v = group['variable_value']
    kmeans = KMeans(n_clusters = 7)

    kmeans.fit(v.values.reshape(-1, 1)) 
    group['cluster'] = kmeans.labels_

    return group


df_out = df_in.groupby(['city', 'season', 'variable']).apply(k_means_on_partition)        

此外，对于pyspark，我将使用pyspark.ml而不是sklearn。