diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index e3ed2f5c22..3d5cf094d3 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -34,6 +34,7 @@ jobs: echo "tag=$tag" >> "$GITHUB_OUTPUT" wheels: + # Keep the platform matrix and maturin settings in sync with verify.yml. name: Build wheel (${{ matrix.platform.name }}) needs: release-tag runs-on: ${{ matrix.platform.runner }} diff --git a/.github/workflows/verify.yml b/.github/workflows/verify.yml index a08a616a8f..cc41b9c9e4 100644 --- a/.github/workflows/verify.yml +++ b/.github/workflows/verify.yml @@ -22,6 +22,7 @@ jobs: performance: ${{ steps.scope.outputs.performance }} format_performance: ${{ steps.scope.outputs.format_performance }} rules_performance: ${{ steps.scope.outputs.rules_performance }} + wheels: ${{ steps.scope.outputs.wheels }} steps: - uses: actions/checkout@v4 with: @@ -41,11 +42,12 @@ jobs: echo "format_performance=$4" echo "rules_performance=$5" echo "checks=$6" + echo "wheels=$7" } >> "${GITHUB_OUTPUT}" } if [[ "${EVENT_NAME}" != "pull_request" ]]; then - write_scope true true true true true true + write_scope true true true true true true true exit 0 fi @@ -54,6 +56,7 @@ jobs: performance=false format_performance=false rules_performance=false + wheels=false broad=false changed=false # A change that only touches the website, root-level Markdown, the README demo GIFs or the @@ -69,6 +72,11 @@ jobs: *.md) ;; *) docs_only=false ;; esac + case "${path}" in + crates/*|Cargo.toml|Cargo.lock|pyproject.toml|.github/workflows/verify.yml) + wheels=true + ;; + esac case "${path}" in crates/*|Cargo.toml|Cargo.lock|pyproject.toml|uv.lock) broad=true @@ -123,7 +131,7 @@ jobs: done < <(git diff --name-only -z "${BASE_SHA}...${HEAD_SHA}") if [[ "${changed}" != "true" ]]; then - write_scope true true true true true true + write_scope true true true true true true true exit 0 fi @@ -143,7 +151,8 @@ jobs: "${performance}" \ "${format_performance}" \ "${rules_performance}" \ - "${checks}" + "${checks}" \ + "${wheels}" lockfile: name: Lockfile consistency @@ -414,7 +423,7 @@ jobs: - kind: dense max_ratio: "1.10" - kind: fresh - max_ratio: "1.30" + max_ratio: "1.10" steps: - uses: actions/checkout@v4 - name: Check out the comparison base @@ -601,6 +610,59 @@ jobs: -m "real_warehouse and postgres" -vv --color=yes -n auto --dist loadfile + wheels: + # Keep the platform matrix and maturin settings in sync with publish.yml. + name: Build wheel (${{ matrix.platform.name }}) + needs: [changes] + if: needs.changes.outputs.wheels == 'true' + runs-on: ${{ matrix.platform.runner }} + timeout-minutes: 45 + strategy: + fail-fast: false + matrix: + platform: + - name: linux-x86_64 + runner: ubuntu-24.04 + target: x86_64 + - name: linux-aarch64 + runner: ubuntu-24.04-arm + target: aarch64 + - name: macos-x86_64 + runner: macos-15-intel + target: x86_64 + - name: macos-arm64 + runner: macos-15 + target: aarch64 + - name: windows-x86_64 + runner: windows-2025 + target: x86_64 + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + - uses: dtolnay/rust-toolchain@master + with: + toolchain: stable + - uses: PyO3/maturin-action@v1 + with: + command: build + args: --release --out dist --features extension-module + target: ${{ matrix.platform.target }} + manylinux: auto + - name: Smoke import the native extension from the built wheel + shell: bash + run: | + python -m venv "${RUNNER_TEMP}/wheel-smoke" + if [[ -x "${RUNNER_TEMP}/wheel-smoke/Scripts/python.exe" ]]; then + smoke_python="${RUNNER_TEMP}/wheel-smoke/Scripts/python.exe" + else + smoke_python="${RUNNER_TEMP}/wheel-smoke/bin/python" + fi + "${smoke_python}" -m pip install --no-deps dist/*.whl + cd "${RUNNER_TEMP}" + "${smoke_python}" -c "import sqlbuild._native" + verify: name: Verify if: always() @@ -618,6 +680,7 @@ jobs: cache_compile_performance, rules_performance, postgres, + wheels, ] runs-on: ubuntu-24.04 steps: @@ -636,6 +699,7 @@ jobs: CACHE_COMPILE_PERFORMANCE_RESULT: ${{ needs.cache_compile_performance.result }} RULES_PERFORMANCE_RESULT: ${{ needs.rules_performance.result }} POSTGRES_RESULT: ${{ needs.postgres.result }} + WHEELS_RESULT: ${{ needs.wheels.result }} run: | if [[ "${CHANGES_RESULT}" != "success" || "${LOCKFILE_RESULT}" != "success" ]]; then echo "Required classification or lockfile checks did not pass" >&2 @@ -655,7 +719,8 @@ jobs: "${COMPILE_PERFORMANCE_RATIO_RESULT}" \ "${CACHE_COMPILE_PERFORMANCE_RESULT}" \ "${RULES_PERFORMANCE_RESULT}" \ - "${POSTGRES_RESULT}"; do + "${POSTGRES_RESULT}" \ + "${WHEELS_RESULT}"; do if [[ "${result}" != "success" && "${result}" != "skipped" ]]; then echo "A selected verification suite did not pass: ${result}" >&2 exit 1 diff --git a/Cargo.lock b/Cargo.lock index 386148b44f..2eb564b990 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -255,6 +255,15 @@ version = "0.2.189" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2" +[[package]] +name = "libmimalloc-sys" +version = "0.1.49" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a45a52f43e1c16f667ccfe4dd8c85b7f7c204fd5e3bf46c5b0db9a5c3c0b8e9" +dependencies = [ + "cc", +] + [[package]] name = "linux-raw-sys" version = "0.12.1" @@ -282,6 +291,15 @@ dependencies = [ "autocfg", ] +[[package]] +name = "mimalloc" +version = "0.1.52" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d4139bb28d14ad1facf21d5eb8825051b326e172d216b39f6d31df53cc97862" +dependencies = [ + "libmimalloc-sys", +] + [[package]] name = "nix" version = "0.27.1" @@ -603,6 +621,7 @@ version = "0.126.5" dependencies = [ "fensu-policy", "globset", + "mimalloc", "polyglot-sql", "pyo3", "rayon", diff --git a/Cargo.toml b/Cargo.toml index 53d956f36e..1b058e560b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,6 +13,8 @@ rust-version = "1.88" [workspace.dependencies] fensu-policy = "=0.22.0" globset = "0.4" +# Python loads the extension with dlopen, so mimalloc must not use initial-exec TLS. +mimalloc = { version = "0.1", default-features = false, features = ["local_dynamic_tls"] } proc-macro2 = { version = "1", features = ["span-locations"] } pyo3 = { version = "0.26", features = ["abi3-py312"] } polyglot-sql = { git = "https://github.com/chio-labs/polyglot.git", rev = "25ad011c49f2c7cd48f4427abb8fc0a6af2e4120", default-features = false, features = [ diff --git a/crates/sqlbuild-rules-native/Cargo.toml b/crates/sqlbuild-rules-native/Cargo.toml index c1e8429a4b..eebb1245c8 100644 --- a/crates/sqlbuild-rules-native/Cargo.toml +++ b/crates/sqlbuild-rules-native/Cargo.toml @@ -17,6 +17,7 @@ extension-module = ["pyo3/extension-module"] [dependencies] fensu-policy.workspace = true globset.workspace = true +mimalloc.workspace = true polyglot-sql.workspace = true pyo3.workspace = true rayon.workspace = true diff --git a/crates/sqlbuild-rules-native/src/bindings/_helpers/allocator.rs b/crates/sqlbuild-rules-native/src/bindings/_helpers/allocator.rs new file mode 100644 index 0000000000..69a2eb13cd --- /dev/null +++ b/crates/sqlbuild-rules-native/src/bindings/_helpers/allocator.rs @@ -0,0 +1,4 @@ +//! Process-wide allocator for the many short-lived syntax trees built across analysis threads. + +#[global_allocator] +static GLOBAL_ALLOCATOR: mimalloc::MiMalloc = mimalloc::MiMalloc; diff --git a/crates/sqlbuild-rules-native/src/bindings/_helpers/mod.rs b/crates/sqlbuild-rules-native/src/bindings/_helpers/mod.rs index 579b19c7b5..6539c78531 100644 --- a/crates/sqlbuild-rules-native/src/bindings/_helpers/mod.rs +++ b/crates/sqlbuild-rules-native/src/bindings/_helpers/mod.rs @@ -1,2 +1,3 @@ +mod allocator; pub(crate) mod functions; pub(crate) mod panics; diff --git a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/planning.rs b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/planning.rs index 404be12990..0d0cb061af 100644 --- a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/planning.rs +++ b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/planning.rs @@ -22,7 +22,7 @@ use crate::compiler::_helpers::sql_tests::markers::{ replace_dbt_ref_markers, replace_named_markers, }; use crate::compiler::_helpers::sql_tests::rendering::{ - AssertionStep, ChainStep, RenderRequest, render_comparison_sql, render_dialect, + AssertionStep, ChainStep, RenderRequest, StatementCache, render_comparison_sql, render_dialect, rendered_chain_steps, }; use crate::constants::{TABLE_FUNCTION_TEST_MODE, UDF_TEST_MODE}; @@ -206,6 +206,7 @@ struct ProjectContext { analysis_templates: Arc, patterns: SqlTestPatterns, render_dialect: Arc, + rendered_statements: Arc, } #[derive(Clone)] @@ -614,6 +615,7 @@ pub(crate) fn plan_and_render_json(request_json: &str) -> Result analysis_templates: Arc::new(Mutex::new(HashMap::new())), patterns: SqlTestPatterns::new()?, render_dialect, + rendered_statements: Arc::default(), }; let workers = request.workers.clamp(1, MAX_WORKERS); let pool = rayon::ThreadPoolBuilder::new() @@ -631,8 +633,13 @@ pub(crate) fn plan_and_render_json(request_json: &str) -> Result let mut planned = plan_test(test, &context)?; let planning_ns = planning_start.elapsed().as_nanos(); let rendering_start = Instant::now(); - let sql = render_sql - .then(|| render_comparison_sql(&planned.request, &context.render_dialect)); + let sql = render_sql.then(|| { + render_comparison_sql( + &planned.request, + &context.render_dialect, + Some(&context.rendered_statements), + ) + }); if !include_plan { planned.request.chain.clear(); planned.request.assertions.clear(); diff --git a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/rendering.rs b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/rendering.rs index d10eabed1e..2c021c79e9 100644 --- a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/rendering.rs +++ b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/rendering.rs @@ -1,6 +1,8 @@ //! Deterministic comparison SQL rendering for planned SQL-native tests. -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; +use std::hash::{BuildHasher, RandomState}; +use std::sync::{Arc, Mutex, OnceLock, PoisonError}; use polyglot_sql::{Dialect, DialectType, Expression}; @@ -14,6 +16,7 @@ use serde::{Deserialize, Serialize}; const DEFAULT_WORKERS: usize = 4; const MAX_WORKERS: usize = 4; const WORKER_STACK_BYTES: usize = 16 * 1024 * 1024; +const STATEMENT_CACHE_SQL_BYTES: usize = 1024 * 1024; #[derive(Debug, Deserialize)] #[serde(rename_all = "camelCase")] @@ -97,15 +100,55 @@ struct RenderResponse { sql: String, } +/// Bounded parses of repeated SQL reused across one single-dialect render batch. +#[derive(Default)] +pub(crate) struct StatementCache { + entries: Mutex, +} + +#[derive(Default)] +struct StatementEntries { + seen: HashSet, + hasher: RandomState, + parsed: HashMap>>>, + parsed_sql_bytes: usize, +} + +impl StatementCache { + /// Admit SQL on its second use, within a byte budget, so unrepeated SQL is never retained. + fn entry(&self, sql: &str) -> Option>>> { + let mut entries = self.entries.lock().unwrap_or_else(PoisonError::into_inner); + if let Some(entry) = entries.parsed.get(sql) { + return Some(Arc::clone(entry)); + } + let fingerprint = entries.hasher.hash_one(sql); + if entries.seen.insert(fingerprint) + || entries.parsed_sql_bytes + sql.len() > STATEMENT_CACHE_SQL_BYTES + { + return None; + } + entries.parsed_sql_bytes += sql.len(); + Some(Arc::clone( + entries.parsed.entry(sql.to_string()).or_default(), + )) + } +} + struct RenderCteState<'a> { dialect: &'a Dialect, + statements: Option<&'a StatementCache>, lifted: Vec<(String, String)>, name_counts: CteSuffixCounter, namespace: CteNamespace, } impl<'a> RenderCteState<'a> { - fn new(dialect: &'a Dialect, chain: &[ChainStep], assertions: &[AssertionStep]) -> Self { + fn new( + dialect: &'a Dialect, + statements: Option<&'a StatementCache>, + chain: &[ChainStep], + assertions: &[AssertionStep], + ) -> Self { let mut namespace = CteNamespace::default(); for step in chain { namespace.reserve(&step.resolved_sql); @@ -126,6 +169,7 @@ impl<'a> RenderCteState<'a> { } Self { dialect, + statements, lifted: Vec::new(), name_counts: CteSuffixCounter::default(), namespace, @@ -136,9 +180,12 @@ impl<'a> RenderCteState<'a> { if !enabled || leading_with_prefix_end(sql).is_none() { return sql.to_string(); } - let Some((step_ctes, body_sql)) = - split_top_level_with(sql, self.dialect, namespace.then_some(&mut self.namespace)) - else { + let Some((step_ctes, body_sql)) = split_top_level_with( + sql, + self.dialect, + self.statements, + namespace.then_some(&mut self.namespace), + ) else { return sql.to_string(); }; if step_ctes.iter().any(|(name, body)| { @@ -202,26 +249,20 @@ impl<'a> RenderCteState<'a> { fn definition(&mut self, name: &str, sql: &str, enabled: bool) -> String { let scoped = if enabled && leading_with_prefix_end(sql).is_some() { let (protected, identifiers) = protect_backtick_identifiers(sql); - match self.dialect.parse(&protected) { - Ok(mut expressions) => (|| { - if expressions.len() != 1 { - return None; - } - let expression = match self - .namespace - .rewrite(expressions.pop()?, self.dialect.dialect_type()) - { - Ok(expression) => expression, - Err(_) => return None, - }; - let sql = match self.dialect.generate(&expression) { - Ok(sql) => sql, - Err(_) => return None, - }; - Some(restore_backtick_identifiers(&sql, &identifiers)) - })(), - Err(_) => None, - } + parse_statement(&protected, self.dialect, self.statements).and_then(|expression| { + let expression = match self + .namespace + .rewrite(expression, self.dialect.dialect_type()) + { + Ok(expression) => expression, + Err(_) => return None, + }; + let sql = match self.dialect.generate(&expression) { + Ok(sql) => sql, + Err(_) => return None, + }; + Some(restore_backtick_identifiers(&sql, &identifiers)) + }) } else { None }; @@ -233,8 +274,12 @@ impl<'a> RenderCteState<'a> { for (name, sql) in std::mem::take(&mut self.lifted) { if enabled && leading_with_prefix_end(&sql).is_some() - && let Some((ctes, body)) = - split_top_level_with(&sql, self.dialect, Some(&mut self.namespace)) + && let Some((ctes, body)) = split_top_level_with( + &sql, + self.dialect, + self.statements, + Some(&mut self.namespace), + ) { definitions.extend( ctes.iter() @@ -289,7 +334,7 @@ pub(crate) fn render_json(request_json: &str) -> Result { .stack_size(WORKER_STACK_BYTES) .build() .map_err(|error| error.to_string())?; - let dialects: HashMap, Dialect> = request + let dialects: HashMap, (Dialect, StatementCache)> = request .requests .iter() .map(|item| item.sql_analysis_dialect.clone()) @@ -297,15 +342,18 @@ pub(crate) fn render_json(request_json: &str) -> Result { .into_iter() .map(|name| { let dialect = render_dialect(name.as_deref()); - (name, dialect) + (name, (dialect, StatementCache::default())) }) .collect(); let responses: Vec = pool.install(|| { request .requests .into_par_iter() - .map(|request| RenderResponse { - sql: render_comparison_sql(&request, &dialects[&request.sql_analysis_dialect]), + .map(|request| { + let (dialect, statements) = &dialects[&request.sql_analysis_dialect]; + RenderResponse { + sql: render_comparison_sql(&request, dialect, Some(statements)), + } }) .collect() }); @@ -333,7 +381,7 @@ fn render_difference_sample_sql(request: &DifferenceSampleRequest, dialect: &Dia let Some(expected_input) = step.expected_cte_sql.as_deref() else { return String::new(); }; - let mut cte_state = RenderCteState::new(dialect, std::slice::from_ref(step), &[]); + let mut cte_state = RenderCteState::new(dialect, None, std::slice::from_ref(step), &[]); let actual_sql = cte_state.actual_step_sql(step, request.sql_analysis_enabled); let expected_sql = cte_state.expected_step_sql(step, expected_input, request.sql_analysis_enabled); @@ -370,11 +418,16 @@ fn compared_projection(step: &ChainStep) -> String { .map_or_else(|| "*".to_string(), |columns| columns.join(", ")) } -pub(crate) fn render_comparison_sql(request: &RenderRequest, dialect: &Dialect) -> String { +pub(crate) fn render_comparison_sql( + request: &RenderRequest, + dialect: &Dialect, + statements: Option<&StatementCache>, +) -> String { if request.chain.is_empty() && request.assertions.is_empty() { return String::new(); } - let mut cte_state = RenderCteState::new(dialect, &request.chain, &request.assertions); + let mut cte_state = + RenderCteState::new(dialect, statements, &request.chain, &request.assertions); let mut comparison_ctes: Vec = Vec::new(); let mut select_parts: Vec = Vec::new(); let rendered_steps = rendered_chain_steps(&request.chain, &request.assertions); @@ -481,20 +534,30 @@ fn existing_cte<'a>(lifted: &'a [(String, String)], name: &str) -> Option<&'a (S .find(|(existing, _)| existing.eq_ignore_ascii_case(name)) } +/// Parse one statement, reusing a batch's earlier parse of the same SQL when a cache is given. +fn parse_statement( + sql: &str, + dialect: &Dialect, + statements: Option<&StatementCache>, +) -> Option { + let parse = || match dialect.parse(sql) { + Ok(mut parsed) if parsed.len() == 1 => parsed.pop(), + _ => None, + }; + match statements.and_then(|statements| statements.entry(sql)) { + Some(cached) => cached.get_or_init(parse).clone(), + None => parse(), + } +} + fn split_top_level_with( sql: &str, dialect: &Dialect, + statements: Option<&StatementCache>, namespace: Option<&mut CteNamespace>, ) -> Option<(Vec<(String, String)>, String)> { let (protected, identifiers) = protect_backtick_identifiers(sql); - let mut statements = match dialect.parse(&protected) { - Ok(statements) => statements, - Err(_) => return None, - }; - if statements.len() != 1 { - return None; - } - let mut expression = statements.pop()?; + let mut expression = parse_statement(&protected, dialect, statements)?; if let Some(namespace) = namespace { expression = match namespace.rewrite(expression, dialect.dialect_type()) { Ok(expression) => expression, diff --git a/crates/sqlbuild-rules-native/src/compiler/tests/helpers.rs b/crates/sqlbuild-rules-native/src/compiler/tests/helpers.rs index 5673f32d9d..d54709cac0 100644 --- a/crates/sqlbuild-rules-native/src/compiler/tests/helpers.rs +++ b/crates/sqlbuild-rules-native/src/compiler/tests/helpers.rs @@ -1442,6 +1442,38 @@ pub(super) fn quoted_names_keep_bindings() -> bool { && rendered.contains("JOIN __sqb_cte_1 AS b") && rendered.contains("'final' AS label") } +pub(super) fn repeated_model_sql_renders_like_separate_batches() -> bool { + let step = |reserved: &str| { + json!({ + "sqlAnalysisDialect": "duckdb", + "chain": [{ + "modelName": "orders", + "resolvedSql": "WITH final AS (SELECT 2 AS order_id) SELECT final.order_id FROM final", + "expectedCteSql": format!("SELECT 2 AS order_id {reserved}") + }] + }) + }; + let render = |requests: Vec| -> Vec { + let response = crate::compiler::_helpers::sql_tests::rendering::render_json( + &json!({"requests": requests}).to_string(), + ) + .expect("render batch"); + serde_json::from_str::>(&response).expect("render response") + }; + let plain = step(""); + let reserved = step("-- __sqb_cte_0"); + let separate: Vec = [&plain, &reserved, &plain] + .into_iter() + .flat_map(|request| render(vec![request.clone()])) + .collect(); + let batched = render(vec![plain.clone(), reserved.clone(), plain]); + separate == batched + && batched[0] != batched[1] + && batched[0]["sql"] + .as_str() + .is_some_and(|sql| sql.contains("__sqb_cte_0 AS")) +} + pub(super) fn model_cte_names_are_isolated_across_dialects() -> bool { [("duckdb", "final"), ("snowflake", "final"), ("postgres", "final"), ("tsql", "final"), ("bigquery", "`final`"), ("databricks", "`final`")].iter().all(|(dialect, name)| { let request = json!({ diff --git a/crates/sqlbuild-rules-native/src/compiler/tests/test_sql_test_rendering.rs b/crates/sqlbuild-rules-native/src/compiler/tests/test_sql_test_rendering.rs index ae54b27aea..880425d85d 100644 --- a/crates/sqlbuild-rules-native/src/compiler/tests/test_sql_test_rendering.rs +++ b/crates/sqlbuild-rules-native/src/compiler/tests/test_sql_test_rendering.rs @@ -7,7 +7,7 @@ use crate::compiler::tests::helpers::{ }; use crate::compiler::tests::helpers::{ bounded_names_are_unique, model_cte_names_are_isolated_across_dialects, - quoted_names_keep_bindings, + quoted_names_keep_bindings, repeated_model_sql_renders_like_separate_batches, }; use crate::compiler::tests::test_types::SqlTestRenderingTestCase; @@ -19,6 +19,11 @@ fn given_sql_rendering_cases_when_rendering_native_batches_then_expected_behavio run: model_cte_names_are_isolated_across_dialects, expected_success: true, }, + SqlTestRenderingTestCase { + description: "repeated model SQL in one batch renders like separate batches", + run: repeated_model_sql_renders_like_separate_batches, + expected_success: true, + }, SqlTestRenderingTestCase { description: "long and reserved CTE names are bounded, unique and deterministic", run: bounded_names_are_unique, 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 19c0d2fcfb..2169948894 100644 --- a/crates/sqlbuild-rules-native/src/query_analysis/_helpers/engine.rs +++ b/crates/sqlbuild-rules-native/src/query_analysis/_helpers/engine.rs @@ -9,6 +9,8 @@ use polyglot_sql::{ use rayon::iter::{IntoParallelIterator, ParallelIterator}; use serde::{Deserialize, Serialize}; use std::collections::{HashMap, HashSet}; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Mutex, PoisonError}; use crate::constants::{SQL_WILDCARD, UNKNOWN_SQL_TYPE, VARCHAR_SQL_TYPE}; @@ -430,13 +432,7 @@ pub(crate) fn analyze_project_compact_with_catalog( if batch.is_empty() { break; } - let analysis_groups: Vec = pool.install(|| { - batch - .into_par_iter() - .map(|work| analyze_compact_query_work(work, catalog)) - .collect() - }); - for group in analysis_groups { + for group in analyze_largest_first(&pool, batch, catalog) { validations.push(group.validation.transpose()?); for (projection_index, analysis) in group.projections { accumulator.compact_analysis(projection_index, analysis)?; @@ -454,6 +450,39 @@ 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. +fn analyze_largest_first( + pool: &rayon::ThreadPool, + batch: Vec, + catalog: Option<&crate::semantic_validation::models::ProjectCatalog>, +) -> Vec { + let mut order: Vec = (0..batch.len()).collect(); + order.sort_by_key(|index| std::cmp::Reverse(batch[*index].query.sql.len())); + let slots: Vec>> = batch + .into_iter() + .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))); + } + } + finished + }); + let mut results: Vec> = slots.iter().map(|_| None).collect(); + for (index, result) in finished.into_iter().flatten() { + results[index] = Some(result); + } + results.into_iter().flatten().collect() +} + fn analyze_compact_query_work( work: CompactQueryWork, catalog: Option<&crate::semantic_validation::models::ProjectCatalog>, 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 0641838a6e..2b346c861d 100644 --- a/crates/sqlbuild-rules-native/src/query_analysis/tests/helpers.rs +++ b/crates/sqlbuild-rules-native/src/query_analysis/tests/helpers.rs @@ -335,6 +335,51 @@ pub(crate) fn interleaved_query_templates_preserve_template_order() -> bool { names == vec!["second_value", "first_value", "second_value"] } +pub(crate) fn larger_later_queries_keep_request_order() -> bool { + let response: Value = serde_json::from_str( + &analyze_project_compact_json( + &json!({ + "queries": [ + {"sql": "SELECT 1 AS first_value", "dialect": "duckdb"}, + {"sql": "SELECT 2 AS second_value, 3 AS third_value", "dialect": "duckdb"}, + {"sql": "SELECT 4 AS fourth_value, 5 AS fifth_value, 6 AS sixth_value", "dialect": "duckdb"} + ], + "templates": [ + {"queryIndex": 0}, + {"queryIndex": 1}, + {"queryIndex": 2} + ], + "projections": [ + {"templateIndex": 0}, + {"templateIndex": 1}, + {"templateIndex": 2} + ], + "workers": 2 + }) + .to_string(), + ) + .expect("test assumption must hold"), + ) + .expect("test assumption must hold"); + + response["strings"] + == json!([ + "first_value", + "INT", + "second_value", + "third_value", + "fourth_value", + "fifth_value", + "sixth_value" + ]) + && response["templates"] + == json!([ + [[0], false, false], + [[1, 2], false, false], + [[3, 4, 5], false, false] + ]) +} + pub(crate) fn canonical_queries_reuse_semantics_and_project_resources() -> bool { let response: Value = serde_json::from_str( &analyze_project_compact_json( 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 4438c97314..051e702212 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 @@ -7,7 +7,7 @@ use crate::query_analysis::tests::helpers::{ bound_lineage_survives_type_recovery, canonical_queries_reuse_semantics_and_project_resources, combined_queries_preserve_standalone_binding, compact_project_query_interns_repeated_lineage_strings, - interleaved_query_templates_preserve_template_order, + 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, @@ -81,6 +81,11 @@ fn given_compact_query_cases_when_analyzing_projects_then_expected_behavior_hold run: interleaved_query_templates_preserve_template_order, expected_success: true, }, + CompactQueryAnalysisTestCase { + description: "larger later queries keep request order", + run: larger_later_queries_keep_request_order, + expected_success: true, + }, ]; for test_case in test_cases { diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/compile/test_fresh_process_compile_cache_performance.py b/tests/e2e/src/sqlbuild/cli/commands/main/compile/test_fresh_process_compile_cache_performance.py index d6029a242e..8ec922c485 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/compile/test_fresh_process_compile_cache_performance.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/compile/test_fresh_process_compile_cache_performance.py @@ -30,6 +30,7 @@ # Writing the cache during a cold compile must stay a small same-runner cost over --no-cache. _MAX_CACHE_WRITE_CPU_OVERHEAD_RATIO: float = 1.10 _MAX_CACHE_WRITE_WALL_OVERHEAD_RATIO: float = 1.15 +_MAX_LARGE_CACHE_WRITE_CPU_OVERHEAD_RATIO: float = 1.15 @pytest.mark.performance @@ -119,7 +120,7 @@ expected_edit_max_wall_seconds=37.5, expected_max_warm_to_cold_ratio=_MAX_WARM_TO_COLD_RATIO, expected_max_edit_to_cold_ratio=_MAX_EDIT_TO_COLD_RATIO, - expected_max_cache_write_cpu_overhead_ratio=_MAX_CACHE_WRITE_CPU_OVERHEAD_RATIO, + expected_max_cache_write_cpu_overhead_ratio=_MAX_LARGE_CACHE_WRITE_CPU_OVERHEAD_RATIO, expected_max_cache_write_wall_overhead_ratio=_MAX_CACHE_WRITE_WALL_OVERHEAD_RATIO, expected_max_rss_bytes=2 * _GIB, expected_max_cache_bytes=320 * _MIB,