benchmark.py

#!/usr/bin/env python3
"""
Main benchmark script to compare Python vs Cython performance.

This script benchmarks the generation of a 1 million row pandas DataFrame
and writing it to a parquet file with snappy compression.
"""
import os
import sys
import time
from typing import Dict, List, Tuple

# Import implementations
import python_impl
import cython_impl


def run_multiple_benchmarks(num_iterations: int = 5, num_rows: int = 1_000_000) -> Dict[str, List[Tuple[float, float]]]:
    """
    Run benchmarks multiple times to get average results.
    
    Args:
        num_iterations: Number of times to run each benchmark
        num_rows: Number of rows in the DataFrame
        
    Returns:
        Dictionary with results for each implementation
    """
    results = {
        'python': [],
        'cython': []
    }
    
    print(f"Running benchmarks with {num_rows:,} rows...")
    print(f"Each benchmark will be run {num_iterations} times for accuracy.\n")
    
    for i in range(num_iterations):
        print(f"Iteration {i + 1}/{num_iterations}")
        
        # Python benchmark
        print("  Running Pure Python implementation...")
        py_gen_time, py_write_time = python_impl.run_benchmark(
            num_rows, f'output_python_{i}.parquet'
        )
        results['python'].append((py_gen_time, py_write_time))
        print(f"    Generation: {py_gen_time:.4f}s, Writing: {py_write_time:.4f}s, Total: {py_gen_time + py_write_time:.4f}s")
        
        # Cython benchmark
        print("  Running Cython implementation...")
        cy_gen_time, cy_write_time = cython_impl.run_benchmark_cython(
            num_rows, f'output_cython_{i}.parquet'
        )
        results['cython'].append((cy_gen_time, cy_write_time))
        print(f"    Generation: {cy_gen_time:.4f}s, Writing: {cy_write_time:.4f}s, Total: {cy_gen_time + cy_write_time:.4f}s")
        print()
    
    return results


def calculate_statistics(results: Dict[str, List[Tuple[float, float]]]) -> Dict[str, Dict[str, float]]:
    """
    Calculate average statistics from benchmark results.
    
    Args:
        results: Raw benchmark results
        
    Returns:
        Dictionary with calculated statistics
    """
    stats = {}
    
    for impl_name, times_list in results.items():
        gen_times = [t[0] for t in times_list]
        write_times = [t[1] for t in times_list]
        total_times = [t[0] + t[1] for t in times_list]
        
        stats[impl_name] = {
            'avg_gen_time': sum(gen_times) / len(gen_times),
            'avg_write_time': sum(write_times) / len(write_times),
            'avg_total_time': sum(total_times) / len(total_times),
            'min_total_time': min(total_times),
            'max_total_time': max(total_times),
        }
    
    return stats


def print_results(stats: Dict[str, Dict[str, float]]) -> None:
    """
    Print formatted benchmark results.
    
    Args:
        stats: Calculated statistics
    """
    print("=" * 80)
    print("BENCHMARK RESULTS")
    print("=" * 80)
    print()
    
    print("Pure Python Implementation:")
    print(f"  Average Generation Time:  {stats['python']['avg_gen_time']:.4f}s")
    print(f"  Average Writing Time:     {stats['python']['avg_write_time']:.4f}s")
    print(f"  Average Total Time:       {stats['python']['avg_total_time']:.4f}s")
    print(f"  Min Total Time:           {stats['python']['min_total_time']:.4f}s")
    print(f"  Max Total Time:           {stats['python']['max_total_time']:.4f}s")
    print()
    
    print("Cython Implementation:")
    print(f"  Average Generation Time:  {stats['cython']['avg_gen_time']:.4f}s")
    print(f"  Average Writing Time:     {stats['cython']['avg_write_time']:.4f}s")
    print(f"  Average Total Time:       {stats['cython']['avg_total_time']:.4f}s")
    print(f"  Min Total Time:           {stats['cython']['min_total_time']:.4f}s")
    print(f"  Max Total Time:           {stats['cython']['max_total_time']:.4f}s")
    print()
    
    print("Performance Comparison:")
    gen_speedup = stats['python']['avg_gen_time'] / stats['cython']['avg_gen_time']
    write_speedup = stats['python']['avg_write_time'] / stats['cython']['avg_write_time']
    total_speedup = stats['python']['avg_total_time'] / stats['cython']['avg_total_time']
    
    print(f"  Generation Speedup:       {gen_speedup:.2f}x")
    print(f"  Writing Speedup:          {write_speedup:.2f}x")
    print(f"  Total Speedup:            {total_speedup:.2f}x")
    print()
    
    if total_speedup > 1:
        improvement = ((total_speedup - 1) * 100)
        print(f"  Cython is {improvement:.1f}% faster than Pure Python")
    else:
        degradation = ((1 - total_speedup) * 100)
        print(f"  Cython is {degradation:.1f}% slower than Pure Python")
    print()
    print("=" * 80)


def cleanup_output_files() -> None:
    """Remove generated parquet files."""
    import glob
    for pattern in ['output_python_*.parquet', 'output_cython_*.parquet']:
        for filepath in glob.glob(pattern):
            try:
                os.remove(filepath)
            except OSError:
                pass


def main():
    """Main entry point for the benchmark."""
    print("=" * 80)
    print("Python vs Cython Benchmark")
    print("DataFrame Generation and Parquet Writing Performance")
    print("=" * 80)
    print()
    
    # Configuration
    num_iterations = 5
    num_rows = 1_000_000
    
    # Run benchmarks
    results = run_multiple_benchmarks(num_iterations, num_rows)
    
    # Calculate statistics
    stats = calculate_statistics(results)
    
    # Print results
    print_results(stats)
    
    # Cleanup
    print("Cleaning up output files...")
    cleanup_output_files()
    print("Done!")


if __name__ == '__main__':
    main()