Обучение и применение нескольких изотонических регрессий параллельно

Question

Я пишу некоторый код, который изотонически сглаживает количество кривых, которые у меня есть в кадре данных, используя пакет изотонической регрессии pyspark.

Я написал отвратительно медленный цикл, который делает то, что я хочу, но я не на той стадии, когда это огромное узкое место и нуждается в улучшении.Код выглядит следующим образом:

l = [('x', 'a', 1, 2.1),
     ('x', 'a', 2, 3.0),
     ('x', 'a', 3, 3.6),
     ('x', 'a', 4, 3.4),
     ('x', 'a', 5, 4.0),
     ('x', 'a', 6, 5.0),
     ('x', 'b', 1, 1.1),
     ('x', 'b', 2, 1.5),
     ('x', 'b', 3, 1.4),
     ('x', 'b', 4, 1.6),
     ('x', 'b', 5, 1.8),
     ('x', 'b', 6, 2.0),
     ('y', 'a', 1, 10.0),
     ('y', 'a', 2, 10.1),
     ('y', 'a', 3, 11.2),
     ('y', 'a', 4, 10.2),
     ('y', 'a', 5, 12.0),
     ('y', 'a', 6, 13.0)]
df = sqlContext.createDataFrame(l, ['category_1', 'category_2', 'markdown', 'uplift'])
df = df.withColumn('markdown', df["markdown"].cast('double'))

isomodel = IsotonicRegression(featuresCol='markdown', labelCol='uplift', predictionCol='smoothed_uplift')
unique_refs = df.select('category_1', 'category_2').distinct()
unique_refs = unique_refs.toPandas()

for row in range(len(unique_refs)):
        cat_1 = unique_refs.loc[row,'category_1']
        cat_2 = unique_refs.loc[row,'category_2']
        print("processing ", cat_1, cat_2)
        temp_df = df.filter((df['category_1']==cat_1)&\
                              (df['category_2']==cat_2))
        print(temp_df.show())
        ir_model = isomodel.fit(temp_df)
        if row==0:
            smoothed_df = ir_model.transform(temp_df)
        else:
            smoothed_df_temp = ir_model.transform(temp_df)
            smoothed_df = smoothed_df.union(smoothed_df_temp)

Результат, который я получаю (и все еще хочу получить), выглядит следующим образом:

+----------+----------+--------+------+---------------+
|category_1|category_2|markdown|uplift|smoothed_uplift|
+----------+----------+--------+------+---------------+
|         x|         b|     1.0|   1.1|            1.1|
|         x|         b|     2.0|   1.5|           1.45|
|         x|         b|     3.0|   1.4|           1.45|
|         x|         b|     4.0|   1.6|            1.6|
|         x|         b|     5.0|   1.8|            1.8|
|         x|         b|     6.0|   2.0|            2.0|
|         y|         a|     1.0|  10.0|           10.0|
|         y|         a|     2.0|  10.1|           10.1|
|         y|         a|     3.0|  11.2|           10.7|
|         y|         a|     4.0|  10.2|           10.7|
|         y|         a|     5.0|  12.0|           12.0|
|         y|         a|     6.0|  13.0|           13.0|
|         x|         a|     1.0|   2.1|            2.1|
|         x|         a|     2.0|   3.0|            3.0|
|         x|         a|     3.0|   3.6|            3.5|
|         x|         a|     4.0|   3.4|            3.5|
|         x|         a|     5.0|   4.0|            4.0|
|         x|         a|     6.0|   5.0|            5.0|
+----------+----------+--------+------+---------------+

Я знаю, что должен быть хороший способсделать это в искре, но я не знаю, с чего начать!