Fix sort with indices mapping (#5587)

mariosasko · web-flow · commit d57fdcf2c811 · 2023-02-28T18:21:57.000+01:00
Fix sort with indices mapping
diff --git a/src/datasets/arrow_dataset.py b/src/datasets/arrow_dataset.py
@@ -3986,7 +3986,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,
         )
 

Original file line number	Diff line number	Diff line change
`@@ -3986,7 +3986,7 @@ def sort(`
`3986`	`3986`
`3987`	`3987`	`sort_table = query_table(`
`3988`	`3988`	`table=self._data,`
`3989`		`- key=range(self._data.num_rows),`
	`3989`	`+ key=slice(0, len(self)),`
`3990`	`3990`	`indices=self._indices if self._indices is not None else None,`
`3991`	`3991`	`)`
`3992`	`3992`