Keep Pandas datatypes when using pyarrow dtypes (#459)

Author: mallport

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "dapla-toolbelt-pseudo"
-version = "4.3.3"
+version = "4.4.0"
 description = "Pseudonymization extensions for Dapla"
 authors = ["Dapla Developers <dapla-platform-developers@ssb.no>"]
 license = "MIT"

src/dapla_pseudo/v1/baseclasses.py

@@ -10,6 +10,7 @@
 from concurrent.futures import ThreadPoolExecutor
 from datetime import date
 
+import pandas as pd
 import polars as pl
 from dapla_metadata.datasets.core import Datadoc
 
@@ -65,6 +66,7 @@ def _execute_pseudo_operation(
         custom_keyset: PseudoKeyset | str | None = None,
         target_custom_keyset: PseudoKeyset | str | None = None,  # used in repseudo
         target_rules: list[PseudoRule] | None = None,  # used in repseudo
+        schema: pd.Series | pl.Schema | None = None,
     ) -> Result:
         if self._dataset is None:
             raise ValueError("No dataset has been provided.")
@@ -92,6 +94,7 @@ def _execute_pseudo_operation(
                 for pseudo_rule in (target_rules if target_rules else rules)
             ],
             user_provided_metadata=self._user_provided_metadata,
+            schema=schema,
         )
 
     def _pseudonymize_field(

src/dapla_pseudo/v1/depseudo.py

@@ -26,17 +26,20 @@ class Depseudonymize:
     """
 
     dataset: pl.DataFrame
+    schema: pd.Series | pl.Schema
 
     @staticmethod
     def from_pandas(dataframe: pd.DataFrame) -> "Depseudonymize._Depseudonymizer":
         """Initialize a depseudonymization request from a pandas DataFrame."""
         Depseudonymize.dataset = pl.from_pandas(dataframe)
+        Depseudonymize.schema = dataframe.dtypes
         return Depseudonymize._Depseudonymizer()
 
     @staticmethod
     def from_polars(dataframe: pl.DataFrame) -> "Depseudonymize._Depseudonymizer":
         """Initialize a depseudonymization request from a polars DataFrame."""
         Depseudonymize.dataset = dataframe
+        Depseudonymize.schema = dataframe.schema
         return Depseudonymize._Depseudonymizer()
 
     class _Depseudonymizer(_BasePseudonymizer):
@@ -92,7 +95,7 @@ def run(
             )
 
             result = super()._execute_pseudo_operation(
-                self.rules, timeout, custom_keyset
+                self.rules, timeout, custom_keyset, schema=Depseudonymize.schema
             )
             return result
 

src/dapla_pseudo/v1/pseudo.py

@@ -26,6 +26,7 @@ class Pseudonymize:
     """
 
     dataset: pl.DataFrame
+    schema: pd.Series | pl.Schema
 
     @staticmethod
     def from_pandas(dataframe: pd.DataFrame) -> "Pseudonymize._Pseudonymizer":
@@ -38,6 +39,7 @@ def from_pandas(dataframe: pd.DataFrame) -> "Pseudonymize._Pseudonymizer":
             _Pseudonymizer: An instance of the _Pseudonymizer class.
         """
         Pseudonymize.dataset = pl.from_pandas(dataframe)
+        Pseudonymize.schema = dataframe.dtypes
         return Pseudonymize._Pseudonymizer()
 
     @staticmethod
@@ -51,6 +53,7 @@ def from_polars(dataframe: pl.DataFrame) -> "Pseudonymize._Pseudonymizer":
             _Pseudonymizer: An instance of the _Pseudonymizer class.
         """
         Pseudonymize.dataset = dataframe
+        Pseudonymize.schema = dataframe.schema
         return Pseudonymize._Pseudonymizer()
 
     class _Pseudonymizer(_BasePseudonymizer):
@@ -60,7 +63,9 @@ class _Pseudonymizer(_BasePseudonymizer):
         metadata: Datadoc | None = None
 
         def __init__(
-            self, rules: list[PseudoRule] | None = None, metadata: Datadoc | None = None
+            self,
+            rules: list[PseudoRule] | None = None,
+            metadata: Datadoc | None = None,
         ) -> None:
             """Initialize the class."""
             if rules is None:
@@ -109,14 +114,17 @@ def run(
                 hierarchical=hierarchical,
                 user_provided_metadata=self.metadata,
             )
-
             result = super()._execute_pseudo_operation(
-                self.rules, timeout, custom_keyset
+                self.rules, timeout, custom_keyset, schema=Pseudonymize.schema
             )
             return result
 
     class _PseudoFuncSelector(_BaseRuleConstructor):
-        def __init__(self, fields: list[str], metadata: Datadoc | None) -> None:
+        def __init__(
+            self,
+            fields: list[str],
+            metadata: Datadoc | None,
+        ) -> None:
             self._fields = fields
             self._metadata = metadata
             super().__init__(fields)

src/dapla_pseudo/v1/repseudo.py

@@ -25,17 +25,20 @@ class Repseudonymize:
     """
 
     dataset: pl.DataFrame
+    schema: pd.Series | pl.Schema
 
     @staticmethod
     def from_pandas(dataframe: pd.DataFrame) -> "Repseudonymize._Repseudonymizer":
         """Initialize a pseudonymization request from a pandas DataFrame."""
         Repseudonymize.dataset = pl.from_pandas(dataframe)
+        Repseudonymize.schema = dataframe.dtypes
         return Repseudonymize._Repseudonymizer()
 
     @staticmethod
     def from_polars(dataframe: pl.DataFrame) -> "Repseudonymize._Repseudonymizer":
         """Initialize a pseudonymization request from a polars DataFrame."""
         Repseudonymize.dataset = dataframe
+        Repseudonymize.schema = dataframe.schema
         return Repseudonymize._Repseudonymizer()
 
     class _Repseudonymizer(_BasePseudonymizer):
@@ -110,6 +113,7 @@ def run(
                 custom_keyset=source_custom_keyset,
                 target_custom_keyset=target_custom_keyset,
                 timeout=timeout,
+                schema=Repseudonymize.schema,
             )
             return result
 

src/dapla_pseudo/v1/result.py

@@ -30,11 +30,13 @@ def __init__(
         pseudo_operation: PseudoOperation | None = None,
         targeted_columns: list[str] | None = None,
         user_provided_metadata: Datadoc | None = None,
+        schema: pd.Series | pl.Schema | None = None,
     ) -> None:
         """Initialise a PseudonymizationResult."""
         self._pseudo_data: pl.DataFrame = pseudo_response.data
         self._metadata: dict[str, dict[str, list[Any]]] = {}
         self._datadoc: Datadoc | MetadataContainer
+        self._schema = schema
 
         datadoc_fields: list[Variable] = []
         datadoc_paths: list[str | None] = []
@@ -144,7 +146,13 @@ def to_pandas(self, **kwargs: Any) -> pd.DataFrame:
         """
         match self._pseudo_data:
             case pl.DataFrame() as df:
-                return df.to_pandas()
+                pandas_df = df.to_pandas()
+                if isinstance(
+                    self._schema, pd.Series
+                ):  # Apply original schema if available
+                    pandas_df = pandas_df.astype(self._schema)
+
+                return pandas_df
             case _ as invalid_pseudo_data:
                 raise ValueError(f"Invalid response type: {type(invalid_pseudo_data)}")
 

tests/conftest.py

@@ -278,6 +278,22 @@ def df_personer_sid_fnr() -> pl.DataFrame:
     )
 
 
+@pytest_cases.fixture()
+def pandas_diverse_datatypes() -> pd.DataFrame:
+    JSON_FILE = "tests/data/diverse_datatypes.json"
+    return pd.read_json(
+        JSON_FILE,
+        dtype={
+            "string_field": "string[pyarrow]",
+            "int_field": "Int64[pyarrow]",
+            "float_field": "Float64[pyarrow]",
+            "date_pseudonymized": "datetime64[s]",
+            "bool_field": "boolean[pyarrow]",
+        },
+        dtype_backend="pyarrow",
+    )
+
+
 @pytest_cases.fixture()
 def single_field_response() -> MagicMock:
     mock_response = MagicMock()

tests/data/diverse_datatypes.json

@@ -0,0 +1,16 @@
+[
+    {
+        "fnr": "11854898347",
+        "number": 1,
+        "floater": 1.0,
+        "date_pseudonymized": "2023-10-27T10:30:00",
+        "boolean": true
+    },
+    {
+        "fnr": "11852228347",
+        "number": 10,
+        "floater": -1.0,
+        "date_pseudonymized": "2013-10-27T10:30:00",
+        "boolean": false
+    }
+]

tests/v1/integration/test_result.py

@@ -138,3 +138,20 @@ def test_pseudonymize_input_output_funcs(
         case "polars":
             df_polars = result.to_polars()
             pl_assert_frame_equal(df_polars, df_personer_fnr_daead_encrypted)
+
+
+@pytest.mark.usefixtures("setup")
+@integration_test()
+def test_pseudonymize_with_arrow_dtypes(
+    pandas_diverse_datatypes: pd.DataFrame,
+) -> None:
+    """This test ensures that datatypes are retained when converting internally to and from Pandas."""
+    result = (
+        Pseudonymize.from_pandas(pandas_diverse_datatypes)
+        .on_fields("fnr")
+        .with_default_encryption()
+        .run()
+    )
+
+    df_result = result.to_pandas()
+    assert df_result.dtypes.equals(pandas_diverse_datatypes.dtypes)