Repository: kedro-org/kedro
Pull request: Implement support for schema in saving dask parquet (#1736)
Signed-off-by: Aivin V. Solatorio <avsolatorio@gmail.com>

## Description
Resolves: https://github.com/kedro-org/kedro/issues/1736

## Development notes
This change allows the use of the schema argument in the `dask.to_parquet` API from kedro's `dask.ParquetDataSet`.

A custom parser for the schema is implemented. The parser supports all kinds of schema declaration accepted by the underlying `dask.to_parquet` API.

The documentation was updated to show an example of the grammar for defining the schema in the `catalog.yml` 

<!-- What have you changed, and how has this been tested? -->
The change requires the parsing of the `_save_args` for the `schema` key and handles the transformation of the fields to a `pyarrow.DataType` or `pyarrow.Schema` accordingly.

Tests have also been written for this change.

## Checklist

- [x] Read the [contributing](https://github.com/kedro-org/kedro/blob/main/CONTRIBUTING.md) guidelines
- [ ] Opened this PR as a 'Draft Pull Request' if it is work-in-progress
- [x] Updated the documentation to reflect the code changes
- [x] Added a description of this change in the [`RELEASE.md`](https://github.com/kedro-org/kedro/blob/main/RELEASE.md) file
- [x] Added tests to cover my changes


<a href="https://gitpod.io/#https://github.com/kedro-org/kedro/pull/1746"><img src="https://gitpod.io/button/open-in-gitpod.svg"/></a>



Problem statement:
Support for defining a pyarrow schema for dask.ParquetDataSet in catalog.yml
## Description
I am working on a pipeline that outputs a dask ParquetDataSet. The inferred schema is not accurate, so I need to specify the schema manually. I don't think the current implementation supports the definition of complex schema, e.g., nested list of int64.

## Context
This change will allow users to be able to have more expressivity in the kind of dataset they can output. This is significantly important when the automatic inference of the schema fails.

## Possible Implementation
I have implemented a local change that worked. It checks for the presence of a schema key in the save_args, then assumes that the items are pyarrow definitions of the expected field types for the columns.

I can do a PR to add this.

## Possible Alternatives
I haven't seen any other alternatives or workaround. I would be interested to know if this can be done in any other way.

Hints from issue discussion:
We'd very much appreciate a PR @avsolatorio ! This sounds like a useful addition.
> We'd very much appreciate a PR @avsolatorio ! This sounds like a useful addition.

PR opened for this @datajoely : https://github.com/kedro-org/kedro/pull/1746 :)

Commit message: Implement support for schema in saving dask parquet (#1736)

Signed-off-by: Aivin V. Solatorio <avsolatorio@gmail.com>

Patch to review:
diff --git a/RELEASE.md b/RELEASE.md
index c85f122778..37c291fee3 100644
--- a/RELEASE.md
+++ b/RELEASE.md
@@ -22,6 +22,7 @@
 | `video.VideoDataSet`                 | Read and write video files from a filesystem                               | `kedro.extras.datasets.video` |
 | `video.video_dataset.SequenceVideo`  | Create a video object from an iterable sequence to use with `VideoDataSet` | `kedro.extras.datasets.video` |
 | `video.video_dataset.GeneratorVideo` | Create a video object from a generator to use with `VideoDataSet`          | `kedro.extras.datasets.video` |
+* Implement support for a functional definition of schema in `dask.ParquetDataSet` to work with the `dask.to_parquet` API.
 
 ## Bug fixes and other changes
 * Fixed `kedro micropkg pull` for packages on PyPI.
diff --git a/kedro/extras/datasets/dask/parquet_dataset.py b/kedro/extras/datasets/dask/parquet_dataset.py
index 2148aa30a1..4f29c799f5 100644
--- a/kedro/extras/datasets/dask/parquet_dataset.py
+++ b/kedro/extras/datasets/dask/parquet_dataset.py
@@ -6,6 +6,7 @@
 
 import dask.dataframe as dd
 import fsspec
+import triad
 
 from kedro.io.core import AbstractDataSet, get_protocol_and_path
 
@@ -46,7 +47,7 @@ class ParquetDataSet(AbstractDataSet[dd.DataFrame, dd.DataFrame]):
             >>> import dask.dataframe as dd
             >>>
             >>> data = pd.DataFrame({'col1': [1, 2], 'col2': [4, 5],
-            >>>                      'col3': [5, 6]})
+            >>>                      'col3': [[5, 6], [7, 8]]})
             >>> ddf = dd.from_pandas(data, npartitions=2)
             >>>
             >>> data_set = ParquetDataSet(
@@ -63,6 +64,29 @@ class ParquetDataSet(AbstractDataSet[dd.DataFrame, dd.DataFrame]):
             >>> reloaded = data_set.load()
             >>>
             >>> assert ddf.compute().equals(reloaded.compute())
+
+    The output schema can also be explicitly specified using Triad's grammar.
+    This is processed to map specific columns into pyarrow field types or schema.
+
+    References:
+    https://triad.readthedocs.io/en/latest/api/triad.collections.html#module-triad.collections.schema
+    https://arrow.apache.org/docs/python/api/datatypes.html
+
+        .. code-block:: yaml
+
+            >>> parquet_dataset:
+            >>>   type: dask.ParquetDataSet
+            >>>   filepath: "s3://bucket_name/path/to/folder"
+            >>>   credentials:
+            >>>     client_kwargs:
+            >>>       aws_access_key_id: YOUR_KEY
+            >>>       aws_secret_access_key: "YOUR SECRET"
+            >>>   save_args:
+            >>>     compression: GZIP
+            >>>     schema:
+            >>>       col1: [int32]
+            >>>       col2: [int32]
+            >>>       col3: [[int32]]
     """
 
     DEFAULT_LOAD_ARGS = {}  # type: Dict[str, Any]
@@ -128,8 +152,63 @@ def _load(self) -> dd.DataFrame:
         )
 
     def _save(self, data: dd.DataFrame) -> None:
+        self._process_schema()
         data.to_parquet(self._filepath, storage_options=self.fs_args, **self._save_args)
 
+    def _process_schema(self) -> None:
+        """This method processes the schema in the catalog.yml or the API, if provided.
+        This assumes that the schema is specified using Triad's grammar for
+        schema definition.
+
+        When the value of the `schema` variable is a string, it is assumed that
+        it corresponds to the full schema specification for the data.
+
+        Alternatively, if the `schema` is specified as a dictionary, then only the
+        columns that are specified will be strictly mapped to a field type. The other
+        unspecified columns, if present, will be inferred from the data.
+
+        This method converts the Triad-parsed schema into a pyarrow schema.
+        The output directly supports Dask's specifications for providing a schema
+        when saving to a parquet file.
+
+        Note that if a `pa.Schema` object is passed directly in the `schema` argument, no
+        processing will be done. Additionally, the behavior when passing a `pa.Schema`
+        object is assumed to be consistent with how Dask sees it. That is, it should fully
+        define the  schema for all fields.
+        """
+        schema = self._save_args.get("schema")
+
+        if isinstance(schema, dict):
+            # The schema may contain values of different types, e.g., pa.DataType, Python types,
+            # strings, etc. The latter requires a transformation, then we use triad handle all
+            # other value types.
+
+            # Create a schema from values that triad can handle directly
+            triad_schema = triad.Schema(
+                {k: v for k, v in schema.items() if not isinstance(v, str)}
+            )
+
+            # Handle the schema keys that are represented as string and add them to the triad schema
+            triad_schema.update(
+                triad.Schema(
+                    ",".join(
+                        [
+                            ":".join([k, v])
+                            for k, v in schema.items()
+                            if isinstance(v, str)
+                        ]
+                    )
+                )
+            )
+
+            # Update the schema argument with the normalized schema
+            self._save_args["schema"].update(
+                {col: field.type for col, field in triad_schema.items()}
+            )
+
+        elif isinstance(schema, str):
+            self._save_args["schema"] = triad.Schema(schema).pyarrow_schema
+
     def _exists(self) -> bool:
         protocol = get_protocol_and_path(self._filepath)[0]
         file_system = fsspec.filesystem(protocol=protocol, **self.fs_args)
diff --git a/setup.py b/setup.py
index 665e2c98ec..e6e9c0f11d 100644
--- a/setup.py
+++ b/setup.py
@@ -53,7 +53,7 @@ def _collect_requirements(requires):
 
 api_require = {"api.APIDataSet": ["requests~=2.20"]}
 biosequence_require = {"biosequence.BioSequenceDataSet": ["biopython~=1.73"]}
-dask_require = {"dask.ParquetDataSet": ["dask[complete]~=2021.10"]}
+dask_require = {"dask.ParquetDataSet": ["dask[complete]~=2021.10", "triad>=0.6.7, <1.0"]}
 geopandas_require = {
     "geopandas.GeoJSONDataSet": ["geopandas>=0.6.0, <1.0", "pyproj~=3.0"]
 }
diff --git a/test_requirements.txt b/test_requirements.txt
index f58ae352aa..0a8ef425a1 100644
--- a/test_requirements.txt
+++ b/test_requirements.txt
@@ -56,5 +56,6 @@ SQLAlchemy~=1.2
 tables~=3.6.0; platform_system == "Windows" and python_version<'3.9'
 tables~=3.6; platform_system != "Windows"
 tensorflow~=2.0
+triad>=0.6.7, <1.0
 trufflehog~=2.1
 xlsxwriter~=1.0
diff --git a/tests/extras/datasets/dask/test_parquet_dataset.py b/tests/extras/datasets/dask/test_parquet_dataset.py
index 6a07b87236..ca78ad1956 100644
--- a/tests/extras/datasets/dask/test_parquet_dataset.py
+++ b/tests/extras/datasets/dask/test_parquet_dataset.py
@@ -155,8 +155,69 @@ def test_load_extra_params(self, s3_data_set, load_args):
     )
     def test_save_extra_params(self, s3_data_set, save_args):
         """Test overriding the default save arguments."""
+        s3_data_set._process_schema()
+        assert s3_data_set._save_args.get("schema") is None
+
         for key, value in save_args.items():
             assert s3_data_set._save_args[key] == value
 
         for key, value in s3_data_set.DEFAULT_SAVE_ARGS.items():
             assert s3_data_set._save_args[key] == value
+
+    @pytest.mark.parametrize(
+        "save_args",
+        [{"schema": {"col1": "[[int64]]", "col2": "string"}}],
+        indirect=True,
+    )
+    def test_save_extra_params_schema_dict(self, s3_data_set, save_args):
+        """Test setting the schema as dictionary of pyarrow column types
+        in save arguments."""
+
+        for key, value in save_args["schema"].items():
+            assert s3_data_set._save_args["schema"][key] == value
+
+        s3_data_set._process_schema()
+
+        for field in s3_data_set._save_args["schema"].values():
+            assert isinstance(field, pa.DataType)
+
+    @pytest.mark.parametrize(
+        "save_args",
+        [
+            {
+                "schema": {
+                    "col1": "[[int64]]",
+                    "col2": "string",
+                    "col3": float,
+                    "col4": pa.int64(),
+                }
+            }
+        ],
+        indirect=True,
+    )
+    def test_save_extra_params_schema_dict_mixed_types(self, s3_data_set, save_args):
+        """Test setting the schema as dictionary of mixed value types
+        in save arguments."""
+
+        for key, value in save_args["schema"].items():
+            assert s3_data_set._save_args["schema"][key] == value
+
+        s3_data_set._process_schema()
+
+        for field in s3_data_set._save_args["schema"].values():
+            assert isinstance(field, pa.DataType)
+
+    @pytest.mark.parametrize(
+        "save_args",
+        [{"schema": "c1:[int64],c2:int64"}],
+        indirect=True,
+    )
+    def test_save_extra_params_schema_str_schema_fields(self, s3_data_set, save_args):
+        """Test setting the schema as string pyarrow schema (list of fields)
+        in save arguments."""
+
+        assert s3_data_set._save_args["schema"] == save_args["schema"]
+
+        s3_data_set._process_schema()
+
+        assert isinstance(s3_data_set._save_args["schema"], pa.Schema)


Code review comments:
diff_hunk: @@ -128,8 +152,63 @@ def _load(self) -> dd.DataFrame:
         )
 
     def _save(self, data: dd.DataFrame) -> None:
+        self._process_schema()
         data.to_parquet(self._filepath, storage_options=self.fs_args, **self._save_args)
 
+    def _process_schema(self) -> None:
+        """This method processes the schema in the catalog.yml or the API, if provided.
+        This assumes that the schema is specified using Triad's grammar for
+        schema definition.
+
+        When the value of the `schema` variable is a string, it is assumed that
+        it corresponds to the full schema specification for the data.
+
+        Alternatively, if the `schema` is specified as a dictionary, then only the
+        columns that are specified will be strictly mapped to a field type. The other
+        unspecified columns, if present, will be inferred from the data.
+
+        This method converts the Triad-parsed schema into a pyarrow schema.
+        The output directly supports Dask's specifications for providing a schema
+        when saving to a parquet file.
+
+        Note that if a `pa.Schema` object is passed directly in the `schema` argument, no
+        processing will be done. Additionally, the behavior when passing a `pa.Schema`
+        object is assumed to be consistent with how Dask sees it. That is, it should fully
+        define the  schema for all fields.
+        """
+        schema = self._save_args.get("schema")
+
+        if isinstance(schema, dict):
+            # The schema may contain values of different types, e.g., pa.DataType, Python types,
+            # strings, etc. The latter requires a transformation, then we use triad handle all
+            # other value types.
+
+            # Create a schema from values that triad can handle directly
+            triad_schema = triad.Schema(
+                {k: v for k, v in schema.items() if not isinstance(v, str)}
+            )
+
+            # Handle the schema keys that are represented as string and add them to the triad schema
+            triad_schema.update(
+                triad.Schema(
+                    ",".join(
+                        [
+                            ":".join([k, v])
path: kedro/extras/datasets/dask/parquet_dataset.py
text: @user1:
```suggestion
                            f"{k}:{v}"
```

Nit: I think this should be faster (not terribly important), but it also feels more readable to me. `k + ":" + v` is also fine, I guess.