additory 0.1.0a1__py3-none-any.whl

additory/augment/augmentor.py

@@ -0,0 +1,653 @@
+"""
+Data Augmentation Engine - Polars-Only Architecture
+
+Provides functionality to augment dataframes by adding synthetic rows
+based on existing data patterns.
+
+Architecture:
+1. Detect input format (pandas/polars/cuDF)
+2. Convert to Polars via Arrow bridge (if needed)
+3. Process augmentation in Polars
+4. Convert back to original format via Arrow bridge
+"""
+
+from typing import Union, Optional, Any, Dict, Literal
+import random
+
+from additory.common.backend import detect_backend, to_polars, from_polars
+from additory.common.exceptions import ValidationError, AugmentError
+from additory.common.validation import validate_dataframe
+from additory.common.sample_data import get_sample_dataset
+from additory.augment.strategies import (
+    parse_strategy_dict,
+    get_column_strategy,
+    apply_increment_strategy,
+    apply_choice_strategy,
+    apply_range_strategy,
+    parse_strategy_params
+)
+
+
+def _validate_generative_strategies(strategy_dict: Dict[str, str]) -> None:
+    """
+    Validate that all strategies are generative (not augmentative).
+    
+    Generative strategies can create data from scratch:
+    - increment (with start parameter)
+    - range
+    - choice
+    - choice_list
+    
+    Augmentative strategies require existing data:
+    - auto (random sampling)
+    - forecast (time series)
+    - seasonal (time series)
+    - smote (synthetic minority oversampling)
+    
+    Args:
+        strategy_dict: Dictionary mapping column names to strategy specs
+        
+    Raises:
+        ValidationError: If any augmentative strategies are found
+    """
+    augmentative_strategies = ["auto", "forecast", "seasonal", "smote"]
+    
+    invalid_columns = []
+    
+    for col, strategy_spec in strategy_dict.items():
+        if col == "__default__":
+            continue
+            
+        # Get the base strategy name (before any parameters)
+        strategy_name = strategy_spec.split(":")[0].strip()
+        
+        if strategy_name in augmentative_strategies:
+            invalid_columns.append((col, strategy_name))
+    
+    if invalid_columns:
+        error_lines = [
+            f"Create mode requires generative strategies. Found augmentative strategies:"
+        ]
+        for col, strat in invalid_columns:
+            error_lines.append(f"  - Column '{col}': '{strat}'")
+        
+        error_lines.append("")
+        error_lines.append("Valid generative strategies:")
+        error_lines.append("  - increment (with start parameter)")
+        error_lines.append("  - range:min-max")
+        error_lines.append("  - choice:[value1,value2,...]")
+        error_lines.append("  - choice_list:list_name")
+        
+        raise ValidationError("\n".join(error_lines))
+
+
+def _detect_mode(df: Any) -> Literal["augment", "create", "sample"]:
+    """
+    Detect the augmentation mode based on the df parameter.
+    
+    Three modes are supported:
+    1. "augment" - Augment an existing DataFrame (default)
+    2. "create" - Create data from scratch using "@new" sentinel
+    3. "sample" - Load and optionally augment sample dataset using "@sample" sentinel
+    
+    Args:
+        df: Input parameter (DataFrame or sentinel string)
+        
+    Returns:
+        Mode string: "augment", "create", or "sample"
+        
+    Raises:
+        ValidationError: If df is an invalid string (not "@new" or "@sample")
+    """
+    # Check for sentinel values
+    if isinstance(df, str):
+        if df == "@new":
+            return "create"
+        elif df == "@sample":
+            return "sample"
+        else:
+            # Provide helpful error messages
+            if df.lower() in ["new", "create"]:
+                raise ValidationError(
+                    f"Invalid input: '{df}'. Did you mean '@new'? "
+                    "Use '@new' to create data from scratch."
+                )
+            elif df.lower() in ["sample", "samples"]:
+                raise ValidationError(
+                    f"Invalid input: '{df}'. Did you mean '@sample'? "
+                    "Use '@sample' to load sample dataset."
+                )
+            else:
+                raise ValidationError(
+                    f"Invalid string input: '{df}'. "
+                    "Expected a DataFrame, '@new' (create mode), or '@sample' (sample mode)."
+                )
+    
+    # If not a string, assume it's a DataFrame (augment mode)
+    return "augment"
+
+
+def _parse_n_rows(n_rows: Union[int, str], df_length: int) -> int:
+    """
+    Parse n_rows parameter to get actual number of rows to generate.
+    
+    Args:
+        n_rows: Number of rows (int), percentage ("50%"), or multiplier ("2x")
+        df_length: Length of the input dataframe
+        
+    Returns:
+        Actual number of rows to generate
+        
+    Raises:
+        ValidationError: If n_rows format is invalid
+    """
+    if isinstance(n_rows, int):
+        if n_rows <= 0:
+            raise ValidationError("n_rows must be positive")
+        return n_rows
+    
+    if isinstance(n_rows, str):
+        n_rows = n_rows.strip()
+        
+        # Handle percentage: "50%"
+        if n_rows.endswith("%"):
+            try:
+                percentage = float(n_rows[:-1])
+                if percentage <= 0:
+                    raise ValidationError("Percentage must be positive")
+                return max(1, int(df_length * percentage / 100))
+            except ValueError:
+                raise ValidationError(f"Invalid percentage format: {n_rows}")
+        
+        # Handle multiplier: "2x"
+        if n_rows.endswith("x"):
+            try:
+                multiplier = float(n_rows[:-1])
+                if multiplier <= 0:
+                    raise ValidationError("Multiplier must be positive")
+                return max(1, int(df_length * multiplier))
+            except ValueError:
+                raise ValidationError(f"Invalid multiplier format: {n_rows}")
+        
+        raise ValidationError(
+            f"Invalid n_rows format: {n_rows}. "
+            "Use int (5), percentage ('50%'), or multiplier ('2x')"
+        )
+    
+    raise ValidationError(f"n_rows must be int or str, got {type(n_rows)}")
+
+
+def _augment_polars_engine(df_polars: Any, n_rows: int, strategy_dict: Dict[str, str], seed: Optional[int]) -> Any:
+    """
+    Augment Polars DataFrame with strategy support.
+    
+    This is the core augmentation engine that processes all dataframes
+    (pandas, polars, cuDF) after conversion to Polars format.
+    
+    Args:
+        df_polars: Input Polars DataFrame
+        n_rows: Number of rows to generate
+        strategy_dict: Column-specific strategies
+        seed: Random seed for reproducibility
+        
+    Returns:
+        Augmented Polars DataFrame
+    """
+    import polars as pl
+    
+    # Get column names
+    columns = df_polars.columns
+    
+    # Check if any column uses non-auto strategy
+    has_custom_strategy = any(
+        not get_column_strategy(col, strategy_dict).startswith("auto")
+        for col in columns
+    )
+    
+    if not has_custom_strategy:
+        # Simple random sampling (original behavior)
+        if seed is not None:
+            sampled = df_polars.sample(n=n_rows, with_replacement=True, seed=seed)
+        else:
+            sampled = df_polars.sample(n=n_rows, with_replacement=True)
+    else:
+        # Build new rows column by column
+        new_data = {}
+        
+        for col in columns:
+            col_strategy = get_column_strategy(col, strategy_dict)
+            
+            if col_strategy.startswith("increment"):
+                # Generate incremented values
+                new_values = apply_increment_strategy(
+                    df_polars, col, col_strategy, n_rows
+                )
+                new_data[col] = new_values
+            elif col_strategy.startswith("range"):
+                # Parse range parameters
+                strategy_name, params = parse_strategy_params(col_strategy)
+                
+                if "min" not in params or "max" not in params:
+                    raise ValidationError(
+                        f"Range strategy for column '{col}' requires min and max parameters. "
+                        f"Use format: 'range:min-max' (e.g., 'range:18-65')"
+                    )
+                
+                # Generate range values
+                new_values = apply_range_strategy(
+                    min_val=params["min"],
+                    max_val=params["max"],
+                    n_rows=n_rows,
+                    seed=seed
+                )
+                new_data[col] = new_values
+            elif col_strategy.startswith("choice"):
+                # Generate choice values
+                new_values = apply_choice_strategy(
+                    col_strategy, n_rows, seed
+                )
+                new_data[col] = new_values
+            elif col_strategy.startswith("forecast"):
+                # Import here to avoid circular dependency
+                from additory.augment.strategies import apply_forecast_strategy
+                
+                # Generate forecasted values
+                new_values = apply_forecast_strategy(
+                    df_polars, col, col_strategy, n_rows, seed
+                )
+                
+                # Cast to match original column dtype if needed
+                original_dtype = df_polars[col].dtype
+                if original_dtype.is_integer():
+                    # Round and convert to int for integer columns
+                    new_values = [int(round(v)) for v in new_values]
+                
+                new_data[col] = new_values
+            elif col_strategy.startswith(("normal", "uniform", "skewed_left", "skewed_right", "beta", "gamma", "exponential", "kde")):
+                # Import here to avoid circular dependency
+                from additory.augment.strategies import apply_distribution_strategy
+                
+                # Generate distribution values
+                new_values = apply_distribution_strategy(
+                    df_polars, col, col_strategy, n_rows, seed
+                )
+                
+                # Cast to match original column dtype if needed
+                original_dtype = df_polars[col].dtype
+                if original_dtype.is_integer():
+                    # Round and convert to int for integer columns
+                    new_values = [int(round(v)) for v in new_values]
+                
+                new_data[col] = new_values
+            else:
+                # Random sampling for this column (auto)
+                if seed is not None:
+                    sampled_col = df_polars.select(col).sample(n=n_rows, with_replacement=True, seed=seed)
+                else:
+                    sampled_col = df_polars.select(col).sample(n=n_rows, with_replacement=True)
+                new_data[col] = sampled_col[col].to_list()
+        
+        sampled = pl.DataFrame(new_data)
+    
+    # Concatenate original and new rows
+    result = pl.concat([df_polars, sampled])
+    
+    return result
+
+
+def _create_from_scratch_engine(
+    n_rows: int,
+    strategy_dict: Dict[str, str],
+    seed: Optional[int]
+) -> Any:
+    """
+    Create DataFrame from scratch using generative strategies.
+    
+    This engine generates data column by column without requiring
+    an existing DataFrame. All strategies must be generative.
+    
+    Generative strategies (supported):
+    - increment (with start parameter)
+    - range
+    - choice
+    - choice_list
+    
+    Augmentative strategies (NOT supported):
+    - auto (requires existing data)
+    - forecast (requires time series)
+    - seasonal (requires time series)
+    - smote (requires existing data)
+    
+    Args:
+        n_rows: Number of rows to generate
+        strategy_dict: Column-specific strategies (all must be generative)
+        seed: Random seed for reproducibility
+        
+    Returns:
+        Polars DataFrame with generated data
+        
+    Raises:
+        ValidationError: If any augmentative strategies are found
+    
+    Examples:
+        >>> # Simple create with increment and range
+        >>> result = _create_from_scratch_engine(
+        ...     n_rows=10,
+        ...     strategy_dict={
+        ...         "id": "increment:start=1",
+        ...         "age": "range:18-65"
+        ...     },
+        ...     seed=42
+        ... )
+        >>> result.shape
+        (10, 2)
+        
+        >>> # Create with mixed strategies
+        >>> result = _create_from_scratch_engine(
+        ...     n_rows=100,
+        ...     strategy_dict={
+        ...         "id": "increment:start=1",
+        ...         "emp_id": "increment:start=1:pattern=EMP_[001]",
+        ...         "age": "range:18-65",
+        ...         "status": "choice:[Active,Inactive]",
+        ...         "department": "choice_list:departments"
+        ...     },
+        ...     seed=42
+        ... )
+        >>> result.shape
+        (100, 5)
+    """
+    import polars as pl
+    
+    # Validate all strategies are generative
+    _validate_generative_strategies(strategy_dict)
+    
+    # Build data column by column
+    new_data = {}
+    
+    for col, col_strategy in strategy_dict.items():
+        if col == "__default__":
+            continue
+            
+        if col_strategy.startswith("increment"):
+            # Parse parameters for increment strategy
+            strategy_name, params = parse_strategy_params(col_strategy)
+            
+            # Generate incremented values (create mode)
+            new_values = apply_increment_strategy(
+                df_polars=None,  # No DataFrame in create mode
+                column=col,
+                strategy_spec=col_strategy,
+                n_rows=n_rows,
+                params=params
+            )
+            new_data[col] = new_values
+            
+        elif col_strategy.startswith("range"):
+            # Parse range parameters
+            strategy_name, params = parse_strategy_params(col_strategy)
+            
+            if "min" not in params or "max" not in params:
+                raise ValidationError(
+                    f"Range strategy for column '{col}' requires min and max parameters. "
+                    f"Use format: 'range:min-max' (e.g., 'range:18-65')"
+                )
+            
+            # Generate range values
+            new_values = apply_range_strategy(
+                min_val=params["min"],
+                max_val=params["max"],
+                n_rows=n_rows,
+                seed=seed
+            )
+            new_data[col] = new_values
+            
+        elif col_strategy.startswith("choice"):
+            # Generate choice values
+            new_values = apply_choice_strategy(
+                col_strategy, n_rows, seed
+            )
+            new_data[col] = new_values
+            
+        else:
+            raise ValidationError(
+                f"Unknown or unsupported strategy for column '{col}': '{col_strategy}'"
+            )
+    
+    # Build Polars DataFrame from generated columns
+    result = pl.DataFrame(new_data)
+    
+    return result
+
+
+def augment(
+    df: Any,
+    n_rows: Union[int, str] = 5,
+    strategy: Union[str, Dict[str, str]] = "auto",
+    seed: Optional[int] = None,
+    output_format: str = "pandas"
+) -> Any:
+    """
+    Augment a dataframe by adding synthetic rows based on existing data.
+    
+    Uses Polars-only architecture:
+    1. Detect input format (pandas/polars/cuDF)
+    2. Convert to Polars via Arrow bridge (if needed)
+    3. Process augmentation in Polars
+    4. Convert back to original format via Arrow bridge
+    
+    This function adds new rows to a dataframe using various strategies:
+    - "auto": Random sampling from existing values (default)
+    - "increment": Increment numeric or pattern-based values
+    - "range:min-max": Random integers within range
+    - "choice:[...]": Random selection from inline list
+    - "choice_list:name": Random selection from registered/built-in list
+    - "forecast:method": Time series forecasting (linear, polynomial, exponential, seasonal)
+    - "normal": Normal distribution generation
+    - "uniform": Uniform distribution generation
+    - "skewed_left/skewed_right": Skewed distribution generation
+    - "smote": Synthetic Minority Over-sampling Technique
+    
+    Args:
+        df: Input dataframe (pandas, polars, or cudf), or sentinel:
+            - DataFrame: Augment mode (add rows to existing data)
+            - "@new": Create mode (generate data from scratch)
+            - "@sample": Sample mode (load sample dataset)
+        n_rows: Number of rows to add. Can be:
+            - int: Exact number (e.g., 5)
+            - str percentage: Percentage of current size (e.g., "50%")
+            - str multiplier: Multiple of current size (e.g., "2x")
+        strategy: Augmentation strategy. Can be:
+            - str: "auto" (applies to all columns)
+            - dict: Column-specific strategies, e.g.:
+                {
+                    "id": "increment",
+                    "emp_id": "increment:EMP_[001]_ID",
+                    "age": "range:18-65",
+                    "status": "choice:[Active,Inactive,Pending]",
+                    "bank": "choice_list:banks",
+                    "sales": "forecast:seasonal:period=12",
+                    "score": "normal:mean=75:std=10",
+                    "income": "skewed_right:skewness=1.5"
+                }
+                Unlisted columns default to "auto"
+        seed: Random seed for reproducibility. If None, results will vary.
+        output_format: Output format for create/sample modes. Options:
+            - "pandas": Return pandas DataFrame (default)
+            - "polars": Return Polars DataFrame
+            - "cudf": Return cuDF DataFrame
+            Note: In augment mode (with DataFrame input), output format
+            matches input format and this parameter is ignored.
+    
+    Returns:
+        Augmented dataframe with original + new rows (same type as input)
+    
+    Raises:
+        ValidationError: If input validation fails
+        AugmentError: If augmentation fails
+    
+    Examples:
+        >>> # Add 5 rows with random sampling (default)
+        >>> df_aug = add.augment(df)
+        
+        >>> # Increment numeric ID column
+        >>> df_aug = add.augment(df, strategy={"id": "increment"})
+        
+        >>> # Forecast sales with seasonal pattern
+        >>> df_aug = add.augment(df, n_rows=24, strategy={
+        ...     "sales": "forecast:seasonal:period=12"
+        ... })
+        
+        >>> # Generate from normal distribution
+        >>> df_aug = add.augment(df, n_rows=100, strategy={
+        ...     "age": "normal:mean=35:std=10",
+        ...     "score": "uniform:min=0:max=100"
+        ... })
+        
+        >>> # Mixed strategies
+        >>> df_aug = add.augment(df, n_rows=100, strategy={
+        ...     "id": "increment",
+        ...     "age": "range:18-65",
+        ...     "status": "choice:[Active,Inactive]",
+        ...     "sales": "forecast:linear",
+        ...     "score": "normal:auto"
+        ... })
+        
+        >>> # Create data from scratch (returns pandas by default)
+        >>> df_new = add.augment("@new", n_rows=50, strategy={
+        ...     "id": "increment:start=1",
+        ...     "age": "range:18-65",
+        ...     "status": "choice:[Active,Inactive]"
+        ... })
+    """
+    # Detect mode
+    mode = _detect_mode(df)
+    
+    # Validate output_format parameter
+    valid_formats = ["pandas", "polars", "cudf"]
+    if output_format not in valid_formats:
+        raise ValidationError(
+            f"Invalid output_format: '{output_format}'. "
+            f"Must be one of: {', '.join(valid_formats)}"
+        )
+    
+    # Parse and validate strategy
+    try:
+        strategy_dict = parse_strategy_dict(strategy)
+    except ValidationError as e:
+        raise ValidationError(f"Invalid strategy parameter: {e}")
+    
+    # Handle create mode
+    if mode == "create":
+        # Validate create mode requirements
+        if not isinstance(strategy, dict):
+            raise ValidationError(
+                "Create mode requires a strategy dict with column definitions. "
+                "Example: strategy={'id': 'increment:start=1', 'age': 'range:18-65'}"
+            )
+        
+        if not strategy or len(strategy) == 0:
+            raise ValidationError(
+                "Create mode requires at least one column in strategy dict"
+            )
+        
+        if not isinstance(n_rows, int):
+            raise ValidationError(
+                f"Create mode requires n_rows to be an integer, got {type(n_rows).__name__}. "
+                "Percentage ('50%') and multiplier ('2x') formats are not supported in create mode."
+            )
+        
+        if n_rows <= 0:
+            raise ValidationError("n_rows must be positive")
+        
+        try:
+            # Generate data from scratch
+            result_polars = _create_from_scratch_engine(n_rows, strategy_dict, seed)
+            
+            # Convert to requested output format
+            result_df = from_polars(result_polars, output_format)
+            
+            # Memory cleanup
+            del result_polars
+            import gc
+            gc.collect()
+            
+            return result_df
+            
+        except Exception as e:
+            if isinstance(e, (ValidationError, AugmentError)):
+                raise
+            raise AugmentError(f"Create mode failed: {e}")
+    
+    # Handle sample mode
+    if mode == "sample":
+        # Load sample dataset
+        try:
+            df = get_sample_dataset("augment", "sample", "clean")
+        except Exception as e:
+            raise ValidationError(f"Failed to load sample dataset: {e}")
+        
+        # Continue to augment mode with loaded sample
+        # (will use output_format at the end)
+    
+    # Augment mode (original behavior)
+    # Validate input dataframe
+    validate_dataframe(df, "df")
+    
+    # Check minimum size - require at least 3 rows for meaningful augmentation
+    MIN_ROWS = 3
+    df_length = len(df)
+    if df_length < MIN_ROWS:
+        raise ValidationError(
+            f"Minimum {MIN_ROWS} rows required for augmentation. "
+            f"Current size: {df_length}"
+        )
+    
+    # Parse n_rows
+    try:
+        actual_n_rows = _parse_n_rows(n_rows, df_length)
+    except ValidationError as e:
+        raise ValidationError(f"Invalid n_rows parameter: {e}")
+    
+    # Detect input backend (for augment mode, use input format; for sample mode, use output_format)
+    if mode == "sample":
+        input_backend = output_format
+    else:
+        input_backend = detect_backend(df)
+    
+    # Augment using Polars-only architecture
+    try:
+        # 1. Convert to Polars via Arrow bridge
+        df_polars = to_polars(df, input_backend if mode != "sample" else "polars")
+        
+        # Memory cleanup: delete original if converted
+        if mode != "sample" and input_backend != 'polars':
+            del df
+            import gc
+            gc.collect()
+        
+        # 2. Process augmentation in Polars
+        result_polars = _augment_polars_engine(df_polars, actual_n_rows, strategy_dict, seed)
+        
+        # Memory cleanup: delete intermediate Polars DataFrame
+        del df_polars
+        import gc
+        gc.collect()
+        
+        # 3. Convert back to target format
+        # In augment mode: match input format
+        # In sample mode: use output_format parameter
+        target_backend = output_format if mode == "sample" else input_backend
+        result_df = from_polars(result_polars, target_backend)
+        
+        # Final memory cleanup
+        del result_polars
+        import gc
+        gc.collect()
+        
+        return result_df
+        
+    except Exception as e:
+        if isinstance(e, (ValidationError, AugmentError)):
+            raise
+        raise AugmentError(f"Augmentation failed: {e}")