feat: Adding preview support to yaml datasets (#718)

* Adding preview support to yaml datasets Signed-off-by: Lukas Innig <[email protected]> * added test Signed-off-by: Lukas Innig <[email protected]> * fix the test Signed-off-by: Lukas Innig <[email protected]> * formatting Signed-off-by: Lukas Innig <[email protected]> * Update pyspark Signed-off-by: Ankita Katiyar <[email protected]> * Add release notes Signed-off-by: Ankita Katiyar <[email protected]> --------- Signed-off-by: Lukas Innig <[email protected]> Signed-off-by: Ankita Katiyar <[email protected]> Co-authored-by: Ankita Katiyar <[email protected]> Co-authored-by: Ankita Katiyar <[email protected]>
kedro-org · Jun 11, 2024 · 8c15f03 · 8c15f03
1 parent 6b213c2
commit 8c15f03
Show file tree

Hide file tree

Showing 4 changed files with 43 additions and 1 deletion.
diff --git a/kedro-datasets/RELEASE.md b/kedro-datasets/RELEASE.md
@@ -9,6 +9,12 @@
 | `langchain.ChatCohereDataset`       | A dataset for loading a ChatCohere langchain model.       | `kedro_datasets_experimental.langchain` |
 | `langchain.OpenAIEmbeddingsDataset` | A dataset for loading a OpenAIEmbeddings langchain model. | `kedro_datasets_experimental.langchain` |
 | `langchain.ChatOpenAIDataset`       | A dataset for loading a ChatOpenAI langchain model.       | `kedro_datasets_experimental.langchain` |
+* Extended preview feature to `yaml.YAMLDataset`.
+
+## Community contributions
+
+Many thanks to the following Kedroids for contributing PRs to this release:
+* [Lukas Innig](https://github.com/derluke)
 
 
 # Release 3.0.1

diff --git a/kedro-datasets/kedro_datasets/yaml/yaml_dataset.py b/kedro-datasets/kedro_datasets/yaml/yaml_dataset.py
@@ -1,8 +1,10 @@
 """``YAMLDataset`` loads/saves data from/to a YAML file using an underlying
 filesystem (e.g.: local, S3, GCS). It uses PyYAML to handle the YAML file.
 """
+
 from __future__ import annotations
 
+import json
 from copy import deepcopy
 from pathlib import PurePosixPath
 from typing import Any
@@ -17,6 +19,8 @@
     get_protocol_and_path,
 )
 
+from kedro_datasets._typing import JSONPreview
+
 
 class YAMLDataset(AbstractVersionedDataset[dict, dict]):
     """``YAMLDataset`` loads/saves data from/to a YAML file using an underlying
@@ -157,3 +161,14 @@ def _invalidate_cache(self) -> None:
         """Invalidate underlying filesystem caches."""
         filepath = get_filepath_str(self._filepath, self._protocol)
         self._fs.invalidate_cache(filepath)
+
+    def preview(self) -> JSONPreview:
+        """
+        Generate a preview of the YAML dataset with a specified number of items.
+
+        Returns:
+            A string representing the YAML data for previewing.
+        """
+        data = self._load()
+
+        return JSONPreview(json.dumps(data))
diff --git a/kedro-datasets/pyproject.toml b/kedro-datasets/pyproject.toml
@@ -226,7 +226,7 @@ test = [
     "pyarrow>=7.0; python_version >= '3.11'",  # Adding to avoid numpy build errors
     "pyodbc~=5.0",
     "pyproj~=3.0",
-    "pyspark>=3.0, <3.4; python_version < '3.11'",
+    "pyspark>=3.0; python_version < '3.11'",
     "pyspark>=3.4; python_version >= '3.11'",
     "pytest-cov~=3.0",
     "pytest-mock>=1.7.1, <2.0",

diff --git a/kedro-datasets/tests/yaml/test_yaml_dataset.py b/kedro-datasets/tests/yaml/test_yaml_dataset.py
@@ -1,7 +1,10 @@
+import inspect
+import json
 from pathlib import Path, PurePosixPath
 
 import pandas as pd
 import pytest
+import yaml
 from fsspec.implementations.http import HTTPFileSystem
 from fsspec.implementations.local import LocalFileSystem
 from gcsfs import GCSFileSystem
@@ -207,3 +210,21 @@ def test_versioning_existing_dataset(
         Path(yaml_dataset._filepath.as_posix()).unlink()
         versioned_yaml_dataset.save(dummy_data)
         assert versioned_yaml_dataset.exists()
+
+    def test_preview(self, yaml_dataset, dummy_data):
+        """Test the preview method."""
+        yaml_dataset.save(dummy_data)
+        preview_data = yaml_dataset.preview()
+
+        # Load the data directly for comparison
+        with yaml_dataset._fs.open(yaml_dataset._get_load_path(), mode="r") as fs_file:
+            full_data = yaml.safe_load(fs_file)
+
+        expected_data = json.dumps(full_data)
+
+        assert (
+            preview_data == expected_data
+        ), "The preview data does not match the expected data."
+        assert (
+            inspect.signature(yaml_dataset.preview).return_annotation == "JSONPreview"
+        )