From 0f7f35732d1857783be60e0ef37ab8828560e2d9 Mon Sep 17 00:00:00 2001 From: Vidit Patankar Date: Wed, 8 Jul 2026 13:24:37 +0530 Subject: [PATCH] Fix Dataset.take raising IndexError when taking more elements than available Dataset.take(n) was implemented as self.select(range(n)), so when n exceeds the dataset length it raised IndexError instead of returning the whole dataset. This is inconsistent with IterableDataset.take, which clamps and returns all available elements (the two were meant to be aligned, see #6813); Dataset.skip already clamps correctly. Clamp n to len(self) so take returns the whole dataset when asked for more. --- src/datasets/arrow_dataset.py | 5 ++++- tests/test_arrow_dataset.py | 17 +++++++++++++++++ 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/src/datasets/arrow_dataset.py b/src/datasets/arrow_dataset.py index 59451a640e6..4d8d0a56df0 100644 --- a/src/datasets/arrow_dataset.py +++ b/src/datasets/arrow_dataset.py @@ -4699,7 +4699,10 @@ def take(self, n: int) -> "Dataset": 'text': 'the gorgeously elaborate continuation of " the lord of the rings " trilogy is so huge that a column of words cannot adequately describe co-writer/director peter jackson\'s expanded vision of j . r . r . tolkien\'s middle-earth .'}] ``` """ - return self.select(range(n)) + # Clamp `n` to the dataset length so taking more elements than available + # returns the whole dataset instead of raising an IndexError, matching + # the behavior of `IterableDataset.take`. + return self.select(range(min(n, len(self)))) @transmit_format @fingerprint_transform(inplace=False, ignore_kwargs=["load_from_cache_file", "indices_cache_file_name"]) diff --git a/tests/test_arrow_dataset.py b/tests/test_arrow_dataset.py index 18a8a6038fe..1f00c66aa93 100644 --- a/tests/test_arrow_dataset.py +++ b/tests/test_arrow_dataset.py @@ -2885,6 +2885,23 @@ def test_shard(self, in_memory): with dset.shard(num_shards=3, index=0) as dset_sharded_formatted: self.assertEqual(dset_sharded_formatted.format["type"], "numpy") + def test_take(self, in_memory): + with tempfile.TemporaryDirectory() as tmp_dir: + with self._create_dummy_dataset(in_memory, tmp_dir) as dset: + n = len(dset) + with dset.take(2) as taken: + self.assertEqual(len(taken), 2) + self.assertEqual(taken["filename"], dset["filename"][:2]) + with dset.take(n) as taken: + self.assertEqual(len(taken), n) + # taking more elements than available returns the whole dataset instead + # of raising an IndexError, matching the behavior of IterableDataset.take + with dset.take(n + 10) as taken: + self.assertEqual(len(taken), n) + self.assertEqual(taken["filename"], dset["filename"]) + with dset.take(0) as taken: + self.assertEqual(len(taken), 0) + def test_flatten_indices(self, in_memory): with tempfile.TemporaryDirectory() as tmp_dir: with self._create_dummy_dataset(in_memory, tmp_dir) as dset: