feat: implement parallel AST extraction using ProcessPoolExecutor with benchmarking support

graphify/extract.py

@@ -3417,7 +3417,167 @@ def _check_tree_sitter_version() -> None:
         )
 
 
-def extract(paths: list[Path], cache_root: Path | None = None) -> dict:
+_DISPATCH: dict[str, Any] = {
+    ".py": extract_python,
+    ".js": extract_js,
+    ".jsx": extract_js,
+    ".mjs": extract_js,
+    ".ts": extract_js,
+    ".tsx": extract_js,
+    ".go": extract_go,
+    ".rs": extract_rust,
+    ".java": extract_java,
+    ".c": extract_c,
+    ".h": extract_c,
+    ".cpp": extract_cpp,
+    ".cc": extract_cpp,
+    ".cxx": extract_cpp,
+    ".hpp": extract_cpp,
+    ".rb": extract_ruby,
+    ".cs": extract_csharp,
+    ".kt": extract_kotlin,
+    ".kts": extract_kotlin,
+    ".scala": extract_scala,
+    ".php": extract_php,
+    ".swift": extract_swift,
+    ".lua": extract_lua,
+    ".toc": extract_lua,
+    ".zig": extract_zig,
+    ".ps1": extract_powershell,
+    ".ex": extract_elixir,
+    ".exs": extract_elixir,
+    ".m": extract_objc,
+    ".mm": extract_objc,
+    ".jl": extract_julia,
+    ".vue": extract_js,
+    ".svelte": extract_js,
+    ".dart": extract_dart,
+    ".v": extract_verilog,
+    ".sv": extract_verilog,
+    ".sql": extract_sql,
+}
+
+
+def _get_extractor(path: Path) -> Any | None:
+    """Return the correct extractor function for a file, or None if unsupported."""
+    if path.name.endswith(".blade.php"):
+        return extract_blade
+    return _DISPATCH.get(path.suffix)
+
+
+def _extract_single_file(args: tuple) -> tuple[int, dict]:
+    """Worker function for parallel extraction. Runs in a subprocess.
+
+    Must be at module level (not a closure) so it can be pickled by
+    ProcessPoolExecutor.
+
+    Args:
+        args: (index, path_str, cache_root_str) tuple
+
+    Returns:
+        (index, result_dict) so results can be placed back in order.
+    """
+    idx, path_str, cache_root_str = args
+    path = Path(path_str)
+    cache_root = Path(cache_root_str)
+
+    # Check cache first (avoid re-extraction)
+    cached = load_cached(path, cache_root)
+    if cached is not None:
+        return idx, cached
+
+    extractor = _get_extractor(path)
+    if extractor is None:
+        return idx, {"nodes": [], "edges": []}
+
+    result = extractor(path)
+    if "error" not in result:
+        save_cached(path, result, cache_root)
+    return idx, result
+
+
+def _extract_parallel(
+    uncached_work: list[tuple[int, Path]],
+    per_file: list[dict | None],
+    effective_root: Path,
+    max_workers: int | None,
+    total_files: int,
+) -> None:
+    """Extract uncached files in parallel using ProcessPoolExecutor."""
+    import concurrent.futures
+
+    if max_workers is None:
+        max_workers = min(os.cpu_count() or 4, len(uncached_work), 8)
+
+    root_str = str(effective_root)
+    work_items = [(idx, str(path), root_str) for idx, path in uncached_work]
+
+    done_count = 0
+    _PROGRESS_INTERVAL = 100
+    with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as pool:
+        futures = {
+            pool.submit(_extract_single_file, item): item[0] for item in work_items
+        }
+        for future in concurrent.futures.as_completed(futures):
+            idx, result = future.result()
+            per_file[idx] = result
+            done_count += 1
+            if (
+                total_files >= _PROGRESS_INTERVAL
+                and done_count % _PROGRESS_INTERVAL == 0
+            ):
+                print(
+                    f"  AST extraction: {done_count}/{len(uncached_work)} uncached files "
+                    f"({done_count * 100 // len(uncached_work)}%) [{max_workers} workers]",
+                    flush=True,
+                )
+    if total_files >= _PROGRESS_INTERVAL:
+        print(
+            f"  AST extraction: {total_files}/{total_files} files (100%) [{max_workers} workers]",
+            flush=True,
+        )
+
+
+def _extract_sequential(
+    uncached_work: list[tuple[int, Path]],
+    per_file: list[dict | None],
+    effective_root: Path,
+    total_files: int,
+) -> None:
+    """Extract uncached files sequentially (fallback for small batches)."""
+    _PROGRESS_INTERVAL = 100
+    for work_idx, (idx, path) in enumerate(uncached_work):
+        if (
+            total_files >= _PROGRESS_INTERVAL
+            and work_idx % _PROGRESS_INTERVAL == 0
+            and work_idx > 0
+        ):
+            print(
+                f"  AST extraction: {work_idx}/{len(uncached_work)} uncached files ({work_idx * 100 // len(uncached_work)}%)",
+                flush=True,
+            )
+        extractor = _get_extractor(path)
+        if extractor is None:
+            per_file[idx] = {"nodes": [], "edges": []}
+            continue
+        result = extractor(path)
+        if "error" not in result:
+            save_cached(path, result, effective_root)
+        per_file[idx] = result
+    if total_files >= _PROGRESS_INTERVAL:
+        print(f"  AST extraction: {total_files}/{total_files} files (100%)", flush=True)
+
+
+_PARALLEL_THRESHOLD = 20
+
+
+def extract(
+    paths: list[Path],
+    cache_root: Path | None = None,
+    *,
+    parallel: bool = True,
+    max_workers: int | None = None,
+) -> dict:
     """Extract AST nodes and edges from a list of code files.
 
     Two-pass process:
@@ -3430,9 +3590,11 @@ def extract(paths: list[Path], cache_root: Path | None = None) -> dict:
         cache_root: explicit root for graphify-out/cache/ (overrides the
             inferred common path prefix). Pass Path('.') when running on a
             subdirectory so the cache stays at ./graphify-out/cache/.
+        parallel: if True and there are >= _PARALLEL_THRESHOLD uncached files,
+            use ProcessPoolExecutor for multi-core extraction.
+        max_workers: max subprocess count. Defaults to min(cpu_count, 8).
     """
     _check_tree_sitter_version()
-    per_file: list[dict] = []
 
     # Infer a common root for cache keys (use first diverging segment, not sum of all matches)
     try:
@@ -3453,68 +3615,36 @@ def extract(paths: list[Path], cache_root: Path | None = None) -> dict:
         root = Path(".")
     root = root.resolve()
 
-    _DISPATCH: dict[str, Any] = {
+    effective_root = cache_root or root
-        ".py": extract_python,
-        ".js": extract_js,
-        ".jsx": extract_js,
-        ".mjs": extract_js,
-        ".ts": extract_js,
-        ".tsx": extract_js,
-        ".go": extract_go,
-        ".rs": extract_rust,
-        ".java": extract_java,
-        ".c": extract_c,
-        ".h": extract_c,
-        ".cpp": extract_cpp,
-        ".cc": extract_cpp,
-        ".cxx": extract_cpp,
-        ".hpp": extract_cpp,
-        ".rb": extract_ruby,
-        ".cs": extract_csharp,
-        ".kt": extract_kotlin,
-        ".kts": extract_kotlin,
-        ".scala": extract_scala,
-        ".php": extract_php,
-        ".swift": extract_swift,
-        ".lua": extract_lua,
-        ".toc": extract_lua,
-        ".zig": extract_zig,
-        ".ps1": extract_powershell,
-        ".ex": extract_elixir,
-        ".exs": extract_elixir,
-        ".m": extract_objc,
-        ".mm": extract_objc,
-        ".jl": extract_julia,
-        ".vue": extract_js,
-        ".svelte": extract_js,
-        ".dart": extract_dart,
-        ".v": extract_verilog,
-        ".sv": extract_verilog,
-        ".sql": extract_sql,
-    }
-
     total = len(paths)
-    _PROGRESS_INTERVAL = 100
+
+    # Phase 1: separate cached hits from uncached work
+    per_file: list[dict | None] = [None] * total
+    uncached_work: list[tuple[int, Path]] = []
+
     for i, path in enumerate(paths):
-        if total >= _PROGRESS_INTERVAL and i % _PROGRESS_INTERVAL == 0 and i > 0:
+        if _get_extractor(path) is None:
-            print(f"  AST extraction: {i}/{total} files ({i * 100 // total}%)", flush=True)
+            per_file[i] = {"nodes": [], "edges": []}
-        # .blade.php must be checked before suffix lookup since Path.suffix returns .php
-        if path.name.endswith(".blade.php"):
-            extractor = extract_blade
-        else:
-            extractor = _DISPATCH.get(path.suffix)
-        if extractor is None:
             continue
-        cached = load_cached(path, cache_root or root)
+        cached = load_cached(path, effective_root)
         if cached is not None:
-            per_file.append(cached)
+            per_file[i] = cached
             continue
-        result = extractor(path)
+        uncached_work.append((i, path))
-        if "error" not in result:
+
-            save_cached(path, result, cache_root or root)
+    # Phase 2: extract uncached files (parallel or sequential)
-        per_file.append(result)
+    if uncached_work:
-    if total >= _PROGRESS_INTERVAL:
+        if parallel and len(uncached_work) >= _PARALLEL_THRESHOLD:
-        print(f"  AST extraction: {total}/{total} files (100%)", flush=True)
+            _extract_parallel(
+                uncached_work, per_file, effective_root, max_workers, total
+            )
+        else:
+            _extract_sequential(uncached_work, per_file, effective_root, total)
+
+    # Fill any remaining None slots (shouldn't happen, but defensive)
+    for i in range(total):
+        if per_file[i] is None:
+            per_file[i] = {"nodes": [], "edges": []}
 
     all_nodes: list[dict] = []
     all_edges: list[dict] = []

tests/bench_extract.py

@@ -0,0 +1,178 @@
+#!/usr/bin/env python3
+"""Benchmark: sequential vs parallel AST extraction.
+
+Usage:
+    python tests/bench_extract.py [path-to-repo]
+
+Defaults to the current directory if no path is given.
+Clears the AST cache between runs so every file is re-extracted.
+
+Example output:
+    === Graphify AST Extraction Benchmark ===
+    Files:        1,247
+    Languages:    Python (412), TypeScript (389), Go (201), ...
+
+    Sequential:   4.32s (8,934 nodes, 12,456 edges)
+    Parallel (8): 1.28s (8,934 nodes, 12,456 edges)
+
+    Speedup:      3.38x
+    Results:      ✓ identical
+"""
+
+from __future__ import annotations
+
+import sys
+import time
+from collections import Counter
+from pathlib import Path
+
+# Ensure the project root is importable
+_project_root = Path(__file__).resolve().parent.parent
+if str(_project_root) not in sys.path:
+    sys.path.insert(0, str(_project_root))
+
+from graphify.extract import extract, collect_files
+from graphify.cache import clear_cache
+
+
+def _count_by_ext(paths: list[Path]) -> dict[str, int]:
+    """Count files by extension."""
+    counter: Counter[str] = Counter()
+    for p in paths:
+        ext = p.suffix.lower()
+        counter[ext] += 1
+    return dict(counter.most_common())
+
+
+_EXT_NAMES: dict[str, str] = {
+    ".py": "Python",
+    ".js": "JavaScript",
+    ".jsx": "JSX",
+    ".mjs": "MJS",
+    ".ts": "TypeScript",
+    ".tsx": "TSX",
+    ".go": "Go",
+    ".rs": "Rust",
+    ".java": "Java",
+    ".c": "C",
+    ".h": "C Header",
+    ".cpp": "C++",
+    ".cc": "C++",
+    ".cxx": "C++",
+    ".hpp": "C++ Header",
+    ".rb": "Ruby",
+    ".cs": "C#",
+    ".kt": "Kotlin",
+    ".kts": "Kotlin Script",
+    ".scala": "Scala",
+    ".php": "PHP",
+    ".swift": "Swift",
+    ".lua": "Lua",
+    ".toc": "Lua TOC",
+    ".zig": "Zig",
+    ".ps1": "PowerShell",
+    ".ex": "Elixir",
+    ".exs": "Elixir Script",
+    ".m": "Obj-C",
+    ".mm": "Obj-C++",
+    ".jl": "Julia",
+    ".vue": "Vue",
+    ".svelte": "Svelte",
+    ".dart": "Dart",
+    ".v": "Verilog",
+    ".sv": "SystemVerilog",
+    ".sql": "SQL",
+}
+
+
+def _format_languages(ext_counts: dict[str, int]) -> str:
+    parts = []
+    for ext, count in ext_counts.items():
+        name = _EXT_NAMES.get(ext, ext)
+        parts.append(f"{name} ({count})")
+    return ", ".join(parts)
+
+
+def _run_extraction(
+    paths: list[Path],
+    cache_root: Path,
+    parallel: bool,
+    max_workers: int | None = None,
+) -> tuple[float, int, int]:
+    """Run extraction, return (elapsed_seconds, node_count, edge_count)."""
+    clear_cache(cache_root)
+    t0 = time.perf_counter()
+    result = extract(
+        paths, cache_root=cache_root, parallel=parallel, max_workers=max_workers
+    )
+    elapsed = time.perf_counter() - t0
+    nodes = len(result.get("nodes", []))
+    edges = len(result.get("edges", []))
+    return elapsed, nodes, edges
+
+
+def main() -> None:
+    target = Path(sys.argv[1]) if len(sys.argv) > 1 else Path(".")
+    target = target.resolve()
+
+    if not target.exists():
+        print(f"Error: {target} does not exist", file=sys.stderr)
+        sys.exit(1)
+
+    print("=== Graphify AST Extraction Benchmark ===\n")
+    print(f"Scanning {target} ...", flush=True)
+
+    paths = collect_files(target)
+    if not paths:
+        print("No extractable files found.", file=sys.stderr)
+        sys.exit(1)
+
+    ext_counts = _count_by_ext(paths)
+    print(f"Files:        {len(paths):,}")
+    print(f"Languages:    {_format_languages(ext_counts)}")
+    print()
+
+    cache_root = target if target.is_dir() else target.parent
+
+    # Workers count (same logic as _extract_parallel)
+    import os
+
+    workers = min(os.cpu_count() or 4, len(paths), 8)
+
+    # Run sequential
+    print("Running sequential extraction...", flush=True)
+    seq_time, seq_nodes, seq_edges = _run_extraction(paths, cache_root, parallel=False)
+    print(f"Sequential:   {seq_time:.2f}s ({seq_nodes:,} nodes, {seq_edges:,} edges)")
+
+    # Run parallel
+    print(f"\nRunning parallel extraction ({workers} workers)...", flush=True)
+    par_time, par_nodes, par_edges = _run_extraction(
+        paths, cache_root, parallel=True, max_workers=workers
+    )
+    print(
+        f"Parallel ({workers}): {par_time:.2f}s ({par_nodes:,} nodes, {par_edges:,} edges)"
+    )
+
+    # Results
+    print()
+    if seq_time > 0:
+        speedup = seq_time / par_time if par_time > 0 else float("inf")
+        print(f"Speedup:      {speedup:.2f}x")
+    print(f"Workers:      {workers} (auto-detected)")
+
+    # Validate correctness
+    if seq_nodes == par_nodes and seq_edges == par_edges:
+        print("Results:      ✓ identical (node count, edge count match)")
+    else:
+        print("Results:      ✗ MISMATCH!")
+        print(f"  Sequential: {seq_nodes} nodes, {seq_edges} edges")
+        print(f"  Parallel:   {par_nodes} nodes, {par_edges} edges")
+        sys.exit(1)
+
+    # Clean up cache after benchmark
+    clear_cache(cache_root)
+    print("\nCache cleared after benchmark.")
+
+
+if __name__ == "__main__":
+    main()