resolves the ValueError: Unable to avoid copy while creating an array (#7831)

ArjunJagdale · lhoestq · web-flow · commit 9e5b0e6f1e2f · 2025-10-28T17:10:42.000+01:00
* Fix argument passing in stratified shuffle split

NumPy 2.0 changed the behavior of the `copy=False` parameter to be stricter. When `train_test_split` converted Arrow arrays to NumPy format for stratification, it triggered this error for non-contiguous arrays. Using `np.asarray()` allows copying when necessary, which is the recommended migration path per NumPy 2.0 documentation.

* make style

---------

Co-authored-by: Quentin Lhoest &lt;42851186+lhoestq@users.noreply.github.com&gt;
diff --git a/src/datasets/arrow_dataset.py b/src/datasets/arrow_dataset.py
@@ -4868,7 +4868,7 @@ def train_test_split(
                 try:
                     train_indices, test_indices = next(
                         stratified_shuffle_split_generate_indices(
-                            self.with_format("numpy")[stratify_by_column], n_train, n_test, rng=generator
+                            np.asarray(self.with_format("numpy")[stratify_by_column]), n_train, n_test, rng=generator
                         )
                     )
                 except Exception as error:

Original file line number	Diff line number	Diff line change
`@@ -4868,7 +4868,7 @@ def train_test_split(`
`4868`	`4868`	`try:`
`4869`	`4869`	`train_indices, test_indices = next(`
`4870`	`4870`	`stratified_shuffle_split_generate_indices(`
`4871`		`- self.with_format("numpy")[stratify_by_column], n_train, n_test, rng=generator`
	`4871`	`+ np.asarray(self.with_format("numpy")[stratify_by_column]), n_train, n_test, rng=generator`
`4872`	`4872`	`)`
`4873`	`4873`	`)`
`4874`	`4874`	`except Exception as error:`