From 39b45424b9135d07cc70a8a82f2605ed10ea9821 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Tue, 29 Sep 2026 03:05:10 +0100 Subject: [PATCH 1/2] perf(lineage): analyse each traced model against only its referenced relations Rich column lineage rebuilt the Polyglot schema for the whole project for every analysed model and handed all of it to analyze_query, so a column trace cost (models traced x project schema). Build the Polyglot table entries once per invocation and pass each model only the relations it references. Output is unchanged. Refs: CHI-433 --- .../compiler/lineage/_helpers/rich_columns.py | 33 ++++++-- .../sqlbuild/compiler/lineage/_test_types.py | 10 +++ .../src/sqlbuild/compiler/lineage/helpers.py | 27 +++++++ .../sqlbuild/compiler/lineage/test_columns.py | 81 +++++++++++++++++++ 4 files changed, 145 insertions(+), 6 deletions(-) diff --git a/src/sqlbuild/compiler/lineage/_helpers/rich_columns.py b/src/sqlbuild/compiler/lineage/_helpers/rich_columns.py index 6d1ac03d5d..2cbd7b89e7 100644 --- a/src/sqlbuild/compiler/lineage/_helpers/rich_columns.py +++ b/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: diff --git a/tests/unit/src/sqlbuild/compiler/lineage/_test_types.py b/tests/unit/src/sqlbuild/compiler/lineage/_test_types.py index 11d8c51746..b1af5052b4 100644 --- a/tests/unit/src/sqlbuild/compiler/lineage/_test_types.py +++ b/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 diff --git a/tests/unit/src/sqlbuild/compiler/lineage/helpers.py b/tests/unit/src/sqlbuild/compiler/lineage/helpers.py index 758d1aa3c1..4e2a123983 100644 --- a/tests/unit/src/sqlbuild/compiler/lineage/helpers.py +++ b/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) diff --git a/tests/unit/src/sqlbuild/compiler/lineage/test_columns.py b/tests/unit/src/sqlbuild/compiler/lineage/test_columns.py index d628fb85c7..630cc235fe 100644 --- a/tests/unit/src/sqlbuild/compiler/lineage/test_columns.py +++ b/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,6 +12,8 @@ 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 ( @@ -18,14 +21,17 @@ 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] From 691c45afb21a802dbea4dc11b82387cc3807d5f4 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Tue, 29 Sep 2026 03:05:10 +0100 Subject: [PATCH 2/2] test: guard column-trace lineage on the shared-dependency benchmark Add rich column traces upstream from the lattice rollup (text) and downstream from the hub (JSON) to the 3,000-model inspection guard with wall-time, peak-RSS and output-line ceilings. Refs: CHI-433 --- .../main/compile/COMPILE_BENCHMARKS.md | 5 ++-- .../test_inspection_command_performance.py | 29 +++++++++++++++++++ 2 files changed, 32 insertions(+), 2 deletions(-) diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/compile/COMPILE_BENCHMARKS.md b/tests/e2e/src/sqlbuild/cli/commands/main/compile/COMPILE_BENCHMARKS.md index 65fb6bec1a..180d6f4be7 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/compile/COMPILE_BENCHMARKS.md +++ b/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. diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/compile/test_inspection_command_performance.py b/tests/e2e/src/sqlbuild/cli/commands/main/compile/test_inspection_command_performance.py index 54f0750776..a7aae1f4a8 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/compile/test_inspection_command_performance.py +++ b/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"),