Блог

У меня есть такой фрейм данных

  --- ------------------------- ------- 
id  |     sentence            | value |
 --- ------------------------- ------- 
  1 |dogs love fruits         |     50|
  1 |Mary is in the kitchen   |     45|
  2 |John eats fruits         |     56|
 --- ------------------------- ------- 

  

У меня есть список слов, и я хотел бы создать несколько столбцов на основе сопоставления строк, результат будет выглядеть следующим образом

 List = ["fruits", "kitchen"]
 --- ------------------------- ------- ----------- ------------ ------------ ------------- 
id  |     sentence            | value |fruits_bool|fruits_value|kitchen_bool|kitchen_value
 --- ------------------------- ------- ----------- ------------ ------------ ------------- 
  1 |dogs love fruits         |     50|1          |50          |0           |0 
  1 |Mary is in the kitchen   |     45|0          |0           |1           |45
  2 |John eats fruits         |     56|1          |56          |0           |0
 --- ------------------------- ------- ----------- ------------ ------------ ------------- 

  

Я могу легко создавать эти столбцы слово в слово, но я хотел бы делать это автоматически на основе списка.
Вот код, который я пробовал

 groupby("id"). 
    pivot('sentence', List). 
    agg(
       F.count('sentence').alias('bool')
       F.max('value').alias('value')
     ). 
    fillna(0)
  

Я получаю пустой фрейм данных.
Каков наилучший способ сделать это?

Вы можете использовать F.contains и преобразовать его в целое число

 df
.withColumn('fruits_bool',F.col('sentence').contains(F.lit('fruits')).cast('integer'))
.withColumn('fruits_value',F.col('fruits_bool')*F.col('value'))
.withColumn('kitchen_bool',F.col('sentence').contains(F.lit('kitchen')).cast('integer'))
.withColumn('kitchen_value',F.col('kitchen_bool')*F.col('value')).show()
  

вывод :

  --- -------------------- ----- ----------- ------------ ------------ ------------- 
| id|            sentence|value|fruits_bool|fruits_value|kitchen_bool|kitchen_value|
 --- -------------------- ----- ----------- ------------ ------------ ------------- 
|  1|    dogs love fruits|   50|          1|          50|           0|            0|
|  2|Mary is in the ki...|   45|          0|           0|           1|           45|
|  3|    John eats fruits|   56|          1|          56|           0|            0|
 --- -------------------- ----- ----------- ------------ ------------ ------------- 
  

Вероятно, что-то вроде этого :

 from pyspark.sql import functions as F

df.select(
    "*",
    *(
        F.when(F.lit(val).isin(F.col("sentence")), 1).oterhwise(0).alias(f"{val}_bool")
        for val in List
    ),
    *(
        F.when(F.lit(val).isin(F.col("sentence")), F.col("value"))
        .oterhwise(0)
        .alias(f"{val}_value")
        for val in List
    ),
)