Блог

Я хочу обработать каждую строку фрейма данных. Здесь столбец feat содержит много элементов с форматом idx:value . Я хочу сохранить idx то, что мне нужно.

Например, я хочу сохранить idx=1 или idx=5 .

 df = spark.createDataFrame([("u1","1:a 2:k 5:c 6:i"),("u2","2:k 4:p 5:b 6:k")],["id","feat"])
  

Input:

  --- --------------- 
| id|           feat|
 --- --------------- 
| u1|1:a 2:k 5:c 6:i|
| u2|2:k 4:p 5:b 6:k|
 --- --------------- 
  

Expected :

  --- --------------- 
| id|           feat|
 --- --------------- 
| u1|1:a 5:c        |
| u2|5:b            |
 --- --------------- 
  

Вот моя попытка с несколькими функциями.

 df = spark.createDataFrame([("u1","1:a 2:k 5:c 6:i"),("u2","2:k 4:p 5:b 6:k")],["id","feat"])

select_idx=[1,5]

df.withColumn('feat', f.explode(f.split('feat', ' '))) 
  .withColumn('feat', f.split('feat', ':')) 
  .filter(f.col('feat')[0].isin(select_idx)) 
  .withColumn('feat', f.concat_ws(':', 'feat')) 
  .groupBy('id').agg(f.collect_list('feat').alias('feat')) 
  .withColumn('feat', f.concat_ws(' ', 'feat')) 
  .show(10, False)

 --- ------- 
|id |feat   |
 --- ------- 
|u1 |1:a 5:c|
|u2 |5:b    |
 --- -------