Hotfix/arff (#1388)

* Allow skipping parquet download through environment variable

* Allow skip of parquet file, fix bug if no pq file is returned

* Declare the environment file in config.py

Author: PGijsbers

openml/config.py

@@ -23,6 +23,7 @@
 file_handler: logging.handlers.RotatingFileHandler | None = None
 
 OPENML_CACHE_DIR_ENV_VAR = "OPENML_CACHE_DIR"
+OPENML_SKIP_PARQUET_ENV_VAR = "OPENML_SKIP_PARQUET"
 
 
 class _Config(TypedDict):

openml/datasets/dataset.py

@@ -3,6 +3,7 @@
 
 import gzip
 import logging
+import os
 import pickle
 import re
 import warnings
@@ -17,6 +18,7 @@
 import xmltodict
 
 from openml.base import OpenMLBase
+from openml.config import OPENML_SKIP_PARQUET_ENV_VAR
 from openml.exceptions import PyOpenMLError
 
 from .data_feature import OpenMLDataFeature
@@ -358,8 +360,10 @@ def _download_data(self) -> None:
         # import required here to avoid circular import.
         from .functions import _get_dataset_arff, _get_dataset_parquet
 
-        if self._parquet_url is not None:
-            self.parquet_file = str(_get_dataset_parquet(self))
+        skip_parquet = os.environ.get(OPENML_SKIP_PARQUET_ENV_VAR, "false").casefold() == "true"
+        if self._parquet_url is not None and not 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:
             self.data_file = str(_get_dataset_arff(self))
 

openml/datasets/functions.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import logging
+import os
 import warnings
 from collections import OrderedDict
 from pathlib import Path
@@ -20,6 +21,7 @@
 
 import openml._api_calls
 import openml.utils
+from openml.config import OPENML_SKIP_PARQUET_ENV_VAR
 from openml.exceptions import (
     OpenMLHashException,
     OpenMLPrivateDatasetError,
@@ -560,20 +562,22 @@ def get_dataset(  # noqa: C901, PLR0912
         if download_qualities:
             qualities_file = _get_dataset_qualities_file(did_cache_dir, dataset_id)
 
-        if "oml:parquet_url" in description and download_data:
+        parquet_file = None
+        skip_parquet = os.environ.get(OPENML_SKIP_PARQUET_ENV_VAR, "false").casefold() == "true"
+        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(
                     description,
                     download_all_files=download_all_files,
                 )
             except urllib3.exceptions.MaxRetryError:
                 parquet_file = None
-        else:
-            parquet_file = None
 
         arff_file = None
         if parquet_file is None and download_data:
-            logger.warning("Failed to download parquet, fallback on ARFF.")
+            if download_parquet:
+                logger.warning("Failed to download parquet, fallback on ARFF.")
             arff_file = _get_dataset_arff(description)
 
         remove_dataset_cache = False