cursus.steps.scripts.missing_value_imputation

Missing Value Imputation Processing Script

This script handles missing value imputation for tabular data using simple statistical methods. It supports both training mode (fit and transform) and inference mode (transform only). Follows the same pattern as risk_table_mapping.py for consistency.

find_input_shards(input_dir, log_func)

Find all input shards in directory.

Searches for various shard formats (CSV, JSON, Parquet with/without compression).

Parameters:

• input_dir (str) – Directory containing input shards

• log_func (Callable) – Logging function

Returns:

Sorted list of shard paths

Raises:

RuntimeError – If no shards found in input directory

Return type:

List[Path]

find_split_shards(input_dir, split_name, log_func)

Find all input shards in a specific split subdirectory.

Used when input data is organized as: input_dir/

train/part-00000.csv, part-00001.csv, … val/part-00000.csv, part-00001.csv, … test/part-00000.csv, part-00001.csv, …

Parameters:

• input_dir (str) – Base input directory

• split_name (str) – Split subdirectory name (“train”, “val”, “test”, etc.)

• log_func (Callable) – Logging function

Returns:

Sorted list of shard paths from the split subdirectory

Raises:

RuntimeError – If split subdirectory or shards not found

Return type:

List[Path]

extract_shard_number(shard_path)

Extract shard number from filename like part-00042.csv.

Handles various formats: - part-00042.csv → 42 - part-00042.csv.gz → 42 - part-00042.parquet → 42 - part-00042.snappy.parquet → 42

Parameters:

shard_path (Path) – Path to shard file

Returns:

Integer shard number

Raises:

ValueError – If shard number cannot be extracted

Return type:

int

Example

>>> extract_shard_number(Path("part-00042.csv"))
42
>>> extract_shard_number(Path("part-00001.csv.gz"))
1

write_shard_file(df, output_path, output_format)

Write a DataFrame to a shard file in the specified format.

Creates parent directories if needed.

Parameters:

• df (DataFrame) – DataFrame to write

• output_path (Path) – Full path for output file (including filename)

• output_format (str) – Format to write (‘csv’, ‘tsv’, or ‘parquet’)

Raises:

ValueError – If output_format is not supported

aggregate_shard_results(results, job_type)

Aggregate statistics from parallel shard processing.

Parameters:

• results (List[Dict[str, int]]) – List of statistics dictionaries from each shard

• job_type (str) – Type of job (‘training’, ‘validation’, etc.’)

Returns:

Dictionary with total row counts per split

Return type:

Dict[str, int]

detect_shard_format(shard_path)

Auto-detect output format from input shard filename.

Mirrors batch mode’s format preservation behavior.

Parameters:

shard_path (Path) – Path to a shard file

Returns:

‘csv’, ‘tsv’, or ‘parquet’

Return type:

Format string

Example

>>> detect_shard_format(Path("part-00001.csv"))
'csv'
>>> detect_shard_format(Path("part-00001.parquet"))
'parquet'

collect_imputation_statistics_pass1(all_shards, signature_columns, label_field, imputation_config, log_func)

Pass 1: Collect imputation statistics from training shards.

Memory-efficient incremental aggregation: - Numeric columns: Accumulate sum + count → compute mean - Categorical/Text columns: Collect all non-null values → compute mode

Parameters:

• all_shards (List[Path]) – List of all input shard paths

• signature_columns (List[str] | None) – Optional column names for CSV/TSV files

• label_field (str) – Name of label column to exclude from imputation

• imputation_config (Dict[str, Any]) – Imputation configuration dictionary

• log_func (Callable) – Logging function

Returns:

Dictionary mapping column names to imputation values Format: {column_name: imputation_value} (XGBoost compatible)

Return type:

Dict[str, Any]

process_shard_end_to_end_imputation(args)

Process single shard: read → apply imputation → write.

Stateless per-shard processing using global impute_dict from Pass 1. Preserves 1:1 shard mapping (input shard number → output shard number).

Parameters:

• args (tuple) – Tuple of (shard_path, shard_num, global_context, output_base, signature_columns, output_format)

• contain (global_context must)

• "impute_dict" (-) – Dictionary of imputation values

• "split_name" (-) – Which split this shard belongs to (“train”, “val”, “test”, etc.)

Returns:

Statistics dict with row count for this split Format: {“train”: 1000} or {“val”: 200} or {“validation”: 500}

Return type:

Dict[str, int]

Example

Input: train/part-00042.csv Output: train/part-00042.csv (imputed)

process_streaming_mode_imputation(input_dir, output_dir, signature_columns, job_type, label_field, imputation_config, max_workers, model_artifacts_input_dir=None, model_artifacts_output_dir=None, logger=None)

Streaming mode for missing value imputation with train/val/test subdirectories.

Two-pass architecture: - Pass 1: Collect imputation statistics from training shards only - Pass 2: Apply imputations per split in parallel

Auto-detects output format from input shards (mirrors batch mode behavior).

Input structure (training mode):

input_dir/

train/part-00000.csv, part-00001.csv, … val/part-00000.csv, part-00001.csv, … test/part-00000.csv, part-00001.csv, …

Output structure (training mode):

output_dir/

train/part-00000.csv, part-00001.csv, … (imputed, same format) val/part-00000.csv, part-00001.csv, … (imputed, same format) test/part-00000.csv, part-00001.csv, … (imputed, same format)

Parameters:

• input_dir (str) – Base input directory

• output_dir (str) – Base output directory

• signature_columns (List[str] | None) – Optional column names for CSV/TSV

• job_type (str) – ‘training’, ‘validation’, ‘testing’, ‘calibration’

• label_field (str) – Label column to exclude

• imputation_config (Dict[str, Any]) – Imputation configuration

• max_workers (int | None) – Number of parallel workers

• model_artifacts_input_dir (str | None) – Input model artifacts directory

• model_artifacts_output_dir (str | None) – Output model artifacts directory

• logger (Callable | None) – Logging function

Returns:

Dictionary with total row counts per split

Return type:

Dict[str, int]

load_split_data(job_type, input_dir)

Load data according to job_type with automatic format detection.

For ‘training’: Loads data from train, test, and val subdirectories For others: Loads single job_type split

Returns:

Dictionary with DataFrames and detected format stored in ‘format’ key

Return type:

Dict[str, DataFrame]

save_output_data(job_type, output_dir, data_dict)

Save processed data according to job_type, preserving input format.

For ‘training’: Saves data to train, test, and val subdirectories For others: Saves to single job_type directory

analyze_missing_values(df)

Comprehensive missing value analysis for imputation planning.

validate_imputation_data(df, label_field, exclude_columns=None)

Validate data for imputation processing.

load_imputation_config(environ_vars)

Load imputation configuration from environment variables.

get_pandas_na_values()

Get set of values that pandas interprets as NA/NULL.

validate_text_fill_value(value)

Validate that a text fill value won’t be interpreted as NA by pandas.

detect_column_type(df, column, config)

Enhanced data type detection for imputation strategy selection.

class ImputationStrategyManager(config)

Bases: object

Enhanced strategy manager supporting numerical, text, and categorical data types.

get_strategy_for_column(df, column)

Enhanced strategy selection supporting text and categorical types.

class SimpleImputationEngine(strategy_manager, label_field)

Bases: object

Core engine for simple statistical imputation methods.

fit(df)

Fit imputation parameters on training data.

transform(df)

Apply fitted imputation to data.

fit_transform(df)

Fit imputation parameters and transform data in one step.

get_imputation_summary()

Get comprehensive summary of imputation process.

save_imputation_artifacts(imputation_engine, imputation_config, output_path)

Save imputation artifacts to the specified output path.

Output format matches XGBoost training’s impute_dict.pkl format: A simple dictionary mapping column names to imputation values.

Parameters:

• imputation_engine (SimpleImputationEngine) – SimpleImputationEngine instance with fitted parameters

• imputation_config (Dict[str, Any]) – Imputation configuration dictionary

• output_path (Path) – Path to save artifacts to

load_imputation_parameters(imputation_params_path)

Load imputation parameters from a pickle file.

Expected format (XGBoost training compatible): Simple dict mapping column names to imputation values: {column: value}

Parameters:

imputation_params_path (Path) – Path to the imputation parameters file

Returns:

imputation_value}

Return type:

Dictionary of imputation parameters {column_name

process_data(data_dict, label_field, job_type, imputation_config, imputation_parameters=None)

Core data processing logic for missing value imputation.

Parameters:

• data_dict (Dict[str, DataFrame]) – Dictionary of dataframes keyed by split name

• label_field (str) – Target column name

• job_type (str) – Type of job (training, validation, testing, calibration)

• imputation_config (Dict[str, Any]) – Imputation configuration dictionary

• imputation_parameters (Dict | None) – Pre-fitted imputation parameters (simple dict {column: value})

Returns:

• Dictionary of imputed dataframes

• SimpleImputationEngine instance with fitted parameters

Return type:

Tuple containing

generate_imputation_report(imputation_engine, missing_analysis, validation_report, output_dir)

Generate comprehensive imputation report with statistics and insights.

calculate_imputation_quality_metrics(imputation_summary)

Calculate quality metrics for imputation process.

generate_imputation_recommendations(imputation_summary, missing_analysis)

Generate actionable recommendations based on imputation analysis.

copy_existing_artifacts(src_dir, dst_dir)

Copy all existing model artifacts from previous processing steps.

This enables the parameter accumulator pattern where each step: 1. Copies artifacts from previous steps 2. Adds its own artifacts 3. Passes all artifacts to the next step

Parameters:

• src_dir (str) – Source directory containing existing artifacts

• dst_dir (str) – Destination directory to copy artifacts to

generate_imputation_text_summary(report)

Generate human-readable text summary of imputation process.

internal_main(job_type, input_dir, output_dir, imputation_config, label_field, model_artifacts_input_dir=None, model_artifacts_output_dir=None, enable_true_streaming=False, max_workers=None, load_data_func=<function load_split_data>, save_data_func=<function save_output_data>)

Main logic for missing value imputation, handling both training and inference modes.

Supports two modes: - Batch mode (default): Loads entire splits into memory - Streaming mode: Processes shards in parallel (memory-efficient)

Parameters:

• job_type (str) – Type of job (training, validation, testing, calibration)

• input_dir (str) – Input directory for data

• output_dir (str) – Output directory for processed data

• imputation_config (Dict[str, Any]) – Imputation configuration dictionary

• label_field (str) – Target column name

• model_artifacts_input_dir (str | None) – Directory containing model artifacts from previous steps

• model_artifacts_output_dir (str | None) – Directory to save model artifacts for next steps

• enable_true_streaming (bool) – Enable streaming mode (default: False)

• signature_columns – Optional column names for streaming mode CSV/TSV

• output_format – Output format for streaming mode (csv, tsv, parquet)

• max_workers (int | None) – Number of parallel workers for streaming mode

• load_data_func (Callable) – Function to load data (for dependency injection in tests)

• save_data_func (Callable) – Function to save data (for dependency injection in tests)

Returns:

• Dictionary of imputed dataframes (empty in streaming mode)

• SimpleImputationEngine instance with fitted parameters (or None in streaming)

Return type:

Tuple containing

main(input_paths, output_paths, environ_vars, job_args=None)

Standardized main entry point for missing value imputation script.

Parameters:

• input_paths (Dict[str, str]) – Dictionary of input paths with logical names - “data_input”: Input data directory (from tabular_preprocessing) - “model_artifacts_input”: Model artifacts from previous steps (standardized)

• output_paths (Dict[str, str]) – Dictionary of output paths with logical names - “data_output”: Output directory for imputed data - “model_artifacts_output”: Model artifacts output for next steps (standardized)

• environ_vars (Dict[str, str]) – Dictionary of environment variables

• job_args (Namespace | None) – Command line arguments containing job_type

Returns:

• Dictionary of imputed dataframes

• SimpleImputationEngine instance with fitted parameters

Return type:

Tuple containing