Easier CustomDataset Creation

[lancechua]

Description

IO read write functions typically follow this signature:

def load(path: Union[Path, str], **kwars) -> obj:
    ...

def save(obj, path: Union[Path, str], **kwars) -> None:
    ...

Creating custom datasets should ideally be as easy as supplying load / save function(s) that follow this function signature.

Context

kedro supports an extensive range of datasets, but it is not exhaustive. Popular libraries used by relatively niche communities like xarray and arviz aren't currently supported.

Beyond these examples, unofficially adding support for more obscure datasets would be easier.

Initially, I was looking to implement something like this and asked in the Discord chat if this pattern made sense. Then, @datajoely suggested I open a feature request.

Possible Implementation

We can consider a Dataset class factory. Maybe GenericDataset with a class method .create_custom_dataset(name: str, load: callable, save: callable).

Usage would look something like xarrayNetCDF = GenericDataset("xarrayNetCDF", xarray.open_dataset, lambda x, path, **kwargs: x.to_netcdf(path, **kwargs)). Entries can be added to the data catalog yaml just as with any other custom dataset implementation.

Possible Alternatives

• LambdaDataset is very similar but the load, and save are hard coded in the implementation, and cannot be parameterized in the data catalog, as far as I'm aware
• Subclassing AbstractDataset is an option, but this feature request seeks to reduce boilerplate when defining new datasets
• Adding xarray support #1346 officially requires implementing nuances like cloud file storage, partitioned datasets, lazy loading, etc.

[datajoely]

Hi @lancechua thank you for raising this! We have discussed this as a team and have had a good conversation regarding ways we can proceed.

You will see that this issue has now been referenced in the kedro-org/kedro#1778 super issue which is tracking all of the ways we plan to improve the DataCatalog at a low level.

Your last point:

Adding xarray support https://github.com/kedro-org/kedro-plugins/issues/165 officially requires implementing nuances like cloud file storage, partitioned datasets, lazy loading, etc.

Is the most interesting, a lot of the paint-points to mention have emerged organically. Versioning, fsspec etc. have all arrived in Kedro way later than the original AbstractDataSet class definition, implementing a new DataSet today requires (probably too much) copy and pasting.

The team are going to have a go at prototyping what this could look like, so you may be interested in subscribing to parent issue to contribute to these design decisions. The initial thinking reflected on the rather flat structure we have today where everything inherits from AbstractDataSet (or AbstractVersionedDataSet) and we could utilise OO concepts much better and make implementation/testing/overriding much simpler.

[deepyaman]

Is the most interesting, a lot of the paint-points to mention have emerged organically. Versioning, fsspec etc. have all arrived in Kedro way later than the original AbstractDataSet class definition, implementing a new DataSet today requires (probably too much) copy and pasting.

I feel like too much copy-pasting has been an issue for a really long time. 😅

I would be keen to break down the points of duplication (e.g. default argument handling, fsspec-related boilerplate, etc.) and add mixins for each (as in kedro-org/kedro#15), because that way we can make progress on resolving these pain points without having to limit ourselves to a universal solution that works for every possible dataset.

[astrojuanlu]

Collecting some thoughts from our latest Tech Design session about this topic:

Recent developments

• Linked issues: https://github.com/kedro-org/kedro/issues/1778, https://github.com/kedro-org/kedro/issues/2409
• Some recent experiences trying to add "weird" datasets
  • My own Kaggle dataset https://github.com/astrojuanlu/kedro-kaggle-dataset
  • FITS dataset using Astropy, because SunPy and Astropy https://github.com/b4rlw/sunspot-segmentation/blob/master/src/sunspots/extras/datasets/sunpy_datasets.py
  • "It is also worth noting that a HDU cannot be selected within the node logic, because even though the FITS file metadata might be loaded, the data inside a particular HDU cannot be loaded unless the corresponding file is still open, and by the time a dataset arrives at a node, it is not."
  • HDF dataset, recently reopened https://github.com/kedro-org/kedro-plugins/issues/240

General thoughts

• There are 3 dimensions to take into account
  • Where data is
  • What format is
  • What memory representation
• Ideally a design would decouple the three, but this is not practical, and since Kedro offloads data loading to other libraries, the dataset is subject to the abstractions by those other libraries and how they solve the decoupling
• The old design was very painful and using fsspec greatly improved it https://docs.kedro.org/en/0.14.3/kedro.io.html#data-sets
• There is some sort of decoupling with pandas.GenericDataSet https://docs.kedro.org/en/latest/kedro_datasets.pandas.GenericDataSet.html but it cannot be easily achieved with everything

Current issues

• Lots of copy-pasting
  • @deepyaman's comment right above https://github.com/kedro-org/kedro/issues/1936#issuecomment-1297213053
  • "Relatedly, why is _describe an abstract method at all? People just have to write way too much code when writing a custom dataset"
  • There's a suggestion to add a fsspec mixin https://github.com/kedro-org/kedro-plugins/issues/200
    • But multiple inheritance brings some problems
    • AbstractDataSet base class contains lots of logic, can't be easily turned into a protocol https://github.com/kedro-org/kedro/issues/2409#issuecomment-1490596283
• Unclear how to support "lazy" datasets
  • Case in point: HDF, FITS, others
  • Datasets are eager, so if the data has to be lazily loaded, users have to jump through some hoops
• Somewhat unclear separation of "format" and "transport"
  • Not really an issue of the current design, but we could do a better job at communicating it, or just assume that people will do whatever works for them
• Unclear if full reproducibility is encouraged
  • Case in point: Kaggle
  • Another: "My experience loading COVID data https://github.com/deepyaman/inauditus/blob/develop/refresh-data I just write a script to download the data from an external source, and then treat it as a pandas.CSVDataSet. There are other problems with doing the “expensive” data download in a load function, e.g. that we’re loading from a remote source each time when we really should just have it in a downloaded source."
• Unclear interaction with versioning
  • "And it is more subtle about versioning - it actually works only if you use a specific function. The newly ManagedTableDataset which use the built-in versioning actually doesn’t work like this."
• Creating a custom dataset requires a full Kedro installation, which some users consider too heavy
  • The topic of https://github.com/kedro-org/kedro/issues/2409

Other links and ideas

• Any redesign should be grounded on user research and user needs
• @antonymilne suggested making dataset parameters keyword-only, which enables non-breaking refactoring post-0.19.0
  • If we think this is worth doing, we should open an issue and prioritize it sooner rather than later
• GetInData "pipeline dataset" https://github.com/getindata/kedro-azureml/blob/develop/kedro_azureml/datasets/pipeline_dataset.py

[astrojuanlu]

A user that quite likely got confused when defining their own dataset, most likely because they aren't copy-pasting enough https://www.linen.dev/s/kedro/t/12646523/hi-folks-i-m-trying-to-create-a-custom-kedro-dataset-inherit#ac22d4cb-dd72-4072-905a-f9c5cb352205

[yetudada]

I'm also just dropping data from some of our surveys here, the question is "What is the hardest part of working with the Kedro framework?" (Kedro - User Research Study)

• IO abstraction, I found its bit tricky if I want to implement any new IO module
• when I have to work on non-popular data sets such as Excel
• Writing custom IO datasets. Specially Incremental Data Ingestion. Finding rationale behind different kedro layer (raw, stage, int etc) specifically when to use which and when to skip.
• Custom datasets perhaps

[astrojuanlu]

Interesting use case: a user trying to perform data access with a highly customized wget command https://linen-slack.kedro.org/t/14156294/i-m-migrating-to-kedro-for-a-bunch-of-my-projects-for-one-of#a3d857e1-7e31-48f9-98de-439a936786c8

This is not too different from my KaggleDataset attempt. Maybe there should be a "generic dataset" that allows people to just type a subprocess that will take care of leaving some files in a directory.

[datajoely]

Channelling my inner @idanov this does harm reproducibility - personally I think it's a good thing, same for conditional nodes, another things we've held off doing for the same reason.

[astrojuanlu]

This is a bit of a tangent, but I think it's worth addressing.

The reality is that people who need "weird" things will do them, whether are supported by Kedro or not. This not only applies to datasets: I've seen people do very weird things with Kedro.

Speaking of datasets, the options our users have are:

• Put the "non reproducible" part to a separate step, and build the rest as a Kedro project. (What @deepyaman did in his project, and what the user I referred to in my previous comment declared they would do). I contend this is bad DX.
• Hack around their own brittle dataset that does this. Which is a difficult thing to do in any case (see this issue). So not only this is bad DX, but also they're probably shooting themselves in the foot and we're not helping.
• Find an alternative way, assuming it's possible. This is good for our users enlightenment, but still screams "you're doing it wrong", which is something that should be handled with care from the Kedro side (with, at the very least, good docs that are easily discoverable). I'm on the fence on whether this is bad DX or not - frameworks definitely impose some constraints. But should they have escape hatches?
• Say "Kedro is not for me" and try something else. This is bad for Kedro.

There are 2 stances we can take:

1. Consider the bad DX a feature. Friction means the users will try to do something else. Or, put another way, what they're trying to do is out of scope for Kedro so we refuse to implement it ourselves.
2. Consider the bad DX a bug. Friction means we should improve the situation for these users, while helping them do the right thing as much as we can, but offering options to do all their data processing within Kedro.

I'm firmly in camp (2).

[datajoely]

Completely aligned on #2 I think it's related to the wider rigidity questions we have in the framework vs lib debate.

FYI - how prefect do conditionals here looks neat https://github.com/PrefectHQ/prefect/discussions/3545#discussioncomment-106575

[astrojuanlu]

Another user who wants to read a TSV file contained in a .tar.gz https://stackoverflow.com/q/76994906/554319

pandas does not support compressed archives with more than one file, so this requires an intermediate step.

[astrojuanlu]

After implementing some new datasets that will hopefully be open sourced soon, I'm adding some more reflections here.

I think that the difficult part lies in three specific points:

1. Having to overload private methods of an obscure class (that then don't get documented, see any dataset in https://docs.kedro.org/en/0.18.13/kedro_datasets.html),
2. The fsspec boilerplate, and
3. The versioning stuff (#1979, #2355)

(1) was hinted by @lancechua's original comment at the top of this issue, and I agree that we should move towards an architecture that makes load and save public methods. Moreover, I'm advocating to favor composition over inheritance.

(2) and (3) are trickier though. What's interesting is that one can envision datasets that don't refer to specific files on disk and that are not versioned - and these ones are actually quite easy to implement. As such, adding the path to load as @lancechua suggested might be a too strong assumption.

Another reason why this is hard is the naming itself: some things we name "datasets" in Kedro are not really a /ˈdā-tə-ˌset/ in the strict sense of the word, but more like "artifacts" or "connectors". I think there is an opportunity here to collect some of the "weird" use cases we already have and come up with better terminology.

[lancechua]

(1) was hinted by @lancechua's original comment at the top of this issue, and I agree that we should move towards an architecture that makes load and save public methods. Moreover, I'm advocating to favor composition over inheritance.

Can't remember who linked it in the issues here somewhere, but after reading the Python subclassing redux article, I also think composition is the better way to go by injecting load and save callables as dependencies.

(2) and (3) are trickier though. What's interesting is that one can envision datasets that don't refer to specific files on disk and that are not versioned - and these ones are actually quite easy to implement. As such, adding the path to load as @lancechua suggested might be a too strong assumption.

Also agree. The callables can be fully specified by load_kwargs (and potentially load_args) so there isn't a hard requirement for the function signature.

From what I remember about the implementation of versioning, the filename becomes a directory, with timestamp sub-directories. A separate versioner dependency can be injected to redirect the load and save callables to the path of the correct version. This would make the path parameter a hard requirement, but we can always just pass None for "pathless" datasets. The no versioning case would just return the path as is.

A use case that would be good, but potentially annoying to implement, is a cache for expensive "load" operations like web pulls or API requests. Instead of repeating the expensive request, the flow with subsequent reads would be:

expensive first request -> in memory -> save to disk -> read cache from disk -> in memory -> ...

The cache can be a dispatcher that takes both the original load callable, and a load cache callable. cache itself must still be a valid load callable. Perhaps we namespace original and cache load params or create a separate load_cache_kwargs parameter?

The load callable dispatched would depend on:

• dataset parameters (e.g. use_cache=True)
• cache validity (e.g. based on hash of params).

The no caching case would just always dispatch to the original load callable.

P.S. Haven't worked on kedro datasets for a while now, but hopefully the above still makes sense.

[astrojuanlu]

Thanks a lot for getting back @lancechua! Yes, Hynek has influenced me a lot on this as well. Hope we can prioritize this in 2024.

[astrojuanlu]

More evidence of different artifact types: Kubeflow https://www.kubeflow.org/docs/components/pipelines/v2/data-types/artifacts/#artifact-types (via @deepyaman)

[astrojuanlu]

Another problem with the current design: when doing post-mortem debugging on datasets, one gets thrown to AbstractDataset.load:

https://github.com/kedro-org/kedro/blob/4d186ef38adb3a06d64bb35aefdc18af93dc0422/kedro/io/core.py#L191-L201

[datajoely]

Now Python 3.7 is EOL is there any advantage to using Protocol types as an alternative/addition to the abstract classes?

[astrojuanlu]

I proposed it in another issue and @antonymilne made some good remarks https://github.com/kedro-org/kedro/issues/2409#issuecomment-1490989279 I could be wrong here and in any case I defer the implementation details to the engineering team, but I don't think this can be done without "breaking" the current catalog + abstractdataset model. This would essentially be a "catalog 2.0", which is exciting but also kind of scary.

[astrojuanlu]

A user confused about what AbstractDataset does https://linen-slack.kedro.org/t/16006502/hey-team-i-have-query-regarding-custom-datasets-we-are-imple#6bdb0d80-0f60-41a1-a104-f11fd3877f08

[datajoely]

On that last part - we're well overdue a native PyTorch dataset

[astrojuanlu]

Another thing to note: some of our underlying datasets have their own mechanisms of reading remote data that don't rely on fsspec https://github.com/pola-rs/polars/pull/13044 plus sometimes our "copy paste" fsspec logic is really obscure and easy to get wrong https://github.com/kedro-org/kedro-plugins/pull/360#issuecomment-1834634344

Making our "fsspec adapter" (not a mixin please) more explicit will make life easier for custom dataset authors and probably for us.

[astrojuanlu]

Progress towards making load and save public:

Edit: Actually, thinking about it more, should just make load and save abstract instead of _load and _save, to make the new interface clear. Even if just _load and _save is being implemented, you mentioned it's going into __init_subclass__, and load and save will get created there.

Originally posted by @deepyaman in https://github.com/kedro-org/kedro/issues/3920#issuecomment-2191616499

[astrojuanlu]

Also, interesting perspective on object stores vs filesystems: https://docs.rs/object_store/latest/object_store/#why-not-a-filesystem-interface

This design provides the following advantages:

• All operations are atomic, and readers cannot observe partial and/or failed writes
• Methods map directly to object store APIs, providing both efficiency and predictability
• Abstracts away filesystem and operating system specific quirks, ensuring portability
• Allows for functionality not native to filesystems, such as operation preconditions and atomic multipart uploads

Originally shared by @MatthiasRoels in https://github.com/kedro-org/kedro-plugins/issues/625#issuecomment-2195327128