I have a very wide dataframe with label columns. I want to run a logistic regression for each column independenly. I'm trying to find the most efficient way to run this in parallel.
+----------+--------+--------+--------+-----+------------+
| features | label1 | label2 | label3 | ... | label30000 |
+----------+--------+--------+--------+-----+------------+
My initial thought was to use ThreadPoolExecutor
, get result for each column, and join:
extract_prob = udf(lambda x: float(x[1]), FloatType())
def lr_for_column(argm):
col_name = argm[0]
test_res = argm[1]
lr = LogisticRegression(featuresCol="features", labelCol=col_name, regParam=0.1)
lrModel = lr.fit(tfidf)
res = lrModel.transform(test_tfidf)
test_res = test_res.join(res.select('id', 'probability'), on="id")
test_res = test_res.withColumn(col_name, extract_prob('probability')).drop("probability")
return test_res.select('id', col_name)
with futures.ThreadPoolExecutor(max_workers=100) as executor:
future_results = [executor.submit(lr_for_column, [colname, test_res]) for colname in list_of_label_columns]
futures.wait(future_results)
for future in future_results:
test_res = test_res.join(future.result(), on="id")
but this method is not very performant. Is there a faster way to do this?