huggingface / huggingface/datasets

Dataset.take(n) raises IndexError when n > len(dataset), but IterableDataset.take(n) returns what it has

Open
#8,483 3 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
22k
Forks
3.4k
Avg merge
5d 7h
Merged PRs (30d)
17

Description

### Describe the bug

`Dataset.take(n)` raises when asked for more elements than the dataset has, while `IterableDataset.take(n)` returns the elements it has. Both are documented identically as "Create a new [...] with only the first `n` elements".

```python
>>> from datasets import Dataset
>>> ds = Dataset.from_dict({"a": [1, 2, 3]})

>>> ds.take(3)
Dataset({features: ['a'], num_rows: 3})
>>> ds.take(4)
IndexError: Index 3 out of range for dataset of size 3.

>>> len(list(ds.to_iterable_dataset().take(4)))
3
```

So the same call against the same data succeeds or raises depending only on which of the two classes you hold.

### Cause

`take` is `self.select(range(n))`. For `n > len(self)` that's a contiguous range, so `select` takes the fast path and calls `_select_contiguous(0, n)`, where

```python
_check_valid_indices_value(start + length - 1, len(self))
```

is `_check_valid_indices_value(3, 3)` for `take(4)` on 3 rows, which raises.

### Expected behavior

I'd expect `take` to behave like every other "first n" API — `itertools.islice`, list slicing, `IterableDataset.take` — and return `min(n, len(ds))` elements rather than raise.

That said, this is a semantics decision rather than an obvious slip, which is why I'm filing it rather than sending a patch: unlike the `skip` case below, the current behaviour is at least *self-consistent* (it raises for every `n > len`), so someone may be relying on it as a bounds assertion. If you'd prefer to keep it strict, then the two docstrings and `IterableDataset.take` are what should change instead, so the pair stop disagreeing.

### Related

I've opened #8482 for the neighbouring `Dataset.skip` bug, which is a clearer defect — there, `skip(len(ds))` raises while `skip(len(ds) + 1)` returns an empty dataset, so it isn't even monotonic. That fix is deliberately scoped to zero-length slices and leaves `take` behaving exactly as it does today, so the two don't overlap; whichever way you want `take` to go can be decided independently.

Happy to send the `take` PR too if you tell me which direction you'd like.

### Environment info

- `datasets` 4.5.1.dev0 (`main`)
- Python 3.11.9, Windows

Contributor guide

Open the contributing guide

Research direction

Start by comparing Dataset.take with IterableDataset.take and tracing Dataset.take through select, _select_contiguous, and _check_valid_indices_value. Confirm the intended behavior for n greater than the dataset length with maintainers, then add regression coverage and make the relevant documentation consistent with the decision.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
data
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.