Выполнение PySpark замедляется на каждой итерации

Question

Я очень новичок в PySpark.Я написал фрагмент кода для вменения 7 переменных-предикторов.Я использую столбец описания, преобразованный в tfidf, для обучения модели NB и прогнозирования NaN.

to_predict_cols = ['col1', 'col2', ... 'col7']
current_x = ['features_tfidf']
for iter_no, to_predict_col in enumerate(to_predict_cols):
    temp_data_imputed_level1 = data_imputed_level1.select(["_c0"] + current_x + [to_predict_col])

    # create test
    test = temp_data_imputed_level1.where(col(to_predict_col).isNull()).select(["_c0"] + current_x)

    # create train
    train = temp_data_imputed_level1.dropna()

    # create label indexing of the response
    label_stringIndexer = StringIndexer(inputCol = to_predict_col, outputCol = "label")
    stringIndexer_model = label_stringIndexer.fit(train)
    train = stringIndexer_model.transform(train)
    print("label encoding.. [OK]")

    # fit the NB classifier
    nb = NaiveBayes(smoothing=1.0, modelType="multinomial")
    model = nb.fit(train)
    print("model building.. [OK]")

    # predict output
    predictions = model.transform(test)

    # inverse label encoding of the predictions
    labelReverse = IndexToString(inputCol='prediction', outputCol='labeled_prediction', labels=stringIndexer_model.labels)
    predictions = labelReverse.transform(predictions)
    print("model predictions.. [OK]")

Я использую некоторые объединения, чтобы сохранить прогнозы, которые я не включил.Но поскольку цикл идет в 3-4-й итерации, он становится слишком медленным с самого начала, т. Е. Объединения не влияют на скорость.Также данные кэшируются перед выполнением цикла.

Я попытался sc._jvm.System.gc() и удалил все переменные, инициированные вручную.Но все же не повезло.

Я на сервере с 256 ГБ оперативной памяти и 48 ядрами.В настоящее время я не связан ни с какими кластерами.Размер данных составляет 1,5 ГБ.

Итак, есть идеи / предложения о том, чего мне не хватает или что нужно сделать, чтобы ускорить процесс?Любая помощь приветствуется.Заранее спасибо.