присвоение баллов на основе условий в фрейме данных Pyspark

Question

У меня есть фрейм данных в Pyspark, созданный с использованием ниже.

df = sqlContext.createDataFrame(
    [(1,'Y','Y',0,0,0,2,'Y','N','Y','Y'),
     (2,'N','Y',2,1,2,3,'N','Y','Y','N'),
     (3,'Y','N',3,1,0,0,'N','N','N','N'),
     (4,'N','Y',5,0,1,0,'N','N','N','Y'),
     (5,'Y','N',2,2,0,1,'Y','N','N','Y'),
     (6,'Y','Y',0,0,3,6,'Y','N','Y','N'),
     (7,'N','N',1,1,3,4,'N','Y','N','Y'),
     (8,'Y','Y',1,1,2,0,'Y','Y','N','N')
    ],
    ('id', 'compatible', 'product', 'ios', 'pc', 'other', 'devices', 'customer', 'subscriber', 'circle', 'smb')
)

df.show

+---+----------+-------+---+---+-----+-------+--------+----------+------+---+
| id|compatible|product|ios| pc|other|devices|customer|subscriber|circle|smb|
+---+----------+-------+---+---+-----+-------+--------+----------+------+---+
|  1|         Y|      Y|  0|  0|    0|      2|       Y|         N|     Y|  Y|
|  2|         N|      Y|  2|  1|    2|      3|       N|         Y|     Y|  N|
|  3|         Y|      N|  3|  1|    0|      0|       N|         N|     N|  N|
|  4|         N|      Y|  5|  0|    1|      0|       N|         N|     N|  Y|
|  5|         Y|      N|  2|  2|    0|      1|       Y|         N|     N|  Y|
|  6|         Y|      Y|  0|  0|    3|      6|       Y|         N|     Y|  N|
|  7|         N|      N|  1|  1|    3|      4|       N|         Y|     N|  Y|
|  8|         Y|      Y|  1|  1|    2|      0|       Y|         Y|     N|  N|
+---+----------+-------+---+---+-----+-------+--------+----------+------+---+

Теперь из приведенного выше фрейма данных я хочу создать новый столбец на основе некоторых условий.

1) Когда столбец compatible равен Y, тогда

2) если product, customer, subscriber, circle, smb значение столбца = Y присвоить значение = 10 иначе 0

3) если sum of ios, pc, other столбцов больше, чем 4, тогда присвойте значение = 10 else 0

4) если столбец devices больше 4, тогда присвойте значение = 10 else 0

Затем суммируйте все вышеперечисленное values и заполните столбец score в ферме данных pyspark

Вывод, который я хотел бы получить, ниже.

+---+----------+-------+---+---+-----+-------+--------+----------+------+---+-----+
| id|compatible|product|ios| pc|other|devices|customer|subscriber|circle|smb|score|
+---+----------+-------+---+---+-----+-------+--------+----------+------+---+-----+
|  1|         Y|      Y|  0|  0|    0|      2|       Y|         N|     Y|  Y|   50|
|  2|         N|      Y|  2|  1|    2|      3|       N|         Y|     Y|  N|    0|
|  3|         Y|      N|  3|  1|    0|      0|       N|         N|     N|  N|    0|
|  4|         N|      Y|  5|  0|    1|      0|       N|         N|     N|  Y|    0|
|  5|         Y|      N|  2|  2|    0|      1|       Y|         N|     N|  Y|   30|
|  6|         Y|      Y|  0|  0|    3|      6|       Y|         N|     Y|  N|   40| 
|  7|         N|      N|  1|  1|    3|      4|       N|         Y|     N|  Y|    0|
|  8|         Y|      Y|  1|  1|    2|      0|       Y|         Y|     N|  N|   30|
+---+----------+-------+---+---+-----+-------+--------+----------+------+---+-----+

я пробовал как ниже

df1 = df.where(f.col('compatible') == 'Y').\
    withColumn('score', f.when(f.col('product') == 'Y', 10) +
               f.when(f.col('ios') + f.col('pc') + f.col('other') > 4, 10) + f.when(f.col('devices') > 0, 10) +
               f.when(f.col('customer') == 'Y', 10) + f.when(f.col('subscriber') == 'Y', 10) +
               f.when(f.col('circle') == 'Y', 10) + f.when(f.col('smb') == 'Y', 10).otherwise(0))

Вывод, который я получил ниже

+---+----------+-------+---+---+-----+-------+--------+----------+------+---+-----+
| id|compatible|product|ios| pc|other|devices|customer|subscriber|circle|smb|score|
+---+----------+-------+---+---+-----+-------+--------+----------+------+---+-----+
|  1|         Y|      Y|  0|  0|    0|      2|       Y|         N|     Y|  Y| null|
|  3|         Y|      N|  3|  1|    0|      0|       N|         N|     N|  N| null|
|  5|         Y|      N|  2|  2|    0|      1|       Y|         N|     N|  Y| null|
|  6|         Y|      Y|  0|  0|    3|      6|       Y|         N|     Y|  N| null|
|  8|         Y|      Y|  1|  1|    2|      0|       Y|         Y|     N|  N| null|
+---+----------+-------+---+---+-----+-------+--------+----------+------+---+-----+

Как мне добиться того, чего я хочу?

Answer 1

После when/otherwise условия должны удовлетворить ваше требование

df.withColumn('score',
              f.when(df['compatible'] == 'Y',
                     f.when(df['product'] == 'Y', 10).otherwise(0) +
                     f.when(df['customer'] == 'Y', 10).otherwise(0) +
                     f.when(df['subscriber'] == 'Y', 10).otherwise(0) +
                     f.when(df['circle'] == 'Y', 10).otherwise(0) +
                     f.when(df['smb'] == 'Y', 10).otherwise(0) +
                     f.when((df['ios'] + df['pc'] + df['other']) > 4, 10).otherwise(0) +
                     f.when(df['devices'] > 4, 10).otherwise(0)
                ).otherwise(0))\
    .show(truncate=False)

, который должен дать вам

+---+----------+-------+---+---+-----+-------+--------+----------+------+---+-----+
|id |compatible|product|ios|pc |other|devices|customer|subscriber|circle|smb|score|
+---+----------+-------+---+---+-----+-------+--------+----------+------+---+-----+
|1  |Y         |Y      |0  |0  |0    |2      |Y       |N         |Y     |Y  |40   |
|2  |N         |Y      |2  |1  |2    |3      |N       |Y         |Y     |N  |0    |
|3  |Y         |N      |3  |1  |0    |0      |N       |N         |N     |N  |0    |
|4  |N         |Y      |5  |0  |1    |0      |N       |N         |N     |Y  |0    |
|5  |Y         |N      |2  |2  |0    |1      |Y       |N         |N     |Y  |20   |
|6  |Y         |Y      |0  |0  |3    |6      |Y       |N         |Y     |N  |40   |
|7  |N         |N      |1  |1  |3    |4      |N       |Y         |N     |Y  |0    |
|8  |Y         |Y      |1  |1  |2    |0      |Y       |Y         |N     |N  |30   |
+---+----------+-------+---+---+-----+-------+--------+----------+------+---+-----+

вы можете сделать его модульным как

def firstCondition(dataframe):
    return f.when(dataframe['product'] == 'Y', 10).otherwise(0) + \
           f.when(dataframe['customer'] == 'Y', 10).otherwise(0) + \
           f.when(dataframe['subscriber'] == 'Y', 10).otherwise(0) + \
           f.when(dataframe['circle'] == 'Y', 10).otherwise(0) + \
           f.when(dataframe['smb'] == 'Y', 10).otherwise(0)

def secondCondition(dataframe):
    return f.when((dataframe['ios'] + dataframe['pc'] + dataframe['other']) > 4, 10).otherwise(0)

def thirdCondition(dataframe):
    return f.when(dataframe['devices'] > 4, 10).otherwise(0)

df.withColumn('score',
              f.when(df['compatible'] == 'Y', firstCondition(df) + secondCondition(df) + thirdCondition(df)).otherwise(0))\
    .show(truncate=False)

Надеюсь, ответ полезен