Accept numpy integers as train_size / test_size in train_test_split

`Dataset.train_test_split()` decides between "number of rows" and "fraction"
with `isinstance(test_size, int)` / `isinstance(test_size, float)`, and rejects
anything else:

    ds.train_test_split(test_size=np.int64(2))
    # ValueError: Invalid value for test_size: 2 of type <class 'numpy.int64'>

numpy floats are subclasses of `float` so they already work, but numpy integers
are not subclasses of `int`, which makes the behaviour inconsistent for values
that come out of the same numpy computation.

Convert numpy integers to `int` before the checks.

Author: LeSingh1

src/datasets/arrow_dataset.py

@@ -5086,6 +5086,13 @@ def train_test_split(
         if test_size is None and train_size is None:
             test_size = 0.25
 
+        # numpy integers are not instances of `int` (contrary to numpy floats, which are instances of `float`),
+        # so they wouldn't pass the checks below
+        if isinstance(test_size, np.integer):
+            test_size = int(test_size)
+        if isinstance(train_size, np.integer):
+            train_size = int(train_size)
+
         # Safety checks similar to scikit-learn's ones.
         # (adapted from https://github.com/scikit-learn/scikit-learn/blob/fd237278e895b42abe8d8d09105cbb82dc2cbba7/sklearn/model_selection/_split.py#L1750)
         n_samples = len(self)

tests/test_arrow_dataset.py

@@ -3694,6 +3694,17 @@ def test_concatenate_datasets(dataset_type, axis, expected_shape, dataset_dict,
     assert_arrow_metadata_are_synced_with_dataset_features(dataset)
 
 
+@pytest.mark.parametrize("dtype", ["int32", "int64", "uint8"])
+def test_train_test_split_with_numpy_integer_sizes(dtype):
+    dataset = Dataset.from_dict({"col_1": list(range(10))})
+    split = dataset.train_test_split(test_size=getattr(np, dtype)(2), seed=42)
+    assert len(split["test"]) == 2
+    assert len(split["train"]) == 8
+    split = dataset.train_test_split(train_size=getattr(np, dtype)(4), seed=42)
+    assert len(split["train"]) == 4
+    assert len(split["test"]) == 6
+
+
 def test_concatenate_datasets_new_columns():
     dataset1 = Dataset.from_dict({"col_1": ["a", "b", "c"]})
     dataset2 = Dataset.from_dict({"col_1": ["d", "e", "f"], "col_2": [True, False, True]})