feat(datasets): Add rioxarray and RasterDataset

[tgoelles]

Description

• Added GeoTiffDataset to read and write geotiff files with rioxarray and work with xarray in kedro. This is also related to issue Xarray NetCDFDataSet #165

Development notes

Added kedro_datasets/xarray/geotiff_dataset.py and a very basic framework for testing in /kedro-datasets/tests/xarray/test_geotiff_dataset.py using a sample file from rioxarray. I could not run the test suit locally or on github due to the complex environment. So more tests are certainly needed.
I think it would also be easy to make a preview method.

Checklist

• Opened this PR as a 'Draft Pull Request' if it is work-in-progress
• Updated the documentation to reflect the code changes
• Added a description of this change in the relevant RELEASE.md file
• Added tests to cover my changes

[astrojuanlu]

Thanks @tgoelles for this pull request 🚀

I could not run the test suit locally or on github due to the complex environment.

What instructions did you follow? Would love to understand how can we make the process easier. It's definitely a challenge to install all the dependencies of all datasets, but we could try to make running the tests of only 1 dataset easier.

[tgoelles]

Thanks @tgoelles for this pull request 🚀

I could not run the test suit locally or on github due to the complex environment.

What instructions did you follow? Would love to understand how can we make the process easier. It's definitely a challenge to install all the dependencies of all datasets, but we could try to make running the tests of only 1 dataset easier.

I did not follow any instruction. Is there one? Yes It would be great to just test one dataset and not on windows at the moment with the github actions. VS code debugging does not work for some reason: discovery error. Also if I want to run it in the terminal I get this:

"no tests ran" even when I runnning yaml for example

[tgoelles]

A basic test suite now works. There I need help, but the basic functionality works and I guess the rest can be added by someone who knows kedro better

[merelcht]

Thanks for this contribution @tgoelles! I've left some minor comments, but I gather you'll need some help to finish this in terms of tests?

It might take some time for the team to help out, but I'll make sure your PR is not forgotten.

[astrojuanlu]

Just one broad question: are there any parts of this dataset that are GeoTIFF-specific, and that prevent it to work on generic n-dimensional labeled arrays from xarray?

[astrojuanlu]

I did not follow any instruction. Is there one? Yes It would be great to just test one dataset and not on windows at the moment with the github actions. VS code debugging does not work for some reason: discovery error. Also if I want to run it in the terminal I get this:

After you commented this I realized that this plugin didn't have a dedicated CONTRIBUTING.md (#364) and that's fixed now: https://github.com/kedro-org/kedro-plugins/blob/main/kedro-datasets/CONTRIBUTING.md#unit-tests-100-coverage-pytest-pytest-cov

This however runs all the tests. You can try pytest -k xarray to only run those related to xarray. Let me know if that helps.

[tgoelles]

Just one broad question: are there any parts of this dataset that are GeoTIFF-specific, and that prevent it to work on generic n-dimensional labeled arrays from xarray?

It uses rioxarray to handle GeoTiffs: https://corteva.github.io/rioxarray/stable/getting_started/getting_started.html

[RobertoPalomares]

Why the rioxarray dependency is not >=0.15.0 as above?

[RobertoPalomares]

why the is rtol needed? It is causing issues with resolution or transform parameter. Pure curiosity

[nickofca-mck]

Looks good overall. Would be curious about the dependency on rioxarray 0.15.0. rioxarray 0.14.0+ dropped support of python 3.8 (https://corteva.github.io/rioxarray/stable/history.html#id2), so it could constrain 3.7 & 3.8 users (which base kedro supports for now).

[nickofca-mck]

rioxarray 0.14.0+ dropped support of python 3.8 (https://corteva.github.io/rioxarray/stable/history.html#id2)

[astrojuanlu]

kedro 0.19.0 will drop support for Python 3.7. About Python 3.8, as per the current policy it will continue supporting it, but it's up for discussion kedro-org/kedro#2815

[astrojuanlu]

kedro-datasets will follow NEP 29, hence you can focus on Python 3.9+

[astrojuanlu]

Hello @tgoelles! We'd love to get this PR in. How can we help?

[tgoelles]

Hello @tgoelles! We'd love to get this PR in. How can we help?

I am currently super busy with other things. But I think its ready and someone needs to see that all the checks are Ok. I might have time in 2 weeks again for this.

[merelcht]

I've resolved the merge conflicts for this PR and done some small cleanups, but there's still more things to do:

1. The doctest is failing for the example provided in GeoTiffDataset. This is the python example inside the docstring.
2. Test coverage is not at 100% for GeoTiffDataset:

Name                                                         Stmts   Miss  Cover   Missing
------------------------------------------------------------------------------------------
kedro_datasets/xarray/geotiff_dataset.py     42      6    86%   100, 103, 117-118, 132-133

@tgoelles Can you finish this or do you want someone to take over?

[astrojuanlu]

cc @riley-brady FYI :)

[riley-brady]

Just some initial comments from scanning through this!

We have a highly specialized implementation here: https://github.com/McK-Internal/CRA-data_pipelines/blob/dev/src/cra_data_pipelines/datasets/geotiff.py.

I would think about saving/loading from S3, because I think geotiff might suffer from the same thing as NetCDF and not allow streaming from cloud buckets (unlike COGs). So you might need to do the temporary file trick that we did with the NetCDF data.

There's also some preprocessing that needs to be done for concatenating a set of globbed files.

[riley-brady]

Wondering if this should be GeoTIFFDataset to follow conventional naming for GeoTIFF. For example, our NetCDF PR was NetCDFDataset.

[riley-brady]

Wondering on organization here. We return all netcdf objects as an xarray dataset here: https://github.com/kedro-org/kedro-plugins/pull/360/files. But organize it under kedro_datasets/netcdf/.

NetCDFs don't have to be loaded as xarray, but it's the most convenient platform IMO. It could also be loaded in with the NetCDF package or iris as a cube object.

Should the current NetCDFDataset object be moved to an xarray folder for clarity? Or should this PR just be housed under kedro_datasets/geotiff?

[riley-brady]

Might want to follow this commit: 210e4ed and replace these instances of get_filepath_str. It seemed to have caused an issue with tests breaking in the NetCDF PR.

[riley-brady]

Might want to add a detector for wildcard globbing and set up an _is_multifile like in #360. It's pretty common to have a set of GeoTIFF layers that need to be concatenated on load.

[riley-brady]

rioxarray for some reason ignores all tags inside a geotiff and doesn't assign them appropriately . We deal with this by dual-loading with rasterio.

            with rasterio.open(load_path) as data:
                # Pull in relevant tags.
                tags = data.tags()
            data = rioxarray.open_rasterio(load_path, **self._load_args).squeeze()
            data.attrs.update(tags)

[merelcht]

I've been doing some small bits fixing the doc string and added some missing tests, but while doing this I also realised I'm probably not the right person to finish the implementation. I'm not familiar at all with GeoTiff and its purpose. I noticed the dataset as it stands now doesn't have the credentials argument and is maybe missing some other bits when I compare it to e.g. NetCDFDataset that was merged yesterday.

Either someone with more knowledge about this data type should finish this PR or we can take this as our first "experimental" contribution (see #517) and get it in that way. WDYT @astrojuanlu?

[BielStela]

Hello! just a note: I think the name GeoTiffDataset is a bit misleading because rioxarray wraps rastertio.open() which can open any kind of raster dataset. It can open any format supported by gdal raster drivers. So maybe something in the line of RasterDataset would be more strict and avoid possible confusions.

[merelcht]

Hi @tgoelles, I'm aware that this PR has been open for quite a while and has had a lot of comments. I'm not sure the Kedro team can help at this point to get the dataset to a full ready state. However, we're about to launch our new experimental dataset contribution model, which basically means you can contribute datasets that are more experimental and don't have to have full test coverage etc here https://github.com/kedro-org/kedro-plugins/tree/main/kedro-datasets/kedro_datasets_experimental.

I think the GeoTiffDataset might be a good candidate to get contributed in that way. What do you think?

[tgoelles]

I think the GeoTiffDataset might be a good candidate to get contributed in that way. What do you think?

Yes I think that would be a possible way. Also as @BielStela pointed out it is for geospatial raster data. So there is also an overlap with the already now existing netcdf format.

I could be like this. RasterDataset => rioxarray.
Then we could also think about enforcing that a coordinate reference system (CRS) is used. Or at least make an option for this. Also this needs to work well together with other geospatial datasets like netcdf and geopandas.

Yes this needs more experimenting, but it would introduce a huge list of new file formats which can be used.

Will you move it to the experimental folder? I could then rework it with the new concept of dealing with raster data in general and not just GeoTiff. I also gained a lot of experience working with these kind of data recently.

[BielStela]

I'm happy to give a hand with this if needed @tgoelles :)

[merelcht]

@tgoelles I've moved it! There's some issues with the docs, which I'll continue to fix, but you should be able to work on the code 🙂

[riley-brady]

In light of the recent comments I would say this should be /raster/raster_dataset.py with this named RasterDataset. Given this is more than just xarray, which is just being used as a frontend to load in tiff datasets.

[riley-brady]

It would also be interesting to think about whether this supports COGs directly or needs an explicit kwarg to get those loaded in

[riley-brady]

Replicating a comment from earlier, but I've found that rioxarray can't handle loading in any tags associated with the raster layer. There might be a better solution, but we do

import rasterio

with rasterio.open(load_path) as data:
    # Pull in relevant tags like units and return period.
    tags = data.tags()
data = rioxarray.open_rasterio(load_path, **self._load_args).squeeze()
data.attrs.update(tags)

[riley-brady]

It's also nice to support multiple files if you want to concatenate many individual layers into a multi-dimensional xarray object. See how we do it for NetCDF here: https://github.com/kedro-org/kedro-plugins/blob/main/kedro-datasets/kedro_datasets/netcdf/netcdf_dataset.py.

self._is_multifile = (
        True if "*" in str(PurePosixPath(self._filepath).stem) else False
    )

if self._is_multifile:
  # Note that we require an explicit "concat_dim" entry in the dataset
  # catalog. There might be a better way to do this though.
  concat_dim = self._load_args.get("concat_dim", None)
  
  remote_files = fs.glob(load_path)
  if concat_dim:
      preprocess = partial(
          self._preproc_multiple_geotiff, concat_dim=concat_dim
      )
  data = xr.open_mfdataset(
      remote_files,
      combine_attrs="drop_conflicts",
      engine="rasterio",
      **self._load_args,
  ).sortby(coords)

def _preproc_multiple_geotiff(self, x: xr.DataArray, concat_dim: str):
    """Helper function for open_mfdataset to set up coordinates for concatenated
       geotiffs.

   `concat_dim` is a required argument in the DataCatalog and should also be an
    attribute/tag on the geotiff for labeling the coordinate.

    Args:
        x: Individual tif DataArray being concatenated (handled automatically by
            `open_mfdataset`).
        concat_dim: Name of the dimension to concatenate over.

    Returns:
        Individual DataArray slice to be concatenated with proper formatting,
        coordinates, attributes, etc.
    """
    # Squeeze out `band_data`
    x = x.squeeze()

    # Pull individual filepath so we can use `rasterio` to get the tag
    # attributes.
    path = x.encoding["source"]
    with rasterio.open(path) as data:
        # Pull in relevant tags like units and return period.
        tags = data.tags()
    
    # Note here that this makes this tedious as well. Need to have `concat_dim`
    # in the catalog as well as a tag in the tif to label the axis. I guess you
    # could not require this and would just leave the xarray coordinate unlabeled
    if concat_dim not in tags:
        raise ValueError(
            f"{concat_dim} needs to be a tag in geotiff files to help with "
            + "coordinate axis labeling."
        )
    x.attrs.update(tags)

    x.expand_dims(concat_dim)
    x = x.assign_coords({concat_dim: tags[concat_dim]})
    return x

[tgoelles]

I don't know why this pandas check failed. it is the same code as in the main branch. I tried to solve the DCO, following the instructions of it with git rebase HEAD~228 --signoff, but this than leads to a huge amount of conflicts. there must be an easier way out of this. Also I am afraid to make mistakes there.

[astrojuanlu]

I tried to solve the DCO, following the instructions of it with git rebase HEAD~228 --signoff, but this than leads to a huge amount of conflicts. there must be an easier way out of this. Also I am afraid to make mistakes there.

You're not the first one...

I see maintainers are allowed to push to this branch, so I'll try to get you some help from the team.

[astrojuanlu]

I did notice one thing that didn't work but this PR has been open for very long already so I don't want to block it.

I'm approving it and leaving it to the team to decide what to do with it.

Thansk a lot @tgoelles and also to all the people who reviewed this 🙏🏼

[astrojuanlu]

I tried using the dataset with a remote path and it didn't work:

In [17]: ds = GeoTIFFDataset(filepath="https://download.osgeo.org/geotiff/samples/GeogToWGS84GeoKey/GeogToWGS84GeoKey5.tif")

In [18]: ds._load()
---------------------------------------------------------------------------
CPLE_OpenFailedError                      Traceback (most recent call last)
File rasterio/_base.pyx:310, in rasterio._base.DatasetBase.__init__()

File rasterio/_base.pyx:221, in rasterio._base.open_dataset()

File rasterio/_err.pyx:221, in rasterio._err.exc_wrap_pointer()

CPLE_OpenFailedError: download.osgeo.org/geotiff/samples/GeogToWGS84GeoKey/GeogToWGS84GeoKey5.tif: No such file or directory

During handling of the above exception, another exception occurred:

RasterioIOError                           Traceback (most recent call last)
Cell In[18], line 1
----> 1 ds._load()

File ~/Projects/QuantumBlackLabs/Kedro/kedro-plugins/kedro-datasets/kedro_datasets_experimental/rioxarray/geotiff_dataset.py:127, in GeoTIFFDataset._load(self)
    125 def _load(self) -> xarray.DataArray:
    126     load_path = self._get_load_path().as_posix()
--> 127     with rasterio.open(load_path) as data:
    128         tags = data.tags()
    129     data = rxr.open_rasterio(load_path, **self._load_args)

File ~/Projects/QuantumBlackLabs/Kedro/kedro/.venv/lib/python3.11/site-packages/rasterio/env.py:451, in ensure_env_with_credentials.<locals>.wrapper(*args, **kwds)
    448     session = DummySession()
    450 with env_ctor(session=session):
--> 451     return f(*args, **kwds)

File ~/Projects/QuantumBlackLabs/Kedro/kedro/.venv/lib/python3.11/site-packages/rasterio/__init__.py:304, in open(fp, mode, driver, width, height, count, crs, transform, dtype, nodata, sharing, **kwargs)
    301 path = _parse_path(raw_dataset_path)
    303 if mode == "r":
--> 304     dataset = DatasetReader(path, driver=driver, sharing=sharing, **kwargs)
    305 elif mode == "r+":
    306     dataset = get_writer_for_path(path, driver=driver)(
    307         path, mode, driver=driver, sharing=sharing, **kwargs
    308     )

File rasterio/_base.pyx:312, in rasterio._base.DatasetBase.__init__()

RasterioIOError: download.osgeo.org/geotiff/samples/GeogToWGS84GeoKey/GeogToWGS84GeoKey5.tif: No such file or directory

But rasterio knows how to deal with remote paths:

In [15]: with rasterio.open("https://download.osgeo.org/geotiff/samples/GeogToWGS84GeoKey/GeogToWGS84GeoKey5.tif") as data:
    ...:     tags = data.tags()
    ...: 

In [16]: tags
Out[16]: 
{'AREA_OR_POINT': 'Area',
 'TIFFTAG_ARTIST': '',
 'TIFFTAG_DATETIME': '2008:03:01 10:28:18',
 'TIFFTAG_RESOLUTIONUNIT': '2 (pixels/inch)',
 'TIFFTAG_SOFTWARE': 'Paint Shop Pro 8.0',
 'TIFFTAG_XRESOLUTION': '300',
 'TIFFTAG_YRESOLUTION': '300'}

so maybe there's something wrong in how the load path is handled here.

[noklam]

@astrojuanlu Use get_filepath_str instead

The reason is that self._get_load_path doesn't handle protocols correctly. I suggest we merge this now and handle this separately as I think it probably affect more than 1 dataset

[astrojuanlu]

I suggest we merge this now and handle this separately as I think it probably affect more than 1 dataset

I agree!

[astrojuanlu]

CI failure is unrelated #740 (comment)

[astrojuanlu]

This is merged thanks again everyone for your patience!