[FIX] Cap listing request batches at 1000

openml/utils/_openml.py

@@ -27,6 +27,8 @@
 import openml._api_calls
 import openml.exceptions
 
+MAX_LISTING_BATCH_SIZE = 1_000
+
 # Avoid import cycles: https://mypy.readthedocs.io/en/latest/common_issues.html#import-cycles
 if TYPE_CHECKING:
     from collections.abc import Callable, Iterable, Mapping, Sequence, Sized
@@ -251,7 +253,7 @@ def _list_all(  # noqa: C901
     *,
     limit: int | None = None,
     offset: int | None = None,
-    batch_size: int | None = 10_000,
+    batch_size: int | None = MAX_LISTING_BATCH_SIZE,
 ) -> list[_SizedT]:
     """Helper to handle paged listing requests.
 
@@ -265,7 +267,7 @@ def _list_all(  # noqa: C901
         Call listing, e.g. list_evaluations. Takes two positional
         arguments: batch_size and offset.
     batch_size : int, optional
-        The batch size to use for the listing call.
+        The batch size to use for the listing call. Requests are capped at 1,000 items.
     offset : int, optional
         The initial offset to use for the listing call.
     limit : int, optional
@@ -281,7 +283,7 @@ def _list_all(  # noqa: C901
     results: list[_SizedT] = []
 
     offset = offset if offset is not None else 0
-    batch_size = batch_size if batch_size is not None else 10_000
+    batch_size = batch_size if batch_size is not None else MAX_LISTING_BATCH_SIZE
 
     LIMIT = limit
     BATCH_SIZE_ORIG = batch_size
@@ -292,6 +294,8 @@ def _list_all(  # noqa: C901
     if not isinstance(BATCH_SIZE_ORIG, int):
         raise ValueError(f"'batch_size' should be an integer but got {BATCH_SIZE_ORIG}")
 
+    BATCH_SIZE_ORIG = min(BATCH_SIZE_ORIG, MAX_LISTING_BATCH_SIZE)
+
     if (LIMIT is not None) and (not isinstance(LIMIT, int)) and (not np.isinf(LIMIT)):
         raise ValueError(f"'limit' should be an integer or inf but got {LIMIT}")
 

tests/test_utils/test_listing_pagination.py

@@ -0,0 +1,200 @@
+# License: BSD 3-Clause
+# mypy: disable-error-code=untyped-decorator
+from __future__ import annotations
+
+import unittest.mock
+from typing import Any, Literal
+
+import pandas as pd
+import pytest
+
+import openml
+
+
+@pytest.mark.parametrize("batch_size", [None, 10_000])
+def test_list_all_caps_batch_size(batch_size: int | None) -> None:
+    calls: list[tuple[int, int]] = []
+
+    def listing_call(size: int, offset: int) -> list[int]:
+        calls.append((size, offset))
+        return list(range(offset, min(offset + size, 2_500)))
+
+    kwargs = {} if batch_size is None else {"batch_size": batch_size}
+    batches = openml.utils._list_all(listing_call, limit=2_500, **kwargs)
+
+    assert calls == [(1_000, 0), (1_000, 1_000), (500, 2_000)]
+    assert sum(len(batch) for batch in batches) == 2_500
+
+
+@pytest.mark.parametrize(
+    ("batch_size", "limit", "offset", "expected_calls"),
+    [
+        (None, 2_500, 125, [(1_000, 125), (1_000, 1_125), (500, 2_125)]),
+        (10_000, 2_500, 125, [(1_000, 125), (1_000, 1_125), (500, 2_125)]),
+        (None, 250, 125, [(250, 125)]),
+        (10_000, 2_000, 125, [(1_000, 125), (1_000, 1_125)]),
+        (400, 1_050, 125, [(400, 125), (400, 525), (250, 925)]),
+    ],
+)
+def test_list_all_preserves_limit_and_offset(
+    batch_size: int | None,
+    limit: int,
+    offset: int,
+    expected_calls: list[tuple[int, int]],
+) -> None:
+    listing_call = unittest.mock.Mock(
+        side_effect=lambda size, start: list(range(start, start + size)),
+    )
+
+    batches = openml.utils._list_all(
+        listing_call,
+        batch_size=batch_size,
+        limit=limit,
+        offset=offset,
+    )
+
+    assert listing_call.call_args_list == [unittest.mock.call(*call) for call in expected_calls]
+    assert [item for batch in batches for item in batch] == list(range(offset, offset + limit))
+
+
+@pytest.mark.parametrize("limit", [None, 5_000])
+@pytest.mark.parametrize("available", [0, 1_000, 1_250])
+def test_list_all_stops_when_results_end(limit: int | None, available: int) -> None:
+    offset = 125
+    listing_call = unittest.mock.Mock(
+        side_effect=lambda size, start: list(range(start, min(start + size, offset + available))),
+    )
+
+    batches = openml.utils._list_all(
+        listing_call,
+        batch_size=10_000,
+        limit=limit,
+        offset=offset,
+    )
+
+    expected_calls = [unittest.mock.call(1_000, offset)]
+    if available >= 1_000:
+        expected_calls.append(unittest.mock.call(1_000, offset + 1_000))
+    assert listing_call.call_args_list == expected_calls
+    assert [item for batch in batches for item in batch] == list(range(offset, offset + available))
+
+
+def test_list_all_stops_on_no_result() -> None:
+    listing_call = unittest.mock.Mock(
+        side_effect=[list(range(125, 1_125)), openml.exceptions.OpenMLServerNoResult("No results")],
+    )
+
+    batches = openml.utils._list_all(listing_call, batch_size=10_000, offset=125)
+
+    assert listing_call.call_args_list == [
+        unittest.mock.call(1_000, 125),
+        unittest.mock.call(1_000, 1_125),
+    ]
+    assert batches == [list(range(125, 1_125))]
+
+
+def test_list_datasets_pages_preserve_filters() -> None:
+    def listing_call(size: int, offset: int, **filters: Any) -> pd.DataFrame:
+        assert filters["tag"] == "pagination-test"
+        assert filters["status"] == "active"
+        assert filters["data_name"] == "example"
+        return pd.DataFrame({"did": range(offset, offset + size)})
+
+    with unittest.mock.patch(
+        "openml.datasets.functions._list_datasets",
+        side_effect=listing_call,
+    ) as page_fetch:
+        datasets = openml.datasets.list_datasets(
+            size=2_500,
+            offset=125,
+            tag="pagination-test",
+            status="active",
+            data_name="example",
+        )
+
+    assert [call.args for call in page_fetch.call_args_list] == [
+        (1_000, 125),
+        (1_000, 1_125),
+        (500, 2_125),
+    ]
+    assert isinstance(datasets, pd.DataFrame)
+    assert datasets["did"].tolist() == list(range(125, 2_625))
+
+
+def test_list_runs_pages_preserve_filters() -> None:
+    def listing_call(size: int, offset: int, **filters: Any) -> pd.DataFrame:
+        assert filters["tag"] == "pagination-test"
+        assert filters["task"] == [42]
+        return pd.DataFrame({"run_id": range(offset, offset + size)})
+
+    with unittest.mock.patch.object(
+        openml._backend.run,
+        "list",
+        side_effect=listing_call,
+    ) as page_fetch:
+        runs = openml.runs.list_runs(size=2_500, offset=125, tag="pagination-test", task=[42])
+
+    assert [call.args for call in page_fetch.call_args_list] == [
+        (1_000, 125),
+        (1_000, 1_125),
+        (500, 2_125),
+    ]
+    assert isinstance(runs, pd.DataFrame)
+    assert runs["run_id"].tolist() == list(range(125, 2_625))
+
+
+@pytest.mark.parametrize("output_format", ["object", "dataframe"])
+def test_list_evaluations_pages_preserve_output(
+    output_format: Literal["object", "dataframe"],
+) -> None:
+    def listing_call(
+        size: int, offset: int, **filters: Any
+    ) -> list[openml.evaluations.OpenMLEvaluation]:
+        assert filters["function"] == "predictive_accuracy"
+        assert filters["tag"] == "pagination-test"
+        return [
+            openml.evaluations.OpenMLEvaluation(
+                run_id=run_id,
+                task_id=42,
+                setup_id=1,
+                flow_id=1,
+                flow_name="example",
+                data_id=1,
+                data_name="example",
+                function="predictive_accuracy",
+                upload_time="2026-01-01",
+                uploader=1,
+                uploader_name="example",
+                value=0.8,
+                values=None,
+            )
+            for run_id in range(offset, offset + size)
+        ]
+
+    with unittest.mock.patch.object(
+        openml._backend.evaluation,
+        "list",
+        side_effect=listing_call,
+    ) as page_fetch:
+        evaluations = openml.evaluations.list_evaluations(
+            "predictive_accuracy",
+            size=2_500,
+            offset=125,
+            tag="pagination-test",
+            output_format=output_format,
+        )
+
+    assert [call.args for call in page_fetch.call_args_list] == [
+        (1_000, 125),
+        (1_000, 1_125),
+        (500, 2_125),
+    ]
+    if output_format == "dataframe":
+        assert isinstance(evaluations, pd.DataFrame)
+        assert evaluations["run_id"].tolist() == list(range(125, 2_625))
+    else:
+        assert isinstance(evaluations, dict)
+        assert list(evaluations) == list(range(125, 2_625))
+        assert all(
+            isinstance(item, openml.evaluations.OpenMLEvaluation) for item in evaluations.values()
+        )