diff --git a/src/datasets/arrow_dataset.py b/src/datasets/arrow_dataset.py
index affe1fa4470..f2731c83963 100644
--- a/src/datasets/arrow_dataset.py
+++ b/src/datasets/arrow_dataset.py
@@ -3994,7 +3994,7 @@ def sort(
 
         sort_table = query_table(
             table=self._data,
-            key=range(self._data.num_rows),
+            key=slice(0, len(self)),
             indices=self._indices if self._indices is not None else None,
         )