Sort pandas DataFrame first by subject ID then pick their corresponding values from other columns

Question

NOTES=[]
hadm=[]
for i in sorted(list(notes_df.SUBJECT_ID.unique())):
    #print(i)
    NOTES.append(list(notes_df[notes_df.SUBJECT_ID == i].TEXT))
    hadm.append(list(notes_df[notes_df.SUBJECT_ID == i].HADM_ID.astype(float)))
notes_1_df = pd.DataFrame(list(zip(sorted(list(notes_df.SUBJECT_ID.unique())),hadm,[item for sublist in NOTES for item in sublist])), columns =['SUBJECT_ID','HADM_ID','NOTES'])
notes_1_df.HADM_ID=notes_1_df.HADM_ID.astype(str)

The above code works fine to generate a new Dataframe with sorted values of subject ID and its corresponding values of other column. I used lists to store values and then used them to create a Dataframe.

But it is slow. How to make it fast?

notes_df structure:

ROW_ID|SUBJECT_ID|HADM_ID|TEXT
1      4           89     Here is the text
2      23          433    Here is the text,and so on
3      65          1212   Here is the text,and 
4      914         2212   Here is the text,and so on
5      23          112    Here is the text,and so on
6      23          773    Here is the text,and so on
7      65          1210   Here is the text,and so 
8      23          1212   Here is the text,and so on

Answer 1

What makes your code slow is the repeated calls to unique, and manually shunting all entries into lists repeatedly. Ideally you want to do it all in pandas for the most speed.

If I understand it correctly, you want to group by the subject ID, collect all hadm IDs and texts. In that case, you can just use pandas.DataFrame.groupby and pandas.DataFrame.aggregate to achieve (almost) the same result:

notes_1_df = notes_df.drop(columns=["ROW_ID"]) \
                     .groupby("SUBJECT_ID") \
                     .aggregate(list) \
                     .reset_index()

Which directly produces this output:

   SUBJECT_ID                HADM_ID                                               TEXT
0           4                   [89]                                 [Here is the text]
1          23  [433, 112, 773, 1212]  [Here is the text,and so on, Here is the text,...
2          65           [1212, 1210]    [Here is the text,and, Here is the text,and so]
3         914                 [2212]                       [Here is the text,and so on]

Whereas your code produces:

   SUBJECT_ID                        HADM_ID                       NOTES
0           4                         [89.0]            Here is the text
1          23  [433.0, 112.0, 773.0, 1212.0]  Here is the text,and so on
2          65               [1212.0, 1210.0]  Here is the text,and so on
3         914                       [2212.0]  Here is the text,and so on

This differs mainly in the notes. Here your code is a bit weird. Instead of doing what you did, let's just pick the first text from each subject for now.

I am assuming that having integers for the hadm ID is fine, otherwise you can be more specific and supply two different functions to aggregate, in which case you don't even need the drop anymore:

notes_1_df = notes_df.groupby("SUBJECT_ID") \
                     .aggregate({"HADM_ID": lambda x: list(map(float, x)),
                                 "TEXT": "first"}) \
                     .rename(columns={"TEXT": "NOTES"}) \
                     .reset_index()

   SUBJECT_ID                        HADM_ID                       NOTES
0           4                         [89.0]            Here is the text
1          23  [433.0, 112.0, 773.0, 1212.0]  Here is the text,and so on
2          65               [1212.0, 1210.0]        Here is the text,and
3         914                       [2212.0]  Here is the text,and so on

---

To see why I think your code produces weird results, let's replace the text with unique texts:

df["TEXT"] = "Text from subject " + df.SUBJECT_ID.astype(str) + ", hadm " + df.HADM_ID.astype(str)

Then my (last) code produces:

   SUBJECT_ID                        HADM_ID                             NOTES
0           4                         [89.0]      Text from subject 4, hadm 89
1          23  [433.0, 112.0, 773.0, 1212.0]    Text from subject 23, hadm 433
2          65               [1212.0, 1210.0]   Text from subject 65, hadm 1212
3         914                       [2212.0]  Text from subject 914, hadm 2212

Where each text is the first text from that actual subject. In contrast, your code produces:

   SUBJECT_ID                        HADM_ID                           NOTES
0           4                         [89.0]    Text from subject 4, hadm 89
1          23  [433.0, 112.0, 773.0, 1212.0]  Text from subject 23, hadm 433
2          65               [1212.0, 1210.0]  Text from subject 23, hadm 112
3         914                       [2212.0]  Text from subject 23, hadm 773

Note how the texts do not correspond to the same subjects anymore!