Instead of relying on list_dataset_elements_per_split, we should use peewee queries.
list_dataset_elements_per_split