[ENH] Make Parquet downloads configurable

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"]

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:

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(

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
+    assert dataset.data_file is None  # is alias for arff path

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: