Создание нового столбца и строк в pyspark для датафреймов

Question

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

id_1    id_desc    cat_1    cat_2
111      ask        ele     phone
222      ask hr     ele     phone
333      ask hr dk  ele     phone
444      askh       ele     phone

Если cat_1, cat_2 одинаковы для нескольких id_1, эта ассоциация должна быть зафиксирована как новый столбец.

Нужен вывод, подобный этому,

id_1    id_desc        cat_1    cat_2   id_2
111      ask             ele    phone   222
111      ask             ele    phone   333
111      ask             ele    phone   444
222      ask hr          ele    phone   111
222      ask hr          ele    phone   333
222      ask hr          ele    phone   444
333      ask hr dk       ele    phone   111
333      ask hr dk       ele    phone   222
333      ask hr dk       ele    phone   444

как это сделать в python?

Answer 1

Я не смог придумать ничего особенно изящного, но это должно было сделать работу:

import pandas as pd
import numpy as np

df = pd.DataFrame([[111, 'ask', 'ele', 'phone'], 
                   [222, 'ask_hr', 'ele', 'phone'], 
                   [333, 'ask_hr_dk', 'ele', 'phone'], 
                   [444, 'askh', 'ele', 'phone']], 
                   columns=['id_1', 'id_desc', 'cat_1', 'cat_2'])

grouped = df.groupby(by=['cat_1', 'cat_2'])  # group by the columns you want to be identical

data = []  # a list to store all unique groups

# In your example, this loop is not needed, but this generalizes to more than 1 pair
# of cat_1 and cat_2 values
for group in grouped.groups:  
    n_rows = grouped.get_group(group).shape[0]  # how many unique id's in a group
    all_data = np.tile(grouped.get_group(group).values, (n_rows, 1))  # tile the data n_row times
    ids = np.repeat(grouped.get_group(group)['id_1'].values, n_rows)  # repeat the ids n_row times
    data += [np.c_[all_data, ids]]  # concat the two sets of data and add to list

df_2 = pd.DataFrame(np.concatenate(data), columns=['id_1', 'id_desc', 'cat_1', 'cat_2', 'id_2'])

Основная идея - сгруппировать ваши данные по cat_1 и cat_2столбцы (используя groupby), используйте np.tile, чтобы создавать копии каждой группы столько раз, сколько в группе уникальных значений id_1, и объединять результатс уникальными значениями id_1 (по одному значению на группу данных).

Если вы не хотите, чтобы id_1 всегда был таким же, как id_2, просто выберите строки, в которых они не совпадают.match:

df_2 = df_2[df_2['id_1'] != df_2['id_2']] 

И если вы хотите, чтобы они были отсортированы по id_1:

df_2.sort_values('id_1', inplace=True)