diff --git a/crates/sqlbuild-rules-native/src/bindings/_helpers/functions.rs b/crates/sqlbuild-rules-native/src/bindings/_helpers/functions.rs index 1ff8a22929..fd4aabb6c4 100644 --- a/crates/sqlbuild-rules-native/src/bindings/_helpers/functions.rs +++ b/crates/sqlbuild-rules-native/src/bindings/_helpers/functions.rs @@ -515,6 +515,7 @@ pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { .get_type::(), )?; module.add_class::()?; + module.add_class::()?; module.add_class::()?; module.add_function(wrap_pyfunction!(normalize_analysis_sql, module)?)?; module.add_function(wrap_pyfunction!(normalize_analysis_sqls, module)?)?; diff --git a/crates/sqlbuild-rules-native/src/query_analysis/_helpers/engine.rs b/crates/sqlbuild-rules-native/src/query_analysis/_helpers/engine.rs index 98955b002b..1f95a3bc07 100644 --- a/crates/sqlbuild-rules-native/src/query_analysis/_helpers/engine.rs +++ b/crates/sqlbuild-rules-native/src/query_analysis/_helpers/engine.rs @@ -323,6 +323,24 @@ pub(crate) fn analyze_project_compact_with_catalog( request_json: &str, catalog: Option<&crate::semantic_validation::models::ProjectCatalog>, ) -> Result { + analyze_resolved_compact_batch(resolve_compact_batch(request_json, catalog)?, catalog) +} + +/// Resolve catalog schemas now and return the analysis to run later without catalog tables. +pub(crate) fn prepare_project_compact_with_catalog( + request_json: &str, + catalog: &crate::semantic_validation::models::ProjectCatalog, +) -> Result { + let request = resolve_compact_batch(request_json, Some(catalog))?; + Ok(Box::new(move |view| { + analyze_resolved_compact_batch(request, Some(view)) + })) +} + +fn resolve_compact_batch( + request_json: &str, + catalog: Option<&crate::semantic_validation::models::ProjectCatalog>, +) -> Result { let mut request: CompactProjectAnalysisBatchRequest = serde_json::from_str(request_json).map_err(|error| error.to_string())?; for query in &mut request.queries { @@ -349,6 +367,14 @@ pub(crate) fn analyze_project_compact_with_catalog( } } } + Ok(request) +} + +/// Analyse a resolved batch; a catalog supplies options, probes, and its compile-wide pool. +fn analyze_resolved_compact_batch( + request: CompactProjectAnalysisBatchRequest, + catalog: Option<&crate::semantic_validation::models::ProjectCatalog>, +) -> Result { let workers = request.workers.clamp(1, MAX_WORKERS); if request .templates @@ -400,12 +426,16 @@ pub(crate) fn analyze_project_compact_with_catalog( }) .collect(); let sql_bytes: usize = request.queries.iter().map(|query| query.sql.len()).sum(); - let analysis_workers = workers.min(unique_query_count.max(1)); - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(analysis_workers) - .stack_size(ANALYSIS_WORKER_STACK_BYTES) - .build() - .map_err(|error| error.to_string())?; + let pool = match catalog { + Some(catalog) => catalog.analysis_pool()?, + None => std::sync::Arc::new( + rayon::ThreadPoolBuilder::new() + .num_threads(workers.min(unique_query_count.max(1))) + .stack_size(ANALYSIS_WORKER_STACK_BYTES) + .build() + .map_err(|error| error.to_string())?, + ), + }; let query_work: Vec = request .queries .into_iter() @@ -450,7 +480,7 @@ pub(crate) fn analyze_project_compact_with_catalog( serde_json::to_string(&response).map_err(|error| error.to_string()) } -/// Analyse a batch largest SQL first so no late wide query becomes its critical path. +/// Analyse largest SQL first on spawned drainers rather than a broadcast to every pool thread. fn analyze_largest_first( pool: &rayon::ThreadPool, batch: Vec, @@ -463,21 +493,33 @@ fn analyze_largest_first( .map(|work| Mutex::new(Some(work))) .collect(); let next = AtomicUsize::new(0); - let finished: Vec> = pool.broadcast(|_| { - let mut finished: Vec<(usize, CompiledQueryWorkResult)> = Vec::new(); - while let Some(index) = order.get(next.fetch_add(1, Ordering::Relaxed)).copied() { - let work = slots[index] - .lock() - .unwrap_or_else(PoisonError::into_inner) - .take(); - if let Some(work) = work { - finished.push((index, analyze_compact_query_work(work, catalog))); - } + let finished: Mutex> = Mutex::new(Vec::new()); + let drainers = pool.current_num_threads().min(slots.len()); + pool.scope(|scope| { + for _ in 0..drainers { + scope.spawn(|_| { + let mut drained: Vec<(usize, CompiledQueryWorkResult)> = Vec::new(); + while let Some(index) = order.get(next.fetch_add(1, Ordering::Relaxed)).copied() { + let work = slots[index] + .lock() + .unwrap_or_else(PoisonError::into_inner) + .take(); + if let Some(work) = work { + drained.push((index, analyze_compact_query_work(work, catalog))); + } + } + finished + .lock() + .unwrap_or_else(PoisonError::into_inner) + .extend(drained); + }); } - finished }); let mut results: Vec> = slots.iter().map(|_| None).collect(); - for (index, result) in finished.into_iter().flatten() { + for (index, result) in finished + .into_inner() + .unwrap_or_else(PoisonError::into_inner) + { results[index] = Some(result); } results.into_iter().flatten().collect() @@ -497,7 +539,7 @@ fn analyze_compact_query_work( && let Some(validation) = result.validation.take() { let local_probes = crate::semantic_validation::models::FunctionProbes::default(); - let probes = catalog.map_or(&local_probes, |catalog| &catalog.function_probes); + let probes = catalog.map_or(&local_probes, |catalog| catalog.function_probes.as_ref()); result.validation = Some(validation.and_then(|value| { crate::semantic_validation::main::map_diagnostics::map_diagnostics( &sql, dialect, value, probes, diff --git a/crates/sqlbuild-rules-native/src/query_analysis/main/analyze_project_catalog.rs b/crates/sqlbuild-rules-native/src/query_analysis/main/analyze_project_catalog.rs deleted file mode 100644 index d01c6710a2..0000000000 --- a/crates/sqlbuild-rules-native/src/query_analysis/main/analyze_project_catalog.rs +++ /dev/null @@ -1,6 +0,0 @@ -pub(crate) fn analyze_project_compact_with_catalog( - request_json: &str, - catalog: &crate::semantic_validation::models::ProjectCatalog, -) -> Result { - crate::query_analysis::engine::analyze_project_compact_with_catalog(request_json, Some(catalog)) -} diff --git a/crates/sqlbuild-rules-native/src/query_analysis/main/mod.rs b/crates/sqlbuild-rules-native/src/query_analysis/main/mod.rs index 9a89ac6347..196001e275 100644 --- a/crates/sqlbuild-rules-native/src/query_analysis/main/mod.rs +++ b/crates/sqlbuild-rules-native/src/query_analysis/main/mod.rs @@ -1,4 +1,4 @@ pub(crate) mod analyze; pub(crate) mod analyze_project; -pub(crate) mod analyze_project_catalog; pub(crate) mod analyze_project_compact; +pub(crate) mod prepare_project_catalog; diff --git a/crates/sqlbuild-rules-native/src/query_analysis/main/prepare_project_catalog.rs b/crates/sqlbuild-rules-native/src/query_analysis/main/prepare_project_catalog.rs new file mode 100644 index 0000000000..af2090a29b --- /dev/null +++ b/crates/sqlbuild-rules-native/src/query_analysis/main/prepare_project_catalog.rs @@ -0,0 +1,6 @@ +pub(crate) fn prepare_project_compact_with_catalog( + request_json: &str, + catalog: &crate::semantic_validation::models::ProjectCatalog, +) -> Result { + crate::query_analysis::engine::prepare_project_compact_with_catalog(request_json, catalog) +} diff --git a/crates/sqlbuild-rules-native/src/query_analysis/tests/helpers.rs b/crates/sqlbuild-rules-native/src/query_analysis/tests/helpers.rs index 2b346c861d..8f53fef98e 100644 --- a/crates/sqlbuild-rules-native/src/query_analysis/tests/helpers.rs +++ b/crates/sqlbuild-rules-native/src/query_analysis/tests/helpers.rs @@ -578,3 +578,64 @@ pub(crate) fn borrowed_facts_preserve_union_dependencies_and_join_nullability() ); true } + +pub(crate) fn catalog_views_reuse_one_analysis_pool() -> bool { + let catalog = crate::semantic_validation::models::ProjectCatalog::with_options( + polyglot_sql::DialectType::DuckDB, + polyglot_sql::SchemaValidationOptions::default(), + false, + ); + let first = catalog.analysis_pool().expect("test assumption must hold"); + let second = catalog.analysis_pool().expect("test assumption must hold"); + let from_view = catalog + .analysis_view() + .analysis_pool() + .expect("test assumption must hold"); + std::sync::Arc::ptr_eq(&first, &second) + && std::sync::Arc::ptr_eq(&first, &from_view) + && first.current_num_threads() == 4 +} + +pub(crate) fn concurrent_batches_match_serial_batches() -> bool { + let requests: Vec = (0..12) + .map(|index| { + json!({ + "queries": (0..=index % 5).map(|query| json!({ + "sql": format!( + "WITH ranked AS (SELECT {query} AS order_id, {index} AS amount) \ + SELECT order_id, amount * {query} AS amount FROM ranked" + ), + "dialect": "duckdb" + })).collect::>(), + "templates": (0..=index % 5).map(|query| json!({ + "queryIndex": query, "recoverCteFacts": true + })).collect::>(), + "projections": (0..=index % 5).map(|query| json!({ + "templateIndex": query + })).collect::>(), + "workers": 4 + }) + .to_string() + }) + .collect(); + let serial: Vec = requests + .iter() + .map(|request| analyze_project_compact_json(request).expect("test assumption must hold")) + .collect(); + let concurrent: Vec = std::thread::scope(|scope| { + let handles: Vec<_> = requests + .iter() + .map(|request| scope.spawn(move || analyze_project_compact_json(request))) + .collect(); + handles + .into_iter() + .map(|handle| { + handle + .join() + .expect("test assumption must hold") + .expect("test assumption must hold") + }) + .collect() + }); + concurrent == serial +} diff --git a/crates/sqlbuild-rules-native/src/query_analysis/tests/test_query_analysis.rs b/crates/sqlbuild-rules-native/src/query_analysis/tests/test_query_analysis.rs index 051e702212..528968835e 100644 --- a/crates/sqlbuild-rules-native/src/query_analysis/tests/test_query_analysis.rs +++ b/crates/sqlbuild-rules-native/src/query_analysis/tests/test_query_analysis.rs @@ -5,10 +5,10 @@ use crate::query_analysis::tests::helpers::{ borrowed_facts_preserve_named_outputs_and_terminal_sources, borrowed_facts_preserve_union_dependencies_and_join_nullability, bound_lineage_survives_type_recovery, canonical_queries_reuse_semantics_and_project_resources, - combined_queries_preserve_standalone_binding, + catalog_views_reuse_one_analysis_pool, combined_queries_preserve_standalone_binding, compact_project_query_interns_repeated_lineage_strings, - interleaved_query_templates_preserve_template_order, larger_later_queries_keep_request_order, - native_compatibility_types_preserve_result_semantics, + concurrent_batches_match_serial_batches, interleaved_query_templates_preserve_template_order, + larger_later_queries_keep_request_order, native_compatibility_types_preserve_result_semantics, repeated_project_facts_intern_complete_facts, unannotated_types_require_complete_nested_evidence, widening_aggregates_require_compatibility_recovery, @@ -86,6 +86,16 @@ fn given_compact_query_cases_when_analyzing_projects_then_expected_behavior_hold run: larger_later_queries_keep_request_order, expected_success: true, }, + CompactQueryAnalysisTestCase { + description: "catalog views reuse one analysis pool", + run: catalog_views_reuse_one_analysis_pool, + expected_success: true, + }, + CompactQueryAnalysisTestCase { + description: "concurrent batches match serial batches", + run: concurrent_batches_match_serial_batches, + expected_success: true, + }, ]; for test_case in test_cases { 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 3ece2a48b9..37ff225df0 100644 --- a/crates/sqlbuild-rules-native/src/semantic_validation/_helpers/catalog.rs +++ b/crates/sqlbuild-rules-native/src/semantic_validation/_helpers/catalog.rs @@ -3,8 +3,11 @@ use crate::bindings::main::compiler_error::compiler_error; use crate::bindings::types::CompilerDetach; use std::collections::HashMap; +use std::sync::{Arc, Mutex, PoisonError}; -use crate::semantic_validation::models::{CatalogInput, Columns, FunctionProbes, ProjectCatalog}; +use crate::semantic_validation::models::{ + CatalogInput, Columns, CompactAnalysisJob, FunctionProbes, ProjectCatalog, +}; use crate::semantic_validation::types::{BindingRequest, DiagnosticRow, Relations}; use crate::semantic_validation::{ _helpers::{diagnostics, identifiers}, @@ -21,6 +24,9 @@ use pyo3::pymethods; use pyo3::types::{PyBytes, PyDict, PyDictMethods}; use rayon::iter::{IntoParallelIterator, ParallelIterator}; +const ANALYSIS_WORKERS: usize = 4; +const ANALYSIS_WORKER_STACK_BYTES: usize = 16 * 1024 * 1024; + impl<'py> FromPyObject<'py> for Columns { fn extract_bound(value: &Bound<'py, PyAny>) -> PyResult { value @@ -129,7 +135,8 @@ impl ProjectCatalog { tables: HashMap::new(), overrides: Vec::new(), analysis_tables: HashMap::new(), - function_probes: FunctionProbes::default(), + function_probes: Arc::new(FunctionProbes::default()), + analysis_pool: Arc::default(), }; catalog.update_relations(relations); Ok(catalog) @@ -149,7 +156,8 @@ impl ProjectCatalog { tables: self.tables.clone(), overrides: Vec::new(), analysis_tables: self.analysis_tables.clone(), - function_probes: FunctionProbes::default(), + function_probes: Arc::new(FunctionProbes::default()), + analysis_pool: Arc::clone(&self.analysis_pool), }; catalog.update_relations(relations); catalog @@ -202,16 +210,21 @@ impl ProjectCatalog { serde_json::to_string(&payloads).map_err(|error| PyValueError::new_err(error.to_string())) } - fn analyze_compact<'py>( - &self, - py: Python<'py>, - payload: &[u8], - ) -> PyResult> { + /// Resolve a compact batch now; the job runs later while the catalog keeps changing. + fn prepare_compact(&self, py: Python<'_>, payload: &[u8]) -> PyResult { let text = std::str::from_utf8(payload) .map_err(|error| PyValueError::new_err(error.to_string()))?; - let result = py.compiler_detach(|| crate::query_analysis::main::analyze_project_catalog::analyze_project_compact_with_catalog(text, self)) + let analysis = py + .compiler_detach(|| { + crate::query_analysis::main::prepare_project_catalog::prepare_project_compact_with_catalog( + text, self, + ) + }) .map_err(compiler_error)?; - Ok(PyBytes::new(py, result.as_bytes())) + Ok(CompactAnalysisJob { + analysis: Mutex::new(Some(analysis)), + catalog: self.analysis_view(), + }) } fn binding_results( @@ -220,11 +233,7 @@ impl ProjectCatalog { requests: Vec, ) -> PyResult>> { py.compiler_detach(|| { - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(4.min(requests.len().max(1))) - .stack_size(16 * 1024 * 1024) - .build() - .map_err(|error| error.to_string())?; + let pool = self.analysis_pool()?; pool.install(|| { requests .into_par_iter() @@ -246,7 +255,58 @@ impl ProjectCatalog { } } +#[pymethods] +impl CompactAnalysisJob { + fn run<'py>(&self, py: Python<'py>) -> PyResult> { + let analysis = self + .analysis + .lock() + .unwrap_or_else(PoisonError::into_inner) + .take() + .ok_or_else(|| PyValueError::new_err("compact analysis job already ran"))?; + let result = py + .compiler_detach(|| analysis(&self.catalog)) + .map_err(compiler_error)?; + Ok(PyBytes::new(py, result.as_bytes())) + } +} + impl ProjectCatalog { + /// Return the options-only view a resolved analysis batch validates with. + pub(crate) fn analysis_view(&self) -> Self { + Self { + dialect: self.dialect, + options: self.options.clone(), + quoted_ignore_case: self.quoted_ignore_case, + tables: HashMap::new(), + overrides: Vec::new(), + analysis_tables: HashMap::new(), + function_probes: Arc::clone(&self.function_probes), + analysis_pool: Arc::clone(&self.analysis_pool), + } + } + + /// Return this compile's analysis pool, building it on first use. + pub(crate) fn analysis_pool(&self) -> Result, String> { + let mut pool = self + .analysis_pool + .lock() + .unwrap_or_else(PoisonError::into_inner); + if let Some(pool) = pool.as_ref() { + return Ok(Arc::clone(pool)); + } + let built = Arc::new( + rayon::ThreadPoolBuilder::new() + .num_threads(ANALYSIS_WORKERS) + .stack_size(ANALYSIS_WORKER_STACK_BYTES) + .thread_name(|index| format!("sqlbuild-analysis-{index}")) + .build() + .map_err(|error| error.to_string())?, + ); + *pool = Some(Arc::clone(&built)); + Ok(built) + } + pub(crate) fn with_options( dialect: DialectType, options: SchemaValidationOptions, @@ -259,7 +319,8 @@ impl ProjectCatalog { tables: HashMap::new(), overrides: Vec::new(), analysis_tables: HashMap::new(), - function_probes: FunctionProbes::default(), + function_probes: Arc::new(FunctionProbes::default()), + analysis_pool: Arc::default(), } } pub(crate) fn analysis_schema( diff --git a/crates/sqlbuild-rules-native/src/semantic_validation/models.rs b/crates/sqlbuild-rules-native/src/semantic_validation/models.rs index e185093a05..46afc66c34 100644 --- a/crates/sqlbuild-rules-native/src/semantic_validation/models.rs +++ b/crates/sqlbuild-rules-native/src/semantic_validation/models.rs @@ -1,11 +1,12 @@ //! Compile-owned native catalog, mapping state and Python-boundary requests. -use crate::semantic_validation::types::{Expansion, ProbeKey, Relations}; +use crate::semantic_validation::types::{Expansion, PreparedCompactAnalysis, ProbeKey, Relations}; use polyglot_sql::validation::SchemaTable; use polyglot_sql::{DialectType, SchemaValidationOptions}; use pyo3::{FromPyObject, pyclass}; +use rayon::ThreadPool; use std::collections::HashMap; -use std::sync::Mutex; +use std::sync::{Arc, Mutex}; #[derive(Default, Debug)] pub(crate) struct Columns(pub(crate) Vec<(String, Option)>); @@ -53,7 +54,16 @@ pub(crate) struct ProjectCatalog { pub(super) tables: HashMap, pub(super) overrides: Vec>, pub(super) analysis_tables: HashMap, - pub(crate) function_probes: FunctionProbes, + 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>>>, +} + +/// One resolved compact analysis batch that runs without borrowing its project catalog. +#[pyclass(module = "sqlbuild._native", frozen)] +pub(crate) struct CompactAnalysisJob { + pub(super) analysis: Mutex>, + pub(super) catalog: ProjectCatalog, } #[pyclass(module = "sqlbuild._native")] diff --git a/crates/sqlbuild-rules-native/src/semantic_validation/types.rs b/crates/sqlbuild-rules-native/src/semantic_validation/types.rs index 4a105c3cd7..1e171ebe37 100644 --- a/crates/sqlbuild-rules-native/src/semantic_validation/types.rs +++ b/crates/sqlbuild-rules-native/src/semantic_validation/types.rs @@ -1,4 +1,4 @@ -use crate::semantic_validation::models::Columns; +use crate::semantic_validation::models::{Columns, ProjectCatalog}; use std::collections::HashMap; pub(crate) type Relations = HashMap; @@ -14,3 +14,6 @@ pub(crate) type DiagnosticRow = ( ); pub(crate) type Expansion = (usize, usize, usize, usize); pub(crate) type ProbeKey = (polyglot_sql::DialectType, String, usize); +/// Compact analysis whose catalog schemas are resolved; it runs with an options-only catalog. +pub(crate) type PreparedCompactAnalysis = + Box Result + Send>; diff --git a/scripts/cold_compile_performance/_helpers/random_dag_project.py b/scripts/cold_compile_performance/_helpers/random_dag_project.py new file mode 100644 index 0000000000..7ecb629896 --- /dev/null +++ b/scripts/cold_compile_performance/_helpers/random_dag_project.py @@ -0,0 +1,153 @@ +"""Seeded DuckDB projects whose models bind against shapes inferred from their producers.""" + +from __future__ import annotations + +import random +from collections.abc import Callable +from pathlib import Path + +from scripts.cold_compile_performance.models import RandomDagProject + +_RANDOM_DAG_ROOTS: int = 3 +_RANDOM_DAG_ROOT_SHARE: float = 0.1 +_RANDOM_DAG_FEATURE_SHARE: float = 0.1 +_RANDOM_DAG_WINDOW: int = 12 +_RANDOM_DAG_SOURCES: tuple[tuple[str, tuple[str, ...]], ...] = ( + ("raw_orders", ("id", "quantity", "status")), + ("raw_customers", ("id", "region")), +) + + +def write_random_dag_project(*, project_dir: Path, project: RandomDagProject) -> tuple[str, ...]: + """Write seeded models over inferred shapes; return model names in topological order.""" + + rng: random.Random = random.Random(project.seed) + project_dir.mkdir(parents=True, exist_ok=True) + rules: str = ", ".join(f'"{code}"' for code in project.rules) + (project_dir / "sqlbuild_project.toml").write_text( + f'name = "orders"\nadapter = "duckdb"\n[rules]\nselect = [{rules}]\n', encoding="utf-8" + ) + (project_dir / "sources").mkdir(exist_ok=True) + (project_dir / "sources" / "raw.yml").write_text( + "sources:\n" + " - name: raw_orders\n description: Test source raw_orders.\n" + " expression: SELECT 1 AS id, 2 AS quantity, 'open' AS status\n" + " - name: raw_customers\n description: Test source raw_customers.\n" + " expression: SELECT 1 AS id, 'north' AS region\n", + encoding="utf-8", + ) + models: Path = project_dir / "models" + models.mkdir(exist_ok=True) + columns: dict[str, tuple[str, ...]] = {} + names: list[str] = [] + for index in range(project.model_count): + name: str = f"orders_{index:03d}" + options: str = "" + body: str + output: tuple[str, ...] + if index < _RANDOM_DAG_ROOTS or rng.random() < _RANDOM_DAG_ROOT_SHARE: + source, output = rng.choice(_RANDOM_DAG_SOURCES) + body = f'SELECT * FROM __source("{source}")' + else: + body, output = _random_dag_query( + rng=rng, index=index, names=names, columns=columns, errors=project.errors + ) + if project.run_ids and rng.random() < _RANDOM_DAG_FEATURE_SHARE: + options += ', tags ["${CTX:run.id}"]' + if project.analysis_opt_outs and rng.random() < _RANDOM_DAG_FEATURE_SHARE: + options += ", sql_analysis false" + (models / f"{name}.sql").write_text( + f'MODEL (description "Test model {name}.", materialized view{options});\n{body}\n', + encoding="utf-8", + ) + columns[name] = output + names.append(name) + if project.missing_reference: + (models / "orders_missing.sql").write_text( + 'MODEL (description "Test model orders_missing.", materialized view);\n' + 'SELECT id FROM __ref("orders_absent")\n', + encoding="utf-8", + ) + if project.cycle: + for name, upstream in ( + ("orders_loop_a", "orders_loop_b"), + ("orders_loop_b", "orders_loop_a"), + ): + (models / f"{name}.sql").write_text( + f'MODEL (description "Test model {name}.", materialized view);\n' + f'SELECT id FROM __ref("{upstream}")\n', + encoding="utf-8", + ) + return tuple(names) + + +def _random_dag_query( + *, + rng: random.Random, + index: int, + names: list[str], + columns: dict[str, tuple[str, ...]], + errors: bool, +) -> tuple[str, tuple[str, ...]]: + upstream: str = rng.choice(names[max(0, index - _RANDOM_DAG_WINDOW) :]) + other: str = rng.choice(names) + if errors and rng.random() < _RANDOM_DAG_FEATURE_SHARE: + return ( + f'SELECT id, missing_{index} FROM __ref("{upstream}")', + ("id", f"missing_{index}"), + ) + return rng.choice(_query_builders())( + index=index, upstream=upstream, other=other, columns=columns + ) + + +def _star_query( + *, index: int, upstream: str, other: str, columns: dict[str, tuple[str, ...]] +) -> tuple[str, tuple[str, ...]]: + del index, other + return f'SELECT * FROM __ref("{upstream}")', columns[upstream] + + +def _projection_query( + *, index: int, upstream: str, other: str, columns: dict[str, tuple[str, ...]] +) -> tuple[str, tuple[str, ...]]: + del other + return ( + f"SELECT id, {columns[upstream][-1]} AS kept_{index}, " + f'CAST(id AS BIGINT) * {index} AS value_{index} FROM __ref("{upstream}")', + ("id", f"kept_{index}", f"value_{index}"), + ) + + +def _join_query( + *, index: int, upstream: str, other: str, columns: dict[str, tuple[str, ...]] +) -> tuple[str, tuple[str, ...]]: + return ( + f"SELECT a.*, b.{columns[other][-1]} AS joined_{index} " + f'FROM __ref("{upstream}") AS a LEFT JOIN __ref("{other}") AS b ON a.id = b.id', + (*columns[upstream], f"joined_{index}"), + ) + + +def _union_query( + *, index: int, upstream: str, other: str, columns: dict[str, tuple[str, ...]] +) -> tuple[str, tuple[str, ...]]: + del index, columns + return ( + f'SELECT id FROM __ref("{upstream}") UNION ALL SELECT id FROM __ref("{other}")', + ("id",), + ) + + +def _aggregate_query( + *, index: int, upstream: str, other: str, columns: dict[str, tuple[str, ...]] +) -> tuple[str, tuple[str, ...]]: + del other, columns + return ( + f'SELECT id, COUNT(*) AS row_count_{index} FROM __ref("{upstream}") GROUP BY id', + ("id", f"row_count_{index}"), + ) + + +def _query_builders() -> tuple[Callable[..., tuple[str, tuple[str, ...]]], ...]: + return (_star_query, _projection_query, _join_query, _union_query, _aggregate_query) diff --git a/scripts/cold_compile_performance/models.py b/scripts/cold_compile_performance/models.py index 9118eaf908..566a6ca7d8 100644 --- a/scripts/cold_compile_performance/models.py +++ b/scripts/cold_compile_performance/models.py @@ -13,3 +13,17 @@ class VariedProjectStatistics: leaf_models: int array_models: int window_models: int + + +@dataclass(frozen=True) +class RandomDagProject: + """One seeded DuckDB project whose models bind against inferred upstream shapes.""" + + seed: int + model_count: int + errors: bool = False + run_ids: bool = False + analysis_opt_outs: bool = False + missing_reference: bool = False + cycle: bool = False + rules: tuple[str, ...] = () diff --git a/src/sqlbuild/compiler/compile/_helpers/analysis/compact.py b/src/sqlbuild/compiler/compile/_helpers/analysis/compact.py index cdd4dc7b7c..145dd4de11 100644 --- a/src/sqlbuild/compiler/compile/_helpers/analysis/compact.py +++ b/src/sqlbuild/compiler/compile/_helpers/analysis/compact.py @@ -3,7 +3,10 @@ from __future__ import annotations import logging -from collections.abc import Sequence +import threading +from collections.abc import Iterator, Sequence +from contextlib import contextmanager +from contextvars import ContextVar, Token from dataclasses import replace from typing import Any, cast @@ -176,14 +179,47 @@ SqlBindingResult, SqlSchemaValidationRequest, ) -from sqlbuild.compiler.sql_analysis.types import NativeQueryAnalysisModule +from sqlbuild.compiler.sql_analysis.types import ( + NativeCompactAnalysisJob, + NativeProjectCatalog, + NativeQueryAnalysisModule, +) from sqlbuild.diagnostics.main.log_debug_event import log_debug_event _DEBUG_LOGGER: logging.Logger = logging.getLogger("sqlbuild.compile") +_ANALYSIS_LOCK: ContextVar[threading.Condition | None] = ContextVar( + "sqlbuild_model_analysis_lock", default=None +) _COMPACT_ANALYSIS_WORKERS: int = 4 _NATIVE_LEGACY_FALLBACK: str = "native project type recovery requires legacy fallback" +@contextmanager +def serialized_analysis(lock: threading.Condition) -> Iterator[None]: + """Mark the held `lock` as the guard of every Python analysis step in this context.""" + + token: Token[threading.Condition | None] = _ANALYSIS_LOCK.set(lock) + try: + yield + finally: + _ANALYSIS_LOCK.reset(token) + + +@contextmanager +def native_section() -> Iterator[None]: + """Release the analysis lock around native work that reads no shared analysis state.""" + + lock: threading.Condition | None = _ANALYSIS_LOCK.get() + if lock is None: + yield + return + lock.release() + try: + yield + finally: + _ = lock.acquire() + + def infer_columns_with_sql_analysis( *, query_sql: str, @@ -889,31 +925,35 @@ def _run_compact_analysis_batch(*, preparation: CompactBatchPreparation) -> obje workers: int = _COMPACT_ANALYSIS_WORKERS if preparation.binding_catalog is not None: - return orjson.loads( - preparation.binding_catalog.native.analyze_compact( - orjson.dumps( - { - "queries": preparation.queries, - "templates": preparation.templates, - "projections": preparation.projections, - "workers": workers, - } - ) - ) - ) - return orjson.loads( - cast(NativeQueryAnalysisModule, _native).analyze_project_queries_compact_json( + job: NativeCompactAnalysisJob = cast( + NativeProjectCatalog, preparation.binding_catalog.native + ).prepare_compact( orjson.dumps( { "queries": preparation.queries, "templates": preparation.templates, "projections": preparation.projections, "workers": workers, - }, - option=orjson.OPT_SORT_KEYS, - ).decode() + } + ) + ) + with native_section(): + return orjson.loads(job.run()) + request_json: str = orjson.dumps( + { + "queries": preparation.queries, + "templates": preparation.templates, + "projections": preparation.projections, + "workers": workers, + }, + option=orjson.OPT_SORT_KEYS, + ).decode() + with native_section(): + return orjson.loads( + cast(NativeQueryAnalysisModule, _native).analyze_project_queries_compact_json( + request_json + ) ) - ) def _attach_compiled_bindings( diff --git a/src/sqlbuild/compiler/compile/_helpers/assembly/binding_waves.py b/src/sqlbuild/compiler/compile/_helpers/assembly/binding_waves.py index 8922440b40..b084a8d591 100644 --- a/src/sqlbuild/compiler/compile/_helpers/assembly/binding_waves.py +++ b/src/sqlbuild/compiler/compile/_helpers/assembly/binding_waves.py @@ -3,15 +3,10 @@ from __future__ import annotations from collections.abc import Callable -from dataclasses import replace -from graphlib import CycleError, TopologicalSorter +from functools import partial from sqlbuild.adapter.contract.models import ExpressionInferenceProfile -from sqlbuild.compiler.compile._helpers.analysis.cache import model_analysis_output_signature -from sqlbuild.compiler.compile._helpers.assembly.semantic_shapes import ( - binding_relation_names, - published_model_shape, -) +from sqlbuild.compiler.compile.classes.binding_dataflow import BindingDataflow from sqlbuild.compiler.compile.models import CompileModelInput, PolyglotAnalysisResult from sqlbuild.compiler.compile.models import ( ModelSqlAnalysis as _ModelSqlAnalysis, @@ -22,6 +17,10 @@ from sqlbuild.compiler.lineage.types import InferredNullability from sqlbuild.compiler.references.types import SqlReferenceKind +_DATAFLOW_WORKERS: int = 4 +_DATAFLOW_BATCH_MIN: int = 8 +_DATAFLOW_BATCH_LIMIT: int = 64 + def analyze_binding_waves( *, @@ -36,97 +35,58 @@ def analyze_binding_waves( analyze: Callable[..., tuple[_ModelSqlAnalysis, ...]], complete: Callable[..., tuple[_ModelSqlAnalysis, ...]], ) -> tuple[tuple[_ModelSqlAnalysis, ...], dict[str, PolyglotAnalysisResult]]: - by_name: dict[str, _ModelSqlAnalysisRequest] = dict(zip(names, requests, strict=True)) - available_names: frozenset[str] = frozenset(by_name) - dependencies: dict[str, tuple[str, ...]] = {} - for name, request in by_name.items(): - dependencies[name] = referenced_model_names( - model_input=request.model_input, available_names=available_names - ) - sorter: TopologicalSorter[str] = TopologicalSorter(dependencies) - try: - sorter.prepare() - except CycleError: - return complete( - requests=requests, - analyses=analyze( - requests=requests, - cached_analyses=cached, - column_types_by_table=types, - column_nullability_by_table=nullability, - ), - complete_binding_schemas=shapes, - ), cached - complete_shapes: dict[str, dict[str, str]] = dict(shapes) - available_types: dict[str, dict[str, str]] = dict(types) - available_nullability: dict[str, dict[str, InferredNullability]] = dict(nullability) - reusable: dict[str, PolyglotAnalysisResult] = dict(cached) - changed: set[str] = set() - results: dict[str, _ModelSqlAnalysis] = {} - while sorter.is_active(): - ready: tuple[str, ...] = tuple(sorter.get_ready()) - wave: list[_ModelSqlAnalysisRequest] = [] - for name in ready: - request: _ModelSqlAnalysisRequest = by_name[name] - if any(parent in changed for parent in dependencies[name]): - changed.add(name) - if request.cache_key is not None: - reusable.pop(request.cache_key, None) - bindings: dict[str, dict[str, str]] | None = ( - None - if request.binding_schema is None - else { - relation: complete_shapes.get(relation, {}) - for relation in binding_relation_names(request.model_input.references) - } - ) - wave.append(replace(request, binding_schema=bindings)) - wave_requests: tuple[_ModelSqlAnalysisRequest, ...] = tuple(wave) - wave_names: set[str] = set(ready) - for request in wave_requests: - wave_names.update(binding_relation_names(request.model_input.references)) - analyses: tuple[_ModelSqlAnalysis, ...] = complete( - requests=wave_requests, - analyses=analyze( - requests=wave_requests, - cached_analyses=reusable, - column_types_by_table=available_types, - column_nullability_by_table=available_nullability, - ), - complete_binding_schemas={ - name: complete_shapes[name] for name in wave_names if name in complete_shapes - }, - supplied_relations=frozenset( - name - for name in wave_names - if name in complete_shapes - and available_types.get(name, {}).keys() == complete_shapes[name].keys() - ), + """Analyze one topological level at a time, the reference for dataflow scheduling.""" + + return BindingDataflow( + requests=requests, + names=names, + cached=cached, + previous_signatures=previous_signatures, + shapes=shapes, + types=types, + nullability=nullability, + profile=profile, + analyze=analyze, + complete=complete, + ).analyze_waves() + + +def analyze_binding_dataflow( + *, + requests: tuple[_ModelSqlAnalysisRequest, ...], + names: tuple[str, ...], + cached: dict[str, PolyglotAnalysisResult], + previous_signatures: dict[str, str], + shapes: dict[str, dict[str, str]], + types: dict[str, dict[str, str]], + nullability: dict[str, dict[str, InferredNullability]], + profile: ExpressionInferenceProfile, + analyze: Callable[..., tuple[_ModelSqlAnalysis, ...]], + complete: Callable[..., tuple[_ModelSqlAnalysis, ...]], +) -> tuple[tuple[_ModelSqlAnalysis, ...], dict[str, PolyglotAnalysisResult]]: + """Analyze each model once its producers are final; replay waves after any error.""" + + create: Callable[[], BindingDataflow] = partial( + BindingDataflow, + requests=requests, + names=names, + cached=cached, + previous_signatures=previous_signatures, + shapes=shapes, + types=types, + nullability=nullability, + profile=profile, + analyze=analyze, + complete=complete, + ) + result: tuple[tuple[_ModelSqlAnalysis, ...], dict[str, PolyglotAnalysisResult]] | None = ( + create().analyze_dataflow( + workers=_DATAFLOW_WORKERS, + batch_min=_DATAFLOW_BATCH_MIN, + batch_limit=_DATAFLOW_BATCH_LIMIT, ) - for name, request, result in zip(ready, wave_requests, analyses, strict=True): - results[name] = result - analysis: PolyglotAnalysisResult = result.polyglot_analysis - if model_analysis_output_signature(analysis) != previous_signatures.get(name): - changed.add(name) - required: frozenset[str] = binding_relation_names(request.model_input.references) - star_known: bool = not analysis.has_star or ( - analysis.star_resolved and required <= complete_shapes.keys() - ) - if analysis.columns and star_known: - shape: dict[str, str] = published_model_shape( - sql=result.cleaned_sql or request.query_sql, - columns={column.name: column.type or "UNKNOWN" for column in analysis.columns}, - inputs={table: complete_shapes.get(table, {}) for table in required}, - profile=profile, - config_values=request.model_input.config.values, - ) - complete_shapes.setdefault(name, shape) - available_types.setdefault(name, shape) - available_nullability.setdefault( - name, dict.fromkeys(shape, InferredNullability.UNKNOWN) - ) - sorter.done(*ready) - return tuple(results[name] for name in names), reusable + ) + return result if result is not None else create().analyze_waves() def downstream_model_names( @@ -147,18 +107,3 @@ def downstream_model_names( downstream_names.add(downstream_name) pending.append(downstream_name) return downstream_names - - -def referenced_model_names( - *, model_input: CompileModelInput, available_names: frozenset[str] | None = None -) -> tuple[str, ...]: - return tuple( - sorted( - { - reference.ref_name - for reference in model_input.references - if reference.ref_kind == SqlReferenceKind.REF - and (available_names is None or reference.ref_name in available_names) - } - ) - ) diff --git a/src/sqlbuild/compiler/compile/_helpers/assembly/project.py b/src/sqlbuild/compiler/compile/_helpers/assembly/project.py index fefad914c0..0556b04bd9 100644 --- a/src/sqlbuild/compiler/compile/_helpers/assembly/project.py +++ b/src/sqlbuild/compiler/compile/_helpers/assembly/project.py @@ -45,13 +45,10 @@ get_authored_binding_location, query_line_offset, ) -from sqlbuild.compiler.compile._helpers.assembly.binding_waves import analyze_binding_waves +from sqlbuild.compiler.compile._helpers.assembly.binding_waves import analyze_binding_dataflow from sqlbuild.compiler.compile._helpers.assembly.binding_waves import ( downstream_model_names as _downstream_model_names, ) -from sqlbuild.compiler.compile._helpers.assembly.binding_waves import ( - referenced_model_names as _referenced_model_names, -) from sqlbuild.compiler.compile._helpers.assembly.native_declarations import ( known_declared_types, known_function_names, @@ -66,6 +63,9 @@ from sqlbuild.compiler.compile._helpers.assembly.semantic_shapes import ( build_declared_column_types as _build_column_types_by_table, ) +from sqlbuild.compiler.compile._helpers.assembly.semantic_shapes import ( + referenced_model_names as _referenced_model_names, +) from sqlbuild.compiler.compile._helpers.assembly.targets import ( build_model_relation_target, build_seed_relation_target, @@ -795,7 +795,7 @@ def _analyze_model_sql_in_parallel( ) analyses: tuple[_ModelSqlAnalysis, ...] if dependency_ordered: - analyses, cached_analyses = analyze_binding_waves( + analyses, cached_analyses = analyze_binding_dataflow( requests=requests, names=tuple(_model_name(request.model_input) for request in requests), cached=cached_analyses, diff --git a/src/sqlbuild/compiler/compile/_helpers/assembly/semantic_shapes.py b/src/sqlbuild/compiler/compile/_helpers/assembly/semantic_shapes.py index c0216b843e..1fb6963cb3 100644 --- a/src/sqlbuild/compiler/compile/_helpers/assembly/semantic_shapes.py +++ b/src/sqlbuild/compiler/compile/_helpers/assembly/semantic_shapes.py @@ -309,3 +309,18 @@ def _expression_source_shape( columns={column.name: column.type or "UNKNOWN" for column in analysis.columns}, inputs={}, ) + + +def referenced_model_names( + *, model_input: CompileModelInput, available_names: frozenset[str] | None = None +) -> tuple[str, ...]: + return tuple( + sorted( + { + reference.ref_name + for reference in model_input.references + if reference.ref_kind == SqlReferenceKind.REF + and (available_names is None or reference.ref_name in available_names) + } + ) + ) diff --git a/src/sqlbuild/compiler/compile/classes/binding_dataflow.py b/src/sqlbuild/compiler/compile/classes/binding_dataflow.py new file mode 100644 index 0000000000..0909ba58c7 --- /dev/null +++ b/src/sqlbuild/compiler/compile/classes/binding_dataflow.py @@ -0,0 +1,314 @@ +"""Model analysis state that publishes producer shapes as each analysis completes.""" + +from __future__ import annotations + +import contextvars +import heapq +import sys +import threading +import time +from collections.abc import Callable +from dataclasses import replace +from graphlib import CycleError, TopologicalSorter + +from sqlbuild.adapter.contract.models import ExpressionInferenceProfile +from sqlbuild.compiler.compile._helpers.analysis.cache import model_analysis_output_signature +from sqlbuild.compiler.compile._helpers.analysis.compact import serialized_analysis +from sqlbuild.compiler.compile._helpers.assembly.semantic_shapes import ( + binding_relation_names, + published_model_shape, + referenced_model_names, +) +from sqlbuild.compiler.compile.models import ( + ModelSqlAnalysis, + ModelSqlAnalysisRequest, + PolyglotAnalysisResult, +) +from sqlbuild.compiler.lineage.types import InferredNullability +from sqlbuild.presentation.main.transient_line_coordinator import ( + shared_transient_line_coordinator, +) + +_WORKER_POLL_SECONDS: float = 0.05 +_WORKER_START_GRACE_SECONDS: float = 1.0 +_INTERRUPT_NOTICE: str = "Interrupted; finishing in-flight model analysis...\n" + + +class BindingDataflow: + """Analyze each model with exactly the producer shapes and cache state it depends on.""" + + def __init__( + self, + *, + requests: tuple[ModelSqlAnalysisRequest, ...], + names: tuple[str, ...], + cached: dict[str, PolyglotAnalysisResult], + previous_signatures: dict[str, str], + shapes: dict[str, dict[str, str]], + types: dict[str, dict[str, str]], + nullability: dict[str, dict[str, InferredNullability]], + profile: ExpressionInferenceProfile, + analyze: Callable[..., tuple[ModelSqlAnalysis, ...]], + complete: Callable[..., tuple[ModelSqlAnalysis, ...]], + ) -> None: + self._requests: tuple[ModelSqlAnalysisRequest, ...] = requests + self._names: tuple[str, ...] = names + self._cached: dict[str, PolyglotAnalysisResult] = cached + self._by_name: dict[str, ModelSqlAnalysisRequest] = dict(zip(names, requests, strict=True)) + available_names: frozenset[str] = frozenset(names) + self._dependencies: dict[str, tuple[str, ...]] = { + name: referenced_model_names( + model_input=request.model_input, available_names=available_names + ) + for name, request in self._by_name.items() + } + self._previous_signatures: dict[str, str] = previous_signatures + self._complete_shapes: dict[str, dict[str, str]] = dict(shapes) + self._available_types: dict[str, dict[str, str]] = dict(types) + self._available_nullability: dict[str, dict[str, InferredNullability]] = dict(nullability) + self._reusable: dict[str, PolyglotAnalysisResult] = dict(cached) + self._profile: ExpressionInferenceProfile = profile + self._analyze: Callable[..., tuple[ModelSqlAnalysis, ...]] = analyze + self._complete: Callable[..., tuple[ModelSqlAnalysis, ...]] = complete + self._changed: set[str] = set() + self._results: dict[str, ModelSqlAnalysis] = {} + self._ready: list[tuple[int, int, str]] = [] + self._pending: dict[str, int] = {} + self._dependents: dict[str, list[str]] = {} + self._failure: BaseException | None = None + self._cancelled: bool = False + self._condition: threading.Condition = threading.Condition(threading.Lock()) + + def analyze_waves( + self, + ) -> tuple[tuple[ModelSqlAnalysis, ...], dict[str, PolyglotAnalysisResult]]: + """Analyze one topological level at a time, the reference schedule.""" + + sorter: TopologicalSorter[str] = TopologicalSorter(self._dependencies) + try: + sorter.prepare() + except CycleError: + return self._analyze_cycle(), self._cached + while sorter.is_active(): + ready: tuple[str, ...] = tuple(sorter.get_ready()) + self._analyze_ready(ready) + sorter.done(*ready) + return tuple(self._results[name] for name in self._names), self._reusable + + def analyze_dataflow( + self, *, workers: int, batch_min: int, batch_limit: int + ) -> tuple[tuple[ModelSqlAnalysis, ...], dict[str, PolyglotAnalysisResult]] | None: + """Analyze models as their producers finish; return None after an analysis error.""" + + try: + order: tuple[str, ...] = tuple(TopologicalSorter(self._dependencies).static_order()) + except CycleError: + return self._analyze_cycle(), self._cached + heights: dict[str, int] = {} + self._dependents = {name: [] for name in self._by_name} + for name in self._by_name: + for parent in self._dependencies[name]: + self._dependents[parent].append(name) + for name in reversed(order): + heights[name] = 1 + max((heights[child] for child in self._dependents[name]), default=0) + priorities: dict[str, tuple[int, int]] = { + name: (-heights[name], position) for position, name in enumerate(self._by_name) + } + self._pending = {name: len(self._dependencies[name]) for name in self._by_name} + self._ready = [ + (*priorities[name], name) for name in self._by_name if not self._pending[name] + ] + heapq.heapify(self._ready) + share: tuple[int, int, int] = (workers, batch_min, batch_limit) + with serialized_analysis(self._condition): + context: contextvars.Context = contextvars.copy_context() + threads: list[threading.Thread] = [ + threading.Thread( + target=context.copy().run, + args=(self._run_worker,), + kwargs={"priorities": priorities, "share": share}, + name=f"sqlbuild-model-analysis-{index}", + daemon=True, + ) + for index in range(max(1, workers)) + ] + self._wait_for_workers(threads) + if self._failure is not None and not isinstance(self._failure, Exception): + raise self._failure + if self._failure is not None: + return None + return tuple(self._results[name] for name in self._names), self._reusable + + def _wait_for_workers(self, threads: list[threading.Thread]) -> None: + """Join workers lock-free; a first interrupt cancels them, repeats wait for their exit.""" + + attempted: list[threading.Thread] = [] + interrupt: BaseException | None = None + try: + for thread in threads: + attempted.append(thread) + thread.start() + except BaseException as error: # noqa: BLE001 - re-raised after workers stop + interrupt = self._cancel(interrupt=None, error=error) + pending: list[threading.Thread] = list(attempted) + unstarted_deadline: float = time.monotonic() + _WORKER_START_GRACE_SECONDS + while pending: + try: + pending = [ + thread + for thread in pending + if self._worker_running(thread=thread, unstarted_deadline=unstarted_deadline) + ] + except BaseException as error: # noqa: BLE001 - re-raised after workers stop + interrupt = self._cancel(interrupt=interrupt, error=error) + if interrupt is not None: + raise interrupt + + def _cancel(self, *, interrupt: BaseException | None, error: BaseException) -> BaseException: + if interrupt is not None: + return interrupt + self._cancelled = True + if not isinstance(error, Exception): + shared_transient_line_coordinator().write_persistent( + stream=sys.stderr, text=_INTERRUPT_NOTICE + ) + return error + + @staticmethod + def _worker_running(*, thread: threading.Thread, unstarted_deadline: float) -> bool: + """Join briefly; a thread whose OS thread never started is dropped after a grace period.""" + + if thread.ident is None: + time.sleep(_WORKER_POLL_SECONDS / 10) + return time.monotonic() < unstarted_deadline + try: + thread.join(_WORKER_POLL_SECONDS) + except RuntimeError: + return True + return thread.is_alive() + + def _run_worker( + self, *, priorities: dict[str, tuple[int, int]], share: tuple[int, int, int] + ) -> None: + with self._condition: + try: + self._drain_ready(priorities=priorities, share=share) + except BaseException as error: # noqa: BLE001 - re-raised or replayed in waves + self._record_failure(error) + self._condition.notify_all() + + def _record_failure(self, error: BaseException) -> None: + """Keep the first failure, except that a non-Exception interrupt always wins.""" + + if self._failure is None or ( + isinstance(self._failure, Exception) and not isinstance(error, Exception) + ): + self._failure = error + + def _drain_ready( + self, *, priorities: dict[str, tuple[int, int]], share: tuple[int, int, int] + ) -> None: + workers, batch_min, batch_limit = share + while ( + not self._cancelled + and self._failure is None + and len(self._results) < len(self._by_name) + ): + if not self._ready: + _ = self._condition.wait(_WORKER_POLL_SECONDS) + continue + count: int = min( + len(self._ready), + max(batch_min, -(-len(self._ready) // workers)), + batch_limit, + ) + batch: tuple[str, ...] = tuple(heapq.heappop(self._ready)[2] for _ in range(count)) + self._analyze_ready(batch) + for name in batch: + for child in self._dependents[name]: + self._pending[child] -= 1 + if not self._pending[child]: + heapq.heappush(self._ready, (*priorities[child], child)) + self._condition.notify_all() + + def _analyze_cycle(self) -> tuple[ModelSqlAnalysis, ...]: + return self._complete( + requests=self._requests, + analyses=self._analyze( + requests=self._requests, + cached_analyses=self._cached, + column_types_by_table=self._available_types, + column_nullability_by_table=self._available_nullability, + ), + complete_binding_schemas=self._complete_shapes, + ) + + def _analyze_ready(self, ready: tuple[str, ...]) -> None: + wave: list[ModelSqlAnalysisRequest] = [] + for name in ready: + request: ModelSqlAnalysisRequest = self._by_name[name] + if any(parent in self._changed for parent in self._dependencies[name]): + self._changed.add(name) + if request.cache_key is not None: + self._reusable.pop(request.cache_key, None) + bindings: dict[str, dict[str, str]] | None = ( + None + if request.binding_schema is None + else { + relation: self._complete_shapes.get(relation, {}) + for relation in binding_relation_names(request.model_input.references) + } + ) + wave.append(replace(request, binding_schema=bindings)) + wave_requests: tuple[ModelSqlAnalysisRequest, ...] = tuple(wave) + wave_names: set[str] = set(ready) + for request in wave_requests: + wave_names.update(binding_relation_names(request.model_input.references)) + analyses: tuple[ModelSqlAnalysis, ...] = self._complete( + requests=wave_requests, + analyses=self._analyze( + requests=wave_requests, + cached_analyses=self._reusable, + column_types_by_table=self._available_types, + column_nullability_by_table=self._available_nullability, + ), + complete_binding_schemas={ + name: self._complete_shapes[name] + for name in wave_names + if name in self._complete_shapes + }, + supplied_relations=frozenset( + name + for name in wave_names + if name in self._complete_shapes + and self._available_types.get(name, {}).keys() == self._complete_shapes[name].keys() + ), + ) + for name, request, result in zip(ready, wave_requests, analyses, strict=True): + self._results[name] = result + self._publish(name=name, request=request, result=result) + + def _publish( + self, *, name: str, request: ModelSqlAnalysisRequest, result: ModelSqlAnalysis + ) -> None: + analysis: PolyglotAnalysisResult = result.polyglot_analysis + if model_analysis_output_signature(analysis) != self._previous_signatures.get(name): + self._changed.add(name) + required: frozenset[str] = binding_relation_names(request.model_input.references) + star_known: bool = not analysis.has_star or ( + analysis.star_resolved and required <= self._complete_shapes.keys() + ) + if not analysis.columns or not star_known: + return + shape: dict[str, str] = published_model_shape( + sql=result.cleaned_sql or request.query_sql, + columns={column.name: column.type or "UNKNOWN" for column in analysis.columns}, + inputs={table: self._complete_shapes.get(table, {}) for table in required}, + profile=self._profile, + config_values=request.model_input.config.values, + ) + self._complete_shapes.setdefault(name, shape) + self._available_types.setdefault(name, shape) + self._available_nullability.setdefault( + name, dict.fromkeys(shape, InferredNullability.UNKNOWN) + ) diff --git a/src/sqlbuild/compiler/sql_analysis/types.py b/src/sqlbuild/compiler/sql_analysis/types.py index d30921b1f2..f490eb9c05 100644 --- a/src/sqlbuild/compiler/sql_analysis/types.py +++ b/src/sqlbuild/compiler/sql_analysis/types.py @@ -9,6 +9,10 @@ type NativeBindingRequest = tuple[str, list[tuple[str, bool]], dict[str, Mapping[str, str]]] +class NativeCompactAnalysisJob(Protocol): + def run(self) -> bytes: ... + + class NativeProjectCatalog(Protocol): def inferred_schema( self, *, sql: str, columns: Mapping[str, str], inputs: Mapping[str, Mapping[str, str]] @@ -21,7 +25,7 @@ def update_analysis( self, relations: Mapping[str, tuple[Mapping[str, str], Mapping[str, str]]] ) -> None: ... def register_override(self, relations: Mapping[str, Mapping[str, str]]) -> int: ... - def analyze_compact(self, payload: bytes) -> bytes: ... + def prepare_compact(self, payload: bytes) -> NativeCompactAnalysisJob: ... def update_relations(self, relations: Mapping[str, Mapping[str, str]]) -> None: ... def binding_results( self, requests: list[NativeBindingRequest] diff --git a/tests/integration/src/sqlbuild/compiler/pipeline/_test_types.py b/tests/integration/src/sqlbuild/compiler/pipeline/_test_types.py index b5dd9c4a9a..2642c01a53 100644 --- a/tests/integration/src/sqlbuild/compiler/pipeline/_test_types.py +++ b/tests/integration/src/sqlbuild/compiler/pipeline/_test_types.py @@ -3,6 +3,7 @@ from collections.abc import Callable from dataclasses import dataclass, field +from scripts.cold_compile_performance.models import RandomDagProject from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter @@ -389,3 +390,90 @@ class FunctionArgumentTypeCase: description: str projection: str expected_codes: tuple[str, ...] = () + + +type CompileOutcome = tuple[int, dict[str, object], dict[str, bytes]] + + +@dataclass(frozen=True) +class DataflowOracleCase: + description: str + project: RandomDagProject + reshaped_steps: tuple[tuple[int, ...], ...] + expected_exit_codes: tuple[int, ...] + compile_args: tuple[str, ...] = () + + +@dataclass(frozen=True) +class DataflowFixtureCase: + description: str + fixture: str + expected_exit_code: int + + +@dataclass(frozen=True) +class DataflowDenseCase: + description: str + model_count: int + expected_exit_codes: tuple[int, ...] + + +@dataclass(frozen=True) +class DataflowBuildCase: + description: str + project: RandomDagProject + expected_compile_exit_code: int + expected_build_exit_code: int + + +@dataclass(frozen=True) +class DataflowScheduleCase: + description: str + workers: int + batch_limit: int + max_delay_seconds: float + expected_cold_exit_code: int = 1 + expected_edit_exit_code: int = 1 + + +@dataclass(frozen=True) +class DataflowFailureCase: + description: str + failing_models: tuple[str, ...] + expected_message: str + + +@dataclass(frozen=True) +class DataflowPoolCase: + description: str + compiles: int + minimum_batches: int + expected_pool_threads: int + + +@dataclass(frozen=True) +class DataflowInterruptCase: + description: str + expected_live_workers: tuple[str, ...] + expected_overlaps: tuple[str, ...] + expected_thread_errors: int + expected_notices: int + + +@dataclass(frozen=True) +class DataflowInterruptAfterFaultCase: + description: str + failing_model: str + interrupted_model: str + expected_wave_replays: int + + +@dataclass(frozen=True) +class DataflowStartFailureCase: + description: str + failing_start: int + interrupt_after_start: bool + expected_error_type: type[BaseException] + expected_error: str + expected_live_workers: tuple[str, ...] + expected_notices: int diff --git a/tests/integration/src/sqlbuild/compiler/pipeline/helpers.py b/tests/integration/src/sqlbuild/compiler/pipeline/helpers.py index 906227cb77..e0b685df6a 100644 --- a/tests/integration/src/sqlbuild/compiler/pipeline/helpers.py +++ b/tests/integration/src/sqlbuild/compiler/pipeline/helpers.py @@ -1,7 +1,13 @@ from __future__ import annotations +import itertools import json +import random +import signal +import threading +import time from collections.abc import Callable +from contextlib import suppress from pathlib import Path from typing import Any @@ -9,7 +15,13 @@ 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 from sqlbuild.compiler.compile._helpers.analysis import compact +from sqlbuild.compiler.compile._helpers.analysis.compact import native_section +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.classes.binding_dataflow import BindingDataflow from sqlbuild.compiler.compile.main._build_compile_inputs import build_compile_inputs from sqlbuild.compiler.compile.models import ( CompactBatchPreparation, @@ -27,7 +39,11 @@ from sqlbuild.compiler.pipeline.models import CompilePipelineOptions, CompilePipelineResult from sqlbuild.runtime.contracts.models import ConnectionHooks from sqlbuild.sql_values.types import CollectionRendering -from tests.integration.src.sqlbuild.compiler.pipeline._test_types import SharedBindingQueryCase +from tests.integration.src.sqlbuild.compiler.pipeline._test_types import ( + CompileOutcome, + DataflowScheduleCase, + SharedBindingQueryCase, +) _SCHEMA_FIXTURE_PATH: Path = ( Path(__file__).resolve().parents[5] / "fixtures" / "dbt_manifest_v12_schema.json" @@ -37,6 +53,12 @@ "staged_orders": "SELECT 1 AS customer_id, 'books' AS category, 10 AS amount", "wide_orders": "SELECT 1 AS customer_id, 2 AS books, 3 AS games", } +_ANALYSIS_THREAD_NAME: str = "sqlbuild-analys" +_MODEL_ANALYSIS_THREAD_PREFIX: str = "sqlbuild-model-analysis-" +_HOLD_SECONDS: float = 0.3 +_REACQUIRE_DELAY_SECONDS: float = 0.05 +_INTERRUPT_TIMEOUT_SECONDS: float = 30.0 +_INTERRUPT_DELAY_SECONDS: float = 0.15 _SEMANTIC_BINDING_PROJECT_TOML: str = 'name = "semantic_binding"\nadapter = "duckdb"\n' _AUDIT_FACTORY_ADAPTER_CONTEXT: CompileAdapterContext = CompileAdapterContext( value_renderer=DuckDbAdapter(), @@ -222,3 +244,272 @@ def apply_warehouse_setup_sql(*, project_dir: Path, statements: tuple[str, ...]) connection.execute(statement) finally: connection.close() + + +def reshape_random_dag_model(*, project_dir: Path, name: str) -> None: + """Add one leading output column, changing the model's published shape signature.""" + + path: Path = project_dir / "models" / f"{name}.sql" + header, body = path.read_text(encoding="utf-8").split("\n", 1) + path.write_text( + f"{header}\nSELECT CAST(7 AS INTEGER) AS reshaped, *\n" + f"FROM ({body.strip()}) AS reshaped_input\n", + encoding="utf-8", + ) + + +def compile_outcome( + *, project_dir: Path, args: tuple[str, ...], capsys: pytest.CaptureFixture[str] +) -> CompileOutcome: + """Compile through the CLI; return the exit code, timing-free JSON, and artifacts.""" + + 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) + compiled: Path = project_dir / "target" / "compiled" + return ( + exit_code, + payload, + { + path.relative_to(compiled).as_posix(): path.read_bytes() + for path in compiled.rglob("*.sql") + }, + ) + + +def reshape_models(*, project_dir: Path, names: tuple[str, ...], indexes: tuple[int, ...]) -> None: + """Reshape the models at `indexes` of the generated topological name order.""" + + for index in indexes: + reshape_random_dag_model(project_dir=project_dir, name=names[index]) + + +def use_wave_analysis(monkeypatch: pytest.MonkeyPatch) -> None: + """Analyze one topological level at a time, the reference schedule.""" + + monkeypatch.setattr(project_assembly, "analyze_binding_dataflow", analyze_binding_waves) + + +def perturb_dataflow_schedule( + *, monkeypatch: pytest.MonkeyPatch, case: DataflowScheduleCase, seed: int +) -> None: + """Change worker count and batch size, and delay native completions at random.""" + + rng: random.Random = random.Random(seed) + original: Callable[..., object] = compact._run_compact_analysis_batch + + def delayed(*, preparation: CompactBatchPreparation) -> object: + result: object = original(preparation=preparation) + with native_section(): + time.sleep(rng.random() * case.max_delay_seconds) + return result + + monkeypatch.setattr(binding_waves, "_DATAFLOW_WORKERS", case.workers) + monkeypatch.setattr(binding_waves, "_DATAFLOW_BATCH_MIN", 1) + monkeypatch.setattr(binding_waves, "_DATAFLOW_BATCH_LIMIT", case.batch_limit) + monkeypatch.setattr(compact, "_run_compact_analysis_batch", delayed) + + +def fail_model_analysis(*, monkeypatch: pytest.MonkeyPatch, model_names: tuple[str, ...]) -> None: + """Raise from analysis of a batch holding a named model, as an internal fault would.""" + + original: Callable[..., tuple[Any, ...]] = project_assembly._analyze_model_sql_requests + failing_names: frozenset[str] = frozenset(model_names) + + def failing(**kwargs: Any) -> tuple[Any, ...]: + batch_names: set[str] = { + request.model_input.model_file.file_path.stem for request in kwargs["requests"] + } + for name in sorted(batch_names & failing_names)[:1]: + _raise_analysis_fault(name) + return original(**kwargs) + + monkeypatch.setattr(project_assembly, "_analyze_model_sql_requests", failing) + + +def _raise_analysis_fault(name: str) -> None: + raise RuntimeError(f"analysis fault in {name}") + + +def analysis_pool_thread_ids(monkeypatch: pytest.MonkeyPatch) -> list[set[str]]: + """Record native analysis threads now, then again after each compact batch.""" + + original: Callable[..., object] = compact._run_compact_analysis_batch + observed: list[set[str]] = [_analysis_thread_ids()] + + def traced(*, preparation: CompactBatchPreparation) -> object: + result: object = original(preparation=preparation) + observed.append(_analysis_thread_ids()) + return result + + monkeypatch.setattr(compact, "_run_compact_analysis_batch", traced) + return observed + + +def _analysis_thread_ids() -> set[str]: + threads_by_name: dict[str, set[str]] = {} + for task in Path("/proc/self/task").iterdir(): + with suppress(OSError): + name: str = (task / "comm").read_text(encoding="utf-8").strip() + threads_by_name.setdefault(name, set()).add(task.name) + return threads_by_name.get(_ANALYSIS_THREAD_NAME, set()) + + +def interrupt_reacquiring_analysis(*, monkeypatch: pytest.MonkeyPatch) -> list[str]: + """Send SIGINT while a worker thread holds the analysis lock and others reacquire it. + + Return overlaps, which name any batch that entered Python analysis while another held it. + """ + + original: Callable[..., object] = compact._run_compact_analysis_batch + holder: threading.Event = threading.Event() + owner: list[str] = [] + overlaps: list[str] = [] + worker_calls: itertools.count[int] = itertools.count() + hold_calls: dict[bool, Callable[[], int]] = { + False: lambda: next(worker_calls), + True: lambda: -1, + } + + def guarded(*, preparation: CompactBatchPreparation) -> object: + overlaps.extend(owner) + owner.append(threading.current_thread().name) + call: int = hold_calls[threading.current_thread() is threading.main_thread()]() + {0: _hold_analysis_lock}.get(call, _skip_hold)(holder) + owner.clear() + result: object = original(preparation=preparation) + with native_section(): + time.sleep(_REACQUIRE_DELAY_SECONDS) + return result + + def interrupt() -> None: + _ = holder.wait(_INTERRUPT_TIMEOUT_SECONDS) + time.sleep(_INTERRUPT_DELAY_SECONDS) + signal.pthread_kill(threading.main_thread().ident or 0, signal.SIGINT) + + monkeypatch.setattr(binding_waves, "_DATAFLOW_BATCH_MIN", 1) + monkeypatch.setattr(binding_waves, "_DATAFLOW_BATCH_LIMIT", 1) + monkeypatch.setattr(compact, "_run_compact_analysis_batch", guarded) + threading.Thread(target=interrupt, daemon=True).start() + return overlaps + + +def _hold_analysis_lock(holder: threading.Event) -> None: + holder.set() + time.sleep(_HOLD_SECONDS) + + +def _skip_hold(holder: threading.Event) -> None: + del holder + + +def interrupt_after_analysis_fault( + *, monkeypatch: pytest.MonkeyPatch, failing_model: str, interrupted_model: str +) -> list[int]: + """Fail one batch, then interrupt a batch that started beside it; return wave replays.""" + + original: Callable[..., tuple[Any, ...]] = project_assembly._analyze_model_sql_requests + actions: dict[str, Callable[[str], None]] = { + failing_model: _fail_after_native_delay, + interrupted_model: _interrupt_after_native_delay, + } + replays: list[int] = [] + analyze_waves: Callable[..., Any] = BindingDataflow.analyze_waves + + def failing(**kwargs: Any) -> tuple[Any, ...]: + batch_names: set[str] = { + request.model_input.model_file.file_path.stem for request in kwargs["requests"] + } + for name in sorted(batch_names & actions.keys()): + actions[name](name) + return original(**kwargs) + + def replay(self: BindingDataflow) -> Any: + replays.append(1) + return analyze_waves(self) + + monkeypatch.setattr(binding_waves, "_DATAFLOW_BATCH_MIN", 1) + monkeypatch.setattr(binding_waves, "_DATAFLOW_BATCH_LIMIT", 1) + monkeypatch.setattr(project_assembly, "_analyze_model_sql_requests", failing) + monkeypatch.setattr(BindingDataflow, "analyze_waves", replay) + return replays + + +def _fail_after_native_delay(name: str) -> None: + with native_section(): + time.sleep(_REACQUIRE_DELAY_SECONDS) + raise RuntimeError(f"analysis fault in {name}") + + +def _interrupt_after_native_delay(name: str) -> None: + with native_section(): + time.sleep(_HOLD_SECONDS) + raise KeyboardInterrupt(name) + + +def fail_worker_start( + *, monkeypatch: pytest.MonkeyPatch, failing_start: int, interrupt_after_start: bool +) -> list[str]: + """Fail one analysis worker start; return the names of analysis batches' threads. + + The failing start either never starts its thread or starts it and is then interrupted. + """ + + original_start: Callable[[threading.Thread], None] = threading.Thread.start + starts: itertools.count[int] = itertools.count() + batches: list[str] = trace_analysis_threads(monkeypatch) + + def start_then_interrupt(thread: threading.Thread) -> None: + original_start(thread) + raise KeyboardInterrupt(thread.name) + + failing: Callable[[threading.Thread], None] = { + False: _refuse_start, + True: start_then_interrupt, + }[interrupt_after_start] + + def counted_start(thread: threading.Thread) -> None: + {failing_start: failing}.get(next(starts), original_start)(thread) + + def start(thread: threading.Thread) -> None: + {True: counted_start, False: original_start}[ + thread.name.startswith(_MODEL_ANALYSIS_THREAD_PREFIX) + ](thread) + + monkeypatch.setattr(binding_waves, "_DATAFLOW_BATCH_MIN", 1) + monkeypatch.setattr(binding_waves, "_DATAFLOW_BATCH_LIMIT", 1) + monkeypatch.setattr(threading.Thread, "start", start) + return batches + + +def _refuse_start(thread: threading.Thread) -> None: + raise RuntimeError(f"cannot start {thread.name}") + + +def trace_analysis_threads(monkeypatch: pytest.MonkeyPatch) -> list[str]: + """Record the thread that runs each compact analysis batch, slowed to overlap batches.""" + + original: Callable[..., object] = compact._run_compact_analysis_batch + threads: list[str] = [] + + def traced(*, preparation: CompactBatchPreparation) -> object: + threads.append(threading.current_thread().name) + with native_section(): + time.sleep(_REACQUIRE_DELAY_SECONDS) + return original(preparation=preparation) + + monkeypatch.setattr(compact, "_run_compact_analysis_batch", traced) + return threads + + +def live_model_analysis_threads() -> list[str]: + """Return model analysis worker threads that are still alive.""" + + names_by_kind: dict[bool, list[str]] = {} + for thread in threading.enumerate(): + names_by_kind.setdefault(thread.name.startswith(_MODEL_ANALYSIS_THREAD_PREFIX), []).append( + thread.name + ) + return sorted(names_by_kind.get(True, [])) diff --git a/tests/integration/src/sqlbuild/compiler/pipeline/test_dataflow_analysis_oracle.py b/tests/integration/src/sqlbuild/compiler/pipeline/test_dataflow_analysis_oracle.py new file mode 100644 index 0000000000..7e278c7872 --- /dev/null +++ b/tests/integration/src/sqlbuild/compiler/pipeline/test_dataflow_analysis_oracle.py @@ -0,0 +1,215 @@ +"""Dependency-driven model analysis must compile exactly what wave-by-wave analysis compiles.""" + +import shutil +from pathlib import Path + +import pytest + +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.cli.commands.main.entrypoint.entry import main +from sqlbuild.compiler.compile.main import _build_compile_inputs +from tests.integration.src.sqlbuild.compiler.pipeline._test_types import ( + CompileOutcome, + DataflowBuildCase, + DataflowDenseCase, + DataflowFixtureCase, + DataflowOracleCase, +) +from tests.integration.src.sqlbuild.compiler.pipeline.helpers import ( + compile_outcome, + reshape_models, + use_wave_analysis, +) + +_REPOSITORY_ROOT: Path = Path(__file__).resolve().parents[6] + + +@pytest.mark.parametrize( + "test_case", + [ + DataflowOracleCase( + description="inferred shapes", + project=RandomDagProject(seed=11, model_count=48), + reshaped_steps=((), (), (16,), (24, 40), ()), + expected_exit_codes=(0, 0, 0, 0, 0), + ), + DataflowOracleCase( + description="binding errors mid-graph", + project=RandomDagProject(seed=23, model_count=48, errors=True), + reshaped_steps=((), (), (16,), (24, 40), ()), + expected_exit_codes=(1, 1, 1, 1, 1), + ), + DataflowOracleCase( + description="run ids and analysis opt-outs", + project=RandomDagProject(seed=37, model_count=40, run_ids=True, analysis_opt_outs=True), + reshaped_steps=((), (), (13,), (20,), ()), + expected_exit_codes=(0, 0, 0, 0, 0), + ), + DataflowOracleCase( + description="built-in rules", + project=RandomDagProject(seed=41, model_count=32, rules=("SQBR",)), + reshaped_steps=((), (10,), ()), + expected_exit_codes=(1, 1, 1), + ), + DataflowOracleCase( + description="analysis cache disabled", + project=RandomDagProject(seed=53, model_count=48, errors=True), + reshaped_steps=((), (16,)), + expected_exit_codes=(1, 1), + compile_args=("--no-cache",), + ), + DataflowOracleCase( + description="missing reference", + project=RandomDagProject(seed=61, model_count=24, missing_reference=True), + reshaped_steps=((), (8,)), + expected_exit_codes=(1, 1), + ), + DataflowOracleCase( + description="reference cycle", + project=RandomDagProject(seed=67, model_count=24, cycle=True), + reshaped_steps=((), (8,)), + expected_exit_codes=(0, 0), + ), + ], + ids=lambda case: case.description, +) +def test_given_random_dag_edit_sequence_when_compiling_then_dataflow_matches_waves( + test_case: DataflowOracleCase, + tmp_path: Path, + capsys: pytest.CaptureFixture[str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + waves_dir: Path = tmp_path / "waves" + dataflow_dir: Path = tmp_path / "dataflow" + names: tuple[str, ...] = write_random_dag_project( + project_dir=waves_dir, project=test_case.project + ) + _ = write_random_dag_project(project_dir=dataflow_dir, project=test_case.project) + exit_codes: list[int] = [] + for step, indexes in enumerate(test_case.reshaped_steps): + monkeypatch.setattr( + _build_compile_inputs, + "resolve_run_id", + lambda *, selected_run_id, run_id=f"run-{step}": selected_run_id or run_id, + ) + reshape_models(project_dir=waves_dir, names=names, indexes=indexes) + reshape_models(project_dir=dataflow_dir, names=names, indexes=indexes) + with monkeypatch.context() as waves_patch: + use_wave_analysis(waves_patch) + expected: CompileOutcome = compile_outcome( + project_dir=waves_dir, args=test_case.compile_args, capsys=capsys + ) + actual: CompileOutcome = compile_outcome( + project_dir=dataflow_dir, args=test_case.compile_args, capsys=capsys + ) + uncached: CompileOutcome = compile_outcome( + project_dir=dataflow_dir, args=("--no-cache",), capsys=capsys + ) + exit_codes.append(actual[0]) + + assert actual == expected, f"step {step} differs from wave analysis" + assert actual == uncached, f"step {step} differs from an uncached compile" + + assert tuple(exit_codes) == test_case.expected_exit_codes + + +@pytest.mark.parametrize( + "test_case", + [ + DataflowFixtureCase("waffle shop fixture", "tests/e2e/fixtures/waffle_shop", 0), + DataflowFixtureCase("waffle shop example", "website/examples/waffle-shop", 0), + DataflowFixtureCase("tidy shop example", "website/examples/tidy-shop", 0), + DataflowFixtureCase("hero shop example", "website/examples/hero-shop", 0), + ], + ids=lambda case: case.description, +) +def test_given_repository_fixture_when_compiling_cold_then_dataflow_matches_waves( + test_case: DataflowFixtureCase, + tmp_path: Path, + capsys: pytest.CaptureFixture[str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + project_dir: Path = tmp_path / "project" + _ = shutil.copytree( + _REPOSITORY_ROOT / test_case.fixture, + project_dir, + ignore=shutil.ignore_patterns("target", "*.duckdb"), + ) + monkeypatch.setattr( + _build_compile_inputs, + "resolve_run_id", + lambda *, selected_run_id: selected_run_id or "run-0", + ) + with monkeypatch.context() as waves_patch: + use_wave_analysis(waves_patch) + expected: CompileOutcome = compile_outcome(project_dir=project_dir, args=(), capsys=capsys) + shutil.rmtree(project_dir / "target") + actual: CompileOutcome = compile_outcome(project_dir=project_dir, args=(), capsys=capsys) + + assert actual == expected + assert actual[0] == test_case.expected_exit_code + + +@pytest.mark.parametrize( + "test_case", + [DataflowDenseCase("dense joins unions macros and functions", 64, (0, 0))], + ids=lambda case: case.description, +) +def test_given_dense_fixture_when_compiling_cold_and_warm_then_dataflow_matches_waves( + test_case: DataflowDenseCase, + tmp_path: Path, + capsys: pytest.CaptureFixture[str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + waves_dir: Path = tmp_path / "waves" + dataflow_dir: Path = tmp_path / "dataflow" + write_dense_compile_project(project_dir=waves_dir, model_count=test_case.model_count) + write_dense_compile_project(project_dir=dataflow_dir, model_count=test_case.model_count) + with monkeypatch.context() as waves_patch: + use_wave_analysis(waves_patch) + expected: tuple[CompileOutcome, CompileOutcome] = ( + compile_outcome(project_dir=waves_dir, args=(), capsys=capsys), + compile_outcome(project_dir=waves_dir, args=(), capsys=capsys), + ) + actual: tuple[CompileOutcome, CompileOutcome] = ( + compile_outcome(project_dir=dataflow_dir, args=(), capsys=capsys), + compile_outcome(project_dir=dataflow_dir, args=(), capsys=capsys), + ) + + assert actual == expected + assert (actual[0][0], actual[1][0]) == test_case.expected_exit_codes + + +@pytest.mark.parametrize( + "test_case", + [DataflowBuildCase("inferred shapes", RandomDagProject(seed=79, model_count=24), 0, 0)], + ids=lambda case: case.description, +) +def test_given_random_dag_when_building_on_duckdb_then_every_model_materializes( + test_case: DataflowBuildCase, tmp_path: Path, capsys: pytest.CaptureFixture[str] +) -> None: + project_dir: Path = tmp_path / "orders" + names: tuple[str, ...] = write_random_dag_project( + project_dir=project_dir, project=test_case.project + ) + config: Path = project_dir / "sqlbuild_project.toml" + _ = config.write_text( + config.read_text(encoding="utf-8") + + f'[connection]\ndatabase = "{project_dir / "orders.duckdb"}"\n', + encoding="utf-8", + ) + + compiled: int = main(["--project-dir", str(project_dir), "compile", "--json"]) + compile_output: str = capsys.readouterr().out + built: int = main(["--project-dir", str(project_dir), "build", "--no-tests", "--no-audits"]) + build_output: str = "".join(capsys.readouterr()) + + assert compiled == test_case.expected_compile_exit_code, compile_output + assert built == test_case.expected_build_exit_code, build_output + assert all(name in build_output for name in names) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-n", "auto", "--dist", "loadfile"])) diff --git a/tests/integration/src/sqlbuild/compiler/pipeline/test_dataflow_analysis_schedule.py b/tests/integration/src/sqlbuild/compiler/pipeline/test_dataflow_analysis_schedule.py new file mode 100644 index 0000000000..4fb2220e19 --- /dev/null +++ b/tests/integration/src/sqlbuild/compiler/pipeline/test_dataflow_analysis_schedule.py @@ -0,0 +1,250 @@ +"""Dataflow analysis output must not depend on worker count, batching, or completion order.""" + +import threading +import time +from pathlib import Path + +import pytest +from _pytest.capture import CaptureResult + +from scripts.cold_compile_performance._helpers.random_dag_project import write_random_dag_project +from scripts.cold_compile_performance.models import RandomDagProject +from sqlbuild.cli.commands.main.entrypoint.entry import main +from tests.integration.src.sqlbuild.compiler.pipeline._test_types import ( + CompileOutcome, + DataflowFailureCase, + DataflowInterruptAfterFaultCase, + DataflowInterruptCase, + DataflowPoolCase, + DataflowScheduleCase, + DataflowStartFailureCase, +) +from tests.integration.src.sqlbuild.compiler.pipeline.helpers import ( + analysis_pool_thread_ids, + compile_outcome, + fail_model_analysis, + fail_worker_start, + interrupt_after_analysis_fault, + interrupt_reacquiring_analysis, + live_model_analysis_threads, + perturb_dataflow_schedule, + reshape_models, + use_wave_analysis, +) + +_PROJECT: RandomDagProject = RandomDagProject(seed=97, model_count=40, errors=True) +_EDITED: tuple[int, ...] = (10, 25) +_REPEATS: tuple[int, ...] = (0, 1) +_SETTLE_SECONDS: float = 0.3 +_INTERRUPT_NOTICE: str = "Interrupted; finishing in-flight model analysis..." + + +@pytest.mark.parametrize( + "test_case", + [ + DataflowScheduleCase("one worker with whole ready sets", 1, 64, 0.0), + DataflowScheduleCase("two workers with single models", 2, 1, 0.004), + DataflowScheduleCase("four workers with small batches", 4, 3, 0.004), + DataflowScheduleCase("four workers with wide batches", 4, 64, 0.002), + ], + ids=lambda case: case.description, +) +def test_given_perturbed_schedule_when_compiling_repeatedly_then_output_is_identical( + test_case: DataflowScheduleCase, + tmp_path: Path, + capsys: pytest.CaptureFixture[str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + reference_dir: Path = tmp_path / "reference" + names: tuple[str, ...] = write_random_dag_project(project_dir=reference_dir, project=_PROJECT) + with monkeypatch.context() as waves_patch: + use_wave_analysis(waves_patch) + cold: CompileOutcome = compile_outcome(project_dir=reference_dir, args=(), capsys=capsys) + reshape_models(project_dir=reference_dir, names=names, indexes=_EDITED) + edit: CompileOutcome = compile_outcome(project_dir=reference_dir, args=(), capsys=capsys) + for repeat in _REPEATS: + project_dir: Path = tmp_path / f"dataflow-{repeat}" + _ = write_random_dag_project(project_dir=project_dir, project=_PROJECT) + with monkeypatch.context() as schedule_patch: + perturb_dataflow_schedule(monkeypatch=schedule_patch, case=test_case, seed=repeat) + actual_cold: CompileOutcome = compile_outcome( + project_dir=project_dir, args=(), capsys=capsys + ) + reshape_models(project_dir=project_dir, names=names, indexes=_EDITED) + actual_edit: CompileOutcome = compile_outcome( + project_dir=project_dir, args=(), capsys=capsys + ) + + assert actual_cold == cold, f"cold compile differs in repeat {repeat}" + assert actual_edit == edit, f"edited compile differs in repeat {repeat}" + assert actual_cold[0] == test_case.expected_cold_exit_code + assert actual_edit[0] == test_case.expected_edit_exit_code + + +@pytest.mark.parametrize( + "test_case", + [ + DataflowFailureCase("one failing model", ("orders_020",), "analysis fault in orders_020"), + DataflowFailureCase( + "an upstream failure wins", ("orders_035", "orders_004"), "analysis fault in orders_004" + ), + ], + ids=lambda case: case.description, +) +def test_given_analysis_fault_when_completion_order_varies_then_first_wave_error_is_raised( + test_case: DataflowFailureCase, + tmp_path: Path, + capsys: pytest.CaptureFixture[str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + project_dir: Path = tmp_path / "orders" + _ = write_random_dag_project(project_dir=project_dir, project=_PROJECT) + args: list[str] = ["--project-dir", str(project_dir), "compile", "--json", "--no-cache"] + fail_model_analysis(monkeypatch=monkeypatch, model_names=test_case.failing_models) + with monkeypatch.context() as waves_patch: + use_wave_analysis(waves_patch) + with pytest.raises(RuntimeError) as expected: + _ = main(args) + _ = capsys.readouterr() + perturb_dataflow_schedule( + monkeypatch=monkeypatch, + case=DataflowScheduleCase("four workers with delays", 4, 3, 0.004), + seed=len(test_case.failing_models), + ) + with pytest.raises(RuntimeError) as actual: + _ = main(args) + + assert str(actual.value) == str(expected.value) + assert str(actual.value) == test_case.expected_message + + +@pytest.mark.parametrize( + "test_case", + [DataflowPoolCase("two multi-batch compiles", 2, 3, 4)], + ids=lambda case: case.description, +) +def test_given_multi_batch_compiles_when_analyzing_then_each_compile_reuses_one_native_pool( + test_case: DataflowPoolCase, + tmp_path: Path, + capsys: pytest.CaptureFixture[str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + project_dir: Path = tmp_path / "orders" + _ = write_random_dag_project(project_dir=project_dir, project=_PROJECT) + for _ in range(test_case.compiles): + observed: list[set[str]] = analysis_pool_thread_ids(monkeypatch) + _ = compile_outcome(project_dir=project_dir, args=("--no-cache",), capsys=capsys) + + started: set[str] = set().union(*observed[1:]) - observed[0] + + assert len(observed) > test_case.minimum_batches + assert len(started) == test_case.expected_pool_threads + + +@pytest.mark.parametrize( + "test_case", + [DataflowInterruptCase("interrupt while workers reacquire the lock", (), (), 0, 1)], + ids=lambda case: case.description, +) +def test_given_sigint_while_workers_contend_for_lock_when_analyzing_then_interrupt_propagates_once( + test_case: DataflowInterruptCase, + tmp_path: Path, + capsys: pytest.CaptureFixture[str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + project_dir: Path = tmp_path / "orders" + _ = write_random_dag_project(project_dir=project_dir, project=_PROJECT) + thread_errors: list[threading.ExceptHookArgs] = [] + monkeypatch.setattr(threading, "excepthook", thread_errors.append) + overlaps: list[str] = interrupt_reacquiring_analysis(monkeypatch=monkeypatch) + + with pytest.raises(KeyboardInterrupt): + _ = main(["--project-dir", str(project_dir), "compile", "--json", "--no-cache"]) + captured: CaptureResult[str] = capsys.readouterr() + + assert tuple(live_model_analysis_threads()) == test_case.expected_live_workers + assert tuple(overlaps) == test_case.expected_overlaps + assert len(thread_errors) == test_case.expected_thread_errors + assert captured.err.count(_INTERRUPT_NOTICE) == test_case.expected_notices + assert _INTERRUPT_NOTICE not in captured.out + + +@pytest.mark.parametrize( + "test_case", + [DataflowInterruptAfterFaultCase("interrupt after a fault", "orders_000", "orders_001", 0)], + ids=lambda case: case.description, +) +def test_given_interrupt_after_analysis_fault_when_analyzing_then_interrupt_wins_without_replay( + test_case: DataflowInterruptAfterFaultCase, + tmp_path: Path, + capsys: pytest.CaptureFixture[str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + project_dir: Path = tmp_path / "orders" + _ = write_random_dag_project(project_dir=project_dir, project=_PROJECT) + replays: list[int] = interrupt_after_analysis_fault( + monkeypatch=monkeypatch, + failing_model=test_case.failing_model, + interrupted_model=test_case.interrupted_model, + ) + + with pytest.raises(KeyboardInterrupt, match=test_case.interrupted_model): + _ = main(["--project-dir", str(project_dir), "compile", "--json", "--no-cache"]) + _ = capsys.readouterr() + + assert len(replays) == test_case.expected_wave_replays + + +@pytest.mark.parametrize( + "test_case", + [ + DataflowStartFailureCase( + description="third worker cannot start", + failing_start=2, + interrupt_after_start=False, + expected_error_type=RuntimeError, + expected_error="cannot start", + expected_live_workers=(), + expected_notices=0, + ), + DataflowStartFailureCase( + description="interrupt after the first worker starts", + failing_start=0, + interrupt_after_start=True, + expected_error_type=KeyboardInterrupt, + expected_error="sqlbuild-model-analysis-0", + expected_live_workers=(), + expected_notices=1, + ), + ], + ids=lambda case: case.description, +) +def test_given_worker_start_failure_when_analyzing_then_started_workers_stop_before_raising( + test_case: DataflowStartFailureCase, + tmp_path: Path, + capsys: pytest.CaptureFixture[str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + project_dir: Path = tmp_path / "orders" + _ = write_random_dag_project(project_dir=project_dir, project=_PROJECT) + batches: list[str] = fail_worker_start( + monkeypatch=monkeypatch, + failing_start=test_case.failing_start, + interrupt_after_start=test_case.interrupt_after_start, + ) + + with pytest.raises(test_case.expected_error_type, match=test_case.expected_error): + _ = main(["--project-dir", str(project_dir), "compile", "--json", "--no-cache"]) + live_at_failure: list[str] = live_model_analysis_threads() + batches_at_failure: int = len(batches) + captured: CaptureResult[str] = capsys.readouterr() + time.sleep(_SETTLE_SECONDS) + + assert tuple(live_at_failure) == test_case.expected_live_workers + assert len(batches) == batches_at_failure + assert captured.err.count(_INTERRUPT_NOTICE) == test_case.expected_notices + assert _INTERRUPT_NOTICE not in captured.out + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-n", "auto", "--dist", "loadfile"])) diff --git a/tests/unit/src/sqlbuild/compiler/sql_analysis/classes/_test_types.py b/tests/unit/src/sqlbuild/compiler/sql_analysis/classes/_test_types.py new file mode 100644 index 0000000000..baa1aa4dd1 --- /dev/null +++ b/tests/unit/src/sqlbuild/compiler/sql_analysis/classes/_test_types.py @@ -0,0 +1,27 @@ +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class CompactJobSnapshotCase: + description: str + query_sql: str + prepared_types: dict[str, str] + updated_types: dict[str, str] + expected_changed_by_update: bool + + +@dataclass(frozen=True) +class CompactJobRerunCase: + description: str + query_sql: str + expected_error: str + + +@dataclass(frozen=True) +class ConcurrentCompactJobCase: + description: str + job_count: int + prepared_types: dict[str, str] + expected_matches_serial: bool diff --git a/tests/unit/src/sqlbuild/compiler/sql_analysis/classes/helpers.py b/tests/unit/src/sqlbuild/compiler/sql_analysis/classes/helpers.py new file mode 100644 index 0000000000..85d6dc703d --- /dev/null +++ b/tests/unit/src/sqlbuild/compiler/sql_analysis/classes/helpers.py @@ -0,0 +1,58 @@ +from __future__ import annotations + +import orjson + +from sqlbuild.compiler.sql_analysis.classes.binding_catalog import BindingCatalog + + +def orders_catalog(types: dict[str, str]) -> BindingCatalog: + """Return a DuckDB catalog whose analysis shape for `orders` has the given types.""" + + catalog: BindingCatalog = BindingCatalog( + dialect="duckdb", + quoted_ignore_case=False, + known_functions=(), + known_types=(), + relations={}, + ) + retype_orders(catalog=catalog, types=types) + return catalog + + +def retype_orders(*, catalog: BindingCatalog, types: dict[str, str]) -> None: + """Replace the analysis shape for `orders`.""" + + catalog.prepare_analysis( + types={"orders": types}, nullability={"orders": dict.fromkeys(types, "unknown")} + ) + + +def orders_payload(query_sql: str) -> bytes: + """Encode one compact analysis batch over the catalog relation `orders`.""" + + return orjson.dumps( + { + "queries": [ + { + "sql": query_sql, + "dialect": "duckdb", + "analysis_references": [["orders", "orders"]], + } + ], + "templates": [ + { + "queryIndex": 0, + "recoverCteFacts": True, + "references": {"orders": {"resourceType": "model", "resourceName": "orders"}}, + } + ], + "projections": [{"templateIndex": 0, "resourceNames": {"orders": "orders"}}], + "workers": 4, + } + ) + + +def analyze_now(*, catalog: BindingCatalog, payload: bytes) -> bytes: + """Prepare and run one compact batch against the catalog's current shapes.""" + + return catalog.native.prepare_compact(payload).run() diff --git a/tests/unit/src/sqlbuild/compiler/sql_analysis/classes/test_binding_catalog.py b/tests/unit/src/sqlbuild/compiler/sql_analysis/classes/test_binding_catalog.py new file mode 100644 index 0000000000..e26a8f5d6f --- /dev/null +++ b/tests/unit/src/sqlbuild/compiler/sql_analysis/classes/test_binding_catalog.py @@ -0,0 +1,111 @@ +"""Prepared compact analysis jobs keep their catalog snapshot and run beside catalog updates.""" + +from concurrent.futures import Future, ThreadPoolExecutor + +import pytest + +from sqlbuild.compiler.sql_analysis.classes.binding_catalog import BindingCatalog +from sqlbuild.compiler.sql_analysis.types import NativeCompactAnalysisJob +from tests.unit.src.sqlbuild.compiler.sql_analysis.classes._test_types import ( + CompactJobRerunCase, + CompactJobSnapshotCase, + ConcurrentCompactJobCase, +) +from tests.unit.src.sqlbuild.compiler.sql_analysis.classes.helpers import ( + analyze_now, + orders_catalog, + orders_payload, + retype_orders, +) + + +@pytest.mark.parametrize( + "test_case", + [ + CompactJobSnapshotCase( + description="retyped input", + query_sql="SELECT id, amount + 1 AS amount FROM orders", + prepared_types={"id": "INTEGER", "amount": "DOUBLE"}, + updated_types={"id": "VARCHAR", "amount": "VARCHAR"}, + expected_changed_by_update=True, + ), + CompactJobSnapshotCase( + description="star over a widened input", + query_sql="SELECT * FROM orders", + prepared_types={"id": "INTEGER"}, + updated_types={"id": "INTEGER", "status": "VARCHAR"}, + expected_changed_by_update=True, + ), + ], + ids=lambda case: case.description, +) +def test_given_prepared_job_when_catalog_changes_before_run_then_uses_prepared_shapes( + test_case: CompactJobSnapshotCase, +) -> None: + catalog: BindingCatalog = orders_catalog(test_case.prepared_types) + payload: bytes = orders_payload(test_case.query_sql) + expected: bytes = analyze_now(catalog=catalog, payload=payload) + job: NativeCompactAnalysisJob = catalog.native.prepare_compact(payload) + retype_orders(catalog=catalog, types=test_case.updated_types) + + assert job.run() == expected + assert (analyze_now(catalog=catalog, payload=payload) != expected) is ( + test_case.expected_changed_by_update + ) + + +@pytest.mark.parametrize( + "test_case", + [CompactJobRerunCase("finished job", "SELECT id FROM orders", "already ran")], + ids=lambda case: case.description, +) +def test_given_finished_job_when_running_again_then_rejects_the_rerun( + test_case: CompactJobRerunCase, +) -> None: + catalog: BindingCatalog = orders_catalog({"id": "INTEGER"}) + job: NativeCompactAnalysisJob = catalog.native.prepare_compact( + orders_payload(test_case.query_sql) + ) + _ = job.run() + + with pytest.raises(ValueError, match=test_case.expected_error): + _ = job.run() + + +@pytest.mark.parametrize( + "test_case", + [ + ConcurrentCompactJobCase( + description="eight jobs beside repeated retyping", + job_count=8, + prepared_types={"id": "INTEGER", "amount": "DOUBLE"}, + expected_matches_serial=True, + ) + ], + ids=lambda case: case.description, +) +def test_given_running_jobs_when_catalog_updates_concurrently_then_each_matches_serial_run( + test_case: ConcurrentCompactJobCase, +) -> None: + catalog: BindingCatalog = orders_catalog(test_case.prepared_types) + payloads: list[bytes] = [ + orders_payload(f"SELECT id, amount * {index} AS amount FROM orders") + for index in range(test_case.job_count) + ] + expected: list[bytes] = [analyze_now(catalog=catalog, payload=payload) for payload in payloads] + jobs: list[NativeCompactAnalysisJob] = [ + catalog.native.prepare_compact(payload) for payload in payloads + ] + with ThreadPoolExecutor(max_workers=test_case.job_count) as executor: + futures: list[Future[bytes]] = [executor.submit(job.run) for job in jobs] + for index in range(test_case.job_count): + retype_orders( + catalog=catalog, types={"id": "BIGINT", "amount": f"DECIMAL(18, {index})"} + ) + actual: list[bytes] = [future.result() for future in futures] + + assert (actual == expected) is test_case.expected_matches_serial + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-vv"]))