perf(lineage): analyse traced columns against only referenced relations

src/sqlbuild/compiler/lineage/_helpers/rich_columns.py

@@ -85,6 +85,7 @@ def build_rich_project_column_lineage(
         return None
 
     schema: dict[str, dict[str, str]] = _build_schema_mapping(project)
+    polyglot_tables: dict[str, dict[str, object]] = _polyglot_schema_tables(schema)
     model_results: dict[str, ModelColumnLineage] = {}
     collapsed_edges: list[ColumnLineageEdge] = []
 
@@ -94,6 +95,7 @@ def build_rich_project_column_lineage(
         result: ModelColumnLineage | None = _build_polyglot_model_column_lineage(
             model=model,
             schema=schema,
+            polyglot_tables=polyglot_tables,
             dialect=dialect,
         )
         if result is None:
@@ -123,6 +125,7 @@ def _build_polyglot_model_column_lineage(
     *,
     model: CompiledModel,
     schema: dict[str, dict[str, str]],
+    polyglot_tables: dict[str, dict[str, object]],
     dialect: str | None,
 ) -> ModelColumnLineage | None:
     normalized_sql: str
@@ -137,7 +140,10 @@ def _build_polyglot_model_column_lineage(
             normalized_sql,
             {
                 "dialect": _polyglot_dialect(dialect),
-                "schema": _polyglot_schema(schema),
+                "schema": _referenced_polyglot_schema(
+                    polyglot_tables=polyglot_tables,
+                    physical_resources=physical_resources,
+                ),
             },
         )
     except polyglot_module.PolyglotError as error:
@@ -292,14 +298,29 @@ def _star_expanded_columns(analysis: dict[str, Any]) -> frozenset[str]:
     return frozenset(expanded)
 
 
-def _polyglot_schema(schema: dict[str, dict[str, str]]) -> dict[str, object]:
-    tables: list[dict[str, object]] = []
-    for table_name, columns in sorted(schema.items()):
+def _polyglot_schema_tables(schema: dict[str, dict[str, str]]) -> dict[str, dict[str, object]]:
+    tables: dict[str, dict[str, object]] = {}
+    for table_name, columns in schema.items():
         table_columns: list[dict[str, str]] = []
         for column_name, column_type in sorted(columns.items()):
             table_columns.append({"name": column_name, "type": column_type or "UNKNOWN"})
-        tables.append({"name": table_name, "columns": table_columns})
-    return {"tables": tables}
+        tables[table_name] = {"name": table_name, "columns": table_columns}
+    return tables
+
+
+def _referenced_polyglot_schema(
+    *,
+    polyglot_tables: dict[str, dict[str, object]],
+    physical_resources: tuple[PhysicalResource, ...],
+) -> dict[str, object]:
+    referenced: set[str] = {resource.physical_name for resource in physical_resources}
+    return {
+        "tables": [
+            polyglot_tables[table_name]
+            for table_name in sorted(referenced)
+            if table_name in polyglot_tables
+        ]
+    }
 
 
 def _polyglot_dialect(dialect: str | TypeDialect | None) -> PolyglotAnalysisDialect:

tests/e2e/src/sqlbuild/cli/commands/main/compile/COMPILE_BENCHMARKS.md

@@ -64,8 +64,9 @@ fans the last layer back in. Every lattice model is reachable along exponentiall
 compile guards do not enable the lattice, so their fingerprints and budgets are unchanged.
 
 After one compile and warm lineage and scope caches, the guard runs `sqb lineage` downstream from
-the hub and upstream from the rollup as text and JSON, `sqb dag --json` and `sqb scope --json`,
-each in a fresh process. Every command
+the hub and upstream from the rollup as text and JSON, rich column traces upstream from
+`shared_orders_rollup.amount` and downstream from `shared_orders_hub.amount`, `sqb dag --json` and
+`sqb scope --json`, each in a fresh process. Every command
 has a wall-time ceiling, a peak-RSS ceiling and an output-line bound linear in the model count, so
 a renderer or graph walk that re-expands shared nodes per path fails on size before it can hide
 behind a loose timing budget. It runs in the 3,000-model fresh-process compile job.

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

@@ -101,6 +101,35 @@
             expected_max_wall_seconds=3.0,
             expected_max_rss_bytes=128 * _MIB,
         ),
+        InspectionCommandPerformanceGuardTestCase(
+            description="models_3000_lineage_column_text_upstream_from_fan_in_rollup",
+            sqb_args=("lineage", f"{SHARED_DIAMOND_ROLLUP}.amount", "--direction", "upstream"),
+            expected_fragments=(
+                f"Column trace  {SHARED_DIAMOND_ROLLUP}.amount  upstream",
+                "shared_orders_l15_s00.amount (cast)",
+            ),
+            expected_max_output_lines=_MODEL_COUNT,
+            expected_max_wall_seconds=8.0,
+            expected_max_rss_bytes=512 * _MIB,
+        ),
+        InspectionCommandPerformanceGuardTestCase(
+            description="models_3000_lineage_column_json_downstream_from_fan_out_hub",
+            sqb_args=(
+                "lineage",
+                f"{SHARED_DIAMOND_HUB}.amount",
+                "--direction",
+                "downstream",
+                "--format",
+                "json",
+            ),
+            expected_fragments=(
+                '"direction": "downstream"',
+                f'"resource_name": "{SHARED_DIAMOND_ROLLUP}"',
+            ),
+            expected_max_output_lines=20 * _MODEL_COUNT,
+            expected_max_wall_seconds=8.0,
+            expected_max_rss_bytes=512 * _MIB,
+        ),
         InspectionCommandPerformanceGuardTestCase(
             description="models_3000_dag_json",
             sqb_args=("dag", "--json"),

tests/unit/src/sqlbuild/compiler/lineage/_test_types.py

@@ -31,6 +31,16 @@ class SqlAnalysisDisabledLineageTestCase:
     expected_result_is_none: bool
 
 
+@dataclass(frozen=True)
+class RichLineageAnalysisScalingTestCase:
+    description: str
+    chain_length: int
+    unrelated_model_counts: tuple[int, ...]
+    expected_schema_builds: int
+    expected_analysis_calls: int
+    expected_schema_tables_per_call: tuple[tuple[str, ...], ...]
+
+
 @dataclass(frozen=True)
 class ExpectedCountTestCase:
     description: str

tests/unit/src/sqlbuild/compiler/lineage/helpers.py

@@ -90,3 +90,30 @@ def make_compiled_seed(*, name: str, columns: tuple[str, ...]) -> CompiledSeed:
 
 def edge_label(source_name: str, source_column: str, target_name: str, target_column: str) -> str:
     return f"{source_name}.{source_column}->{target_name}.{target_column}"
+
+
+def make_chain_models(
+    *, chain_length: int, unrelated_model_count: int
+) -> tuple[CompiledModel, ...]:
+    root: CompiledModel = make_compiled_model(
+        name="orders_0",
+        query_sql="SELECT 1 AS order_id, 10 AS amount",
+        inferred_columns=("order_id", "amount"),
+    )
+    chain: tuple[CompiledModel, ...] = tuple(
+        make_compiled_model(
+            name=f"orders_{index}",
+            query_sql=f'SELECT order_id, amount + 1 AS amount FROM __ref("orders_{index - 1}")',
+            inferred_columns=("order_id", "amount"),
+        )
+        for index in range(1, chain_length)
+    )
+    unrelated: tuple[CompiledModel, ...] = tuple(
+        make_compiled_model(
+            name=f"inventory_{index}",
+            query_sql="SELECT 1 AS product_id, 5 AS quantity",
+            inferred_columns=("product_id", "quantity"),
+        )
+        for index in range(unrelated_model_count)
+    )
+    return (root, *chain, *unrelated)

tests/unit/src/sqlbuild/compiler/lineage/test_columns.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import replace
+from typing import Any
 
 import pytest
 
@@ -11,21 +12,26 @@
     CompiledProject,
 )
 from sqlbuild.compiler.compile.types import CompiledResourceType
+from sqlbuild.compiler.lineage._helpers import rich_columns
+from sqlbuild.compiler.lineage._helpers.columns import _build_schema_mapping
 from sqlbuild.compiler.lineage.main.columns import build_project_column_lineage
 from sqlbuild.compiler.lineage.models import ColumnLineage, ModelColumnLineage, ProjectColumnLineage
 from sqlbuild.compiler.lineage.types import (
     ColumnLineageConfidence,
     ColumnLineageMode,
     ColumnTransformKind,
 )
+from sqlbuild.compiler.sql_analysis.main.import_polyglot_sql import import_polyglot_sql
 from tests.unit.src.sqlbuild.compiler.lineage._test_types import (
     ColumnLineageAnalyzerTestCase,
     ExpectedCountTestCase,
     ProjectLineageGraphTestCase,
+    RichLineageAnalysisScalingTestCase,
     SqlAnalysisDisabledLineageTestCase,
 )
 from tests.unit.src.sqlbuild.compiler.lineage.helpers import (
     edge_label,
+    make_chain_models,
     make_compiled_model,
     make_compiled_project,
     make_compiled_seed,
@@ -653,3 +659,78 @@ def test_given_sql_analysis_disabled_when_building_column_lineage_then_returns_n
     result: ProjectColumnLineage | None = build_project_column_lineage(project=project)
 
     assert (result is None) is test_case.expected_result_is_none
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    [
+        RichLineageAnalysisScalingTestCase(
+            description="builds the schema once and analyses each traced model against its refs",
+            chain_length=4,
+            unrelated_model_counts=(0, 250),
+            expected_schema_builds=1,
+            expected_analysis_calls=4,
+            expected_schema_tables_per_call=(
+                (),
+                ("__sqlbuild_model__orders_0",),
+                ("__sqlbuild_model__orders_1",),
+                ("__sqlbuild_model__orders_2",),
+            ),
+        )
+    ],
+    ids=lambda case: case.description,
+)
+def test_given_unrelated_models_when_tracing_rich_lineage_then_analysis_work_stays_linear(
+    test_case: RichLineageAnalysisScalingTestCase,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    polyglot_module: Any = import_polyglot_sql()
+    analyze_query: Any = polyglot_module.analyze_query
+    schema_builds: list[CompiledProject] = []
+    analysed_tables: list[tuple[str, ...]] = []
+
+    def counting_schema_mapping(project: CompiledProject) -> dict[str, dict[str, str]]:
+        schema_builds.append(project)
+        return _build_schema_mapping(project)
+
+    def recording_analyze_query(sql: str, options: dict[str, Any]) -> object:
+        analysed_tables.append(tuple(table["name"] for table in options["schema"]["tables"]))
+        return analyze_query(sql, options)
+
+    monkeypatch.setattr(rich_columns, "_build_schema_mapping", counting_schema_mapping)
+    monkeypatch.setattr(polyglot_module, "analyze_query", recording_analyze_query)
+    chain_names: frozenset[str] = frozenset(
+        f"orders_{index}" for index in range(test_case.chain_length)
+    )
+    edges_by_project_size: list[tuple[str, ...]] = []
+    for unrelated_model_count in test_case.unrelated_model_counts:
+        schema_builds.clear()
+        analysed_tables.clear()
+        project: CompiledProject = make_compiled_project(
+            models=make_chain_models(
+                chain_length=test_case.chain_length,
+                unrelated_model_count=unrelated_model_count,
+            )
+        )
+
+        result: ProjectColumnLineage | None = build_project_column_lineage(
+            project=project, model_names=chain_names
+        )
+
+        assert result is not None
+        assert len(schema_builds) == test_case.expected_schema_builds
+        assert len(analysed_tables) == test_case.expected_analysis_calls
+        assert tuple(analysed_tables) == test_case.expected_schema_tables_per_call
+        edges_by_project_size.append(
+            tuple(
+                edge_label(
+                    edge.source.resource_name,
+                    edge.source.column_name,
+                    edge.target.resource_name,
+                    edge.target.column_name,
+                )
+                for edge in result.edges
+            )
+        )
+    assert len(set(edges_by_project_size)) == 1
+    assert "orders_2.amount->orders_3.amount" in edges_by_project_size[0]