fix: correct batch count and non-terminating loop in stratified_batches - #3036

Issue Link / Problem Description

• No existing issue tracked this; describing it here per the template.
• SingleMetricAnnotation.stratified_batches() computes the number of batches with ceil/floor inverted (src/ragas/dataset_schema.py:792-796 on main): it uses ceil(total_samples / batch_size) when drop_last_batch=True and floor(...) when False — the opposite of the documented behavior ("If True, drops the last batch if it has fewer samples than batch_size" → only complete batches fit → floor) and of what the non-stratified batch() method does. Consequences on main:
  • drop_last_batch=True never terminates. The while len(all_batches) < num_batches loop (src/ragas/dataset_schema.py:808) only appends a batch when it reaches batch_size; with ceil demanding more complete batches than the data can fill (e.g. 20 samples, batch_size=3, or a class running out mid-way at batch_size=2), the loop spins forever.
  • drop_last_batch=False silently drops samples. This is the mode the genetic optimizer uses (src/ragas/optimizers/genetic.py:282): 20 samples with batch_size=2 return 19; and when batch_size exceeds the dataset size, floor() yields 0 batches and every sample is dropped (returns []).
  • replace=True crashes: it calls random.choices() on an exhausted class's empty list and raises IndexError.

Reproduce on main:

from ragas.dataset_schema import SampleAnnotation, SingleMetricAnnotation

def make(mo):
    return SampleAnnotation(metric_input={"user_input": ""}, metric_output=mo, prompts={}, is_accepted=True)

ds = SingleMetricAnnotation(name="metric", samples=[make(1) for _ in range(11)] + [make(0) for _ in range(9)])

ds.stratified_batches(batch_size=3, stratify_key="metric_output", drop_last_batch=True)  # hangs forever
len([s for b in ds.stratified_batches(batch_size=2, stratify_key="metric_output") for s in b])  # 19, not 20

Changes Made

• src/ragas/dataset_schema.py (only stratified_batches):
  • Compute num_batches the right way round: floor when drop_last_batch=True, ceil when False — mirroring batch().
  • Guard the zero-batch case (batch_size > dataset size with drop_last_batch=True) → returns [] instead of dividing/looping over zero batches.
  • Replace the quota-based while loop (which consumed samples for batches it then refused to append) with a round-robin dealing of each class's samples across the batches: every sample lands in a batch when the partial last batch is kept, and only complete batches are returned when it is dropped. Each batch still gets its proportional share of every class, and batches are still shuffled.
  • replace=True now tops up incomplete batches by resampling from the samples already in the batch, instead of calling random.choices() on an exhausted class (which raised IndexError). The replace docstring line is updated to match what the code now does.
• tests/unit/test_dataset_schema.py: four regression tests (see Testing).

Testing

How to Test

• Automated tests added/updated
• Manual testing steps:
  1. n/a — covered by the automated regression tests

New tests in tests/unit/test_dataset_schema.py (all fail on unpatched main, pass with this fix):

• test_stratified_batches_without_drop_last_batch_uses_all_samples — 20 samples (11/9 classes), batch_size=2: all 20 samples batched exactly once in 10 batches (main returns 19).
• test_stratified_batches_with_drop_last_batch_terminates — runs in a subprocess with a timeout since main hangs: 20 samples with batch_size=2 → 10 complete batches; batch_size=3 → 6 complete batches (partial tail dropped); batch_size=25 → [] (no complete batch fits).
• test_stratified_batches_with_fewer_samples_than_batch_size — batch_size=25 over 20 samples with drop_last_batch=False → one batch with all 20 samples (main returns [], dropping everything).
• test_stratified_batches_with_replace_fills_batches — batch_size=3, replace=True → 7 batches all filled to 3 (main returns 6 batches of 2, or raises IndexError).

Commands and results (venv with ragas installed editable; same as make test / CI):

• New tests on unpatched main: pytest tests/unit/test_dataset_schema.py -k stratified_batches → 4 failed, 18 deselected (AssertionError: assert 19 == 20; Failed: stratified_batches(drop_last_batch=True) did not terminate; assert 0 == 1; assert 6 == 7)
• New tests with this fix: → 4 passed
• Module: pytest tests/unit/test_dataset_schema.py → 22 passed (incl. the pre-existing test_loader_batch, which also exercises stratified_batches)
• CI essential suite: pytest tests/unit -k "not slow" --dist loadfile -n auto → 3 failed, 772 passed, 105 skipped, 8 errors; the 3 failures + 8 errors are pre-existing on unpatched main with identical test IDs (missing optional deps sacrebleu/haystack, test_prechunked_generation, test_experiment)
• Lint: ruff format --check on the two touched files → already formatted; ruff check src docs tests → All checks passed!
• Type: pyright src/ragas/dataset_schema.py → 0 errors, 0 warnings, 0 informations

References

• Related issues: none found (searched open issues/PRs; no report of this bug)
• Related PRs: fix: ensure stratified sample returns exactly n samples #3034 touches the same file but fixes sample() (a different method) — no overlap with this change
• Documentation: SingleMetricAnnotation.stratified_batches docstring in src/ragas/dataset_schema.py
• External references:

Screenshots/Examples (if applicable)

n/a