[Feature request] Indexing datasets by a customly-defined id field to enable random access dataset items via the id

[Yu-Shi]

Feature request

Some datasets may contain an id-like field, for example the id field in wikimedia/wikipedia and the _id field in BeIR/dbpedia-entity. HF datasets support efficient random access via row, but not via this kinds of id fields. I wonder if it is possible to add support for indexing by a custom "id-like" field to enable random access via such ids. The ids may be numbers or strings.

Motivation

In some cases, especially during inference/evaluation, I may want to find out the item that has a specified id, defined by the dataset itself.

For example, in a typical re-ranking setting in information retrieval, the user may want to re-rank the set of candidate documents of each query. The input is usually presented in a TREC-style run file, with the following format:

<qid> Q0 <docno> <rank> <score> <tag>

The re-ranking program should be able to fetch the queries and documents according to the <qid> and <docno>, which are the original id defined in the query/document datasets. To accomplish this, I have to iterate over the whole HF dataset to get the mapping from real ids to row ids every time I start the program, which is time-consuming. Thus I want HF dataset to provide options for users to index by a custom id column, not by row.

Your contribution

I'm not an expert in this project and I'm afraid that I'm not able to make contributions on the code.

[lhoestq]

You can simply use a python dict as index:

>>> from datasets import load_dataset
>>> ds = load_dataset("BeIR/dbpedia-entity", "corpus", split="corpus")
>>> index = {key: idx for idx, key in enumerate(ds["_id"])}
>>> ds[index["<dbpedia:Pikachu>"]]
{'_id': '<dbpedia:Pikachu>',
 'title': 'Pikachu',
 'text': 'Pikachu (Japanese: ピカチュウ) are a fictional species of Pokémon.  Pokémon are fictional creatures that appear in an assortment of comic books, animated movies and television shows, video games, and trading card games licensed by The Pokémon Company, a Japanese corporation. The Pikachu design was conceived by Ken Sugimori.'}

[Yu-Shi]

Thanks for your reply. Yes, I can do that, but it is time-consuming to do that every time I launch the program (some datasets are extremely big). HF Datasets has a nice feature to support instant data loading and efficient random access via row ids. I'm curious if this beneficial feature could be further extended to custom data columns.

[davidmrau]

+1 on the issue I think it would be extremely useful

[GeeYangML]

+1. This could be very useful.

[ruth-ann]

+1 - currently having to manually implement this

[davidmrau]

If anyone has an idea how to do this in the right way (perhaps @albertvillanova ?) I would be happy to implement it

[SwayStar123]

This would be very helpful to implement aspect ratio bucketing for image and video datasets