Enhance SimpleImputer and SklearnWrapper

DashAI/back/converters/scikit_learn/simple_imputer.py

@@ -1,3 +1,5 @@
+from typing import TYPE_CHECKING, Union
+
 from sklearn.impute import SimpleImputer as SimpleImputerOperation
 
 from DashAI.back.converters.category.basic_preprocessing import (
@@ -20,6 +22,9 @@
 from DashAI.back.types.dashai_data_type import DashAIDataType
 from DashAI.back.types.value_types import Float, Integer
 
+if TYPE_CHECKING:
+    from DashAI.back.dataloaders.classes.dashai_dataset import DashAIDataset
+
 
 class SimpleImputerSchema(BaseSchema):
     """Schema for configuring the SimpleImputer converter.
@@ -103,9 +108,15 @@ class SimpleImputer(
 
     Columns with all-missing values are handled according to the
     ``keep_empty_features`` flag. When ``add_indicator=True``, a
-    ``MissingIndicator`` binary matrix is stacked onto the output. All
-    output columns are typed as ``Float64`` in DashAI regardless of the
-    original column type.
+    ``MissingIndicator`` binary matrix is stacked onto the output.
+
+    Output typing preserves the original column type whenever the strategy
+    does not force a fractional result: ``"most_frequent"`` and
+    ``"constant"`` never perform arithmetic, so the source type (Integer,
+    Float, or Categorical) is kept. For ``"mean"``/``"median"`` the computed
+    per-column statistic is inspected, and an originally-Integer column
+    stays ``Integer`` if that statistic happens to be a whole number (e.g. a
+    median over an odd count of integers); otherwise it becomes ``Float64``.
 
     Wraps ``sklearn.impute.SimpleImputer``.
 
@@ -170,20 +181,73 @@ def __init__(self, **kwargs):
         """
         super().__init__(**kwargs)
 
+    def fit(
+        self, x: "DashAIDataset", y: Union["DashAIDataset", None] = None
+    ) -> "SimpleImputer":
+        """Fit the imputer, remembering input types and column order.
+
+        These are needed by ``get_output_type`` to preserve the original
+        column type instead of always coercing to ``Float64``.
+
+        Parameters
+        ----------
+        x : DashAIDataset
+            The input dataset to fit the imputer on.
+        y : DashAIDataset, optional
+            Ignored; present for API consistency.
+
+        Returns
+        -------
+        SimpleImputer
+            The fitted imputer instance (self).
+        """
+        if hasattr(x, "types") and x.types is not None:
+            self._input_types = dict(x.types)
+        self._input_columns = {name: idx for idx, name in enumerate(x.column_names)}
+        return super().fit(x, y)
+
     def get_output_type(self, column_name: str = None) -> DashAIDataType:
         """Return the DashAI data type produced by this converter for a column.
 
         Parameters
         ----------
         column_name : str, optional
-            Not used; all output columns share the
-            same type. Defaults to None.
+            The name of the output column. Defaults to None.
 
         Returns
         -------
         DashAIDataType
-            A Float type backed by ``pyarrow.float64()``.
+            ``Integer`` for the binary ``MissingIndicator`` columns appended
+            when ``add_indicator=True``. Otherwise, the original column type
+            for ``"most_frequent"``/``"constant"`` (no arithmetic is performed
+            on the values), or for ``"mean"``/``"median"``, ``Integer`` if the
+            source column was an Integer and the computed statistic is a
+            whole number — otherwise a Float type backed by
+            ``pyarrow.float64()``.
         """
         import pyarrow as pa
 
+        if column_name and str(column_name).startswith("missingindicator_"):
+            return Integer(arrow_type=pa.int64())
+
+        input_types = getattr(self, "_input_types", None)
+        input_type = input_types.get(column_name) if input_types else None
+
+        if self.strategy in ("most_frequent", "constant"):
+            if input_type is not None:
+                return input_type
+            return Float(arrow_type=pa.float64())
+
+        if isinstance(input_type, Integer):
+            columns = getattr(self, "_input_columns", None)
+            statistics = getattr(self, "statistics_", None)
+            if (
+                columns is not None
+                and statistics is not None
+                and column_name in columns
+            ):
+                value = statistics[columns[column_name]]
+                if float(value).is_integer():
+                    return Integer(arrow_type=pa.int64())
+
         return Float(arrow_type=pa.float64())

DashAI/back/converters/sklearn_wrapper.py

@@ -10,6 +10,39 @@
     from DashAI.back.dataloaders.classes.dashai_dataset import DashAIDataset
 
 
+def _normalize_missing(df: object) -> object:
+    """Normalize missing-value sentinels to ``np.nan`` before handing off to sklearn.
+
+    ``DashAIDataset.to_pandas()`` converts Arrow nulls in string/Categorical
+    columns to Python ``None``, while numeric columns get actual float
+    ``nan``. Several sklearn transformers (notably ``SimpleImputer``, whose
+    default ``missing_values=np.nan`` masks via a ``value != value`` check)
+    only recognize float ``nan`` as missing: ``None != None`` is ``False``,
+    so ``None`` entries are silently treated as real values and never
+    imputed, even though ``pandas``/Arrow both consider them missing.
+    ``DataFrame.notna()`` correctly flags both, so this unifies the
+    representation without altering any non-missing value or dtype.
+
+    Parameters
+    ----------
+    df : object
+        The value returned by ``DashAIDataset.to_pandas()``. Only
+        ``pandas.DataFrame`` instances are normalized; anything else is
+        returned unchanged.
+
+    Returns
+    -------
+    object
+        The normalized DataFrame, or ``df`` unchanged if it isn't one.
+    """
+    import numpy as np
+    import pandas as pd
+
+    if not isinstance(df, pd.DataFrame):
+        return df
+    return df.where(df.notna(), np.nan)
+
+
 class SklearnWrapper(BaseConverter, metaclass=ABCMeta):
     """Abstract mixin that adapts scikit-learn transformers to the DashAI converter API.
 
@@ -85,49 +118,60 @@ def fit(
             If no scikit-learn class with a `fit` method is found
             in the MRO.
         """
-        x_pandas = x.to_pandas() if hasattr(x, "to_pandas") else x
-        y_pandas = y.to_pandas() if y is not None and hasattr(y, "to_pandas") else y
-
-        # Detect whether the underlying sklearn estimator needs a target.
-        # sklearn >= 1.6 moved tags from ``_get_tags`` to ``__sklearn_tags__``,
-        # so we consult both for backward/forward compatibility.
-        requires_y = False
-        if hasattr(self, "__sklearn_tags__"):
-            try:
-                tags = self.__sklearn_tags__()
-                target_tags = getattr(tags, "target_tags", None)
-                if target_tags is not None:
-                    requires_y = bool(getattr(target_tags, "required", False))
-            except Exception:
-                requires_y = False
-        if not requires_y and hasattr(self, "_get_tags"):
-            with contextlib.suppress(Exception):
-                requires_y = bool(self._get_tags().get("requires_y", False))
-
-        if requires_y and y is None:
-            raise ValueError("This transformer requires y for fitting")
-
-        sklearn_cls = next(
-            (
-                cls
-                for cls in type(self).__mro__
-                if "sklearn" in cls.__module__
-                and "DashAI" not in cls.__module__
-                and "fit" in cls.__dict__
-            ),
-            None,
-        )
+        try:
+            if hasattr(x, "to_pandas"):
+                x_pandas = _normalize_missing(x.to_pandas())
+                self._fit_input_cache = (x, x_pandas)
+            else:
+                x_pandas = x
+                self._fit_input_cache = None
+            y_pandas = (
+                _normalize_missing(y.to_pandas())
+                if y is not None and hasattr(y, "to_pandas")
+                else y
+            )
 
-        if sklearn_cls is None:
-            raise RuntimeError(
-                "No sklearn class with a 'fit' method found in the MRO. "
-                "Ensure that your transformer inherits from a valid sklearn class."
+            requires_y = False
+            if hasattr(self, "__sklearn_tags__"):
+                try:
+                    tags = self.__sklearn_tags__()
+                    target_tags = getattr(tags, "target_tags", None)
+                    if target_tags is not None:
+                        requires_y = bool(getattr(target_tags, "required", False))
+                except Exception:
+                    requires_y = False
+            if not requires_y and hasattr(self, "_get_tags"):
+                with contextlib.suppress(Exception):
+                    requires_y = bool(self._get_tags().get("requires_y", False))
+
+            if requires_y and y is None:
+                raise ValueError("This transformer requires y for fitting")
+
+            sklearn_cls = next(
+                (
+                    cls
+                    for cls in type(self).__mro__
+                    if "sklearn" in cls.__module__
+                    and "DashAI" not in cls.__module__
+                    and "fit" in cls.__dict__
+                ),
+                None,
             )
-        fit_method = sklearn_cls.__dict__["fit"]
-        if requires_y or y_pandas is not None:
-            fit_method(self, x_pandas, y_pandas)
-        else:
-            fit_method(self, x_pandas)
+
+            if sklearn_cls is None:
+                raise RuntimeError(
+                    "No sklearn class with a 'fit' method found in the MRO. "
+                    "Ensure that your transformer inherits from a valid sklearn "
+                    "class."
+                )
+            fit_method = sklearn_cls.__dict__["fit"]
+            if requires_y or y_pandas is not None:
+                fit_method(self, x_pandas, y_pandas)
+            else:
+                fit_method(self, x_pandas)
+        except Exception:
+            self._fit_input_cache = None
+            raise
 
         return self
 
@@ -165,7 +209,16 @@ def transform(
 
         from DashAI.back.dataloaders.classes.dashai_dataset import to_dashai_dataset
 
-        x_pandas = x.to_pandas() if hasattr(x, "to_pandas") else x
+        cached = getattr(self, "_fit_input_cache", None)
+        if cached is not None and cached[0] is x:
+            x_pandas = cached[1]
+        elif hasattr(x, "to_pandas"):
+            x_pandas = _normalize_missing(x.to_pandas())
+        else:
+            x_pandas = x
+        # Drop the reference now that it's been consumed (or wasn't a hit),
+        # so the cached DataFrame doesn't outlive this transform call.
+        self._fit_input_cache = None
 
         sklearn_cls = next(
             (

DashAI/back/job/converter_job.py

@@ -57,14 +57,23 @@ def _rebuild_dataset_with_transformed_columns(
     original_columns = base.column_names
     transformed_cols = transformed.column_names
 
-    removed_cols = [col for col in scope_column_names if col not in transformed_cols]
-    replacement_cols = [col for col in scope_column_names if col in transformed_cols]
-    new_cols = [col for col in transformed_cols if col not in scope_column_names]
+    transformed_cols_set = set(transformed_cols)
+    scope_column_names_set = set(scope_column_names)
+
+    removed_cols = [
+        col for col in scope_column_names if col not in transformed_cols_set
+    ]
+    replacement_cols = [
+        col for col in scope_column_names if col in transformed_cols_set
+    ]
+    new_cols = [col for col in transformed_cols if col not in scope_column_names_set]
+
+    removed_cols_set = set(removed_cols)
 
     new_columns_order = []
     seen_cols = set()
     for col in original_columns:
-        if col in removed_cols:
+        if col in removed_cols_set:
             continue
         if col not in seen_cols:
             new_columns_order.append(col)
@@ -83,10 +92,10 @@ def _rebuild_dataset_with_transformed_columns(
 
     updated_arrays = {}
     for col in replacement_cols:
-        if col in transformed.arrow_table.column_names:
+        if col in transformed_cols_set:
             updated_arrays[col] = transformed.arrow_table[col]
     for col, unique_col in col_name_mapping.items():
-        if col in transformed.arrow_table.column_names:
+        if col in transformed_cols_set:
             updated_arrays[unique_col] = transformed.arrow_table[col]
 
     updated_types = base.types.copy()
@@ -345,10 +354,11 @@ def instantiate_converters(
                         )
                         if scope_rows_indexes:
                             y_dataset_fit = y_dataset_fit.select(scope_rows_indexes)
-
-                        y_full_transform = loaded_dataset.select_columns(
-                            [target_column_name]
-                        )
+                            y_full_transform = loaded_dataset.select_columns(
+                                [target_column_name]
+                            )
+                        else:
+                            y_full_transform = y_dataset_fit
 
                     X_dataset_fit = loaded_dataset.select_columns(scope_column_names)
 
@@ -370,7 +380,14 @@ def instantiate_converters(
                             f"Error fitting converter {converter_name}: {e}"
                         ) from e
 
-                    X_full_transform = loaded_dataset.select_columns(scope_column_names)
+                    if scope_rows_indexes:
+                        X_full_transform = loaded_dataset.select_columns(
+                            scope_column_names
+                        )
+                    else:
+                        # Same reuse as above: no row-level fit scope means
+                        # X_dataset_fit already covers the full transform scope.
+                        X_full_transform = X_dataset_fit
 
                     try:
                         transformed_dataset = converter_instance.transform(