PySpark с для цикла

Question

Мне нужно перебрать кадр данных в PySpark 16 раз.

Без PySpark задача выполняется быстро, но с помощьюPySpark выдается сообщение об ошибке, возвращаемое при запуске функции get_unidades_parecidas:

WARN  TaskSetManager:66 - Stage 1341 contains a task of very large size (1000 KB). The maximum recommended task size is 100 KB.`

Я думаю, что ошибка вызвана объявлением цикла for вфункция get_unidades_mesma_rua, в строке 8.

Схема df:

StructType(List(StructField(code,LongType,true),StructField(tipo,StringType,true),StructField(sale_value,LongType,true),StructField(address_district,StringType,true),StructField(address_number,StringType,true),StructField(areaPrivate,LongType,true),StructField(address_place,StringType,true),StructField(bedrooms,LongType,true),StructField(parkingSpaces,LongType,true),StructField(sale_per_area,DoubleType,true)))

target - это строка из df

def get_unidades_parecidas(df, target):        
    backup_maior_amostra = SQLcontext.createDataFrame(sc.emptyRDD(), StructType([]))
    lista_de_possilidades_quarto_vaga = get_lista_possiblidades_quarto_vaga(target.bedrooms, target.parkingSpaces)

    for possibilidade in lista_de_possilidades_quarto_vaga:
        controle_de_area = 0.1
        quarto_minimo, quarto_maximo, vaga_minima, vaga_maxima = possibilidade
        while controle_de_area <= 0.3:
            area_minima = target.areaPrivate * (1 - controle_de_area)
            area_maxima = target.areaPrivate * (1 + controle_de_area)
            df_ = df.filter((df.areaPrivate.between(area_minima, area_maxima)) &
                   (df.bedrooms.between(quarto_minimo, quarto_maximo)) &
                   (df.parkingSpaces.between(vaga_minima, vaga_maxima)))
            if df_.count() >= 10:
                return [[area_minima, area_maxima, quarto_minimo, quarto_maximo, vaga_minima, vaga_maxima], df_]
            else:
                controle_de_area += 0.1
                if df_.count() > backup_maior_amostra.count():
                    backup_maior_amostra = df_


    return [[area_minima, area_maxima, quarto_minimo, quarto_maximo, vaga_minima, vaga_maxima],
            backup_maior_amostra] # nenhuma amostra com mais de 10 elementos



def get_lista_possiblidades_quarto_vaga(quarto, vaga):

    quarto_menos_1 = quarto - 1
    quarto_mais_1 =  quarto + 1
    vaga_menos_1 = vaga - 1
    vaga_mais_1 = vaga + 1

    lista_de_possibilidades_quarto_vaga = [(quarto, quarto, vaga, vaga), # a melhor possibilidade

                                           (quarto, quarto_mais_1, vaga, vaga),
                                           (quarto_menos_1, quarto, vaga, vaga),
                                           (quarto, quarto, vaga_menos_1, vaga),
                                           (quarto, quarto, vaga, vaga_mais_1),

                                           (quarto_menos_1, quarto_mais_1, vaga, vaga),
                                           (quarto, quarto, vaga_menos_1, vaga_mais_1),
                                           (quarto_menos_1, quarto, vaga_menos_1, vaga),
                                           (quarto_menos_1, quarto, vaga, vaga_mais_1),
                                           (quarto, quarto_mais_1, vaga_menos_1, vaga),
                                           (quarto, quarto_mais_1, vaga, vaga_mais_1),

                                           (quarto, quarto_mais_1, vaga_menos_1, vaga_mais_1),
                                           (quarto_menos_1, quarto, vaga_menos_1, vaga_mais_1),                                             
                                           (quarto_menos_1, quarto_mais_1, vaga, vaga_mais_1),
                                           (quarto_menos_1, quarto_mais_1, vaga_menos_1, vaga),

                                           (quarto_menos_1, quarto_mais_1, vaga_menos_1, vaga_mais_1)] # pior possiblidade
    return lista_de_possibilidades_quarto_vaga