Добавьте новый столбец в фрейм данных на основе значений в других 2 столбцах (требуется Pyspark)

Question

введите описание изображения здесь

Я хочу добавить столбец с именем «Joint_pred_x» (x = 0,1,2) на основе двух значений в «nb_pred_x» и «svm_pred_x», добавить 0, если nb = 1, svm = 1; добавить 1, если nb = 1, svm = 0; добавьте 2, если nb = 0, svm = 1; добавьте 3, если nb = 0, svm = 0. Я думаю, что withcolumn может выполнить эту работу, но меня смущает условный logi c. Решение должно быть только pyspark, заранее спасибо

Answer 1

Вы можете использовать оператор case.

+---------+---------+---------+----------+----------+----------+
|nb_pred_0|nb_pred_1|nb_pred_2|svm_pred_0|svm_pred_1|svm_pred_2|
+---------+---------+---------+----------+----------+----------+
|0.0      |1.0      |0.0      |0.0       |1.0       |0.0       |
+---------+---------+---------+----------+----------+----------+


from pyspark.sql.functions import expr

for i in range(0, 3):
    
    index = str(i)
    
    df = df.withColumn('joint_pred_' + index, expr(f'''
            CASE 
                WHEN {p1}_pred_{index} == 1 and {p2}_pred_{index} == 1 THEN 0
                WHEN {p1}_pred_{index} == 1 and {p2}_pred_{index} == 0 THEN 1
                WHEN {p1}_pred_{index} == 0 and {p2}_pred_{index} == 1 THEN 2
                WHEN {p1}_pred_{index} == 0 and {p2}_pred_{index} == 0 THEN 3
            END
        '''))

df.show(10, False)

+---------+---------+---------+----------+----------+----------+------------+------------+------------+
|nb_pred_0|nb_pred_1|nb_pred_2|svm_pred_0|svm_pred_1|svm_pred_2|joint_pred_0|joint_pred_1|joint_pred_2|
+---------+---------+---------+----------+----------+----------+------------+------------+------------+
|0.0      |1.0      |0.0      |0.0       |1.0       |0.0       |3           |0           |3           |
+---------+---------+---------+----------+----------+----------+------------+------------+------------+