ds-agent-cli 0.1.0

package/src/tools/data_cleaning.py

@@ -0,0 +1,614 @@
+"""
+Data Cleaning Tools
+Tools for handling missing values, outliers, and data type issues.
+"""
+
+import polars as pl
+import numpy as np
+from typing import Dict, Any, List, Optional
+from pathlib import Path
+import sys
+import os
+
+# Add parent directory to path for imports
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from ds_agent.utils.polars_helpers import (
+    load_dataframe,
+    save_dataframe,
+    get_numeric_columns,
+    get_categorical_columns,
+    get_datetime_columns,
+    detect_id_columns,
+)
+from ds_agent.utils.validation import (
+    validate_file_exists,
+    validate_file_format,
+    validate_dataframe,
+    validate_columns_exist,
+)
+
+
+def clean_missing_values(file_path: str, strategy, 
+                        output_path: str, threshold: float = 0.4) -> Dict[str, Any]:
+    """
+    Handle missing values using appropriate strategies with smart threshold-based column dropping.
+    
+    Args:
+        file_path: Path to CSV or Parquet file
+        strategy: Either "auto" (string) to automatically decide strategies for all columns,
+                 or a dictionary mapping column names to strategies 
+                 ('median', 'mean', 'mode', 'forward_fill', 'drop')
+        output_path: Path to save cleaned dataset
+        threshold: For "auto" strategy, drop columns with missing % > threshold (default: 0.4 = 40%)
+        
+    Returns:
+        Dictionary with cleaning report
+        
+    Auto Strategy Behavior:
+        1. Drop columns with >threshold missing (default 40%)
+        2. Impute numeric columns with median
+        3. Impute categorical columns with mode
+        4. Forward-fill for time series columns
+    """
+    # Validation
+    validate_file_exists(file_path)
+    validate_file_format(file_path)
+    
+    # Load data
+    df = load_dataframe(file_path)
+    validate_dataframe(df)
+    
+    # Get column type information
+    numeric_cols = get_numeric_columns(df)
+    categorical_cols = get_categorical_columns(df)
+    datetime_cols = get_datetime_columns(df)
+    id_cols = detect_id_columns(df)
+    
+    report = {
+        "original_rows": len(df),
+        "original_columns": len(df.columns),
+        "columns_dropped": [],
+        "columns_processed": {},
+        "rows_dropped": 0,
+        "threshold_used": threshold
+    }
+    
+    # Handle string strategy modes
+    if isinstance(strategy, str):
+        if strategy == "auto":
+            # Step 1: Identify and drop high-missing columns (>threshold)
+            cols_to_drop = []
+            for col in df.columns:
+                null_count = df[col].null_count()
+                null_pct = null_count / len(df) if len(df) > 0 else 0
+                
+                if null_pct > threshold:
+                    cols_to_drop.append(col)
+                    report["columns_dropped"].append({
+                        "column": col,
+                        "missing_percentage": round(null_pct * 100, 2),
+                        "reason": f"Missing >{threshold*100}% of values"
+                    })
+            
+            # Drop high-missing columns
+            if cols_to_drop:
+                df = df.drop(cols_to_drop)
+                print(f"🗑️  Dropped {len(cols_to_drop)} columns with >{threshold*100}% missing:")
+                for col_info in report["columns_dropped"]:
+                    print(f"    - {col_info['column']} ({col_info['missing_percentage']}% missing)")
+            
+            # Step 2: Build strategy for remaining columns
+            strategy = {}
+            for col in df.columns:
+                if df[col].null_count() > 0:
+                    if col in id_cols:
+                        strategy[col] = "drop"  # Drop rows with missing IDs
+                    elif col in datetime_cols:
+                        strategy[col] = "forward_fill"  # Forward fill for time series
+                    elif col in numeric_cols:
+                        strategy[col] = "median"  # Median for numeric (robust to outliers)
+                    elif col in categorical_cols:
+                        strategy[col] = "mode"  # Mode for categorical
+                    else:
+                        strategy[col] = "mode"  # Default to mode
+            
+            print(f"🔧 Auto-detected strategies for {len(strategy)} remaining columns with missing values")
+        
+        elif strategy in ["median", "mean", "mode", "forward_fill", "drop"]:
+            # Apply same strategy to all columns with missing values
+            strategy_dict = {}
+            for col in df.columns:
+                if df[col].null_count() > 0:
+                    strategy_dict[col] = strategy
+            strategy = strategy_dict
+            print(f"🔧 Applying '{list(strategy_dict.values())[0] if strategy_dict else strategy}' strategy to {len(strategy_dict)} columns with missing values")
+        
+        elif strategy in ["iterative", "mice"]:
+            # MICE / Iterative Imputation using sklearn IterativeImputer
+            # This handles ALL numeric columns at once (multivariate imputation)
+            print(f"🔧 Applying Iterative (MICE) imputation to numeric columns...")
+            try:
+                from sklearn.experimental import enable_iterative_imputer  # noqa: F401
+                from sklearn.impute import IterativeImputer
+                from sklearn.linear_model import BayesianRidge
+                import pandas as pd
+                
+                # Identify numeric columns with missing values
+                numeric_cols_with_nulls = [
+                    col for col in numeric_cols if df[col].null_count() > 0
+                ]
+                
+                if not numeric_cols_with_nulls:
+                    print("   ℹ️ No numeric columns with missing values for MICE imputation")
+                else:
+                    # Convert numeric columns to pandas for IterativeImputer
+                    df_pd = df.select(numeric_cols).to_pandas()
+                    
+                    # Fit and transform
+                    imputer = IterativeImputer(
+                        estimator=BayesianRidge(),
+                        max_iter=10,
+                        random_state=42,
+                        missing_values=float('nan')
+                    )
+                    imputed_data = imputer.fit_transform(df_pd)
+                    
+                    # Replace columns back in Polars DataFrame
+                    for i, col_name in enumerate(numeric_cols):
+                        df = df.with_columns(
+                            pl.Series(col_name, imputed_data[:, i])
+                        )
+                    
+                    for col_name in numeric_cols_with_nulls:
+                        report["columns_processed"][col_name] = {
+                            "status": "success",
+                            "strategy": "iterative_mice",
+                            "nulls_before": int(df[col_name].null_count()),  # Should be 0 now
+                            "nulls_after": 0
+                        }
+                    
+                    print(f"   ✅ MICE imputed {len(numeric_cols_with_nulls)} numeric columns using {len(numeric_cols)} features")
+                
+                # Handle remaining non-numeric columns with mode
+                for col in df.columns:
+                    if df[col].null_count() > 0 and col not in numeric_cols:
+                        mode_val = df[col].drop_nulls().mode().first()
+                        if mode_val is not None:
+                            df = df.with_columns(
+                                pl.col(col).fill_null(mode_val).alias(col)
+                            )
+                            report["columns_processed"][col] = {
+                                "status": "success",
+                                "strategy": "mode (non-numeric fallback)",
+                                "nulls_before": int(df[col].null_count()),
+                                "nulls_after": 0
+                            }
+                
+            except ImportError:
+                return {
+                    "success": False,
+                    "error": "IterativeImputer requires scikit-learn >= 1.4. Install with: pip install scikit-learn>=1.4",
+                    "error_type": "MissingDependency"
+                }
+            
+            # Skip per-column processing for MICE (already handled above)
+            strategy = {}
+        
+        else:
+            return {
+                "success": False,
+                "error": f"Invalid strategy '{strategy}'. Use 'auto', 'median', 'mean', 'mode', 'forward_fill', 'drop', 'iterative', 'mice', or provide a dictionary.",
+                "error_type": "ValueError"
+            }
+    
+    # Process each column based on strategy
+    for col, strat in strategy.items():
+        if col not in df.columns:
+            report["columns_processed"][col] = {
+                "status": "error",
+                "message": f"Column not found (may have been dropped)"
+            }
+            continue
+        
+        null_count_before = df[col].null_count()
+        
+        if null_count_before == 0:
+            report["columns_processed"][col] = {
+                "status": "skipped",
+                "message": "No missing values"
+            }
+            continue
+        
+        # Don't impute ID columns - drop rows instead
+        if col in id_cols and strat != "drop":
+            report["columns_processed"][col] = {
+                "status": "skipped",
+                "message": "ID column - not imputed (use 'drop' to remove rows)"
+            }
+            continue
+        
+        # Apply strategy
+        try:
+            rows_before = len(df)
+            
+            if strat == "median":
+                if col in numeric_cols:
+                    median_val = df[col].median()
+                    df = df.with_columns(
+                        pl.col(col).fill_null(median_val).alias(col)
+                    )
+                    report["columns_processed"][col] = {
+                        "status": "success",
+                        "strategy": "median",
+                        "nulls_before": int(null_count_before),
+                        "nulls_after": int(df[col].null_count()),
+                        "fill_value": float(median_val)
+                    }
+                else:
+                    report["columns_processed"][col] = {
+                        "status": "error",
+                        "message": "Cannot use median on non-numeric column"
+                    }
+                    continue
+            
+            elif strat == "mean":
+                if col in numeric_cols:
+                    mean_val = df[col].mean()
+                    df = df.with_columns(
+                        pl.col(col).fill_null(mean_val).alias(col)
+                    )
+                    report["columns_processed"][col] = {
+                        "status": "success",
+                        "strategy": "mean",
+                        "nulls_before": int(null_count_before),
+                        "nulls_after": int(df[col].null_count()),
+                        "fill_value": float(mean_val)
+                    }
+                else:
+                    report["columns_processed"][col] = {
+                        "status": "error",
+                        "message": "Cannot use mean on non-numeric column"
+                    }
+                    continue
+            
+            elif strat == "mode":
+                mode_val = df[col].drop_nulls().mode().first()
+                if mode_val is not None:
+                    df = df.with_columns(
+                        pl.col(col).fill_null(mode_val).alias(col)
+                    )
+                    report["columns_processed"][col] = {
+                        "status": "success",
+                        "strategy": "mode",
+                        "nulls_before": int(null_count_before),
+                        "nulls_after": int(df[col].null_count()),
+                        "fill_value": str(mode_val)
+                    }
+            
+            elif strat == "forward_fill":
+                df = df.with_columns(
+                    pl.col(col).forward_fill().alias(col)
+                )
+                report["columns_processed"][col] = {
+                    "status": "success",
+                    "strategy": "forward_fill",
+                    "nulls_before": int(null_count_before),
+                    "nulls_after": int(df[col].null_count())
+                }
+            
+            elif strat == "drop":
+                df = df.filter(pl.col(col).is_not_null())
+                rows_after = len(df)
+                report["columns_processed"][col] = {
+                    "status": "success",
+                    "strategy": "drop",
+                    "nulls_before": int(null_count_before),
+                    "rows_dropped": rows_before - rows_after
+                }
+            
+            else:
+                report["columns_processed"][col] = {
+                    "status": "error",
+                    "message": f"Unknown strategy: {strat}"
+                }
+                continue
+        
+        except Exception as e:
+            report["columns_processed"][col] = {
+                "status": "error",
+                "message": str(e)
+            }
+    
+    report["final_rows"] = len(df)
+    report["final_columns"] = len(df.columns)
+    report["rows_dropped"] = report["original_rows"] - report["final_rows"]
+    report["columns_dropped_count"] = len(report["columns_dropped"])
+    
+    # Save cleaned dataset
+    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
+    save_dataframe(df, output_path)
+    report["output_path"] = output_path
+    
+    # Summary message
+    report["message"] = f"Cleaned {report['original_rows']} rows → {report['final_rows']} rows. "
+    report["message"] += f"Dropped {report['columns_dropped_count']} columns. "
+    report["message"] += f"Processed {len([c for c in report['columns_processed'].values() if c['status'] == 'success'])} columns."
+    
+    return report
+
+
+def handle_outliers(file_path: str, strategy: str, columns: List[str], 
+                   output_path: str) -> Dict[str, Any]:
+    """
+    Detect and handle outliers in numeric columns.
+    
+    Args:
+        file_path: Path to CSV or Parquet file
+        strategy: Method to handle outliers ('clip', 'cap', 'winsorize', 'remove')
+        columns: List of columns to check, or ['all'] for all numeric columns
+        output_path: Path to save cleaned dataset
+        
+    Returns:
+        Dictionary with outlier handling report
+    """
+    # Validation
+    validate_file_exists(file_path)
+    validate_file_format(file_path)
+    
+    # Load data
+    df = load_dataframe(file_path)
+    validate_dataframe(df)
+    
+    # Determine which columns to process
+    numeric_cols = get_numeric_columns(df)
+    
+    if columns == ["all"]:
+        target_cols = numeric_cols
+    else:
+        # Filter to only existing numeric columns (auto-skip dropped columns)
+        target_cols = []
+        for col in columns:
+            if col not in df.columns:
+                print(f"⚠️  Skipping '{col}' - column was dropped in previous step")
+                continue
+            if col not in numeric_cols:
+                print(f"⚠️  Skipping '{col}' - not numeric")
+                continue
+            target_cols.append(col)
+        
+        # If no valid columns remain, return early
+        if not target_cols:
+            return {
+                "success": False,
+                "error": f"None of the requested columns exist in the dataset. Available numeric columns: {', '.join(numeric_cols[:20])}",
+                "error_type": "ValueError"
+            }
+    
+    report = {
+        "original_rows": len(df),
+        "strategy": strategy,
+        "columns_processed": {}
+    }
+    
+    # Process each column
+    for col in target_cols:
+        col_data = df[col].drop_nulls()
+        
+        if len(col_data) == 0:
+            report["columns_processed"][col] = {
+                "status": "skipped",
+                "message": "All values are null"
+            }
+            continue
+        
+        # Calculate IQR bounds
+        q1 = col_data.quantile(0.25)
+        q3 = col_data.quantile(0.75)
+        iqr = q3 - q1
+        
+        lower_bound = q1 - 1.5 * iqr
+        upper_bound = q3 + 1.5 * iqr
+        
+        # Count outliers
+        outliers_mask = (df[col] < lower_bound) | (df[col] > upper_bound)
+        outlier_count = outliers_mask.sum()
+        
+        if outlier_count == 0:
+            report["columns_processed"][col] = {
+                "status": "skipped",
+                "message": "No outliers detected"
+            }
+            continue
+        
+        # Apply strategy
+        if strategy == "clip" or strategy == "cap":
+            # Clip/cap values to bounds
+            df = df.with_columns(
+                pl.col(col).clip(lower_bound, upper_bound).alias(col)
+            )
+        
+        elif strategy == "winsorize":
+            # Winsorize: cap at 1st and 99th percentiles
+            p1 = col_data.quantile(0.01)
+            p99 = col_data.quantile(0.99)
+            df = df.with_columns(
+                pl.col(col).clip(p1, p99).alias(col)
+            )
+        
+        elif strategy == "remove":
+            # Remove rows with outliers
+            df = df.filter(~outliers_mask)
+        
+        report["columns_processed"][col] = {
+            "status": "success",
+            "outliers_detected": int(outlier_count),
+            "bounds": {
+                "lower": float(lower_bound),
+                "upper": float(upper_bound)
+            }
+        }
+    
+    report["final_rows"] = len(df)
+    report["rows_dropped"] = report["original_rows"] - report["final_rows"]
+    
+    # Save cleaned dataset
+    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
+    save_dataframe(df, output_path)
+    report["output_path"] = output_path
+    
+    return report
+
+
+def fix_data_types(file_path: str, type_mapping: Optional[Dict[str, str]] = None,
+                  output_path: str = None) -> Dict[str, Any]:
+    """
+    Auto-detect and fix incorrect data types.
+    
+    Args:
+        file_path: Path to CSV or Parquet file
+        type_mapping: Optional dictionary mapping columns to target types
+                     ('int', 'float', 'string', 'date', 'bool', 'category')
+                     Use 'auto' or None for automatic detection
+        output_path: Path to save dataset with fixed types
+        
+    Returns:
+        Dictionary with type fixing report
+    """
+    # Validation
+    validate_file_exists(file_path)
+    validate_file_format(file_path)
+    
+    # Load data
+    df = load_dataframe(file_path)
+    validate_dataframe(df)
+    
+    if type_mapping is None or type_mapping == {"auto": "auto"}:
+        type_mapping = {}
+    
+    report = {
+        "columns_processed": {}
+    }
+    
+    for col in df.columns:
+        original_dtype = str(df[col].dtype)
+        
+        # Get target type from mapping or auto-detect
+        if col in type_mapping and type_mapping[col] != "auto":
+            target_type = type_mapping[col]
+        else:
+            # Auto-detect target type
+            target_type = _auto_detect_type(df[col])
+        
+        if target_type is None:
+            report["columns_processed"][col] = {
+                "status": "skipped",
+                "original_dtype": original_dtype,
+                "message": "Could not auto-detect type"
+            }
+            continue
+        
+        # Try to convert
+        try:
+            if target_type == "int":
+                df = df.with_columns(
+                    pl.col(col).cast(pl.Int64, strict=False).alias(col)
+                )
+            elif target_type == "float":
+                df = df.with_columns(
+                    pl.col(col).cast(pl.Float64, strict=False).alias(col)
+                )
+            elif target_type == "string":
+                df = df.with_columns(
+                    pl.col(col).cast(pl.Utf8).alias(col)
+                )
+            elif target_type == "date":
+                df = df.with_columns(
+                    pl.col(col).str.strptime(pl.Date, "%Y-%m-%d", strict=False).alias(col)
+                )
+            elif target_type == "bool":
+                df = df.with_columns(
+                    pl.col(col).cast(pl.Boolean, strict=False).alias(col)
+                )
+            elif target_type == "category":
+                df = df.with_columns(
+                    pl.col(col).cast(pl.Categorical).alias(col)
+                )
+            
+            new_dtype = str(df[col].dtype)
+            
+            report["columns_processed"][col] = {
+                "status": "success",
+                "original_dtype": original_dtype,
+                "new_dtype": new_dtype,
+                "target_type": target_type
+            }
+        
+        except Exception as e:
+            report["columns_processed"][col] = {
+                "status": "error",
+                "original_dtype": original_dtype,
+                "target_type": target_type,
+                "message": str(e)
+            }
+    
+    # Save dataset
+    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
+    save_dataframe(df, output_path)
+    report["output_path"] = output_path
+    
+    return report
+
+
+def _auto_detect_type(series: pl.Series) -> Optional[str]:
+    """
+    Auto-detect appropriate type for a series.
+    
+    Args:
+        series: Polars series
+        
+    Returns:
+        Detected type string or None
+    """
+    # Already correct type
+    if series.dtype in pl.NUMERIC_DTYPES:
+        return None
+    
+    if series.dtype in [pl.Date, pl.Datetime]:
+        return None
+    
+    # Try to detect from string values
+    if series.dtype == pl.Utf8:
+        sample = series.drop_nulls().head(100)
+        
+        if len(sample) == 0:
+            return None
+        
+        # Check for boolean
+        unique_vals = set(str(v).lower() for v in sample.to_list())
+        if unique_vals.issubset({'true', 'false', '1', '0', 'yes', 'no', 't', 'f'}):
+            return "bool"
+        
+        # Check for numeric
+        try:
+            sample.cast(pl.Float64)
+            # Check if all are integers
+            if all('.' not in str(v) for v in sample.to_list() if v is not None):
+                return "int"
+            return "float"
+        except:
+            pass
+        
+        # Check for date
+        try:
+            sample.str.strptime(pl.Date, "%Y-%m-%d", strict=False)
+            return "date"
+        except:
+            pass
+        
+        # Check if should be categorical (low cardinality)
+        n_unique = series.n_unique()
+        if n_unique < len(series) * 0.5 and n_unique < 100:
+            return "category"
+    
+    return None