perf(compile): keep native analysis parallel at 10k models

src/sqlbuild/compiler/compile/_helpers/analysis/compact.py

@@ -180,8 +180,7 @@
 from sqlbuild.diagnostics.main.log_debug_event import log_debug_event
 
 _DEBUG_LOGGER: logging.Logger = logging.getLogger("sqlbuild.compile")
-_LARGE_COMPACT_PROJECT_MODELS: int = 10_000
-_LARGE_COMPACT_SQL_BYTES: int = 48 * 1024 * 1024
+_COMPACT_ANALYSIS_WORKERS: int = 4
 _NATIVE_LEGACY_FALLBACK: str = "native project type recovery requires legacy fallback"
 
 
@@ -886,14 +885,9 @@ def _prepare_compact_analysis_batch(
 
 
 def _run_compact_analysis_batch(*, preparation: CompactBatchPreparation) -> object:
-    """Bound resident analysis heaps before the isolated custom-rule host starts."""
+    """Analyse every project size with the full worker pool; the rules phase sets peak memory."""
 
-    large_project: bool = (
-        len(preparation.projections) >= _LARGE_COMPACT_PROJECT_MODELS
-        and sum(len(sql.encode("utf-8")) for sql in preparation.cleaned_sql)
-        >= _LARGE_COMPACT_SQL_BYTES
-    )
-    workers: int = 1 if large_project else 4
+    workers: int = _COMPACT_ANALYSIS_WORKERS
     if preparation.binding_catalog is not None:
         return orjson.loads(
             preparation.binding_catalog.native.analyze_compact(

src/sqlbuild/rule_engine/_helpers/engine/native.py

@@ -5,7 +5,7 @@
 import inspect
 import json
 from concurrent.futures import Executor, Future, ThreadPoolExecutor
-from dataclasses import asdict, replace
+from dataclasses import asdict, dataclass, replace
 from decimal import Decimal
 from pathlib import Path
 from typing import Any, cast
@@ -54,6 +54,13 @@
 from sqlbuild.sql_values.types import SqlValueKind
 
 
+@dataclass(frozen=True)
+class _PreparedNativeRequest:
+    identity: str | None
+    reused: str | None
+    request_text: str | None
+
+
 def evaluate_native(
     *,
     project: CompiledProject,
@@ -74,31 +81,20 @@ def evaluate_native(
         if custom_payloads is None
         else custom_payloads
     )
-    request: dict[str, object] = {
-        "version": RULES_NATIVE_API_VERSION,
-        "project_dir": str(project_dir.resolve()),
-        "dialect": dialect,
-        "config": _config_payload(config),
-        "models": _model_payloads(
-            project=project,
-            dialect=dialect,
-            include_type_proof=any(
-                _rule_selected(config=config, code=code) for code in TYPE_PROOF_RULE_CODES
-            ),
+    prepared: _PreparedNativeRequest = _prepare_native_request(
+        request_json=_serialise_native_request(
+            _native_request(
+                project=project,
+                config=config,
+                project_dir=project_dir,
+                dialect=dialect,
+                initial_findings=initial_findings,
+                custom_payloads=payloads,
+            )
         ),
-        "sql_tests": _sql_test_payloads(project),
-        "sql_scenarios": _sql_scenario_payloads(project),
-        "public_enums": [
-            _enum_payload(declaration) for declaration in project.public_enums.values()
-        ],
-        "public_constants": [
-            _constant_payload(declaration) for declaration in project.public_constants.values()
-        ],
-        "scope_index": scope_metadata_projection(index=project.scope_index),
-        "initial_findings": [_finding_payload(finding) for finding in initial_findings],
-        "defer_suppressions": True,
-        "custom_rules": payloads,
-    }
+        project_dir=project_dir,
+        cache_enabled=config.cache.enabled,
+    )
     with ThreadPoolExecutor(max_workers=1, thread_name_prefix="sqlbuild-custom-rules") as pool:
         custom_future: Future[CustomRulesOutcome] | None = (
             custom_outcome
@@ -114,15 +110,14 @@ def evaluate_native(
                 verify_determinism=verify_determinism,
             )
         )
+        reused: bool = prepared.reused is not None
         try:
-            response_json: str
-            reused: bool
-            response_json, reused = _evaluate_request(
-                request=request, project_dir=project_dir, cache_enabled=config.cache.enabled
+            response: object = orjson.loads(
+                _evaluate_request(prepared=prepared, project_dir=project_dir)
             )
-            response: object = orjson.loads(response_json)
         except (ValueError, TypeError) as error:
             raise RulesError(str(error)) from error
+        del prepared
         custom: CustomRulesOutcome = (
             CustomRulesOutcome(findings=(), cache_hits=0, cache_misses=0, custom_ms=0)
             if custom_future is None
@@ -159,6 +154,42 @@ def evaluate_native(
     )
 
 
+def _native_request(
+    *,
+    project: CompiledProject,
+    config: RulesConfig,
+    project_dir: Path,
+    dialect: str,
+    initial_findings: tuple[Finding, ...],
+    custom_payloads: list[dict[str, object]],
+) -> dict[str, object]:
+    return {
+        "version": RULES_NATIVE_API_VERSION,
+        "project_dir": str(project_dir.resolve()),
+        "dialect": dialect,
+        "config": _config_payload(config),
+        "models": _model_payloads(
+            project=project,
+            dialect=dialect,
+            include_type_proof=any(
+                _rule_selected(config=config, code=code) for code in TYPE_PROOF_RULE_CODES
+            ),
+        ),
+        "sql_tests": _sql_test_payloads(project),
+        "sql_scenarios": _sql_scenario_payloads(project),
+        "public_enums": [
+            _enum_payload(declaration) for declaration in project.public_enums.values()
+        ],
+        "public_constants": [
+            _constant_payload(declaration) for declaration in project.public_constants.values()
+        ],
+        "scope_index": scope_metadata_projection(index=project.scope_index),
+        "initial_findings": [_finding_payload(finding) for finding in initial_findings],
+        "defer_suppressions": True,
+        "custom_rules": custom_payloads,
+    }
+
+
 def start_custom_rules(
     *,
     executor: Executor,
@@ -206,19 +237,42 @@ def _selected_custom_rules(
     return tuple(rule for rule in catalogue if rule.custom and rule.code in selected)
 
 
-def _evaluate_request(
-    *, request: dict[str, object], project_dir: Path, cache_enabled: bool
-) -> tuple[str, bool]:
-    request_json: bytes = orjson.dumps(request, option=orjson.OPT_SORT_KEYS, default=str)
-    if not cache_enabled:
-        return _native.evaluate_json(request_json.decode()), False
-    identity: str = native_request_identity(request_json)
-    reused: str | None = read_native_response(project_dir=project_dir, identity=identity)
-    if reused is not None:
-        return reused, True
-    response: str = _native.evaluate_json(request_json.decode())
-    write_native_response(project_dir=project_dir, identity=identity, response=response)
-    return response, False
+def _serialise_native_request(request: dict[str, object]) -> bytes:
+    """Encode the native request, reporting values the encoder rejects as rules errors."""
+
+    try:
+        return orjson.dumps(request, option=orjson.OPT_SORT_KEYS, default=str)
+    except orjson.JSONEncodeError as error:
+        raise RulesError(str(error)) from error
+
+
+def _prepare_native_request(
+    *, request_json: bytes, project_dir: Path, cache_enabled: bool
+) -> _PreparedNativeRequest:
+    """Resolve the memo and decode once; the caller's bytes die when this returns."""
+
+    identity: str | None = native_request_identity(request_json) if cache_enabled else None
+    reused: str | None = (
+        read_native_response(project_dir=project_dir, identity=identity)
+        if identity is not None
+        else None
+    )
+    return _PreparedNativeRequest(
+        identity=identity,
+        reused=reused,
+        request_text=request_json.decode() if reused is None else None,
+    )
+
+
+def _evaluate_request(*, prepared: _PreparedNativeRequest, project_dir: Path) -> str:
+    if prepared.reused is not None:
+        return prepared.reused
+    response: str = _native.evaluate_json(cast(str, prepared.request_text))
+    if prepared.identity is not None:
+        write_native_response(
+            project_dir=project_dir, identity=prepared.identity, response=response
+        )
+    return response
 
 
 def finalize_native_findings(

tests/e2e/src/sqlbuild/cli/commands/main/compile/test_dense_compile_performance.py

@@ -62,7 +62,7 @@
             "dense_models_10000_all_rules",
             10000,
             145.0,
-            4 * _GIB,
+            7 * _GIB // 2,
             "913a08c2277304efacf75e293728767a3fa416f6b32f5488dca6ab5baea80668",
         ),
     ),

tests/unit/src/sqlbuild/compiler/compile/_helpers/analysis/_test_types.py

@@ -14,3 +14,11 @@ class DynamicPivotAnalysisTestCase:
     expected_family_types: tuple[str | None, ...] = ("DECIMAL(12,2)",)
     expected_input_relations: tuple[str, ...] = ()
     aggregate: str = "MAX"
+
+
+@dataclass(frozen=True)
+class CompactAnalysisWorkersTestCase:
+    description: str
+    model_count: int
+    sql_bytes_per_model: int
+    expected_workers: int

tests/unit/src/sqlbuild/compiler/compile/_helpers/analysis/test_compact_workers.py

@@ -0,0 +1,58 @@
+from __future__ import annotations
+
+import orjson
+import pytest
+
+from sqlbuild.compiler.compile._helpers.analysis import compact
+from sqlbuild.compiler.compile.models import CompactBatchPreparation
+from tests.unit.src.sqlbuild.compiler.compile._helpers.analysis._test_types import (
+    CompactAnalysisWorkersTestCase,
+)
+
+_MIB: int = 1024 * 1024
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    (
+        CompactAnalysisWorkersTestCase(
+            description="small project",
+            model_count=3,
+            sql_bytes_per_model=64,
+            expected_workers=4,
+        ),
+        CompactAnalysisWorkersTestCase(
+            description="dense ten thousand model project",
+            model_count=10_000,
+            sql_bytes_per_model=(64 * _MIB) // 10_000,
+            expected_workers=4,
+        ),
+    ),
+    ids=lambda case: case.description,
+)
+def test_given_project_size_when_running_compact_analysis_then_native_uses_all_workers(
+    test_case: CompactAnalysisWorkersTestCase,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    captured: list[dict[str, object]] = []
+
+    def analyze(request_json: str) -> str:
+        captured.append(orjson.loads(request_json))
+        return "{}"
+
+    monkeypatch.setattr(compact._native, "analyze_project_queries_compact_json", analyze)
+    sql: str = "SELECT 1 AS order_id" + " " * (test_case.sql_bytes_per_model - 20)
+    preparation: CompactBatchPreparation = CompactBatchPreparation(
+        cleaned_sql=(sql,) * test_case.model_count,
+        queries=({"sql": "SELECT 1 AS order_id", "dialect": "duckdb"},),
+        templates=({"queryIndex": 0},),
+        projections=({"templateIndex": 0},) * test_case.model_count,
+    )
+
+    _ = compact._run_compact_analysis_batch(preparation=preparation)
+
+    assert [request["workers"] for request in captured] == [test_case.expected_workers]
+
+
+if __name__ == "__main__":
+    raise SystemExit(pytest.main([__file__, "-n", "auto", "--dist", "loadfile"]))

tests/unit/src/sqlbuild/rule_engine/_helpers/engine/_test_types.py

@@ -202,3 +202,20 @@ class NativeBuildIdentityTestCase:
     description: str
     expected_rebuilt_evaluations: int
     expected_warm_evaluations: int
+
+
+@dataclass(frozen=True)
+class NativeRequestBuildErrorTestCase:
+    """One programming error raised while building the native rules request."""
+
+    description: str
+    expected_error: Exception
+
+
+@dataclass(frozen=True)
+class NativeRequestEncodeErrorTestCase:
+    """One request value the JSON encoder rejects."""
+
+    description: str
+    rejected_value: object
+    expected_message: str