Data Cleaning in Machine Learning

import pandas as pd
import numpy as np
from sklearn.impute import KNNImputer
from sklearn.ensemble import IsolationForest

def clean_dataframe(df: pd.DataFrame, config: dict) -> pd.DataFrame:
    """Production-grade data cleaning pipeline.
    
    Args:
        df: Raw input DataFrame
        config: Cleaning configuration dict
    Returns:
        Cleaned DataFrame with audit log
    """
    audit_log = []
    original_shape = df.shape
    
    # ---- Step 1: Remove exact duplicates ----
    n_dupes = df.duplicated(subset=config.get('dedup_columns')).sum()
    df = df.drop_duplicates(subset=config.get('dedup_columns'), keep='first')
    audit_log.append(f"Removed {n_dupes} exact duplicates")
    
    # ---- Step 2: Handle missing values per column ----
    for col, strategy in config.get('missing_strategies', {}).items():
        if col not in df.columns:
            continue
        n_missing = df[col].isna().sum()
        if strategy == 'mean':
            df[col] = df[col].fillna(df[col].mean())
        elif strategy == 'median':
            df[col] = df[col].fillna(df[col].median())
        elif strategy == 'mode':
            df[col] = df[col].fillna(df[col].mode()[0])
        elif strategy == 'drop':
            df = df.dropna(subset=[col])
        elif strategy == 'ffill':
            df[col] = df[col].ffill()
        elif strategy == 'constant':
            fill_val = config.get('fill_values', {}).get(col, 0)
            df[col] = df[col].fillna(fill_val)
        audit_log.append(f"Imputed {n_missing} missing values in '{col}' using {strategy}")
    
    # ---- Step 3: KNN imputation for remaining numeric NaNs ----
    numeric_cols = df.select_dtypes(include=[np.number]).columns.tolist()
    remaining_nans = df[numeric_cols].isna().sum().sum()
    if remaining_nans > 0:
        imputer = KNNImputer(n_neighbors=5, weights='distance')
        df[numeric_cols] = imputer.fit_transform(df[numeric_cols])
        audit_log.append(f"KNN-imputed {remaining_nans} remaining numeric NaNs")
    
    # ---- Step 4: Outlier detection with Isolation Forest ----
    outlier_cols = config.get('outlier_columns', numeric_cols)
    if len(outlier_cols) > 0:
        iso = IsolationForest(
            contamination=config.get('contamination', 0.05),
            random_state=42
        )
        outlier_mask = iso.fit_predict(df[outlier_cols]) == -1
        n_outliers = outlier_mask.sum()
        action = config.get('outlier_action', 'flag')
        if action == 'remove':
            df = df[~outlier_mask]
        elif action == 'flag':
            df['is_outlier'] = outlier_mask
        audit_log.append(f"Detected {n_outliers} outliers, action: {action}")
    
    # ---- Step 5: Text normalization ----
    text_cols = config.get('text_columns', [])
    for col in text_cols:
        if col in df.columns:
            df[col] = (
                df[col]
                .astype(str)
                .str.strip()
                .str.lower()
                .str.replace(r'\s+', ' ', regex=True)
            )
    
    # ---- Step 6: Type enforcement ----
    for col, dtype in config.get('enforce_types', {}).items():
        if col in df.columns:
            df[col] = pd.to_numeric(df[col], errors='coerce') if dtype == 'numeric' else df[col].astype(dtype)
    
    print(f"Cleaning complete: {original_shape} -> {df.shape}")
    for entry in audit_log:
        print(f"  - {entry}")
    
    return df


# --- Usage ---
config = {
    'dedup_columns': ['user_id', 'timestamp'],
    'missing_strategies': {
        'age': 'median',
        'income': 'mean',
        'city': 'mode',
        'notes': 'constant'
    },
    'fill_values': {'notes': 'N/A'},
    'outlier_columns': ['age', 'income', 'transaction_amount'],
    'contamination': 0.03,
    'outlier_action': 'flag',
    'text_columns': ['name', 'address', 'notes'],
    'enforce_types': {'age': 'numeric', 'income': 'numeric'}
}

df_clean = clean_dataframe(df_raw, config)

import pandas as pd
import numpy as np
from sklearn.experimental import enable_iterative_imputer  # noqa
from sklearn.impute import IterativeImputer
from sklearn.ensemble import RandomForestRegressor

def mice_impute(df: pd.DataFrame, numeric_cols: list, n_iterations: int = 10) -> pd.DataFrame:
    """Apply MICE imputation using Random Forest as the estimator.
    
    MICE iteratively models each feature with missing values as a function
    of the other features, cycling through all incomplete features until
    convergence.
    
    Args:
        df: DataFrame with missing values
        numeric_cols: Columns to impute
        n_iterations: Number of MICE iterations (default 10)
    Returns:
        DataFrame with imputed values
    """
    # Store original missing mask for auditing
    missing_mask = df[numeric_cols].isna()
    missing_rate = missing_mask.mean()
    print("Missing rates before imputation:")
    print(missing_rate[missing_rate > 0].to_string())
    
    # Configure MICE with MissForest-style estimator
    mice = IterativeImputer(
        estimator=RandomForestRegressor(
            n_estimators=100,
            max_depth=10,
            random_state=42,
            n_jobs=-1
        ),
        max_iter=n_iterations,
        random_state=42,
        verbose=1
    )
    
    # Fit and transform
    imputed_values = mice.fit_transform(df[numeric_cols])
    df_imputed = df.copy()
    df_imputed[numeric_cols] = imputed_values
    
    # Audit: compare imputed distributions
    for col in numeric_cols:
        if missing_mask[col].any():
            orig_mean = df[col].mean()
            imputed_mean = df_imputed[col].mean()
            drift = abs(imputed_mean - orig_mean) / (orig_mean + 1e-8)
            print(f"{col}: original mean={orig_mean:.3f}, "
                  f"imputed mean={imputed_mean:.3f}, drift={drift:.2%}")
    
    return df_imputed


# --- Usage ---
numeric_features = ['age', 'income', 'credit_score', 'tenure_months']
df_clean = mice_impute(df_raw, numeric_features, n_iterations=10)

import pandas as pd
import recordlinkage
from recordlinkage.preprocessing import clean as rl_clean

def fuzzy_dedup(df: pd.DataFrame, match_cols: list, threshold: float = 0.85) -> pd.DataFrame:
    """Fuzzy deduplication using record linkage with blocking.
    
    Uses sorted neighborhood blocking to avoid O(n^2) comparisons,
    then applies string similarity scoring to identify near-duplicates.
    
    Args:
        df: DataFrame with potential duplicates
        match_cols: Columns to use for matching
        threshold: Similarity threshold (0-1) for considering a match
    Returns:
        Deduplicated DataFrame
    """
    # Preprocess: clean and standardize match columns
    df_clean = df.copy()
    for col in match_cols:
        if df_clean[col].dtype == 'object':
            df_clean[f"{col}_clean"] = rl_clean(df_clean[col])
    
    clean_cols = [f"{c}_clean" if df_clean[c].dtype == 'object' else c for c in match_cols]
    
    # Build candidate pairs using sorted neighborhood blocking
    # This reduces O(n^2) to approximately O(n * window_size)
    indexer = recordlinkage.Index()
    indexer.sortedneighbourhood(left_on=clean_cols[0], window=5)
    candidate_pairs = indexer.index(df_clean)
    print(f"Generated {len(candidate_pairs)} candidate pairs from {len(df)} records")
    
    # Compare candidate pairs
    comparer = recordlinkage.Compare()
    for col in clean_cols:
        if df_clean[col].dtype == 'object' or col.endswith('_clean'):
            comparer.string(col, col, method='jarowinkler', label=col)
        else:
            comparer.numeric(col, col, label=col)
    
    features = comparer.compute(candidate_pairs, df_clean)
    
    # Score and threshold
    features['avg_score'] = features.mean(axis=1)
    matches = features[features['avg_score'] >= threshold]
    print(f"Found {len(matches)} duplicate pairs above threshold {threshold}")
    
    # Resolve: keep first occurrence, mark rest for removal
    to_remove = set()
    for idx_a, idx_b in matches.index:
        if idx_a not in to_remove:
            to_remove.add(idx_b)
    
    df_deduped = df.drop(index=list(to_remove))
    print(f"Removed {len(to_remove)} fuzzy duplicates: {len(df)} -> {len(df_deduped)}")
    
    return df_deduped


# --- Usage: Dedup an Indian e-commerce product catalog ---
df_products = pd.DataFrame({
    'product_name': ['Samsung Galaxy S24 Ultra', 'samsung galaxy s24ultra',
                     'iPhone 15 Pro Max', 'Apple iPhone 15 ProMax 256GB'],
    'brand': ['Samsung', 'samsung', 'Apple', 'Apple'],
    'price': [1_29_999, 1_29_999, 1_59_900, 1_59_900]
})

df_clean = fuzzy_dedup(df_products, match_cols=['product_name', 'brand'], threshold=0.80)

import great_expectations as gx

def validate_cleaned_data(df, suite_name="post_cleaning_suite"):
    """Validate cleaned data using Great Expectations.
    
    Defines expectations for data quality after cleaning
    and returns a validation report.
    """
    context = gx.get_context()
    
    # Create a data source from our DataFrame
    data_source = context.data_sources.add_pandas(name="cleaned_data")
    data_asset = data_source.add_dataframe_asset(name="df")
    batch_definition = data_asset.add_batch_definition_whole_dataframe("batch")
    batch = batch_definition.get_batch(batch_parameters={"dataframe": df})
    
    # Create expectation suite
    suite = context.suites.add(
        gx.ExpectationSuite(name=suite_name)
    )
    
    # ---- Define expectations ----
    # No nulls in critical columns
    suite.add_expectation(
        gx.expectations.ExpectColumnValuesToNotBeNull(column="user_id")
    )
    suite.add_expectation(
        gx.expectations.ExpectColumnValuesToNotBeNull(column="timestamp")
    )
    
    # Value ranges
    suite.add_expectation(
        gx.expectations.ExpectColumnValuesToBeBetween(
            column="age", min_value=0, max_value=120
        )
    )
    suite.add_expectation(
        gx.expectations.ExpectColumnValuesToBeBetween(
            column="transaction_amount", min_value=0, max_value=10_00_000  # INR 10 lakh
        )
    )
    
    # Uniqueness
    suite.add_expectation(
        gx.expectations.ExpectColumnValuesToBeUnique(column="user_id")
    )
    
    # Column completeness threshold
    suite.add_expectation(
        gx.expectations.ExpectColumnValuesToNotBeNull(
            column="email", mostly=0.95  # Allow 5% nulls
        )
    )
    
    # Run validation
    validation_definition = context.validation_definitions.add(
        gx.ValidationDefinition(
            name="post_clean_validation",
            data=batch_definition,
            suite=suite
        )
    )
    result = validation_definition.run(batch_parameters={"dataframe": df})
    
    if result.success:
        print("All quality checks PASSED")
    else:
        print("Quality checks FAILED:")
        for r in result.results:
            if not r.success:
                print(f"  FAIL: {r.expectation_config.type} "
                      f"on '{r.expectation_config.kwargs.get('column', 'N/A')}'")
    
    return result


# --- Usage ---
result = validate_cleaned_data(df_clean)

# data_cleaning_config.yaml
# Configuration for the data cleaning pipeline

pipeline:
  name: user_transactions_cleaning
  version: "2.1.0"
  schedule: "0 */4 * * *"  # Every 4 hours

profiling:
  tool: ydata-profiling
  output_path: s3://data-quality/profiles/
  alert_on:
    missing_rate_threshold: 0.30
    duplicate_rate_threshold: 0.05

missing_values:
  default_strategy: median
  per_column:
    user_id: drop          # Cannot impute identifiers
    age: knn               # Correlated with other demographics
    income: mice            # Complex dependencies
    city: mode              # Categorical
    email: constant         # Fill with '[email protected]'
    transaction_amount: median
  knn_params:
    n_neighbors: 5
    weights: distance
  mice_params:
    max_iter: 10
    estimator: random_forest

deduplication:
  exact:
    columns: [user_id, timestamp]
    keep: first
  fuzzy:
    enabled: true
    columns: [name, email, phone]
    method: jaro_winkler
    threshold: 0.88
    blocking_key: city     # Only compare within same city

outliers:
  method: isolation_forest
  contamination: 0.03
  columns: [age, income, transaction_amount]
  action: flag             # flag | clip | remove
  clip_bounds:
    age: [0, 120]
    income: [0, 50000000]  # INR 5 crore
    transaction_amount: [0, 1000000]  # INR 10 lakh

normalization:
  text_columns:
    - name: [strip, lower, unicode_nfc]
    - address: [strip, lower, normalize_whitespace]
    - city: [strip, title_case]
  date_columns:
    - timestamp: ISO8601
    - created_at: ISO8601
  phone_columns:
    - phone: E164_IN       # Indian phone format

validation:
  tool: great_expectations
  suite: post_cleaning_v2
  fail_action: quarantine  # quarantine | warn | block
  expectations:
    - column: user_id
      type: not_null
    - column: age
      type: between
      min: 0
      max: 120
    - column: email
      type: not_null
      mostly: 0.95