benchmark.py

"""
Side-by-side benchmark: datafog v5 (Rust) vs datafog v4 (Python).

Usage:
    # In a venv with the NEW package (this repo):
    pip install .
    python benchmark.py --new

    # In a separate venv with the OLD package (from PyPI):
    pip install datafog
    python benchmark.py --old

    # Or let it auto-detect:
    python benchmark.py
"""

import argparse
import time
import statistics
import sys


# ---------------------------------------------------------------------------
# Test data
# ---------------------------------------------------------------------------

SINGLE_TEXT = (
    "Contact John Doe at john.doe@example.com or call (555) 123-4567. "
    "His SSN is 123-45-6789 and credit card 4111111111111111. "
    "He lives at 123 Main St, New York, NY 10001. "
    "His IP address is 192.168.1.1 and his birthday is 01/15/1985."
)

BATCH_TEXTS = [
    "Email support@company.com for help.",
    "Call us at 800-555-0199 or (212) 555-1234.",
    "SSN: 078-05-1120, DOB: 03/14/1992.",
    "Card ending in 4111111111111111, ZIP 90210.",
    "Server at 10.0.0.1, also try 172.16.0.1.",
    "Reach jane.smith@bigcorp.co.uk for details.",
    "Backup SSN 219-09-9999, CC 378282246310005.",
    "No PII in this sentence at all.",
    "Mixed: bob@test.org, 555.867.5309, 94107-1234.",
    "Patient DOB 12/25/1970, MRN not in regex.",
]

LARGE_TEXT = SINGLE_TEXT * 100  # ~22 KB

ANONYMIZE_TEXT = (
    "Please contact alice@example.com or bob@test.org. "
    "Phone: (555) 111-2222. SSN: 219-09-9999. "
    "Credit card: 4111111111111111. IP: 192.168.0.1. "
    "DOB: 06/15/1990. ZIP: 60601."
)


# ---------------------------------------------------------------------------
# Timing helper
# ---------------------------------------------------------------------------

def bench(fn, rounds=100, warmup=5):
    """Run fn() for warmup rounds, then measure `rounds` iterations."""
    for _ in range(warmup):
        fn()

    times = []
    for _ in range(rounds):
        t0 = time.perf_counter_ns()
        fn()
        t1 = time.perf_counter_ns()
        times.append((t1 - t0) / 1_000)  # nanoseconds -> microseconds

    return {
        "median_us": statistics.median(times),
        "mean_us": statistics.mean(times),
        "min_us": min(times),
        "p95_us": sorted(times)[int(len(times) * 0.95)],
        "rounds": rounds,
    }


def fmt(result):
    """Format a bench result as a compact string."""
    return (
        f"median={result['median_us']:>10.1f} us  "
        f"mean={result['mean_us']:>10.1f} us  "
        f"min={result['min_us']:>10.1f} us  "
        f"p95={result['p95_us']:>10.1f} us"
    )


def print_section(label, result):
    print(f"  {label}:")
    print(f"    {fmt(result)}\n")


# ---------------------------------------------------------------------------
# NEW package (v5, Rust-backed)
# ---------------------------------------------------------------------------

def run_new(rounds):
    from datafog import DataFog, anonymize_text

    fog = DataFog()

    print(f"\n{'='*70}")
    print(f"  datafog v5 (Rust)")
    print(f"{'='*70}\n")

    # Single text detection (reusing engine instance)
    print_section(
        f"detect single ({len(SINGLE_TEXT)} B)",
        bench(lambda: fog.detect(SINGLE_TEXT), rounds=rounds),
    )

    # Batch detection (10 texts)
    print_section(
        f"detect batch x10 ({sum(len(t) for t in BATCH_TEXTS)} B total)",
        bench(lambda: fog.detect_batch(BATCH_TEXTS), rounds=rounds),
    )

    # Large text detection
    print_section(
        f"detect large ({len(LARGE_TEXT)} B)",
        bench(lambda: fog.detect(LARGE_TEXT), rounds=rounds),
    )

    # Anonymize redact
    print_section(
        f"anonymize redact ({len(ANONYMIZE_TEXT)} B)",
        bench(lambda: fog.anonymize(ANONYMIZE_TEXT, method="redact"), rounds=rounds),
    )

    # Anonymize hash
    print_section(
        f"anonymize hash ({len(ANONYMIZE_TEXT)} B)",
        bench(lambda: fog.anonymize(ANONYMIZE_TEXT, method="hash"), rounds=rounds),
    )

    # Correctness
    entities = fog.detect(SINGLE_TEXT)
    types_found = sorted(set(e["type"] for e in entities))
    print(f"  entities found: {types_found}")
    print(f"  count: {len(entities)}")


# ---------------------------------------------------------------------------
# OLD package (v4, pure Python)
# ---------------------------------------------------------------------------

def run_old(rounds):
    from datafog import RegexAnnotator, Anonymizer, AnonymizerType, TextService

    annotator = RegexAnnotator()
    service = TextService(engine="regex")

    print(f"\n{'='*70}")
    print(f"  datafog v4 (Python)")
    print(f"{'='*70}\n")

    # Single text detection
    print_section(
        f"detect single ({len(SINGLE_TEXT)} B)",
        bench(lambda: annotator.annotate(SINGLE_TEXT), rounds=rounds),
    )

    # Batch detection (10 texts)
    print_section(
        f"detect batch x10 ({sum(len(t) for t in BATCH_TEXTS)} B total)",
        bench(lambda: service.batch_annotate_text_sync(BATCH_TEXTS), rounds=rounds),
    )

    # Large text detection
    print_section(
        f"detect large ({len(LARGE_TEXT)} B)",
        bench(lambda: annotator.annotate(LARGE_TEXT), rounds=rounds),
    )

    # Anonymize redact
    def do_redact():
        _, ann = annotator.annotate_with_spans(ANONYMIZE_TEXT)
        anon = Anonymizer(anonymizer_type=AnonymizerType.REDACT)
        anon.anonymize(ANONYMIZE_TEXT, [ann])

    print_section(
        f"anonymize redact ({len(ANONYMIZE_TEXT)} B)",
        bench(do_redact, rounds=rounds),
    )

    # Anonymize hash
    def do_hash():
        _, ann = annotator.annotate_with_spans(ANONYMIZE_TEXT)
        anon = Anonymizer(anonymizer_type=AnonymizerType.HASH)
        anon.anonymize(ANONYMIZE_TEXT, [ann])

    print_section(
        f"anonymize hash ({len(ANONYMIZE_TEXT)} B)",
        bench(do_hash, rounds=rounds),
    )

    # Correctness
    result = annotator.annotate(SINGLE_TEXT)
    types_found = sorted(k for k, v in result.items() if v)
    total = sum(len(v) for v in result.values())
    print(f"  entities found: {types_found}")
    print(f"  count: {total}")


# ---------------------------------------------------------------------------
# Auto-detect
# ---------------------------------------------------------------------------

def detect_version():
    try:
        from datafog import has_ner_support
        return "new"
    except ImportError:
        pass
    try:
        from datafog import RegexAnnotator
        return "old"
    except ImportError:
        pass
    return None


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(
        description="Benchmark datafog v5 (Rust) vs v4 (Python)"
    )
    group = parser.add_mutually_exclusive_group()
    group.add_argument("--new", action="store_true", help="Benchmark the new Rust-backed package")
    group.add_argument("--old", action="store_true", help="Benchmark the old pure-Python package")
    parser.add_argument("--rounds", type=int, default=200, help="Timed iterations per test (default: 200)")
    args = parser.parse_args()

    if args.new:
        version = "new"
    elif args.old:
        version = "old"
    else:
        version = detect_version()
        if version is None:
            print("Error: no datafog package found. Use --new or --old.", file=sys.stderr)
            sys.exit(1)
        print(f"Auto-detected: {version} package")

    if version == "new":
        run_new(args.rounds)
    else:
        run_old(args.rounds)


if __name__ == "__main__":
    main()