Subtable timestamps in python pandas

Question

I have a dataframe like:

    time_stamp           id          next_timestamp
0   2010-04-16 11:57:52  string_1    NaT
1   2010-04-16 12:06:16  string_1    NaT
2   2010-04-16 12:40:53  string_2    NaT

I want to fill next_timestamp column, with the next time_stamp that id has (if it exists).

Result would be something like:

    time_stamp           id          next_timestamp
0   2010-04-16 11:57:52  string_1    2010-04-16 12:06:16
1   2010-04-16 12:06:16  string_1    NaT
2   2010-04-16 12:40:53  string_2    NaT

My code at the moment:

for row in df.index:
    row_time_stamp = df.time_stamp[row]
    id_array = df.id[row]
    df_temp = df.loc[(df['time_stamp'] >= row_time_stamp) & \
                     (df['time_stamp'] <= row_time_stamp + datetime.timedelta(days=7))]
    try:
        next_id_msg = df_temp.loc[(df_temp['id'] == str(id_array))].time_stamp.min()
        df['next_timestamp'][row] = next_id_msg
    except IndexError:
          df['next_timestamp'][row] = pd.NaT  

The problem is that my my df is 50+ million rows long, and setting up a temp table for every row is not a good pattern.

Please help me out with a better pattern.

Answer 1

you can use groupby.transform

df["next_timestamp"] = df.groupby("id")["time_stamp"].transform(
    lambda x: x.shift(-1)
)