Fix Dataset.take raising IndexError when taking more elements than available

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"])

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: