Интерполировать столбец DataFrame и сортировать по другому столбцу в PySpark или Pandas

Question

Учитывая следующий DataFrame, нам нужно интерполировать my_column значения из примера и использовать их как отдельные столбцы, а затем сортировать по int_column значениям, которые принадлежат каждому столбцу some_id в порядке убывания. Пример:

+--------------------+-----------+------------------+
|          some_id   | my_column |      int_column  |
+--------------------+-----------+------------------+
|xx1                 |id_1       |           3      |
|xx1                 |id_2       |           4      |
|xx1                 |id_3       |           5      |
|xx2                 |id_1       |           6      |
|xx2                 |id_2       |           1      |
|xx2                 |id_3       |           3      |
|xx3                 |id_1       |           4      |
|xx3                 |id_2       |           8      |
|xx3                 |id_3       |           9      |
|xx4                 |id_1       |           1      |
+--------------------+-----------+------------------+

Ожидаемый результат:

+--------------------+-----------+------------------+
|          id_1      | id_2      |      id_3        |
+--------------------+-----------+------------------+
| [xx4, 1]           |[xx2, 1]   |[xx2, 3]          |
| [xx1, 3]           |[xx1, 4]   |[xx1, 5]          |
| [xx3, 4]           |[xx3, 8]   |[xx3, 9]          |
| [xx2, 6]           |null       |null              |
+--------------------+-----------+------------------+

Как видите, для id_1 самое низкое число в int_column равно 1 в конце DataFrame, и оно принадлежит xx4 из столбца some_id, следующее значение равно 3, 4, и 6, каждый из которых принадлежит к xx1, xx3 и xx2 соответственно.

Есть какие-нибудь указатели на то, как подойти к этой проблеме? Можно использовать PySpark или Pandas.

Код для воспроизведения входных данных:

import pandas as pd

data = {'some_id': ['xx1', 'xx1', 'xx1', 'xx2', 'xx2', 'xx2', 'xx3', 'xx3', 'xx3', 'xx4'], \
        'my_column' : ['id_1', 'id_2', 'id_3', 'id_1', 'id_2', 'id_3', 'id_1', 'id_2', 'id_3', 'id_1'],\
       'int_column' : [3, 4, 5, 6 , 1, 3, 4, 8, 9, 1]}

df = pd.DataFrame.from_dict(data)

Answer 1

Вот решение в pyspark.

Сначала определите Window, чтобы разделить на my_column и упорядочить по int_column. Мы определим порядок, используя pyspark.sql.functions.row_number() для этого раздела.

from pyspark.sql import Window
import pyspark.sql.functions as f
w = Window.partitionBy("my_column").orderBy("int_column")
df.withColumn("order", f.row_number().over(w)).sort("order").show()
#+-------+---------+----------+-----+
#|some_id|my_column|int_column|order|
#+-------+---------+----------+-----+
#|    xx4|     id_1|         1|    1|
#|    xx2|     id_2|         1|    1|
#|    xx2|     id_3|         3|    1|
#|    xx1|     id_2|         4|    2|
#|    xx1|     id_1|         3|    2|
#|    xx1|     id_3|         5|    2|
#|    xx3|     id_2|         8|    3|
#|    xx3|     id_3|         9|    3|
#|    xx3|     id_1|         4|    3|
#|    xx2|     id_1|         6|    4|
#+-------+---------+----------+-----+

Обратите внимание, что (xx4, 1) находится в первом ряду после сортировки по order, как вы объяснили.

Теперь вы можете группировать по order и pivot фрейм данных на my_column. Для этого требуется функция агрегирования, поэтому я буду использовать pyspark.sql.functions.first(), потому что я предполагаю, что на order существует только одна пара (some_id, int_column). Затем просто отсортируйте по order и опустите этот столбец, чтобы получить желаемый результат:

df.withColumn("order", f.row_number().over(w))\
    .groupBy("order")\
    .pivot("my_column")\
    .agg(f.first(f.array([f.col("some_id"), f.col("int_column")])))\
    .sort("order")\
    .drop("order")\
    .show(truncate=False)
#+--------+--------+--------+
#|id_1    |id_2    |id_3    |
#+--------+--------+--------+
#|[xx4, 1]|[xx2, 1]|[xx2, 3]|
#|[xx1, 3]|[xx1, 4]|[xx1, 5]|
#|[xx3, 4]|[xx3, 8]|[xx3, 9]|
#|[xx2, 6]|null    |null    |
#+--------+--------+--------+

Answer 2

Нам нужен вспомогательный ключ, создайте с помощью cumcount, затем мы используем groupby + apply (эта часть точно так же, как pivot, или вы можете использовать pivot_table или crosstab)

df=df.assign(key=df.groupby('my_column').cumcount())
df.groupby(['key','my_column']).apply(lambda x : list(zip(x['some_id'],x['int_column']))[0]).unstack()
Out[378]: 
my_column      id_1      id_2      id_3
key                                    
0          (xx1, 3)  (xx1, 4)  (xx1, 5)
1          (xx2, 6)  (xx2, 1)  (xx2, 3)
2          (xx3, 4)  (xx3, 8)  (xx3, 9)
3          (xx4, 1)      None      None

При использовании pivot + sort_values

df=df.sort_values('int_column').assign(key=df.groupby('my_column').cumcount())
df['Value']=list(zip(df['some_id'],df['int_column']))
s=df.pivot(index='key',columns='my_column',values='Value')
s
Out[397]: 
my_column      id_1      id_2      id_3
key                                    
0          (xx4, 1)  (xx2, 1)  (xx2, 3)
1          (xx1, 3)  (xx1, 4)  (xx1, 5)
2          (xx3, 4)  (xx3, 8)  (xx3, 9)
3          (xx2, 6)      None      None