huggingface / huggingface/datasets
Dataset.take(n) raises IndexError when n > len(dataset), but IterableDataset.take(n) returns what it has
- 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
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