cython_impl.pyx

# cython: language_level=3
# cython: boundscheck=False
# cython: wraparound=False
# cython: cdivision=True

"""Cython implementation for DataFrame generation and parquet writing."""
import pandas as pd
import numpy as np
cimport numpy as cnp
from libc.math cimport sqrt


def generate_dataframe_cython(int num_rows=1_000_000):
    """
    Generate a pandas DataFrame with random data using Cython optimizations.
    
    Args:
        num_rows: Number of rows to generate (default: 1 million)
        
    Returns:
        pandas DataFrame with multiple columns of different types
    """
    cdef int i
    cdef cnp.ndarray[cnp.int64_t, ndim=1] ids
    cdef cnp.ndarray[cnp.float64_t, ndim=1] values1
    cdef cnp.ndarray[cnp.float64_t, ndim=1] values2
    cdef cnp.ndarray[cnp.uint8_t, ndim=1, cast=True] flags
    
    # Pre-allocate arrays
    ids = np.arange(num_rows, dtype=np.int64)
    values1 = np.empty(num_rows, dtype=np.float64)
    values2 = np.empty(num_rows, dtype=np.float64)
    flags = np.empty(num_rows, dtype=np.uint8)
    
    # Fill arrays with computed values
    for i in range(num_rows):
        values1[i] = <double>i * 2.5
        values2[i] = sqrt(<double>i)
        flags[i] = <unsigned char>(i % 2)
    
    # Create category array
    categories = [f'cat_{i % 10}' for i in range(num_rows)]
    
    data = {
        'id': ids,
        'value1': values1,
        'value2': values2,
        'category': categories,
        'flag': flags.astype(bool),
    }
    
    return pd.DataFrame(data)


def write_parquet_cython(df, filename: str):
    """
    Write DataFrame to parquet file with snappy compression.
    
    Args:
        df: DataFrame to write
        filename: Output parquet file path
    """
    df.to_parquet(filename, compression='snappy', engine='pyarrow', index=False)


def run_benchmark_cython(int num_rows=1_000_000, str output_file='output_cython.parquet'):
    """
    Run the complete benchmark: generate DataFrame and write to parquet.
    
    Args:
        num_rows: Number of rows to generate
        output_file: Output parquet file path
        
    Returns:
        Tuple of (generation_time, write_time) in seconds
    """
    import time
    
    # Time DataFrame generation
    start = time.perf_counter()
    df = generate_dataframe_cython(num_rows)
    gen_time = time.perf_counter() - start
    
    # Time parquet writing
    start = time.perf_counter()
    write_parquet_cython(df, output_file)
    write_time = time.perf_counter() - start
    
    return gen_time, write_time