Как воспроизвести функцию окна Spark с помощью collect_list в кадрах данных Pandas?

#pandas #apache-spark #pyspark #window

Вопрос:

У меня есть начальный фрейм данных

 df1 =   --- --- ---  | A| B| C|  --- --- ---  | 1| 1| 10| | 1| 2| 11| | 1| 2| 12| | 3| 1| 13| | 2| 1| 14| | 2| 1| 15| | 2| 1| 16| | 4| 1| 17| | 4| 2| 18| | 4| 3| 19| | 4| 4| 19| | 4| 5| 20| | 4| 5| 20|  --- --- ---    

Используя pyspark, я закодировал фрейм данных с оконной функцией, используя collect_list функцию, с учетом столбца группировки 'A' и с учетом столбца 'B' , отсортированного для создания столбца с кумулятивными списками

 spec = Window.partitionBy('A').orderBy('B') df1 = df1.withColumn('D',collect_list('C').over(spec))  df1.orderBy('A','B').show()   --- --- --- ------------------------  |A |B |C |D |  --- --- --- ------------------------  |1 |1 |10 |[10] | |1 |2 |11 |[10, 11, 12] | |1 |2 |12 |[10, 11, 12] | |2 |1 |14 |[14, 15, 16] | |2 |1 |15 |[14, 15, 16] | |2 |1 |16 |[14, 15, 16] | |3 |1 |13 |[13] | |4 |1 |17 |[17] | |4 |2 |18 |[17, 18] | |4 |3 |19 |[17, 18, 19] | |4 |4 |19 |[17, 18, 19, 19] | |4 |5 |20 |[17, 18, 19, 19, 20, 20]| |4 |5 |20 |[17, 18, 19, 19, 20, 20]|  --- --- --- ------------------------     

Можно ли выполнить тот же расчет с помощью фрейма данных Pandas?

Я попытался использовать «обычный» код на python, но, вероятно, есть способ сделать это более непосредственно.

Ответ №1:

Один из способов решения этой проблемы в pandas-использовать две группы, т. е. Сначала сгруппировать фрейм данных по столбцу A , затем для каждой группы применить пользовательскую функцию collect_list , которая, в свою очередь, группирует ввод по столбцу B и суммирует столбец C с помощью list

 def collect_list(g):  return g.groupby('B')['C'].agg(list).cumsum()  df.sort_values(['A', 'B']).merge(  df.groupby('A').apply(collect_list).reset_index(name='D'))  

 A B C D 0 1 1 10 [10] 1 1 2 11 [10, 11, 12] 2 1 2 12 [10, 11, 12] 4 2 1 14 [14, 15, 16] 5 2 1 15 [14, 15, 16] 6 2 1 16 [14, 15, 16] 3 3 1 13 [13] 7 4 1 17 [17] 8 4 2 18 [17, 18] 9 4 3 19 [17, 18, 19] 10 4 4 19 [17, 18, 19, 19] 11 4 5 20 [17, 18, 19, 19, 20, 20] 12 4 5 20 [17, 18, 19, 19, 20, 20]