Method train_test_split

python/ray/data/dataset.py:2614–2681

Materialize and split the dataset into train and test subsets. Examples: >>> import ray >>> ds = ray.data.range(8) >>> train, test = ds.train_test_split(test_size=0.25) >>> train.take_batch() {'id': array([0, 1, 2, 3, 4, 5])}

(
        self,
        test_size: Union[int, float],
        *,
        shuffle: bool = False,
        seed: Optional[int] = None,
        stratify: Optional[str] = None,
    )

    @ConsumptionAPI
    @PublicAPI(api_group=SMJ_API_GROUP)
    def train_test_split(
        self,
        test_size: Union[int, float],
        *,
        shuffle: bool = False,
        seed: Optional[int] = None,
        stratify: Optional[str] = None,
    ) -> Tuple["MaterializedDataset", "MaterializedDataset"]:
        """Materialize and split the dataset into train and test subsets.

        Examples:

            >>> import ray
            >>> ds = ray.data.range(8)
            >>> train, test = ds.train_test_split(test_size=0.25)
            >>> train.take_batch()
            {'id': array([0, 1, 2, 3, 4, 5])}
            >>> test.take_batch()
            {'id': array([6, 7])}

        Args:
            test_size: If float, should be between 0.0 and 1.0 and represent the
                proportion of the dataset to include in the test split. If int,
                represents the absolute number of test samples. The train split
                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: