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

[shashvat-singham]

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".

>>> 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

_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