Issue Link / Problem Description

• No existing issue tracked this; describing it here per the template.
• SingleMetricAnnotation.sample(n, stratify_key=...) can return more than n samples, violating the documented contract ("A subset of the dataset with n samples").

Root cause: src/ragas/dataset_schema.py:719 computes each class's count as int(np.round(proportion * n)). With two balanced classes and an odd n where the half rounds up (e.g. n=3 → np.round(0.5 * 3) = 2 per class), the per-class counts sum to n + 1. The rounding-discrepancy loop at src/ragas/dataset_schema.py:726 only tops up shortfalls (while len(sampled_indices) < n); it never trims an overshoot. So on a balanced dataset sample(3, stratify_key=...) returns 4 samples, sample(19, ...) returns 20, etc.

Reproduce on main:

from ragas.dataset_schema import SingleMetricAnnotation

ds = SingleMetricAnnotation(name="metric", samples=[...])  # 10 samples with metric_output=1, 10 with 0
print(len(ds.sample(3, stratify_key="metric_output")))  # 4 — expected 3

Changes Made

• src/ragas/dataset_schema.py: after the shortfall top-up loop, trim sampled_indices down to n when the rounded per-class counts overshot, so the stratified path returns exactly n samples as documented. (Non-stratified path already used random.sample(..., n) and was correct.)

Testing

How to Test

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

Added test_loader_sample_stratified_returns_exactly_n in tests/unit/test_dataset_schema.py, parametrized over n ∈ [1, 2, 3, 4, 5, 7, 9, 11, 15, 19] on a balanced 10/10 dataset — covering both the overshoot cases (3, 7, 11, 15, 19) and the exact/shortfall cases.

Commands and results:

• New test on unpatched main: pytest tests/unit/test_dataset_schema.py -k stratified_returns_exactly_n → 5 failed, 5 passed (e.g. assert 4 == 3 … assert 20 == 19)
• New test with this fix: → 10 passed
• Module tests: pytest tests/unit/test_dataset_schema.py → 28 passed
• Full unit suite (as CI runs it): pytest --nbmake tests/unit → 779 passed, 105 skipped, 3 failed, 8 errors — those 3 failures + 8 errors are pre-existing on unpatched main (missing optional deps: sacrebleu, haystack, …) and unrelated to this change
• Lint: ruff check src tests docs → All checks passed!; ruff format --check on the two touched files → already formatted
• Type: pyright src/ragas/dataset_schema.py → 0 errors, 0 warnings

References

• Related issues: none found (searched open issues/PRs; no report of this bug)
• Documentation: SingleMetricAnnotation.sample docstring in src/ragas/dataset_schema.py
• External references:

Screenshots/Examples (if applicable)

n/a