diff --git a/openml/utils/_openml.py b/openml/utils/_openml.py index 125dcd08b..e41c3fb90 100644 --- a/openml/utils/_openml.py +++ b/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}") diff --git a/tests/test_utils/test_listing_pagination.py b/tests/test_utils/test_listing_pagination.py new file mode 100644 index 000000000..4884c1a39 --- /dev/null +++ b/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() + )