From f4a3d627129de32ef6e9b795fa85ef17c8ae78ce Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Mon, 5 Oct 2026 15:51:52 +0100 Subject: [PATCH] perf(compile): analyze models as soon as their producers finish Replace wave-by-wave model analysis with a dependency-driven scheduler. A model is analyzed once every producer shape it binds against is final, so a slow model no longer holds back unrelated models in the next topological level. Worker threads serialize Python preparation and projection under one lock, while self-contained native analysis releases it, so the next ready batch is prepared while earlier batches run natively. The calling thread never takes that lock: it only starts and joins the workers, and on an interrupt or start failure it cancels and joins every started worker before re-raising. Each compile catalog now owns one native analysis pool that every compact analysis and binding batch reuses, instead of building a pool per call. Compact analysis splits into a catalog-resolving prepare step and a run step that no longer borrows the catalog, so catalog updates for later batches proceed while earlier batches run. Spawned drainer tasks replace the pool broadcast, so a batch no longer needs every pool thread to finish. Every model sees exactly the producer shapes, cache entries, and changed upstream signatures it saw under wave scheduling, so compile output does not depend on completion order. An ordinary analysis error replays the wave schedule to raise the same first error; an interrupt or other non-Exception failure always wins and propagates without a replay. --- .../src/bindings/_helpers/functions.rs | 1 + .../src/query_analysis/_helpers/engine.rs | 82 +++-- .../main/analyze_project_catalog.rs | 6 - .../src/query_analysis/main/mod.rs | 2 +- .../main/prepare_project_catalog.rs | 6 + .../src/query_analysis/tests/helpers.rs | 61 ++++ .../tests/test_query_analysis.rs | 16 +- .../semantic_validation/_helpers/catalog.rs | 93 +++++- .../src/semantic_validation/models.rs | 16 +- .../src/semantic_validation/types.rs | 5 +- .../_helpers/random_dag_project.py | 153 +++++++++ scripts/cold_compile_performance/models.py | 14 + .../compile/_helpers/analysis/compact.py | 80 +++-- .../_helpers/assembly/binding_waves.py | 169 ++++------ .../compile/_helpers/assembly/project.py | 10 +- .../_helpers/assembly/semantic_shapes.py | 15 + .../compile/classes/binding_dataflow.py | 314 ++++++++++++++++++ src/sqlbuild/compiler/sql_analysis/types.py | 6 +- .../sqlbuild/compiler/pipeline/_test_types.py | 88 +++++ .../src/sqlbuild/compiler/pipeline/helpers.py | 293 +++++++++++++++- .../pipeline/test_dataflow_analysis_oracle.py | 215 ++++++++++++ .../test_dataflow_analysis_schedule.py | 250 ++++++++++++++ .../sql_analysis/classes/_test_types.py | 27 ++ .../compiler/sql_analysis/classes/helpers.py | 58 ++++ .../classes/test_binding_catalog.py | 111 +++++++ 25 files changed, 1902 insertions(+), 189 deletions(-) delete mode 100644 crates/sqlbuild-rules-native/src/query_analysis/main/analyze_project_catalog.rs create mode 100644 crates/sqlbuild-rules-native/src/query_analysis/main/prepare_project_catalog.rs create mode 100644 scripts/cold_compile_performance/_helpers/random_dag_project.py create mode 100644 src/sqlbuild/compiler/compile/classes/binding_dataflow.py create mode 100644 tests/integration/src/sqlbuild/compiler/pipeline/test_dataflow_analysis_oracle.py create mode 100644 tests/integration/src/sqlbuild/compiler/pipeline/test_dataflow_analysis_schedule.py create mode 100644 tests/unit/src/sqlbuild/compiler/sql_analysis/classes/_test_types.py create mode 100644 tests/unit/src/sqlbuild/compiler/sql_analysis/classes/helpers.py create mode 100644 tests/unit/src/sqlbuild/compiler/sql_analysis/classes/test_binding_catalog.py 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"]))