Support versioning of the underlying dataset with PartitionedDataset

RELEASE.md

@@ -18,6 +18,7 @@
 ## Bug fixes and other changes
 * Fixed `TypeError` when converting dict inputs to a node made from a wrapped `partial` function.
 * `PartitionedDataSet` improvements:
+  - Supported versioning of the underlying dataset.
   - Supported passing arguments to the underlying filesystem.
 * Improved handling of non-ASCII word characters in dataset names.
   - For example, a dataset named `jalapeño` will be accessible as `DataCatalog.datasets.jalapeño` rather than `DataCatalog.datasets.jalape__o`.

kedro/io/partitioned_data_set.py

@@ -31,6 +31,7 @@
 """
 import operator
 from copy import deepcopy
+from pathlib import PurePosixPath
 from typing import Any, Callable, Dict, List, Tuple, Type, Union
 from urllib.parse import urlparse
 from warnings import warn
@@ -58,6 +59,18 @@
 S3_PROTOCOLS = ("s3", "s3a", "s3n")
 
 
+def _grandparent(path: str) -> str:
+    path_obj = PurePosixPath(path)
+    grandparent = path_obj.parents[1]
+    if grandparent.name != path_obj.name:
+        last_three_parts = path_obj.relative_to(*path_obj.parts[:-3])
+        raise DataSetError(
+            f"`{path}` is not a well-formed versioned path ending with "
+            f"`filename/timestamp/filename` (got `{last_three_parts}`)."
+        )
+    return str(grandparent)
+
+
 class PartitionedDataSet(AbstractDataSet):
     # pylint: disable=too-many-instance-attributes,protected-access
     """``PartitionedDataSet`` loads and saves partitioned file-like data using the
@@ -144,9 +157,6 @@ def __init__(  # pylint: disable=too-many-arguments
                 the filesystem implementation.
             fs_args: Extra arguments to pass into underlying filesystem class constructor
                 (e.g. `{"project": "my-project"}` for ``GCSFileSystem``)
-
-        Raises:
-            DataSetError: If versioning is enabled for the underlying dataset.
         """
         # pylint: disable=import-outside-toplevel
         from fsspec.utils import infer_storage_options  # for performance reasons
@@ -160,13 +170,6 @@ def __init__(  # pylint: disable=too-many-arguments
 
         dataset = dataset if isinstance(dataset, dict) else {"type": dataset}
         self._dataset_type, self._dataset_config = parse_dataset_definition(dataset)
-        if VERSION_KEY in self._dataset_config:
-            raise DataSetError(
-                "`{}` does not support versioning of the underlying dataset. "
-                "Please remove `{}` flag from the dataset definition.".format(
-                    self.__class__.__name__, VERSIONED_FLAG_KEY
-                )
-            )
 
         self._credentials, dataset_credentials = _split_credentials(credentials)
         if dataset_credentials:
@@ -217,7 +220,7 @@ def _normalized_path(self) -> str:
     @cachedmethod(cache=operator.attrgetter("_partition_cache"))
     def _list_partitions(self) -> List[str]:
         return [
-            path
+            _grandparent(path) if self._dataset_config.get(VERSION_KEY) else path
             for path in self._filesystem.find(self._normalized_path, **self._load_args)
             if path.endswith(self._filename_suffix)
         ]

tests/io/test_partitioned_dataset.py

@@ -52,7 +52,7 @@ def partitioned_data_pandas():
 
 @pytest.fixture
 def local_csvs(tmp_path, partitioned_data_pandas):
-    local_dir = Path(str(tmp_path / "csvs"))
+    local_dir = tmp_path / "csvs"
     local_dir.mkdir()
 
     for k, data in partitioned_data_pandas.items():
@@ -62,6 +62,11 @@ def local_csvs(tmp_path, partitioned_data_pandas):
     return local_dir
 
 
+@pytest.fixture
+def filepath_csvs(tmp_path):
+    return str(tmp_path / "csvs")
+
+
 LOCAL_DATASET_DEFINITION = [
     "pandas.CSVDataSet",
     "kedro.extras.datasets.pandas.CSVDataSet",
@@ -286,17 +291,56 @@ def test_invalid_dataset_config(self, dataset_config, error_pattern):
     @pytest.mark.parametrize(
         "dataset_config",
         [
-            {"type": CSVDataSet, "versioned": True},
-            {"type": "pandas.CSVDataSet", "versioned": True},
+            {**ds_config, "versioned": True}
+            for ds_config in LOCAL_DATASET_DEFINITION
+            if isinstance(ds_config, dict)
         ],
     )
-    def test_versioned_dataset_not_allowed(self, dataset_config):
-        pattern = (
-            "`PartitionedDataSet` does not support versioning of the underlying "
-            "dataset. Please remove `versioned` flag from the dataset definition."
+    @pytest.mark.parametrize(
+        "suffix,expected_num_parts", [("", 5), (".csv", 3), ("p4", 1)]
+    )
+    def test_versioned_dataset_save_and_load(
+        self,
+        filepath_csvs,
+        dataset_config,
+        suffix,
+        expected_num_parts,
+        partitioned_data_pandas,
+    ):
+        """Test that saved and reloaded data matches the original one for
+        the versioned data set."""
+        PartitionedDataSet(filepath_csvs, dataset_config).save(partitioned_data_pandas)
+
+        pds = PartitionedDataSet(filepath_csvs, dataset_config, filename_suffix=suffix)
+        loaded_partitions = pds.load()
+
+        assert len(loaded_partitions.keys()) == expected_num_parts
+        for partition_id, load_func in loaded_partitions.items():
+            df = load_func()
+            assert_frame_equal(df, partitioned_data_pandas[partition_id + suffix])
+            if suffix:
+                assert not partition_id.endswith(suffix)
+
+    def test_malformed_versioned_path(self, tmp_path):
+        local_dir = tmp_path / "files"
+        local_dir.mkdir()
+
+        path = local_dir / "path/to/folder/new/partition/version/partition/file"
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.write_text("content")
+
+        pds = PartitionedDataSet(
+            str(local_dir / "path/to/folder"),
+            {"type": "pandas.CSVDataSet", "versioned": True},
         )
-        with pytest.raises(DataSetError, match=re.escape(pattern)):
-            PartitionedDataSet(str(Path.cwd()), dataset_config)
+
+        pattern = re.escape(
+            f"`{path.as_posix()}` is not a well-formed versioned path ending with "
+            "`filename/timestamp/filename` (got `version/partition/"
+            "file`)."
+        )
+        with pytest.raises(DataSetError, match=pattern):
+            pds.load()
 
     def test_no_partitions(self, tmpdir):
         pds = PartitionedDataSet(str(tmpdir), "pandas.CSVDataSet")
@@ -328,7 +372,7 @@ def test_no_partitions(self, tmpdir):
     def test_filepath_arg_warning(self, pds_config, filepath_arg):
         pattern = (
             f"`{filepath_arg}` key must not be specified in the dataset definition as it "
-            f"will be overwritten by partition path"
+            "will be overwritten by partition path"
         )
         with pytest.warns(UserWarning, match=re.escape(pattern)):
             PartitionedDataSet(**pds_config)