feat: preserver dtypes across pandas and polars

Adopting narwhals, it's important that dtypes are preserved across
backends - a value that's NaN in pandas should be NaN in polars,
datetimes should remain datetimes, etc. .

This is achieved by
- using narwhals dtypes support as much as possible
- explicitly casting dtypes after transforming values when necessary

Author: philipp-jung

tab_err/_utils.py

@@ -1,6 +1,15 @@
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from narwhals.typing import IntoDType
+
 import random
+import warnings
+from typing import Any
 
 import narwhals as nw
 import numpy as np
@@ -60,7 +69,7 @@ def check_data_emptiness(data: nw.DataFrame) -> None:
 
 def is_string_dtype(series: nw.Series) -> bool:
     """Check if a series has a string dtype."""
-    return series.dtype == nw.String or series.dtype == nw.Object
+    return series.dtype in {nw.String, nw.Object}
 
 
 def is_numeric_dtype(series: nw.Series) -> bool:
@@ -78,22 +87,22 @@ def is_datetime_dtype(series: nw.Series) -> bool:
     return series.dtype == nw.Datetime
 
 
-def select_string_columns(data: nw.DataFrame) -> list[str]:
+def select_string_columns(data: nw.DataFrame) -> list[str | int]:
     """Select columns with string dtype."""
     return [col for col in data.columns if is_string_dtype(data[col])]
 
 
-def select_numeric_columns(data: nw.DataFrame) -> list[str]:
+def select_numeric_columns(data: nw.DataFrame) -> list[str | int]:
     """Select columns with numeric dtype."""
     return [col for col in data.columns if is_numeric_dtype(data[col])]
 
 
-def select_datetime_columns(data: nw.DataFrame) -> list[str]:
+def select_datetime_columns(data: nw.DataFrame) -> list[str | int]:
     """Select columns with datetime dtype."""
     return [col for col in data.columns if is_datetime_dtype(data[col])]
 
 
-def select_numeric_or_datetime_columns(data: nw.DataFrame) -> list[str]:
+def select_numeric_or_datetime_columns(data: nw.DataFrame) -> list[str | int]:
     """Select columns with numeric or datetime dtype."""
     return [col for col in data.columns if is_numeric_dtype(data[col]) or is_datetime_dtype(data[col])]
 
@@ -106,3 +115,32 @@ def create_empty_boolean_mask(data: nw.DataFrame) -> nw.DataFrame:
         dict.fromkeys(data.columns, mask_values),
         backend=nw.get_native_namespace(data),
     )
+
+
+def cast_series_like(series: nw.Series, like: nw.Series, column: int | str) -> nw.Series:
+    """Cast series to the dtype of 'like' when possible, otherwise keep original."""
+    if series.dtype == like.dtype:
+        return series
+    dtype: IntoDType = like.dtype
+
+    try:
+        return series.cast(dtype)
+    except Exception as exc:  # noqa: BLE001
+        msg = f"Failed to cast column {column} to dtype {like.dtype}: {exc}. Keeping inferred dtype."
+        warnings.warn(msg, stacklevel=2)
+        return series
+
+
+def _values_to_list(values: Sequence[Any] | np.ndarray) -> list[Any]:
+    """Normalize values into a list for nw.new_series."""
+    if isinstance(values, np.ndarray):
+        return values.tolist()
+    return list(values)
+
+
+def new_series_like(data: nw.DataFrame, column: int | str, values: Sequence[Any] | np.ndarray) -> nw.Series:
+    """Create a new series for 'column' and cast it back to the original dtype."""
+    col_name = get_column_str(data, column)
+    original = get_column(data, column)
+    series = nw.new_series(col_name, _values_to_list(values), backend=nw.get_native_namespace(data))
+    return cast_series_like(series, original, column)

tab_err/api/high_level.py

@@ -7,7 +7,7 @@
 
 from tab_err import ErrorMechanism, ErrorType, error_mechanism, error_type
 from tab_err._error_model import ErrorModel
-from tab_err._utils import check_data_emptiness, check_error_rate, create_empty_boolean_mask, seed_randomness_and_get_generator
+from tab_err._utils import check_data_emptiness, check_error_rate, seed_randomness_and_get_generator
 from tab_err.api import MidLevelConfig, mid_level
 
 if TYPE_CHECKING:
@@ -124,7 +124,7 @@ def _build_column_mechanism_dictionary(
         msg = "Possible conflict in error mechanisms to apply. Set at least on of: error_mechanisms_to_exclude or error_mechanisms_to_include to None."
         raise ValueError(msg)
 
-    columns_mechanisms = {}
+    columns_mechanisms: dict[int | str, list[ErrorMechanism]] = {}
 
     if error_mechanisms_to_include is not None and error_mechanisms_to_exclude is None:  # Include specified
         if not all(issubclass(type(cls), ErrorMechanism) for cls in error_mechanisms_to_include):  # Check input
@@ -177,7 +177,7 @@ def _build_column_number_of_models_dictionary(
     Returns:
         dict[int | str, int]: A dictionary mapping from column names to the number of error models to apply to that column.
     """
-    column_num_models = {}
+    column_num_models: dict[int | str, int] = {}
 
     for column in data.columns:
         column_num_models[column] = len(column_types[column]) * len(column_mechanisms[column])
@@ -202,7 +202,7 @@ def create_errors(  # noqa: PLR0913
     """Creates errors in a given DataFrame, at a rate of *approximately* max_error_rate.
 
     Args:
-        data (IntoDataFrame): The DataFrame to create errors in. Supports pandas, Polars, and other narwhals-compatible backends.
+        data (IntoDataFrame): The DataFrame to create errors in. Supports pandas, Polars, and (experimental) other narwhals-compatible backends.
         error_rate (float): The maximum error rate to be introduced to each column in the DataFrame.
         n_error_models_per_column (int, optional): The number of valid error models to apply to each column. Defaults to 1.
         error_types_to_include (list[ErrorType] | None, optional): A list of the error types to be included when building error models. Defaults to None.
@@ -231,7 +231,6 @@ def create_errors(  # noqa: PLR0913
 
     # Set Up Data
     data_copy = data_nw.clone()
-    error_mask = create_empty_boolean_mask(data_nw)
 
     # Build Dictionaries
     col_type = _build_column_type_dictionary(
@@ -272,5 +271,5 @@ def create_errors(  # noqa: PLR0913
         raise ValueError(msg)
 
     # Create Errors & Return (mid_level handles native conversion)
-    dirty_data, error_mask = mid_level.create_errors(nw.to_native(data_copy), config)
-    return dirty_data, error_mask
+    dirty_data_native, error_mask_native = mid_level.create_errors(nw.to_native(data_copy), config)
+    return dirty_data_native, error_mask_native

tab_err/error_type/_add_delta.py

@@ -1,11 +1,13 @@
 from __future__ import annotations
 
 import warnings
+from typing import TYPE_CHECKING
 
-import narwhals as nw
+if TYPE_CHECKING:
+    import narwhals as nw
 import numpy as np
 
-from tab_err._utils import get_column, get_column_str, is_datetime_dtype, is_numeric_dtype, select_numeric_or_datetime_columns
+from tab_err._utils import get_column, is_datetime_dtype, is_numeric_dtype, new_series_like, select_numeric_or_datetime_columns
 
 from ._error_type import ErrorType
 
@@ -39,7 +41,6 @@ def _apply(self: AddDelta, data: nw.DataFrame, error_mask: nw.DataFrame, column:
         Returns:
             nw.Series: The data column, 'column', after AddDelta errors at the locations specified by 'error_mask' are introduced.
         """
-        col_name = get_column_str(data, column)
         series = get_column(data, column)
         series_mask = get_column(error_mask, column)
         was_datetime = False
@@ -62,13 +63,14 @@ def _apply(self: AddDelta, data: nw.DataFrame, error_mask: nw.DataFrame, column:
             mean_val = np.nanmean(data_arr)
             std_val = np.nanstd(data_arr)
             random_choice = self._random_generator.choice(data_arr[~np.isnan(data_arr)])
-            self.config.add_delta_value = (random_choice - mean_val) / std_val if std_val != 0 else 0
+            delta_value = (random_choice - mean_val) / std_val if std_val != 0 else 0
+        else:
+            delta_value = self.config.add_delta_value
 
-        # Apply delta where mask is True
-        data_arr[mask_arr] += self.config.add_delta_value
+        data_arr[mask_arr] += delta_value
 
         if was_datetime:
             # Convert back to datetime (from seconds)
             data_arr = (data_arr * 10**9).astype("int64").astype("datetime64[ns]")
 
-        return nw.new_series(col_name, data_arr.tolist(), backend=nw.get_native_namespace(data))
+        return new_series_like(data, column, data_arr)

tab_err/error_type/_category_swap.py

@@ -3,9 +3,8 @@
 import random
 
 import narwhals as nw
-import numpy as np
 
-from tab_err._utils import get_column, get_column_str
+from tab_err._utils import get_column, new_series_like
 
 from ._error_type import ErrorType
 
@@ -40,7 +39,7 @@ def _check_type(data: nw.DataFrame, column: int | str) -> None:
 
     def _get_valid_columns(self: CategorySwap, data: nw.DataFrame) -> list[str | int]:
         """Checks which columns are categorical and returns the indices of those with two or more categories."""
-        valid_columns = []
+        valid_columns: list[str | int] = []
         for col_name in data.columns:
             series = get_column(data, col_name)
 
@@ -65,7 +64,6 @@ def _apply(self: CategorySwap, data: nw.DataFrame, error_mask: nw.DataFrame, col
         Returns:
             nw.Series: The data column, 'column', after CategorySwap errors at the locations specified by 'error_mask' are introduced.
         """
-        col_name = get_column_str(data, column)
         series = get_column(data, column)
         series_mask = get_column(error_mask, column)
 
@@ -84,16 +82,16 @@ def sample_label(old_label: str) -> str:
 
         elif self.config.mislabel_weighing == "frequency":
             # Calculate frequency weights
-            value_counts = {}
+            value_counts: dict[str, int] = {}
             for val in data_arr:
                 if val not in value_counts:
                     value_counts[val] = 0
                 value_counts[val] += 1
 
             def sample_label(old_label: str) -> str:
                 choices = [x for x in categories if x != old_label]
-                weights = [value_counts.get(x, 1) for x in choices]
-                total = sum(weights)
+                weights: list[float] = [float(value_counts.get(x, 1)) for x in choices]
+                total = float(sum(weights))
                 weights = [w / total for w in weights]
                 return random.choices(choices, weights=weights, k=1)[0]
         else:
@@ -105,4 +103,4 @@ def sample_label(old_label: str) -> str:
             if mask_arr[i]:
                 data_arr[i] = sample_label(data_arr[i])
 
-        return nw.new_series(col_name, data_arr.tolist(), backend=nw.get_native_namespace(data))
+        return new_series_like(data, column, data_arr)

tab_err/error_type/_error_type.py

@@ -3,13 +3,12 @@
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Any
 
-import narwhals as nw
-
 from tab_err._utils import seed_randomness_and_get_generator
 
 from ._config import ErrorTypeConfig
 
 if TYPE_CHECKING:
+    import narwhals as nw
     import numpy as np
 
 

tab_err/error_type/_extraneous.py

@@ -2,13 +2,14 @@
 
 import string
 import warnings
+from typing import TYPE_CHECKING
 
-import narwhals as nw
-
-from tab_err._utils import get_column, get_column_str, select_string_columns
+from tab_err._utils import get_column, new_series_like, select_string_columns
 
 from ._error_type import ErrorType
 
+if TYPE_CHECKING:
+    import narwhals as nw
 
 class Extraneous(ErrorType):
     """Adds Extraneous strings around the values in a column."""
@@ -45,7 +46,6 @@ def _apply(self: Extraneous, data: nw.DataFrame, error_mask: nw.DataFrame, colum
         Returns:
             nw.Series: The data column, 'column', after Extraneous errors at the locations specified by 'error_mask' are introduced.
         """
-        col_name = get_column_str(data, column)
         series = get_column(data, column)
         series_mask = get_column(error_mask, column)
 
@@ -69,4 +69,4 @@ def _apply(self: Extraneous, data: nw.DataFrame, error_mask: nw.DataFrame, colum
                 val = data_arr[i]
                 data_arr[i] = self.config.extraneous_value_template.format(value=val)
 
-        return nw.new_series(col_name, data_arr.tolist(), backend=nw.get_native_namespace(data))
+        return new_series_like(data, column, data_arr)

tab_err/error_type/_missing.py

@@ -1,12 +1,13 @@
 from __future__ import annotations
 
-import narwhals as nw
-import numpy as np
+from typing import TYPE_CHECKING, Union, cast
 
-from tab_err._utils import get_column, get_column_str, is_string_dtype, select_string_columns
+from tab_err._utils import get_column, is_string_dtype, new_series_like, select_string_columns
 
 from ._error_type import ErrorType
 
+if TYPE_CHECKING:
+    import narwhals as nw
 
 class MissingValue(ErrorType):
     """Insert missing values into a column.
@@ -22,7 +23,9 @@ def _check_type(data: nw.DataFrame, column: int | str) -> None:
 
     def _get_valid_columns(self: MissingValue, data: nw.DataFrame) -> list[str | int]:
         """If the config missing value is None, returns all columns. Otherwise, only the columns with string type."""
-        return data.columns if self.config.missing_value is None else select_string_columns(data)
+        if self.config.missing_value is None:
+            return cast("list[Union[str, int]]", list(data.columns))
+        return select_string_columns(data)
 
     def _apply(self: MissingValue, data: nw.DataFrame, error_mask: nw.DataFrame, column: int | str) -> nw.Series:
         """Applies the MissingValue ErrorType to a column of data.
@@ -35,7 +38,6 @@ def _apply(self: MissingValue, data: nw.DataFrame, error_mask: nw.DataFrame, col
         Returns:
             nw.Series: The data column, 'column', after MissingValue errors at the locations specified by 'error_mask' are introduced.
         """
-        col_name = get_column_str(data, column)
         series = get_column(data, column)
         series_mask = get_column(error_mask, column)
 
@@ -57,4 +59,4 @@ def _apply(self: MissingValue, data: nw.DataFrame, error_mask: nw.DataFrame, col
             data_arr[mask_arr] = missing_val
 
         # Create new series with the modified data
-        return nw.new_series(col_name, data_arr.tolist(), backend=nw.get_native_namespace(data))
+        return new_series_like(data, column, data_arr)

tab_err/error_type/_mistype.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import narwhals as nw
-import numpy as np
 
 from tab_err._utils import get_column, get_column_str
 
@@ -60,9 +59,7 @@ def _apply(self: Mistype, data: nw.DataFrame, error_mask: nw.DataFrame, column:
                 target_dtype = "object"
             elif current_dtype == nw.Int64:
                 target_dtype = "float64"
-            elif current_dtype == nw.Float64:
-                target_dtype = "int64"
-            elif current_dtype == nw.Boolean:
+            elif current_dtype in {nw.Float64, nw.Boolean}:
                 target_dtype = "int64"
             elif current_dtype.is_integer():
                 target_dtype = "float64"

tab_err/error_type/_mojibake.py

@@ -1,13 +1,14 @@
 from __future__ import annotations
 
 import random
+from typing import TYPE_CHECKING
 
-import narwhals as nw
-
-from tab_err._utils import get_column, get_column_str, is_string_dtype, select_string_columns
+from tab_err._utils import get_column, is_string_dtype, new_series_like, select_string_columns
 
 from ._error_type import ErrorType
 
+if TYPE_CHECKING:
+    import narwhals as nw
 
 class Mojibake(ErrorType):
     """Inserts mojibake into a column containing strings."""
@@ -53,7 +54,6 @@ def _apply(self: Mojibake, data: nw.DataFrame, error_mask: nw.DataFrame, column:
             "iso-8859-2": top10 - {"iso-8859-2", "windows-1250", "iso-8859-1", "windows-1252"},
         }
 
-        col_name = get_column_str(data, column)
         series = get_column(data, column)
         series_mask = get_column(error_mask, column)
 
@@ -77,4 +77,4 @@ def _apply(self: Mojibake, data: nw.DataFrame, error_mask: nw.DataFrame, column:
                 if val is not None and isinstance(val, str):
                     data_arr[i] = val.encode(encoding_sender, errors="ignore").decode(encoding_receiver, errors="ignore")
 
-        return nw.new_series(col_name, data_arr.tolist(), backend=nw.get_native_namespace(data))
+        return new_series_like(data, column, data_arr)