6367318f455aa8c27b6341c9b98794351dfd168e,examples/pytorch/pinsage/data_utils.py,,train_test_split_by_time,#Any#Any#Any#,9

Before Change

    df["test_mask"] = np.zeros((len(df),), dtype=np.bool)
    df = df.sort_values([item, timestamp])
    for track_id in df[item].unique():
        idx = (df[item] == track_id).to_numpy().nonzero()[0]
        idx = df.index[idx]
        if len(idx) > 1:
            df.loc[idx[-1], "train_mask"] = False
            df.loc[idx[-1], "test_mask"] = True
        if len(idx) > 2:

After Change

            df.iloc[-2, -2] = True
        return df
    df = df.groupby(item).apply(train_test_split).compute(scheduler="processes").sort_index()
    print(df[df[item] == df[item].unique()[0]].sort_values(timestamp))
    return df["train_mask"].to_numpy().nonzero()[0], \
           df["val_mask"].to_numpy().nonzero()[0], \
           df["test_mask"].to_numpy().nonzero()[0]

Instances