Source code for cursus.processing.numerical.numerical_imputation_processor

"""
Numerical Imputation Processor - Single Column Architecture

This processor performs imputation on a SINGLE numerical column.
Follows the single-column architecture pattern for real-time inference pipelines.

For batch processing of multiple columns, use one processor per column.
"""

from typing import Any, Union, Optional
import pandas as pd
import numpy as np
import pickle as pkl
import json
import logging
from pathlib import Path

from ..processors import Processor

# Setup logger
logger = logging.getLogger(__name__)


[docs] class NumericalVariableImputationProcessor(Processor): """ A processor that performs imputation on a SINGLE numerical variable/column. Designed for real-time inference pipelines where each processor handles one column and processors can be chained with >> operator. For batch processing of multiple columns, use one processor per column. Examples: >>> # Create processor for single column >>> proc = NumericalImputationProcessor( ... column_name='age', ... strategy='mean' ... ) >>> >>> # Fit on training data >>> proc.fit(train_df['age']) >>> >>> # Process single value (real-time inference) >>> imputed_value = proc.process(None) # Returns mean value >>> >>> # Transform Series or DataFrame >>> imputed_series = proc.transform(test_df['age']) """ def __init__( self, column_name: str, imputation_value: Optional[Union[int, float]] = None, strategy: Optional[str] = None, ): """ Initialize numerical imputation processor. Args: column_name: Name of the column to impute (single column) imputation_value: Pre-computed imputation value (for inference) strategy: Strategy for fitting ('mean', 'median', 'mode') Required if imputation_value is None Raises: ValueError: If column_name is empty or invalid ValueError: If both imputation_value and strategy are None """ super().__init__() # Initialize base Processor self.processor_name = "numerical_imputation_processor" self.function_name_list = ["fit", "process", "transform"] # Validate column_name if not isinstance(column_name, str) or not column_name: raise ValueError("column_name must be a non-empty string") self.column_name = column_name self.strategy = strategy # Set fitted state based on whether we have an imputation value self.is_fitted = imputation_value is not None if imputation_value is not None: self._validate_imputation_value(imputation_value) self.imputation_value = imputation_value else: self.imputation_value = None # If no imputation_value provided, strategy is required for fitting if strategy is None: raise ValueError("Either imputation_value or strategy must be provided")
[docs] def get_name(self) -> str: """Return processor name for base class compatibility.""" return self.processor_name
def _validate_imputation_value(self, value: Any) -> None: """ Validate that imputation value is numeric. Args: value: Value to validate Raises: ValueError: If value is not numeric """ if not isinstance(value, (int, float, np.number)): raise ValueError( f"Imputation value must be numeric, got {type(value)} for column '{self.column_name}'" ) if pd.isna(value): # A NaN imputation value imputes NaN — i.e. silently does nothing while # appearing to fill, corrupting the contract. Reject it outright. raise ValueError( f"Imputation value cannot be NaN for column '{self.column_name}'" )
[docs] def fit( self, X: Union[pd.Series, pd.DataFrame], y: Optional[pd.Series] = None ) -> "NumericalVariableImputationProcessor": """ Fit imputation value on a Series (single column). Args: X: Series (preferred) or DataFrame with column_name y: Ignored (for sklearn compatibility) Returns: self (for method chaining) Raises: ValueError: If column_name not found in DataFrame ValueError: If strategy is unknown """ # Extract Series if DataFrame provided if isinstance(X, pd.DataFrame): if self.column_name not in X.columns: raise ValueError(f"Column '{self.column_name}' not found in DataFrame") data = X[self.column_name] else: data = X # Calculate imputation value based on strategy if data.isna().all(): self.imputation_value = 0.0 # Default for all-NaN columns logger.warning( f"Column '{self.column_name}' has all NaN values, using 0.0 as imputation value" ) elif self.strategy == "mean": self.imputation_value = float(data.mean()) elif self.strategy == "median": self.imputation_value = float(data.median()) elif self.strategy == "mode": mode_values = data.mode() # pandas .mode() returns all tied modes sorted ascending; take the # smallest for a DETERMINISTIC choice (log when there is a tie so the # selection is transparent rather than silently arbitrary). if len(mode_values) > 1: logger.debug( "Column '%s' has %d tied modes %s; using the smallest (%s).", self.column_name, len(mode_values), list(mode_values), mode_values[0], ) self.imputation_value = ( float(mode_values[0]) if len(mode_values) > 0 else 0.0 ) else: raise ValueError(f"Unknown strategy: {self.strategy}") self.is_fitted = True logger.info( f"Fitted {self.column_name}: strategy={self.strategy}, " f"imputation_value={self.imputation_value:.4f}" ) return self
[docs] def process(self, input_value: Union[int, float, Any]) -> Union[int, float]: """ Process a SINGLE numerical value for this column. This method is called by __call__ (inherited from base Processor). It handles single-value processing for real-time inference. Args: input_value: Single value to impute if missing Returns: Imputed value (or original if not missing) Raises: RuntimeError: If processor not fitted """ if not self.is_fitted: raise RuntimeError( f"Processor for column '{self.column_name}' must be fitted before processing" ) # Handle missing values if pd.isna(input_value): return self.imputation_value return input_value
[docs] def transform( self, X: Union[pd.Series, pd.DataFrame, Any] ) -> Union[pd.Series, pd.DataFrame, float]: """ Transform data using the fitted imputation value. Args: X: Series, DataFrame, or single value Returns: Imputed data in same format as input Raises: RuntimeError: If processor not fitted ValueError: If column_name not found in DataFrame Performance optimized: Uses fast path for single-value Series. """ if not self.is_fitted: raise RuntimeError( f"Processor for column '{self.column_name}' must be fitted before transforming" ) # Handle Series if isinstance(X, pd.Series): # Fast path for single-value Series (10-100x faster) if len(X) == 1: val = X.iloc[0] result = self.process(val) return pd.Series([result], index=X.index) # Batch path for multiple values return X.fillna(self.imputation_value) # Handle DataFrame elif isinstance(X, pd.DataFrame): if self.column_name not in X.columns: raise ValueError(f"Column '{self.column_name}' not found in DataFrame") # Fast path for single-row DataFrame if len(X) == 1: df = X.copy() val = df[self.column_name].iloc[0] df[self.column_name] = self.process(val) return df # Batch path for multiple rows df = X.copy() df[self.column_name] = df[self.column_name].fillna(self.imputation_value) return df # Handle single value (delegate to process) else: return self.process(X)
[docs] def get_imputation_value(self) -> Union[int, float]: """ Get the fitted imputation value. Returns: Imputation value for this column Raises: RuntimeError: If processor not fitted """ if not self.is_fitted: raise RuntimeError( f"Processor for column '{self.column_name}' has not been fitted" ) return self.imputation_value
[docs] def set_imputation_value(self, value: Union[int, float]) -> None: """ Set imputation value (for pre-fitted processor). Args: value: Imputation value to use Raises: ValueError: If value is not numeric """ self._validate_imputation_value(value) self.imputation_value = value self.is_fitted = True logger.info( f"Set imputation value for '{self.column_name}': {self.imputation_value:.4f}" )
[docs] def get_params(self) -> dict: """ Get processor parameters (DEPRECATED). Use get_imputation_value() instead for the fitted value. Returns: Dictionary with all parameters """ import warnings warnings.warn( "get_params() is deprecated, use get_imputation_value() instead", DeprecationWarning, stacklevel=2, ) return { "column_name": self.column_name, "imputation_value": self.imputation_value, "strategy": self.strategy, }
[docs] def save_imputation_value(self, output_dir: Union[Path, str]) -> None: """ Save imputation value to disk. Creates two files: 1. {column_name}_impute_value.pkl (for loading) 2. {column_name}_impute_value.json (for human readability) Args: output_dir: Directory to save artifacts to Raises: RuntimeError: If processor not fitted Examples: >>> proc = NumericalImputationProcessor('age', strategy='mean') >>> proc.fit(train_df['age']) >>> proc.save_imputation_value('model_artifacts/') """ if not self.is_fitted: raise RuntimeError( f"Cannot save before fitting processor for column '{self.column_name}'" ) output_dir_path = Path(output_dir) output_dir_path.mkdir(parents=True, exist_ok=True) # Save pickle (for loading) pkl_file = output_dir_path / f"{self.column_name}_impute_value.pkl" with open(pkl_file, "wb") as f: pkl.dump(self.imputation_value, f) # Save JSON (for readability) json_file = output_dir_path / f"{self.column_name}_impute_value.json" with open(json_file, "w") as f: json.dump( { "column_name": self.column_name, "imputation_value": float(self.imputation_value), "strategy": self.strategy, }, f, indent=2, ) logger.info( f"Saved imputation value for '{self.column_name}' to {output_dir_path}" )
[docs] def load_imputation_value(self, filepath: Union[Path, str]) -> None: """ Load imputation value from disk. Args: filepath: Path to pickle file or directory containing it Raises: FileNotFoundError: If file not found ValueError: If value is invalid Examples: >>> proc = NumericalImputationProcessor('age', strategy='mean') >>> proc.load_imputation_value('model_artifacts/') >>> # Or specify exact file >>> proc.load_imputation_value('model_artifacts/age_impute_value.pkl') """ filepath_path = Path(filepath) # Handle directory path if filepath_path.is_dir(): pkl_file = filepath_path / f"{self.column_name}_impute_value.pkl" else: pkl_file = filepath_path if not pkl_file.exists(): raise FileNotFoundError(f"Imputation value file not found: {pkl_file}") # Load value with open(pkl_file, "rb") as f: loaded_value = pkl.load(f) # Validate self._validate_imputation_value(loaded_value) self.imputation_value = loaded_value self.is_fitted = True logger.info( f"Loaded imputation value for '{self.column_name}' from {pkl_file}: {self.imputation_value:.4f}" )
[docs] @classmethod def from_imputation_dict(cls, imputation_dict: dict) -> dict: """ Create processors from script output (impute_dict.pkl). This factory method simplifies creating multiple processors from the dictionary format used by missing_value_imputation.py script. Args: imputation_dict: Dictionary mapping column names to imputation values Format: {column_name: imputation_value} Returns: Dictionary mapping column names to fitted processors Raises: TypeError: If imputation_dict is not a dictionary ValueError: If column name is not a string Examples: >>> with open("impute_dict.pkl", "rb") as f: ... impute_dict = pkl.load(f) >>> processors = NumericalImputationProcessor.from_imputation_dict(impute_dict) >>> # Use processors in pipeline >>> for col, proc in processors.items(): ... dataset.add_pipeline(col, proc) """ if not isinstance(imputation_dict, dict): raise TypeError("imputation_dict must be a dictionary") processors = {} for col, val in imputation_dict.items(): if not isinstance(col, str): raise ValueError(f"Column name must be string, got {type(col)}") proc = cls(column_name=col, imputation_value=val) processors[col] = proc logger.info(f"Created {len(processors)} processors from imputation dict") return processors
[docs] @classmethod def from_script_artifacts( cls, artifacts_dir: Union[Path, str], filename: str = "impute_dict.pkl" ) -> dict: """ Load processors from script output directory. Looks for impute_dict.pkl in the specified directory and creates processors from it. Args: artifacts_dir: Directory containing impute_dict.pkl filename: Name of the imputation dict file (default: impute_dict.pkl) Returns: Dictionary mapping column names to fitted processors Raises: FileNotFoundError: If impute_dict.pkl not found Examples: >>> processors = NumericalImputationProcessor.from_script_artifacts( ... "model_artifacts/" ... ) >>> for col, proc in processors.items(): ... dataset.add_pipeline(col, proc) """ artifacts_path = Path(artifacts_dir) impute_dict_file = artifacts_path / filename if not impute_dict_file.exists(): raise FileNotFoundError(f"{filename} not found in {artifacts_path}") with open(impute_dict_file, "rb") as f: impute_dict = pkl.load(f) logger.info(f"Loaded imputation dict from {impute_dict_file}") return cls.from_imputation_dict(impute_dict)