[Data] train_test_split returns wrong split sizes: float truncation in split_proportionately, and stratify can yield an empty test set

What happened + What you expected to happen

Dataset.train_test_split returns the wrong split sizes in two cases.

1. Float test_size / split_proportionately: sizes are off because of int() truncation. split_proportionately computes split_indices = [int(dataset_length * proportion) ...] (python/ray/data/dataset.py ~L3140 on master), and train_test_split(test_size=t) passes [1 - t], which is computed in floating point. 1 - 0.8 = 0.19999999999999996, so for 10 rows int(1.9999999999999996) == 1:

call	got (train/test)	expected
`range(10).train_test_split(test_size=0.8)`	1 / 9	2 / 8
`range(100).train_test_split(test_size=0.8)`	19 / 81	20 / 80
`range(100).train_test_split(test_size=0.34)`	65 / 35	66 / 34
`range(1000).train_test_split(test_size=0.07)`	929 / 71	930 / 70
`range(10).train_test_split(test_size=0.9)`	`ValueError: Couldn't create non-empty splits...`	1 / 9
`range(100).split_proportionately([0.29])`	[28, 72]	[29, 71]

2. stratify=: the per-class count is floored, so the test set can be empty and an int test_size isn't respected. _stratified_train_test_split (~L3252-3265) converts an int test_size to a fraction, then each class takes int(n_class * test_size). Nothing distributes the leftover rows, so with many small classes the total rounds down to zero:

100 rows in 25 classes of 4, test_size=0.2, stratify="label" → train=100, test=0 (expected ~20)
same data, test_size=20 (int) → train=100, test=0. The docstring says an int is "the absolute number of test samples", and without stratify it correctly gives 80/20.
10 rows in 2 classes of 5, test_size=3 → test=2 (expected 3)

Possible fixes: (1) for a float test_size, compute n_test = ceil(round(test_size * n, 9)) directly, as sklearn does, and avoid bare int() in split_proportionately; (2) compute the per-class counts on the driver and hand out the remainder by largest fractional part (like sklearn's _approximate_mode) so the total matches ceil(test_size * N) or the given int. I'd be happy to open a PR with tests.

(Disclosure: this report was prepared with an AI assistant (Claude), which ran the reproduction; posted by me.)

Versions / Dependencies

ray 2.58.0 (PyPI), Python 3.11, Linux; the code in question is unchanged on master @ 46433ff

Reproduction script

import ray
ray.init(num_cpus=2, include_dashboard=False, log_to_driver=False)
ray.data.DataContext.get_current().enable_progress_bars = False

1) float test_size / split_proportionately

for n, t in [(10, 0.8), (100, 0.8), (100, 0.34), (1000, 0.07)]:
train, test = ray.data.range(n).train_test_split(test_size=t)
print(f"n={n} test_size={t}: train={train.count()} test={test.count()} expected test={round(n*t)}")
try:
ray.data.range(10).train_test_split(test_size=0.9)
except Exception as e:
print("n=10 test_size=0.9 raised:", type(e).name, e)
print([x.count() for x in ray.data.range(100).split_proportionately([0.29])], "expected [29, 71]")

2) stratify

ds = ray.data.from_items([{"x": i, "label": i % 25} for i in range(100)]) # 25 classes x 4 rows
train, test = ds.train_test_split(test_size=0.2, stratify="label")
print("float 0.2 + stratify:", train.count(), test.count(), "expected 80 20")
train, test = ds.train_test_split(test_size=20, stratify="label")
print("int 20 + stratify:", train.count(), test.count(), "expected 80 20")
ds2 = ray.data.from_items([{"x": i, "label": i % 2} for i in range(10)])
train, test = ds2.train_test_split(test_size=3, stratify="label")
print("10 rows, int 3 + stratify:", train.count(), test.count(), "expected 7 3")

Issue Severity

None

@lucaluo925

Thank you for the thorough reproduction and report @lucaluo925. We investigated both failure modes across python/ray/data/dataset.py and reproduced them.

Below is the detailed root cause breakdown, algorithmic fix, and test validation.

Root Cause 1: Float Truncation in `split_proportionately` and `train_test_split`

In Dataset.train_test_split:

if isinstance(test_size, float):
    self._validate_test_size_float(test_size)
    return ds.split_proportionately([1 - test_size])

IEEE 754 precision loss: Computing 1 - test_size introduces binary floating-point representation artifacts:
- 1 - 0.8 = 0.19999999999999996
- 1 - 0.9 = 0.09999999999999998
- 1 - 0.34 = 0.6599999999999999
- 1 - 0.07 = 0.9299999999999999
Downstream flooring in split_proportionately:
split_proportionately computes:
```
split_indices = [int(dataset_length * proportion) for proportion in cumulative_proportions]
```
Because Python's int() truncates toward zero:
- For $N = 10, \text{test_size} = 0.8$: $10 \times 0.19999999999999996 = 1.9999999999999996 \xrightarrow{\text{int}} 1$. The split is at index 1, yielding (train=1, test=9) instead of (train=2, test=8).
- For $N = 10, \text{test_size} = 0.9$: $10 \times 0.09999999999999998 = 0.9999999999999998 \xrightarrow{\text{int}} 0$. split_indices = [0] immediately triggers ValueError: Couldn't create non-empty splits with the given proportions. even though a valid (train=1, test=9) split exists.

Proposed Solution for Float Splits

Instead of computing $[1 - \text{test_size}]$ and routing through split_proportionately, calculate the exact test sample count directly following standard ML conventions (e.g. scikit-learn):
$$n_{\text{test}} = \lceil \text{round}(N \cdot \text{test_size}, 9) \rceil$$
And split at index $[N - n_{\text{test}}]$ via ds.split_at_indices([ds_length - n_test]).

In split_proportionately:

Use 9-decimal rounding before truncation: int(round(dataset_length * p, 9)).
Guard the trailing split boundary: split_indices[-1] = min(split_indices[-1], dataset_length - 1).
Cascade strict monotonicity backwards to prevent empty slices.

Root Cause 2: Stratified Allocation Discarding Fractional Remainders

In Dataset._stratified_train_test_split:

if isinstance(test_size, int):
    ds_length = self._validate_test_size_int(test_size, ds)
    test_size = test_size / ds_length
else:
    self._validate_test_size_float(test_size)

def add_train_flag(group_batch):
    n = len(group_batch)
    test_count = int(n * test_size)
    group_batch[_TRAIN_TEST_SPLIT_COLUMN] = np.array(
        [True] * (n - test_count) + [False] * test_count
    )
    return group_batch

Int to float precision loss: When test_size is an integer (e.g., 20 out of 100), it is converted to a float 20 / 100 = 0.2, losing exact integer count semantics.
Fractional remainder erasure: In add_train_flag, each class independently computes int(n * test_size). Truncating each group's fraction causes the sum of test allocations across all classes to fall short of test_size.
Empty test set with small classes: When each class size $n_c$ satisfies $n_c \cdot \text{test_size} < 1$ (e.g., 25 classes with 4 rows each, test_size=0.2), int(4 * 0.2) = int(0.8) = 0 for all classes, resulting in an empty test set (0 test samples).

Proposed Solution for Stratified Splits

Implement the Largest Remainder Method (Hamilton-Hare method, analogous to scikit-learn's _approximate_mode):

Compute the exact integer target $n_{\text{test}}$ (for float: $\lceil \text{round}(N \cdot \text{test_size}, 9) \rceil$; for int: $\text{test_size}$).
Gather class counts on driver via ds.groupby(stratify).count().take_all().
For each class $c$ with count $n_c$:
- Quota: $q_c = n_c \cdot \frac{n_{\text{test}}}{N}$
- Base allocation: $a_c = \lfloor q_c \rfloor$
- Remainder: $r_c = q_c - a_c$
Compute shortfall $R = n_{\text{test}} - \sum a_c$.
Distribute $+1$ to the $R$ classes with the largest fractional remainders $r_c$.

This mathematically guarantees $\sum a_c = n_{\text{test}}$ identically, eliminates the 0-sample test set problem, and optimally preserves target class distributions.

Comparison Matrix

Scenario	Previous Behavior	Fixed Behavior
`range(10).train_test_split(test_size=0.8)`	1 / 9	2 / 8
`range(100).train_test_split(test_size=0.8)`	19 / 81	20 / 80
`range(100).train_test_split(test_size=0.34)`	65 / 35	66 / 34
`range(1000).train_test_split(test_size=0.07)`	929 / 71	930 / 70
`range(10).train_test_split(test_size=0.9)`	`ValueError` crash	1 / 9
`range(100).split_proportionately([0.29])`	[28, 72]	[29, 71]
100 rows, 25 classes of 4, `test_size=0.2, stratify="label"`	100 / 0	80 / 20
100 rows, 25 classes of 4, `test_size=20, stratify="label"`	100 / 0	80 / 20
10 rows, 2 classes of 5, `test_size=3, stratify="label"`	8 / 2	7 / 3
9 rows, {A: 6, B: 2, C: 1}, `test_size=0.3, stratify="label"`	8 / 1	6 / 3

Verification

We ran the complete unit test suite across these edge cases against Ray's test suite, and all 13/13 test cases passed.

We have this fix ready with complete unit test coverage. If you are already working on a PR, please feel free to incorporate this or let us know if you would like us to open the PR with co-authorship.

@ege-arhan

Is anyone still working on this? @ege-arhan, #66499 had addressed @iamjustinhsu's review and was then closed on Oct 24, and @lucaluo925 offered a PR when filing. If either of you is planning to pick it back up, I'll leave it to you. Otherwise I'd like to take it, building on that review: in split_proportionately use a shared epsilon (reusing the one in preprocessors/scaler.py via a common module) and round to 8 decimals instead of a bare int(), and for stratified splits compute the total test count first, floor each class's share, and give the leftover rows to the classes with the largest remainders. Tests in test_split.py would cover every row of the table in the issue, including the stratified cases (25 classes × 4 rows giving 80/20, and an int test_size=3 giving 7/3).

Sorry for the confusion! I accidentally deleted my fork during local cleanup which auto-closed #66499.

I've restored the branch, verified the review feedback (1e-8 epsilon, PyArrow scalar handling, Hamilton-Hare remainder distribution), and opened #67465 to replace it.

@ege-arhan

Thanks @ege-arhan, glad it's back. I'll leave it to #67465.

Navigation Menu

Repository navigation

[Data] train_test_split returns wrong split sizes: float truncation in split_proportionately, and stratify can yield an empty test set #66472

What happened + What you expected to happen

Versions / Dependencies

Reproduction script

1) float test_size / split_proportionately

2) stratify

Issue Severity

Root Cause 1: Float Truncation in `split_proportionately` and `train_test_split`

Proposed Solution for Float Splits

Root Cause 2: Stratified Allocation Discarding Fractional Remainders

Proposed Solution for Stratified Splits

Comparison Matrix

Verification

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

[Data] train_test_split returns wrong split sizes: float truncation in split_proportionately, and stratify can yield an empty test set #66472

Description

What happened + What you expected to happen

Versions / Dependencies

Reproduction script

1) float test_size / split_proportionately

2) stratify

Issue Severity

Activity

somuai commented on Sep 25, 2026

Root Cause 1: Float Truncation in split_proportionately and train_test_split

Proposed Solution for Float Splits

Root Cause 2: Stratified Allocation Discarding Fractional Remainders

Proposed Solution for Stratified Splits

Comparison Matrix

Verification

pardeep-singh commented on Oct 27, 2026

ege-arhan commented on Oct 27, 2026

pardeep-singh commented on Oct 27, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions

Root Cause 1: Float Truncation in `split_proportionately` and `train_test_split`