Iterating over values of a column in the IterableDataset
Opened by TopCoder2K (closed)

### Feature request

I would like to be able to iterate (and re-iterate if needed) over a column of an `IterableDataset` instance. The following example shows the supposed API:
```python
def gen():
    yield {"text": "Good", "label": 0}
    yield {"text": "Bad", "label": 1}

ds = IterableDataset.from_generator(gen)
texts = ds["text"]

for v in texts:
    print(v)  # Prints "Good" and "Bad"

for v in texts:
    print(v)  # Prints "Good" and "Bad" again
```

### Motivation

In the real world problems, huge NNs like Transformer are not always the best option, so there is a need to conduct experiments with different methods. While 🤗Datasets is perfectly adapted to 🤗Transformers, it may be inconvenient when being used with other libraries. The ability to retrieve a particular column is the case (e.g., gensim's FastText [requires](https://radimrehurek.com/gensim/models/fasttext.html#gensim.models.fasttext.FastText.train) only lists of strings, not dictionaries).
While there are ways to achieve the desired functionality, they are not good ([forum](https://discuss.huggingface.co/t/how-to-iterate-over-values-of-a-column-in-the-iterabledataset/135649)). It would be great if there was a built-in solution.

### Your contribution

Theoretically, I can submit a PR, but I have very little knowledge of the internal structure of 🤗Datasets, so some help may be needed.
Moreover, I can only work on weekends, since I have a full-time job. However, the feature does not seem to be popular, so there is no need to implement it as fast as possible.

Comments: I'd be in favor of that ! I saw many people implementing their own iterables that wrap a dataset just to iterate on a single column, that would make things more practical.

Kinda related: https://github.com/huggingface/datasets/issues/5847

(For anyone's information, I'm going on vacation for the next 3 weeks, so the work is postponed. If anyone can implement this feature within the next 4 weeks, go ahead :) )

UPD from 04/06/25:
I'm planning to start work on the feature in early May.

#self-assign

# Preliminary discussion

Ideally, I would like to be able to operate on a column with [map](https://huggingface.co/docs/datasets/package_reference/main_classes#datasets.IterableDataset.map), [filter](https://huggingface.co/docs/datasets/package_reference/main_classes#datasets.IterableDataset.filter), [batch](https://huggingface.co/docs/datasets/package_reference/main_classes#datasets.IterableDataset.batch) and probably some other `IterableDataset`'s methods, however, the same results can be achieved by using the methods on an `IterableDataset` object and utilizing `__getitem__()` afterwards. Thus, one may not support these methods at first and try to make the implementation as simple as possible.

# Implementation

Based on the preliminary discussion, one can do the following:
```python
class IterableColumn:
    def __init__(self, dataset: "IterableDataset", column_name: str):
        self.dataset = dataset
        self.column_name = column_name

    def __iter__(self) -> Iterator[Any]:
        for example in self.dataset:
            yield example[self.column_name]


class IterableDataset(DatasetInfoMixin):
    ...
    def __getitem__(self, column_name: str) -> IterableColumn:
        return IterableColumn(self, column_name)
    ...
```

# Testing

It works as expected in our simple test:
```python
def gen():
    yield {"text": "Good", "label": 0}
    yield {"text": "Bad", "label": 1}

ds = IterableDataset.from_generator(gen)

texts = ds["text"]  # `texts` is an IterableColumn object
for v in texts:
    print(v)  # Prints "Good" and "Bad"
for v in texts:
    print(v)  # Prints "Good" and "Bad" again
```

# Questions

1. What do you think about the implementation, @lhoestq?
2. How to properly test the implementation? I've found [test_iterable_dataset.py](https://github.com/huggingface/datasets/blob/main/tests/test_iterable_dataset.py) but 1) I haven't found any guidelines for testing, 2) the script tests a lot of things while I'd like to test only my feature.

Sounds great !

Regarding testing, it's actually possible to have your test function in test_iterable_dataset.py, which you can run using

```python
pytest tests/test_iterable_dataset.py::my_function
```

> Regarding testing, it's actually possible to have your test function in test_iterable_dataset.py, which you can run using

I hoped not to run `pip install -e ".[dev]"`, but your answer implies that I should. The problem is that I was unable to install the dependencies with Python 3.13 due to `tensorflow` and  with Python 3.11-3.12 due to "there are no versions of pyav" [¬º-°]¬ Therefore, I had to test in a separate script file to avoid importing optional dependencies. Anyway, I've opened a PR: https://github.com/huggingface/datasets/pull/7564. Please, take a look (there are questions about the documentation).

Moreover, I want to note that `make style` and `pre-commit` give different results for `test_iterable_dataset.py` (and a couple of files). Example:
```python
    assert skip_ex_iterable.shuffle_data_sources(np.random.default_rng(42)) is skip_ex_iterable, (
        "skip examples makes the shards order fixed"
    )
```
vs
```python
    assert (
        skip_ex_iterable.shuffle_data_sources(np.random.default_rng(42)) is skip_ex_iterable
    ), "skip examples makes the shards order fixed"
```
 ¯\\_(ツ)_/¯

> Kinda related: https://github.com/huggingface/datasets/issues/5847

I had forgotten about this, but I've looked at it by now. [This comment](https://github.com/huggingface/datasets/issues/5847#issuecomment-1549799951) implies that `IterableColumn` should support chained indexing, so thank you for pointing this out! Did you mean anything else by referencing the issue?

> I hoped not to run pip install -e ".[dev]", but your answer implies that I should. The problem is that I was unable to install the dependencies with Python 3.13 due to tensorflow and with Python 3.11-3.12 due to "there are no versions of pyav" [¬º-°]¬ Therefore, I had to test in a separate script file to avoid importing optional dependencies. Anyway, I've opened a PR: https://github.com/huggingface/datasets/pull/7564. Please, take a look (there are questions about the documentation).

we try to not not require optional dependencies when running tests, so you can try running the tests only with `pytest`, `pytest-datadir` and `pytest-xdist`

> I had forgotten about this, but I've looked at it by now. https://github.com/huggingface/datasets/issues/5847#issuecomment-1549799951 implies that IterableColumn should support chained indexing, so thank you for pointing this out! Did you mean anything else by referencing the issue?

No I simply referenced the issue because it will enable `pipe(ds["column_name"])`, but no need to support nested fields access in a first step - we can see that later as it's uncommon and would add complexity to the contribution

> we try to not not require optional dependencies when running tests, so you can try running the tests only with `pytest`, `pytest-datadir` and `pytest-xdist`

Understood. If it's necessary to run the tests again, I'll try to install only the mentioned libraries, thank you!

> No I simply referenced the issue because it will enable pipe(ds["column_name"]), but no need to support nested fields access in a first step - we can see that later as it's uncommon and would add complexity to the contribution

Ah, I see. Anyway, I've already implemented chained indexing (it was relatively easy).

@lhoestq, could you please take a look at the PR and answer [questions](https://github.com/huggingface/datasets/pull/7564#issuecomment-2863391781) there?

> so you can try running the tests only with pytest, pytest-datadir and pytest-xdist

Yes, they are sufficient. There was one more problem with Python 3.12 and `distutils` that were removed, but I just downgraded to 3.11 and successfully ran `test_iterable_dataset.py`.

@lhoestq, could you write in the [discussion](https://discuss.huggingface.co/t/how-to-iterate-over-values-of-a-column-in-the-iterabledataset/135649) for people coming there from the Internet that the feature has been implemented? I could do it by myself but the topic is closed to me.

done, thanks you !