write_benchmark.py

#!/usr/bin/env python3
"""
Benchmark script for Ladybug write performance.

Compares performance of:
1. INSERT via CREATE statements (node tables)
2. INSERT via CREATE statements (rel tables)
3. Bulk loading via COPY from Parquet (node tables)
4. Bulk loading via COPY from Parquet (rel tables)

Usage:
    python write_benchmark.py [--num-rows NUM_ROWS] [--batch-size BATCH_SIZE] [--num-threads NUM_THREADS]
"""

import argparse
import os
import shutil
import sys
import time
from pathlib import Path
from typing import Dict, List, Tuple

try:
    import psutil
except ImportError:
    psutil = None

import pyarrow as pa
import pyarrow.parquet as pq

# Import Ladybug
try:
    import real_ladybug as lb
except ImportError:
    lb = None


class BenchmarkResults:
    """Store and format benchmark results."""
    
    def __init__(self):
        self.results = []
    
    def add_result(
        self,
        name: str,
        operation: str,
        table_type: str,
        num_rows: int,
        duration_ms: float,
        throughput: float,
        memory_mb: float
    ):
        """Add a benchmark result."""
        self.results.append({
            'name': name,
            'operation': operation,
            'table_type': table_type,
            'num_rows': num_rows,
            'duration_ms': duration_ms,
            'throughput': throughput,
            'memory_mb': memory_mb
        })
    
    def print_results(self):
        """Print formatted results."""
        print("\n" + "=" * 100)
        print("BENCHMARK RESULTS")
        print("=" * 100)
        print(f"{'Test Name':<35} {'Type':<8} {'Rows':<12} {'Time (ms)':<12} "
              f"{'Throughput (rows/s)':<22} {'Memory (MB)':<12}")
        print("-" * 100)
        
        for result in self.results:
            print(f"{result['name']:<35} {result['table_type']:<8} "
                  f"{result['num_rows']:<12,} {result['duration_ms']:<12,.2f} "
                  f"{result['throughput']:<22,.2f} {result['memory_mb']:<12,.2f}")
        
        print("=" * 100)
        
        # Print comparison summary
        print("\nPERFORMANCE COMPARISON:")
        print("-" * 60)
        
        # Group by table type
        node_results = [r for r in self.results if r['table_type'] == 'NODE']
        rel_results = [r for r in self.results if r['table_type'] == 'REL']
        
        for results_group, table_type in [(node_results, 'NODE'), (rel_results, 'REL')]:
            if not results_group:
                continue
            
            print(f"\n{table_type} Tables:")
            insert_result = next((r for r in results_group if r['operation'] == 'INSERT'), None)
            parquet_result = next((r for r in results_group if r['operation'] == 'PARQUET'), None)
            
            if insert_result and parquet_result:
                speedup = parquet_result['throughput'] / insert_result['throughput']
                print(f"  INSERT:  {insert_result['throughput']:,.2f} rows/s")
                print(f"  PARQUET: {parquet_result['throughput']:,.2f} rows/s")
                print(f"  Speedup: {speedup:.2f}x (Parquet is {speedup:.2f}x faster)")
        
        print("\n")


def measure_memory() -> float:
    """Get current memory usage in MB."""
    if psutil is None:
        return 0.0
    process = psutil.Process()
    return process.memory_info().rss / (1024 * 1024)


def benchmark_node_insert(
    conn,  # lb.Connection
    num_rows: int,
    batch_size: int
) -> Tuple[float, float]:
    """Benchmark INSERT operations for node tables."""
    print(f"\nBenchmarking NODE table INSERT ({num_rows:,} rows)...")
    
    # Create node table
    conn.execute("""
        CREATE NODE TABLE BenchmarkNode(
            id INT64,
            name STRING,
            value DOUBLE,
            timestamp INT64,
            is_active BOOLEAN,
            tags STRING[],
            PRIMARY KEY (id)
        )
    """)
    
    mem_start = measure_memory()
    start_time = time.time()
    
    # Insert in batches using UNWIND
    for batch_start in range(0, num_rows, batch_size):
        batch_end = min(batch_start + batch_size, num_rows)
        batch_data = [
            {
                'id': i,
                'name': f'Node_{i}',
                'value': float(i) * 1.5,
                'timestamp': 1700000000 + i,
                'is_active': i % 2 == 0,
                'tags': [f'tag_{i % 10}', f'tag_{i % 20}']
            }
            for i in range(batch_start, batch_end)
        ]
        
        conn.execute("""
            UNWIND $batch AS row
            CREATE (:BenchmarkNode {
                id: row.id,
                name: row.name,
                value: row.value,
                timestamp: row.timestamp,
                is_active: row.is_active,
                tags: row.tags
            })
        """, {'batch': batch_data})
        
        if (batch_start // batch_size) % 10 == 0:
            print(f"  Inserted {batch_end:,} / {num_rows:,} rows...")
    
    end_time = time.time()
    mem_end = measure_memory()
    
    duration_ms = (end_time - start_time) * 1000
    memory_mb = mem_end - mem_start
    
    # Verify row count
    result = conn.execute("MATCH (n:BenchmarkNode) RETURN COUNT(n) AS count")
    count = result.get_next()[0]
    assert count == num_rows, f"Expected {num_rows} rows, got {count}"
    
    print(f"  ✓ Completed in {duration_ms:.2f} ms")
    print(f"  ✓ Throughput: {num_rows / (duration_ms / 1000):.2f} rows/s")
    
    return duration_ms, memory_mb


def benchmark_node_parquet(
    conn,  # lb.Connection
    num_rows: int,
    tmp_dir: Path
) -> Tuple[float, float]:
    """Benchmark bulk loading from Parquet for node tables."""
    print(f"\nBenchmarking NODE table PARQUET bulk load ({num_rows:,} rows)...")
    
    # Create node table
    conn.execute("""
        CREATE NODE TABLE BenchmarkNodeParquet(
            id INT64,
            name STRING,
            value DOUBLE,
            timestamp INT64,
            is_active BOOLEAN,
            tags STRING[],
            PRIMARY KEY (id)
        )
    """)
    
    # Generate Parquet file
    parquet_file = tmp_dir / "nodes.parquet"
    print(f"  Generating Parquet file...")
    
    table = pa.Table.from_arrays([
        pa.array(range(num_rows), type=pa.int64()),
        pa.array([f'Node_{i}' for i in range(num_rows)], type=pa.string()),
        pa.array([float(i) * 1.5 for i in range(num_rows)], type=pa.float64()),
        pa.array([1700000000 + i for i in range(num_rows)], type=pa.int64()),
        pa.array([i % 2 == 0 for i in range(num_rows)], type=pa.bool_()),
        pa.array([[f'tag_{i % 10}', f'tag_{i % 20}'] for i in range(num_rows)], 
                 type=pa.list_(pa.string()))
    ], names=['id', 'name', 'value', 'timestamp', 'is_active', 'tags'])
    
    pq.write_table(table, parquet_file)
    
    mem_start = measure_memory()
    start_time = time.time()
    
    # Load from Parquet
    conn.execute(f'COPY BenchmarkNodeParquet FROM "{parquet_file}"')
    
    end_time = time.time()
    mem_end = measure_memory()
    
    duration_ms = (end_time - start_time) * 1000
    memory_mb = mem_end - mem_start
    
    # Verify row count
    result = conn.execute("MATCH (n:BenchmarkNodeParquet) RETURN COUNT(n) AS count")
    count = result.get_next()[0]
    assert count == num_rows, f"Expected {num_rows} rows, got {count}"
    
    print(f"  ✓ Completed in {duration_ms:.2f} ms")
    print(f"  ✓ Throughput: {num_rows / (duration_ms / 1000):.2f} rows/s")
    
    return duration_ms, memory_mb


def benchmark_rel_insert(
    conn,  # lb.Connection
    num_rows: int,
    batch_size: int,
    num_nodes: int
) -> Tuple[float, float]:
    """Benchmark INSERT operations for rel tables."""
    print(f"\nBenchmarking REL table INSERT ({num_rows:,} relationships)...")
    
    # Create source and target node tables
    conn.execute("""
        CREATE NODE TABLE SourceNode(
            id INT64,
            PRIMARY KEY (id)
        )
    """)
    
    conn.execute("""
        CREATE NODE TABLE TargetNode(
            id INT64,
            PRIMARY KEY (id)
        )
    """)
    
    # Create rel table
    conn.execute("""
        CREATE REL TABLE BenchmarkRel(
            FROM SourceNode TO TargetNode,
            weight DOUBLE,
            label STRING
        )
    """)
    
    # Insert nodes
    print(f"  Creating {num_nodes:,} source and target nodes...")
    for batch_start in range(0, num_nodes, batch_size):
        batch_end = min(batch_start + batch_size, num_nodes)
        conn.execute(f"""
            UNWIND RANGE({batch_start}, {batch_end - 1}) AS x
            CREATE (:SourceNode {{id: x}})
        """)
        conn.execute(f"""
            UNWIND RANGE({batch_start}, {batch_end - 1}) AS x
            CREATE (:TargetNode {{id: x}})
        """)
    
    mem_start = measure_memory()
    start_time = time.time()
    
    # Insert relationships in batches
    for batch_start in range(0, num_rows, batch_size):
        batch_end = min(batch_start + batch_size, num_rows)
        batch_data = [
            {
                'src': i % num_nodes,
                'dst': (i + 1) % num_nodes,
                'weight': float(i) * 0.5,
                'label': f'rel_{i % 100}'
            }
            for i in range(batch_start, batch_end)
        ]
        
        conn.execute("""
            UNWIND $batch AS row
            MATCH (s:SourceNode {id: row.src}), (t:TargetNode {id: row.dst})
            CREATE (s)-[:BenchmarkRel {weight: row.weight, label: row.label}]->(t)
        """, {'batch': batch_data})
        
        if (batch_start // batch_size) % 10 == 0:
            print(f"  Inserted {batch_end:,} / {num_rows:,} relationships...")
    
    end_time = time.time()
    mem_end = measure_memory()
    
    duration_ms = (end_time - start_time) * 1000
    memory_mb = mem_end - mem_start
    
    # Verify relationship count
    result = conn.execute("MATCH ()-[r:BenchmarkRel]->() RETURN COUNT(r) AS count")
    count = result.get_next()[0]
    assert count == num_rows, f"Expected {num_rows} relationships, got {count}"
    
    print(f"  ✓ Completed in {duration_ms:.2f} ms")
    print(f"  ✓ Throughput: {num_rows / (duration_ms / 1000):.2f} rows/s")
    
    return duration_ms, memory_mb


def benchmark_rel_parquet(
    conn,  # lb.Connection
    num_rows: int,
    num_nodes: int,
    batch_size: int,
    tmp_dir: Path
) -> Tuple[float, float]:
    """Benchmark bulk loading from Parquet for rel tables."""
    print(f"\nBenchmarking REL table PARQUET bulk load ({num_rows:,} relationships)...")
    
    # Create source and target node tables
    conn.execute("""
        CREATE NODE TABLE SourceNodeParquet(
            id INT64,
            PRIMARY KEY (id)
        )
    """)
    
    conn.execute("""
        CREATE NODE TABLE TargetNodeParquet(
            id INT64,
            PRIMARY KEY (id)
        )
    """)
    
    # Create rel table
    conn.execute("""
        CREATE REL TABLE BenchmarkRelParquet(
            FROM SourceNodeParquet TO TargetNodeParquet,
            weight DOUBLE,
            label STRING
        )
    """)
    
    # Insert nodes
    print(f"  Creating {num_nodes:,} source and target nodes...")
    for batch_start in range(0, num_nodes, batch_size):
        batch_end = min(batch_start + batch_size, num_nodes)
        conn.execute(f"""
            UNWIND RANGE({batch_start}, {batch_end - 1}) AS x
            CREATE (:SourceNodeParquet {{id: x}})
        """)
        conn.execute(f"""
            UNWIND RANGE({batch_start}, {batch_end - 1}) AS x
            CREATE (:TargetNodeParquet {{id: x}})
        """)
    
    # Generate Parquet file
    parquet_file = tmp_dir / "rels.parquet"
    print(f"  Generating Parquet file...")
    
    table = pa.Table.from_arrays([
        pa.array([i % num_nodes for i in range(num_rows)], type=pa.int64()),
        pa.array([(i + 1) % num_nodes for i in range(num_rows)], type=pa.int64()),
        pa.array([float(i) * 0.5 for i in range(num_rows)], type=pa.float64()),
        pa.array([f'rel_{i % 100}' for i in range(num_rows)], type=pa.string())
    ], names=['from', 'to', 'weight', 'label'])
    
    pq.write_table(table, parquet_file)
    
    mem_start = measure_memory()
    start_time = time.time()
    
    # Load from Parquet
    conn.execute(f'COPY BenchmarkRelParquet FROM "{parquet_file}"')
    
    end_time = time.time()
    mem_end = measure_memory()
    
    duration_ms = (end_time - start_time) * 1000
    memory_mb = mem_end - mem_start
    
    # Verify relationship count
    result = conn.execute("MATCH ()-[r:BenchmarkRelParquet]->() RETURN COUNT(r) AS count")
    count = result.get_next()[0]
    assert count == num_rows, f"Expected {num_rows} relationships, got {count}"
    
    print(f"  ✓ Completed in {duration_ms:.2f} ms")
    print(f"  ✓ Throughput: {num_rows / (duration_ms / 1000):.2f} rows/s")
    
    return duration_ms, memory_mb


def run_benchmarks(
    num_rows: int,
    batch_size: int,
    num_threads: int,
    db_path: str
):
    """Run all benchmarks."""
    # Check if dependencies are available
    if lb is None:
        print("Error: real_ladybug module not found. Please build and install it first.")
        print("Run: make GEN=Ninja python")
        sys.exit(1)
    
    if psutil is None:
        print("Warning: psutil module not found. Memory measurements will be unavailable.")
        print("Install with: pip install psutil")
    
    print("=" * 100)
    print("LADYBUG WRITE PERFORMANCE BENCHMARK")
    print("=" * 100)
    print(f"Configuration:")
    print(f"  Number of rows:      {num_rows:,}")
    print(f"  Batch size:          {batch_size:,}")
    print(f"  Number of threads:   {num_threads}")
    print(f"  Database path:       {db_path}")
    if psutil:
        print(f"  System memory:       {psutil.virtual_memory().total / (1024**3):.2f} GB")
        print(f"  Available CPUs:      {psutil.cpu_count()}")
    print("=" * 100)
    
    # Create temp directory for parquet files
    tmp_dir = Path("/tmp/ladybug_benchmark")
    if tmp_dir.exists():
        shutil.rmtree(tmp_dir)
    tmp_dir.mkdir(parents=True, exist_ok=True)
    
    results = BenchmarkResults()
    
    try:
        # Test 1: Node INSERT
        if os.path.exists(db_path):
            os.remove(db_path)
        db = lb.Database(db_path, buffer_pool_size=512 * 1024 * 1024)
        conn = lb.Connection(db, num_threads=num_threads)
        
        duration_ms, memory_mb = benchmark_node_insert(conn, num_rows, batch_size)
        throughput = num_rows / (duration_ms / 1000)
        results.add_result(
            "Node INSERT (CREATE statements)",
            "INSERT",
            "NODE",
            num_rows,
            duration_ms,
            throughput,
            memory_mb
        )
        
        conn.close()
        db.close()
        
        # Test 2: Node Parquet
        if os.path.exists(db_path):
            os.remove(db_path)
        db = lb.Database(db_path, buffer_pool_size=512 * 1024 * 1024)
        conn = lb.Connection(db, num_threads=num_threads)
        
        duration_ms, memory_mb = benchmark_node_parquet(conn, num_rows, tmp_dir)
        throughput = num_rows / (duration_ms / 1000)
        results.add_result(
            "Node COPY FROM Parquet",
            "PARQUET",
            "NODE",
            num_rows,
            duration_ms,
            throughput,
            memory_mb
        )
        
        conn.close()
        db.close()
        
        # Test 3: Rel INSERT
        # Use fewer relationships and nodes for rel benchmarks
        num_rels = min(num_rows, 100000)
        num_nodes = min(10000, num_rows // 10)
        
        if os.path.exists(db_path):
            os.remove(db_path)
        db = lb.Database(db_path, buffer_pool_size=512 * 1024 * 1024)
        conn = lb.Connection(db, num_threads=num_threads)
        
        duration_ms, memory_mb = benchmark_rel_insert(conn, num_rels, batch_size, num_nodes)
        throughput = num_rels / (duration_ms / 1000)
        results.add_result(
            "Rel INSERT (CREATE statements)",
            "INSERT",
            "REL",
            num_rels,
            duration_ms,
            throughput,
            memory_mb
        )
        
        conn.close()
        db.close()
        
        # Test 4: Rel Parquet
        if os.path.exists(db_path):
            os.remove(db_path)
        db = lb.Database(db_path, buffer_pool_size=512 * 1024 * 1024)
        conn = lb.Connection(db, num_threads=num_threads)
        
        duration_ms, memory_mb = benchmark_rel_parquet(conn, num_rels, num_nodes, batch_size, tmp_dir)
        throughput = num_rels / (duration_ms / 1000)
        results.add_result(
            "Rel COPY FROM Parquet",
            "PARQUET",
            "REL",
            num_rels,
            duration_ms,
            throughput,
            memory_mb
        )
        
        conn.close()
        db.close()
        
    finally:
        # Cleanup
        if tmp_dir.exists():
            shutil.rmtree(tmp_dir)
        if os.path.exists(db_path):
            os.remove(db_path)
    
    # Print results
    results.print_results()


def main():
    """Main entry point."""
    parser = argparse.ArgumentParser(
        description="Benchmark Ladybug write performance",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument(
        "--num-rows",
        type=int,
        default=100000,
        help="Number of rows to insert for node benchmarks"
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=1000,
        help="Batch size for INSERT operations"
    )
    parser.add_argument(
        "--num-threads",
        type=int,
        default=4,
        help="Number of threads for database connection"
    )
    parser.add_argument(
        "--db-path",
        type=str,
        default="/tmp/ladybug_write_benchmark.lbug",
        help="Path to benchmark database"
    )
    
    args = parser.parse_args()
    
    run_benchmarks(
        num_rows=args.num_rows,
        batch_size=args.batch_size,
        num_threads=args.num_threads,
        db_path=args.db_path
    )


if __name__ == "__main__":
    main()