ds-agent-cli 0.1.0

package/src/tools/data_type_conversion.py

@@ -0,0 +1,268 @@
+"""
+Advanced data type conversion tools for handling tricky type issues.
+"""
+
+import polars as pl
+from pathlib import Path
+from typing import Dict, Any, List, Optional
+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
+)
+from ds_agent.utils.validation import (
+    validate_file_exists,
+    validate_file_format,
+    validate_dataframe
+)
+
+
+def force_numeric_conversion(
+    file_path: str,
+    columns: List[str],
+    output_path: str,
+    errors: str = "coerce"
+) -> Dict[str, Any]:
+    """
+    Force convert columns to numeric type, even if they're detected as strings/objects.
+    
+    This is crucial for datasets where numeric columns are stored as strings with 
+    formatting issues (commas, spaces, currency symbols, etc.).
+    
+    Args:
+        file_path: Path to CSV or Parquet file
+        columns: List of column names to force convert, or ["all"] for all non-ID columns
+        output_path: Path to save converted dataset
+        errors: How to handle conversion errors:
+               - "coerce": Invalid values become null (default)
+               - "raise": Raise error on invalid values
+               
+    Returns:
+        Dictionary with conversion report and statistics
+    """
+    # Validation
+    validate_file_exists(file_path)
+    validate_file_format(file_path)
+    
+    # Load data
+    df = load_dataframe(file_path)
+    validate_dataframe(df)
+    
+    original_types = {col: str(df[col].dtype) for col in df.columns}
+    
+    # Determine which columns to convert
+    if columns == ["all"]:
+        # Auto-detect: skip ID columns, already-numeric columns, and actual text columns
+        id_keywords = ['id', 'key', 'code', 'name', 'description', 'text', 'comment', 'notes']
+        target_columns = []
+        
+        for col in df.columns:
+            # Skip if already numeric
+            if df[col].dtype in [pl.Int64, pl.Int32, pl.Float64, pl.Float32]:
+                continue
+            
+            # Skip if looks like an ID or text column
+            col_lower = col.lower()
+            if any(keyword in col_lower for keyword in id_keywords):
+                continue
+            
+            # Only attempt conversion if column looks numeric
+            # Sample first 100 non-null values to check if they're numeric-like
+            sample_values = df[col].drop_nulls().head(100).to_list()
+            if len(sample_values) == 0:
+                continue
+            
+            numeric_like_count = 0
+            for val in sample_values[:min(50, len(sample_values))]:  # Check first 50 samples
+                val_str = str(val).replace(",", "").replace(" ", "").replace("$", "").replace("€", "").strip()
+                
+                # Check if it looks like a number (digits, decimal point, minus sign)
+                if val_str.replace(".", "").replace("-", "").replace("+", "").replace("e", "").replace("E", "").isdigit():
+                    numeric_like_count += 1
+                # Also check for percentage-like values
+                elif val_str.endswith("%") and val_str[:-1].replace(".", "").isdigit():
+                    numeric_like_count += 1
+            
+            # Only include if >70% of samples look numeric
+            if len(sample_values) > 0 and (numeric_like_count / min(50, len(sample_values))) > 0.7:
+                target_columns.append(col)
+                print(f"🔍 '{col}': Detected as numeric-like ({numeric_like_count}/{min(50, len(sample_values))} samples)")
+            else:
+                print(f"⏭️ '{col}': Skipping (appears to be text, not numeric)")
+    else:
+        target_columns = columns
+    
+    print(f"🔢 Force converting {len(target_columns)} columns to numeric...")
+    
+    # Track conversion results
+    conversion_report = {
+        "successful_conversions": [],
+        "failed_conversions": [],
+        "null_values_introduced": {}
+    }
+    
+    # Convert each column
+    for col in target_columns:
+        if col not in df.columns:
+            print(f"⚠️ Column '{col}' not found, skipping")
+            conversion_report["failed_conversions"].append(col)
+            continue
+        
+        try:
+            # Get original null count
+            original_nulls = df[col].null_count()
+            
+            # Try to convert to numeric
+            # First, clean the column if it's a string (remove commas, spaces, etc.)
+            if df[col].dtype == pl.Utf8:
+                # Remove common non-numeric characters
+                df = df.with_columns([
+                    pl.col(col)
+                    .str.replace_all(",", "")  # Remove commas
+                    .str.replace_all(" ", "")  # Remove spaces
+                    .str.replace_all("$", "")  # Remove dollar signs
+                    .str.replace_all("€", "")  # Remove euro signs
+                    .str.replace_all("%", "")  # Remove percent signs
+                    .str.strip_chars()  # Strip whitespace
+                    .alias(col)
+                ])
+            
+            # Now convert to float
+            if errors == "coerce":
+                df = df.with_columns([
+                    pl.col(col).cast(pl.Float64, strict=False).alias(col)
+                ])
+            else:
+                df = df.with_columns([
+                    pl.col(col).cast(pl.Float64, strict=True).alias(col)
+                ])
+            
+            # Check how many nulls were introduced
+            new_nulls = df[col].null_count()
+            nulls_introduced = new_nulls - original_nulls
+            
+            conversion_report["successful_conversions"].append(col)
+            conversion_report["null_values_introduced"][col] = int(nulls_introduced)
+            
+            if nulls_introduced > 0:
+                print(f"✅ '{col}': Converted to numeric ({nulls_introduced} values became null)")
+            else:
+                print(f"✅ '{col}': Converted to numeric (no data loss)")
+                
+        except Exception as e:
+            print(f"❌ '{col}': Conversion failed - {str(e)}")
+            conversion_report["failed_conversions"].append(col)
+    
+    # Save converted dataset
+    save_dataframe(df, output_path)
+    
+    new_types = {col: str(df[col].dtype) for col in df.columns}
+    
+    return {
+        "status": "success",
+        "message": f"Force converted {len(conversion_report['successful_conversions'])} columns to numeric",
+        "output_path": output_path,
+        "conversion_report": conversion_report,
+        "type_changes": {
+            col: {"from": original_types[col], "to": new_types[col]}
+            for col in conversion_report["successful_conversions"]
+        },
+        "total_successful": len(conversion_report["successful_conversions"]),
+        "total_failed": len(conversion_report["failed_conversions"]),
+        "total_nulls_introduced": sum(conversion_report["null_values_introduced"].values())
+    }
+
+
+def smart_type_inference(
+    file_path: str,
+    output_path: str,
+    aggressive: bool = True
+) -> Dict[str, Any]:
+    """
+    Intelligently infer and fix data types for all columns.
+    
+    This tool goes beyond basic type detection and tries to understand the
+    semantic meaning of each column to assign the correct type.
+    
+    Args:
+        file_path: Path to CSV or Parquet file
+        output_path: Path to save dataset with fixed types
+        aggressive: If True, attempts aggressive type conversion (force numeric on ambiguous columns)
+        
+    Returns:
+        Dictionary with type inference report
+    """
+    # Validation
+    validate_file_exists(file_path)
+    validate_file_format(file_path)
+    
+    # Load data
+    df = load_dataframe(file_path)
+    validate_dataframe(df)
+    
+    original_types = {col: str(df[col].dtype) for col in df.columns}
+    type_changes = {}
+    
+    print(f"🧠 Performing smart type inference on {len(df.columns)} columns...")
+    
+    for col in df.columns:
+        current_type = df[col].dtype
+        
+        # Skip if already numeric
+        if current_type in [pl.Int64, pl.Int32, pl.Float64, pl.Float32]:
+            continue
+        
+        # If it's a string column, try to infer the correct type
+        if current_type == pl.Utf8:
+            sample_values = df[col].drop_nulls().head(100).to_list()
+            
+            if len(sample_values) == 0:
+                continue
+            
+            # Try to detect if it's actually numeric
+            numeric_count = 0
+            for val in sample_values:
+                # Clean and test
+                cleaned = str(val).replace(",", "").replace(" ", "").replace("$", "").strip()
+                try:
+                    float(cleaned)
+                    numeric_count += 1
+                except:
+                    pass
+            
+            # If >80% of values are numeric, convert to numeric
+            if numeric_count / len(sample_values) > 0.8:
+                print(f"🔢 '{col}': Detected as numeric ({numeric_count}/{len(sample_values)} samples)")
+                
+                # Clean and convert
+                df = df.with_columns([
+                    pl.col(col)
+                    .str.replace_all(",", "")
+                    .str.replace_all(" ", "")
+                    .str.replace_all("$", "")
+                    .str.replace_all("€", "")
+                    .str.strip_chars()
+                    .cast(pl.Float64, strict=False)
+                    .alias(col)
+                ])
+                
+                type_changes[col] = {"from": "Utf8", "to": "Float64", "reason": "numeric_pattern_detected"}
+    
+    # Save dataset
+    save_dataframe(df, output_path)
+    
+    return {
+        "status": "success",
+        "message": f"Smart type inference completed, changed {len(type_changes)} columns",
+        "output_path": output_path,
+        "type_changes": type_changes,
+        "original_types": original_types,
+        "new_types": {col: str(df[col].dtype) for col in df.columns}
+    }

package/src/tools/data_wrangling.py

@@ -0,0 +1,433 @@
+"""
+Data Wrangling Tools
+Tools for merging, concatenating, and manipulating multiple datasets.
+"""
+
+import polars as pl
+import numpy as np
+from typing import Dict, Any, List, Optional, Literal
+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,
+)
+from ds_agent.utils.validation import (
+    validate_file_exists,
+    validate_file_format,
+    validate_dataframe,
+)
+
+
+def merge_datasets(
+    left_path: str,
+    right_path: str,
+    output_path: str,
+    how: Literal["inner", "left", "right", "outer", "cross"] = "inner",
+    on: Optional[str] = None,
+    left_on: Optional[str] = None,
+    right_on: Optional[str] = None,
+    suffix: str = "_right"
+) -> Dict[str, Any]:
+    """
+    Merge two datasets using various join strategies (SQL-like join operations).
+    
+    This function performs database-style joins on two datasets, similar to SQL JOIN operations.
+    Supports inner, left, right, outer, and cross joins.
+    
+    Args:
+        left_path: Path to left dataset (CSV or Parquet)
+        right_path: Path to right dataset (CSV or Parquet)
+        output_path: Path to save merged dataset
+        how: Join type - "inner", "left", "right", "outer", or "cross"
+            - "inner": Only rows with matching keys in both datasets
+            - "left": All rows from left, matching rows from right (nulls if no match)
+            - "right": All rows from right, matching rows from left (nulls if no match)
+            - "outer": All rows from both (nulls where no match)
+            - "cross": Cartesian product (all combinations)
+        on: Column name to join on (if same in both datasets)
+        left_on: Column name in left dataset (if different from right)
+        right_on: Column name in right dataset (if different from left)
+        suffix: Suffix to add to duplicate column names from right dataset (default: "_right")
+    
+    Returns:
+        Dictionary with merge report including:
+        - success: bool
+        - output_path: str
+        - left_rows: int
+        - right_rows: int
+        - result_rows: int
+        - merge_type: str
+        - join_columns: dict
+        - duplicate_columns: list (columns that got suffixed)
+    
+    Examples:
+        >>> # Simple join on same column name
+        >>> merge_datasets(
+        ...     "customers.csv", 
+        ...     "orders.csv",
+        ...     "merged.csv",
+        ...     how="left",
+        ...     on="customer_id"
+        ... )
+        
+        >>> # Join on different column names
+        >>> merge_datasets(
+        ...     "products.csv",
+        ...     "sales.csv",
+        ...     "product_sales.csv",
+        ...     how="inner",
+        ...     left_on="product_id",
+        ...     right_on="prod_id"
+        ... )
+    """
+    try:
+        # Validation
+        validate_file_exists(left_path)
+        validate_file_exists(right_path)
+        validate_file_format(left_path)
+        validate_file_format(right_path)
+        
+        # Load datasets
+        left_df = load_dataframe(left_path)
+        right_df = load_dataframe(right_path)
+        
+        validate_dataframe(left_df)
+        validate_dataframe(right_df)
+        
+        left_rows = len(left_df)
+        right_rows = len(right_df)
+        
+        # Determine join columns
+        if on:
+            # Same column name in both datasets
+            join_left_on = on
+            join_right_on = on
+            
+            # Validate column exists
+            if on not in left_df.columns:
+                return {
+                    "success": False,
+                    "error": f"Column '{on}' not found in left dataset. Available: {left_df.columns}"
+                }
+            if on not in right_df.columns:
+                return {
+                    "success": False,
+                    "error": f"Column '{on}' not found in right dataset. Available: {right_df.columns}"
+                }
+        elif left_on and right_on:
+            # Different column names
+            join_left_on = left_on
+            join_right_on = right_on
+            
+            # Validate columns exist
+            if left_on not in left_df.columns:
+                return {
+                    "success": False,
+                    "error": f"Column '{left_on}' not found in left dataset. Available: {left_df.columns}"
+                }
+            if right_on not in right_df.columns:
+                return {
+                    "success": False,
+                    "error": f"Column '{right_on}' not found in right dataset. Available: {right_df.columns}"
+                }
+        else:
+            return {
+                "success": False,
+                "error": "Must specify either 'on' (same column name) or both 'left_on' and 'right_on' (different names)"
+            }
+        
+        # Check for duplicate column names (excluding join columns)
+        left_cols = set(left_df.columns)
+        right_cols = set(right_df.columns)
+        duplicate_cols = list((left_cols & right_cols) - {join_left_on, join_right_on})
+        
+        # Perform merge
+        merged_df = left_df.join(
+            right_df,
+            left_on=join_left_on,
+            right_on=join_right_on,
+            how=how,
+            suffix=suffix
+        )
+        
+        result_rows = len(merged_df)
+        
+        # Save result
+        Path(output_path).parent.mkdir(parents=True, exist_ok=True)
+        save_dataframe(merged_df, output_path)
+        
+        # Build report
+        report = {
+            "success": True,
+            "output_path": output_path,
+            "left_file": Path(left_path).name,
+            "right_file": Path(right_path).name,
+            "left_rows": left_rows,
+            "right_rows": right_rows,
+            "result_rows": result_rows,
+            "result_columns": len(merged_df.columns),
+            "merge_type": how,
+            "join_columns": {
+                "left": join_left_on,
+                "right": join_right_on
+            },
+            "duplicate_columns": duplicate_cols,
+            "rows_added": result_rows - left_rows if how in ["left", "inner"] else None,
+            "message": f"Successfully merged {left_rows:,} rows with {right_rows:,} rows using {how} join → {result_rows:,} rows"
+        }
+        
+        # Add warnings
+        if how == "inner" and result_rows < min(left_rows, right_rows):
+            report["warning"] = f"Inner join reduced data: only {result_rows:,} of {min(left_rows, right_rows):,} rows had matches"
+        elif how == "outer" and result_rows > left_rows + right_rows:
+            report["warning"] = "Outer join created duplicate rows - check for many-to-many relationships"
+        
+        if duplicate_cols:
+            report["note"] = f"{len(duplicate_cols)} column(s) were suffixed with '{suffix}': {', '.join(duplicate_cols)}"
+        
+        return report
+    
+    except Exception as e:
+        return {
+            "success": False,
+            "error": str(e),
+            "error_type": type(e).__name__
+        }
+
+
+def concat_datasets(
+    file_paths: List[str],
+    output_path: str,
+    axis: Literal["vertical", "horizontal"] = "vertical",
+    ignore_index: bool = True
+) -> Dict[str, Any]:
+    """
+    Concatenate multiple datasets vertically (stack rows) or horizontally (add columns).
+    
+    Args:
+        file_paths: List of file paths to concatenate (CSV or Parquet)
+        output_path: Path to save concatenated dataset
+        axis: "vertical" to stack rows (union), "horizontal" to add columns side-by-side
+        ignore_index: If True, reset index after concatenation (default: True)
+    
+    Returns:
+        Dictionary with concatenation report including:
+        - success: bool
+        - output_path: str
+        - input_files: int
+        - result_rows: int
+        - result_cols: int
+        - axis: str
+    
+    Examples:
+        >>> # Stack multiple CSV files (union)
+        >>> concat_datasets(
+        ...     ["jan_sales.csv", "feb_sales.csv", "mar_sales.csv"],
+        ...     "q1_sales.csv",
+        ...     axis="vertical"
+        ... )
+        
+        >>> # Combine datasets side-by-side (add columns)
+        >>> concat_datasets(
+        ...     ["features.csv", "labels.csv"],
+        ...     "full_dataset.csv",
+        ...     axis="horizontal"
+        ... )
+    """
+    try:
+        # Validation
+        if not file_paths or len(file_paths) < 2:
+            return {
+                "success": False,
+                "error": "Must provide at least 2 files to concatenate"
+            }
+        
+        for fp in file_paths:
+            validate_file_exists(fp)
+            validate_file_format(fp)
+        
+        # Load all datasets
+        dfs = []
+        file_info = []
+        
+        for fp in file_paths:
+            df = load_dataframe(fp)
+            validate_dataframe(df)
+            dfs.append(df)
+            file_info.append({
+                "file": Path(fp).name,
+                "rows": len(df),
+                "columns": len(df.columns)
+            })
+        
+        # Perform concatenation
+        if axis == "vertical":
+            # Stack rows (union) - requires same columns
+            result = pl.concat(dfs, how="vertical")
+        else:  # horizontal
+            # Add columns side-by-side - requires same number of rows
+            row_counts = [len(df) for df in dfs]
+            if len(set(row_counts)) > 1:
+                return {
+                    "success": False,
+                    "error": f"Horizontal concatenation requires same number of rows. Got: {row_counts}",
+                    "file_info": file_info
+                }
+            result = pl.concat(dfs, how="horizontal")
+        
+        # Save result
+        Path(output_path).parent.mkdir(parents=True, exist_ok=True)
+        save_dataframe(result, output_path)
+        
+        return {
+            "success": True,
+            "output_path": output_path,
+            "input_files": len(file_paths),
+            "file_info": file_info,
+            "result_rows": len(result),
+            "result_cols": len(result.columns),
+            "axis": axis,
+            "message": f"Successfully concatenated {len(file_paths)} files ({axis}) → {len(result):,} rows × {len(result.columns)} columns"
+        }
+    
+    except Exception as e:
+        return {
+            "success": False,
+            "error": str(e),
+            "error_type": type(e).__name__
+        }
+
+
+def reshape_dataset(
+    file_path: str,
+    output_path: str,
+    operation: Literal["pivot", "melt", "transpose"],
+    **kwargs
+) -> Dict[str, Any]:
+    """
+    Reshape dataset using pivot, melt, or transpose operations.
+    
+    Args:
+        file_path: Path to CSV or Parquet file
+        output_path: Path to save reshaped dataset
+        operation: "pivot" (wide format), "melt" (long format), or "transpose"
+        **kwargs: Operation-specific parameters
+            For pivot: index, columns, values, aggregate_function
+            For melt: id_vars, value_vars, var_name, value_name
+    
+    Returns:
+        Dictionary with reshape report
+    
+    Examples:
+        >>> # Pivot: wide format
+        >>> reshape_dataset(
+        ...     "sales_long.csv",
+        ...     "sales_wide.csv",
+        ...     operation="pivot",
+        ...     index="date",
+        ...     columns="product",
+        ...     values="sales"
+        ... )
+        
+        >>> # Melt: long format
+        >>> reshape_dataset(
+        ...     "sales_wide.csv",
+        ...     "sales_long.csv",
+        ...     operation="melt",
+        ...     id_vars=["date"],
+        ...     value_vars=["product_a", "product_b"],
+        ...     var_name="product",
+        ...     value_name="sales"
+        ... )
+    """
+    try:
+        # Validation
+        validate_file_exists(file_path)
+        validate_file_format(file_path)
+        
+        # Load data
+        df = load_dataframe(file_path)
+        validate_dataframe(df)
+        
+        original_shape = (len(df), len(df.columns))
+        
+        # Perform operation
+        if operation == "pivot":
+            # Pivot to wide format
+            index = kwargs.get("index")
+            columns = kwargs.get("columns")
+            values = kwargs.get("values")
+            
+            if not all([index, columns, values]):
+                return {
+                    "success": False,
+                    "error": "Pivot requires: index, columns, values parameters"
+                }
+            
+            result = df.pivot(
+                index=index,
+                columns=columns,
+                values=values
+            )
+        
+        elif operation == "melt":
+            # Melt to long format
+            id_vars = kwargs.get("id_vars")
+            value_vars = kwargs.get("value_vars")
+            var_name = kwargs.get("var_name", "variable")
+            value_name = kwargs.get("value_name", "value")
+            
+            if not id_vars:
+                return {
+                    "success": False,
+                    "error": "Melt requires: id_vars parameter"
+                }
+            
+            result = df.melt(
+                id_vars=id_vars,
+                value_vars=value_vars,
+                variable_name=var_name,
+                value_name=value_name
+            )
+        
+        elif operation == "transpose":
+            # Transpose rows and columns
+            result = df.transpose()
+        
+        else:
+            return {
+                "success": False,
+                "error": f"Unknown operation: {operation}. Use 'pivot', 'melt', or 'transpose'"
+            }
+        
+        # Save result
+        Path(output_path).parent.mkdir(parents=True, exist_ok=True)
+        save_dataframe(result, output_path)
+        
+        return {
+            "success": True,
+            "output_path": output_path,
+            "operation": operation,
+            "original_shape": {
+                "rows": original_shape[0],
+                "columns": original_shape[1]
+            },
+            "result_shape": {
+                "rows": len(result),
+                "columns": len(result.columns)
+            },
+            "message": f"Successfully {operation}ed dataset: {original_shape[0]}×{original_shape[1]} → {len(result)}×{len(result.columns)}"
+        }
+    
+    except Exception as e:
+        return {
+            "success": False,
+            "error": str(e),
+            "error_type": type(e).__name__
+        }