From b9c36a43a9f480712a054a88483e71651870da04 Mon Sep 17 00:00:00 2001 From: Jayant-kernel Date: Wed, 23 Sep 2026 19:18:01 +0530 Subject: [PATCH] [ENH] Make Parquet downloads configurable --- openml/_config.py | 10 +++++++- openml/datasets/dataset.py | 6 +---- openml/datasets/functions.py | 6 ++--- tests/test_datasets/test_dataset_functions.py | 24 ++++++++++++++++++- tests/test_openml/test_config.py | 18 ++++++++++++-- 5 files changed, 51 insertions(+), 13 deletions(-) diff --git a/openml/_config.py b/openml/_config.py index d6850b7b0..f804e5e5a 100644 --- a/openml/_config.py +++ b/openml/_config.py @@ -119,6 +119,7 @@ class OpenMLConfig: retry_policy: Literal["human", "robot"] = "human" connection_n_retries: int = 5 show_progress: bool = False + skip_parquet: bool = False @property def server(self) -> str: @@ -333,6 +334,12 @@ def set_retry_policy( ), ) + def should_skip_parquet(self) -> bool: + """Return whether dataset downloads should prefer ARFF over Parquet.""" + return self.skip_parquet or ( + os.environ.get(self.OPENML_SKIP_PARQUET_ENV_VAR, "false").casefold() == "true" + ) + def _handle_xdg_config_home_backwards_compatibility(self, xdg_home: str) -> Path: config_dir = Path(xdg_home) / "openml" @@ -402,7 +409,7 @@ def _parse_config(self, config_file: str | Path) -> dict[str, Any]: config_file_.seek(0) config.read_file(config_file_) configuration = dict(config.items("FAKE_SECTION")) - for boolean_field in ["avoid_duplicate_runs", "show_progress"]: + for boolean_field in ["avoid_duplicate_runs", "show_progress", "skip_parquet"]: if isinstance(config["FAKE_SECTION"][boolean_field], str): configuration[boolean_field] = config["FAKE_SECTION"].getboolean(boolean_field) # type: ignore return configuration # type: ignore @@ -440,6 +447,7 @@ def _setup(self, config: dict[str, Any] | None = None) -> None: avoid_duplicate_runs=config["avoid_duplicate_runs"], retry_policy=config["retry_policy"], connection_n_retries=int(config["connection_n_retries"]), + skip_parquet=config["skip_parquet"], ) if "server" in config: self._config.server = config["server"] diff --git a/openml/datasets/dataset.py b/openml/datasets/dataset.py index 59d6205ba..4e9e0db20 100644 --- a/openml/datasets/dataset.py +++ b/openml/datasets/dataset.py @@ -3,7 +3,6 @@ import gzip import logging -import os import pickle import re import warnings @@ -375,10 +374,7 @@ def _download_data(self) -> None: # import required here to avoid circular import. from .functions import _get_dataset_arff, _get_dataset_parquet - skip_parquet = ( - os.environ.get(openml.config.OPENML_SKIP_PARQUET_ENV_VAR, "false").casefold() == "true" - ) - if self._parquet_url is not None and not skip_parquet: + if self._parquet_url is not None and not openml.config.should_skip_parquet(): parquet_file = _get_dataset_parquet(self) self.parquet_file = None if parquet_file is None else str(parquet_file) if self.parquet_file is None: diff --git a/openml/datasets/functions.py b/openml/datasets/functions.py index 7faaf5693..730ac95eb 100644 --- a/openml/datasets/functions.py +++ b/openml/datasets/functions.py @@ -3,7 +3,6 @@ from __future__ import annotations import logging -import os import warnings from collections import OrderedDict from functools import partial @@ -504,10 +503,9 @@ def get_dataset( # noqa: C901, PLR0912 qualities_file = _get_dataset_qualities_file(did_cache_dir, dataset_id) parquet_file = None - skip_parquet = ( - os.environ.get(openml.config.OPENML_SKIP_PARQUET_ENV_VAR, "false").casefold() == "true" + download_parquet = ( + "oml:parquet_url" in description and not openml.config.should_skip_parquet() ) - download_parquet = "oml:parquet_url" in description and not skip_parquet if download_parquet and (download_data or download_all_files): try: parquet_file = _get_dataset_parquet( diff --git a/tests/test_datasets/test_dataset_functions.py b/tests/test_datasets/test_dataset_functions.py index 80b0b4215..c490418d2 100644 --- a/tests/test_datasets/test_dataset_functions.py +++ b/tests/test_datasets/test_dataset_functions.py @@ -1982,6 +1982,28 @@ def test__get_dataset_parquet_not_cached(): assert path.is_file(), "_get_dataset_parquet returns path to real file" +def test_download_data_skips_parquet_when_configured(monkeypatch): + monkeypatch.delenv(openml.config.OPENML_SKIP_PARQUET_ENV_VAR, raising=False) + previous_value = openml.config.skip_parquet + openml.config.skip_parquet = True + dataset = mock.Mock(_parquet_url="https://example.com/dataset.pq", parquet_file=None) + + try: + with ( + mock.patch("openml.datasets.functions._get_dataset_parquet") as get_parquet, + mock.patch( + "openml.datasets.functions._get_dataset_arff", return_value=Path("dataset.arff") + ) as get_arff, + ): + OpenMLDataset._download_data(dataset) + + get_parquet.assert_not_called() + get_arff.assert_called_once_with(dataset) + assert dataset.data_file == str(Path("dataset.arff")) + finally: + openml.config.skip_parquet = previous_value + + def test_read_features_from_xml_with_whitespace() -> None: from openml.datasets.dataset import _read_features @@ -2005,4 +2027,4 @@ def test_get_dataset_parquet(requests_mock, test_files_directory, test_server_v1 assert dataset._parquet_url is not None assert dataset.parquet_file is not None assert os.path.isfile(dataset.parquet_file) - assert dataset.data_file is None # is alias for arff path \ No newline at end of file + assert dataset.data_file is None # is alias for arff path diff --git a/tests/test_openml/test_config.py b/tests/test_openml/test_config.py index 941af9f1c..0d41c6242 100644 --- a/tests/test_openml/test_config.py +++ b/tests/test_openml/test_config.py @@ -87,8 +87,9 @@ def test_get_config_as_dict(self): _config["connection_n_retries"] = 20 _config["retry_policy"] = "robot" _config["show_progress"] = False + _config["skip_parquet"] = False assert isinstance(config, dict) - assert len(config) == 8 + assert len(config) == 9 self.assertDictEqual(config, _config) def test_setup_with_config(self): @@ -102,6 +103,7 @@ def test_setup_with_config(self): _config["retry_policy"] = "human" _config["connection_n_retries"] = 100 _config["show_progress"] = False + _config["skip_parquet"] = True orig_config = openml.config.get_config_as_dict() openml.config._setup(_config) updated_config = openml.config.get_config_as_dict() @@ -172,7 +174,7 @@ def test_configuration_file_not_overwritten_on_load(): def test_configuration_loads_booleans(tmp_path): - config_file_content = "avoid_duplicate_runs=true\nshow_progress=false" + config_file_content = "avoid_duplicate_runs=true\nshow_progress=false\nskip_parquet=true" tmp_file = tmp_path / "config" with tmp_file.open("w") as config_file: config_file.write(config_file_content) @@ -181,6 +183,18 @@ def test_configuration_loads_booleans(tmp_path): # Explicit test to avoid truthy/falsy modes of other types assert read_config["avoid_duplicate_runs"] is True assert read_config["show_progress"] is False + assert read_config["skip_parquet"] is True + + +def test_should_skip_parquet_uses_configuration(monkeypatch): + monkeypatch.delenv(openml.config.OPENML_SKIP_PARQUET_ENV_VAR, raising=False) + previous_value = openml.config.skip_parquet + openml.config.skip_parquet = True + + try: + assert openml.config.should_skip_parquet() is True + finally: + openml.config.skip_parquet = previous_value def test_openml_cache_dir_env_var(tmp_path: Path) -> None: