Repository: kedro-org/kedro

Issue:
Render concrete parameter/return types for dataset `save`/`load`
## Description
The parameter type for a dataset's `save` is `_DI`, and the return type for `load` is `_DO`. This doesn't make any sense to a user (and barely makes sense to me, as a Kedro developer, because I saw the typing PR go in and have that context). I'd rather it not display any types, if getting the correct concrete types is too difficult.

<img width="1436" alt="image" src="https://user-images.githubusercontent.com/14007150/211978326-bd1cced3-03d3-4f4b-bc88-4c48269bf159.png">

Discussion:
datajoely: Love this 
Galileo-Galilei: This is a brilliant idea, I am jealous of not having it myself ;) I always have to explain to users who create a custom dataset that they should implement ``_load`` and ``_save`` instead of ``load`` and ``save``, and some less advanced users struggle to get the point and this would likely help. 

One small caveat is that I don't know if this would have unintended consequences on ``MlflowArtifactDataset`` which also wraps the methods (I *think* it shouldn't, but I'd rather check)
deepyaman: > One small caveat is that I don't know if this would have unintended consequences on `MlflowArtifactDataset` which also wraps the methods (I _think_ it shouldn't, but I'd rather check)

Will be great if you can check that! I did guard against wrapping datasets that inherit `load` or `save` multiple times, but `MlflowArtifactDataset` implementation looks quite different, so I'll defer to you. :)
astrojuanlu: I can't comment much on the implementation for now but [I've been a long advocate of making `load` and `save` public](https://github.com/kedro-org/kedro/issues/1936#issuecomment-1717214883), so big +1 from me!

We should be mindful how this affects the way people define custom datasets. Will they still need to implement `_load` and `_save`? If so, do we need to adjust our docs? And finally, is it too early to think of a world in which `_load` and `_save` are actually no longer necessary?
deepyaman: > We should be mindful how this affects the way people define custom datasets. Will they still need to implement `_load` and `_save`? If so, do we need to adjust our docs? And finally, is it too early to think of a world in which `_load` and `_save` are actually no longer necessary?

I'd included some information in the [development notes]I(https://github.com/kedro-org/kedro/pull/3920#issue-2332651112) above, but let me copy them below:

I have updated the documentation on extending a dataset, but I have intentionally not explained that "core" datasets (in Kedro and Kedro-Datasets) should use the following pattern to work across older Kedro versions:

```python
class MyDataset(...):
    def _load(...) -> ...:
        ...

    load = _load

    def _save(...) -> ...:
        ...

    save = _save
```

If we want to take things slow, in the next minor release of Kedro (`0.20.0`), we can start raising `DeprecationWarning`s for datasets that don't implement `load` or `save`. In `0.21.0`, we can drop support for datasets that don't define `load` or `save` (and just have the older `_load` or `_save`).

The reason I don't want to introduce this in the docs is that it will confuse new users who want to write local datasets, and it's a very small change to make before maintainers add datasets to the core.
Galileo-Galilei: Hi, sorry there is something I don't get. I am trying to do this (inside ``spaceflights-pandas`` starter, this code is copy pasted from tour example in the docstring):

```python
from pathlib import Path

import pandas as pd
from kedro.io import AbstractDataset


class MyOwnDataset(AbstractDataset):
    def __init__(self, filepath):
        self._filepath = Path(filepath)

    def load(self) -> pd.DataFrame:
        return pd.read_csv(self._filepath)

    def save(self, df: pd.DataFrame) -> None:
        df.to_csv(str(self._filepath))

    def _exists(self) -> bool:
        return Path(self._filepath.as_posix()).exists()

    def _describe(self):
        return dict(param1=self._filepath)


ds = MyOwnDataset(
    filepath=(Path(__file__).parents[1] / "data/01_raw/companies.csv").as_posix(),
)

ds.load()
``` 
 And I am getting the error : 

```bash 
TypeError: Can't instantiate abstract class MyOwnDataset with abstract methods _load, _save
```

I've double checked and I have the correct developement version installed. while debugging I am entering the ``__init_subclass`` method though :confused: so I don't really get if it's how it is intended to be used. I don't see anywhere in the code where you monkeypatch ``_load`` and ````_save in case they are not defined in the custom dataset, but maybe I am missing something. 

 
deepyaman: > ```shell
> TypeError: Can't instantiate abstract class MyOwnDataset with abstract methods _load, _save
> ```
> 
> I've double checked and I have the correct developement version installed. while debugging I am entering the `__init_subclass` method though 😕 so I don't really get if it's how it is intended to be used. I don't see anywhere in the code where you monkeypatch `_load` and ````_save in case they are not defined in the custom dataset, but maybe I am missing something.

Ah, this is probably right... let me try to take a look and patch that. Can probably just throw in some dummy method, since it shouldn't actually be hit.
deepyaman: @Galileo-Galilei The specific (massive) bug you pointed out should be resolved, and, in trying to figure out how to best resolve it, I think I've also made the code MUCH cleaner. :)

There's no duplicate load/save handling logic, and the `__init_subclass__`:

1. Aliases `_load` to `load`, `_save` to `save`, if it is defined for the class.
2. Does as before (apply the load wrapper to the `load`, save wrapper to the `save`).

There's no more, what do we do if `_load` is defined, and how that is different from what we do when `save` is defined. Also, `load` and `save` are marked abstract, which makes sense.
merelcht: I don't really like how this now adds inconsistency between the abstract methods that need to be implemented. `save` and `load` vs. `_describe`. 
deepyaman: > I don't really like how this now adds inconsistency between the abstract methods that need to be implemented. `save` and `load` vs. `_describe`.

I'd be willing to make `_describe` public in a follow-up PR? No objection from my side, if that's desirable; I see no real reason `_describe` should be a private API.

However, I think that `load` and `save` are publicly exposed already, whereas `_describe` is not, which is what drives this.
Galileo-Galilei: Still need to deep dive in the implementation, but unfortunately this breaks kedro-mlflow ``MlflowArtifactDataset``  :confused: because of this line: 

https://github.com/Galileo-Galilei/kedro-mlflow/blob/64b8e94e1dafa02d979e7753dab9b9dfd4d7341c/kedro_mlflow/io/artifacts/mlflow_artifact_dataset.py#L65C17-L65C36

I still need to think about the best workaround and check if i can just call ``__super__.save`` but this generates double logging so I am not sure this would be retrocompatible.  
Galileo-Galilei: I am very annoyed about this. 

I understand what is going on and I can't find a way to make the change backward compatible with kedro-mlflow, unless if make a huge code duplication to define MlflowArtifactDataset with a condition on kedro version once this is realeased, and I really don't like it. 

### Context 

If one has read the code of ``MLflowArtifactDataset``, he may have been very surprised because it is very different of standard dataset. I'll try to explain the rationale below. 

To create a dataset that logs automatically in mlflow on ``save``, several design were envisioned: 

- *naive*: create a dataset for each underlying format and ask user to replace their original code with my imlmentation : ``MLflowArtifact<CSV/Parquet/Whatever>Dataset`` This is definitely a no go because maitenance burden would be increasingly high, I would never ensure full consistency with ``kedro-datasets`` and it would not support custom datasets. 
- *hook based*: something similar to this
    ```yaml
    my_data:
        type: pandas.CSVDataset
        filepath: /path/to/data.csv
        metadata:
            kedro-mlflow:
                log_as_artifact: true
    ```
    I never really considered this for an obvious reason : the metadata key did not exist back then :). However, I don't like this 
    design because it ties a lot the dataset to the ``catalog`` and the ``session``: the logging will be triggered in a 
    ``after_node_run`` hook. As a consequence: 
   - it does not work interactively (e.g. in a notebook), or if you use the catalog as a standalone component
   - it makes testing much harder
   - it is less backward compatible and satble over time because it is is impacted by session object modification

- *method decorators*: in the old kedro versions (<=0.15), I had the possibility to decorate dataset method with somethin like this (I don't remember the exact syntax and I can't find doc before 0.18, likely something we should fix): 
    ```yaml
    my_data:
        type: pandas.CSVDataset
        filepath: /path/to/data.csv
        decorator: 
            - save: kedro_mlflow.log_in_mlflow" a custom function
    ```
    This was not very flexible and replaced with hooks, but this is very useful in this specific situation. 
    I decided to emulate this behaviour with a class: 
    ```yaml
    my_data:
        type: kedro_mlflow.io.artifacts.MlflowArtifactDataset
        dataset:
            type: pandas.CSVDataset
            filepath: /path/to/data.csv
    ```
    The syntax is very simple and compatible with all kedro versions (until this PR ;-) ) because it only assumes the existence of ``_load`` and ``_save``

### How does it work

The key idea of ``MlflowArtifactDataset`` is to use the ``__new__`` method to fake inheritance : we create on the fly a new class that inherits from the target dataset (``CSVDataset`` in above example) and decorates its ``_save`` method. 

Finally, I create an instance of this new class which will be the one accessible  for the end user. 

If I modify the ``_load`` and ``_save`` methods of this fake class to ``load`` and ``save``, it works smoothly with above PR; however this is not backward compatible because old datasets implementation without the changes in this PR will break (I'll modify their original load and save, which is *not* what I want to do). 

deepyaman: > If I modify the `_load` and `_save` methods of this fake class to `load` and `save`, it works smoothly with above PR; however this is not backward compatible because old datasets implementation without the changes in this PR will break (I'll modify their original load and save, which is _not_ what I want to do).

If I understand your comments and the line of code linked above correctly, the issue is that:

1. It's looking for the `_load` and `_save` method on the dataset.
2. It's implementing it's own wrapper on `_load` and `_save`.

Ideally, your wrapper class should implement `_load_wrapper` and `_save_wrapper` (e.g. see how `AbstractVersionedDataset` has redefined it).

Would it be sufficient if, for now, we didn't remove the `_load` or `_save` method on existing datasets? In this PR, we alias existing `_load` and `_save` to `load` and `save` automatically, so any datasets that still define the "old" versions will work fine:

```python
        if hasattr(cls, "_load") and not cls._load.__qualname__.startswith("Abstract"):
            cls.load = cls._load  # type: ignore[method-assign]

        if hasattr(cls, "_save") and not cls._save.__qualname__.startswith("Abstract"):
            cls.save = cls._save  # type: ignore[method-assign]
```

As a result, we can hold off on making these renames to a future Kedro version (0.20.0?), as well as maintain the existing guidance of defining `_load` and `_save` (as opposed to `load` and `save`) in the docs.
deepyaman: > I understand what is going on and I can't find a way to make the change backward compatible with kedro-mlflow, unless if make a huge code duplication to define MlflowArtifactDataset with a condition on kedro version once this is realeased, and I really don't like it.

@Galileo-Galilei With this change, you can detect whether the logic is wrapped (like in the current `load`) or not wrapped (as in the current `_load`) from the `__loadwrapped__` or `__savewrapped__` attribute. For example: 

```python
                cls.load
                if not getattr(cls.load, "__loadwrapped__", False)
                else cls.load.__wrapped__
```

can be used to get the underlying load logic. This can be used to help make things backwards-compatible, I think.
deepyaman: @Galileo-Galilei I went ahead and reverted the `_load` -> `load` and `_save` -> `save` renames for core datasets (the backwards-compatible "legacy" path is still exposing them publicly). I think this should prevent breaking Kedro-MLFlows use of the internal attributes for now. If that works, I think we should be able to get this in.

In the meantime, we can find a solution for Kedro-MLFlow to look at `load.__wrapped__` (if defined), and start removing the legacy `_load`/`_save` methods.
Galileo-Galilei: Hi @deepyaman, I finally got a chance to take a stab at it, sorry for the delay and thanks for your patience 🙏 

## :white_check_mark: It's good to go ...
I've tested interactively and with the CLI / for old kedro versions and this branch / with and without underscored methods, and it seems to work in all cases 🥳 I need to make a release to support the new version though. Basically I just need to replace ``super._load()`` by ``super().load() if getattr(cls.load, "__loadwrapped__", False) else super()._load()`` in the wrapper class.

## :truck: ... but we still need to agree about migration strategy 

I'd be inclined to do the following: 
- Add a ``DeprecationWarning`` to ``_load`` and ``_save`` to suggest the new behaviour to users
- update the documentation to remove ``_load and ``_save`` and example, but only keep a warning admonition somewhere (to  avoid blowing users mind if they find a working legacy code with underscored method and no mention anywhere in the doc nor in the code :exploding_head:) ? 
- definitely remove it in 0.20, and I'll break kedro-mlflow in the meantime so user have to pay the migration cost once for all

## Bonus: code for interactive tests

```python
from pathlib import Path

import pandas as pd
from kedro.io import AbstractDataset
from kedro_mlflow.io.artifacts import MlflowArtifactDataset

# test 1 : new way without save and load


class MyOwnDataset_without_underscore_methods(AbstractDataset):
    def __init__(self, filepath):
        self._filepath = Path(filepath)

    def load(self) -> pd.DataFrame:
        return pd.read_csv(self._filepath)

    def save(self, df: pd.DataFrame) -> None:
        df.to_csv(str(self._filepath))

    # _load = load
    # _save = save

    def _exists(self) -> bool:
        return Path(self._filepath.as_posix()).exists()

    def _describe(self):
        return dict(param1=self._filepath)


ds1 = MyOwnDataset_without_underscore_methods(
    filepath=(Path(__file__).parents[1] / "data/01_raw/companies.csv").as_posix(),
)

companies = ds1.load()


mlflow_ds1 = MlflowArtifactDataset(
    dataset=dict(
        type=MyOwnDataset_without_underscore_methods,
        filepath=(Path(__file__).parents[1] / "data/01_raw/truc.csv").as_posix(),
    ),
    artifact_path="truc",
)


mlflow_ds1.save(companies)


# test 2 : old way


class MyOwnDataset_with_underscore_methods(AbstractDataset):
    def __init__(self, filepath):
        self._filepath = Path(filepath)

    def _load(self) -> pd.DataFrame:
        return pd.read_csv(self._filepath)

    def _save(self, df: pd.DataFrame) -> None:
        df.to_csv(str(self._filepath))

    # _load = load
    # _save = save

    def _exists(self) -> bool:
        return Path(self._filepath.as_posix()).exists()

    def _describe(self):
        return dict(param1=self._filepath)


ds = MyOwnDataset_with_underscore_methods(
    filepath=(Path(__file__).parents[1] / "data/01_raw/companies.csv").as_posix(),
)

companies = ds.load()


mlflow_ds2 = MlflowArtifactDataset(
    dataset=dict(
        type=MyOwnDataset_with_underscore_methods,
        filepath=(Path(__file__).parents[1] / "data/01_raw/truc.csv").as_posix(),
    ),
    artifact_path="truc",
)


mlflow_ds2.save(companies)
```
deepyaman: @Galileo-Galilei Thanks for taking another look!

> I need to make a release to support the new version though. Basically I just need to replace `super._load()` by `super().load() if getattr(cls.load, "__loadwrapped__", False) else super()._load()` in the wrapper class.

Do you need this even if nobody defines a public `load`? I thought `_load` should always be available for you after my most-recent commits, until people start defining `load` and not `_load`.

> ## 🚚 ... but we still need to agree about migration strategy
> I'd be inclined to do the following:
> 
> * Add a `DeprecationWarning` to `_load` and `_save` to suggest the new behaviour to users
> * update the documentation to remove `_load and `_save`` and example, but only keep a warning admonition somewhere (to  avoid blowing users mind if they find a working legacy code with underscored method and no mention anywhere in the doc nor in the code 🤯) ?
> * definitely remove it in 0.20, and I'll break kedro-mlflow in the meantime so user have to pay the migration cost once for all

Since I've reverted the guidance to use `_load` and `_save` in the docs, I think we can figure out migration strategy before you break Kedro-MLflow; for now, anybody looking at existing datasets as an example, or reading the docs, should not get the idea to define public `load` and `save` themselves.

I think we need to wait to issue a deprecation warning until Kedro-Datasets is updated; don't want people getting deprecation warnings they can't do anything about.

I'm good with the goal of removing it in 0.20, though!
Galileo-Galilei: > Do you need this even if nobody defines a public load? I thought _load should always be available for you after my most-recent commits, until people start defining load and not _load

My first code example shows that people now can define a public ``load`` instead of ``_load`` and ``hasttr(ds1, "_load")`` returns ``False``. I need this line of code incase people start defining only the public methods. 

> for now, anybody looking at existing datasets as an example, or reading the docs, should not get the idea to define public load and save themselves.

Once again, I don't understand : suggesting people to define ``load`` and ``save`` is  the very goal of this PR, isn't it? We want people to start defining themselves the public method and not the private ones, all above discussion is about backward compatibilty, not about the fact that it is what we want users to lend towards. 
deepyaman: > Once again, I don't understand : suggesting people to define `load` and `save` is the very goal of this PR, isn't it? We want people to start defining themselves the public method and not the private ones, all above discussion is about backward compatibilty, not about the fact that it is what we want users to lend towards.

The reason I started looking into this was actually to resolve #2199, which is achieved even without explicitly defining `load` and `save`. This is achieved already by the "legacy" path handling in this PR.

Allowing people to define public `load` and `save` was a nice result of the solution I landed on. :) I 100% think we should move towards that; I just want to be cognizant of causing the least breakages possible. If this means giving some more time before breaking Kedro-MLflow, I was open to it. (I'm also happy if you want to rip off the bandaid and do the breaking change sooner, up to you!)
Galileo-Galilei: I don't consider the current solution to be a breaking change from the ``kedro-mlflow`` point of view : when kedro adds a feature, I have to release a minor version to support it. Since we already enable people to define only load and save and this is the desirable behaviour, I think we should already emphasize in the documentation that it is "the kedro way" to define a custom dataset, just to ensure "a path to the least breakage" in 0.20 when we will force it. 

Completely happy to merge it as is 🚀, and postpone removing ``_load`` and ``_save`` from the abstract class on 0.20 so I'll break kedro-mlflow simultaneously. We should open an issue to keep track of it!

Patch:
diff --git a/RELEASE.md b/RELEASE.md
index f3eeec10..32a47371 100644
--- a/RELEASE.md
+++ b/RELEASE.md
@@ -1,6 +1,7 @@
 # Upcoming Release 0.19.7
 
 ## Major features and improvements
+* Exposed `load` and `save` publicly for each dataset in the core `kedro` library, and enabled other datasets to do the same. If a dataset doesn't expose `load` or `save` publicly, Kedro will fall back to using `_load` or `_save`, respectively.
 * Kedro commands are now lazily loaded to add performance gains when running Kedro commands.
 
 ## Bug fixes and other changes
diff --git a/docs/source/data/how_to_create_a_custom_dataset.md b/docs/source/data/how_to_create_a_custom_dataset.md
index 5ebe50a8..01ad199f 100644
--- a/docs/source/data/how_to_create_a_custom_dataset.md
+++ b/docs/source/data/how_to_create_a_custom_dataset.md
@@ -312,7 +312,7 @@ To add versioning support to the new dataset we need to extend the
  {py:class}`~kedro.io.AbstractVersionedDataset` to:
 
 * Accept a `version` keyword argument as part of the constructor
-* Adapt the `_save` and `_load` method to use the versioned data path obtained from `_get_save_path` and `_get_load_path` respectively
+* Adapt the `_load` and `_save` method to use the versioned data path obtained from `_get_load_path` and `_get_save_path` respectively
 
 The following amends the full implementation of our basic `ImageDataset`. It now loads and saves data to and from a versioned subfolder (`data/01_raw/pokemon-images-and-types/images/images/pikachu.png/<version>/pikachu.png` with `version` being a datetime-formatted string `YYYY-MM-DDThh.mm.ss.sssZ` by default):
 
diff --git a/kedro/io/cached_dataset.py b/kedro/io/cached_dataset.py
index 2ca0bac2..c4424c34 100644
--- a/kedro/io/cached_dataset.py
+++ b/kedro/io/cached_dataset.py
@@ -71,7 +71,7 @@ class CachedDataset(AbstractDataset):
                 "The argument type of 'dataset' should be either a dict/YAML "
                 "representation of the dataset, or the actual dataset object."
             )
-        self._cache = MemoryDataset(copy_mode=copy_mode)
+        self._cache = MemoryDataset(copy_mode=copy_mode)  # type: ignore[abstract]
         self.metadata = metadata
 
     def _release(self) -> None:
diff --git a/kedro/io/core.py b/kedro/io/core.py
index 53c2000e..fc48e536 100644
--- a/kedro/io/core.py
+++ b/kedro/io/core.py
@@ -13,7 +13,7 @@ import sys
 import warnings
 from collections import namedtuple
 from datetime import datetime, timezone
-from functools import partial
+from functools import partial, wraps
 from glob import iglob
 from operator import attrgetter
 from pathlib import Path, PurePath, PurePosixPath
@@ -22,6 +22,7 @@ from urllib.parse import urlsplit
 
 from cachetools import Cache, cachedmethod
 from cachetools.keys import hashkey
+from typing_extensions import Self
 
 from kedro.utils import load_obj
 
@@ -74,6 +75,7 @@ _DO = TypeVar("_DO")
 
 class AbstractDataset(abc.ABC, Generic[_DI, _DO]):
     """``AbstractDataset`` is the base class for all data set implementations.
+
     All data set implementations should extend this abstract class
     and implement the methods marked as abstract.
     If a specific dataset implementation cannot be used in conjunction with
@@ -93,10 +95,10 @@ class AbstractDataset(abc.ABC, Generic[_DI, _DO]):
         >>>         self._param1 = param1
         >>>         self._param2 = param2
         >>>
-        >>>     def _load(self) -> pd.DataFrame:
+        >>>     def load(self) -> pd.DataFrame:
         >>>         return pd.read_csv(self._filepath)
         >>>
-        >>>     def _save(self, df: pd.DataFrame) -> None:
+        >>>     def save(self, df: pd.DataFrame) -> None:
         >>>         df.to_csv(str(self._filepath))
         >>>
         >>>     def _exists(self) -> bool:
@@ -178,57 +180,6 @@ class AbstractDataset(abc.ABC, Generic[_DI, _DO]):
     def _logger(self) -> logging.Logger:
         return logging.getLogger(__name__)
 
-    def load(self) -> _DO:
-        """Loads data by delegation to the provided load method.
-
-        Returns:
-            Data returned by the provided load method.
-
-        Raises:
-            DatasetError: When underlying load method raises error.
-
-        """
-
-        self._logger.debug("Loading %s", str(self))
-
-        try:
-            return self._load()
-        except DatasetError:
-            raise
-        except Exception as exc:
-            # This exception handling is by design as the composed data sets
-            # can throw any type of exception.
-            message = (
-                f"Failed while loading data from data set {str(self)}.\n{str(exc)}"
-            )
-            raise DatasetError(message) from exc
-
-    def save(self, data: _DI) -> None:
-        """Saves data by delegation to the provided save method.
-
-        Args:
-            data: the value to be saved by provided save method.
-
-        Raises:
-            DatasetError: when underlying save method raises error.
-            FileNotFoundError: when save method got file instead of dir, on Windows.
-            NotADirectoryError: when save method got file instead of dir, on Unix.
-        """
-
-        if data is None:
-            raise DatasetError("Saving 'None' to a 'Dataset' is not allowed")
-
-        try:
-            self._logger.debug("Saving %s", str(self))
-            self._save(data)
-        except DatasetError:
-            raise
-        except (FileNotFoundError, NotADirectoryError):
-            raise
-        except Exception as exc:
-            message = f"Failed while saving data to data set {str(self)}.\n{str(exc)}"
-            raise DatasetError(message) from exc
-
     def __str__(self) -> str:
         # TODO: Replace with __repr__ implementation in 0.20.0 release.
         def _to_str(obj: Any, is_root: bool = False) -> str:
@@ -257,6 +208,85 @@ class AbstractDataset(abc.ABC, Generic[_DI, _DO]):
 
         return f"{type(self).__name__}({_to_str(self._describe(), True)})"
 
+    @classmethod
+    def _load_wrapper(cls, load_func: Callable[[Self], _DO]) -> Callable[[Self], _DO]:
+        """Decorate `load_func` with logging and error handling code."""
+
+        @wraps(load_func)
+        def load(self: Self) -> _DO:
+            self._logger.debug("Loading %s", str(self))
+
+            try:
+                return load_func(self)
+            except DatasetError:
+                raise
+            except Exception as exc:
+                # This exception handling is by design as the composed data sets
+                # can throw any type of exception.
+                message = (
+                    f"Failed while loading data from data set {str(self)}.\n{str(exc)}"
+                )
+                raise DatasetError(message) from exc
+
+        load.__annotations__["return"] = load_func.__annotations__.get("return")
+        load.__loadwrapped__ = True  # type: ignore[attr-defined]
+        return load
+
+    @classmethod
+    def _save_wrapper(
+        cls, save_func: Callable[[Self, _DI], None]
+    ) -> Callable[[Self, _DI], None]:
+        """Decorate `save_func` with logging and error handling code."""
+
+        @wraps(save_func)
+        def save(self: Self, data: _DI) -> None:
+            if data is None:
+                raise DatasetError("Saving 'None' to a 'Dataset' is not allowed")
+
+            try:
+                self._logger.debug("Saving %s", str(self))
+                save_func(self, data)
+            except (DatasetError, FileNotFoundError, NotADirectoryError):
+                raise
+            except Exception as exc:
+                message = (
+                    f"Failed while saving data to data set {str(self)}.\n{str(exc)}"
+                )
+                raise DatasetError(message) from exc
+
+        save.__annotations__["data"] = save_func.__annotations__.get("data", Any)
+        save.__annotations__["return"] = save_func.__annotations__.get("return")
+        save.__savewrapped__ = True  # type: ignore[attr-defined]
+        return save
+
+    def __init_subclass__(cls, **kwargs: Any) -> None:
+        """Decorate the `load` and `save` methods provided by the class.
+
+        If `_load` or `_save` are defined, alias them as a prerequisite.
+
+        """
+        super().__init_subclass__(**kwargs)
+
+        if hasattr(cls, "_load") and not cls._load.__qualname__.startswith("Abstract"):
+            cls.load = cls._load  # type: ignore[method-assign]
+
+        if hasattr(cls, "_save") and not cls._save.__qualname__.startswith("Abstract"):
+            cls.save = cls._save  # type: ignore[method-assign]
+
+        if hasattr(cls, "load") and not cls.load.__qualname__.startswith("Abstract"):
+            cls.load = cls._load_wrapper(  # type: ignore[assignment]
+                cls.load
+                if not getattr(cls.load, "__loadwrapped__", False)
+                else cls.load.__wrapped__  # type: ignore[attr-defined]
+            )
+
+        if hasattr(cls, "save") and not cls.save.__qualname__.startswith("Abstract"):
+            cls.save = cls._save_wrapper(  # type: ignore[assignment]
+                cls.save
+                if not getattr(cls.save, "__savewrapped__", False)
+                else cls.save.__wrapped__  # type: ignore[attr-defined]
+            )
+
     def _pretty_repr(self, object_description: dict[str, Any]) -> str:
         str_keys = []
         for arg_name, arg_descr in object_description.items():
@@ -276,17 +306,37 @@ class AbstractDataset(abc.ABC, Generic[_DI, _DO]):
         return self._pretty_repr(self._describe())
 
     @abc.abstractmethod
-    def _load(self) -> _DO:
+    def load(self) -> _DO:
+        """Loads data by delegation to the provided load method.
+
+        Returns:
+            Data returned by the provided load method.
+
+        Raises:
+            DatasetError: When underlying load method raises error.
+
+        """
         raise NotImplementedError(
             f"'{self.__class__.__name__}' is a subclass of AbstractDataset and "
-            f"it must implement the '_load' method"
+            f"it must implement the 'load' method"
         )
 
     @abc.abstractmethod
-    def _save(self, data: _DI) -> None:
+    def save(self, data: _DI) -> None:
+        """Saves data by delegation to the provided save method.
+
+        Args:
+            data: the value to be saved by provided save method.
+
+        Raises:
+            DatasetError: when underlying save method raises error.
+            FileNotFoundError: when save method got file instead of dir, on Windows.
+            NotADirectoryError: when save method got file instead of dir, on Unix.
+
+        """
         raise NotImplementedError(
             f"'{self.__class__.__name__}' is a subclass of AbstractDataset and "
-            f"it must implement the '_save' method"
+            f"it must implement the 'save' method"
         )
 
     @abc.abstractmethod
@@ -509,7 +559,9 @@ def _local_exists(local_filepath: str) -> bool:  # SKIP_IF_NO_SPARK
 class AbstractVersionedDataset(AbstractDataset[_DI, _DO], abc.ABC):
     """
     ``AbstractVersionedDataset`` is the base class for all versioned data set
-    implementations. All data sets that implement versioning should extend this
+    implementations.
+
+    All data sets that implement versioning should extend this
     abstract class and implement the methods marked as abstract.
 
     Example:
@@ -526,11 +578,11 @@ class AbstractVersionedDataset(AbstractDataset[_DI, _DO], abc.ABC):
         >>>         self._param1 = param1
         >>>         self._param2 = param2
         >>>
-        >>>     def _load(self) -> pd.DataFrame:
+        >>>     def load(self) -> pd.DataFrame:
         >>>         load_path = self._get_load_path()
         >>>         return pd.read_csv(load_path)
         >>>
-        >>>     def _save(self, df: pd.DataFrame) -> None:
+        >>>     def save(self, df: pd.DataFrame) -> None:
         >>>         save_path = self._get_save_path()
         >>>         df.to_csv(str(save_path))
         >>>
@@ -652,34 +704,43 @@ class AbstractVersionedDataset(AbstractDataset[_DI, _DO], abc.ABC):
     def _get_versioned_path(self, version: str) -> PurePosixPath:
         return self._filepath / version / self._filepath.name
 
-    def load(self) -> _DO:
-        return super().load()
-
-    def save(self, data: _DI) -> None:
-        self._version_cache.clear()
-        save_version = self.resolve_save_version()  # Make sure last save version is set
-        try:
-            super().save(data)
-        except (FileNotFoundError, NotADirectoryError) as err:
-            # FileNotFoundError raised in Win, NotADirectoryError raised in Unix
-            _default_version = "YYYY-MM-DDThh.mm.ss.sssZ"
-            raise DatasetError(
-                f"Cannot save versioned dataset '{self._filepath.name}' to "
-                f"'{self._filepath.parent.as_posix()}' because a file with the same "
-                f"name already exists in the directory. This is likely because "
-                f"versioning was enabled on a dataset already saved previously. Either "
-                f"remove '{self._filepath.name}' from the directory or manually "
-                f"convert it into a versioned dataset by placing it in a versioned "
-                f"directory (e.g. with default versioning format "
-                f"'{self._filepath.as_posix()}/{_default_version}/{self._filepath.name}"
-                f"')."
-            ) from err
+    @classmethod
+    def _save_wrapper(
+        cls, save_func: Callable[[Self, _DI], None]
+    ) -> Callable[[Self, _DI], None]:
+        """Decorate `save_func` with logging and error handling code."""
+
+        @wraps(save_func)
+        def save(self: Self, data: _DI) -> None:
+            self._version_cache.clear()
+            save_version = (
+                self.resolve_save_version()
+            )  # Make sure last save version is set
+            try:
+                super()._save_wrapper(save_func)(self, data)
+            except (FileNotFoundError, NotADirectoryError) as err:
+                # FileNotFoundError raised in Win, NotADirectoryError raised in Unix
+                _default_version = "YYYY-MM-DDThh.mm.ss.sssZ"
+                raise DatasetError(
+                    f"Cannot save versioned dataset '{self._filepath.name}' to "
+                    f"'{self._filepath.parent.as_posix()}' because a file with the same "
+                    f"name already exists in the directory. This is likely because "
+                    f"versioning was enabled on a dataset already saved previously. Either "
+                    f"remove '{self._filepath.name}' from the directory or manually "
+                    f"convert it into a versioned dataset by placing it in a versioned "
+                    f"directory (e.g. with default versioning format "
+                    f"'{self._filepath.as_posix()}/{_default_version}/{self._filepath.name}"
+                    f"')."
+                ) from err
+
+            load_version = self.resolve_load_version()
+            if load_version != save_version:
+                warnings.warn(
+                    _CONSISTENCY_WARNING.format(save_version, load_version, str(self))
+                )
+                self._version_cache.clear()
 
-        load_version = self.resolve_load_version()
-        if load_version != save_version:
-            warnings.warn(
-                _CONSISTENCY_WARNING.format(save_version, load_version, str(self))
-            )
+        return save
 
     def exists(self) -> bool:
         """Checks whether a data set's output already exists by calling
diff --git a/kedro/io/data_catalog.py b/kedro/io/data_catalog.py
index 5b95c124..6c70cd03 100644
--- a/kedro/io/data_catalog.py
+++ b/kedro/io/data_catalog.py
@@ -732,7 +732,7 @@ class DataCatalog:
             if isinstance(feed_dict[dataset_name], AbstractDataset):
                 dataset = feed_dict[dataset_name]
             else:
-                dataset = MemoryDataset(data=feed_dict[dataset_name])
+                dataset = MemoryDataset(data=feed_dict[dataset_name])  # type: ignore[abstract]
 
             self.add(dataset_name, dataset, replace)
 
diff --git a/kedro/io/lambda_dataset.py b/kedro/io/lambda_dataset.py
index 4ed3a7a2..126bef66 100644
--- a/kedro/io/lambda_dataset.py
+++ b/kedro/io/lambda_dataset.py
@@ -49,21 +49,21 @@ class LambdaDataset(AbstractDataset):
 
         return descr
 
-    def _save(self, data: Any) -> None:
-        if not self.__save:
+    def load(self) -> Any:
+        if not self.__load:
             raise DatasetError(
-                "Cannot save to data set. No 'save' function "
+                "Cannot load data set. No 'load' function "
                 "provided when LambdaDataset was created."
             )
-        self.__save(data)
+        return self.__load()
 
-    def _load(self) -> Any:
-        if not self.__load:
+    def save(self, data: Any) -> None:
+        if not self.__save:
             raise DatasetError(
-                "Cannot load data set. No 'load' function "
+                "Cannot save to data set. No 'save' function "
                 "provided when LambdaDataset was created."
             )
-        return self.__load()
+        self.__save(data)
 
     def _exists(self) -> bool:
         if not self.__exists:
diff --git a/kedro/io/memory_dataset.py b/kedro/io/memory_dataset.py
index b3b08e1f..30b6e19f 100644
--- a/kedro/io/memory_dataset.py
+++ b/kedro/io/memory_dataset.py
@@ -57,7 +57,7 @@ class MemoryDataset(AbstractDataset):
         self.metadata = metadata
         self._EPHEMERAL = True
         if data is not _EMPTY:
-            self._save(data)
+            self.save.__wrapped__(self, data)  # type: ignore[attr-defined]
 
     def _load(self) -> Any:
         if self._data is _EMPTY:
diff --git a/pyproject.toml b/pyproject.toml
index 383b7aa3..0aef9e1c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -32,6 +32,7 @@ dependencies = [
     "rich>=12.0,<14.0",
     "rope>=0.21,<2.0",  # subject to LGPLv3 license
     "toml>=0.10.0",
+    "typing_extensions>=4.0",
     "graphlib_backport>=1.0.0; python_version < '3.9'",
 ]
 keywords = [


Test patch:
diff --git a/tests/io/test_core.py b/tests/io/test_core.py
index c2f07a10..ebff20de 100644
--- a/tests/io/test_core.py
+++ b/tests/io/test_core.py
@@ -309,10 +309,18 @@ class TestCoreFunctions:
         dataset, _ = parse_dataset_definition(config)
         assert dataset is LambdaDataset
 
-    def test_test_parse_dataset_definition_with_python_class_type(self):
+    def test_parse_dataset_definition_with_python_class_type(self):
         config = {"type": MyDataset}
         parse_dataset_definition(config)
 
+    def test_load_and_save_are_wrapped_once(self):
+        assert not getattr(
+            MyOtherVersionedDataset.load.__wrapped__, "__loadwrapped__", False
+        )
+        assert not getattr(
+            MyOtherVersionedDataset.save.__wrapped__, "__savewrapped__", False
+        )
+
 
 class TestAbstractVersionedDataset:
     def test_version_str_repr(self, load_version, save_version):
@@ -431,3 +439,135 @@ class TestAbstractVersionedDataset:
 
         my_versioned_dataset._release()
         assert my_versioned_dataset._version_cache.currsize == 0
+
+
+class MyLegacyDataset(AbstractDataset):
+    def __init__(self, filepath="", save_args=None, fs_args=None, var=None):
+        self._filepath = PurePosixPath(filepath)
+        self.save_args = save_args
+        self.fs_args = fs_args
+        self.var = var
+
+    def _describe(self):
+        return {"filepath": self._filepath, "var": self.var}
+
+    def _exists(self) -> bool:
+        return Path(self._filepath.as_posix()).exists()
+
+    def _load(self):
+        return pd.read_csv(self._filepath)
+
+    def _save(self, data: str) -> None:
+        with open(self._filepath, mode="w") as file:
+            file.write(data)
+
+
+class MyLegacyVersionedDataset(AbstractVersionedDataset[str, str]):
+    def __init__(  # noqa: PLR0913
+        self,
+        filepath: str,
+        version: Version = None,
+    ) -> None:
+        _fs_args: dict[Any, Any] = {}
+        _fs_args.setdefault("auto_mkdir", True)
+        protocol, path = get_protocol_and_path(filepath, version)
+
+        self._protocol = protocol
+        self._fs = fsspec.filesystem(self._protocol, **_fs_args)
+
+        super().__init__(
+            filepath=PurePosixPath(path),
+            version=version,
+            exists_function=self._fs.exists,
+            glob_function=self._fs.glob,
+        )
+
+    def _describe(self) -> dict[str, Any]:
+        return dict(filepath=self._filepath, version=self._version)
+
+    def _load(self) -> str:
+        load_path = get_filepath_str(self._get_load_path(), self._protocol)
+
+        with self._fs.open(load_path, mode="r") as fs_file:
+            return fs_file.read()
+
+    def _save(self, data: str) -> None:
+        save_path = get_filepath_str(self._get_save_path(), self._protocol)
+
+        with self._fs.open(save_path, mode="w") as fs_file:
+            fs_file.write(data)
+
+    def _exists(self) -> bool:
+        try:
+            load_path = get_filepath_str(self._get_load_path(), self._protocol)
+        except DatasetError:
+            return False
+
+        return self._fs.exists(load_path)
+
+
+@pytest.fixture
+def my_legacy_dataset(filepath_versioned, save_args, fs_args):
+    return MyLegacyDataset(
+        filepath=filepath_versioned, save_args=save_args, fs_args=fs_args
+    )
+
+
+@pytest.fixture
+def my_legacy_versioned_dataset(filepath_versioned, load_version, save_version):
+    return MyLegacyVersionedDataset(
+        filepath=filepath_versioned, version=Version(load_version, save_version)
+    )
+
+
+class TestLegacyLoadAndSave:
+    def test_saving_none(self, my_legacy_dataset):
+        """Check the error when attempting to save the dataset without
+        providing the data"""
+        pattern = r"Saving 'None' to a 'Dataset' is not allowed"
+        with pytest.raises(DatasetError, match=pattern):
+            my_legacy_dataset.save(None)
+
+    def test_saving_invalid_data(self, my_legacy_dataset, dummy_data):
+        pattern = r"Failed while saving data to data set"
+        with pytest.raises(DatasetError, match=pattern):
+            my_legacy_dataset.save(pd.DataFrame())
+
+    @pytest.mark.parametrize(
+        "load_version", ["2019-01-01T23.59.59.999Z"], indirect=True
+    )
+    @pytest.mark.parametrize(
+        "save_version", ["2019-01-02T00.00.00.000Z"], indirect=True
+    )
+    def test_save_version_warning(
+        self, my_legacy_versioned_dataset, load_version, save_version, dummy_data
+    ):
+        """Check the warning when saving to the path that differs from
+        the subsequent load path."""
+        pattern = (
+            f"Save version '{save_version}' did not match "
+            f"load version '{load_version}' for "
+            r"MyLegacyVersionedDataset\(.+\)"
+        )
+        with pytest.warns(UserWarning, match=pattern):
+            my_legacy_versioned_dataset.save(dummy_data)
+
+    def test_versioning_existing_dataset(
+        self, my_legacy_dataset, my_legacy_versioned_dataset, dummy_data
+    ):
+        """Check the error when attempting to save a versioned dataset on top of an
+        already existing (non-versioned) dataset."""
+        my_legacy_dataset.save(dummy_data)
+        assert my_legacy_dataset.exists()
+        assert my_legacy_dataset._filepath == my_legacy_versioned_dataset._filepath
+        pattern = (
+            f"(?=.*file with the same name already exists in the directory)"
+            f"(?=.*{my_legacy_versioned_dataset._filepath.parent.as_posix()})"
+        )
+        with pytest.raises(DatasetError, match=pattern):
+            my_legacy_versioned_dataset.save(dummy_data)
+
+        # Remove non-versioned dataset and try again
+        Path(my_legacy_dataset._filepath.as_posix()).unlink()
+        my_legacy_versioned_dataset.save(dummy_data)
+        assert my_legacy_versioned_dataset.exists()