fix: Allow to pass custom metadata to failure info parquets (#83)

Author: borchero

dataframely/failure.py

@@ -97,7 +97,7 @@ def write_parquet(self, file: str | Path | IO[bytes], **kwargs: Any) -> None:
             Be aware that this method suffers from the same limitations as
             :meth:`Schema.serialize`.
         """
-        metadata = self._build_metadata(**kwargs)
+        metadata, kwargs = self._build_metadata(**kwargs)
         self._df.write_parquet(file, metadata=metadata, **kwargs)
 
     def sink_parquet(
@@ -117,14 +117,16 @@ def sink_parquet(
             Be aware that this method suffers from the same limitations as
             :meth:`Schema.serialize`.
         """
-        metadata = self._build_metadata(**kwargs)
+        metadata, kwargs = self._build_metadata(**kwargs)
         self._lf.sink_parquet(file, metadata=metadata, **kwargs)
 
-    def _build_metadata(self, **kwargs: Any) -> dict[str, Any]:
+    def _build_metadata(
+        self, **kwargs: dict[str, Any]
+    ) -> tuple[dict[str, Any], dict[str, Any]]:
         metadata = kwargs.pop("metadata", {})
         metadata[RULE_METADATA_KEY] = json.dumps(self._rule_columns)
         metadata[SCHEMA_METADATA_KEY] = self.schema.serialize()
-        return metadata
+        return metadata, kwargs
 
     @classmethod
     def read_parquet(

tests/test_failure_info.py

@@ -52,6 +52,18 @@ def test_scan_sink_parquet(tmp_path: Path) -> None:
     assert MySchema.matches(read.schema)
 
 
+def test_write_parquet_custom_metadata(tmp_path: Path) -> None:
+    df = pl.DataFrame(
+        {
+            "a": [4, 5, 6, 6, 7, 8],
+            "b": [1, 2, 3, 4, 5, 6],
+        }
+    )
+    _, failure = MySchema.filter(df)
+    failure.write_parquet(tmp_path / "failure.parquet", metadata={"custom": "test"})
+    assert pl.read_parquet_metadata(tmp_path / "failure.parquet")["custom"] == "test"
+
+
 @pytest.mark.parametrize(
     "read_fn",
     [dy.FailureInfo.read_parquet, dy.FailureInfo.scan_parquet],