back to pyarrow 1.0.0 + raise error if using old pyarrow for parquet read/write

Author: lhoestq

.circleci/config.yml

@@ -18,7 +18,7 @@ jobs:
             - run: pip install pyarrow --upgrade
             - run: HF_SCRIPTS_VERSION=master python -m pytest -sv ./tests/
 
-    run_dataset_script_tests_pyarrow_3:
+    run_dataset_script_tests_pyarrow_1:
         working_directory: ~/datasets
         docker:
             - image: circleci/python:3.6
@@ -29,7 +29,7 @@ jobs:
             - run: source venv/bin/activate
             - run: pip install .[tests]
             - run: pip install -r additional-tests-requirements.txt --no-deps
-            - run: pip install pyarrow==3.0.0
+            - run: pip install pyarrow==1.0.0
             - run: HF_SCRIPTS_VERSION=master python -m pytest -sv ./tests/
 
     run_dataset_script_tests_pyarrow_latest_WIN:
@@ -50,7 +50,7 @@ jobs:
             - run: $env:HF_SCRIPTS_VERSION="master"
             - run: python -m pytest -sv ./tests/
 
-    run_dataset_script_tests_pyarrow_3_WIN:
+    run_dataset_script_tests_pyarrow_1_WIN:
         working_directory: ~/datasets
         executor:
             name: win/default
@@ -64,7 +64,7 @@ jobs:
             - run: "& venv/Scripts/activate.ps1"
             - run: pip install .[tests]
             - run: pip install -r additional-tests-requirements.txt --no-deps
-            - run: pip install pyarrow==3.0.0
+            - run: pip install pyarrow==1.0.0
             - run: $env:HF_SCRIPTS_VERSION="master"
             - run: python -m pytest -sv ./tests/
 

setup.py

@@ -73,9 +73,9 @@
     # We use numpy>=1.17 to have np.random.Generator (Dataset shuffling)
     "numpy>=1.17",
     # Backend and serialization.
-    # Minimum 3.0.0 to support mix of struct and list types in parquet format
+    # Minimum 3.0.0 to support mix of struct and list types in parquet, and batch iterators of parquet data
     # pyarrow 4.0.0 introduced segfault bug, see: https://github.com/huggingface/datasets/pull/2268
-    "pyarrow>=3.0.0,!=4.0.0",
+    "pyarrow>=1.0.0,!=4.0.0",
     # For smart caching dataset processing
     "dill",
     # For performance gains with apache arrow

src/datasets/io/parquet.py

@@ -3,6 +3,7 @@
 
 import pyarrow as pa
 import pyarrow.parquet as pq
+from packaging import version
 
 from .. import Dataset, Features, NamedSplit, config
 from ..formatting import query_table
@@ -22,6 +23,10 @@ def __init__(
         keep_in_memory: bool = False,
         **kwargs,
     ):
+        if version.parse(pa.__version__) < version.parse("3.0.0"):
+            raise ImportError(
+                "PyArrow >= 3.0.0 is required to used the ParquetDatasetReader: pip install --upgrade pyarrow"
+            )
         super().__init__(
             path_or_paths, split=split, features=features, cache_dir=cache_dir, keep_in_memory=keep_in_memory, **kwargs
         )
@@ -66,6 +71,10 @@ def __init__(
         batch_size: Optional[int] = None,
         **parquet_writer_kwargs,
     ):
+        if version.parse(pa.__version__) < version.parse("3.0.0"):
+            raise ImportError(
+                "PyArrow >= 3.0.0 is required to used the ParquetDatasetWriter: pip install --upgrade pyarrow"
+            )
         self.dataset = dataset
         self.path_or_buf = path_or_buf
         self.batch_size = batch_size

src/datasets/packaged_modules/parquet/parquet.py

@@ -5,6 +5,7 @@
 
 import pyarrow as pa
 import pyarrow.parquet as pq
+from packaging import version
 
 import datasets
 
@@ -25,6 +26,10 @@ class Parquet(datasets.ArrowBasedBuilder):
     BUILDER_CONFIG_CLASS = ParquetConfig
 
     def _info(self):
+        if version.parse(pa.__version__) < version.parse("3.0.0"):
+            raise ImportError(
+                "PyArrow >= 3.0.0 is required to used the Parquet dataset builder: pip install --upgrade pyarrow"
+            )
         return datasets.DatasetInfo(features=self.config.features)
 
     def _split_generators(self, dl_manager):

tests/io/test_parquet.py

@@ -4,7 +4,7 @@
 from datasets import Dataset, DatasetDict, Features, NamedSplit, Value
 from datasets.io.parquet import ParquetDatasetReader, ParquetDatasetWriter
 
-from ..utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases
+from ..utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases, require_pyarrow_at_least_3
 
 
 def _check_parquet_dataset(dataset, expected_features):
@@ -16,6 +16,7 @@ def _check_parquet_dataset(dataset, expected_features):
         assert dataset.features[feature].dtype == expected_dtype
 
 
+@require_pyarrow_at_least_3
 @pytest.mark.parametrize("keep_in_memory", [False, True])
 def test_dataset_from_parquet_keep_in_memory(keep_in_memory, parquet_path, tmp_path):
     cache_dir = tmp_path / "cache"
@@ -25,6 +26,7 @@ def test_dataset_from_parquet_keep_in_memory(keep_in_memory, parquet_path, tmp_p
     _check_parquet_dataset(dataset, expected_features)
 
 
+@require_pyarrow_at_least_3
 @pytest.mark.parametrize(
     "features",
     [
@@ -46,6 +48,7 @@ def test_dataset_from_parquet_features(features, parquet_path, tmp_path):
     _check_parquet_dataset(dataset, expected_features)
 
 
+@require_pyarrow_at_least_3
 @pytest.mark.parametrize("split", [None, NamedSplit("train"), "train", "test"])
 def test_dataset_from_parquet_split(split, parquet_path, tmp_path):
     cache_dir = tmp_path / "cache"
@@ -55,6 +58,7 @@ def test_dataset_from_parquet_split(split, parquet_path, tmp_path):
     assert dataset.split == str(split) if split else "train"
 
 
+@require_pyarrow_at_least_3
 @pytest.mark.parametrize("path_type", [str, list])
 def test_dataset_from_parquet_path_type(path_type, parquet_path, tmp_path):
     if issubclass(path_type, str):
@@ -78,6 +82,7 @@ def _check_parquet_datasetdict(dataset_dict, expected_features, splits=("train",
             assert dataset.features[feature].dtype == expected_dtype
 
 
+@require_pyarrow_at_least_3
 @pytest.mark.parametrize("keep_in_memory", [False, True])
 def test_parquet_datasetdict_reader_keep_in_memory(keep_in_memory, parquet_path, tmp_path):
     cache_dir = tmp_path / "cache"
@@ -89,6 +94,7 @@ def test_parquet_datasetdict_reader_keep_in_memory(keep_in_memory, parquet_path,
     _check_parquet_datasetdict(dataset, expected_features)
 
 
+@require_pyarrow_at_least_3
 @pytest.mark.parametrize(
     "features",
     [
@@ -110,6 +116,7 @@ def test_parquet_datasetdict_reader_features(features, parquet_path, tmp_path):
     _check_parquet_datasetdict(dataset, expected_features)
 
 
+@require_pyarrow_at_least_3
 @pytest.mark.parametrize("split", [None, NamedSplit("train"), "train", "test"])
 def test_parquet_datasetdict_reader_split(split, parquet_path, tmp_path):
     if split:
@@ -124,6 +131,7 @@ def test_parquet_datasetdict_reader_split(split, parquet_path, tmp_path):
     assert all(dataset[split].split == split for split in path.keys())
 
 
+@require_pyarrow_at_least_3
 def test_parquer_write(dataset, tmp_path):
     writer = ParquetDatasetWriter(dataset, tmp_path / "foo.parquet")
     assert writer.write() > 0

tests/test_arrow_dataset.py

@@ -27,6 +27,7 @@
 from .utils import (
     assert_arrow_memory_doesnt_increase,
     assert_arrow_memory_increases,
+    require_pyarrow_at_least_3,
     require_s3,
     require_tf,
     require_torch,
@@ -1692,6 +1693,7 @@ def test_to_pandas(self, in_memory):
                     for col_name in dset.column_names:
                         self.assertEqual(len(dset_to_pandas[col_name]), dset.num_rows)
 
+    @require_pyarrow_at_least_3
     def test_to_parquet(self, in_memory):
         with tempfile.TemporaryDirectory() as tmp_dir:
             # File path argument
@@ -2677,6 +2679,7 @@ def _check_parquet_dataset(dataset, expected_features):
         assert dataset.features[feature].dtype == expected_dtype
 
 
+@require_pyarrow_at_least_3
 @pytest.mark.parametrize("keep_in_memory", [False, True])
 def test_dataset_from_parquet_keep_in_memory(keep_in_memory, parquet_path, tmp_path):
     cache_dir = tmp_path / "cache"
@@ -2686,6 +2689,7 @@ def test_dataset_from_parquet_keep_in_memory(keep_in_memory, parquet_path, tmp_p
     _check_parquet_dataset(dataset, expected_features)
 
 
+@require_pyarrow_at_least_3
 @pytest.mark.parametrize(
     "features",
     [
@@ -2707,6 +2711,7 @@ def test_dataset_from_parquet_features(features, parquet_path, tmp_path):
     _check_parquet_dataset(dataset, expected_features)
 
 
+@require_pyarrow_at_least_3
 @pytest.mark.parametrize("split", [None, NamedSplit("train"), "train", "test"])
 def test_dataset_from_parquet_split(split, parquet_path, tmp_path):
     cache_dir = tmp_path / "cache"
@@ -2716,6 +2721,7 @@ def test_dataset_from_parquet_split(split, parquet_path, tmp_path):
     assert dataset.split == str(split) if split else "train"
 
 
+@require_pyarrow_at_least_3
 @pytest.mark.parametrize("path_type", [str, list])
 def test_dataset_from_parquet_path_type(path_type, parquet_path, tmp_path):
     if issubclass(path_type, str):

tests/test_dataset_common.py

@@ -22,7 +22,9 @@
 from typing import List, Optional
 from unittest import TestCase
 
+import pyarrow as pa
 from absl.testing import parameterized
+from packaging import version
 
 from datasets import cached_path, import_main_class, load_dataset, prepare_module
 from datasets.builder import BuilderConfig, DatasetBuilder
@@ -270,7 +272,10 @@ def test_load_real_dataset_all_configs(self, dataset_name):
 
 
 def get_packaged_dataset_names():
-    return [{"testcase_name": x, "dataset_name": x} for x in _PACKAGED_DATASETS_MODULES.keys()]
+    packaged_datasets = [{"testcase_name": x, "dataset_name": x} for x in _PACKAGED_DATASETS_MODULES.keys()]
+    if version.parse(pa.__version) < version.parse("3.0.0"):  # parquet is not supported for pyarrow<3.0.0
+        packaged_datasets = [pd for pd in packaged_datasets if pd["dataset_name"] != "parquet"]
+    return packaged_datasets
 
 
 @parameterized.named_parameters(get_packaged_dataset_names())

tests/test_dataset_dict.py

@@ -16,6 +16,7 @@
 from .utils import (
     assert_arrow_memory_doesnt_increase,
     assert_arrow_memory_increases,
+    require_pyarrow_at_least_3,
     require_s3,
     require_tf,
     require_torch,
@@ -597,6 +598,7 @@ def _check_parquet_datasetdict(dataset_dict, expected_features, splits=("train",
             assert dataset.features[feature].dtype == expected_dtype
 
 
+@require_pyarrow_at_least_3
 @pytest.mark.parametrize("keep_in_memory", [False, True])
 def test_datasetdict_from_parquet_keep_in_memory(keep_in_memory, parquet_path, tmp_path):
     cache_dir = tmp_path / "cache"
@@ -606,6 +608,7 @@ def test_datasetdict_from_parquet_keep_in_memory(keep_in_memory, parquet_path, t
     _check_parquet_datasetdict(dataset, expected_features)
 
 
+@require_pyarrow_at_least_3
 @pytest.mark.parametrize(
     "features",
     [
@@ -627,6 +630,7 @@ def test_datasetdict_from_parquet_features(features, parquet_path, tmp_path):
     _check_parquet_datasetdict(dataset, expected_features)
 
 
+@require_pyarrow_at_least_3
 @pytest.mark.parametrize("split", [None, NamedSplit("train"), "train", "test"])
 def test_datasetdict_from_parquet_split(split, parquet_path, tmp_path):
     if split:

tests/utils.py

@@ -8,6 +8,7 @@
 from unittest.mock import patch
 
 import pyarrow as pa
+from packaging import version
 
 from datasets import config
 
@@ -34,6 +35,19 @@ def parse_flag_from_env(key, default=False):
 _run_packaged_tests = parse_flag_from_env("RUN_PACKAGED", default=True)
 
 
+def require_pyarrow_at_least_3(test_case):
+    """
+    Decorator marking a test that requires PyArrow 3.0.0
+    to allow nested types in parquet, as well as batch iterators of parquet files.
+
+    These tests are skipped when the PyArrow version is outdated.
+
+    """
+    if version.parse(config.PYARROW_VERSION) < version.parse("3.0.0"):
+        test_case = unittest.skip("test requires PyTorch")(test_case)
+    return test_case
+
+
 def require_beam(test_case):
     """
     Decorator marking a test that requires Apache Beam.