#!/usr/bin/env python
"""
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.
"""
import argparse
import os
import sys
import pandas as pd
import json
import pickle as pkl
import traceback
import shutil
import gc
from pathlib import Path
from collections import Counter
from multiprocessing import Pool, cpu_count
from sklearn.impute import SimpleImputer
import logging
from typing import Dict, List, Tuple, Any, Optional, Callable
from datetime import datetime
# Default paths (will be overridden by parameters in main function)
DEFAULT_INPUT_DIR = "/opt/ml/processing/input/data"
DEFAULT_OUTPUT_DIR = "/opt/ml/processing/output"
DEFAULT_MODEL_ARTIFACTS_DIR = "/opt/ml/processing/input/model_artifacts"
# Constants for file paths to ensure consistency between training and inference
# Match XGBoost training output format
IMPUTATION_PARAMS_FILENAME = "impute_dict.pkl"
IMPUTATION_SUMMARY_FILENAME = "imputation_summary.json"
# Set up logging
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s - %(levelname)s - [%(filename)s:%(lineno)d] - %(message)s",
datefmt="%Y-%m-%d %H:%M:%S",
)
logger = logging.getLogger(__name__)
# ============================================================================
# STREAMING MODE UTILITIES (Reused from temporal_split_preprocessing)
# ============================================================================
[docs]
def find_split_shards(
input_dir: str, split_name: str, log_func: Callable
) -> List[Path]:
"""
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, ...
Args:
input_dir: Base input directory
split_name: Split subdirectory name ("train", "val", "test", etc.)
log_func: Logging function
Returns:
Sorted list of shard paths from the split subdirectory
Raises:
RuntimeError: If split subdirectory or shards not found
"""
split_dir = Path(input_dir) / split_name
if not split_dir.exists():
raise RuntimeError(f"Split subdirectory not found: {split_dir}")
patterns = [
"part-*.csv",
"part-*.csv.gz",
"part-*.json",
"part-*.json.gz",
"part-*.parquet",
"part-*.snappy.parquet",
"part-*.parquet.gz",
]
all_shards = sorted([p for pat in patterns for p in split_dir.glob(pat)])
if not all_shards:
raise RuntimeError(f"No shards found in {split_dir}")
log_func(f"[STREAMING] Found {len(all_shards)} shards in {split_name} split")
return all_shards
[docs]
def write_shard_file(df: pd.DataFrame, output_path: Path, output_format: str) -> None:
"""
Write a DataFrame to a shard file in the specified format.
Creates parent directories if needed.
Args:
df: DataFrame to write
output_path: Full path for output file (including filename)
output_format: Format to write ('csv', 'tsv', or 'parquet')
Raises:
ValueError: If output_format is not supported
"""
# Create parent directory if needed
output_path.parent.mkdir(parents=True, exist_ok=True)
if output_format == "csv":
df.to_csv(output_path, index=False)
elif output_format == "tsv":
df.to_csv(output_path, sep="\t", index=False)
elif output_format == "parquet":
df.to_parquet(output_path, index=False)
else:
raise ValueError(f"Unsupported output format: {output_format}")
[docs]
def aggregate_shard_results(
results: List[Dict[str, int]], job_type: str
) -> Dict[str, int]:
"""
Aggregate statistics from parallel shard processing.
Args:
results: List of statistics dictionaries from each shard
job_type: Type of job ('training', 'validation', etc.')
Returns:
Dictionary with total row counts per split
"""
if job_type == "training":
# Training mode: aggregate train/val/test splits
total_stats = {
"train": sum(r.get("train", 0) for r in results),
"val": sum(r.get("val", 0) for r in results),
"test": sum(r.get("test", 0) for r in results),
}
else:
# Single split mode
total_stats = {job_type: sum(r.get(job_type, 0) for r in results)}
return total_stats
# ============================================================================
# STREAMING MODE - PASS 1: COLLECT IMPUTATION STATISTICS
# ============================================================================
def _read_file_to_df(
file_path: Path, column_names: Optional[List[str]] = None
) -> pd.DataFrame:
"""
Read a single file (CSV, TSV, JSON, Parquet) into a DataFrame.
Simplified version for streaming mode - handles common formats.
Args:
file_path: Path to file
column_names: Optional column names (for CSV/TSV files)
Returns:
DataFrame from file
"""
suffix = file_path.suffix.lower()
if suffix == ".csv" or (suffix == ".gz" and file_path.stem.endswith(".csv")):
if column_names:
return pd.read_csv(file_path, names=column_names, header=0)
return pd.read_csv(file_path)
elif suffix == ".tsv" or (suffix == ".gz" and file_path.stem.endswith(".tsv")):
if column_names:
return pd.read_csv(file_path, sep="\t", names=column_names, header=0)
return pd.read_csv(file_path, sep="\t")
elif suffix == ".parquet" or suffix.endswith(".parquet"):
return pd.read_parquet(file_path)
else:
# Default to CSV
return pd.read_csv(file_path)
[docs]
def collect_imputation_statistics_pass1(
all_shards: List[Path],
signature_columns: Optional[List[str]],
label_field: str,
imputation_config: Dict[str, Any],
log_func: Callable,
) -> Dict[str, Any]:
"""
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
Args:
all_shards: List of all input shard paths
signature_columns: Optional column names for CSV/TSV files
label_field: Name of label column to exclude from imputation
imputation_config: Imputation configuration dictionary
log_func: Logging function
Returns:
Dictionary mapping column names to imputation values
Format: {column_name: imputation_value} (XGBoost compatible)
"""
log_func("[PASS1] Collecting imputation statistics from training shards...")
# Step 1: Identify imputable columns from first shard
first_shard = all_shards[0]
df_first = _read_file_to_df(first_shard, signature_columns)
df_first.columns = [col.replace("__DOT__", ".") for col in df_first.columns]
# Get columns to impute (exclude label and configured exclusions)
exclude_cols = [label_field] + imputation_config.get("exclude_columns", [])
imputable_columns = [
col
for col in df_first.columns
if col not in exclude_cols and df_first[col].isnull().any()
]
log_func(f"[PASS1] Found {len(imputable_columns)} columns with missing values")
if not imputable_columns:
log_func("[PASS1] No columns with missing values found")
return {}
# Step 2: Initialize aggregators for each column
column_aggregators = {}
for col in imputable_columns:
# Detect column type
col_type = detect_column_type(df_first, col, imputation_config)
if col_type == "numerical":
column_aggregators[col] = {
"type": "numerical",
"sum": 0.0,
"count": 0,
"dtype": str(df_first[col].dtype),
}
else: # categorical or text
column_aggregators[col] = {
"type": "categorical",
"values": [],
"dtype": str(df_first[col].dtype),
}
del df_first
gc.collect()
# Step 3: Process each shard and aggregate statistics
log_func(f"[PASS1] Processing {len(all_shards)} shards...")
for i, shard_path in enumerate(all_shards):
try:
df = _read_file_to_df(shard_path, signature_columns)
df.columns = [col.replace("__DOT__", ".") for col in df.columns]
# Aggregate statistics for each column
for col in imputable_columns:
if col not in df.columns:
continue
# Get non-null values
non_null_values = df[col].dropna()
if len(non_null_values) == 0:
continue
aggregator = column_aggregators[col]
if aggregator["type"] == "numerical":
# Accumulate sum and count for mean calculation
aggregator["sum"] += non_null_values.sum()
aggregator["count"] += len(non_null_values)
else:
# Collect values for mode calculation
# Memory optimization: sample if too many unique values
if len(aggregator["values"]) < 100000:
aggregator["values"].extend(non_null_values.tolist())
else:
# Already have enough samples, just add unique values
unique_new = non_null_values.unique()
if len(unique_new) < 1000:
aggregator["values"].extend(unique_new.tolist())
del df
gc.collect()
if (i + 1) % 100 == 0:
log_func(f"[PASS1] Processed {i + 1}/{len(all_shards)} shards")
except Exception as e:
log_func(f"[PASS1 WARNING] Failed to read {shard_path.name}: {e}")
continue
# Step 4: Compute final imputation values
log_func("[PASS1] Computing final imputation values...")
impute_dict = {}
for col, aggregator in column_aggregators.items():
try:
if aggregator["type"] == "numerical":
# Compute mean
if aggregator["count"] > 0:
impute_value = aggregator["sum"] / aggregator["count"]
impute_dict[col] = float(impute_value)
log_func(f"[PASS1] {col}: mean = {impute_value:.4f}")
else:
log_func(f"[PASS1] {col}: No non-null values, using 0")
impute_dict[col] = 0.0
else:
# Compute mode (most frequent value)
if aggregator["values"]:
# Use Counter to find most common value
mode_value = Counter(aggregator["values"]).most_common(1)[0][0]
impute_dict[col] = mode_value
log_func(f"[PASS1] {col}: mode = '{mode_value}'")
else:
log_func(f"[PASS1] {col}: No non-null values, using 'Unknown'")
impute_dict[col] = "Unknown"
except Exception as e:
log_func(f"[PASS1 WARNING] Failed to compute imputation for {col}: {e}")
# Use safe defaults
if aggregator["type"] == "numerical":
impute_dict[col] = 0.0
else:
impute_dict[col] = "Unknown"
log_func(
f"[PASS1] Complete! Collected imputation values for {len(impute_dict)} columns"
)
# Estimate memory usage
memory_mb = len(impute_dict) * 50 / 1024 / 1024 # Rough estimate
log_func(f"[PASS1] Map size: ~{memory_mb:.2f} MB")
return impute_dict
# ============================================================================
# STREAMING MODE - PASS 2: PARALLEL PER-SHARD IMPUTATION
# ============================================================================
[docs]
def process_shard_end_to_end_imputation(args: tuple) -> Dict[str, int]:
"""
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).
Args:
args: Tuple of (shard_path, shard_num, global_context,
output_base, signature_columns, output_format)
global_context must contain:
- "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}
Example:
Input: train/part-00042.csv
Output: train/part-00042.csv (imputed)
"""
(
shard_path,
shard_num,
global_context,
output_base,
signature_columns,
output_format,
) = args
try:
# ====================================================================
# STEP 1: Read Single Shard
# ====================================================================
df = _read_file_to_df(shard_path, signature_columns)
df.columns = [col.replace("__DOT__", ".") for col in df.columns]
# ====================================================================
# STEP 2: Apply Imputation (Using Global Context)
# ====================================================================
impute_dict = global_context["impute_dict"]
# Simple fillna operation for each column
for column, impute_value in impute_dict.items():
if column in df.columns:
# Only fill NaN values (preserve existing non-null values)
df[column] = df[column].fillna(impute_value)
# ====================================================================
# STEP 3: Write to Correct Split Folder (Preserving Shard Number)
# ====================================================================
split_name = global_context["split_name"]
stats = {}
if len(df) > 0:
output_path = (
output_base / split_name / f"part-{shard_num:05d}.{output_format}"
)
write_shard_file(df, output_path, output_format)
stats[split_name] = len(df)
else:
stats[split_name] = 0
return stats
except Exception as e:
# Log error but don't crash the entire pool
print(f"[ERROR] Failed to process shard {shard_num} ({shard_path.name}): {e}")
# Return zero stats for this shard
split_name = global_context.get("split_name", "unknown")
return {split_name: 0}
# ============================================================================
# STREAMING MODE - MAIN ORCHESTRATION
# ============================================================================
[docs]
def process_streaming_mode_imputation(
input_dir: str,
output_dir: str,
signature_columns: Optional[List[str]],
job_type: str,
label_field: str,
imputation_config: Dict[str, Any],
max_workers: Optional[int],
model_artifacts_input_dir: Optional[str] = None,
model_artifacts_output_dir: Optional[str] = None,
logger: Optional[Callable] = None,
) -> Dict[str, int]:
"""
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)
Args:
input_dir: Base input directory
output_dir: Base output directory
signature_columns: Optional column names for CSV/TSV
job_type: 'training', 'validation', 'testing', 'calibration'
label_field: Label column to exclude
imputation_config: Imputation configuration
max_workers: Number of parallel workers
model_artifacts_input_dir: Input model artifacts directory
model_artifacts_output_dir: Output model artifacts directory
logger: Logging function
Returns:
Dictionary with total row counts per split
"""
log = logger or print
output_path = Path(output_dir)
# Determine optimal workers
if max_workers is None:
max_workers = min(cpu_count(), 8) # Default to 8 workers
log(f"[STREAMING] Starting streaming mode imputation")
log(f"[STREAMING] Job type: {job_type}")
log(f"[STREAMING] Max workers: {max_workers}")
# ========================================================================
# PASS 1: Collect Imputation Statistics (Training Only)
# ========================================================================
if job_type == "training":
log("[STREAMING] PASS 1: Collecting imputation statistics from train split...")
train_shards = find_split_shards(input_dir, "train", log)
impute_dict = collect_imputation_statistics_pass1(
train_shards, signature_columns, label_field, imputation_config, log
)
# Save imputation artifacts
if model_artifacts_output_dir:
artifacts_path = Path(model_artifacts_output_dir)
artifacts_path.mkdir(parents=True, exist_ok=True)
# Save impute_dict
impute_dict_path = artifacts_path / IMPUTATION_PARAMS_FILENAME
with open(impute_dict_path, "wb") as f:
pkl.dump(impute_dict, f)
log(f"[STREAMING] Saved imputation dictionary to {impute_dict_path}")
else:
# Non-training: Load imputation parameters
if not model_artifacts_input_dir:
raise ValueError(f"model_artifacts_input_dir required for {job_type} mode")
impute_dict_path = Path(model_artifacts_input_dir) / IMPUTATION_PARAMS_FILENAME
if not impute_dict_path.exists():
raise FileNotFoundError(
f"Imputation parameters not found: {impute_dict_path}"
)
log(f"[STREAMING] Loading imputation parameters from {impute_dict_path}")
with open(impute_dict_path, "rb") as f:
impute_dict = pkl.load(f)
log(f"[STREAMING] Loaded {len(impute_dict)} imputation values")
# ========================================================================
# PASS 2: Process Each Split Independently
# ========================================================================
log("[STREAMING] PASS 2: Processing splits in parallel...")
# Determine which splits to process
if job_type == "training":
splits_to_process = ["train", "val", "test"]
else:
splits_to_process = [
job_type
] # Single split (validation, testing, calibration)
total_stats = {}
for split_name in splits_to_process:
log(f"[STREAMING] Processing {split_name} split...")
# Find shards for this split
split_shards = find_split_shards(input_dir, split_name, log)
# Auto-detect format from first shard (mirrors batch mode behavior)
output_format = detect_shard_format(split_shards[0])
log(f"[STREAMING] Detected format: {output_format}")
# Build global context for this split
global_context = {
"split_name": split_name,
"impute_dict": impute_dict,
}
# Prepare arguments for parallel processing
shard_args = [
(
shard,
extract_shard_number(shard),
global_context,
output_path,
signature_columns,
output_format,
)
for shard in split_shards
]
# Process shards in parallel
log(
f"[STREAMING] Processing {len(shard_args)} shards from {split_name} with {max_workers} workers"
)
with Pool(processes=max_workers) as pool:
results = pool.map(process_shard_end_to_end_imputation, shard_args)
# Aggregate results for this split
split_total = sum(r.get(split_name, 0) for r in results)
total_stats[split_name] = split_total
log(f"[STREAMING] Completed {split_name} split: {split_total:,} rows")
log(f"[STREAMING] Complete! Row distribution: {total_stats}")
return total_stats
# --- File I/O Helper Functions with Format Preservation ---
def _detect_file_format(split_dir: Path, split_name: str) -> tuple:
"""
Detect the format of processed data file.
Returns:
Tuple of (file_path, format) where format is 'csv', 'tsv', or 'parquet'
"""
# Try different formats in order of preference
formats = [
(f"{split_name}_processed_data.csv", "csv"),
(f"{split_name}_processed_data.tsv", "tsv"),
(f"{split_name}_processed_data.parquet", "parquet"),
]
for filename, fmt in formats:
file_path = split_dir / filename
if file_path.exists():
return file_path, fmt
raise RuntimeError(
f"No processed data file found in {split_dir}. "
f"Looked for: {[f[0] for f in formats]}"
)
[docs]
def load_split_data(job_type: str, input_dir: str) -> Dict[str, pd.DataFrame]:
"""
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
"""
input_path = Path(input_dir)
result = {}
if job_type == "training":
# For training, we expect data in train/test/val subdirectories
splits = ["train", "test", "val"]
detected_format = None
for split_name in splits:
split_dir = input_path / split_name
file_path, fmt = _detect_file_format(split_dir, split_name)
# Store format from first split (they should all match)
if detected_format is None:
detected_format = fmt
# Read based on format
if fmt == "csv":
df = pd.read_csv(file_path)
elif fmt == "tsv":
df = pd.read_csv(file_path, sep="\t")
elif fmt == "parquet":
df = pd.read_parquet(file_path)
else:
raise RuntimeError(f"Unsupported format: {fmt}")
result[split_name] = df
result["_format"] = detected_format # Store detected format
logger.info(
f"Loaded training data splits (format={detected_format}): "
f"train={result['train'].shape}, test={result['test'].shape}, val={result['val'].shape}"
)
else:
# For other job types, we expect data in a single directory named after job_type
split_dir = input_path / job_type
file_path, detected_format = _detect_file_format(split_dir, job_type)
# Read based on format
if detected_format == "csv":
df = pd.read_csv(file_path)
elif detected_format == "tsv":
df = pd.read_csv(file_path, sep="\t")
elif detected_format == "parquet":
df = pd.read_parquet(file_path)
else:
raise RuntimeError(f"Unsupported format: {detected_format}")
result[job_type] = df
result["_format"] = detected_format # Store detected format
logger.info(f"Loaded {job_type} data (format={detected_format}): {df.shape}")
return result
[docs]
def save_output_data(
job_type: str, output_dir: str, data_dict: Dict[str, pd.DataFrame]
) -> None:
"""
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
"""
output_path = Path(output_dir)
# Extract format from data_dict (stored during load)
output_format = data_dict.get("_format", "csv") # Default to CSV if not found
for split_name, df in data_dict.items():
# Skip the format metadata key
if split_name == "_format":
continue
split_output_dir = output_path / split_name
split_output_dir.mkdir(exist_ok=True, parents=True)
# Save in detected format
if output_format == "csv":
output_file = split_output_dir / f"{split_name}_processed_data.csv"
df.to_csv(output_file, index=False)
elif output_format == "tsv":
output_file = split_output_dir / f"{split_name}_processed_data.tsv"
df.to_csv(output_file, sep="\t", index=False)
elif output_format == "parquet":
output_file = split_output_dir / f"{split_name}_processed_data.parquet"
df.to_parquet(output_file, index=False)
else:
raise RuntimeError(f"Unsupported output format: {output_format}")
logger.info(
f"Saved {split_name} data to {output_file} (format={output_format}), shape: {df.shape}"
)
[docs]
def analyze_missing_values(df: pd.DataFrame) -> Dict[str, Any]:
"""
Comprehensive missing value analysis for imputation planning.
"""
missing_analysis = {
"total_records": len(df),
"columns_with_missing": {},
"missing_patterns": {},
"data_types": {},
"imputation_recommendations": {},
}
for col in df.columns:
missing_count = df[col].isnull().sum()
missing_percentage = (missing_count / len(df)) * 100
if missing_count > 0:
missing_analysis["columns_with_missing"][col] = {
"missing_count": int(missing_count),
"missing_percentage": float(missing_percentage),
"data_type": str(df[col].dtype),
"unique_values": int(df[col].nunique()),
"sample_values": df[col].dropna().head(5).tolist(),
}
# Recommend imputation strategy based on data type and distribution
if pd.api.types.is_numeric_dtype(df[col]):
try:
skewness = df[col].skew()
# skew() can return NaN (e.g. constant column, <3 non-null
# values); treat that as not-skewed and fall back to mean.
if pd.notna(skewness) and abs(skewness) > 1: # Highly skewed
missing_analysis["imputation_recommendations"][col] = "median"
else:
missing_analysis["imputation_recommendations"][col] = "mean"
except (ValueError, TypeError, AttributeError) as e:
logger.warning(
f"Skewness computation failed for column '{col}' "
f"({type(e).__name__}: {e}); defaulting to mean imputation."
)
missing_analysis["imputation_recommendations"][col] = "mean"
else:
missing_analysis["imputation_recommendations"][col] = "mode"
missing_analysis["data_types"][col] = str(df[col].dtype)
# Analyze missing patterns
missing_pattern = df.isnull().sum(axis=1)
missing_analysis["missing_patterns"] = {
"records_with_no_missing": int((missing_pattern == 0).sum()),
"records_with_missing": int((missing_pattern > 0).sum()),
"max_missing_per_record": int(missing_pattern.max()),
"avg_missing_per_record": float(missing_pattern.mean()),
}
return missing_analysis
[docs]
def validate_imputation_data(
df: pd.DataFrame, label_field: str, exclude_columns: List[str] = None
) -> Dict[str, Any]:
"""
Validate data for imputation processing.
"""
exclude_columns = exclude_columns or []
validation_report = {
"is_valid": True,
"errors": [],
"warnings": [],
"imputable_columns": [],
"excluded_columns": exclude_columns.copy(),
}
# Check if label field exists and exclude it from imputation
if label_field in df.columns:
validation_report["excluded_columns"].append(label_field)
else:
validation_report["warnings"].append(
f"Label field '{label_field}' not found in data"
)
# Identify columns suitable for imputation
for col in df.columns:
if col not in validation_report["excluded_columns"]:
if df[col].isnull().any():
validation_report["imputable_columns"].append(col)
if not validation_report["imputable_columns"]:
validation_report["warnings"].append(
"No columns with missing values found for imputation"
)
return validation_report
[docs]
def load_imputation_config(environ_vars: Dict[str, str]) -> Dict[str, Any]:
"""
Load imputation configuration from environment variables.
"""
config = {
"default_numerical_strategy": environ_vars.get(
"DEFAULT_NUMERICAL_STRATEGY", "mean"
),
"default_categorical_strategy": environ_vars.get(
"DEFAULT_CATEGORICAL_STRATEGY", "mode"
),
"default_text_strategy": environ_vars.get("DEFAULT_TEXT_STRATEGY", "mode"),
"numerical_constant_value": float(
environ_vars.get("NUMERICAL_CONSTANT_VALUE", "0")
),
"categorical_constant_value": environ_vars.get(
"CATEGORICAL_CONSTANT_VALUE", "Unknown"
),
"text_constant_value": environ_vars.get("TEXT_CONSTANT_VALUE", "Unknown"),
"categorical_preserve_dtype": environ_vars.get(
"CATEGORICAL_PRESERVE_DTYPE", "true"
).lower()
== "true",
"auto_detect_categorical": environ_vars.get(
"AUTO_DETECT_CATEGORICAL", "true"
).lower()
== "true",
"categorical_unique_ratio_threshold": float(
environ_vars.get("CATEGORICAL_UNIQUE_RATIO_THRESHOLD", "0.1")
),
"validate_fill_values": environ_vars.get("VALIDATE_FILL_VALUES", "true").lower()
== "true",
"column_strategies": {},
"exclude_columns": environ_vars.get("EXCLUDE_COLUMNS", "").split(",")
if environ_vars.get("EXCLUDE_COLUMNS")
else [],
}
# Parse column-specific strategies from environment variables
# Format: COLUMN_STRATEGY_<column_name>=<strategy>
for key, value in environ_vars.items():
if key.startswith("COLUMN_STRATEGY_"):
column_name = key.replace("COLUMN_STRATEGY_", "").lower()
config["column_strategies"][column_name] = value.lower()
return config
[docs]
def get_pandas_na_values() -> set:
"""
Get set of values that pandas interprets as NA/NULL.
"""
# Common pandas NA values to avoid
return {
"N/A",
"NA",
"NULL",
"NaN",
"nan",
"NAN",
"#N/A",
"#N/A N/A",
"#NA",
"-1.#IND",
"-1.#QNAN",
"-NaN",
"-nan",
"1.#IND",
"1.#QNAN",
"<NA>",
"null",
"Null",
"none",
"None",
"NONE",
}
[docs]
def validate_text_fill_value(value: str) -> bool:
"""
Validate that a text fill value won't be interpreted as NA by pandas.
"""
pandas_na_values = get_pandas_na_values()
return value not in pandas_na_values
[docs]
def detect_column_type(df: pd.DataFrame, column: str, config: Dict[str, Any]) -> str:
"""
Enhanced data type detection for imputation strategy selection.
"""
if pd.api.types.is_numeric_dtype(df[column]):
return "numerical"
elif pd.api.types.is_categorical_dtype(df[column]):
return "categorical"
elif df[column].dtype == "object":
if config.get("auto_detect_categorical", True):
# Distinguish between text and categorical based on unique values
non_null_count = df[column].dropna().shape[0]
if non_null_count > 0:
unique_ratio = df[column].nunique() / non_null_count
threshold = config.get("categorical_unique_ratio_threshold", 0.1)
if unique_ratio < threshold:
return "categorical"
return "text"
else:
return "text" # Default for other types
[docs]
class ImputationStrategyManager:
"""
Enhanced strategy manager supporting numerical, text, and categorical data types.
"""
def __init__(self, config: Dict[str, Any]):
self.config = config
self.pandas_na_values = get_pandas_na_values()
[docs]
def get_strategy_for_column(self, df: pd.DataFrame, column: str) -> SimpleImputer:
"""
Enhanced strategy selection supporting text and categorical types.
"""
# Detect column type using enhanced detection
column_type = detect_column_type(df, column, self.config)
# Check if strategy is explicitly configured
if column in self.config.get("column_strategies", {}):
strategy_name = self.config["column_strategies"][column]
return self._create_strategy_from_name(
df, column, column_type, strategy_name
)
# Auto-select based on detected type
if column_type == "numerical":
default_strategy = self.config.get("default_numerical_strategy", "mean")
elif column_type == "categorical":
default_strategy = self.config.get("default_categorical_strategy", "mode")
else: # text
default_strategy = self.config.get("default_text_strategy", "mode")
return self._create_strategy_from_name(
df, column, column_type, default_strategy
)
def _create_strategy_from_name(
self, df: pd.DataFrame, column: str, column_type: str, strategy_name: str
) -> SimpleImputer:
"""
Create appropriate imputation strategy based on column type and strategy name.
"""
if column_type == "numerical":
return self._create_numerical_strategy(strategy_name)
elif column_type == "categorical":
return self._create_categorical_strategy(df, column, strategy_name)
else: # text
return self._create_text_strategy(strategy_name)
def _create_numerical_strategy(self, strategy_name: str) -> SimpleImputer:
"""
Create numerical imputation strategy.
"""
if strategy_name == "mean":
return SimpleImputer(strategy="mean")
elif strategy_name == "median":
return SimpleImputer(strategy="median")
elif strategy_name == "constant":
fill_value = self.config.get("numerical_constant_value", 0)
return SimpleImputer(strategy="constant", fill_value=fill_value)
else:
logger.warning(f"Unknown numerical strategy '{strategy_name}', using mean")
return SimpleImputer(strategy="mean")
def _create_categorical_strategy(
self, df: pd.DataFrame, column: str, strategy_name: str
) -> SimpleImputer:
"""
Create categorical imputation strategy with dtype preservation.
"""
if strategy_name == "mode":
return SimpleImputer(strategy="most_frequent")
elif strategy_name == "constant":
fill_value = self.config.get("categorical_constant_value", "Unknown")
# Validate fill value is pandas-safe
if (
self.config.get("validate_fill_values", True)
and fill_value in self.pandas_na_values
):
logger.warning(
f"Categorical fill value '{fill_value}' may be interpreted as NA by pandas. Using 'Missing' instead."
)
fill_value = "Missing"
return SimpleImputer(strategy="constant", fill_value=fill_value)
else:
logger.warning(
f"Unknown categorical strategy '{strategy_name}', using mode"
)
return SimpleImputer(strategy="most_frequent")
def _create_text_strategy(self, strategy_name: str) -> SimpleImputer:
"""
Create text-specific imputation strategy with pandas-safe values.
"""
if strategy_name == "mode":
return SimpleImputer(strategy="most_frequent")
elif strategy_name == "constant":
fill_value = self.config.get("text_constant_value", "Unknown")
# Validate fill value is pandas-safe
if (
self.config.get("validate_fill_values", True)
and fill_value in self.pandas_na_values
):
logger.warning(
f"Text fill value '{fill_value}' may be interpreted as NA by pandas. Using 'Unknown' instead."
)
fill_value = "Unknown"
return SimpleImputer(strategy="constant", fill_value=fill_value)
elif strategy_name == "empty":
return SimpleImputer(strategy="constant", fill_value="")
else:
logger.warning(f"Unknown text strategy '{strategy_name}', using mode")
return SimpleImputer(strategy="most_frequent")
[docs]
class SimpleImputationEngine:
"""
Core engine for simple statistical imputation methods.
"""
def __init__(self, strategy_manager: ImputationStrategyManager, label_field: str):
self.strategy_manager = strategy_manager
self.label_field = label_field
self.fitted_imputers = {}
self.imputation_statistics = {}
[docs]
def fit(self, df: pd.DataFrame) -> None:
"""
Fit imputation parameters on training data.
"""
logger.info("Fitting imputation parameters on training data")
# Get columns to impute (exclude label and other specified columns)
exclude_cols = [self.label_field] + self.strategy_manager.config.get(
"exclude_columns", []
)
imputable_columns = [
col
for col in df.columns
if col not in exclude_cols and df[col].isnull().any()
]
logger.info(f"Columns to impute: {imputable_columns}")
for column in imputable_columns:
# Get appropriate strategy for this column
imputer = self.strategy_manager.get_strategy_for_column(df, column)
# Fit the imputer on non-null values
column_data = df[[column]]
imputer.fit(column_data)
# Store fitted imputer
self.fitted_imputers[column] = imputer
# Store imputation statistics
self.imputation_statistics[column] = {
"strategy": imputer.strategy,
"fill_value": getattr(imputer, "fill_value", None),
"statistics": getattr(imputer, "statistics_", None),
"missing_count_training": int(df[column].isnull().sum()),
"missing_percentage_training": float(
(df[column].isnull().sum() / len(df)) * 100
),
"data_type": str(df[column].dtype),
}
logger.info(f"Fitted imputer for column '{column}': {imputer.strategy}")
[docs]
def get_imputation_summary(self) -> Dict[str, Any]:
"""
Get comprehensive summary of imputation process.
"""
return {
"fitted_columns": list(self.fitted_imputers.keys()),
"imputation_statistics": self.imputation_statistics,
"last_transformation_log": getattr(self, "last_transformation_log", {}),
"total_imputers": len(self.fitted_imputers),
}
[docs]
def save_imputation_artifacts(
imputation_engine: SimpleImputationEngine,
imputation_config: Dict[str, Any],
output_path: Path,
) -> None:
"""
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.
Args:
imputation_engine: SimpleImputationEngine instance with fitted parameters
imputation_config: Imputation configuration dictionary
output_path: Path to save artifacts to
"""
# Extract simple imputation dictionary matching XGBoost training format
# Format: {column_name: imputation_value}
impute_dict = {}
for column, imputer in imputation_engine.fitted_imputers.items():
# Get the imputation value from the sklearn SimpleImputer
if hasattr(imputer, "statistics_") and imputer.statistics_ is not None:
# For mean/median/mode strategies, use statistics_
value = imputer.statistics_[0]
# Try to convert to float for numeric values, keep as-is for strings
try:
impute_dict[column] = float(value)
except (ValueError, TypeError):
# Keep string values as-is (e.g., categorical mode results)
impute_dict[column] = value
elif hasattr(imputer, "fill_value"):
# For constant strategy, use fill_value
impute_dict[column] = imputer.fill_value
else:
logger.warning(f"Could not extract imputation value for column {column}")
# Save imputation dictionary in XGBoost training format
params_output_path = output_path / IMPUTATION_PARAMS_FILENAME
with open(params_output_path, "wb") as f:
pkl.dump(impute_dict, f)
logger.info(f"Saved imputation dictionary to {params_output_path}")
logger.info(f"Format: {{{list(impute_dict.keys())[:3]}...}} -> values")
logger.info(f"This file can be used as input for non-training jobs")
# Save human-readable summary
summary = imputation_engine.get_imputation_summary()
summary_output_path = output_path / IMPUTATION_SUMMARY_FILENAME
with open(summary_output_path, "w") as f:
json.dump(summary, f, indent=2, default=str)
logger.info(f"Saved imputation summary to {summary_output_path}")
[docs]
def load_imputation_parameters(imputation_params_path: Path) -> Dict:
"""
Load imputation parameters from a pickle file.
Expected format (XGBoost training compatible):
Simple dict mapping column names to imputation values: {column: value}
Args:
imputation_params_path: Path to the imputation parameters file
Returns:
Dictionary of imputation parameters {column_name: imputation_value}
"""
if not imputation_params_path.exists():
raise FileNotFoundError(
f"Imputation parameters file not found: {imputation_params_path}"
)
logger.info(f"Loading imputation parameters from {imputation_params_path}")
with open(imputation_params_path, "rb") as f:
impute_dict = pkl.load(f)
if not isinstance(impute_dict, dict):
raise ValueError(f"Expected dict format, got {type(impute_dict)}")
logger.info(f"Loaded imputation parameters for {len(impute_dict)} columns")
return impute_dict
[docs]
def process_data(
data_dict: Dict[str, pd.DataFrame],
label_field: str,
job_type: str,
imputation_config: Dict[str, Any],
imputation_parameters: Optional[Dict] = None,
) -> Tuple[Dict[str, pd.DataFrame], SimpleImputationEngine]:
"""
Core data processing logic for missing value imputation.
Args:
data_dict: Dictionary of dataframes keyed by split name
label_field: Target column name
job_type: Type of job (training, validation, testing, calibration)
imputation_config: Imputation configuration dictionary
imputation_parameters: Pre-fitted imputation parameters (simple dict {column: value})
Returns:
Tuple containing:
- Dictionary of imputed dataframes
- SimpleImputationEngine instance with fitted parameters
"""
strategy_manager = ImputationStrategyManager(imputation_config)
imputation_engine = SimpleImputationEngine(strategy_manager, label_field)
if job_type == "training":
logger.info(
"Running in 'training' mode: fitting on train data, transforming all splits"
)
# Fit imputation parameters on training data only
imputation_engine.fit(data_dict["train"])
# Transform all splits
transformed_data = {}
for split_name, df in data_dict.items():
# Skip the format metadata key
if split_name == "_format":
transformed_data[split_name] = df # Preserve the format key
continue
df_imputed = imputation_engine.transform(df)
transformed_data[split_name] = df_imputed
logger.info(f"Imputed {split_name} data, shape: {df_imputed.shape}")
else:
# Non-training mode: use simple imputation dict {column: value}
if not imputation_parameters:
raise ValueError(
"For non-training job types, imputation_parameters must be provided"
)
logger.info(
f"Using pre-fitted imputation parameters for {len(imputation_parameters)} columns"
)
# Transform the data using simple fillna with the imputation dict
transformed_data = {}
for split_name, df in data_dict.items():
# Skip the format metadata key
if split_name == "_format":
transformed_data[split_name] = df # Preserve the format key
continue
df_imputed = df.copy()
for column, impute_value in imputation_parameters.items():
if column in df_imputed.columns:
# Only fill NaN values
df_imputed[column] = df_imputed[column].fillna(impute_value)
transformed_data[split_name] = df_imputed
logger.info(f"Imputed {split_name} data, shape: {df_imputed.shape}")
# Create a minimal engine for consistency (won't be used for transformation)
# This is just for returning a consistent interface
imputation_engine.imputation_statistics = {
col: {"strategy": "constant", "fill_value": val}
for col, val in imputation_parameters.items()
}
return transformed_data, imputation_engine
[docs]
def generate_imputation_report(
imputation_engine: SimpleImputationEngine,
missing_analysis: Dict[str, Any],
validation_report: Dict[str, Any],
output_dir: str,
) -> Dict[str, str]:
"""
Generate comprehensive imputation report with statistics and insights.
"""
# Get imputation summary
imputation_summary = imputation_engine.get_imputation_summary()
# Generate comprehensive report
report = {
"timestamp": datetime.utcnow().isoformat(),
"missing_value_analysis": missing_analysis,
"validation_report": validation_report,
"imputation_summary": imputation_summary,
"quality_metrics": calculate_imputation_quality_metrics(imputation_summary),
"recommendations": generate_imputation_recommendations(
imputation_summary, missing_analysis
),
}
# Save JSON report
json_path = os.path.join(output_dir, "imputation_report.json")
with open(json_path, "w") as f:
json.dump(report, f, indent=2, default=str)
# Generate text summary
text_summary = generate_imputation_text_summary(report)
text_path = os.path.join(output_dir, "imputation_summary.txt")
with open(text_path, "w") as f:
f.write(text_summary)
return {"json_report": json_path, "text_summary": text_path}
[docs]
def calculate_imputation_quality_metrics(
imputation_summary: Dict[str, Any],
) -> Dict[str, Any]:
"""
Calculate quality metrics for imputation process.
"""
quality_metrics = {
"total_columns_imputed": len(imputation_summary["fitted_columns"]),
"imputation_coverage": {},
"strategy_distribution": {},
"data_type_coverage": {},
}
# Calculate imputation coverage by column
for column, stats in imputation_summary["imputation_statistics"].items():
quality_metrics["imputation_coverage"][column] = {
"missing_percentage": stats["missing_percentage_training"],
"strategy_used": stats["strategy"],
"data_type": stats["data_type"],
}
# Calculate strategy distribution
strategies = [
stats["strategy"]
for stats in imputation_summary["imputation_statistics"].values()
]
strategy_counts = {}
for strategy in strategies:
strategy_counts[strategy] = strategy_counts.get(strategy, 0) + 1
quality_metrics["strategy_distribution"] = strategy_counts
# Calculate data type coverage
data_types = [
stats["data_type"]
for stats in imputation_summary["imputation_statistics"].values()
]
type_counts = {}
for dtype in data_types:
type_counts[dtype] = type_counts.get(dtype, 0) + 1
quality_metrics["data_type_coverage"] = type_counts
return quality_metrics
[docs]
def generate_imputation_recommendations(
imputation_summary: Dict[str, Any], missing_analysis: Dict[str, Any]
) -> List[str]:
"""
Generate actionable recommendations based on imputation analysis.
"""
recommendations = []
# Check for high missing value percentages
high_missing_columns = []
for column, stats in imputation_summary["imputation_statistics"].items():
if stats["missing_percentage_training"] > 50:
high_missing_columns.append(column)
if high_missing_columns:
recommendations.append(
f"Columns with >50% missing values detected: {high_missing_columns}. "
"Consider investigating data collection issues or using advanced imputation methods."
)
# Check strategy appropriateness
numerical_mode_columns = []
for column, stats in imputation_summary["imputation_statistics"].items():
if "int" in stats["data_type"] or "float" in stats["data_type"]:
if stats["strategy"] == "most_frequent":
numerical_mode_columns.append(column)
if numerical_mode_columns:
recommendations.append(
f"Numerical columns using mode imputation: {numerical_mode_columns}. "
"Consider using mean or median imputation for better statistical properties."
)
# Check for potential data quality issues
total_missing_patterns = missing_analysis["missing_patterns"][
"records_with_missing"
]
total_records = missing_analysis["total_records"]
missing_record_percentage = (total_missing_patterns / total_records) * 100
if missing_record_percentage > 30:
recommendations.append(
f"{missing_record_percentage:.1f}% of records have missing values. "
"Consider investigating systematic data collection issues."
)
# General recommendations
if len(imputation_summary["fitted_columns"]) > 10:
recommendations.append(
"Large number of columns require imputation. Consider feature selection "
"or advanced imputation methods like MICE for better performance."
)
return recommendations
[docs]
def copy_existing_artifacts(src_dir: str, dst_dir: str) -> None:
"""
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
Args:
src_dir: Source directory containing existing artifacts
dst_dir: Destination directory to copy artifacts to
"""
if not src_dir or not os.path.exists(src_dir):
logger.info(f"No existing artifacts to copy from {src_dir}")
return
os.makedirs(dst_dir, exist_ok=True)
copied_count = 0
for filename in os.listdir(src_dir):
src_file = os.path.join(src_dir, filename)
dst_file = os.path.join(dst_dir, filename)
if os.path.isfile(src_file):
shutil.copy2(src_file, dst_file)
copied_count += 1
logger.info(f" Copied existing artifact: {filename}")
logger.info(f"✓ Copied {copied_count} existing artifact(s) to {dst_dir}")
[docs]
def generate_imputation_text_summary(report: Dict[str, Any]) -> str:
"""
Generate human-readable text summary of imputation process.
"""
summary_lines = [
"=" * 60,
"MISSING VALUE IMPUTATION SUMMARY",
"=" * 60,
f"Generated: {report['timestamp']}",
"",
"DATA OVERVIEW:",
f" Total Records: {report['missing_value_analysis']['total_records']:,}",
f" Columns with Missing Values: {len(report['missing_value_analysis']['columns_with_missing'])}",
f" Records with Missing Values: {report['missing_value_analysis']['missing_patterns']['records_with_missing']:,}",
"",
"IMPUTATION RESULTS:",
f" Columns Imputed: {report['quality_metrics']['total_columns_imputed']}",
f" Strategy Distribution: {report['quality_metrics']['strategy_distribution']}",
"",
]
# Add column-specific details
if report["imputation_summary"]["imputation_statistics"]:
summary_lines.append("COLUMN DETAILS:")
for column, stats in report["imputation_summary"][
"imputation_statistics"
].items():
summary_lines.append(
f" {column}: {stats['strategy']} imputation, "
f"{stats['missing_percentage_training']:.1f}% missing"
)
summary_lines.append("")
# Add recommendations
if report["recommendations"]:
summary_lines.append("RECOMMENDATIONS:")
for i, rec in enumerate(report["recommendations"], 1):
summary_lines.append(f" {i}. {rec}")
summary_lines.append("")
summary_lines.append("=" * 60)
return "\n".join(summary_lines)
[docs]
def internal_main(
job_type: str,
input_dir: str,
output_dir: str,
imputation_config: Dict[str, Any],
label_field: str,
model_artifacts_input_dir: Optional[str] = None,
model_artifacts_output_dir: Optional[str] = None,
enable_true_streaming: bool = False,
max_workers: Optional[int] = None,
load_data_func: Callable = load_split_data,
save_data_func: Callable = save_output_data,
) -> Tuple[Dict[str, pd.DataFrame], SimpleImputationEngine]:
"""
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)
Args:
job_type: Type of job (training, validation, testing, calibration)
input_dir: Input directory for data
output_dir: Output directory for processed data
imputation_config: Imputation configuration dictionary
label_field: Target column name
model_artifacts_input_dir: Directory containing model artifacts from previous steps
model_artifacts_output_dir: Directory to save model artifacts for next steps
enable_true_streaming: 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: Number of parallel workers for streaming mode
load_data_func: Function to load data (for dependency injection in tests)
save_data_func: Function to save data (for dependency injection in tests)
Returns:
Tuple containing:
- Dictionary of imputed dataframes (empty in streaming mode)
- SimpleImputationEngine instance with fitted parameters (or None in streaming)
"""
output_path = Path(output_dir)
output_path.mkdir(parents=True, exist_ok=True)
logger.info(f"Using imputation configuration: {imputation_config}")
logger.info(f"Label field: {label_field}")
logger.info(f"Streaming mode: {'ENABLED' if enable_true_streaming else 'DISABLED'}")
# Determine model artifacts output directory
artifacts_output_dir = (
Path(model_artifacts_output_dir)
if model_artifacts_output_dir
else output_path / "model_artifacts"
)
artifacts_output_dir.mkdir(parents=True, exist_ok=True)
# Copy existing artifacts from previous steps (parameter accumulator pattern)
if model_artifacts_input_dir:
copy_existing_artifacts(model_artifacts_input_dir, str(artifacts_output_dir))
# ========================================================================
# STREAMING MODE
# ========================================================================
if enable_true_streaming:
logger.info("=" * 60)
logger.info("STREAMING MODE ENABLED")
logger.info("=" * 60)
# Call streaming mode orchestration (signature_columns=None, files have headers)
# Format is auto-detected from input shards (mirrors batch mode behavior)
stats = process_streaming_mode_imputation(
input_dir=input_dir,
output_dir=output_dir,
signature_columns=None, # Files from tabular_preprocessing have headers
job_type=job_type,
label_field=label_field,
imputation_config=imputation_config,
max_workers=max_workers,
model_artifacts_input_dir=model_artifacts_input_dir,
model_artifacts_output_dir=str(artifacts_output_dir),
logger=logger.info,
)
logger.info(f"Streaming mode complete! Final statistics: {stats}")
# Return empty data dict and None engine (data written to disk)
return {}, None
# ========================================================================
# BATCH MODE (DEFAULT)
# ========================================================================
logger.info("Running in BATCH MODE")
# Load data according to job type
data_dict = load_data_func(job_type, input_dir)
# Load imputation parameters if needed (non-training modes)
imputation_parameters = None
if job_type != "training" and model_artifacts_input_dir:
# Use the consistent filename for loading imputation parameters
imputation_params_path = (
Path(model_artifacts_input_dir) / IMPUTATION_PARAMS_FILENAME
)
if imputation_params_path.exists():
imputation_parameters = load_imputation_parameters(imputation_params_path)
logger.info(
f"Loaded pre-trained imputation parameters from {imputation_params_path}"
)
else:
logger.warning(
f"Imputation parameters not found at {imputation_params_path}"
)
# Process the data
transformed_data, imputation_engine = process_data(
data_dict=data_dict,
label_field=label_field,
job_type=job_type,
imputation_config=imputation_config,
imputation_parameters=imputation_parameters,
)
# Save processed data
save_data_func(job_type, output_dir, transformed_data)
# Save fitted artifacts (only for training jobs)
if job_type == "training":
save_imputation_artifacts(
imputation_engine, imputation_config, artifacts_output_dir
)
# Generate comprehensive report (only for training jobs)
if job_type == "training" and transformed_data:
sample_df = next(iter(transformed_data.values()))
missing_analysis = analyze_missing_values(sample_df)
validation_report = validate_imputation_data(sample_df, label_field)
generate_imputation_report(
imputation_engine, missing_analysis, validation_report, output_dir
)
logger.info("Generated imputation report for training job")
logger.info("Missing value imputation complete.")
return transformed_data, imputation_engine
[docs]
def main(
input_paths: Dict[str, str],
output_paths: Dict[str, str],
environ_vars: Dict[str, str],
job_args: Optional[argparse.Namespace] = None,
) -> Tuple[Dict[str, pd.DataFrame], SimpleImputationEngine]:
"""
Standardized main entry point for missing value imputation script.
Args:
input_paths: 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: 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: Dictionary of environment variables
job_args: Command line arguments containing job_type
Returns:
Tuple containing:
- Dictionary of imputed dataframes
- SimpleImputationEngine instance with fitted parameters
"""
try:
# Extract paths from input parameters - required keys must be present
if "input_data" not in input_paths:
raise ValueError("Missing required input path: input_data")
if "processed_data" not in output_paths:
raise ValueError("Missing required output path: processed_data")
# Extract job_type from args
if job_args is None or not hasattr(job_args, "job_type"):
raise ValueError("job_args must contain job_type parameter")
job_type = job_args.job_type
input_dir = input_paths["input_data"]
output_dir = output_paths["processed_data"]
# Get standardized model artifacts paths
model_artifacts_input_dir = input_paths.get("model_artifacts_input")
model_artifacts_output_dir = output_paths.get("model_artifacts_output")
# Log input/output paths for clarity
logger.info(f"Input data directory: {input_dir}")
logger.info(f"Output directory: {output_dir}")
if model_artifacts_input_dir:
logger.info(f"Model artifacts input directory: {model_artifacts_input_dir}")
logger.info(
f"Expected imputation parameters path: {Path(model_artifacts_input_dir) / IMPUTATION_PARAMS_FILENAME}"
)
if model_artifacts_output_dir:
logger.info(
f"Model artifacts output directory: {model_artifacts_output_dir}"
)
# Load imputation configuration from environment variables
imputation_config = load_imputation_config(environ_vars)
label_field = environ_vars.get("LABEL_FIELD", "target")
# Extract streaming mode configuration
enable_true_streaming = (
environ_vars.get("ENABLE_TRUE_STREAMING", "false").lower() == "true"
)
max_workers_str = environ_vars.get("MAX_WORKERS", "0")
max_workers = int(max_workers_str) if max_workers_str else 0
# Execute the internal main logic
return internal_main(
job_type=job_type,
input_dir=input_dir,
output_dir=output_dir,
imputation_config=imputation_config,
label_field=label_field,
model_artifacts_input_dir=model_artifacts_input_dir,
model_artifacts_output_dir=model_artifacts_output_dir,
enable_true_streaming=enable_true_streaming,
max_workers=max_workers,
)
except Exception as e:
logger.error(f"Error in missing value imputation: {str(e)}")
logger.error(traceback.format_exc())
raise
if __name__ == "__main__":
try:
parser = argparse.ArgumentParser()
parser.add_argument(
"--job_type",
type=str,
required=True,
choices=["training", "validation", "testing", "calibration"],
help="Type of job to perform",
)
args = parser.parse_args()
# Define standard SageMaker paths based on contract
# Separate data and model artifacts into different subfolders
input_paths = {
"input_data": DEFAULT_INPUT_DIR,
}
output_paths = {
"processed_data": DEFAULT_OUTPUT_DIR + "/data",
"model_artifacts_output": DEFAULT_OUTPUT_DIR + "/model_artifacts",
}
# For non-training jobs, add model artifacts input path
if args.job_type != "training":
input_paths["model_artifacts_input"] = DEFAULT_MODEL_ARTIFACTS_DIR
# Environment variables dictionary
environ_vars = {
"LABEL_FIELD": os.environ.get("LABEL_FIELD", "target"),
"DEFAULT_NUMERICAL_STRATEGY": os.environ.get(
"DEFAULT_NUMERICAL_STRATEGY", "mean"
),
"DEFAULT_CATEGORICAL_STRATEGY": os.environ.get(
"DEFAULT_CATEGORICAL_STRATEGY", "mode"
),
"DEFAULT_TEXT_STRATEGY": os.environ.get("DEFAULT_TEXT_STRATEGY", "mode"),
"NUMERICAL_CONSTANT_VALUE": os.environ.get("NUMERICAL_CONSTANT_VALUE", "0"),
"CATEGORICAL_CONSTANT_VALUE": os.environ.get(
"CATEGORICAL_CONSTANT_VALUE", "Unknown"
),
"TEXT_CONSTANT_VALUE": os.environ.get("TEXT_CONSTANT_VALUE", "Unknown"),
"CATEGORICAL_PRESERVE_DTYPE": os.environ.get(
"CATEGORICAL_PRESERVE_DTYPE", "true"
),
"AUTO_DETECT_CATEGORICAL": os.environ.get(
"AUTO_DETECT_CATEGORICAL", "true"
),
"CATEGORICAL_UNIQUE_RATIO_THRESHOLD": os.environ.get(
"CATEGORICAL_UNIQUE_RATIO_THRESHOLD", "0.1"
),
"VALIDATE_FILL_VALUES": os.environ.get("VALIDATE_FILL_VALUES", "true"),
"EXCLUDE_COLUMNS": os.environ.get("EXCLUDE_COLUMNS", ""),
# Streaming mode configuration
"ENABLE_TRUE_STREAMING": os.environ.get("ENABLE_TRUE_STREAMING", "false"),
"MAX_WORKERS": os.environ.get("MAX_WORKERS", "0"),
}
# Add column-specific strategies from environment variables
for key, value in os.environ.items():
if key.startswith("COLUMN_STRATEGY_"):
environ_vars[key] = value
# Execute the main function with standardized inputs
result, _ = main(input_paths, output_paths, environ_vars, args)
logger.info(f"Missing value imputation completed successfully")
sys.exit(0)
except FileNotFoundError as e:
logger.error(f"File not found error: {str(e)}")
sys.exit(1)
except ValueError as e:
logger.error(f"Value error: {str(e)}")
sys.exit(2)
except Exception as e:
logger.error(f"Error in missing value imputation script: {str(e)}")
logger.error(traceback.format_exc())
sys.exit(3)