From 596e775961f46f0997b9a8a2a9a66e5b1148b7c4 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Tue, 6 Oct 2026 00:23:24 +0100 Subject: [PATCH 1/2] perf(compile): batch analysis preparation normalization Analysis preparation normalized each model's stubbed SQL with its own native call, releasing and reacquiring the GIL once per model. Under the analysis worker threads most of that time was spent waiting to reacquire the GIL rather than normalizing. Normalize every batch member in one detached native call instead. A compile catalog runs the batch in parallel on its analysis pool while no analysis batch occupies the pool, and in turn otherwise, so preparing the next ready batch never queues behind a running one. Each member keeps its own result or error in place, so the same first error is raised at the same point as before, and prepared payloads, cache keys, and cache hits are unchanged. --- .../src/bindings/_helpers/functions.rs | 20 +- .../src/bindings/main/catch_compiler_panic.rs | 5 + .../src/bindings/main/mod.rs | 2 + .../bindings/main/normalization_results.rs | 17 ++ .../semantic_validation/_helpers/catalog.rs | 57 ++++- .../main/normalize_batch.rs | 28 ++- .../src/semantic_validation/models.rs | 3 + .../src/semantic_validation/tests.rs | 2 + .../tests/test_normalize_batch.rs | 89 ++++++++ .../semantic_validation/tests/test_types.rs | 7 + .../src/semantic_validation/types.rs | 2 + .../compile/_helpers/analysis/compact.py | 162 +++++++++----- .../compile/_helpers/sharing/binding.py | 89 +++++--- .../main/_normalize_analysis_batch.py | 30 ++- src/sqlbuild/compiler/sql_analysis/types.py | 8 +- .../cli/commands/main/compile/_test_types.py | 12 ++ .../compile/test_dense_batched_preparation.py | 61 ++++++ .../sqlbuild/compiler/pipeline/_test_types.py | 19 ++ .../src/sqlbuild/compiler/pipeline/helpers.py | 153 +++++++++++++- .../test_batched_analysis_preparation.py | 199 ++++++++++++++++++ .../compiler/sql_analysis/main/_test_types.py | 24 +++ .../compiler/sql_analysis/main/helpers.py | 21 ++ .../main/test_normalize_analysis_batch.py | 125 +++++++++++ 23 files changed, 1021 insertions(+), 114 deletions(-) create mode 100644 crates/sqlbuild-rules-native/src/bindings/main/catch_compiler_panic.rs create mode 100644 crates/sqlbuild-rules-native/src/bindings/main/normalization_results.rs create mode 100644 crates/sqlbuild-rules-native/src/semantic_validation/tests/test_normalize_batch.rs create mode 100644 tests/e2e/src/sqlbuild/cli/commands/main/compile/test_dense_batched_preparation.py create mode 100644 tests/integration/src/sqlbuild/compiler/pipeline/test_batched_analysis_preparation.py create mode 100644 tests/unit/src/sqlbuild/compiler/sql_analysis/main/test_normalize_analysis_batch.py diff --git a/crates/sqlbuild-rules-native/src/bindings/_helpers/functions.rs b/crates/sqlbuild-rules-native/src/bindings/_helpers/functions.rs index fd4aabb6c4..267114a012 100644 --- a/crates/sqlbuild-rules-native/src/bindings/_helpers/functions.rs +++ b/crates/sqlbuild-rules-native/src/bindings/_helpers/functions.rs @@ -38,18 +38,24 @@ fn normalize_dialect_sql(py: Python<'_>, sql: &str, dialect: &str) -> PyResult, dialect: &str, - requests: Vec<(String, std::collections::HashMap)>, -) -> PyResult> { - py.compiler_detach(|| { - crate::semantic_validation::main::normalize_batch::normalize_analysis_sqls( - dialect, requests, + requests: Vec, +) -> PyResult>> { + let results = py.compiler_detach(|| { + Ok( + crate::semantic_validation::main::normalize_batch::normalize_analysis_sqls( + dialect, requests, None, + ), ) - }) - .map_err(value_error) + }); + crate::bindings::main::normalization_results::normalization_results( + py, + results.map_err(value_error)?, + ) } #[pyfunction] diff --git a/crates/sqlbuild-rules-native/src/bindings/main/catch_compiler_panic.rs b/crates/sqlbuild-rules-native/src/bindings/main/catch_compiler_panic.rs new file mode 100644 index 0000000000..f87fc5b8ec --- /dev/null +++ b/crates/sqlbuild-rules-native/src/bindings/main/catch_compiler_panic.rs @@ -0,0 +1,5 @@ +pub(crate) fn catch_compiler_panic( + operation: impl FnOnce() -> Result, +) -> Result { + crate::bindings::_helpers::panics::catch_compiler_panic(operation) +} diff --git a/crates/sqlbuild-rules-native/src/bindings/main/mod.rs b/crates/sqlbuild-rules-native/src/bindings/main/mod.rs index 3b8691d5be..9f6415af70 100644 --- a/crates/sqlbuild-rules-native/src/bindings/main/mod.rs +++ b/crates/sqlbuild-rules-native/src/bindings/main/mod.rs @@ -1,3 +1,5 @@ +pub(crate) mod catch_compiler_panic; pub(crate) mod compiler_error; pub(crate) mod compiler_guard; +pub(crate) mod normalization_results; pub mod register; diff --git a/crates/sqlbuild-rules-native/src/bindings/main/normalization_results.rs b/crates/sqlbuild-rules-native/src/bindings/main/normalization_results.rs new file mode 100644 index 0000000000..88cb1a632a --- /dev/null +++ b/crates/sqlbuild-rules-native/src/bindings/main/normalization_results.rs @@ -0,0 +1,17 @@ +use pyo3::prelude::{IntoPyObject, Py, PyAny, PyResult, Python}; + +/// Return normalized SQL in order, carrying each failure as the exception a single call raises. +pub(crate) fn normalization_results( + py: Python<'_>, + results: Vec>, +) -> PyResult>> { + results + .into_iter() + .map(|result| match result { + Ok(sql) => Ok(sql.into_pyobject(py)?.into_any().unbind()), + Err(error) => Ok(crate::bindings::_helpers::panics::compiler_error(error) + .into_value(py) + .into_any()), + }) + .collect() +} diff --git a/crates/sqlbuild-rules-native/src/semantic_validation/_helpers/catalog.rs b/crates/sqlbuild-rules-native/src/semantic_validation/_helpers/catalog.rs index 37ff225df0..a44bde62ff 100644 --- a/crates/sqlbuild-rules-native/src/semantic_validation/_helpers/catalog.rs +++ b/crates/sqlbuild-rules-native/src/semantic_validation/_helpers/catalog.rs @@ -1,14 +1,19 @@ //! Compile-owned schema catalog and native, batched binding requests. use crate::bindings::main::compiler_error::compiler_error; +use crate::bindings::main::normalization_results::normalization_results; use crate::bindings::types::CompilerDetach; +use crate::semantic_validation::main::normalize_batch; use std::collections::HashMap; +use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex, PoisonError}; use crate::semantic_validation::models::{ CatalogInput, Columns, CompactAnalysisJob, FunctionProbes, ProjectCatalog, }; -use crate::semantic_validation::types::{BindingRequest, DiagnosticRow, Relations}; +use crate::semantic_validation::types::{ + BindingRequest, DiagnosticRow, NormalizationRequest, Relations, +}; use crate::semantic_validation::{ _helpers::{diagnostics, identifiers}, main as validation, @@ -19,7 +24,7 @@ use polyglot_sql::{ ValidationSchema, }; use pyo3::exceptions::PyValueError; -use pyo3::prelude::{Bound, FromPyObject, PyAny, PyAnyMethods, PyResult, Python}; +use pyo3::prelude::{Bound, FromPyObject, Py, PyAny, PyAnyMethods, PyResult, Python}; use pyo3::pymethods; use pyo3::types::{PyBytes, PyDict, PyDictMethods}; use rayon::iter::{IntoParallelIterator, ParallelIterator}; @@ -137,6 +142,7 @@ impl ProjectCatalog { analysis_tables: HashMap::new(), function_probes: Arc::new(FunctionProbes::default()), analysis_pool: Arc::default(), + running_analyses: Arc::default(), }; catalog.update_relations(relations); Ok(catalog) @@ -158,6 +164,7 @@ impl ProjectCatalog { analysis_tables: self.analysis_tables.clone(), function_probes: Arc::new(FunctionProbes::default()), analysis_pool: Arc::clone(&self.analysis_pool), + running_analyses: Arc::clone(&self.running_analyses), }; catalog.update_relations(relations); catalog @@ -234,6 +241,7 @@ impl ProjectCatalog { ) -> PyResult>> { py.compiler_detach(|| { let pool = self.analysis_pool()?; + let _running = RunningAnalysis::start(&self.running_analyses); pool.install(|| { requests .into_par_iter() @@ -253,6 +261,44 @@ impl ProjectCatalog { }) .map_err(compiler_error) } + + /// Normalize a preparation batch on the idle analysis pool, else in turn, in one detached call. + fn normalize_analysis_sqls( + &self, + py: Python<'_>, + dialect: &str, + requests: Vec, + ) -> PyResult>> { + let results = py + .compiler_detach(|| { + let pool = (self.running_analyses.load(Ordering::Acquire) == 0) + .then(|| self.analysis_pool()) + .transpose()?; + Ok(normalize_batch::normalize_analysis_sqls( + dialect, + requests, + pool.as_deref(), + )) + }) + .map_err(compiler_error)?; + normalization_results(py, results) + } +} + +/// Count one analysis batch as running on the catalog's pool until dropped. +struct RunningAnalysis<'a>(&'a AtomicUsize); + +impl<'a> RunningAnalysis<'a> { + fn start(running: &'a AtomicUsize) -> Self { + running.fetch_add(1, Ordering::AcqRel); + Self(running) + } +} + +impl Drop for RunningAnalysis<'_> { + fn drop(&mut self) { + self.0.fetch_sub(1, Ordering::AcqRel); + } } #[pymethods] @@ -265,7 +311,10 @@ impl CompactAnalysisJob { .take() .ok_or_else(|| PyValueError::new_err("compact analysis job already ran"))?; let result = py - .compiler_detach(|| analysis(&self.catalog)) + .compiler_detach(|| { + let _running = RunningAnalysis::start(&self.catalog.running_analyses); + analysis(&self.catalog) + }) .map_err(compiler_error)?; Ok(PyBytes::new(py, result.as_bytes())) } @@ -283,6 +332,7 @@ impl ProjectCatalog { analysis_tables: HashMap::new(), function_probes: Arc::clone(&self.function_probes), analysis_pool: Arc::clone(&self.analysis_pool), + running_analyses: Arc::clone(&self.running_analyses), } } @@ -321,6 +371,7 @@ impl ProjectCatalog { analysis_tables: HashMap::new(), function_probes: Arc::new(FunctionProbes::default()), analysis_pool: Arc::default(), + running_analyses: Arc::default(), } } pub(crate) fn analysis_schema( diff --git a/crates/sqlbuild-rules-native/src/semantic_validation/main/normalize_batch.rs b/crates/sqlbuild-rules-native/src/semantic_validation/main/normalize_batch.rs index 5337ae4090..0ad8408745 100644 --- a/crates/sqlbuild-rules-native/src/semantic_validation/main/normalize_batch.rs +++ b/crates/sqlbuild-rules-native/src/semantic_validation/main/normalize_batch.rs @@ -1,19 +1,27 @@ -use std::collections::HashMap; +use rayon::ThreadPool; +use rayon::iter::{IntoParallelIterator, ParallelIterator}; +use crate::bindings::main::catch_compiler_panic::catch_compiler_panic; +use crate::semantic_validation::types::NormalizationRequest; + +/// Normalize every request, keeping each failure in place so callers raise the first in order. pub(crate) fn normalize_analysis_sqls( dialect: &str, - requests: Vec<(String, HashMap)>, -) -> Result, String> { - let mut results: Vec = Vec::with_capacity(requests.len()); - for (sql, placeholders) in requests { - results.push( + requests: Vec, + pool: Option<&ThreadPool>, +) -> Vec> { + let normalize = |(sql, stubs, placeholders): NormalizationRequest| { + catch_compiler_panic(|| { crate::semantic_validation::_helpers::normalization::normalize_analysis_sql( &sql, dialect, - HashMap::new(), + stubs, placeholders, - )?, - ); + ) + }) + }; + match pool { + Some(pool) => pool.install(|| requests.into_par_iter().map(normalize).collect()), + None => requests.into_iter().map(normalize).collect(), } - Ok(results) } diff --git a/crates/sqlbuild-rules-native/src/semantic_validation/models.rs b/crates/sqlbuild-rules-native/src/semantic_validation/models.rs index 46afc66c34..8e94c6ec28 100644 --- a/crates/sqlbuild-rules-native/src/semantic_validation/models.rs +++ b/crates/sqlbuild-rules-native/src/semantic_validation/models.rs @@ -6,6 +6,7 @@ use polyglot_sql::{DialectType, SchemaValidationOptions}; use pyo3::{FromPyObject, pyclass}; use rayon::ThreadPool; use std::collections::HashMap; +use std::sync::atomic::AtomicUsize; use std::sync::{Arc, Mutex}; #[derive(Default, Debug)] @@ -57,6 +58,8 @@ pub(crate) struct ProjectCatalog { pub(crate) function_probes: Arc, /// One analysis pool per compile catalog, built on first use and shared by its views. pub(super) analysis_pool: Arc>>>, + /// Analysis batches running on the pool; preparation only borrows the pool while none run. + pub(super) running_analyses: Arc, } /// One resolved compact analysis batch that runs without borrowing its project catalog. diff --git a/crates/sqlbuild-rules-native/src/semantic_validation/tests.rs b/crates/sqlbuild-rules-native/src/semantic_validation/tests.rs index bd969772d8..418d1ace8a 100644 --- a/crates/sqlbuild-rules-native/src/semantic_validation/tests.rs +++ b/crates/sqlbuild-rules-native/src/semantic_validation/tests.rs @@ -1,5 +1,7 @@ #[path = "tests/test_function_names.rs"] mod function_names; +#[path = "tests/test_normalize_batch.rs"] +mod normalize_batch; #[path = "tests/test_positions.rs"] mod positions; #[path = "tests/test_semantic_validation.rs"] diff --git a/crates/sqlbuild-rules-native/src/semantic_validation/tests/test_normalize_batch.rs b/crates/sqlbuild-rules-native/src/semantic_validation/tests/test_normalize_batch.rs new file mode 100644 index 0000000000..3cdbaceca7 --- /dev/null +++ b/crates/sqlbuild-rules-native/src/semantic_validation/tests/test_normalize_batch.rs @@ -0,0 +1,89 @@ +use std::collections::HashMap; + +use crate::semantic_validation::main::normalize_batch::normalize_analysis_sqls; +use crate::semantic_validation::tests::test_types::NormalizationBatchTestCase; +use crate::semantic_validation::types::NormalizationRequest; + +#[test] +fn given_batch_when_normalizing_on_pool_or_in_turn_then_matches_single_calls_in_order() { + let test_cases = [ + NormalizationBatchTestCase { + description: "stubbed references and placeholders", + dialect: "snowflake", + requests: &[ + ( + r#"SELECT * FROM __ref("orders") JOIN __ref("customers") USING (id)"#, + &[("orders", "__sqb_rel_0"), ("customers", "__sqb_rel_1")], + ), + ("SELECT payload:item[0] AS item FROM inventory", &[]), + ("SELECT 1 LIMIT @@@limit", &[]), + ], + expected_failures: &[], + }, + NormalizationBatchTestCase { + description: "failures stay at their own positions", + dialect: "snowflake", + requests: &[ + ("SELECT 'unterminated FROM orders", &[]), + ("SELECT id FROM products", &[]), + ("SELECT /* unterminated FROM customers", &[]), + ], + expected_failures: &[0, 2], + }, + ]; + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(3) + .build() + .expect("test pool must build"); + + for test_case in test_cases { + let requests = || -> Vec { + test_case + .requests + .iter() + .map(|(sql, stubs)| { + ( + (*sql).to_owned(), + stubs + .iter() + .map(|(name, stub)| ((*name).to_owned(), (*stub).to_owned())) + .collect(), + HashMap::from([("limit".to_owned(), "10".to_owned())]), + ) + }) + .collect() + }; + let expected: Vec> = requests() + .into_iter() + .map(|(sql, stubs, placeholders)| { + crate::semantic_validation::_helpers::normalization::normalize_analysis_sql( + &sql, + test_case.dialect, + stubs, + placeholders, + ) + }) + .collect(); + let failures: Vec = (0..expected.len()) + .filter(|index| expected[*index].is_err()) + .collect(); + + assert_eq!( + failures, test_case.expected_failures, + "{}", + test_case.description + ); + assert_eq!( + normalize_analysis_sqls(test_case.dialect, requests(), Some(&pool)), + expected, + "{}", + test_case.description + ); + assert_eq!( + normalize_analysis_sqls(test_case.dialect, requests(), None), + expected, + "{}", + test_case.description + ); + } +} diff --git a/crates/sqlbuild-rules-native/src/semantic_validation/tests/test_types.rs b/crates/sqlbuild-rules-native/src/semantic_validation/tests/test_types.rs index 030a8221de..521de7be8e 100644 --- a/crates/sqlbuild-rules-native/src/semantic_validation/tests/test_types.rs +++ b/crates/sqlbuild-rules-native/src/semantic_validation/tests/test_types.rs @@ -45,3 +45,10 @@ pub(crate) struct TypeLocationTestCase { pub(crate) findings: &'static [(&'static str, &'static str)], pub(crate) expected_starts: &'static [&'static str], } + +pub(crate) struct NormalizationBatchTestCase { + pub(crate) description: &'static str, + pub(crate) dialect: &'static str, + pub(crate) requests: &'static [(&'static str, &'static [(&'static str, &'static str)])], + pub(crate) expected_failures: &'static [usize], +} diff --git a/crates/sqlbuild-rules-native/src/semantic_validation/types.rs b/crates/sqlbuild-rules-native/src/semantic_validation/types.rs index 1e171ebe37..87ebd9353e 100644 --- a/crates/sqlbuild-rules-native/src/semantic_validation/types.rs +++ b/crates/sqlbuild-rules-native/src/semantic_validation/types.rs @@ -3,6 +3,8 @@ use std::collections::HashMap; pub(crate) type Relations = HashMap; pub(crate) type BindingRequest = (String, Vec<(String, bool)>, Relations); +/// One analysis SQL with its relation stubs and placeholder defaults. +pub(crate) type NormalizationRequest = (String, HashMap, HashMap); pub(crate) type DiagnosticRow = ( String, String, diff --git a/src/sqlbuild/compiler/compile/_helpers/analysis/compact.py b/src/sqlbuild/compiler/compile/_helpers/analysis/compact.py index 145dd4de11..3833027435 100644 --- a/src/sqlbuild/compiler/compile/_helpers/analysis/compact.py +++ b/src/sqlbuild/compiler/compile/_helpers/analysis/compact.py @@ -33,6 +33,7 @@ from sqlbuild.compiler.compile._helpers.sharing.binding import ( binding_query_fields, lineage_reference_map, + normalized_analysis_sql, prepare_binding_queries, relation_stubs_collide, remembered_shared_results, @@ -171,7 +172,9 @@ from sqlbuild.compiler.sql_analysis.main._binding_catalog import create_binding_catalog from sqlbuild.compiler.sql_analysis.main._decode_schema_validation import decode_schema_validation from sqlbuild.compiler.sql_analysis.main._normalize_analysis import normalize_analysis_sql -from sqlbuild.compiler.sql_analysis.main._normalize_analysis_batch import normalize_analysis_sqls +from sqlbuild.compiler.sql_analysis.main._normalize_analysis_batch import ( + normalize_analysis_sql_results, +) from sqlbuild.compiler.sql_analysis.main._schema_validation import get_schema_validations from sqlbuild.compiler.sql_analysis.main.import_polyglot_sql import import_polyglot_sql from sqlbuild.compiler.sql_analysis.models import ( @@ -666,11 +669,18 @@ def _compact_batch_context( if name in inputs.column_nullability_by_table }, ) - normalized_sqls: list[str] = normalize_analysis_sqls( - sqls=inputs.query_sqls, + requests: list[tuple[str, None, dict[str, str] | None]] = [ + (query_sql, None, query_placeholders) + for query_sql, query_placeholders in zip( + inputs.query_sqls, inputs.placeholders, strict=True + ) + ] + results: list[str | Exception] = normalize_analysis_sql_results( dialect=dialect, - placeholders=inputs.placeholders, + requests=requests, + catalog=binding_catalog.native if binding_catalog is not None else None, ) + normalized_sqls: list[str] = [normalized_analysis_sql(result) for result in results] binding_queries: tuple[PreparedBindingQuery | None, ...] = ( prepare_binding_queries( binding_catalog=binding_catalog, @@ -697,7 +707,6 @@ def _prepare_compact_analysis_batch( ) -> CompactBatchPreparation: query_sqls: tuple[str, ...] = inputs.query_sqls references: tuple[tuple[CompileSqlReference, ...], ...] = inputs.references - placeholders: tuple[dict[str, str] | None, ...] = inputs.placeholders column_nullability_by_table: dict[str, dict[str, InferredNullability]] = ( inputs.column_nullability_by_table ) @@ -735,63 +744,40 @@ def _prepare_compact_analysis_batch( ) binding_queries: tuple[PreparedBindingQuery | None, ...] = context.binding_queries shared_query_indexes: set[int] = set() - for ( - query_sql, + members: list[tuple[dict[str, tuple[CompiledResourceType, str]], dict[str, str]]] + stubbed_sqls: dict[int, str | Exception] + members, stubbed_sqls = _member_analysis_stubs( + inputs=inputs, + context=context, + binding_schemas=member_binding_schemas, + ) + for index, ( cleaned_sql, query_references, - query_placeholders, query_recover_cte_facts, binding_schema, binding_query, - ) in zip( - query_sqls, - normalized_sqls, - references, - placeholders, - recover_cte_facts, - member_binding_schemas, - binding_queries, - strict=True, - ): - lineage_references: dict[str, tuple[CompiledResourceType, str]] = lineage_reference_map( - query_references + (lineage_references, canonical_stubs), + ) in enumerate( + zip( + normalized_sqls, + references, + recover_cte_facts, + member_binding_schemas, + binding_queries, + members, + strict=True, ) + ): shared: SharedBindingQuery | None = ( binding_query.shared if binding_query is not None else None ) - stubbed_reference_names: frozenset[str] = ( - frozenset() - if binding_schema is not None - or relation_stubs_collide(cleaned_sql=cleaned_sql, names=lineage_references) - else frozenset(lineage_references) - - _qualified_reference_names( - query_sql=cleaned_sql, - reference_names=lineage_references.keys(), - ) - ) - canonical_stubs: dict[str, str] = ( - shared.stubs - if shared is not None - else { - name: ( - f"{COMPACT_RELATION_STUB_PREFIX}{index}" - if name in stubbed_reference_names - else name - ) - for index, name in enumerate(lineage_references) - } - ) analysis_sql: str = ( shared.sql if shared is not None else cleaned_sql if binding_schema is not None - else _cleaned_analysis_sql( - query_sql=query_sql, - placeholders=query_placeholders, - dialect=dialect, - relation_stubs=canonical_stubs, - ) + else normalized_analysis_sql(stubbed_sqls[index]) ) query: dict[str, object] = { "sql": analysis_sql, @@ -920,6 +906,86 @@ def _prepare_compact_analysis_batch( ) +def _member_analysis_stubs( + *, + inputs: CompactBatchInputs, + context: CompactBatchContext, + binding_schemas: tuple[dict[str, dict[str, str]] | None, ...], +) -> tuple[ + list[tuple[dict[str, tuple[CompiledResourceType, str]], dict[str, str]]], + dict[int, str | Exception], +]: + """Return each member's relation stubs and, in one native batch, its stubbed SQL.""" + + members: list[tuple[dict[str, tuple[CompiledResourceType, str]], dict[str, str]]] = [] + stubbed_requests: dict[int, tuple[str, dict[str, str], dict[str, str] | None]] = {} + for index, ( + query_sql, + cleaned_sql, + query_references, + query_placeholders, + member_schema, + member_query, + ) in enumerate( + zip( + inputs.query_sqls, + context.normalized_sqls, + inputs.references, + inputs.placeholders, + binding_schemas, + context.binding_queries, + strict=True, + ) + ): + member: tuple[dict[str, tuple[CompiledResourceType, str]], dict[str, str]] = ( + _member_relation_stubs( + cleaned_sql=cleaned_sql, + references=query_references, + binding_schema=member_schema, + binding_query=member_query, + ) + ) + members.append(member) + if member_schema is None and (member_query is None or member_query.shared is None): + stubbed_requests[index] = (query_sql, member[1], query_placeholders) + stubbed_sqls: list[str | Exception] = normalize_analysis_sql_results( + dialect=inputs.inference_profile.sql_analysis_dialect, + requests=tuple(stubbed_requests.values()), + catalog=(context.binding_catalog.native if context.binding_catalog is not None else None), + ) + return members, dict(zip(stubbed_requests, stubbed_sqls, strict=True)) + + +def _member_relation_stubs( + *, + cleaned_sql: str, + references: tuple[CompileSqlReference, ...], + binding_schema: dict[str, dict[str, str]] | None, + binding_query: PreparedBindingQuery | None, +) -> tuple[dict[str, tuple[CompiledResourceType, str]], dict[str, str]]: + """Return a member's lineage references and the relation names its analysis SQL uses.""" + + lineage_references: dict[str, tuple[CompiledResourceType, str]] = lineage_reference_map( + references + ) + if binding_query is not None and binding_query.shared is not None: + return lineage_references, binding_query.shared.stubs + stubbed_reference_names: frozenset[str] = ( + frozenset() + if binding_schema is not None + or relation_stubs_collide(cleaned_sql=cleaned_sql, names=lineage_references) + else frozenset(lineage_references) + - _qualified_reference_names( + query_sql=cleaned_sql, + reference_names=lineage_references.keys(), + ) + ) + return lineage_references, { + name: f"{COMPACT_RELATION_STUB_PREFIX}{index}" if name in stubbed_reference_names else name + for index, name in enumerate(lineage_references) + } + + def _run_compact_analysis_batch(*, preparation: CompactBatchPreparation) -> object: """Analyse every project size with the full worker pool; the rules phase sets peak memory.""" diff --git a/src/sqlbuild/compiler/compile/_helpers/sharing/binding.py b/src/sqlbuild/compiler/compile/_helpers/sharing/binding.py index 21ba613ae5..709e95c39a 100644 --- a/src/sqlbuild/compiler/compile/_helpers/sharing/binding.py +++ b/src/sqlbuild/compiler/compile/_helpers/sharing/binding.py @@ -30,7 +30,9 @@ ) from sqlbuild.compiler.compile.types import CompiledResourceType from sqlbuild.compiler.sql_analysis.constants import SQL_QUOTED_IDENTIFIER_DELIMITER -from sqlbuild.compiler.sql_analysis.main._normalize_analysis import normalize_analysis_sql +from sqlbuild.compiler.sql_analysis.main._normalize_analysis_batch import ( + normalize_analysis_sql_results, +) from sqlbuild.compiler.sql_analysis.models import SqlSchemaValidationRequest @@ -108,6 +110,35 @@ def prepare_binding_queries( if prekey is not None and count >= MIN_SHARED_BINDING_QUERY_MEMBERS ) ) + shared_stubs: dict[int, dict[str, str]] = {} + for index, binding_schema in enumerate(binding_schemas): + prekey: str | None = prekeys[index] + if binding_schema is None or prekey is None or prekey not in candidates: + continue + names: list[str] | None = _shared_reference_names( + cleaned_sql=cleaned_sqls[index], + references=inputs.references[index], + binding_schema=binding_schema, + ) + if names is not None: + shared_stubs[index] = { + name: f"{COMPACT_RELATION_STUB_PREFIX}{position}" + for position, name in enumerate(names) + } + shared_sqls: dict[int, str | Exception] = dict( + zip( + shared_stubs, + normalize_analysis_sql_results( + dialect=dialect, + requests=[ + (inputs.query_sqls[index], stubs, inputs.placeholders[index]) + for index, stubs in shared_stubs.items() + ], + catalog=binding_catalog.native, + ), + strict=True, + ) + ) prepared: list[PreparedBindingQuery | None] = [] for index, binding_schema in enumerate(binding_schemas): if binding_schema is None: @@ -120,19 +151,18 @@ def prepare_binding_queries( ) ] )[0] - prekey: str | None = prekeys[index] - shared: SharedBindingQuery | None = None - if prekey is not None and prekey in candidates: - shared = _shared_binding_query( + shared: SharedBindingQuery | None = ( + _shared_binding_query( binding_catalog=binding_catalog, - query=SqlSchemaValidationRequest( - sql=inputs.query_sqls[index], dialect=dialect, schema=binding_schema - ), - cleaned_sql=cleaned_sqls[index], - references=inputs.references[index], - placeholders=inputs.placeholders[index], + sql=normalized_analysis_sql(shared_sqls[index]), + dialect=dialect, + schema=binding_schema, + stubs=shared_stubs[index], overrides=overrides, ) + if index in shared_stubs + else None + ) prepared.append( PreparedBindingQuery(references=binding_references, overrides=overrides, shared=shared) ) @@ -190,6 +220,14 @@ def shared_result_is_exact(*, validation: object, template: object) -> bool: return isinstance(validation, dict) and cast(dict[str, object], validation).get("errors") == [] +def normalized_analysis_sql(result: str | Exception) -> str: + """Return one normalized batch member, raising the error its own normalization raised.""" + + if isinstance(result, Exception): + raise result + return result + + def relation_stubs_collide(*, cleaned_sql: str, names: Iterable[str]) -> bool: """Return whether authored SQL or relation names already use the relation stub prefix.""" @@ -226,29 +264,20 @@ def _shared_reference_names( def _shared_binding_query( *, binding_catalog: Any, - query: SqlSchemaValidationRequest, - cleaned_sql: str, - references: tuple[CompileSqlReference, ...], - placeholders: dict[str, str] | None, + sql: str, + dialect: str | None, + schema: Mapping[str, Mapping[str, str]], + stubs: dict[str, str], overrides: Mapping[str, Mapping[str, str]], -) -> SharedBindingQuery | None: - names: list[str] | None = _shared_reference_names( - cleaned_sql=cleaned_sql, references=references, binding_schema=query.schema - ) - if names is None: - return None - stubs: dict[str, str] = { - name: f"{COMPACT_RELATION_STUB_PREFIX}{index}" for index, name in enumerate(names) - } - sql: str = normalize_analysis_sql( - sql=query.sql, dialect=query.dialect, stubs=stubs, placeholders=placeholders - ) +) -> SharedBindingQuery: + """Identify one query from its SQL normalized with relation `stubs`.""" + analysis_shapes: list[tuple[object, ...]] = [] - for name in names: + for name in stubs: types, nullability = binding_catalog.analysis_shapes.get(name, ({}, {})) analysis_shapes.append((stubs[name], tuple(types.items()), tuple(nullability.items()))) binding_shapes: list[tuple[object, ...]] = [] - for relation, columns in query.schema.items(): + for relation, columns in schema.items(): effective: Mapping[str, str] = ( overrides[relation] if relation in overrides @@ -264,7 +293,7 @@ def _shared_binding_query( stubs=stubs, key=( sql, - query.dialect or "generic", + dialect or "generic", tuple(analysis_shapes), tuple(sorted(binding_shapes)), ), diff --git a/src/sqlbuild/compiler/sql_analysis/main/_normalize_analysis_batch.py b/src/sqlbuild/compiler/sql_analysis/main/_normalize_analysis_batch.py index c6ebd32ad1..326b35422e 100644 --- a/src/sqlbuild/compiler/sql_analysis/main/_normalize_analysis_batch.py +++ b/src/sqlbuild/compiler/sql_analysis/main/_normalize_analysis_batch.py @@ -1,21 +1,31 @@ -"""Normalize a compile batch under one native GIL-release boundary.""" +"""Normalize analysis SQL batches under one native GIL-release boundary.""" +from collections.abc import Mapping, Sequence from typing import cast import sqlbuild._native as _native -from sqlbuild.compiler.sql_analysis.types import NativePositionsModule +from sqlbuild.compiler.sql_analysis.types import ( + NativeNormalizationRequest, + NativePositionsModule, + NativeProjectCatalog, +) -def normalize_analysis_sqls( +def normalize_analysis_sql_results( *, - sqls: tuple[str, ...], dialect: str | None, - placeholders: tuple[dict[str, str] | None, ...], -) -> list[str]: - requests: list[tuple[str, dict[str, str]]] = [ - (sql, defaults or {}) for sql, defaults in zip(sqls, placeholders, strict=True) + requests: Sequence[tuple[str, Mapping[str, str] | None, Mapping[str, str] | None]], + catalog: NativeProjectCatalog | None, +) -> list[str | Exception]: + """Normalize `(sql, stubs, placeholders)` in order, keeping each failure in its place.""" + + if not requests: + return [] + payload: list[NativeNormalizationRequest] = [ + (sql, dict(stubs or {}), dict(placeholders or {})) for sql, stubs, placeholders in requests ] + if catalog is not None: + return catalog.normalize_analysis_sqls(dialect=dialect or "generic", requests=payload) return cast(NativePositionsModule, _native).normalize_analysis_sqls( - dialect=dialect or "generic", - requests=requests, + dialect=dialect or "generic", requests=payload ) diff --git a/src/sqlbuild/compiler/sql_analysis/types.py b/src/sqlbuild/compiler/sql_analysis/types.py index f490eb9c05..7e69e03f36 100644 --- a/src/sqlbuild/compiler/sql_analysis/types.py +++ b/src/sqlbuild/compiler/sql_analysis/types.py @@ -7,6 +7,7 @@ type NativeDiagnosticRow = tuple[str, str, int | None, int | None, int | None, int | None, str] type NativeBindingRequest = tuple[str, list[tuple[str, bool]], dict[str, Mapping[str, str]]] +type NativeNormalizationRequest = tuple[str, dict[str, str], dict[str, str]] class NativeCompactAnalysisJob(Protocol): @@ -30,6 +31,9 @@ def update_relations(self, relations: Mapping[str, Mapping[str, str]]) -> None: def binding_results( self, requests: list[NativeBindingRequest] ) -> list[list[NativeDiagnosticRow]]: ... + def normalize_analysis_sqls( + self, *, dialect: str, requests: list[NativeNormalizationRequest] + ) -> list[str | Exception]: ... class NativeCatalogModule(Protocol): @@ -47,8 +51,8 @@ def position( class NativePositionsModule(Protocol): def normalize_analysis_sqls( - self, *, dialect: str, requests: list[tuple[str, dict[str, str]]] - ) -> list[str]: ... + self, *, dialect: str, requests: list[NativeNormalizationRequest] + ) -> list[str | Exception]: ... def binding_diagnostics( self, *, sql: str, dialect: str, rows: list[NativeDiagnosticRow] ) -> list[NativeDiagnosticRow]: ... diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/compile/_test_types.py b/tests/e2e/src/sqlbuild/cli/commands/main/compile/_test_types.py index a4f6dea82a..b399e498dd 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/compile/_test_types.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/compile/_test_types.py @@ -6,6 +6,8 @@ from dataclasses import dataclass, field from pathlib import Path +from tests.integration.src.sqlbuild.compiler.pipeline._test_types import PreparedCompile + @dataclass(frozen=True) class SourceSemanticBindingCase: @@ -759,3 +761,13 @@ class RenderLoadNoticeTestCase: description: str notice_bytes: int expected_notice: bool + + +@dataclass(frozen=True) +class DenseBatchedPreparationTestCase: + description: str + model_count: int + expected_exit_code: int + + +type DensePreparedCompile = PreparedCompile diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/compile/test_dense_batched_preparation.py b/tests/e2e/src/sqlbuild/cli/commands/main/compile/test_dense_batched_preparation.py new file mode 100644 index 0000000000..fffac8f29e --- /dev/null +++ b/tests/e2e/src/sqlbuild/cli/commands/main/compile/test_dense_batched_preparation.py @@ -0,0 +1,61 @@ +"""Dense batched analysis preparation must match per-model preparation byte for byte.""" + +from pathlib import Path + +import pytest + +from scripts.cold_compile_performance._helpers.dense_project import write_dense_compile_project +from sqlbuild.compiler.compile.constants import COMPACT_RELATION_STUB_PREFIX +from tests.e2e.src.sqlbuild.cli.commands.main.compile._test_types import ( + DenseBatchedPreparationTestCase, + DensePreparedCompile, +) +from tests.integration.src.sqlbuild.compiler.pipeline.helpers import ( + keep_batched_normalization, + prepared_compile, + use_per_model_normalization, +) + + +@pytest.mark.performance +@pytest.mark.parametrize( + "test_case", + (DenseBatchedPreparationTestCase("dense_models_3000", 3000, 0),), + ids=lambda case: case.description, +) +def test_given_dense_project_when_compiling_cold_then_batched_preparation_matches_per_model( + test_case: DenseBatchedPreparationTestCase, + tmp_path: Path, + capsys: pytest.CaptureFixture[str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + batched_dir: Path = tmp_path / "batched" + reference_dir: Path = tmp_path / "reference" + write_dense_compile_project(project_dir=batched_dir, model_count=test_case.model_count) + write_dense_compile_project(project_dir=reference_dir, model_count=test_case.model_count) + + expected: DensePreparedCompile = prepared_compile( + project_dir=reference_dir, + args=(), + capsys=capsys, + monkeypatch=monkeypatch, + normalization=use_per_model_normalization, + ) + actual: DensePreparedCompile = prepared_compile( + project_dir=batched_dir, + args=(), + capsys=capsys, + monkeypatch=monkeypatch, + normalization=keep_batched_normalization, + ) + + assert actual[3] == expected[3] + assert actual[1] == expected[1] + assert actual[2] == expected[2] + assert actual[0] == expected[0] + assert actual[0][0] == test_case.expected_exit_code + assert any(COMPACT_RELATION_STUB_PREFIX.encode() in payload for payload in actual[3]) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-vv"])) diff --git a/tests/integration/src/sqlbuild/compiler/pipeline/_test_types.py b/tests/integration/src/sqlbuild/compiler/pipeline/_test_types.py index 2642c01a53..81a1f14d2d 100644 --- a/tests/integration/src/sqlbuild/compiler/pipeline/_test_types.py +++ b/tests/integration/src/sqlbuild/compiler/pipeline/_test_types.py @@ -2,6 +2,7 @@ from collections.abc import Callable from dataclasses import dataclass, field +from pathlib import Path from scripts.cold_compile_performance.models import RandomDagProject from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter @@ -393,6 +394,9 @@ class FunctionArgumentTypeCase: type CompileOutcome = tuple[int, dict[str, object], dict[str, bytes]] +type PreparedCompile = tuple[ + CompileOutcome, tuple[str, ...], dict[str, object | None], tuple[bytes, ...] +] @dataclass(frozen=True) @@ -477,3 +481,18 @@ class DataflowStartFailureCase: expected_error: str expected_live_workers: tuple[str, ...] expected_notices: int + + +@dataclass(frozen=True) +class BatchedPreparationCase: + description: str + write_project: Callable[[Path], tuple[str, ...]] + reshaped: tuple[int, ...] + expected_exit_code: int + + +@dataclass(frozen=True) +class PerturbedPreparationCase: + description: str + project: RandomDagProject + expected_payloads_match: bool diff --git a/tests/integration/src/sqlbuild/compiler/pipeline/helpers.py b/tests/integration/src/sqlbuild/compiler/pipeline/helpers.py index e0b685df6a..66f7ada616 100644 --- a/tests/integration/src/sqlbuild/compiler/pipeline/helpers.py +++ b/tests/integration/src/sqlbuild/compiler/pipeline/helpers.py @@ -3,16 +3,23 @@ import itertools import json import random +import shutil import signal import threading import time -from collections.abc import Callable +from collections.abc import Callable, Mapping, Sequence from contextlib import suppress from pathlib import Path -from typing import Any +from types import ModuleType +from typing import Any, cast +import orjson import pytest +import sqlbuild._native as _native +from scripts.cold_compile_performance._helpers.dense_project import write_dense_compile_project +from scripts.cold_compile_performance._helpers.random_dag_project import write_random_dag_project +from scripts.cold_compile_performance.models import RandomDagProject from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter from sqlbuild.adapters.duckdb.classes.duckdb_adapter import DuckDbAdapter from sqlbuild.cli.commands.main.entrypoint.entry import main @@ -21,6 +28,7 @@ from sqlbuild.compiler.compile._helpers.assembly import binding_waves from sqlbuild.compiler.compile._helpers.assembly import project as project_assembly from sqlbuild.compiler.compile._helpers.assembly.binding_waves import analyze_binding_waves +from sqlbuild.compiler.compile._helpers.sharing import binding as binding_sharing from sqlbuild.compiler.compile.classes.binding_dataflow import BindingDataflow from sqlbuild.compiler.compile.main._build_compile_inputs import build_compile_inputs from sqlbuild.compiler.compile.models import ( @@ -37,14 +45,25 @@ from sqlbuild.compiler.pipeline.main.compile import run_compile_pipeline from sqlbuild.compiler.pipeline.main.project import compile_project from sqlbuild.compiler.pipeline.models import CompilePipelineOptions, CompilePipelineResult +from sqlbuild.compiler.sql_analysis.main import _normalize_analysis_batch as normalize_batch +from sqlbuild.compiler.sql_analysis.types import NativePositionsModule, NativeProjectCatalog from sqlbuild.runtime.contracts.models import ConnectionHooks from sqlbuild.sql_values.types import CollectionRendering from tests.integration.src.sqlbuild.compiler.pipeline._test_types import ( CompileOutcome, DataflowScheduleCase, + PreparedCompile, SharedBindingQueryCase, ) +_REPOSITORY_ROOT: Path = Path(__file__).resolve().parents[6] +_NORMALIZATION_MODULES: tuple[ModuleType, ...] = (normalize_batch, binding_sharing, compact) +_ANALYSIS_CACHE_METRICS: tuple[str, ...] = ( + "analysis_batch_cache_hits", + "analysis_entry_cache_hits", + "analysis_cache_misses", + "analysis_cache_bypasses", +) _SCHEMA_FIXTURE_PATH: Path = ( Path(__file__).resolve().parents[5] / "fixtures" / "dbt_manifest_v12_schema.json" ) @@ -263,11 +282,19 @@ def compile_outcome( ) -> CompileOutcome: """Compile through the CLI; return the exit code, timing-free JSON, and artifacts.""" + return timed_compile_outcome(project_dir=project_dir, args=args, capsys=capsys)[0] + + +def timed_compile_outcome( + *, project_dir: Path, args: tuple[str, ...], capsys: pytest.CaptureFixture[str] +) -> tuple[CompileOutcome, dict[str, object]]: + """Compile through the CLI; return the timing-free outcome and its compile timings.""" + exit_code: int = main(["--project-dir", str(project_dir), "compile", "--json", *args]) payload: dict[str, object] = json.loads( capsys.readouterr().out.replace(str(project_dir), "") ) - _ = payload.pop("compile_timings", None) + timings: dict[str, object] = cast(dict[str, object], payload.pop("compile_timings", {})) compiled: Path = project_dir / "target" / "compiled" return ( exit_code, @@ -276,7 +303,7 @@ def compile_outcome( path.relative_to(compiled).as_posix(): path.read_bytes() for path in compiled.rglob("*.sql") }, - ) + ), timings def reshape_models(*, project_dir: Path, names: tuple[str, ...], indexes: tuple[int, ...]) -> None: @@ -513,3 +540,121 @@ def live_model_analysis_threads() -> list[str]: thread.name ) return sorted(names_by_kind.get(True, [])) + + +def random_dag_writer(project: RandomDagProject) -> Callable[[Path], tuple[str, ...]]: + return lambda project_dir: write_random_dag_project(project_dir=project_dir, project=project) + + +def fixture_writer(relative: str) -> Callable[[Path], tuple[str, ...]]: + def write(project_dir: Path) -> tuple[str, ...]: + _ = shutil.copytree( + _REPOSITORY_ROOT / relative, + project_dir, + ignore=shutil.ignore_patterns("target", "*.duckdb"), + ) + return () + + return write + + +def dense_writer(model_count: int) -> Callable[[Path], tuple[str, ...]]: + def write(project_dir: Path) -> tuple[str, ...]: + write_dense_compile_project(project_dir=project_dir, model_count=model_count) + return () + + return write + + +def use_per_model_normalization(monkeypatch: pytest.MonkeyPatch) -> None: + """Normalize each batch member with its own native call, the reference preparation.""" + + for module in _NORMALIZATION_MODULES: + monkeypatch.setattr(module, "normalize_analysis_sql_results", _per_model_normalization) + + +def keep_batched_normalization(monkeypatch: pytest.MonkeyPatch) -> None: + """Leave batched normalization in place, the preparation under test.""" + + _ = monkeypatch + + +def perturb_batched_normalization(monkeypatch: pytest.MonkeyPatch) -> None: + """Change the last member of every batched normalization, which comparisons must catch.""" + + batched: Callable[..., list[str | Exception]] = normalize_batch.normalize_analysis_sql_results + + def perturbed(**arguments: Any) -> list[str | Exception]: + results: list[str | Exception] = batched(**arguments) + return [*results[:-1], *(f"{result} " for result in results[-1:])] + + for module in _NORMALIZATION_MODULES: + monkeypatch.setattr(module, "normalize_analysis_sql_results", perturbed) + + +def prepared_compile( + *, + project_dir: Path, + args: tuple[str, ...], + capsys: pytest.CaptureFixture[str], + monkeypatch: pytest.MonkeyPatch, + normalization: Callable[[pytest.MonkeyPatch], None], +) -> PreparedCompile: + """Compile once; return the outcome, cache keys, cache metrics, and native batch payloads.""" + + with monkeypatch.context() as patch: + normalization(patch) + preparations: list[CompactBatchPreparation] = trace_native_compact_batches(patch) + cache_keys: list[str] = _trace_analysis_cache_keys(patch) + outcome: CompileOutcome + timings: dict[str, object] + outcome, timings = timed_compile_outcome(project_dir=project_dir, args=args, capsys=capsys) + return ( + outcome, + tuple(sorted(cache_keys)), + {name: timings.get(name) for name in _ANALYSIS_CACHE_METRICS}, + tuple(_preparation_payload(preparation) for preparation in preparations), + ) + + +def _trace_analysis_cache_keys(monkeypatch: pytest.MonkeyPatch) -> list[str]: + original: Callable[..., str] = project_assembly.model_analysis_cache_key + keys: list[str] = [] + + def traced(**arguments: Any) -> str: + key: str = original(**arguments) + keys.append(key) + return key + + monkeypatch.setattr(project_assembly, "model_analysis_cache_key", traced) + return keys + + +def _per_model_normalization( + *, + dialect: str | None, + requests: Sequence[tuple[str, Mapping[str, str] | None, Mapping[str, str] | None]], + catalog: NativeProjectCatalog | None, +) -> list[str | Exception]: + _ = catalog + native: NativePositionsModule = cast(NativePositionsModule, _native) + return [ + native.normalize_analysis_sqls( + dialect=dialect or "generic", + requests=[(sql, dict(stubs or {}), dict(placeholders or {}))], + )[0] + for sql, stubs, placeholders in requests + ] + + +def _preparation_payload(preparation: CompactBatchPreparation) -> bytes: + return orjson.dumps( + { + "cleaned_sql": preparation.cleaned_sql, + "queries": preparation.queries, + "templates": preparation.templates, + "projections": preparation.projections, + "shared_query_indexes": sorted(preparation.shared_query_indexes), + }, + option=orjson.OPT_SORT_KEYS, + ) diff --git a/tests/integration/src/sqlbuild/compiler/pipeline/test_batched_analysis_preparation.py b/tests/integration/src/sqlbuild/compiler/pipeline/test_batched_analysis_preparation.py new file mode 100644 index 0000000000..c60d1408c3 --- /dev/null +++ b/tests/integration/src/sqlbuild/compiler/pipeline/test_batched_analysis_preparation.py @@ -0,0 +1,199 @@ +"""Batched analysis preparation must match per-model preparation byte for byte.""" + +from pathlib import Path + +import pytest + +from scripts.cold_compile_performance.models import RandomDagProject +from sqlbuild.compiler.compile.main import _build_compile_inputs +from tests.integration.src.sqlbuild.compiler.pipeline._test_types import ( + BatchedPreparationCase, + PerturbedPreparationCase, + PreparedCompile, +) +from tests.integration.src.sqlbuild.compiler.pipeline.helpers import ( + dense_writer, + fixture_writer, + keep_batched_normalization, + perturb_batched_normalization, + prepared_compile, + random_dag_writer, + reshape_models, + use_per_model_normalization, + use_wave_analysis, +) + + +@pytest.mark.parametrize( + "test_case", + [ + BatchedPreparationCase( + "inferred shapes", + random_dag_writer(RandomDagProject(seed=11, model_count=48)), + (16,), + 0, + ), + BatchedPreparationCase( + "binding errors mid-graph", + random_dag_writer(RandomDagProject(seed=23, model_count=48, errors=True)), + (24,), + 1, + ), + BatchedPreparationCase( + "run ids and analysis opt-outs", + random_dag_writer( + RandomDagProject(seed=37, model_count=40, run_ids=True, analysis_opt_outs=True) + ), + (20,), + 0, + ), + BatchedPreparationCase( + "reference cycle", + random_dag_writer(RandomDagProject(seed=67, model_count=24, cycle=True)), + (8,), + 0, + ), + BatchedPreparationCase( + "waffle shop fixture", fixture_writer("tests/e2e/fixtures/waffle_shop"), (), 0 + ), + BatchedPreparationCase( + "tidy shop example", fixture_writer("website/examples/tidy-shop"), (), 0 + ), + BatchedPreparationCase("dense single batch", dense_writer(120), (), 0), + ], + ids=lambda case: case.description, +) +def test_given_project_when_compiling_in_waves_then_batched_preparation_matches_per_model( + test_case: BatchedPreparationCase, + tmp_path: Path, + capsys: pytest.CaptureFixture[str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + batched_dir: Path = tmp_path / "batched" + reference_dir: Path = tmp_path / "reference" + names: tuple[str, ...] = test_case.write_project(batched_dir) + _ = test_case.write_project(reference_dir) + monkeypatch.setattr( + _build_compile_inputs, + "resolve_run_id", + lambda *, selected_run_id: selected_run_id or "run-0", + ) + use_wave_analysis(monkeypatch) + for step, reshaped, args in ( + ("cold", (), ()), + ("warm", (), ()), + ("edit", test_case.reshaped, ()), + ("uncached", (), ("--no-cache",)), + ): + reshape_models(project_dir=batched_dir, names=names, indexes=reshaped) + reshape_models(project_dir=reference_dir, names=names, indexes=reshaped) + expected: PreparedCompile = prepared_compile( + project_dir=reference_dir, + args=args, + capsys=capsys, + monkeypatch=monkeypatch, + normalization=use_per_model_normalization, + ) + actual: PreparedCompile = prepared_compile( + project_dir=batched_dir, + args=args, + capsys=capsys, + monkeypatch=monkeypatch, + normalization=keep_batched_normalization, + ) + + assert actual[3] == expected[3], f"{step} native batch payloads differ" + assert actual[1] == expected[1], f"{step} analysis cache keys differ" + assert actual[2] == expected[2], f"{step} analysis cache hits differ" + assert actual[0] == expected[0], f"{step} compile differs from per-model preparation" + assert actual[0][0] == test_case.expected_exit_code + + +@pytest.mark.parametrize( + "test_case", + [ + BatchedPreparationCase( + "inferred shapes", random_dag_writer(RandomDagProject(seed=13, model_count=48)), (), 0 + ), + BatchedPreparationCase( + "binding errors mid-graph", + random_dag_writer(RandomDagProject(seed=29, model_count=48, errors=True)), + (), + 1, + ), + ], + ids=lambda case: case.description, +) +def test_given_duckdb_project_when_compiling_dataflow_with_and_without_cache_then_matches_per_model( + test_case: BatchedPreparationCase, + tmp_path: Path, + capsys: pytest.CaptureFixture[str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + batched_dir: Path = tmp_path / "batched" + reference_dir: Path = tmp_path / "reference" + _ = test_case.write_project(batched_dir) + _ = test_case.write_project(reference_dir) + monkeypatch.setattr( + _build_compile_inputs, + "resolve_run_id", + lambda *, selected_run_id: selected_run_id or "run-0", + ) + for args in ((), (), ("--no-cache",)): + expected: PreparedCompile = prepared_compile( + project_dir=reference_dir, + args=args, + capsys=capsys, + monkeypatch=monkeypatch, + normalization=use_per_model_normalization, + ) + actual: PreparedCompile = prepared_compile( + project_dir=batched_dir, + args=args, + capsys=capsys, + monkeypatch=monkeypatch, + normalization=keep_batched_normalization, + ) + + assert actual[:3] == expected[:3] + assert actual[0][0] == test_case.expected_exit_code + + +@pytest.mark.parametrize( + "test_case", + [PerturbedPreparationCase("last member", RandomDagProject(seed=11, model_count=24), False)], + ids=lambda case: case.description, +) +def test_given_perturbed_batch_result_when_compiling_then_payloads_differ_from_per_model( + test_case: PerturbedPreparationCase, + tmp_path: Path, + capsys: pytest.CaptureFixture[str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + batched_dir: Path = tmp_path / "batched" + reference_dir: Path = tmp_path / "reference" + _ = random_dag_writer(test_case.project)(batched_dir) + _ = random_dag_writer(test_case.project)(reference_dir) + use_wave_analysis(monkeypatch) + + expected: PreparedCompile = prepared_compile( + project_dir=reference_dir, + args=("--no-cache",), + capsys=capsys, + monkeypatch=monkeypatch, + normalization=use_per_model_normalization, + ) + actual: PreparedCompile = prepared_compile( + project_dir=batched_dir, + args=("--no-cache",), + capsys=capsys, + monkeypatch=monkeypatch, + normalization=perturb_batched_normalization, + ) + + assert len(actual[3]) == len(expected[3]) > 0 + assert (actual[3] == expected[3]) is test_case.expected_payloads_match + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-n", "auto", "--dist", "loadfile"])) diff --git a/tests/unit/src/sqlbuild/compiler/sql_analysis/main/_test_types.py b/tests/unit/src/sqlbuild/compiler/sql_analysis/main/_test_types.py index 159458e9b0..55b239b2fa 100644 --- a/tests/unit/src/sqlbuild/compiler/sql_analysis/main/_test_types.py +++ b/tests/unit/src/sqlbuild/compiler/sql_analysis/main/_test_types.py @@ -1,7 +1,10 @@ from __future__ import annotations +from collections.abc import Callable from dataclasses import dataclass +from sqlbuild.compiler.sql_analysis.types import NativeProjectCatalog + @dataclass(frozen=True) class SemanticDiagnosticMappingCase: @@ -47,3 +50,24 @@ class PolyglotExplicitGuardTestCase: function_depth: int maximum_function_depth: int expected_error_pattern: str + + +@dataclass(frozen=True) +class NormalizationBatchTestCase: + description: str + dialect: str + requests: tuple[tuple[str, dict[str, str], dict[str, str]], ...] + catalog: Callable[[str], NativeProjectCatalog | None] + expected_stubbed_sql: str + + +@dataclass(frozen=True) +class NormalizationBatchFailureTestCase: + description: str + dialect: str + before: tuple[str, ...] + failing: str + after: tuple[str, ...] + catalog: Callable[[str], NativeProjectCatalog | None] + expected_error_type: type[Exception] + expected_result_types: tuple[str, ...] diff --git a/tests/unit/src/sqlbuild/compiler/sql_analysis/main/helpers.py b/tests/unit/src/sqlbuild/compiler/sql_analysis/main/helpers.py index c73c266f3a..efd7dcee2a 100644 --- a/tests/unit/src/sqlbuild/compiler/sql_analysis/main/helpers.py +++ b/tests/unit/src/sqlbuild/compiler/sql_analysis/main/helpers.py @@ -1,8 +1,29 @@ """Helpers for SQL analysis public-entry tests.""" +from sqlbuild.compiler.sql_analysis.classes.binding_catalog import BindingCatalog +from sqlbuild.compiler.sql_analysis.types import NativeProjectCatalog + def nested_coalesce_sql(*, function_depth: int) -> str: expression: str = "value" for _ in range(function_depth): expression = f"COALESCE({expression}, 0)" return f"SELECT {expression} AS resolved_value FROM records" + + +def pool_catalog(dialect: str) -> NativeProjectCatalog: + """Return a compile catalog, which normalizes batches on its analysis pool.""" + + return BindingCatalog( + dialect=dialect, + quoted_ignore_case=False, + known_functions=(), + known_types=(), + relations={}, + ).native + + +def without_catalog(dialect: str) -> None: + """Return no catalog, which normalizes batches in turn.""" + + _ = dialect diff --git a/tests/unit/src/sqlbuild/compiler/sql_analysis/main/test_normalize_analysis_batch.py b/tests/unit/src/sqlbuild/compiler/sql_analysis/main/test_normalize_analysis_batch.py new file mode 100644 index 0000000000..f3de5765e5 --- /dev/null +++ b/tests/unit/src/sqlbuild/compiler/sql_analysis/main/test_normalize_analysis_batch.py @@ -0,0 +1,125 @@ +from __future__ import annotations + +import pytest + +from sqlbuild.compiler.sql_analysis.main._normalize_analysis import normalize_analysis_sql +from sqlbuild.compiler.sql_analysis.main._normalize_analysis_batch import ( + normalize_analysis_sql_results, +) +from tests.unit.src.sqlbuild.compiler.sql_analysis.main._test_types import ( + NormalizationBatchFailureTestCase, + NormalizationBatchTestCase, +) +from tests.unit.src.sqlbuild.compiler.sql_analysis.main.helpers import ( + pool_catalog, + without_catalog, +) + + +@pytest.mark.parametrize( + "test_case", + ( + NormalizationBatchTestCase( + description="stubbed relations and placeholders in turn", + dialect="snowflake", + requests=( + ( + 'SELECT * FROM __ref("orders") JOIN __ref("customers") USING (customer_id)', + {"orders": "__sqb_rel_0", "customers": "__sqb_rel_1"}, + {}, + ), + ("SELECT payload:item[0] AS item FROM inventory", {}, {}), + ("SELECT 1 AS quantity LIMIT @@@limit", {}, {"limit": "10"}), + ), + catalog=without_catalog, + expected_stubbed_sql="SELECT * FROM __sqb_rel_0 JOIN __sqb_rel_1 USING (customer_id)", + ), + NormalizationBatchTestCase( + description="stubbed relations and placeholders on the catalog pool", + dialect="snowflake", + requests=( + ( + 'SELECT * FROM __ref("orders") JOIN __ref("customers") USING (customer_id)', + {"orders": "__sqb_rel_0", "customers": "__sqb_rel_1"}, + {}, + ), + ("SELECT payload:item[0] AS item FROM inventory", {}, {}), + ("SELECT 1 AS quantity LIMIT @@@limit", {}, {"limit": "10"}), + ), + catalog=pool_catalog, + expected_stubbed_sql="SELECT * FROM __sqb_rel_0 JOIN __sqb_rel_1 USING (customer_id)", + ), + ), + ids=lambda case: case.description, +) +def test_given_batch_when_normalizing_then_matches_single_normalizations_in_order( + test_case: NormalizationBatchTestCase, +) -> None: + expected: list[str | Exception] = [ + normalize_analysis_sql( + sql=sql, dialect=test_case.dialect, stubs=stubs, placeholders=placeholders + ) + for sql, stubs, placeholders in test_case.requests + ] + + results: list[str | Exception] = normalize_analysis_sql_results( + dialect=test_case.dialect, + requests=test_case.requests, + catalog=test_case.catalog(test_case.dialect), + ) + + assert results == expected + assert results[0] == test_case.expected_stubbed_sql + + +@pytest.mark.parametrize( + "test_case", + ( + NormalizationBatchFailureTestCase( + description="unterminated quote in turn", + dialect="snowflake", + before=("SELECT id FROM products",), + failing="SELECT 'unterminated FROM orders", + after=("SELECT /* unterminated FROM customers", "SELECT id FROM customers"), + catalog=without_catalog, + expected_error_type=ValueError, + expected_result_types=("str", "ValueError", "ValueError", "str"), + ), + NormalizationBatchFailureTestCase( + description="unterminated quote on the catalog pool", + dialect="snowflake", + before=("SELECT id FROM products",), + failing="SELECT 'unterminated FROM orders", + after=("SELECT /* unterminated FROM customers", "SELECT id FROM customers"), + catalog=pool_catalog, + expected_error_type=ValueError, + expected_result_types=("str", "ValueError", "ValueError", "str"), + ), + ), + ids=lambda case: case.description, +) +def test_given_failing_member_when_normalizing_batch_then_keeps_its_single_call_error_in_place( + test_case: NormalizationBatchFailureTestCase, +) -> None: + with pytest.raises(test_case.expected_error_type) as raised: + _ = normalize_analysis_sql(sql=test_case.failing, dialect=test_case.dialect) + + results: list[str | Exception] = normalize_analysis_sql_results( + dialect=test_case.dialect, + requests=[ + (sql, None, None) for sql in (*test_case.before, test_case.failing, *test_case.after) + ], + catalog=test_case.catalog(test_case.dialect), + ) + first_error: str | Exception = results[len(test_case.before)] + + assert results[: len(test_case.before)] == [ + normalize_analysis_sql(sql=sql, dialect=test_case.dialect) for sql in test_case.before + ] + assert type(first_error) is type(raised.value) + assert str(first_error) == str(raised.value) + assert tuple(type(result).__name__ for result in results) == test_case.expected_result_types + + +if __name__ == "__main__": + pytest.main([__file__, "-vv"]) From aaefad5f21ed4540ea493e03280080875b7f8ddf Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Tue, 6 Oct 2026 01:44:40 +0100 Subject: [PATCH 2/2] perf(compile): bound analysis normalization batches to fixed chunks Normalize each preparation batch in chunks of 256 members so one native call never holds copies of a whole dense batch at once. This keeps warm compile peak memory at its previous level while preserving the single GIL release per chunk. --- .../compiler/sql_analysis/constants.py | 1 + .../main/_normalize_analysis_batch.py | 26 ++++++----- .../compiler/sql_analysis/main/_test_types.py | 9 ++++ .../main/test_normalize_analysis_batch.py | 45 +++++++++++++++++++ 4 files changed, 71 insertions(+), 10 deletions(-) diff --git a/src/sqlbuild/compiler/sql_analysis/constants.py b/src/sqlbuild/compiler/sql_analysis/constants.py index cfeba4649c..87feca62c3 100644 --- a/src/sqlbuild/compiler/sql_analysis/constants.py +++ b/src/sqlbuild/compiler/sql_analysis/constants.py @@ -181,3 +181,4 @@ "motherduck": "duckdb", "sqlserver": "tsql", } +ANALYSIS_NORMALIZATION_CHUNK_SIZE: int = 256 diff --git a/src/sqlbuild/compiler/sql_analysis/main/_normalize_analysis_batch.py b/src/sqlbuild/compiler/sql_analysis/main/_normalize_analysis_batch.py index 326b35422e..a8dc96e619 100644 --- a/src/sqlbuild/compiler/sql_analysis/main/_normalize_analysis_batch.py +++ b/src/sqlbuild/compiler/sql_analysis/main/_normalize_analysis_batch.py @@ -1,9 +1,10 @@ """Normalize analysis SQL batches under one native GIL-release boundary.""" -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Mapping, Sequence from typing import cast import sqlbuild._native as _native +from sqlbuild.compiler.sql_analysis.constants import ANALYSIS_NORMALIZATION_CHUNK_SIZE from sqlbuild.compiler.sql_analysis.types import ( NativeNormalizationRequest, NativePositionsModule, @@ -19,13 +20,18 @@ def normalize_analysis_sql_results( ) -> list[str | Exception]: """Normalize `(sql, stubs, placeholders)` in order, keeping each failure in its place.""" - if not requests: - return [] - payload: list[NativeNormalizationRequest] = [ - (sql, dict(stubs or {}), dict(placeholders or {})) for sql, stubs, placeholders in requests - ] - if catalog is not None: - return catalog.normalize_analysis_sqls(dialect=dialect or "generic", requests=payload) - return cast(NativePositionsModule, _native).normalize_analysis_sqls( - dialect=dialect or "generic", requests=payload + normalize: Callable[..., list[str | Exception]] = ( + catalog.normalize_analysis_sqls + if catalog is not None + else cast(NativePositionsModule, _native).normalize_analysis_sqls ) + results: list[str | Exception] = [] + for start in range(0, len(requests), ANALYSIS_NORMALIZATION_CHUNK_SIZE): + chunk: list[NativeNormalizationRequest] = [ + (sql, dict(stubs or {}), dict(placeholders or {})) + for sql, stubs, placeholders in requests[ + start : start + ANALYSIS_NORMALIZATION_CHUNK_SIZE + ] + ] + results.extend(normalize(dialect=dialect or "generic", requests=chunk)) + return results diff --git a/tests/unit/src/sqlbuild/compiler/sql_analysis/main/_test_types.py b/tests/unit/src/sqlbuild/compiler/sql_analysis/main/_test_types.py index 55b239b2fa..d3b072c465 100644 --- a/tests/unit/src/sqlbuild/compiler/sql_analysis/main/_test_types.py +++ b/tests/unit/src/sqlbuild/compiler/sql_analysis/main/_test_types.py @@ -71,3 +71,12 @@ class NormalizationBatchFailureTestCase: catalog: Callable[[str], NativeProjectCatalog | None] expected_error_type: type[Exception] expected_result_types: tuple[str, ...] + + +@dataclass(frozen=True) +class NormalizationChunkTestCase: + description: str + valid_count: int + failing_index: int + catalog: Callable[[str], NativeProjectCatalog | None] + expected_error_type: type[Exception] diff --git a/tests/unit/src/sqlbuild/compiler/sql_analysis/main/test_normalize_analysis_batch.py b/tests/unit/src/sqlbuild/compiler/sql_analysis/main/test_normalize_analysis_batch.py index f3de5765e5..98e6448a9d 100644 --- a/tests/unit/src/sqlbuild/compiler/sql_analysis/main/test_normalize_analysis_batch.py +++ b/tests/unit/src/sqlbuild/compiler/sql_analysis/main/test_normalize_analysis_batch.py @@ -2,6 +2,7 @@ import pytest +from sqlbuild.compiler.sql_analysis.constants import ANALYSIS_NORMALIZATION_CHUNK_SIZE from sqlbuild.compiler.sql_analysis.main._normalize_analysis import normalize_analysis_sql from sqlbuild.compiler.sql_analysis.main._normalize_analysis_batch import ( normalize_analysis_sql_results, @@ -9,6 +10,7 @@ from tests.unit.src.sqlbuild.compiler.sql_analysis.main._test_types import ( NormalizationBatchFailureTestCase, NormalizationBatchTestCase, + NormalizationChunkTestCase, ) from tests.unit.src.sqlbuild.compiler.sql_analysis.main.helpers import ( pool_catalog, @@ -121,5 +123,48 @@ def test_given_failing_member_when_normalizing_batch_then_keeps_its_single_call_ assert tuple(type(result).__name__ for result in results) == test_case.expected_result_types +@pytest.mark.parametrize( + "test_case", + ( + NormalizationChunkTestCase( + description="failure on a chunk boundary in turn", + valid_count=ANALYSIS_NORMALIZATION_CHUNK_SIZE * 2 + 3, + failing_index=ANALYSIS_NORMALIZATION_CHUNK_SIZE, + catalog=without_catalog, + expected_error_type=ValueError, + ), + NormalizationChunkTestCase( + description="failure on a chunk boundary on the catalog pool", + valid_count=ANALYSIS_NORMALIZATION_CHUNK_SIZE * 2 + 3, + failing_index=ANALYSIS_NORMALIZATION_CHUNK_SIZE, + catalog=pool_catalog, + expected_error_type=ValueError, + ), + ), + ids=lambda case: case.description, +) +def test_given_batch_beyond_one_chunk_when_normalizing_then_keeps_every_member_in_order( + test_case: NormalizationChunkTestCase, +) -> None: + valid: list[str] = [ + f"SELECT {index} AS order_id FROM orders" for index in range(test_case.valid_count) + ] + sqls: list[str] = [ + *valid[: test_case.failing_index], + "SELECT 'unterminated FROM orders", + *valid[test_case.failing_index :], + ] + + results: list[str | Exception] = normalize_analysis_sql_results( + dialect="snowflake", + requests=[(sql, None, None) for sql in sqls], + catalog=test_case.catalog("snowflake"), + ) + + assert len(results) == len(sqls) + assert isinstance(results[test_case.failing_index], test_case.expected_error_type) + assert [*results[: test_case.failing_index], *results[test_case.failing_index + 1 :]] == valid + + if __name__ == "__main__": pytest.main([__file__, "-vv"])