Hi ! Since you used train_test_split, then the new train dataset is now shuffled. Querying a shuffled dataset takes some time since it has to re-order all the elements. So the bigger the dataset, the longer it takes to query all the data of a shuffled column.
If you want more speed, either don’t shuffle (shuffle=False in train_test_split), or use flatten_indices once and for all (it re-order the rows correctly).