Advanced PySpark Design Patterns: Real-World Implementation Examples

 1from abc import ABC, abstractmethod
 2from typing import Dict, Any
 3from pyspark.sql import DataFrame, SparkSession
 4
 5class DataProcessingStrategy(ABC):
 6    """Abstract strategy for data processing."""
 7    
 8    @abstractmethod
 9    def process(self, data: DataFrame) -> DataFrame:
10        pass
11    
12    @abstractmethod
13    def get_processing_info(self) -> Dict[str, Any]:
14        pass
15
16class AggregationStrategy(DataProcessingStrategy):
17    """Strategy for aggregation-based processing."""
18    
19    def __init__(self, group_by_cols: list, agg_cols: Dict[str, str]):
20        self.group_by_cols = group_by_cols
21        self.agg_cols = agg_cols
22    
23    def process(self, data: DataFrame) -> DataFrame:
24        return data.groupBy(self.group_by_cols).agg(self.agg_cols)
25    
26    def get_processing_info(self) -> Dict[str, Any]:
27        return {
28            "strategy_type": "aggregation",
29            "group_by_columns": self.group_by_cols,
30            "aggregation_columns": self.agg_cols
31        }
32
33class FilteringStrategy(DataProcessingStrategy):
34    """Strategy for filtering-based processing."""
35    
36    def __init__(self, filter_condition: str):
37        self.filter_condition = filter_condition
38    
39    def process(self, data: DataFrame) -> DataFrame:
40        return data.filter(self.filter_condition)
41    
42    def get_processing_info(self) -> Dict[str, Any]:
43        return {
44            "strategy_type": "filtering",
45            "filter_condition": self.filter_condition
46        }
47
48class DataProcessor:
49    """Context class that uses different processing strategies."""
50    
51    def __init__(self, strategy: DataProcessingStrategy):
52        self.strategy = strategy
53    
54    def set_strategy(self, strategy: DataProcessingStrategy):
55        """Change the processing strategy at runtime."""
56        self.strategy = strategy
57    
58    def process_data(self, data: DataFrame) -> DataFrame:
59        """Process data using the current strategy."""
60        return self.strategy.process(data)
61    
62    def get_strategy_info(self) -> Dict[str, Any]:
63        """Get information about the current strategy."""
64        return self.strategy.get_processing_info()
65
66# Usage example
67spark = SparkSession.builder.appName("StrategyPattern").getOrCreate()
68
69# Sample data
70data = spark.createDataFrame([
71    (1, "A", 100), (1, "B", 200), (2, "A", 150), (2, "B", 250)
72], ["id", "category", "value"])
73
74# Use aggregation strategy
75agg_strategy = AggregationStrategy(
76    group_by_cols=["id"],
77    agg_cols={"value": "sum", "category": "count"}
78)
79processor = DataProcessor(agg_strategy)
80result = processor.process_data(data)
81result.show()
82
83# Switch to filtering strategy
84filter_strategy = FilteringStrategy("value > 150")
85processor.set_strategy(filter_strategy)
86filtered_result = processor.process_data(data)
87filtered_result.show()

  1from abc import ABC, abstractmethod
  2from typing import Callable, Any
  3from functools import wraps
  4import time
  5import logging
  6
  7class DataTransformation(ABC):
  8    """Abstract base class for data transformations."""
  9    
 10    @abstractmethod
 11    def transform(self, data: DataFrame) -> DataFrame:
 12        pass
 13
 14class BaseTransformation(DataTransformation):
 15    """Base transformation class."""
 16    
 17    def __init__(self, name: str):
 18        self.name = name
 19    
 20    def transform(self, data: DataFrame) -> DataFrame:
 21        # Base transformation logic
 22        return data
 23
 24class TransformationDecorator(DataTransformation):
 25    """Base decorator class."""
 26    
 27    def __init__(self, transformation: DataTransformation):
 28        self._transformation = transformation
 29    
 30    def transform(self, data: DataFrame) -> DataFrame:
 31        return self._transformation.transform(data)
 32
 33class LoggingDecorator(TransformationDecorator):
 34    """Decorator that adds logging functionality."""
 35    
 36    def transform(self, data: DataFrame) -> DataFrame:
 37        start_time = time.time()
 38        row_count_before = data.count()
 39        
 40        logging.info(f"Starting transformation: {self._transformation.name}")
 41        logging.info(f"Input rows: {row_count_before}")
 42        
 43        result = self._transformation.transform(data)
 44        
 45        end_time = time.time()
 46        row_count_after = result.count()
 47        
 48        logging.info(f"Completed transformation: {self._transformation.name}")
 49        logging.info(f"Output rows: {row_count_after}")
 50        logging.info(f"Processing time: {end_time - start_time:.2f} seconds")
 51        
 52        return result
 53
 54class ValidationDecorator(TransformationDecorator):
 55    """Decorator that adds data validation."""
 56    
 57    def __init__(self, transformation: DataTransformation, validation_rules: Dict[str, Callable]):
 58        super().__init__(transformation)
 59        self.validation_rules = validation_rules
 60    
 61    def transform(self, data: DataFrame) -> DataFrame:
 62        # Apply validation rules
 63        for column, validation_func in self.validation_rules.items():
 64            if column in data.columns:
 65                invalid_count = data.filter(~validation_func(data[column])).count()
 66                if invalid_count > 0:
 67                    logging.warning(f"Found {invalid_count} invalid values in column {column}")
 68        
 69        return self._transformation.transform(data)
 70
 71class CachingDecorator(TransformationDecorator):
 72    """Decorator that adds caching functionality."""
 73    
 74    def __init__(self, transformation: DataTransformation, cache_name: str):
 75        super().__init__(transformation)
 76        self.cache_name = cache_name
 77    
 78    def transform(self, data: DataFrame) -> DataFrame:
 79        # Check if data is already cached
 80        cached_data = self._get_cached_data()
 81        if cached_data is not None:
 82            logging.info(f"Using cached data for: {self.cache_name}")
 83            return cached_data
 84        
 85        # Perform transformation and cache result
 86        result = self._transformation.transform(data)
 87        self._cache_data(result)
 88        logging.info(f"Cached data for: {self.cache_name}")
 89        
 90        return result
 91    
 92    def _get_cached_data(self) -> DataFrame:
 93        # Implementation for retrieving cached data
 94        # This could use Spark's cache, external cache, or file system
 95        pass
 96    
 97    def _cache_data(self, data: DataFrame):
 98        # Implementation for caching data
 99        pass
100
101# Usage example
102def is_positive(col):
103    return col > 0
104
105def is_not_null(col):
106    return col.isNotNull()
107
108# Create base transformation
109base_transform = BaseTransformation("data_cleaning")
110
111# Add decorators
112validation_rules = {"value": is_positive, "id": is_not_null}
113validated_transform = ValidationDecorator(base_transform, validation_rules)
114logged_transform = LoggingDecorator(validated_transform)
115cached_transform = CachingDecorator(logged_transform, "cleaned_data")
116
117# Use the decorated transformation
118result = cached_transform.transform(data)

  1from abc import ABC, abstractmethod
  2from typing import List, Optional
  3from dataclasses import dataclass
  4from datetime import datetime
  5
  6@dataclass
  7class PipelineCommand:
  8    """Command object that encapsulates a pipeline operation."""
  9    command_id: str
 10    timestamp: datetime
 11    operation_type: str
 12    parameters: Dict[str, Any]
 13    undo_data: Optional[Dict[str, Any]] = None
 14
 15class PipelineOperation(ABC):
 16    """Abstract base class for pipeline operations."""
 17    
 18    @abstractmethod
 19    def execute(self, data: DataFrame) -> DataFrame:
 20        pass
 21    
 22    @abstractmethod
 23    def undo(self, data: DataFrame) -> DataFrame:
 24        pass
 25    
 26    @abstractmethod
 27    def get_command(self) -> PipelineCommand:
 28        pass
 29
 30class FilterOperation(PipelineOperation):
 31    """Concrete operation for filtering data."""
 32    
 33    def __init__(self, filter_condition: str):
 34        self.filter_condition = filter_condition
 35        self.command_id = f"filter_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
 36    
 37    def execute(self, data: DataFrame) -> DataFrame:
 38        # Store original row count for undo
 39        original_count = data.count()
 40        result = data.filter(self.filter_condition)
 41        
 42        # Store undo information
 43        self.undo_data = {
 44            "original_count": original_count,
 45            "filtered_count": result.count()
 46        }
 47        
 48        return result
 49    
 50    def undo(self, data: DataFrame) -> DataFrame:
 51        # In a real implementation, you might need to store the original data
 52        # or implement a more sophisticated undo mechanism
 53        logging.info(f"Undoing filter operation: {self.filter_condition}")
 54        return data
 55    
 56    def get_command(self) -> PipelineCommand:
 57        return PipelineCommand(
 58            command_id=self.command_id,
 59            timestamp=datetime.now(),
 60            operation_type="filter",
 61            parameters={"filter_condition": self.filter_condition},
 62            undo_data=self.undo_data
 63        )
 64
 65class AggregationOperation(PipelineOperation):
 66    """Concrete operation for aggregating data."""
 67    
 68    def __init__(self, group_by_cols: list, agg_cols: Dict[str, str]):
 69        self.group_by_cols = group_by_cols
 70        self.agg_cols = agg_cols
 71        self.command_id = f"agg_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
 72    
 73    def execute(self, data: DataFrame) -> DataFrame:
 74        # Store original schema for undo
 75        original_schema = data.schema
 76        result = data.groupBy(self.group_by_cols).agg(self.agg_cols)
 77        
 78        # Store undo information
 79        self.undo_data = {
 80            "original_schema": original_schema,
 81            "group_by_columns": self.group_by_cols,
 82            "aggregation_columns": self.agg_cols
 83        }
 84        
 85        return result
 86    
 87    def undo(self, data: DataFrame) -> DataFrame:
 88        # Note: Aggregation undo is complex and might not be fully reversible
 89        logging.warning("Aggregation undo is not fully supported")
 90        return data
 91    
 92    def get_command(self) -> PipelineCommand:
 93        return PipelineCommand(
 94            command_id=self.command_id,
 95            timestamp=datetime.now(),
 96            operation_type="aggregation",
 97            parameters={
 98                "group_by_columns": self.group_by_cols,
 99                "aggregation_columns": self.agg_cols
100            },
101            undo_data=self.undo_data
102        )
103
104class PipelineInvoker:
105    """Invoker class that manages command execution."""
106    
107    def __init__(self):
108        self.command_history: List[PipelineCommand] = []
109        self.undo_stack: List[PipelineOperation] = []
110    
111    def execute_operation(self, operation: PipelineOperation, data: DataFrame) -> DataFrame:
112        """Execute a pipeline operation and store it in history."""
113        result = operation.execute(data)
114        command = operation.get_command()
115        
116        self.command_history.append(command)
117        self.undo_stack.append(operation)
118        
119        logging.info(f"Executed operation: {command.operation_type}")
120        return result
121    
122    def undo_last_operation(self, data: DataFrame) -> DataFrame:
123        """Undo the last executed operation."""
124        if not self.undo_stack:
125            logging.warning("No operations to undo")
126            return data
127        
128        operation = self.undo_stack.pop()
129        command = self.command_history.pop()
130        
131        result = operation.undo(data)
132        logging.info(f"Undid operation: {command.operation_type}")
133        
134        return result
135    
136    def get_command_history(self) -> List[PipelineCommand]:
137        """Get the history of executed commands."""
138        return self.command_history.copy()
139
140# Usage example
141invoker = PipelineInvoker()
142
143# Execute operations
144filter_op = FilterOperation("value > 100")
145data = invoker.execute_operation(filter_op, data)
146
147agg_op = AggregationOperation(["id"], {"value": "sum"})
148data = invoker.execute_operation(agg_op, data)
149
150# Undo last operation
151data = invoker.undo_last_operation(data)
152
153# View command history
154for command in invoker.get_command_history():
155    print(f"{command.timestamp}: {command.operation_type}")

  1from abc import ABC, abstractmethod
  2from typing import List, Dict, Any
  3
  4class DataPipelineTemplate(ABC):
  5    """Template class for data pipeline workflows."""
  6    
  7    def run_pipeline(self, input_data: DataFrame) -> DataFrame:
  8        """Template method that defines the pipeline workflow."""
  9        try:
 10            # Step 1: Validate input
 11            validated_data = self.validate_input(input_data)
 12            
 13            # Step 2: Preprocess data
 14            preprocessed_data = self.preprocess_data(validated_data)
 15            
 16            # Step 3: Apply business logic
 17            processed_data = self.apply_business_logic(preprocessed_data)
 18            
 19            # Step 4: Post-process data
 20            postprocessed_data = self.postprocess_data(processed_data)
 21            
 22            # Step 5: Validate output
 23            final_data = self.validate_output(postprocessed_data)
 24            
 25            # Step 6: Log results
 26            self.log_results(final_data)
 27            
 28            return final_data
 29            
 30        except Exception as e:
 31            self.handle_error(e)
 32            raise
 33    
 34    @abstractmethod
 35    def validate_input(self, data: DataFrame) -> DataFrame:
 36        """Validate input data - must be implemented by subclasses."""
 37        pass
 38    
 39    @abstractmethod
 40    def preprocess_data(self, data: DataFrame) -> DataFrame:
 41        """Preprocess data - must be implemented by subclasses."""
 42        pass
 43    
 44    @abstractmethod
 45    def apply_business_logic(self, data: DataFrame) -> DataFrame:
 46        """Apply business logic - must be implemented by subclasses."""
 47        pass
 48    
 49    @abstractmethod
 50    def postprocess_data(self, data: DataFrame) -> DataFrame:
 51        """Post-process data - must be implemented by subclasses."""
 52        pass
 53    
 54    @abstractmethod
 55    def validate_output(self, data: DataFrame) -> DataFrame:
 56        """Validate output data - must be implemented by subclasses."""
 57        pass
 58    
 59    def log_results(self, data: DataFrame):
 60        """Default implementation for logging results."""
 61        logging.info(f"Pipeline completed successfully. Output rows: {data.count()}")
 62    
 63    def handle_error(self, error: Exception):
 64        """Default implementation for error handling."""
 65        logging.error(f"Pipeline failed with error: {str(error)}")
 66
 67class SalesDataPipeline(DataPipelineTemplate):
 68    """Concrete implementation for sales data processing."""
 69    
 70    def validate_input(self, data: DataFrame) -> DataFrame:
 71        required_columns = ["sale_id", "product_id", "amount", "date"]
 72        missing_columns = [col for col in required_columns if col not in data.columns]
 73        
 74        if missing_columns:
 75            raise ValueError(f"Missing required columns: {missing_columns}")
 76        
 77        # Check for null values in critical columns
 78        null_counts = data.select([F.count(F.when(F.col(c).isNull(), c)).alias(c) for c in required_columns])
 79        null_counts.show()
 80        
 81        return data
 82    
 83    def preprocess_data(self, data: DataFrame) -> DataFrame:
 84        # Remove duplicates
 85        data = data.dropDuplicates(["sale_id"])
 86        
 87        # Convert date column to proper format
 88        data = data.withColumn("date", F.to_date("date"))
 89        
 90        # Add calculated columns
 91        data = data.withColumn("year", F.year("date"))
 92        data = data.withColumn("month", F.month("date"))
 93        
 94        return data
 95    
 96    def apply_business_logic(self, data: DataFrame) -> DataFrame:
 97        # Calculate daily sales totals
 98        daily_sales = data.groupBy("date").agg(
 99            F.sum("amount").alias("daily_total"),
100            F.count("*").alias("daily_transactions")
101        )
102        
103        # Calculate product performance
104        product_performance = data.groupBy("product_id").agg(
105            F.sum("amount").alias("total_revenue"),
106            F.count("*").alias("total_sales"),
107            F.avg("amount").alias("avg_sale_amount")
108        )
109        
110        return data.join(daily_sales, "date", "left").join(product_performance, "product_id", "left")
111    
112    def postprocess_data(self, data: DataFrame) -> DataFrame:
113        # Round monetary values
114        data = data.withColumn("amount", F.round("amount", 2))
115        data = data.withColumn("daily_total", F.round("daily_total", 2))
116        data = data.withColumn("total_revenue", F.round("total_revenue", 2))
117        data = data.withColumn("avg_sale_amount", F.round("avg_sale_amount", 2))
118        
119        # Add performance indicators
120        data = data.withColumn("performance_ratio", 
121                              F.when(F.col("amount") > F.col("avg_sale_amount"), "above_avg")
122                              .otherwise("below_avg"))
123        
124        return data
125    
126    def validate_output(self, data: DataFrame) -> DataFrame:
127        # Check for negative amounts
128        negative_count = data.filter(F.col("amount") < 0).count()
129        if negative_count > 0:
130            logging.warning(f"Found {negative_count} records with negative amounts")
131        
132        # Check for reasonable date ranges
133        min_date = data.agg(F.min("date")).collect()[0][0]
134        max_date = data.agg(F.max("date")).collect()[0][0]
135        logging.info(f"Data date range: {min_date} to {max_date}")
136        
137        return data
138
139class LogDataPipeline(DataPipelineTemplate):
140    """Concrete implementation for log data processing."""
141    
142    def validate_input(self, data: DataFrame) -> DataFrame:
143        required_columns = ["timestamp", "level", "message"]
144        missing_columns = [col for col in required_columns if col not in data.columns]
145        
146        if missing_columns:
147            raise ValueError(f"Missing required columns: {missing_columns}")
148        
149        return data
150    
151    def preprocess_data(self, data: DataFrame) -> DataFrame:
152        # Parse timestamp
153        data = data.withColumn("timestamp", F.to_timestamp("timestamp"))
154        
155        # Extract date components
156        data = data.withColumn("date", F.date("timestamp"))
157        data = data.withColumn("hour", F.hour("timestamp"))
158        
159        # Clean message column
160        data = data.withColumn("message", F.trim("message"))
161        
162        return data
163    
164    def apply_business_logic(self, data: DataFrame) -> DataFrame:
165        # Calculate error rates by hour
166        error_rates = data.groupBy("date", "hour").agg(
167            F.count("*").alias("total_logs"),
168            F.sum(F.when(F.col("level") == "ERROR", 1).otherwise(0)).alias("error_count")
169        ).withColumn("error_rate", F.col("error_count") / F.col("total_logs"))
170        
171        # Calculate level distribution
172        level_distribution = data.groupBy("level").count()
173        
174        return data.join(error_rates, ["date", "hour"], "left")
175    
176    def postprocess_data(self, data: DataFrame) -> DataFrame:
177        # Add severity indicators
178        data = data.withColumn("severity", 
179                              F.when(F.col("level") == "ERROR", "high")
180                              .when(F.col("level") == "WARN", "medium")
181                              .otherwise("low"))
182        
183        return data
184    
185    def validate_output(self, data: DataFrame) -> DataFrame:
186        # Check for reasonable error rates
187        high_error_rate = data.filter(F.col("error_rate") > 0.5).count()
188        if high_error_rate > 0:
189            logging.warning(f"Found {high_error_rate} hours with high error rates")
190        
191        return data
192
193# Usage example
194# Sales pipeline
195sales_pipeline = SalesDataPipeline()
196sales_result = sales_pipeline.run_pipeline(sales_data)
197
198# Log pipeline
199log_pipeline = LogDataPipeline()
200log_result = log_pipeline.run_pipeline(log_data)