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

Code before:
                Unlike :meth:`~Dataset.split_proportionately`, which lets you split a
                dataset into different sizes, :meth:`Dataset.split` splits a dataset
                into approximately equal splits.

            :meth:`Dataset.split_at_indices`
                :meth:`Dataset.split_proportionately` uses this method under the hood.

            :meth:`Dataset.streaming_split`.
                Unlike :meth:`~Dataset.split`, :meth:`~Dataset.streaming_split`
                doesn't materialize the dataset in memory.
        """

        if len(proportions) < 1:
            raise ValueError("proportions must be at least of length 1")
        if sum(proportions) >= 1:
            raise ValueError("proportions must sum to less than 1")
        if any(p <= 0 for p in proportions):
            raise ValueError("proportions must be bigger than 0")

        ds = self
        dataset_length = ds._meta_count()
        if dataset_length is None:
            ds = ds.materialize()
            dataset_length = ds._meta_count()
            if dataset_length is None:
                dataset_length = ds.count()
        cumulative_proportions = np.cumsum(proportions)
        split_indices = [
            int(dataset_length * proportion) for proportion in cumulative_proportions
        ]

        # Ensure each split has at least one element
        subtract = 0
        for i in range(len(split_indices) - 2, -1, -1):
            split_indices[i] -= subtract
            if split_indices[i] == split_indices[i + 1]:
                subtract += 1
                split_indices[i] -= 1
        if any(i <= 0 for i in split_indices):
            raise ValueError(
                "Couldn't create non-empty splits with the given proportions."
            )

        return ds.split_at_indices(split_indices)

    @ConsumptionAPI
    @PublicAPI(api_group=SMJ_API_GROUP)
    def train_test_split(
        self,
        test_size: Union[int, float],
        *,

Review comment (suggestion):
Actually can't we just call `ds.count()` here? Cause materializing defeats the purpose of this optimization

Code after:
            logical_plan = LogicalPlan(
                InputData(input_data=ref_bundles),
                self.context,
            )

            splits.append(
                MaterializedDataset(
                    ExecutionPlan(stats, self.context.copy()),
                    logical_plan,
                )
            )
        return splits

    @ConsumptionAPI
    @PublicAPI(api_group=SMJ_API_GROUP)
    def split_proportionately(
        self, proportions: List[float]
    ) -> List["MaterializedDataset"]:
        """Materialize and split the dataset using proportions.

        A common use case for this is splitting the dataset into train
        and test sets (equivalent to eg. scikit-learn's ``train_test_split``).
        For a higher level abstraction, see :meth:`Dataset.train_test_split`.

        This method splits datasets so that all splits
        always contains at least one row. If that isn't possible,
        an exception is raised.

        This is equivalent to caulculating the indices manually and calling
        :meth:`Dataset.split_at_indices`.

        Examples:
            >>> import ray
            >>> ds = ray.data.range(10)
            >>> d1, d2, d3 = ds.split_proportionately([0.2, 0.5])
            >>> d1.take_batch()
            {'id': array([0, 1])}
            >>> d2.take_batch()
            {'id': array([2, 3, 4, 5, 6])}
            >>> d3.take_batch()
            {'id': array([7, 8, 9])}

        Time complexity: O(num splits)

        Args:
            proportions: List of proportions to split the dataset according to.
                Must sum up to less than 1, and each proportion must be bigger
                than 0.

        Returns:
            The dataset splits.