ds-agent-cli 0.1.0

package/src/tools/data_profiling.py

@@ -0,0 +1,593 @@
+"""
+Data Profiling Tools
+Tools for analyzing and understanding dataset characteristics.
+"""
+
+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,
+    get_numeric_columns,
+    get_categorical_columns,
+    get_datetime_columns,
+    get_column_info,
+    calculate_memory_usage,
+    detect_id_columns,
+)
+from ds_agent.utils.validation import (
+    validate_file_exists,
+    validate_file_format,
+    validate_dataframe,
+)
+
+
+def profile_dataset(file_path: str) -> Dict[str, Any]:
+    """
+    Get comprehensive statistics about a dataset.
+    
+    Args:
+        file_path: Path to CSV or Parquet file
+        
+    Returns:
+        Dictionary with dataset profile including:
+        - shape (rows, columns)
+        - column types
+        - memory usage
+        - null counts
+        - unique values
+        - missing value percentage per column (NEW)
+        - unique value counts per column (NEW)
+        - basic statistics for each column
+    """
+    # Validation
+    validate_file_exists(file_path)
+    validate_file_format(file_path)
+    
+    # Load data
+    df = load_dataframe(file_path)
+    validate_dataframe(df)
+    
+    # Basic info
+    profile = {
+        "file_path": file_path,
+        "shape": {
+            "rows": len(df),
+            "columns": len(df.columns)
+        },
+        "memory_usage": calculate_memory_usage(df),
+        "column_types": {
+            "numeric": get_numeric_columns(df),
+            "categorical": get_categorical_columns(df),
+            "datetime": get_datetime_columns(df),
+            "id_columns": detect_id_columns(df),
+        },
+        "columns": {},
+        "missing_values_per_column": {},  # NEW: Per-column missing %
+        "unique_counts_per_column": {}   # NEW: Per-column unique counts
+    }
+    
+    # Per-column statistics with enhanced missing % and unique counts
+    for col in df.columns:
+        # Get existing column info
+        profile["columns"][col] = get_column_info(df, col)
+        
+        # NEW: Calculate missing value percentage for this column
+        null_count = df[col].null_count()
+        missing_pct = round((null_count / len(df)) * 100, 2) if len(df) > 0 else 0
+        profile["missing_values_per_column"][col] = {
+            "count": int(null_count),
+            "percentage": missing_pct
+        }
+        
+        # NEW: Calculate unique value counts (with dict handling)
+        try:
+            # Try to get unique count directly
+            unique_count = df[col].n_unique()
+            profile["unique_counts_per_column"][col] = int(unique_count)
+        except Exception as e:
+            # If column contains unhashable types (dicts, lists), handle gracefully
+            try:
+                # Convert to string and then count unique
+                unique_count = df[col].cast(pl.Utf8).n_unique()
+                profile["unique_counts_per_column"][col] = int(unique_count)
+            except:
+                profile["unique_counts_per_column"][col] = "N/A (unhashable type)"
+    
+    # Overall statistics
+    total_nulls = sum(df[col].null_count() for col in df.columns)
+    total_cells = len(df) * len(df.columns)
+    
+    profile["overall_stats"] = {
+        "total_cells": total_cells,
+        "total_nulls": total_nulls,
+        "null_percentage": round(total_nulls / total_cells * 100, 2) if total_cells > 0 else 0,
+        "duplicate_rows": df.is_duplicated().sum(),
+        "duplicate_percentage": round(df.is_duplicated().sum() / len(df) * 100, 2) if len(df) > 0 else 0,
+    }
+    
+    return profile
+
+
+def get_smart_summary(file_path: str, n_samples: int = 30) -> Dict[str, Any]:
+    """
+    Enhanced data summary with missing %, unique counts, and safe dict handling.
+    
+    This function provides a smarter, more LLM-friendly summary compared to profile_dataset().
+    It includes per-column missing percentages, unique value counts, and handles
+    dictionary columns gracefully (converts to strings to avoid hashing errors).
+    
+    Args:
+        file_path: Path to CSV or Parquet file
+        n_samples: Number of sample rows to include (default: 30)
+    
+    Returns:
+        Dictionary with comprehensive smart summary including:
+        - Basic shape info
+        - Column data types
+        - Missing value percentage by column (sorted by % descending)
+        - Unique value counts by column
+        - First N sample rows
+        - Descriptive statistics for numeric columns
+        - Safe handling of dictionary/unhashable columns
+    
+    Example:
+        >>> summary = get_smart_summary("data.csv")
+        >>> print(summary["missing_summary"])
+        >>> # Output: [("col_A", 45.2), ("col_B", 12.3), ...]
+    """
+    # Validation
+    validate_file_exists(file_path)
+    validate_file_format(file_path)
+    
+    # Load data
+    df = load_dataframe(file_path)
+    validate_dataframe(df)
+    
+    # Convert dictionary-type columns to strings (prevents unhashable dict errors)
+    for col in df.columns:
+        try:
+            # Try to detect if column might contain dicts/lists
+            sample = df[col].drop_nulls().head(5)
+            if len(sample) > 0:
+                first_val = sample[0]
+                # Check if it's a complex type
+                if isinstance(first_val, (dict, list)):
+                    df = df.with_columns(pl.col(col).cast(pl.Utf8).alias(col))
+        except:
+            # If any error, just continue
+            pass
+    
+    # Calculate missing value statistics (sorted by % descending)
+    missing_stats = []
+    for col in df.columns:
+        null_count = df[col].null_count()
+        null_pct = round((null_count / len(df)) * 100, 2) if len(df) > 0 else 0
+        missing_stats.append({
+            "column": col,
+            "count": int(null_count),
+            "percentage": null_pct
+        })
+    
+    # Sort by percentage descending
+    missing_stats.sort(key=lambda x: x["percentage"], reverse=True)
+    
+    # Calculate unique value counts
+    unique_counts = {}
+    for col in df.columns:
+        try:
+            unique_count = df[col].n_unique()
+            unique_counts[col] = int(unique_count)
+        except:
+            # Fallback for unhashable types
+            try:
+                unique_count = df[col].cast(pl.Utf8).n_unique()
+                unique_counts[col] = int(unique_count)
+            except:
+                unique_counts[col] = "N/A"
+    
+    # Get column data types
+    column_types = {col: str(df[col].dtype) for col in df.columns}
+    
+    # Get sample rows (first n_samples)
+    sample_data = df.head(n_samples).to_dicts()
+    
+    # Get descriptive statistics for numeric columns
+    numeric_cols = get_numeric_columns(df)
+    numeric_stats = {}
+    
+    if numeric_cols:
+        df_numeric = df.select(numeric_cols)
+        # Convert to pandas for describe() functionality
+        df_pd = df_numeric.to_pandas()
+        stats_df = df_pd.describe()
+        numeric_stats = stats_df.to_dict()
+    
+    # Build comprehensive summary
+    summary = {
+        "file_path": file_path,
+        "shape": {
+            "rows": len(df),
+            "columns": len(df.columns)
+        },
+        "column_types": column_types,
+        "missing_summary": missing_stats,  # Sorted by % descending
+        "unique_counts": unique_counts,
+        "sample_data": sample_data,
+        "numeric_statistics": numeric_stats,
+        "memory_usage_mb": calculate_memory_usage(df),
+        "summary_notes": []
+    }
+    
+    # Add helpful notes for LLM
+    high_missing_cols = [item for item in missing_stats if item["percentage"] > 40]
+    if high_missing_cols:
+        summary["summary_notes"].append(
+            f"{len(high_missing_cols)} column(s) have >40% missing values (consider dropping)"
+        )
+    
+    high_cardinality_cols = [col for col, count in unique_counts.items() 
+                            if isinstance(count, int) and count > len(df) * 0.5]
+    if high_cardinality_cols:
+        summary["summary_notes"].append(
+            f"{len(high_cardinality_cols)} column(s) have very high cardinality (>50% unique values)"
+        )
+    
+    return summary
+
+
+def detect_data_quality_issues(file_path: str) -> Dict[str, Any]:
+    """
+    Detect data quality issues in the dataset.
+    
+    Args:
+        file_path: Path to CSV or Parquet file
+        
+    Returns:
+        Dictionary with detected issues organized by severity:
+        - critical: Issues that will break model training
+        - warning: Issues that may affect model performance
+        - info: Observations that may be relevant
+    """
+    # Validation
+    validate_file_exists(file_path)
+    validate_file_format(file_path)
+    
+    # Load data
+    df = load_dataframe(file_path)
+    validate_dataframe(df)
+    
+    issues = {
+        "critical": [],
+        "warning": [],
+        "info": []
+    }
+    
+    # Check for completely null columns
+    for col in df.columns:
+        null_count = df[col].null_count()
+        null_pct = (null_count / len(df)) * 100
+        
+        if null_count == len(df):
+            issues["critical"].append({
+                "type": "all_null_column",
+                "column": col,
+                "message": f"Column '{col}' has all null values"
+            })
+        elif null_pct > 50:
+            issues["warning"].append({
+                "type": "high_null_percentage",
+                "column": col,
+                "null_percentage": round(null_pct, 2),
+                "message": f"Column '{col}' has {round(null_pct, 2)}% null values"
+            })
+        elif null_pct > 10:
+            issues["info"].append({
+                "type": "moderate_null_percentage",
+                "column": col,
+                "null_percentage": round(null_pct, 2),
+                "message": f"Column '{col}' has {round(null_pct, 2)}% null values"
+            })
+    
+    # Check for duplicate rows
+    dup_count = df.is_duplicated().sum()
+    if dup_count > 0:
+        dup_pct = (dup_count / len(df)) * 100
+        severity = "warning" if dup_pct > 10 else "info"
+        issues[severity].append({
+            "type": "duplicate_rows",
+            "count": int(dup_count),
+            "percentage": round(dup_pct, 2),
+            "message": f"Dataset has {dup_count} duplicate rows ({round(dup_pct, 2)}%)"
+        })
+    
+    # Check for outliers in numeric columns using IQR method
+    numeric_cols = get_numeric_columns(df)
+    for col in numeric_cols:
+        col_data = df[col].drop_nulls()
+        if len(col_data) == 0:
+            continue
+        
+        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
+        
+        outliers = ((col_data < lower_bound) | (col_data > upper_bound)).sum()
+        
+        if outliers > 0:
+            outlier_pct = (outliers / len(col_data)) * 100
+            if outlier_pct > 10:
+                issues["warning"].append({
+                    "type": "outliers",
+                    "column": col,
+                    "count": int(outliers),
+                    "percentage": round(outlier_pct, 2),
+                    "bounds": {"lower": float(lower_bound), "upper": float(upper_bound)},
+                    "message": f"Column '{col}' has {outliers} outliers ({round(outlier_pct, 2)}%)"
+                })
+            elif outlier_pct > 1:
+                issues["info"].append({
+                    "type": "outliers",
+                    "column": col,
+                    "count": int(outliers),
+                    "percentage": round(outlier_pct, 2),
+                    "bounds": {"lower": float(lower_bound), "upper": float(upper_bound)},
+                    "message": f"Column '{col}' has {outliers} outliers ({round(outlier_pct, 2)}%)"
+                })
+    
+    # Check for high cardinality in categorical columns
+    categorical_cols = get_categorical_columns(df)
+    for col in categorical_cols:
+        n_unique = df[col].n_unique()
+        cardinality_pct = (n_unique / len(df)) * 100
+        
+        if n_unique > 100 and cardinality_pct > 50:
+            issues["warning"].append({
+                "type": "high_cardinality",
+                "column": col,
+                "unique_values": int(n_unique),
+                "percentage": round(cardinality_pct, 2),
+                "message": f"Column '{col}' has very high cardinality ({n_unique} unique values, {round(cardinality_pct, 2)}%)"
+            })
+    
+    # Check for constant columns (single unique value)
+    for col in df.columns:
+        n_unique = df[col].n_unique()
+        if n_unique == 1:
+            issues["warning"].append({
+                "type": "constant_column",
+                "column": col,
+                "message": f"Column '{col}' has only one unique value (constant)"
+            })
+    
+    # Check for imbalanced datasets (for potential target columns)
+    for col in df.columns:
+        col_data = df[col]
+        n_unique = col_data.n_unique()
+        
+        # Check if this could be a target column (2-20 unique values)
+        if 2 <= n_unique <= 20:
+            value_counts = col_data.value_counts()
+            if len(value_counts) >= 2:
+                max_count = value_counts[value_counts.columns[1]][0]
+                max_pct = (max_count / len(df)) * 100
+                
+                if max_pct > 90:
+                    issues["warning"].append({
+                        "type": "class_imbalance",
+                        "column": col,
+                        "dominant_class_percentage": round(max_pct, 2),
+                        "message": f"Column '{col}' may be imbalanced (dominant class: {round(max_pct, 2)}%)"
+                    })
+    
+    # Summary
+    issues["summary"] = {
+        "total_issues": len(issues["critical"]) + len(issues["warning"]) + len(issues["info"]),
+        "critical_count": len(issues["critical"]),
+        "warning_count": len(issues["warning"]),
+        "info_count": len(issues["info"])
+    }
+    
+    return issues
+
+
+def analyze_correlations(file_path: str, target: Optional[str] = None) -> Dict[str, Any]:
+    """
+    Analyze correlations between features.
+    
+    Args:
+        file_path: Path to CSV or Parquet file
+        target: Optional target column to analyze correlations with
+        
+    Returns:
+        Dictionary with correlation analysis including:
+        - correlation matrix (for numeric columns)
+        - top correlations with target (if specified)
+        - highly correlated feature pairs
+    """
+    # Validation
+    validate_file_exists(file_path)
+    validate_file_format(file_path)
+    
+    # Load data
+    df = load_dataframe(file_path)
+    validate_dataframe(df)
+    
+    numeric_cols = get_numeric_columns(df)
+    
+    if len(numeric_cols) < 2:
+        return {
+            "error": "Dataset must have at least 2 numeric columns for correlation analysis",
+            "numeric_columns_found": len(numeric_cols)
+        }
+    
+    # Select only numeric columns for correlation
+    df_numeric = df.select(numeric_cols)
+    
+    # Calculate correlation matrix using pandas (Polars doesn't have native corr yet)
+    df_pd = df_numeric.to_pandas()
+    corr_matrix = df_pd.corr()
+    
+    result = {
+        "numeric_columns": numeric_cols,
+        "correlation_matrix": corr_matrix.to_dict()
+    }
+    
+    # Find highly correlated pairs (excluding diagonal)
+    high_corr_pairs = []
+    for i in range(len(corr_matrix.columns)):
+        for j in range(i + 1, len(corr_matrix.columns)):
+            col1 = corr_matrix.columns[i]
+            col2 = corr_matrix.columns[j]
+            corr_value = corr_matrix.iloc[i, j]
+            
+            if abs(corr_value) > 0.7:  # High correlation threshold
+                high_corr_pairs.append({
+                    "feature_1": col1,
+                    "feature_2": col2,
+                    "correlation": round(float(corr_value), 4)
+                })
+    
+    # Sort by absolute correlation
+    high_corr_pairs.sort(key=lambda x: abs(x["correlation"]), reverse=True)
+    result["high_correlations"] = high_corr_pairs
+    
+    # If target specified, show top correlations with target
+    if target:
+        if target not in df.columns:
+            result["target_correlations_error"] = f"Target column '{target}' not found"
+        elif target not in numeric_cols:
+            result["target_correlations_error"] = f"Target column '{target}' is not numeric"
+        else:
+            target_corrs = []
+            for col in numeric_cols:
+                if col != target:
+                    corr_value = corr_matrix.loc[target, col]
+                    target_corrs.append({
+                        "feature": col,
+                        "correlation": round(float(corr_value), 4)
+                    })
+            
+            # Sort by absolute correlation
+            target_corrs.sort(key=lambda x: abs(x["correlation"]), reverse=True)
+            result["target_correlations"] = {
+                "target": target,
+                "top_features": target_corrs[:20]  # Top 20
+            }
+    
+    return result
+
+
+def detect_label_errors(
+    file_path: str,
+    target_col: str,
+    features: Optional[List[str]] = None,
+    n_folds: int = 5,
+    output_path: Optional[str] = None
+) -> Dict[str, Any]:
+    """
+    Detect potential label errors in a classification dataset using cleanlab.
+    
+    Uses confident learning to find mislabeled examples by:
+    1. Training cross-validated classifiers
+    2. Computing out-of-sample predicted probabilities
+    3. Identifying labels that disagree with model predictions
+    
+    Args:
+        file_path: Path to dataset
+        target_col: Target/label column name
+        features: Feature columns to use (None = all numeric)
+        n_folds: Number of cross-validation folds
+        output_path: Optional path to save flagged rows
+        
+    Returns:
+        Dictionary with label error analysis results
+    """
+    try:
+        from cleanlab.classification import CleanLearning
+    except ImportError:
+        return {
+            'status': 'error',
+            'message': 'cleanlab not installed. Install with: pip install cleanlab>=2.6'
+        }
+    
+    from sklearn.linear_model import LogisticRegression
+    from sklearn.preprocessing import LabelEncoder
+    
+    validate_file_exists(file_path)
+    validate_file_format(file_path)
+    
+    df = load_dataframe(file_path)
+    validate_dataframe(df)
+    validate_column_exists(df, target_col)
+    
+    print(f"🔍 Detecting label errors in '{target_col}' using cleanlab...")
+    
+    # Get features
+    if features is None:
+        features = get_numeric_columns(df)
+        features = [f for f in features if f != target_col]
+    
+    if not features:
+        return {'status': 'error', 'message': 'No numeric features found for label error detection'}
+    
+    # Convert to pandas/numpy
+    df_pd = df.to_pandas()
+    X = df_pd[features].fillna(0).values
+    y_raw = df_pd[target_col].values
+    
+    # Encode labels
+    le = LabelEncoder()
+    y = le.fit_transform(y_raw)
+    
+    # Use CleanLearning to find label issues
+    cl = CleanLearning(
+        clf=LogisticRegression(max_iter=500, solver='lbfgs', multi_class='auto'),
+        cv_n_folds=n_folds
+    )
+    
+    label_issues = cl.find_label_issues(X, y)
+    
+    # Extract results
+    n_issues = label_issues['is_label_issue'].sum()
+    issue_indices = label_issues[label_issues['is_label_issue']].index.tolist()
+    
+    # Get details for flagged rows
+    flagged_rows = []
+    for idx in issue_indices[:50]:  # Limit to top 50
+        flagged_rows.append({
+            'row_index': int(idx),
+            'current_label': str(y_raw[idx]),
+            'suggested_label': str(le.inverse_transform([label_issues.loc[idx, 'predicted_label']])[0]) if 'predicted_label' in label_issues.columns else 'unknown',
+            'confidence': float(1 - label_issues.loc[idx, 'label_quality']) if 'label_quality' in label_issues.columns else None
+        })
+    
+    print(f"   🚨 Found {n_issues} potential label errors ({n_issues/len(y)*100:.1f}%)")
+    
+    # Save flagged rows
+    if output_path and issue_indices:
+        flagged_df = df_pd.iloc[issue_indices]
+        flagged_df.to_csv(output_path, index=False)
+        print(f"   💾 Flagged rows saved to: {output_path}")
+    
+    return {
+        'status': 'success',
+        'total_samples': len(y),
+        'label_errors_found': int(n_issues),
+        'error_percentage': round(n_issues / len(y) * 100, 2),
+        'flagged_rows': flagged_rows,
+        'n_classes': len(le.classes_),
+        'classes': le.classes_.tolist(),
+        'output_path': output_path,
+        'recommendation': f'Review {n_issues} flagged samples for potential mislabeling' if n_issues > 0 else 'No label errors detected'
+    }