diff --git a/Makefile b/Makefile index ec21b6679..eb71ebbcd 100644 --- a/Makefile +++ b/Makefile @@ -76,11 +76,12 @@ dataset-doctest%: exit 2; \ fi; \ \ - # The ignored datasets below require complicated setup with cloud/database clients or network model download which is overkill for the doctest examples. + # The ignored datasets below require complicated setup with cloud/database clients, optional dependencies (e.g. pyiceberg for iceberg_dataset), or network model download which is overkill for the doctest examples. cd kedro-datasets && pytest kedro_datasets --doctest-modules --doctest-continue-on-failure --no-cov \ --ignore kedro_datasets/huggingface/transformer_pipeline_dataset.py \ --ignore kedro_datasets/pandas/gbq_dataset.py \ --ignore kedro_datasets/partitions/partitioned_dataset.py \ + --ignore kedro_datasets/polars/iceberg_dataset.py \ --ignore kedro_datasets/redis/redis_dataset.py \ --ignore kedro_datasets/snowflake/snowpark_dataset.py \ --ignore kedro_datasets/spark/gbq_dataset.py \ diff --git a/kedro-datasets/RELEASE.md b/kedro-datasets/RELEASE.md index f9847dce2..c42b0c8e7 100755 --- a/kedro-datasets/RELEASE.md +++ b/kedro-datasets/RELEASE.md @@ -3,6 +3,7 @@ ## Major features and improvements - Added support for configuring external Hive table locations in `spark.SparkHiveDataset` through `save_args.path`. - Added standard Kedro versioning support to the experimental `netcdf.NetCDFDataset`, including local and remote (S3) files. +- Added `polars.IcebergDataset` to support loading and saving Apache Iceberg tables using Polars and PyIceberg. ## Breaking changes ## Breaking changes to experimental datasets @@ -11,6 +12,7 @@ - Added `os.PathLike` support for `redis.PickleDataset` keys. ## Community contributions +- [Saurav Gupta](https://github.com/Saurav-Gupta-9741) - [akira-in-tech](https://github.com/akira-in-tech) - [Tanmay Singh](https://github.com/tannnmayy) - [Shizoqua](https://github.com/Shizoqua) diff --git a/kedro-datasets/kedro_datasets/polars/__init__.py b/kedro-datasets/kedro_datasets/polars/__init__.py index c24b0ba2a..c7c465c17 100644 --- a/kedro-datasets/kedro_datasets/polars/__init__.py +++ b/kedro-datasets/kedro_datasets/polars/__init__.py @@ -18,6 +18,13 @@ # https://github.com/pylint-dev/pylint/issues/4300#issuecomment-1043601901 EagerPolarsDataset: Any +try: + from .iceberg_dataset import IcebergDataset +except (ImportError, RuntimeError): + # For documentation builds that might fail due to dependency issues + # https://github.com/pylint-dev/pylint/issues/4300#issuecomment-1043601901 + IcebergDataset: Any + try: from .lazy_polars_dataset import LazyPolarsDataset except (ImportError, RuntimeError): @@ -32,6 +39,7 @@ "eager_polars_dataset": [ "EagerPolarsDataset", ], + "iceberg_dataset": ["IcebergDataset"], "lazy_polars_dataset": ["LazyPolarsDataset"], }, ) diff --git a/kedro-datasets/kedro_datasets/polars/iceberg_dataset.py b/kedro-datasets/kedro_datasets/polars/iceberg_dataset.py new file mode 100644 index 000000000..099cae42e --- /dev/null +++ b/kedro-datasets/kedro_datasets/polars/iceberg_dataset.py @@ -0,0 +1,219 @@ +"""``IcebergDataset`` loads and saves data from/to Apache Iceberg tables. + +Loads via Polars ``scan_iceberg`` when available, with a PyIceberg scan fallback. +Saves data to Iceberg tables via PyIceberg catalog operations. Catalog access is handled through PyIceberg. +""" + +from __future__ import annotations + +from copy import deepcopy +from typing import TYPE_CHECKING, Any + +from kedro.io.core import AbstractDataset, DatasetError + +if TYPE_CHECKING: + import polars as pl + + +class IcebergDataset(AbstractDataset): + """``IcebergDataset`` loads and saves data from/to Apache Iceberg tables. + + Loads via Polars ``scan_iceberg`` when available, with a PyIceberg scan fallback. + Saves data to Iceberg tables via PyIceberg catalog operations. Catalog access is handled through PyIceberg. + + Examples: + Using the [YAML API](https://docs.kedro.org/en/stable/catalog-data/data_catalog_yaml_examples/): + + ```yaml + sales_iceberg_polars: + type: polars.IcebergDataset + table_name: analytics.sales + catalog_name: glue_catalog + catalog_properties: + type: glue + credentials: glue_credentials + load_args: + snapshot_id: 1234567890 + save_args: + mode: overwrite + ``` + + Using the [Python API](https://docs.kedro.org/en/stable/catalog-data/advanced_data_catalog_usage/): + + >>> from kedro_datasets.polars import IcebergDataset # doctest: +SKIP + >>> import polars as pl # doctest: +SKIP + >>> + >>> data = pl.DataFrame({"col1": [1, 2], "col2": [4, 5]}) # doctest: +SKIP + >>> dataset = IcebergDataset( # doctest: +SKIP + ... table_name="default.my_table", + ... catalog_properties={"type": "sql", "uri": "sqlite:///test.db"} + ... ) + >>> dataset.save(data) # doctest: +SKIP + >>> reloaded = dataset.load() # doctest: +SKIP + >>> assert data.equals(reloaded) # doctest: +SKIP + + """ + + DEFAULT_WRITE_MODE = "overwrite" + ACCEPTED_WRITE_MODES = ("overwrite", "append") + + DEFAULT_LOAD_ARGS: dict[str, Any] = {} + DEFAULT_SAVE_ARGS: dict[str, Any] = {"mode": DEFAULT_WRITE_MODE} + + def __init__( # noqa: PLR0913 + self, + *, + table_name: str, + catalog_name: str | None = None, + catalog_properties: dict[str, Any] | None = None, + credentials: dict[str, Any] | None = None, + load_args: dict[str, Any] | None = None, + save_args: dict[str, Any] | None = None, + metadata: dict[str, Any] | None = None, + ) -> None: + """Creates a new instance of ``IcebergDataset`` pointing to an Apache Iceberg table. + + Args: + table_name: Table identifier (e.g. ``"namespace.table_name"`` or ``"table_name"``). + catalog_name: Name of the Iceberg catalog to load. Defaults to None. + catalog_properties: Properties required to instantiate the catalog (e.g. + ``{"type": "rest", "uri": "https://..."}`` or ``{"type": "glue"}``). + Do not pass secrets/tokens here as they are visible in ``_describe()``; + use ``credentials`` instead. + credentials: Authentication credentials or secrets (e.g. tokens, AWS/GCP keys). + These are merged with ``catalog_properties`` when connecting to the catalog + and are safely excluded from ``_describe()`` to avoid leaking secrets. + load_args: Additional scan options passed to ``polars.scan_iceberg`` + (e.g. ``snapshot_id``, ``storage_options``, ``reader_override``, + ``use_metadata_statistics``, ``use_pyiceberg_filter``). + save_args: Additional save options. Supported keys: + - ``mode``: ``"overwrite"`` (default) or ``"append"``. + - ``location``: Explicit storage location for newly created tables. + - ``snapshot_properties``: Custom properties to attach to the snapshot summary. + metadata: Any arbitrary metadata to attach to the dataset. + + Raises: + DatasetError: If an invalid write mode or configuration is provided. + """ + self._table_name = table_name + self._catalog_name = catalog_name + self._catalog_properties = deepcopy(catalog_properties or {}) + self._credentials = deepcopy(credentials or {}) + self._load_args = {**self.DEFAULT_LOAD_ARGS, **(load_args or {})} + self._save_args = {**self.DEFAULT_SAVE_ARGS, **(save_args or {})} + self._metadata = metadata + + write_mode = self._save_args.get("mode", self.DEFAULT_WRITE_MODE) + if write_mode not in self.ACCEPTED_WRITE_MODES: + raise DatasetError( + f"Write mode '{write_mode}' is not supported. " + f"Please use one of {self.ACCEPTED_WRITE_MODES}." + ) + + def _get_catalog(self) -> Any: + try: + import pyiceberg.catalog as pyiceberg_catalog # noqa: PLC0415 + except ImportError as exc: + raise DatasetError( + "PyIceberg is required to use 'polars.IcebergDataset'. " + "Please install it using 'pip install pyiceberg'." + ) from exc + + properties = {**self._catalog_properties, **self._credentials} + return pyiceberg_catalog.load_catalog(self._catalog_name, **properties) + + def _load(self) -> pl.DataFrame: + """Loads data from the Iceberg table into a Polars DataFrame.""" + try: + import polars as pl # noqa: PLC0415 + except ImportError as exc: + raise DatasetError( + "Polars is required to load data using 'polars.IcebergDataset'. " + "Please install it using 'pip install polars'." + ) from exc + + catalog = self._get_catalog() + table = catalog.load_table(self._table_name) + + # Prefer native Polars Iceberg scan when available + if hasattr(pl, "scan_iceberg"): + try: + lazy_df = pl.scan_iceberg(table, **self._load_args) + return lazy_df.collect() + except (NotImplementedError, AttributeError): + # Fall back to PyIceberg scan if native scan is not supported for this table format/engine + pass + + # Fallback: PyIceberg scan -> to_polars or Arrow zero-copy + scan_args = {} + if "snapshot_id" in self._load_args: + scan_args["snapshot_id"] = self._load_args["snapshot_id"] + + scan = table.scan(**scan_args) + if hasattr(scan, "to_polars"): + return scan.to_polars() + + arrow_table = scan.to_arrow() + return pl.from_arrow(arrow_table) + + def _save(self, data: pl.DataFrame | pl.LazyFrame) -> None: + """Saves a Polars DataFrame into the Apache Iceberg table.""" + if hasattr(data, "collect"): + data = data.collect() + + save_args = deepcopy(self._save_args) + mode = save_args.pop("mode", self.DEFAULT_WRITE_MODE) + location = save_args.pop("location", None) + snapshot_properties = save_args.pop("snapshot_properties", None) + + catalog = self._get_catalog() + arrow_table = data.to_arrow() + + # We write via PyIceberg's Arrow interface (append / overwrite) directly + # rather than experimental/unstable polars write_iceberg/sink_iceberg, + # ensuring atomic commits, snapshot isolation, partition spec compliance, + # and schema enforcement across all catalog types (Glue, REST, SQL, etc.). + if not self._exists(): + table = catalog.create_table( + identifier=self._table_name, + schema=arrow_table.schema, + location=location, + **save_args, + ) + table.append(arrow_table, snapshot_properties=snapshot_properties or {}) + else: + table = catalog.load_table(self._table_name) + if mode == "append": + table.append(arrow_table, snapshot_properties=snapshot_properties or {}) + else: + table.overwrite( + arrow_table, snapshot_properties=snapshot_properties or {} + ) + + def _exists(self) -> bool: + """Checks if the Iceberg table exists in the configured catalog.""" + _not_found_errors: tuple[type[Exception], ...] = () + try: + from pyiceberg.exceptions import ( # noqa: PLC0415 + NoSuchNamespaceError, + NoSuchTableError, + ) + + _not_found_errors = (NoSuchTableError, NoSuchNamespaceError) + except ImportError: + pass + + try: + catalog = self._get_catalog() + return catalog.table_exists(self._table_name) + except _not_found_errors: + return False + + def _describe(self) -> dict[str, Any]: + return { + "table_name": self._table_name, + "catalog_name": self._catalog_name, + "catalog_properties": self._catalog_properties, + "load_args": self._load_args, + "save_args": self._save_args, + } diff --git a/kedro-datasets/pyproject.toml b/kedro-datasets/pyproject.toml index 7bb561bca..eadcfa6a8 100644 --- a/kedro-datasets/pyproject.toml +++ b/kedro-datasets/pyproject.toml @@ -185,10 +185,12 @@ plotly = ["kedro-datasets[plotly-htmldataset,plotly-jsondataset,plotly-plotlydat polars-csvdataset = ["kedro-datasets[polars-base]"] polars-eagerpolarsdataset = ["kedro-datasets[polars-base]", "pyarrow>=4.0", "xlsx2csv>=0.8.0", "deltalake>=0.6.2"] +polars-icebergdataset = ["kedro-datasets[polars-base]", "pyiceberg>=0.7.0; python_version < '3.14'", "pyarrow>=6.0"] polars-lazypolarsdataset = ["kedro-datasets[polars-base]", "pyarrow>=4.0", "deltalake>=0.6.2"] polars = [ """kedro-datasets[polars-csvdataset,\ polars-eagerpolarsdataset,\ + polars-icebergdataset,\ polars-lazypolarsdataset]""" ] @@ -326,6 +328,7 @@ test = [ "polars[deltalake,xlsx2csv]>=1.0", "pyarrow>=1.0; python_version < '3.11'", "pyarrow>=7.0; python_version >= '3.11'", # Adding to avoid numpy build errors + "pyiceberg>=0.7.0; python_version < '3.14'", "pyodbc~=5.0", "pyspark>=3.3, <4.0; python_version < '3.11'", "pyspark>=3.4, <4.0; python_version == '3.11'", @@ -481,6 +484,7 @@ omit = [ "kedro_datasets/databricks/*", "kedro_datasets/geopandas/*", # TODO: remove once fiona supports Python 3.14 "kedro_datasets/holoviews/*", + "kedro_datasets/polars/iceberg_dataset.py", # TODO: remove once pyiceberg supports Python 3.14 "kedro_datasets/snowflake/*", "kedro_datasets/spark/spark_hive_dataset.py", "kedro_datasets/tensorflow/*", diff --git a/kedro-datasets/tests/polars/test_iceberg_dataset.py b/kedro-datasets/tests/polars/test_iceberg_dataset.py new file mode 100644 index 000000000..9349ef450 --- /dev/null +++ b/kedro-datasets/tests/polars/test_iceberg_dataset.py @@ -0,0 +1,259 @@ +import sys +from unittest.mock import MagicMock, patch + +import pytest +from kedro.io.core import DatasetError + +from kedro_datasets.polars import IcebergDataset + +_skip_on_314 = pytest.mark.skipif( + sys.version_info >= (3, 14), + reason="PyIceberg does not support Python 3.14", +) + + +@pytest.fixture +def table_name(): + return "default.test_table" + + +@pytest.fixture +def catalog_properties(): + return {"type": "glue"} + + +@pytest.fixture +def dummy_polars_df(): + mock_df = MagicMock() + mock_df.to_arrow.return_value = MagicMock() + mock_lazy = MagicMock() + mock_lazy.collect.return_value = mock_df + mock_df.lazy.return_value = mock_lazy + return mock_df + + +@pytest.fixture +def iceberg_dataset(table_name, catalog_properties): + return IcebergDataset( + table_name=table_name, + catalog_name="glue_catalog", + catalog_properties=catalog_properties, + ) + + +@_skip_on_314 +class TestPolarsIcebergDataset: + def test_invalid_write_mode(self, table_name, catalog_properties): + """Test that initializing with an unsupported write mode raises DatasetError.""" + with pytest.raises(DatasetError, match="Write mode 'invalid' is not supported"): + IcebergDataset( + table_name=table_name, + catalog_properties=catalog_properties, + save_args={"mode": "invalid"}, + ) + + def test_describe(self, iceberg_dataset, table_name, catalog_properties): + """Test the _describe method output.""" + description = iceberg_dataset._describe() + assert description == { + "table_name": table_name, + "catalog_name": "glue_catalog", + "catalog_properties": catalog_properties, + "load_args": {}, + "save_args": {"mode": "overwrite"}, + } + + def test_describe_excludes_credentials(self, table_name, catalog_properties): + """Test that _describe excludes credentials for security.""" + dataset = IcebergDataset( + table_name=table_name, + catalog_name="glue_catalog", + catalog_properties=catalog_properties, + credentials={"user": "test-user", "role_arn": "arn:aws:iam::role/test"}, + ) + description = dataset._describe() + assert "credentials" not in description + assert "test-user" not in str(description) + assert description == { + "table_name": table_name, + "catalog_name": "glue_catalog", + "catalog_properties": catalog_properties, + "load_args": {}, + "save_args": {"mode": "overwrite"}, + } + + def test_credentials_passed_to_catalog(self, table_name, mocker): + """Test that credentials are merged into catalog properties when loading catalog.""" + mock_load_catalog = mocker.patch("pyiceberg.catalog.load_catalog") + dataset = IcebergDataset( + table_name=table_name, + catalog_name="rest_catalog", + catalog_properties={"type": "rest", "uri": "https://catalog.example.com"}, + credentials={"user": "test-user"}, + ) + dataset._get_catalog() + mock_load_catalog.assert_called_once_with( + "rest_catalog", + type="rest", + uri="https://catalog.example.com", + user="test-user", + ) + + def test_exists_true(self, iceberg_dataset, mocker): + """Test _exists when the table exists in the catalog.""" + mock_catalog = MagicMock() + mock_catalog.table_exists.return_value = True + mocker.patch.object(iceberg_dataset, "_get_catalog", return_value=mock_catalog) + + assert iceberg_dataset._exists() is True + mock_catalog.table_exists.assert_called_once_with(iceberg_dataset._table_name) + + def test_exists_false(self, iceberg_dataset, mocker): + """Test _exists returns False when table does not exist.""" + mock_catalog = MagicMock() + mock_catalog.table_exists.return_value = False + mocker.patch.object(iceberg_dataset, "_get_catalog", return_value=mock_catalog) + + assert iceberg_dataset._exists() is False + + def test_exists_surfaces_connection_error(self, iceberg_dataset, mocker): + """Test _exists surfaces real connection or authentication errors.""" + mocker.patch.object( + iceberg_dataset, + "_get_catalog", + side_effect=ConnectionError("Failed to reach catalog"), + ) + with pytest.raises(ConnectionError, match="Failed to reach catalog"): + iceberg_dataset._exists() + + def test_load_native_scan_iceberg(self, iceberg_dataset, dummy_polars_df, mocker): + """Test loading table data via native polars.scan_iceberg.""" + mock_catalog = MagicMock() + mock_table = MagicMock() + mock_catalog.load_table.return_value = mock_table + mocker.patch.object(iceberg_dataset, "_get_catalog", return_value=mock_catalog) + + mock_lazy = MagicMock() + mock_lazy.collect.return_value = dummy_polars_df + mocker.patch("polars.scan_iceberg", return_value=mock_lazy, create=True) + + loaded = iceberg_dataset.load() + assert loaded == dummy_polars_df + mock_catalog.load_table.assert_called_once_with(iceberg_dataset._table_name) + + def test_load_pyiceberg_scan_fallback( + self, iceberg_dataset, dummy_polars_df, mocker + ): + """Test loading table data falls back to PyIceberg scan when scan_iceberg is unavailable/fails.""" + mock_catalog = MagicMock() + mock_table = MagicMock() + mock_scan = MagicMock() + mock_scan.to_polars.return_value = dummy_polars_df + mock_table.scan.return_value = mock_scan + mock_catalog.load_table.return_value = mock_table + + mocker.patch.object(iceberg_dataset, "_get_catalog", return_value=mock_catalog) + mocker.patch( + "polars.scan_iceberg", side_effect=NotImplementedError, create=True + ) + + loaded = iceberg_dataset.load() + assert loaded == dummy_polars_df + mock_catalog.load_table.assert_called_once_with(iceberg_dataset._table_name) + mock_table.scan.assert_called_once() + + def test_load_invalid_load_args_raises_type_error( + self, table_name, catalog_properties, mocker + ): + """Test that invalid load_args raise TypeError and are not silently swallowed.""" + dataset = IcebergDataset( + table_name=table_name, + catalog_properties=catalog_properties, + load_args={"invalid_arg": 123}, + ) + mock_catalog = MagicMock() + mock_table = MagicMock() + mock_catalog.load_table.return_value = mock_table + mocker.patch.object(dataset, "_get_catalog", return_value=mock_catalog) + mocker.patch( + "polars.scan_iceberg", + side_effect=TypeError( + "scan_iceberg() got an unexpected keyword argument 'invalid_arg'" + ), + create=True, + ) + + with pytest.raises(TypeError, match="unexpected keyword argument"): + dataset._load() + + def test_missing_polars_raises_error(self, iceberg_dataset): + """Test that missing polars module raises DatasetError with install message.""" + with patch.dict("sys.modules", {"polars": None}): + with pytest.raises(DatasetError, match="Polars is required"): + iceberg_dataset.load() + + def test_save_new_table(self, iceberg_dataset, dummy_polars_df, mocker): + """Test saving data to a new table (table does not exist yet).""" + mock_catalog = MagicMock() + mock_catalog.table_exists.return_value = False + mock_table = MagicMock() + mock_catalog.create_table.return_value = mock_table + + mocker.patch.object(iceberg_dataset, "_get_catalog", return_value=mock_catalog) + + iceberg_dataset.save(dummy_polars_df) + mock_catalog.create_table.assert_called_once() + mock_table.append.assert_called_once() + + def test_save_overwrite_existing_table( + self, iceberg_dataset, dummy_polars_df, mocker + ): + """Test saving data with overwrite mode to an existing table.""" + mock_catalog = MagicMock() + mock_catalog.table_exists.return_value = True + mock_table = MagicMock() + mock_catalog.load_table.return_value = mock_table + + mocker.patch.object(iceberg_dataset, "_get_catalog", return_value=mock_catalog) + + iceberg_dataset.save(dummy_polars_df) + mock_catalog.load_table.assert_called_once_with(iceberg_dataset._table_name) + mock_table.overwrite.assert_called_once() + + def test_save_lazy_frame(self, iceberg_dataset, dummy_polars_df, mocker): + """Test saving a Polars LazyFrame collects before writing.""" + mock_catalog = MagicMock() + mock_catalog.table_exists.return_value = True + mock_table = MagicMock() + mock_catalog.load_table.return_value = mock_table + + mocker.patch.object(iceberg_dataset, "_get_catalog", return_value=mock_catalog) + + lazy_df = dummy_polars_df.lazy() + iceberg_dataset.save(lazy_df) + mock_table.overwrite.assert_called_once() + + def test_save_append_existing_table( + self, table_name, catalog_properties, dummy_polars_df, mocker + ): + """Test saving data with append mode to an existing table.""" + dataset = IcebergDataset( + table_name=table_name, + catalog_properties=catalog_properties, + save_args={"mode": "append"}, + ) + mock_catalog = MagicMock() + mock_catalog.table_exists.return_value = True + mock_table = MagicMock() + mock_catalog.load_table.return_value = mock_table + + mocker.patch.object(dataset, "_get_catalog", return_value=mock_catalog) + + dataset.save(dummy_polars_df) + mock_table.append.assert_called_once() + + def test_missing_pyiceberg_raises_error(self, iceberg_dataset): + """Test that missing pyiceberg module raises DatasetError with install message.""" + with patch.dict("sys.modules", {"pyiceberg": None, "pyiceberg.catalog": None}): + with pytest.raises(DatasetError, match="PyIceberg is required"): + iceberg_dataset._get_catalog()