API: preprocessing.imputation

skyulf.preprocessing.imputation

Imputation nodes package.

Split from a single 416-LOC module into per-imputer files

_common.py — shared helpers (column resolution, polars fill values, sklearn transform) simple.py — SimpleImputer knn.py — KNNImputer iterative.py — IterativeImputer (MICE)

All public names are re-exported here so existing imports such as from skyulf.preprocessing.imputation import SimpleImputerCalculator continue to work unchanged.

SimpleImputerApplier

Bases: BaseApplier

Apply fitted Simple Imputer fill values to missing values in selected columns.

The calculator artifact records per-column values for mean, median, most_frequent (also accepted as mode), or constant strategies. Missing columns seen during fitting are restored with their stored value.

Source code in skyulf-core/skyulf/preprocessing/imputation/simple.py

class SimpleImputerApplier(BaseApplier):
    """Apply fitted Simple Imputer fill values to missing values in selected columns.

    The calculator artifact records per-column values for ``mean``, ``median``,
    ``most_frequent`` (also accepted as ``mode``), or ``constant`` strategies.
    Missing columns seen during fitting are restored with their stored value.
    """

    @apply_method
    def apply(self, X: Any, _y: Any, params: dict[str, Any]) -> Any:  # pylint: disable=arguments-differ
        return apply_dual_engine(X, params, self._apply_polars, self._apply_pandas)

    @staticmethod
    def _apply_polars(X: Any, _y: Any, params: dict[str, Any]) -> tuple[Any, Any]:
        cols = params.get("columns", [])
        fill_values = params.get("fill_values", {})
        if not cols:
            return X, _y

        exprs: list[Any] = []
        for col in X.columns:
            if col in cols and col in fill_values:
                expr = pl.col(col).fill_null(fill_values[col])
                if X.schema[col].is_float():
                    expr = expr.fill_nan(fill_values[col])
                exprs.append(expr.alias(col))
            else:
                exprs.append(pl.col(col))

        # Restore columns that were present at fit time but missing in input X.
        exprs.extend(
            pl.lit(fill_values[col]).alias(col)
            for col in cols
            if col not in X.columns and col in fill_values
        )

        return X.select(exprs), _y

    @staticmethod
    def _apply_pandas(X: Any, _y: Any, params: dict[str, Any]) -> tuple[Any, Any]:
        cols = params.get("columns", [])
        fill_values = params.get("fill_values", {})
        if not cols:
            return X, _y

        X_out = X.copy()
        for col in cols:
            val = fill_values.get(col)
            if val is None:
                continue
            if col not in X_out.columns:
                X_out[col] = val
            else:
                series = X_out[col]
                # Nullable numeric extension dtypes (Int64...) refuse a float
                # fill value; upcast like the Polars fill does (F-10).
                if (
                    isinstance(val, float)
                    and isinstance(series.dtype, pd.api.extensions.ExtensionDtype)
                    and pd.api.types.is_numeric_dtype(series.dtype)
                ):
                    series = series.astype("float64")
                X_out[col] = series.fillna(val)
        return X_out, _y

SimpleImputerCalculator

Bases: BaseCalculator

Fit per-column values for filling missing data with sklearn-compatible strategies.

Supported strategy values are mean, median, most_frequent, and constant; mode is normalized to most_frequent. Use columns to select columns and fill_value with constant. Mean and median operate on numeric columns, while the other strategies can operate on all selected columns.

Source code in skyulf-core/skyulf/preprocessing/imputation/simple.py

@NodeRegistry.register("SimpleImputer", SimpleImputerApplier)
@node_meta(
    id="SimpleImputer",
    name="Simple Imputer",
    category="Preprocessing",
    description="Imputes missing values using mean, median, or constant.",
    params={"strategy": "mean", "fill_value": None, "columns": []},
    learns_from_data=True,
)
class SimpleImputerCalculator(BaseCalculator):
    """Fit per-column values for filling missing data with sklearn-compatible strategies.

    Supported ``strategy`` values are ``mean``, ``median``, ``most_frequent``,
    and ``constant``; ``mode`` is normalized to ``most_frequent``. Use
    ``columns`` to select columns and ``fill_value`` with ``constant``.
    Mean and median operate on numeric columns, while the other strategies can
    operate on all selected columns.
    """

    def infer_output_schema(
        self, input_schema: SkyulfSchema, config: dict[str, Any]
    ) -> SkyulfSchema:
        # Imputers fill NaNs in place; column set and order are preserved.
        return input_schema

    @fit_method
    def fit(self, X: Any, _y: Any, config: dict[str, Any]) -> SimpleImputerArtifact:  # pylint: disable=arguments-differ
        if user_picked_no_columns(config):
            return {}

        strategy = config.get("strategy", "mean")
        if strategy == "mode":
            strategy = "most_frequent"
        fill_value = config.get("fill_value")

        cols = _resolve_simple_columns(X, config, strategy)
        if not cols:
            return {}

        # Stash resolved-once values into params so dispatched fits don't redo work.
        merged = dict(config)
        merged["_resolved_strategy"] = strategy
        merged["_resolved_cols"] = cols
        merged["_resolved_fill_value"] = fill_value

        return cast(
            SimpleImputerArtifact,
            fit_dual_engine(X, merged, self._fit_polars, self._fit_pandas),
        )

    @staticmethod
    def _fit_polars(X: Any, _y: Any, params: dict[str, Any]) -> dict[str, Any]:
        cols: list[str] = params["_resolved_cols"]
        strategy: str = params["_resolved_strategy"]
        fill_value = params["_resolved_fill_value"]

        fill_values = _compute_polars_fill_values(X, cols, strategy, fill_value)
        missing_counts, total_missing = _polars_missing_counts(X, cols)

        return {
            "type": "simple_imputer",
            "strategy": strategy,
            "fill_values": fill_values,
            "columns": cols,
            "missing_counts": missing_counts,
            "total_missing": total_missing,
        }

    @staticmethod
    def _fit_pandas(X: Any, _y: Any, params: dict[str, Any]) -> dict[str, Any]:
        cols: list[str] = params["_resolved_cols"]
        strategy: str = params["_resolved_strategy"]
        fill_value = params["_resolved_fill_value"]

        # Mean/median: extra safety filter to numeric columns only.
        if strategy in ("mean", "median"):
            numeric = set(detect_numeric_columns(X))
            cols = [c for c in cols if c in numeric]
            if not cols:
                return {}

        imputer = SimpleImputer(strategy=strategy, fill_value=fill_value)
        imputer.fit(X[cols])

        statistics = imputer.statistics_.tolist()
        fill_values = dict(zip(cols, statistics, strict=True))
        missing_counts = X[cols].isnull().sum().to_dict()
        total_missing = int(sum(missing_counts.values()))

        return {
            "type": "simple_imputer",
            "strategy": strategy,
            "fill_values": fill_values,
            "columns": cols,
            "missing_counts": missing_counts,
            "total_missing": total_missing,
        }