[ENH] Make Parquet downloads configurable

openml/_api/resources/dataset.py

@@ -3,7 +3,6 @@
 
 import json
 import logging
-import os
 import pickle
 from collections import OrderedDict
 from pathlib import Path
@@ -89,10 +88,7 @@ def get(
                 qualities_file = self.download_qualities_file(dataset_id)
 
             parquet_file = None
-            skip_parquet = (
-                os.environ.get(openml.config.OPENML_SKIP_PARQUET_ENV_VAR, "false").casefold()
-                == "true"
-            )
+            skip_parquet = 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:
@@ -934,10 +930,7 @@ def get(
                 qualities_file = self.download_qualities_file(dataset_id)
 
             parquet_file = None
-            skip_parquet = (
-                os.environ.get(openml.config.OPENML_SKIP_PARQUET_ENV_VAR, "false").casefold()
-                == "true"
-            )
+            skip_parquet = openml.config.should_skip_parquet()
             download_parquet = "parquet_url" in json_content and not skip_parquet
             if download_parquet and (download_data or download_all_files):
                 try:

openml/_config.py

@@ -120,6 +120,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:
@@ -337,6 +338,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"
 
@@ -406,7 +413,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
@@ -444,6 +451,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

@@ -4,7 +4,6 @@
 import contextlib
 import gzip
 import logging
-import os
 import pickle
 import re
 import warnings
@@ -377,10 +376,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)
 

tests/test_datasets/test_dataset_functions.py

@@ -1936,6 +1936,58 @@ 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
+
+
+@pytest.mark.parametrize("api_version", [1, 2])
+@pytest.mark.parametrize("skip_parquet", [False, True])
+def test_resource_get_respects_skip_parquet(monkeypatch, api_version, skip_parquet):
+    from openml._api.resources.dataset import DatasetV1API, DatasetV2API
+
+    monkeypatch.delenv(openml.config.OPENML_SKIP_PARQUET_ENV_VAR, raising=False)
+    monkeypatch.setattr(openml.config, "skip_parquet", skip_parquet)
+    resource = mock.Mock()
+    resource._http.get.return_value.text = (
+        '<oml:data_set_description xmlns:oml="http://openml.org/openml">'
+        "<oml:parquet_url>https://example.com/data.pq</oml:parquet_url>"
+        "</oml:data_set_description>"
+    )
+    resource._http.get.return_value.json.return_value = {
+        "parquet_url": "https://example.com/data.pq"
+    }
+    resource.download_dataset_parquet.return_value = Path("dataset.pq")
+    resource.download_dataset_arff.return_value = Path("dataset.arff")
+    api = DatasetV1API if api_version == 1 else DatasetV2API
+
+    api.get(resource, 61, download_data=True)
+
+    if skip_parquet:
+        resource.download_dataset_parquet.assert_not_called()
+        resource.download_dataset_arff.assert_called_once()
+    else:
+        resource.download_dataset_parquet.assert_called_once()
+        resource.download_dataset_arff.assert_not_called()
+
+
 def test_read_features_from_xml_with_whitespace() -> None:
     from openml.datasets.dataset import _read_features
 

tests/test_openml/test_config.py

@@ -86,8 +86,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)  # noqa: PT009
 
     def test_setup_with_config(self):
@@ -101,6 +102,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()
@@ -171,7 +173,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)
@@ -180,6 +182,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: