daskdf.py

import dask.dataframe as dd
import pandas as pd
from typing import Callable, Dict
from enum import Enum


class CalculationRule(Enum):
    """Available calculation rules"""
    PERCENTAGE_CHANGE = "percentage_change"
    MOVING_AVERAGE = "moving_average"
    CUMULATIVE_SUM = "cumulative_sum"
    STANDARDIZE = "standardize"
    RANK = "rank"
    CATEGORIZE = "categorize"


class DaskCalculationService:
    """Service for applying calculation rules to Dask DataFrames"""
    
    def __init__(self):
        self.rules: Dict[CalculationRule, Callable] = {
            CalculationRule.PERCENTAGE_CHANGE: self._percentage_change,
            CalculationRule.MOVING_AVERAGE: self._moving_average,
            CalculationRule.CUMULATIVE_SUM: self._cumulative_sum,
            CalculationRule.STANDARDIZE: self._standardize,
            CalculationRule.RANK: self._rank,
            CalculationRule.CATEGORIZE: self._categorize,
        }
    
    def apply_rule(self, ddf: dd.DataFrame, rule: CalculationRule, 
                   column: str, **kwargs) -> dd.DataFrame:
        """
        Apply a calculation rule to a Dask DataFrame
        
        Args:
            ddf: Dask DataFrame
            rule: CalculationRule to apply
            column: Column name to apply the rule to
            **kwargs: Additional parameters for specific rules
        
        Returns:
            Dask DataFrame with new calculated column
        """
        if rule not in self.rules:
            raise ValueError(f"Unknown rule: {rule}")
        
        return self.rules[rule](ddf, column, **kwargs)
    
    def _percentage_change(self, ddf: dd.DataFrame, column: str, 
                          periods: int = 1) -> dd.DataFrame:
        """Calculate percentage change"""
        new_col = f"{column}_pct_change"
        ddf[new_col] = ddf[column].pct_change(periods=periods) * 100
        return ddf
    
    def _moving_average(self, ddf: dd.DataFrame, column: str, 
                       window: int = 3) -> dd.DataFrame:
        """Calculate moving average"""
        new_col = f"{column}_ma_{window}"
        ddf[new_col] = ddf[column].rolling(window=window).mean()
        return ddf
    
    def _cumulative_sum(self, ddf: dd.DataFrame, column: str) -> dd.DataFrame:
        """Calculate cumulative sum"""
        new_col = f"{column}_cumsum"
        ddf[new_col] = ddf[column].cumsum()
        return ddf
    
    def _standardize(self, ddf: dd.DataFrame, column: str) -> dd.DataFrame:
        """Standardize values (z-score)"""
        new_col = f"{column}_standardized"
        mean = ddf[column].mean()
        std = ddf[column].std()
        ddf[new_col] = (ddf[column] - mean) / std
        return ddf
    
    def _rank(self, ddf: dd.DataFrame, column: str, 
              ascending: bool = True) -> dd.DataFrame:
        """Rank values"""
        new_col = f"{column}_rank"
        # Convert to pandas for ranking, then back to dask
        pdf = ddf.compute()
        pdf[new_col] = pdf[column].rank(ascending=ascending)
        return dd.from_pandas(pdf, npartitions=ddf.npartitions)
    
    def _categorize(self, ddf: dd.DataFrame, column: str, 
                   bins: int = 3, labels: list = None) -> dd.DataFrame:
        """Categorize values into bins"""
        new_col = f"{column}_category"
        if labels is None:
            labels = [f"Level_{i+1}" for i in range(bins)]
        
        # Use qcut for quantile-based binning
        ddf[new_col] = ddf[column].map_partitions(
            lambda part: pd.qcut(part, q=bins, labels=labels, duplicates='drop'),
            meta=(new_col, 'object')
        )
        return ddf


# Example usage
if __name__ == "__main__":
    # Create sample data
    data = {
        'date': pd.date_range('2024-01-01', periods=100),
        'sales': range(100, 200),
        'revenue': [x * 1.5 + i * 0.1 for i, x in enumerate(range(100, 200))]
    }
    pdf = pd.DataFrame(data)
    
    # Convert to Dask DataFrame
    ddf = dd.from_pandas(pdf, npartitions=4)
    
    # Initialize service
    service = DaskCalculationService()
    
    # Apply different rules
    print("Applying percentage change...")
    ddf = service.apply_rule(ddf, CalculationRule.PERCENTAGE_CHANGE, 
                             'sales', periods=1)
    
    print("Applying moving average...")
    ddf = service.apply_rule(ddf, CalculationRule.MOVING_AVERAGE, 
                             'revenue', window=5)
    
    print("Applying cumulative sum...")
    ddf = service.apply_rule(ddf, CalculationRule.CUMULATIVE_SUM, 'sales')
    
    print("Applying standardization...")
    ddf = service.apply_rule(ddf, CalculationRule.STANDARDIZE, 'revenue')
    
    print("Applying rank...")
    ddf = service.apply_rule(ddf, CalculationRule.RANK, 'sales')
    
    print("Applying categorization...")
    ddf = service.apply_rule(ddf, CalculationRule.CATEGORIZE, 'revenue', 
                             bins=3, labels=['Low', 'Medium', 'High'])
    
    # Compute and display results
    result = ddf.compute()
    print("\nFirst 10 rows of results:")
    print(result.head(10))
    print("\nColumns:", result.columns.tolist())