Repository: ray-project/ray
File: python/ray/data/dataset.py (Python)

Code before:
            ds = ds.random_shuffle(seed=seed)

        if not isinstance(test_size, (int, float)):
            raise TypeError(f"`test_size` must be int or float got {type(test_size)}.")

        # Validate that shuffle=True and stratify are not both specified
        if shuffle and stratify is not None:
            raise ValueError(
                "Cannot specify both 'shuffle=True' and 'stratify' parameters. "
                "Stratified splitting maintains class proportions and is incompatible with shuffling."
            )

        # Handle stratified splitting
        if stratify is not None:
            return self._stratified_train_test_split(ds, test_size, stratify)

        # Handle non-stratified splitting (existing logic)
        if isinstance(test_size, float):
            self._validate_test_size_float(test_size)
            return ds.split_proportionately([1 - test_size])
        else:
            ds, ds_length = self._get_dataset_length_for_split(ds)
            ds_length = self._validate_test_size_int(test_size, ds, ds_length=ds_length)
            return ds.split_at_indices([ds_length - test_size])

    def _get_dataset_length_for_split(self, ds: "Dataset") -> Tuple["Dataset", int]:
        dataset_length = ds._meta_count()
        if dataset_length is None:
            # Materialize once so split_at_indices() can reuse the computed snapshot.
            # Calling count() first would execute the upstream pipeline for counting,
            # then execute it again for the split.
            ds = ds.materialize()
            dataset_length = ds._meta_count()
            if dataset_length is None:
                dataset_length = ds.count()
        return ds, dataset_length

    def _stratified_train_test_split(
        self, ds: "Dataset", test_size: Union[int, float], stratify: str
    ) -> Tuple["MaterializedDataset", "MaterializedDataset"]:
        """Perform stratified train-test split on the dataset.

        Args:
            ds: The dataset to split.
            test_size: Test size as int or float.
            stratify: Column name to use for stratified sampling.

        Returns:
            Train and test subsets as two MaterializedDatasets.
        """
        # Normalize test_size to float (only materialize if needed)

Review comment (suggestion):
```suggestion
    def _try_count_or_materialize(self, ds: "Dataset") -> Tuple["Dataset", int]:
```

Code after:
                always complements the test split.
            shuffle: Whether or not to globally shuffle the dataset before splitting.
                Defaults to ``False``. This may be a very expensive operation with a
                large dataset.
            seed: Fix the random seed to use for shuffle, otherwise one is chosen
                based on system randomness. Ignored if ``shuffle=False``.
            stratify: Optional column name to use for stratified sampling. If provided,
                the splits will maintain the same proportions of each class in the
                stratify column across both train and test sets.

        Returns:
            Train and test subsets as two ``MaterializedDatasets``.

        .. seealso::

            :meth:`Dataset.split_proportionately`
        """
        ds = self

        if shuffle:
            ds = ds.random_shuffle(seed=seed)

        if not isinstance(test_size, (int, float)):
            raise TypeError(f"`test_size` must be int or float got {type(test_size)}.")

        # Validate that shuffle=True and stratify are not both specified
        if shuffle and stratify is not None:
            raise ValueError(
                "Cannot specify both 'shuffle=True' and 'stratify' parameters. "
                "Stratified splitting maintains class proportions and is incompatible with shuffling."
            )

        # Handle stratified splitting
        if stratify is not None:
            return self._stratified_train_test_split(ds, test_size, stratify)

        # Handle non-stratified splitting (existing logic)
        if isinstance(test_size, float):
            self._validate_test_size_float(test_size)
            return ds.split_proportionately([1 - test_size])
        else:
            ds, ds_length = self._try_count_or_materialize(ds)
            ds_length = self._validate_test_size_int(test_size, ds, ds_length=ds_length)
            return ds.split_at_indices([ds_length - test_size])

    def _try_count_or_materialize(self, ds: "Dataset") -> Tuple["Dataset", int]:
        dataset_length = ds._meta_count()
        if dataset_length is None:
            # Materialize once so split_at_indices() can reuse the computed snapshot.
            # Calling count() first would execute the upstream pipeline for counting,
            # then execute it again for the split.