Skip to content

Commit 9655a4e

Browse files
authored
perf(lineage): analyse traced columns against only referenced relations (#877)
* 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 * 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
1 parent 0e65538 commit 9655a4e

6 files changed

Lines changed: 177 additions & 8 deletions

File tree

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

Lines changed: 27 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,7 @@ def build_rich_project_column_lineage(
8585
return None
8686

8787
schema: dict[str, dict[str, str]] = _build_schema_mapping(project)
88+
polyglot_tables: dict[str, dict[str, object]] = _polyglot_schema_tables(schema)
8889
model_results: dict[str, ModelColumnLineage] = {}
8990
collapsed_edges: list[ColumnLineageEdge] = []
9091

@@ -94,6 +95,7 @@ def build_rich_project_column_lineage(
9495
result: ModelColumnLineage | None = _build_polyglot_model_column_lineage(
9596
model=model,
9697
schema=schema,
98+
polyglot_tables=polyglot_tables,
9799
dialect=dialect,
98100
)
99101
if result is None:
@@ -123,6 +125,7 @@ def _build_polyglot_model_column_lineage(
123125
*,
124126
model: CompiledModel,
125127
schema: dict[str, dict[str, str]],
128+
polyglot_tables: dict[str, dict[str, object]],
126129
dialect: str | None,
127130
) -> ModelColumnLineage | None:
128131
normalized_sql: str
@@ -137,7 +140,10 @@ def _build_polyglot_model_column_lineage(
137140
normalized_sql,
138141
{
139142
"dialect": _polyglot_dialect(dialect),
140-
"schema": _polyglot_schema(schema),
143+
"schema": _referenced_polyglot_schema(
144+
polyglot_tables=polyglot_tables,
145+
physical_resources=physical_resources,
146+
),
141147
},
142148
)
143149
except polyglot_module.PolyglotError as error:
@@ -292,14 +298,29 @@ def _star_expanded_columns(analysis: dict[str, Any]) -> frozenset[str]:
292298
return frozenset(expanded)
293299

294300

295-
def _polyglot_schema(schema: dict[str, dict[str, str]]) -> dict[str, object]:
296-
tables: list[dict[str, object]] = []
297-
for table_name, columns in sorted(schema.items()):
301+
def _polyglot_schema_tables(schema: dict[str, dict[str, str]]) -> dict[str, dict[str, object]]:
302+
tables: dict[str, dict[str, object]] = {}
303+
for table_name, columns in schema.items():
298304
table_columns: list[dict[str, str]] = []
299305
for column_name, column_type in sorted(columns.items()):
300306
table_columns.append({"name": column_name, "type": column_type or "UNKNOWN"})
301-
tables.append({"name": table_name, "columns": table_columns})
302-
return {"tables": tables}
307+
tables[table_name] = {"name": table_name, "columns": table_columns}
308+
return tables
309+
310+
311+
def _referenced_polyglot_schema(
312+
*,
313+
polyglot_tables: dict[str, dict[str, object]],
314+
physical_resources: tuple[PhysicalResource, ...],
315+
) -> dict[str, object]:
316+
referenced: set[str] = {resource.physical_name for resource in physical_resources}
317+
return {
318+
"tables": [
319+
polyglot_tables[table_name]
320+
for table_name in sorted(referenced)
321+
if table_name in polyglot_tables
322+
]
323+
}
303324

304325

305326
def _polyglot_dialect(dialect: str | TypeDialect | None) -> PolyglotAnalysisDialect:

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

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -64,8 +64,9 @@ fans the last layer back in. Every lattice model is reachable along exponentiall
6464
compile guards do not enable the lattice, so their fingerprints and budgets are unchanged.
6565

6666
After one compile and warm lineage and scope caches, the guard runs `sqb lineage` downstream from
67-
the hub and upstream from the rollup as text and JSON, `sqb dag --json` and `sqb scope --json`,
68-
each in a fresh process. Every command
67+
the hub and upstream from the rollup as text and JSON, rich column traces upstream from
68+
`shared_orders_rollup.amount` and downstream from `shared_orders_hub.amount`, `sqb dag --json` and
69+
`sqb scope --json`, each in a fresh process. Every command
6970
has a wall-time ceiling, a peak-RSS ceiling and an output-line bound linear in the model count, so
7071
a renderer or graph walk that re-expands shared nodes per path fails on size before it can hide
7172
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‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -101,6 +101,35 @@
101101
expected_max_wall_seconds=3.0,
102102
expected_max_rss_bytes=128 * _MIB,
103103
),
104+
InspectionCommandPerformanceGuardTestCase(
105+
description="models_3000_lineage_column_text_upstream_from_fan_in_rollup",
106+
sqb_args=("lineage", f"{SHARED_DIAMOND_ROLLUP}.amount", "--direction", "upstream"),
107+
expected_fragments=(
108+
f"Column trace {SHARED_DIAMOND_ROLLUP}.amount upstream",
109+
"shared_orders_l15_s00.amount (cast)",
110+
),
111+
expected_max_output_lines=_MODEL_COUNT,
112+
expected_max_wall_seconds=8.0,
113+
expected_max_rss_bytes=512 * _MIB,
114+
),
115+
InspectionCommandPerformanceGuardTestCase(
116+
description="models_3000_lineage_column_json_downstream_from_fan_out_hub",
117+
sqb_args=(
118+
"lineage",
119+
f"{SHARED_DIAMOND_HUB}.amount",
120+
"--direction",
121+
"downstream",
122+
"--format",
123+
"json",
124+
),
125+
expected_fragments=(
126+
'"direction": "downstream"',
127+
f'"resource_name": "{SHARED_DIAMOND_ROLLUP}"',
128+
),
129+
expected_max_output_lines=20 * _MODEL_COUNT,
130+
expected_max_wall_seconds=8.0,
131+
expected_max_rss_bytes=512 * _MIB,
132+
),
104133
InspectionCommandPerformanceGuardTestCase(
105134
description="models_3000_dag_json",
106135
sqb_args=("dag", "--json"),

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,16 @@ class SqlAnalysisDisabledLineageTestCase:
3131
expected_result_is_none: bool
3232

3333

34+
@dataclass(frozen=True)
35+
class RichLineageAnalysisScalingTestCase:
36+
description: str
37+
chain_length: int
38+
unrelated_model_counts: tuple[int, ...]
39+
expected_schema_builds: int
40+
expected_analysis_calls: int
41+
expected_schema_tables_per_call: tuple[tuple[str, ...], ...]
42+
43+
3444
@dataclass(frozen=True)
3545
class ExpectedCountTestCase:
3646
description: str

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

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -90,3 +90,30 @@ def make_compiled_seed(*, name: str, columns: tuple[str, ...]) -> CompiledSeed:
9090

9191
def edge_label(source_name: str, source_column: str, target_name: str, target_column: str) -> str:
9292
return f"{source_name}.{source_column}->{target_name}.{target_column}"
93+
94+
95+
def make_chain_models(
96+
*, chain_length: int, unrelated_model_count: int
97+
) -> tuple[CompiledModel, ...]:
98+
root: CompiledModel = make_compiled_model(
99+
name="orders_0",
100+
query_sql="SELECT 1 AS order_id, 10 AS amount",
101+
inferred_columns=("order_id", "amount"),
102+
)
103+
chain: tuple[CompiledModel, ...] = tuple(
104+
make_compiled_model(
105+
name=f"orders_{index}",
106+
query_sql=f'SELECT order_id, amount + 1 AS amount FROM __ref("orders_{index - 1}")',
107+
inferred_columns=("order_id", "amount"),
108+
)
109+
for index in range(1, chain_length)
110+
)
111+
unrelated: tuple[CompiledModel, ...] = tuple(
112+
make_compiled_model(
113+
name=f"inventory_{index}",
114+
query_sql="SELECT 1 AS product_id, 5 AS quantity",
115+
inferred_columns=("product_id", "quantity"),
116+
)
117+
for index in range(unrelated_model_count)
118+
)
119+
return (root, *chain, *unrelated)

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

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
from __future__ import annotations
22

33
from dataclasses import replace
4+
from typing import Any
45

56
import pytest
67

@@ -11,21 +12,26 @@
1112
CompiledProject,
1213
)
1314
from sqlbuild.compiler.compile.types import CompiledResourceType
15+
from sqlbuild.compiler.lineage._helpers import rich_columns
16+
from sqlbuild.compiler.lineage._helpers.columns import _build_schema_mapping
1417
from sqlbuild.compiler.lineage.main.columns import build_project_column_lineage
1518
from sqlbuild.compiler.lineage.models import ColumnLineage, ModelColumnLineage, ProjectColumnLineage
1619
from sqlbuild.compiler.lineage.types import (
1720
ColumnLineageConfidence,
1821
ColumnLineageMode,
1922
ColumnTransformKind,
2023
)
24+
from sqlbuild.compiler.sql_analysis.main.import_polyglot_sql import import_polyglot_sql
2125
from tests.unit.src.sqlbuild.compiler.lineage._test_types import (
2226
ColumnLineageAnalyzerTestCase,
2327
ExpectedCountTestCase,
2428
ProjectLineageGraphTestCase,
29+
RichLineageAnalysisScalingTestCase,
2530
SqlAnalysisDisabledLineageTestCase,
2631
)
2732
from tests.unit.src.sqlbuild.compiler.lineage.helpers import (
2833
edge_label,
34+
make_chain_models,
2935
make_compiled_model,
3036
make_compiled_project,
3137
make_compiled_seed,
@@ -653,3 +659,78 @@ def test_given_sql_analysis_disabled_when_building_column_lineage_then_returns_n
653659
result: ProjectColumnLineage | None = build_project_column_lineage(project=project)
654660

655661
assert (result is None) is test_case.expected_result_is_none
662+
663+
664+
@pytest.mark.parametrize(
665+
"test_case",
666+
[
667+
RichLineageAnalysisScalingTestCase(
668+
description="builds the schema once and analyses each traced model against its refs",
669+
chain_length=4,
670+
unrelated_model_counts=(0, 250),
671+
expected_schema_builds=1,
672+
expected_analysis_calls=4,
673+
expected_schema_tables_per_call=(
674+
(),
675+
("__sqlbuild_model__orders_0",),
676+
("__sqlbuild_model__orders_1",),
677+
("__sqlbuild_model__orders_2",),
678+
),
679+
)
680+
],
681+
ids=lambda case: case.description,
682+
)
683+
def test_given_unrelated_models_when_tracing_rich_lineage_then_analysis_work_stays_linear(
684+
test_case: RichLineageAnalysisScalingTestCase,
685+
monkeypatch: pytest.MonkeyPatch,
686+
) -> None:
687+
polyglot_module: Any = import_polyglot_sql()
688+
analyze_query: Any = polyglot_module.analyze_query
689+
schema_builds: list[CompiledProject] = []
690+
analysed_tables: list[tuple[str, ...]] = []
691+
692+
def counting_schema_mapping(project: CompiledProject) -> dict[str, dict[str, str]]:
693+
schema_builds.append(project)
694+
return _build_schema_mapping(project)
695+
696+
def recording_analyze_query(sql: str, options: dict[str, Any]) -> object:
697+
analysed_tables.append(tuple(table["name"] for table in options["schema"]["tables"]))
698+
return analyze_query(sql, options)
699+
700+
monkeypatch.setattr(rich_columns, "_build_schema_mapping", counting_schema_mapping)
701+
monkeypatch.setattr(polyglot_module, "analyze_query", recording_analyze_query)
702+
chain_names: frozenset[str] = frozenset(
703+
f"orders_{index}" for index in range(test_case.chain_length)
704+
)
705+
edges_by_project_size: list[tuple[str, ...]] = []
706+
for unrelated_model_count in test_case.unrelated_model_counts:
707+
schema_builds.clear()
708+
analysed_tables.clear()
709+
project: CompiledProject = make_compiled_project(
710+
models=make_chain_models(
711+
chain_length=test_case.chain_length,
712+
unrelated_model_count=unrelated_model_count,
713+
)
714+
)
715+
716+
result: ProjectColumnLineage | None = build_project_column_lineage(
717+
project=project, model_names=chain_names
718+
)
719+
720+
assert result is not None
721+
assert len(schema_builds) == test_case.expected_schema_builds
722+
assert len(analysed_tables) == test_case.expected_analysis_calls
723+
assert tuple(analysed_tables) == test_case.expected_schema_tables_per_call
724+
edges_by_project_size.append(
725+
tuple(
726+
edge_label(
727+
edge.source.resource_name,
728+
edge.source.column_name,
729+
edge.target.resource_name,
730+
edge.target.column_name,
731+
)
732+
for edge in result.edges
733+
)
734+
)
735+
assert len(set(edges_by_project_size)) == 1
736+
assert "orders_2.amount->orders_3.amount" in edges_by_project_size[0]

0 commit comments

Comments
 (0)