From c2fa7f66061f0f16da269b22b78516cb6680a440 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Fri, 9 Oct 2026 03:36:18 +0100 Subject: [PATCH 1/2] feat(native): enrich analysis inputs natively behind the preview engine The native model analysis session now re-analyses models with their known input shapes itself, deferring to Python only where Python's CTE fact recovery reads the parsed query. It also proves dynamic pivot contracts natively, inside the session and for models it does not analyse, deferring on parse errors and non-ASCII names. --- .../analysis_session/_helpers/dict_walk.rs | 156 ++++ .../_helpers/dynamic_pivot.rs | 675 ++++++++++++++++++ .../analysis_session/_helpers/enrichment.rs | 338 +++++++++ .../assembly/analysis_session/_helpers/mod.rs | 3 + .../analysis_session/_helpers/session.rs | 95 ++- .../assembly/analysis_session/constants.rs | 57 ++ .../src/assembly/analysis_session/main/mod.rs | 1 + .../main/prove_dynamic_contract.rs | 26 + .../src/assembly/analysis_session/models.rs | 52 ++ .../analysis_session/tests/helpers.rs | 91 ++- .../tests/test_analysis_session.rs | 137 +++- .../analysis_session/tests/test_types.rs | 11 + .../_helpers/analysis_session/session.rs | 104 ++- .../analysis_corpus/analysis_session.py | 47 ++ src/sqlbuild/_native.pyi | 37 + .../_helpers/dynamic_contract.py | 52 ++ .../analysis_session/_helpers/session_rows.py | 62 ++ .../classes/native_model_analysis.py | 27 +- .../compiler/analysis_session/constants.py | 3 + .../main/_prove_native_dynamic_contract.py | 31 + .../compiler/analysis_session/types.py | 7 +- .../_helpers/native_stages/assembly.py | 15 +- .../compile/classes/python_model_analysis.py | 34 +- .../test_compiler_differential.py | 11 +- .../compiler/analysis_session/_test_types.py | 2 + .../compiler/analysis_session/helpers.py | 83 ++- .../test_native_model_analysis.py | 8 +- 27 files changed, 2121 insertions(+), 44 deletions(-) create mode 100644 crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/dict_walk.rs create mode 100644 crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/dynamic_pivot.rs create mode 100644 crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/enrichment.rs create mode 100644 crates/sqlbuild-analysis/src/assembly/analysis_session/main/prove_dynamic_contract.rs create mode 100644 src/sqlbuild/compiler/analysis_session/_helpers/dynamic_contract.py create mode 100644 src/sqlbuild/compiler/analysis_session/main/_prove_native_dynamic_contract.py diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/dict_walk.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/dict_walk.rs new file mode 100644 index 0000000000..d11f30c0ae --- /dev/null +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/dict_walk.rs @@ -0,0 +1,156 @@ +//! Python's dict walk over a parsed query's `to_dict()` payload, on the same serde values. + +use serde_json::{Map, Value}; + +use crate::assembly::analysis_session::constants::{ + COLUMN_AST_KIND, DEPENDENCY_FUNCTION_NAMES, FUNCTION_AST_KIND, RENDERED_TYPE_NAMES, + STAR_AST_KIND, TABLE_AST_KIND, +}; + +/// Python's `bool(value)`. +pub(crate) fn truthy(value: Option<&Value>) -> bool { + match value { + Some(Value::Bool(flag)) => *flag, + Some(Value::Null) | None => false, + Some(Value::String(text)) => !text.is_empty(), + Some(Value::Array(items)) => !items.is_empty(), + Some(Value::Object(items)) => !items.is_empty(), + Some(Value::Number(number)) => number.as_f64() != Some(0.0), + } +} + +/// The single key of a one-key object, or "" for anything else. +pub(crate) fn node_key(node: Option<&Value>) -> &str { + let Some(object) = node + .and_then(Value::as_object) + .filter(|object| object.len() == 1) + else { + return ""; + }; + object.keys().next().map_or("", String::as_str) +} + +/// The payload under a node's single key. +pub(crate) fn payload(node: Option<&Value>) -> Option<&Map> { + let key: &str = node_key(node); + node.and_then(|node| node.get(key)) + .and_then(Value::as_object) +} + +/// Python's `_dict_list`: the objects of a list, or nothing. +pub(crate) fn dict_list(value: Option<&Value>) -> Vec<&Value> { + match value.and_then(Value::as_array) { + Some(items) => objects(items), + None => Vec::new(), + } +} + +fn objects(items: &[Value]) -> Vec<&Value> { + items.iter().filter(|item| item.is_object()).collect() +} + +/// Python's `_nested`: follow object keys. +pub(crate) fn nested<'a>(node: Option<&'a Value>, keys: &[&str]) -> Option<&'a Value> { + let mut current: Option<&Value> = node; + for key in keys { + current = current + .and_then(Value::as_object) + .and_then(|object| object.get(*key)); + } + current +} + +/// Python's `_identifier_name`. +pub(crate) fn identifier_name(node: Option<&Value>) -> Option<&str> { + match node { + Some(Value::String(name)) => Some(name), + Some(Value::Object(object)) => object + .get("name") + .and_then(Value::as_str) + .filter(|name| !name.is_empty()), + _ => None, + } +} + +/// Python's `_column_name`. +pub(crate) fn column_name(node: Option<&Value>) -> Option<&str> { + if node_key(node) != COLUMN_AST_KIND { + return None; + } + identifier_name(payload(node)?.get("name")) +} + +/// Python's `_relation_name`: a table, a column or a dependency call's first argument. +pub(crate) fn relation_name(node: Option<&Value>) -> Option<&str> { + let key: &str = node_key(node); + let payload: &Map = payload(node)?; + if key == TABLE_AST_KIND || key == COLUMN_AST_KIND { + return identifier_name(payload.get("name")); + } + let called: &str = payload + .get("name") + .and_then(Value::as_str) + .unwrap_or_default(); + if key == FUNCTION_AST_KIND && DEPENDENCY_FUNCTION_NAMES.contains(&casefold(called).as_str()) { + return column_name(dict_list(payload.get("args")).first().copied()); + } + None +} + +/// Python's `_is_single_wildcard`. +pub(crate) fn is_single_wildcard(value: Option<&Value>) -> bool { + let expressions: Vec<&Value> = dict_list(value); + let [only] = expressions.as_slice() else { + return false; + }; + if node_key(Some(only)) != STAR_AST_KIND { + return false; + } + let Some(star) = only.get(STAR_AST_KIND).and_then(Value::as_object) else { + return false; + }; + !["except", "replace", "rename", "table"] + .iter() + .any(|key| truthy(star.get(*key))) +} + +/// Python's `_render_type` of a cast target. +pub(crate) fn render_type(node: Option<&Value>) -> Option { + let node: &Map = node?.as_object()?; + let raw_type: &str = node + .get("data_type") + .and_then(Value::as_str) + .filter(|raw| !raw.is_empty())?; + let folded: String = casefold(raw_type); + let rendered: String = RENDERED_TYPE_NAMES + .iter() + .find(|(name, _)| *name == folded) + .map_or_else( + || raw_type.replace('_', " ").to_ascii_uppercase(), + |(_, rendered)| (*rendered).to_owned(), + ); + if let Some(length) = python_int(node.get("length")) { + return Some(format!("{rendered}({length})")); + } + let Some(precision) = python_int(node.get("precision")) else { + return Some(rendered); + }; + Some(match python_int(node.get("scale")) { + Some(scale) => format!("{rendered}({precision}, {scale})"), + None => format!("{rendered}({precision})"), + }) +} + +/// Python's `str(value)` for a value `isinstance(value, int)` accepts. +fn python_int(value: Option<&Value>) -> Option { + match value? { + Value::Bool(flag) => Some(if *flag { "True" } else { "False" }.to_owned()), + Value::Number(number) if number.is_i64() || number.is_u64() => Some(number.to_string()), + _ => None, + } +} + +/// Python's `str.casefold` for the ASCII text a pivot proof compares. +pub(crate) fn casefold(text: &str) -> String { + text.to_ascii_lowercase() +} diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/dynamic_pivot.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/dynamic_pivot.rs new file mode 100644 index 0000000000..78f0b3bdba --- /dev/null +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/dynamic_pivot.rs @@ -0,0 +1,675 @@ +//! Python's `analyze_dynamic_column_contract` over the crate's parse of the model SQL. + +use std::collections::{HashMap, HashSet}; + +use polyglot_sql::{ComplexityGuardOptions, Dialect, Expression, ParseOptions}; +use serde_json::{Map, Value, json}; +use sqlbuild_core::panics::main::catch_compiler_panic::catch_compiler_panic; + +use crate::assembly::analysis_session::_helpers::dict_walk::{ + casefold, column_name, dict_list, identifier_name, is_single_wildcard, nested, node_key, + payload, relation_name, render_type, truthy, +}; +use crate::assembly::analysis_session::constants::{ + ALIAS_AST_KIND, CAST_AST_KIND, CTE_AST_KIND, DUCKDB_DIALECT, DYNAMIC_VALUE_SOURCE_KINDS, + MAX_FUNCTION_CALL_DEPTH, MOTHERDUCK_DIALECT, PIVOT_AST_KIND, QUERY_AST_KINDS, SELECT_AST_KIND, + SIMPLIFIED_PIVOT_DIALECTS, SUPPORTED_PIVOT_DIALECTS, TYPE_PRESERVING_AGGREGATES, + UNKNOWN_NULLABILITY, UNKNOWN_TYPE, +}; +use crate::assembly::analysis_session::models::{ + ColumnFact, ContractProof, DynamicFamily, PivotOutcome, +}; +use crate::assembly::analysis_session::types::{Pairs, Shapes}; + +/// Stands in for Python's `{}`: neither has a node key. +static MISSING_EXPRESSION: Value = Value::Null; + +/// The relation facts and declared families a pivot proof reads. +pub(crate) struct PivotFacts<'a> { + pub(crate) dialect: &'a str, + pub(crate) column_types: &'a Shapes, + pub(crate) authoritative_types: &'a Shapes, + pub(crate) column_nullability: &'a Shapes, + pub(crate) families_by_table: &'a [(String, Vec)], +} + +/// One CTE body by folded name; None for a CTE with column aliases. +type CteMap<'a> = HashMap>; + +/// The resolved output boundary: a pivot payload or a passthrough relation. +enum Boundary<'a> { + Pivot(&'a Map), + Passthrough(&'a str), +} + +/// The proof Python computes, or a deferral where its result is not reproduced exactly. +pub(crate) fn pivot_outcome( + sql: &str, + families: &[DynamicFamily], + facts: &PivotFacts<'_>, +) -> PivotOutcome { + if families.is_empty() { + return PivotOutcome::Absent; + } + let folded_dialect: String = casefold(facts.dialect); + if !SUPPORTED_PIVOT_DIALECTS.contains(&folded_dialect.as_str()) { + return failure(format!( + "adapter dialect '{}' does not support compiler-proven dynamic pivots", + facts.dialect + )); + } + if !ascii_inputs(sql, families, facts) { + return PivotOutcome::Deferred; + } + let parse_dialect: &str = if folded_dialect == MOTHERDUCK_DIALECT { + DUCKDB_DIALECT + } else { + folded_dialect.as_str() + }; + let outcome: Result = catch_compiler_panic(|| { + Ok(match parsed(sql, parse_dialect)? { + Some(root) => proof(&root, &folded_dialect, families, facts), + None => PivotOutcome::Deferred, + }) + }); + match outcome { + Ok(outcome) => outcome, + Err(_) => PivotOutcome::Deferred, + } +} + +/// Python's `parse_one(...).to_dict()`, or None where Python reports a parse error. +fn parsed(sql: &str, dialect: &str) -> Result, String> { + let guard: ComplexityGuardOptions = + serde_json::from_value(json!({"maxFunctionCallDepth": MAX_FUNCTION_CALL_DEPTH})) + .map_err(|error| error.to_string())?; + let options: ParseOptions = ParseOptions { + complexity_guard: Some(guard), + }; + let parsed: Result, _> = Dialect::get_by_name(dialect) + .ok_or_else(|| format!("unknown dialect {dialect}"))? + .parse_with_options(sql, &options); + let expressions: Vec = match parsed { + Ok(expressions) => expressions, + Err(_) => return Ok(None), + }; + let [expression] = expressions.as_slice() else { + return Ok(None); + }; + serde_json::to_value(expression) + .map(Some) + .map_err(|error| error.to_string()) +} + +/// Casefolding equals Python's only for ASCII text. +fn ascii_inputs(sql: &str, families: &[DynamicFamily], facts: &PivotFacts<'_>) -> bool { + let shapes = [ + facts.column_types, + facts.authoritative_types, + facts.column_nullability, + ]; + sql.is_ascii() + && families.iter().all(ascii_family) + && facts.families_by_table.iter().all(ascii_families) + && shapes.into_iter().flatten().all(ascii_shape) +} + +fn ascii_families((name, families): &(String, Vec)) -> bool { + name.is_ascii() && families.iter().all(ascii_family) +} + +fn ascii_shape((name, shape): &(String, Pairs)) -> bool { + name.is_ascii() && shape.iter().all(|(column, _)| column.is_ascii()) +} + +fn proof( + root: &Value, + dialect: &str, + families: &[DynamicFamily], + facts: &PivotFacts<'_>, +) -> PivotOutcome { + if !root.is_object() { + return failure("dynamic pivot SQL did not produce a structured query".to_owned()); + } + let ctes: CteMap<'_> = collect_ctes(root); + match resolve_boundary(Some(root), &ctes, &HashSet::new()) { + None => failure( + "output must be a wildcard projection from one supported dynamic pivot or a proven \ + dynamic-family passthrough" + .to_owned(), + ), + Some(Boundary::Passthrough(name)) => passthrough(name, families, facts), + Some(Boundary::Pivot(pivot)) => pivot_proof( + pivot, + &PivotContext { + dialect, + ctes: &ctes, + facts, + bare: node_key(Some(root)) == PIVOT_AST_KIND, + }, + families, + ), + } +} + +struct PivotContext<'a, 'b> { + dialect: &'a str, + ctes: &'a CteMap<'b>, + facts: &'a PivotFacts<'a>, + bare: bool, +} + +fn passthrough(name: &str, families: &[DynamicFamily], facts: &PivotFacts<'_>) -> PivotOutcome { + let upstream: Option<&Vec> = casefold_lookup(facts.families_by_table, name); + if !upstream.is_some_and(|upstream| families_equal(families, upstream)) { + return failure(format!( + "dynamic-family passthrough from '{name}' must redeclare the upstream families exactly" + )); + } + let Some(columns) = external_schema(name, facts.column_types, facts.column_nullability) else { + return failure(format!( + "dynamic-family passthrough source '{name}' has no authoritative schema" + )); + }; + PivotOutcome::Proof(ContractProof { + output_proven: true, + fixed_columns: columns, + families: families + .iter() + .map(|family| (family.name.clone(), Some(family.data_type.clone()))) + .collect(), + input_relations: vec![name.to_owned()], + failure_reason: None, + bare_dynamic_pivot: false, + }) +} + +fn pivot_proof( + pivot: &Map, + context: &PivotContext<'_, '_>, + families: &[DynamicFamily], +) -> PivotOutcome { + if truthy(pivot.get("unpivot")) { + return failure("UNPIVOT does not produce a dynamic column family".to_owned()); + } + let (pivot_columns, aggregates, dynamic, columns_valid) = components(pivot, context.dialect); + if !dynamic { + return failure( + "static pivot values must use ordinary exact column declarations".to_owned(), + ); + } + let [pivot_column] = pivot_columns.as_slice() else { + return failure( + "dynamic column contracts currently require exactly one pivot column".to_owned(), + ); + }; + if !columns_valid { + return failure("dynamic pivot columns must be direct column references".to_owned()); + } + let source: Option<&Value> = pivot.get("this"); + let schema_facts = SchemaFacts { + ctes: context.ctes, + types: context.facts.authoritative_types, + nullability: context.facts.column_nullability, + }; + let Some(source_columns) = relation_schema(source, &schema_facts, &HashSet::new()) else { + return failure( + "dynamic pivot input does not have an authoritative, explicit schema".to_owned(), + ); + }; + let Some(input_relation) = external_relation_name(source, context.ctes, &HashSet::new()) else { + return failure("dynamic pivot input relation could not be resolved".to_owned()); + }; + let mut matched: HashSet = HashSet::new(); + let mut family_proofs: Vec<(String, Option)> = Vec::new(); + for family in families { + if casefold(&family.pivot_column) != casefold(pivot_column) { + return failure(format!( + "dynamic family '{}' declares pivot_column '{}' but the output pivot uses '{}'", + family.name, family.pivot_column, pivot_column + )); + } + let matches: Vec<(usize, Option)> = aggregates + .iter() + .enumerate() + .filter_map(|(index, aggregate)| { + let (name, value, inferred) = aggregate_fact(aggregate, &source_columns); + let same_value: bool = + value.is_some_and(|value| casefold(value) == casefold(&family.value_column)); + (casefold(&name) == casefold(&family.aggregate) && same_value) + .then_some((index, inferred)) + }) + .collect(); + let [(index, inferred)] = matches.as_slice() else { + return failure(format!( + "dynamic family '{}' must match exactly one {}({}) pivot aggregate", + family.name, family.aggregate, family.value_column + )); + }; + if !matched.insert(*index) { + return failure( + "one pivot aggregate cannot satisfy multiple dynamic families".to_owned(), + ); + } + family_proofs.push((family.name.clone(), inferred.clone())); + } + if matched.len() != aggregates.len() { + return failure( + "every dynamic pivot aggregate must have one declared column family".to_owned(), + ); + } + let mut excluded: HashSet = HashSet::from([casefold(pivot_column)]); + excluded.extend(families.iter().map(|family| casefold(&family.value_column))); + let fixed_names: Vec = fixed_names(pivot, context.dialect, &excluded, &source_columns); + let by_name: HashMap = source_columns + .iter() + .map(|column| (casefold(&column.name), column)) + .collect(); + let mut fixed_columns: Vec = Vec::with_capacity(fixed_names.len()); + for name in fixed_names { + let Some(source) = by_name.get(&casefold(&name)) else { + return failure( + "dynamic pivot grouping columns are not present in the pivot input".to_owned(), + ); + }; + fixed_columns.push(ColumnFact { + name, + data_type: source.data_type.clone(), + nullability: source.nullability.clone(), + }); + } + PivotOutcome::Proof(ContractProof { + output_proven: true, + fixed_columns, + families: family_proofs, + input_relations: vec![input_relation.to_owned()], + failure_reason: None, + bare_dynamic_pivot: context.bare, + }) +} + +/// The fixed output names: DuckDB's GROUP BY columns, otherwise the unpivoted input columns. +fn fixed_names( + pivot: &Map, + dialect: &str, + excluded: &HashSet, + source_columns: &[ColumnFact], +) -> Vec { + let group: Option<&Value> = pivot.get("group"); + if SIMPLIFIED_PIVOT_DIALECTS.contains(&dialect) && group.is_some_and(Value::is_object) { + return dict_list(nested(group, &["group", "expressions"])) + .into_iter() + .filter_map(|expression| column_name(Some(expression)).map(str::to_owned)) + .collect(); + } + source_columns + .iter() + .filter(|column| !excluded.contains(&casefold(&column.name))) + .map(|column| column.name.clone()) + .collect() +} + +/// Python's `_pivot_components`: pivot columns, aggregates, dynamic values, columns valid. +fn components<'a>( + pivot: &'a Map, + dialect: &str, +) -> (Vec<&'a str>, Vec<&'a Value>, bool, bool) { + if SIMPLIFIED_PIVOT_DIALECTS.contains(&dialect) && truthy(pivot.get("using")) { + let expressions: Vec<&Value> = dict_list(pivot.get("expressions")); + let columns: Vec<&str> = expressions + .iter() + .filter_map(|expression| column_name(Some(expression))) + .collect(); + let valid: bool = columns.len() == expressions.len(); + return ( + columns, + dict_list(pivot.get("using")), + !truthy(pivot.get("fields")), + valid, + ); + } + let fields: Vec<&Value> = dict_list(pivot.get("fields")); + let columns: Vec<&str> = fields + .iter() + .filter_map(|field| column_name(nested(Some(field), &["in", "this"]))) + .collect(); + let dynamic: bool = fields.iter().any(|field| has_dynamic_values(field)); + let valid: bool = columns.len() == fields.len(); + (columns, dict_list(pivot.get("expressions")), dynamic, valid) +} + +/// Python's `_has_dynamic_value_source` for one pivot field. +fn has_dynamic_values(field: &Value) -> bool { + dict_list(nested(Some(field), &["in", "expressions"])) + .into_iter() + .any(|value| DYNAMIC_VALUE_SOURCE_KINDS.contains(&node_key(Some(value)))) +} + +/// Python's `_aggregate_fact`: aggregate name, value column and type-preserving inferred type. +fn aggregate_fact<'a>( + aggregate: &'a Value, + source_columns: &[ColumnFact], +) -> (String, Option<&'a str>, Option) { + let key: &str = node_key(Some(aggregate)); + let mut name: String = key.to_ascii_uppercase(); + let Some(payload) = aggregate.get(key).and_then(Value::as_object) else { + return (name, None, None); + }; + if let Some(authored) = payload + .get("name") + .and_then(Value::as_str) + .filter(|authored| !authored.is_empty()) + { + name = authored.to_ascii_uppercase(); + } + let preserving: bool = TYPE_PRESERVING_AGGREGATES.contains(&name.as_str()); + let value: Option<&Value> = payload.get("this"); + let mut value_name: Option<&str> = column_name(value); + let mut inferred: Option = None; + let cast: Option<&Map> = value + .filter(|value| node_key(Some(value)) == CAST_AST_KIND) + .and_then(|value| value.get(CAST_AST_KIND)) + .and_then(Value::as_object); + if let Some(cast) = cast { + value_name = column_name(cast.get("this")); + if preserving { + inferred = render_type(cast.get("to")); + } + } + if inferred.is_none() && preserving { + let by_name: HashMap = source_columns + .iter() + .map(|column| (casefold(&column.name), column)) + .collect(); + inferred = value_name + .and_then(|value| by_name.get(&casefold(value))) + .and_then(|column| column.data_type.clone()); + } + (name, value_name, inferred) +} + +fn resolve_boundary<'a>( + node: Option<&'a Value>, + ctes: &CteMap<'a>, + seen: &HashSet, +) -> Option> { + let key: &str = node_key(node); + let payload: &Map = payload(node)?; + if key == PIVOT_AST_KIND { + return Some(Boundary::Pivot(payload)); + } + if QUERY_AST_KINDS.contains(&key) { + return resolve_boundary(payload.get("this"), ctes, seen); + } + if key != SELECT_AST_KIND + || truthy(payload.get("joins")) + || !is_single_wildcard(payload.get("expressions")) + { + return None; + } + let relations: Vec<&Value> = + dict_list(payload.get("from").and_then(|from| from.get("expressions"))); + let [relation] = relations.as_slice() else { + return None; + }; + if node_key(Some(relation)) == PIVOT_AST_KIND { + return relation + .get(PIVOT_AST_KIND) + .and_then(Value::as_object) + .map(Boundary::Pivot); + } + let name: &str = relation_name(Some(relation))?; + let folded: String = casefold(name); + let Some(cte) = ctes.get(&folded) else { + return Some(Boundary::Passthrough(name)); + }; + let cte: &Value = (*cte)?; + if seen.contains(&folded) { + return None; + } + let mut next_seen: HashSet = seen.clone(); + next_seen.insert(folded); + resolve_boundary(Some(cte), ctes, &next_seen) +} + +/// The relation facts schema lookups read. +struct SchemaFacts<'a, 'b> { + ctes: &'a CteMap<'b>, + types: &'a Shapes, + nullability: &'a Shapes, +} + +fn relation_schema( + node: Option<&Value>, + facts: &SchemaFacts<'_, '_>, + seen: &HashSet, +) -> Option> { + let name: &str = relation_name(node)?; + let folded: String = casefold(name); + let Some(cte) = facts.ctes.get(&folded) else { + return external_schema(name, facts.types, facts.nullability); + }; + let cte: &Value = (*cte)?; + if seen.contains(&folded) { + return None; + } + let mut next_seen: HashSet = seen.clone(); + next_seen.insert(folded); + select_schema(Some(cte), facts, &next_seen) +} + +fn external_relation_name<'a>( + node: Option<&'a Value>, + ctes: &CteMap<'a>, + seen: &HashSet, +) -> Option<&'a str> { + let key: &str = node_key(node); + if let Some(payload) = payload(node) { + if QUERY_AST_KINDS.contains(&key) { + return external_relation_name(payload.get("this"), ctes, seen); + } + if key == SELECT_AST_KIND { + if truthy(payload.get("joins")) { + return None; + } + let relations: Vec<&Value> = + dict_list(payload.get("from").and_then(|from| from.get("expressions"))); + let [relation] = relations.as_slice() else { + return None; + }; + return external_relation_name(Some(relation), ctes, seen); + } + } + let name: &str = relation_name(node)?; + let folded: String = casefold(name); + let Some(cte) = ctes.get(&folded) else { + return Some(name); + }; + let cte: &Value = (*cte)?; + if seen.contains(&folded) { + return None; + } + let mut next_seen: HashSet = seen.clone(); + next_seen.insert(folded); + external_relation_name(Some(cte), ctes, &next_seen) +} + +fn select_schema( + node: Option<&Value>, + facts: &SchemaFacts<'_, '_>, + seen: &HashSet, +) -> Option> { + let key: &str = node_key(node); + let payload: &Map = payload(node)?; + if QUERY_AST_KINDS.contains(&key) { + return select_schema(payload.get("this"), facts, seen); + } + if key != SELECT_AST_KIND || truthy(payload.get("joins")) { + return None; + } + let relations: Vec<&Value> = + dict_list(payload.get("from").and_then(|from| from.get("expressions"))); + let [relation] = relations.as_slice() else { + return None; + }; + let source: Vec = relation_schema(Some(relation), facts, seen)?; + let expressions: Option<&Value> = payload.get("expressions"); + if is_single_wildcard(expressions) { + return Some(source); + } + let by_name: HashMap = source + .iter() + .map(|column| (casefold(&column.name), column)) + .collect(); + let mut output: Vec = Vec::new(); + for expression in expressions?.as_array()? { + if !expression.is_object() { + return None; + } + output.push(projected_column(expression, &by_name)?); + } + Some(output) +} + +/// One projection of a CTE body over its single source relation. +fn projected_column( + projection: &Value, + by_name: &HashMap, +) -> Option { + let (output_name, expression) = projection_parts(projection); + let (source_name, cast_type): (Option<&str>, Option) = + if node_key(Some(expression)) == CAST_AST_KIND { + let cast: &Map = expression.get(CAST_AST_KIND)?.as_object()?; + (column_name(cast.get("this")), render_type(cast.get("to"))) + } else { + (column_name(Some(expression)), None) + }; + let source: &ColumnFact = by_name.get(&casefold(source_name?))?; + Some(ColumnFact { + name: output_name?.to_owned(), + data_type: cast_type.or_else(|| source.data_type.clone()), + nullability: source.nullability.clone(), + }) +} + +/// Python's `_projection`: the output name and the aliased expression, `{}` when absent. +fn projection_parts(node: &Value) -> (Option<&str>, &Value) { + if node_key(Some(node)) != ALIAS_AST_KIND { + return (column_name(Some(node)), node); + } + let Some(alias) = node.get(ALIAS_AST_KIND).and_then(Value::as_object) else { + return (None, node); + }; + let expression: &Value = alias + .get("this") + .filter(|expression| expression.is_object()) + .unwrap_or(&MISSING_EXPRESSION); + (identifier_name(alias.get("alias")), expression) +} + +fn collect_ctes(root: &Value) -> CteMap<'_> { + let mut ctes: CteMap<'_> = HashMap::new(); + let Some(with) = payload(Some(root)).and_then(|payload| payload.get("with")) else { + return ctes; + }; + if !with.is_object() { + return ctes; + } + for cte in dict_list(with.get("ctes")) { + let body: &Value = if node_key(Some(cte)) == CTE_AST_KIND { + match cte.get(CTE_AST_KIND) { + Some(body) => body, + None => continue, + } + } else { + cte + }; + let Some(body) = body.as_object() else { + continue; + }; + let Some(name) = identifier_name(body.get("alias")) else { + continue; + }; + let query: Option<&Value> = body.get("this"); + if truthy(body.get("columns")) { + ctes.insert(casefold(name), None); + } else if let Some(query) = query.filter(|query| query.is_object()) { + ctes.insert(casefold(name), Some(query)); + } + } + ctes +} + +/// Python's `_external_schema`: the relation's typed columns with their nullability. +fn external_schema(name: &str, types: &Shapes, nullability: &Shapes) -> Option> { + let types = casefold_lookup(types, name)?; + let mut by_name: HashMap = HashMap::new(); + for (column, value) in casefold_lookup(nullability, name).into_iter().flatten() { + by_name.insert(casefold(column), value); + } + Some( + types + .iter() + .map(|(column, data_type)| ColumnFact { + name: column.clone(), + data_type: (data_type != UNKNOWN_TYPE).then(|| data_type.clone()), + nullability: by_name + .get(&casefold(column)) + .map_or(UNKNOWN_NULLABILITY, |value| value) + .to_owned(), + }) + .collect(), + ) +} + +/// Python's `_casefold_lookup`: the first entry whose folded name matches. +fn casefold_lookup<'a, T>(entries: &'a [(String, T)], key: &str) -> Option<&'a T> { + let requested: String = casefold(key); + entries + .iter() + .find(|(candidate, _)| casefold(candidate) == requested) + .map(|(_, value)| value) +} + +fn families_equal(left: &[DynamicFamily], right: &[DynamicFamily]) -> bool { + left.len() == right.len() + && left + .iter() + .zip(right) + .all(|(left, right)| normalized(left) == normalized(right)) +} + +fn failure(reason: String) -> PivotOutcome { + PivotOutcome::Proof(ContractProof { + output_proven: false, + fixed_columns: Vec::new(), + families: Vec::new(), + input_relations: Vec::new(), + failure_reason: Some(reason), + bare_dynamic_pivot: false, + }) +} + +fn ascii_family(family: &DynamicFamily) -> bool { + [ + &family.name, + &family.pivot_column, + &family.value_column, + &family.aggregate, + &family.data_type, + ] + .iter() + .all(|text| text.is_ascii()) +} + +/// Python's `_families_equal` key for one family. +fn normalized(family: &DynamicFamily) -> (String, String, String, String, String, Option<&str>) { + ( + casefold(&family.name), + casefold(&family.pivot_column), + casefold(&family.value_column), + casefold(&family.aggregate), + casefold(&family.data_type), + family.name_pattern.as_deref(), + ) +} diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/enrichment.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/enrichment.rs new file mode 100644 index 0000000000..2e176ff81e --- /dev/null +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/enrichment.rs @@ -0,0 +1,338 @@ +//! Python's re-analysis of a model with its known input shapes, over the crate's query facts. + +use std::collections::{HashMap, HashSet}; + +use polyglot_sql::{AnalyzeQueryOptions, analyze_query}; +use serde_json::{Map, Value, json}; +use sqlbuild_core::panics::main::catch_compiler_panic::catch_compiler_panic; + +use crate::assembly::analysis_session::_helpers::dict_walk::truthy; +use crate::assembly::analysis_session::constants::{ + CAST_TRANSFORM, CONFIDENCE_HIGH, CONFIDENCE_MEDIUM, CONFIDENCE_UNKNOWN, FILTER_CONTEXT, + MAX_FUNCTION_CALL_DEPTH, NULL_KEYWORD, RESOLVED_SOURCE_CONFIDENCE, SELECT_SHAPE, + SET_OPERATION_SHAPE, TRANSFORM_AGGREGATION, TRANSFORM_CAST, TRANSFORM_CONSTANT, + TRANSFORM_DIRECT, TRANSFORM_EXPRESSION, TRANSFORM_STAR, UNKNOWN_NULLABILITY, UNKNOWN_TYPE, + WILDCARD, +}; +use crate::assembly::analysis_session::models::{ColumnFact, LineageRow, ModelRequest}; +use crate::assembly::analysis_session::types::{Pairs, Shapes}; +use crate::semantic_validation::main::normalize::normalize_analysis_sql; +use crate::semantic_validation::models::NormalizationInput; + +/// One model's enrichment inputs: the model, its known input shapes and the profile facts. +pub(crate) struct EnrichmentInput<'a> { + pub(crate) model: &'a ModelRequest, + pub(crate) input_schemas: &'a Shapes, + pub(crate) dialect: &'a str, + pub(crate) function_return_types: &'a Pairs, +} + +/// Python's successful re-analysis: columns, plain lineage facts and the star flag. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct Enrichment { + pub(crate) columns: Vec, + pub(crate) lineage: Vec, + pub(crate) has_star: bool, +} + +/// One projection's upstream sources and confidence code. +type Upstream = (Vec<(String, String, String)>, u8); + +/// Python's compact re-analysis, or None where Python would take another path. +pub(crate) fn enrichment(input: &EnrichmentInput<'_>) -> Option { + let enrichment: Result, String> = catch_compiler_panic(|| analysed(input)); + enrichment.unwrap_or_default() +} + +/// The re-analysis; Ok(None) where Python takes another path, Err for an unexpected failure. +fn analysed(input: &EnrichmentInput<'_>) -> Result, String> { + let normalized: Result = normalize_analysis_sql(NormalizationInput { + sql: input.model.query_sql.clone(), + dialect: input.dialect.to_owned(), + stubs: HashMap::new(), + placeholders: input.model.placeholders.iter().cloned().collect(), + }); + let cleaned: String = match normalized { + Ok(cleaned) => cleaned, + Err(_) => return Ok(None), + }; + let mut options: Map = Map::new(); + options.insert( + "dialect".to_owned(), + Value::String(input.dialect.to_owned()), + ); + if let Some(schema) = analysis_schema(input) { + options.insert("schema".to_owned(), schema); + } + options.insert( + "complexityGuard".to_owned(), + json!({"maxFunctionCallDepth": MAX_FUNCTION_CALL_DEPTH}), + ); + let options: AnalyzeQueryOptions = + serde_json::from_value(Value::Object(options)).map_err(|error| error.to_string())?; + let analysis = match analyze_query(&cleaned, options) { + Ok(analysis) => analysis, + Err(_) => return Ok(None), + }; + let analysis: Value = serde_json::to_value(analysis).map_err(|error| error.to_string())?; + Ok(projected(input, &analysis)) +} + +/// Python's `_compact_analysis_schema` over the input shapes, nullability unknown. +fn analysis_schema(input: &EnrichmentInput<'_>) -> Option { + let shapes: HashMap<&str, &Pairs> = input + .input_schemas + .iter() + .map(|(name, shape)| (name.as_str(), shape)) + .collect(); + let mut names: Vec<&str> = input + .model + .references + .iter() + .map(|reference| reference.analysis_name.as_str()) + .collect::>() + .into_iter() + .collect(); + names.sort_unstable(); + let mut tables: Vec = Vec::new(); + for name in names { + let Some(shape) = shapes.get(name).filter(|shape| !shape.is_empty()) else { + continue; + }; + let columns: Vec = shape + .iter() + .map(|(column, data_type)| json!({"name": column, "type": data_type})) + .collect(); + tables.push(json!({"name": name, "columns": columns})); + } + (!tables.is_empty()).then(|| json!({"tables": tables})) +} + +/// Python's projection of `analyze_query` facts into columns and lineage. +fn projected(input: &EnrichmentInput<'_>, analysis: &Value) -> Option { + let projections: &Vec = analysis.get("projections")?.as_array()?; + let shape: Option<&str> = analysis.get("shape").and_then(Value::as_str); + if !matches!(shape, Some(SELECT_SHAPE | SET_OPERATION_SHAPE)) || !eligible(projections) { + return None; + } + if has_null_filter(analysis) || recovers_cte_facts(input, analysis) { + return None; + } + let resources: HashMap<&str, (&str, &str)> = input + .model + .lineage_references + .iter() + .map(|(name, kind, resource)| (name.as_str(), (kind.as_str(), resource.as_str()))) + .collect(); + let infer_nullability: bool = shape != Some(SET_OPERATION_SHAPE); + let mut columns: Vec = Vec::new(); + let mut lineage: Vec = Vec::new(); + for projection in projections { + let projection: &Map = projection.as_object()?; + if truthy(projection.get("isStar")) { + continue; + } + let name: &str = text(projection.get("name")); + if name.is_empty() || name == WILDCARD { + continue; + } + columns.push(ColumnFact { + name: name.to_owned(), + data_type: projection_type(projection, input.function_return_types)?, + nullability: projection_nullability(projection, infer_nullability).to_owned(), + }); + let (sources, confidence) = upstream(projection, &resources); + let transform: u8 = transform_code(projection, !sources.is_empty()); + lineage.push(LineageRow { + output_column: name.to_owned(), + confidence_code: if sources.is_empty() && transform != TRANSFORM_CONSTANT { + CONFIDENCE_UNKNOWN + } else { + confidence + }, + transform_code: transform, + sources, + }); + } + let has_star: bool = analysis + .get("starProjections") + .and_then(Value::as_array) + .is_some_and(|stars| !stars.is_empty()); + if has_star && input.model.references.len() == 1 { + let order: HashMap<&str, usize> = declared_order(input); + let fallback: usize = order.len(); + columns.sort_by_key(|column| order.get(column.name.as_str()).copied().unwrap_or(fallback)); + lineage.sort_by_key(|row| { + order + .get(row.output_column.as_str()) + .copied() + .unwrap_or(fallback) + }); + } + Some(Enrichment { + columns, + lineage, + has_star, + }) +} + +/// Whether Python's CTE pass-through recovery reads the parsed query, which stays Python's. +fn recovers_cte_facts(input: &EnrichmentInput<'_>, analysis: &Value) -> bool { + input.model.recover_cte_facts + && truthy(analysis.get("cteFacts")) + && !input.input_schemas.is_empty() +} + +/// Python's `_compact_analysis_is_eligible` for the projections. +fn eligible(projections: &[Value]) -> bool { + for projection in projections { + let Some(projection) = projection.as_object() else { + return false; + }; + let cast: bool = text(projection.get("transformKind")) == CAST_TRANSFORM; + let upstream_empty: bool = projection + .get("upstream") + .and_then(Value::as_array) + .is_none_or(Vec::is_empty); + if cast && upstream_empty { + return false; + } + } + true +} + +/// Whether a filter reads NULL, where Python parses the query for non-null outputs. +fn has_null_filter(analysis: &Value) -> bool { + let Some(uses) = analysis.get("columnUses").and_then(Value::as_array) else { + return false; + }; + uses.iter().filter_map(Value::as_object).any(|value| { + text(value.get("context")) == FILTER_CONTEXT + && text(value.get("expressionSql")) + .to_uppercase() + .contains(NULL_KEYWORD) + }) +} + +/// Python's `_compact_projection_type`; None where the function name is not ASCII. +fn projection_type( + projection: &Map, + function_return_types: &Pairs, +) -> Option> { + let cast_type: &str = text(projection.get("castType")); + if !cast_type.is_empty() && cast_type != UNKNOWN_TYPE { + return Some(Some(cast_type.to_owned())); + } + let function_name: Option<&str> = projection + .get("transformFunction") + .and_then(Value::as_object) + .and_then(|function| function.get("name")) + .and_then(Value::as_str); + if let Some(function_name) = function_name { + if !function_name.is_ascii() { + return None; + } + let upper: String = function_name.to_ascii_uppercase(); + if let Some((_, declared)) = function_return_types + .iter() + .find(|(name, _)| *name == upper) + { + return Some(Some(declared.clone())); + } + } + let type_hint: &str = text(projection.get("typeHint")); + if !type_hint.is_empty() && type_hint != UNKNOWN_TYPE { + return Some(Some(type_hint.to_owned())); + } + Some(None) +} + +fn projection_nullability(projection: &Map, infer: bool) -> &'static str { + if !infer { + return UNKNOWN_NULLABILITY; + } + match text(projection.get("nullability")) { + "non_null" => "non_null", + "nullable" => "nullable", + _ => UNKNOWN_NULLABILITY, + } +} + +/// Python's `_compact_lineage_upstream_columns`. +fn upstream(projection: &Map, resources: &HashMap<&str, (&str, &str)>) -> Upstream { + let Some(values) = projection.get("upstream").and_then(Value::as_array) else { + return (Vec::new(), CONFIDENCE_UNKNOWN); + }; + let mut sources: Vec<(String, String, String)> = Vec::new(); + let mut seen: HashSet<(String, String, String)> = HashSet::new(); + let mut confidence: u8 = CONFIDENCE_HIGH; + for value in values.iter().filter_map(Value::as_object) { + let column: &str = text(value.get("column")); + if column.is_empty() || column == WILDCARD { + continue; + } + let mut source_name: &str = text(value.get("sourceName")); + if source_name.is_empty() { + source_name = text(value.get("table")); + } + if source_name.is_empty() { + confidence = CONFIDENCE_UNKNOWN; + continue; + } + let Some((kind, resource)) = resources.get(source_name) else { + continue; + }; + let resolved: bool = text(value.get("confidence")) == RESOLVED_SOURCE_CONFIDENCE; + let aliased: bool = value.get("sourceAlias").is_some_and(Value::is_string); + if !resolved && !aliased { + confidence = CONFIDENCE_MEDIUM; + } + let key = ( + (*kind).to_owned(), + (*resource).to_owned(), + column.to_owned(), + ); + if seen.insert(key.clone()) { + sources.push(key); + } + } + sources.sort(); + (sources, confidence) +} + +/// Python's `_compact_transform_kind` as a compact code. +fn transform_code(projection: &Map, has_upstream: bool) -> u8 { + match text(projection.get("transformKind")) { + "star" => TRANSFORM_STAR, + "cast" => TRANSFORM_CAST, + "aggregation" => TRANSFORM_AGGREGATION, + "constant" => TRANSFORM_CONSTANT, + _ if !has_upstream => TRANSFORM_CONSTANT, + "direct" => TRANSFORM_DIRECT, + _ => TRANSFORM_EXPRESSION, + } +} + +/// The declared column order of the model's single reference, from its input shape. +fn declared_order<'a>(input: &EnrichmentInput<'a>) -> HashMap<&'a str, usize> { + let Some(reference) = input.model.references.first() else { + return HashMap::new(); + }; + let Some((_, shape)) = input + .input_schemas + .iter() + .find(|(name, _)| *name == reference.analysis_name) + else { + return HashMap::new(); + }; + let mut order: HashMap<&str, usize> = HashMap::new(); + for (index, (column, _)) in shape.iter().enumerate() { + order.entry(column.as_str()).or_insert(index); + } + order +} + +/// Python's `str(value or "")` for the string fields this projection reads. +fn text(value: Option<&Value>) -> &str { + value.and_then(Value::as_str).unwrap_or_default() +} diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/mod.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/mod.rs index f992aa5b95..1ff9c1dda6 100644 --- a/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/mod.rs +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/mod.rs @@ -1,5 +1,8 @@ pub(crate) mod catalog_state; pub(crate) mod compact_batch; +pub(crate) mod dict_walk; +pub(crate) mod dynamic_pivot; +pub(crate) mod enrichment; pub(crate) mod mappings; pub(crate) mod publication; pub(crate) mod schedule; diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/session.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/session.rs index de099cbc93..cf945f7c71 100644 --- a/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/session.rs +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/_helpers/session.rs @@ -2,10 +2,16 @@ use std::collections::HashMap; +use rayon::iter::{IntoParallelRefIterator, ParallelIterator}; + use crate::assembly::analysis_session::_helpers::catalog_state::SessionCatalog; use crate::assembly::analysis_session::_helpers::compact_batch::{ Batch, BatchMember, MemberAnalysis, MemberResult, }; +use crate::assembly::analysis_session::_helpers::dynamic_pivot::{PivotFacts, pivot_outcome}; +use crate::assembly::analysis_session::_helpers::enrichment::{ + Enrichment, EnrichmentInput, enrichment, +}; use crate::assembly::analysis_session::_helpers::mappings::{ ShapeTable, dict_from_pairs, same_keys, same_relations, }; @@ -16,7 +22,7 @@ use crate::assembly::analysis_session::constants::{ }; use crate::assembly::analysis_session::models::{ AnalysisSession, Awaiting, ColumnFact, Deferral, DeferredAnalysis, LineageFacts, ModelAnalysis, - ModelOutcome, ModelRequest, Phase, SessionOutcome, SessionRequest, SessionStep, + ModelOutcome, ModelRequest, Phase, PivotOutcome, SessionOutcome, SessionRequest, SessionStep, }; use crate::assembly::analysis_session::types::{Pairs, Shapes}; use crate::semantic_validation::models::ProjectCatalog; @@ -128,7 +134,12 @@ impl AnalysisSession { .get(&result.model) .ok_or("an enrichment answers no deferral")?; let outcome: &mut ModelOutcome = self.outcome_mut(result.model)?; - outcome.analysis = enriched_analysis(&outcome.analysis, result, pending); + outcome.analysis = enriched_analysis( + &outcome.analysis, + result, + pending, + LineageFacts::PythonEnrichment, + ); } self.phase = Phase::Finish(wave); } @@ -141,6 +152,7 @@ impl AnalysisSession { if !matches!(self.phase, Phase::Done) { return Err("the session has not finished".to_owned()); } + let dynamic_contracts: Vec = self.dynamic_contracts()?; let models: Vec = self .outcomes .into_iter() @@ -151,9 +163,32 @@ impl AnalysisSession { models, schema_additions, analysis_names, + dynamic_contracts, }) } + /// Python's dynamic pivot proof of each model over the relation facts before analysis. + fn dynamic_contracts(&self) -> Result, String> { + let models: &[ModelRequest] = &self.request.models; + if models.iter().all(|model| model.dynamic_families.is_empty()) { + return Ok(vec![PivotOutcome::Absent; models.len()]); + } + let facts = PivotFacts { + dialect: &self.request.dialect, + column_types: &self.request.column_types, + authoritative_types: &self.request.complete_schemas, + column_nullability: &self.request.column_nullability, + families_by_table: &self.request.dynamic_families_by_table, + }; + let pool = self.catalog.native.analysis_pool()?; + Ok(pool.install(|| { + models + .par_iter() + .map(|model| pivot_outcome(&model.pivot_sql, &model.dynamic_families, &facts)) + .collect() + })) + } + fn step(&mut self, deferrals: Vec) -> SessionStep { SessionStep { publications: std::mem::take(&mut self.publications), @@ -251,6 +286,7 @@ impl AnalysisSession { fn complete_wave(&mut self, wave: usize) -> Result { let mut deferrals: Vec = Vec::new(); let mut star_pending_by_model: HashMap = HashMap::new(); + let mut candidates: Vec<(usize, Shapes, bool)> = Vec::new(); for model in self.waves[wave].clone() { let request: &ModelRequest = &self.request.models[model]; let required: Vec = request.required_names.clone(); @@ -275,16 +311,53 @@ impl AnalysisSession { .flatten() .any(|column| column.data_type.is_none()); if !required.is_empty() && (untyped || star_pending) && !has_set_operation { + candidates.push((model, input_schemas, star_pending)); + } + } + let enrichments: Vec> = self.native_enrichments(&candidates)?; + for ((model, input_schemas, star_pending), native) in + candidates.into_iter().zip(enrichments) + { + let Some(native) = native else { star_pending_by_model.insert(model, star_pending); deferrals.push(Deferral::Enrichment { model, input_schemas, }); - } + continue; + }; + let outcome: &mut ModelOutcome = self.outcome_mut(model)?; + outcome.analysis = enriched_analysis( + &outcome.analysis, + native_answer(model, native.columns, native.has_star), + star_pending, + LineageFacts::NativeEnrichment(native.lineage), + ); } Ok((deferrals, star_pending_by_model)) } + /// Python's re-analysis of each candidate natively, None where Python must answer it. + fn native_enrichments( + &self, + candidates: &[(usize, Shapes, bool)], + ) -> Result>, String> { + if candidates.is_empty() { + return Ok(Vec::new()); + } + let inputs: Vec> = candidates + .iter() + .map(|(model, input_schemas, _)| EnrichmentInput { + model: &self.request.models[*model], + input_schemas, + dialect: &self.request.dialect, + function_return_types: &self.request.function_return_types, + }) + .collect(); + let pool = self.catalog.native.analysis_pool()?; + Ok(pool.install(|| inputs.par_iter().map(enrichment).collect())) + } + fn input_schemas(&self, required: &[String]) -> Shapes { let mut schemas: Shapes = Vec::with_capacity(required.len()); for name in required { @@ -470,6 +543,7 @@ fn enriched_analysis( analysis: &ModelAnalysis, enriched: DeferredAnalysis, star_pending: bool, + lineage: LineageFacts, ) -> ModelAnalysis { let mut recovered: HashMap = HashMap::new(); for column in enriched @@ -489,7 +563,7 @@ fn enriched_analysis( ModelAnalysis { analysis_succeeded: enriched.analysis_succeeded, columns: enriched.columns, - lineage: LineageFacts::PythonEnrichment, + lineage, has_star: enriched.has_star, star_resolved: enriched.star_resolved, binding_diagnostics: analysis.binding_diagnostics.clone(), @@ -508,3 +582,16 @@ fn enriched_analysis( merged.star_resolved = merged.star_resolved || star_expanded; merged } + +/// A native re-analysis as Python's successful answer. +fn native_answer(model: usize, columns: Vec, has_star: bool) -> DeferredAnalysis { + DeferredAnalysis { + model, + analysis_succeeded: true, + columns: Some(columns), + has_star, + star_resolved: false, + binding_diagnostics: Vec::new(), + binding_validated: false, + } +} diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/constants.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/constants.rs index 735e39dbc8..d0d6fa5520 100644 --- a/crates/sqlbuild-analysis/src/assembly/analysis_session/constants.rs +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/constants.rs @@ -38,3 +38,60 @@ pub(crate) const QUOTED_IDENTIFIER_DELIMITER: char = '"'; /// Python's names for a legacy-analysis and an enrichment deferral. pub(crate) const DEFERRAL_ANALYSIS: &str = "analysis"; pub(crate) const DEFERRAL_ENRICHMENT: &str = "enrichment"; +/// Python's `ColumnTransformKind` compact codes. +pub(crate) const TRANSFORM_DIRECT: u8 = 0; +pub(crate) const TRANSFORM_CAST: u8 = 1; +pub(crate) const TRANSFORM_EXPRESSION: u8 = 2; +pub(crate) const TRANSFORM_AGGREGATION: u8 = 3; +pub(crate) const TRANSFORM_STAR: u8 = 4; +pub(crate) const TRANSFORM_CONSTANT: u8 = 5; +/// Python's `ColumnLineageConfidence` compact codes. +pub(crate) const CONFIDENCE_UNKNOWN: u8 = 0; +pub(crate) const CONFIDENCE_HIGH: u8 = 1; +pub(crate) const CONFIDENCE_MEDIUM: u8 = 2; +/// The `analyze_query` facts Python's re-analysis reads. +pub(crate) const SELECT_SHAPE: &str = "select"; +pub(crate) const SET_OPERATION_SHAPE: &str = "set_operation"; +pub(crate) const CAST_TRANSFORM: &str = "cast"; +pub(crate) const FILTER_CONTEXT: &str = "filter"; +pub(crate) const NULL_KEYWORD: &str = "NULL"; +pub(crate) const RESOLVED_SOURCE_CONFIDENCE: &str = "resolved"; +pub(crate) const WILDCARD: &str = "*"; +/// SQLBuild's `analyze_query` complexity guard. +pub(crate) const MAX_FUNCTION_CALL_DEPTH: usize = 512; +/// The dialects whose dynamic pivots the compiler proves, and those with DuckDB's PIVOT form. +pub(crate) const SUPPORTED_PIVOT_DIALECTS: [&str; 3] = ["duckdb", "motherduck", "snowflake"]; +pub(crate) const SIMPLIFIED_PIVOT_DIALECTS: [&str; 2] = ["duckdb", "motherduck"]; +pub(crate) const DUCKDB_DIALECT: &str = "duckdb"; +pub(crate) const MOTHERDUCK_DIALECT: &str = "motherduck"; +/// The `to_dict()` node keys Python's dynamic pivot proof walks. +pub(crate) const PIVOT_AST_KIND: &str = "pivot"; +pub(crate) const SELECT_AST_KIND: &str = "select"; +pub(crate) const QUERY_AST_KINDS: [&str; 2] = ["query", "subquery"]; +pub(crate) const CAST_AST_KIND: &str = "cast"; +pub(crate) const COLUMN_AST_KIND: &str = "column"; +pub(crate) const TABLE_AST_KIND: &str = "table"; +pub(crate) const FUNCTION_AST_KIND: &str = "function"; +pub(crate) const ALIAS_AST_KIND: &str = "alias"; +pub(crate) const CTE_AST_KIND: &str = "cte"; +pub(crate) const STAR_AST_KIND: &str = "star"; +pub(crate) const DYNAMIC_VALUE_SOURCE_KINDS: [&str; 3] = ["pivot_any", "select", "query"]; +pub(crate) const DEPENDENCY_FUNCTION_NAMES: [&str; 3] = ["__ref", "__source", "__seed"]; +pub(crate) const TYPE_PRESERVING_AGGREGATES: [&str; 3] = ["ANY_VALUE", "MAX", "MIN"]; +/// Python's `_render_type` names for cast targets. +pub(crate) const RENDERED_TYPE_NAMES: [(&str, &str); 12] = [ + ("big_int", "BIGINT"), + ("bool", "BOOLEAN"), + ("boolean", "BOOLEAN"), + ("decimal", "DECIMAL"), + ("double", "DOUBLE"), + ("float", "FLOAT"), + ("int", "INT"), + ("integer", "INT"), + ("text", "TEXT"), + ("timestamp", "TIMESTAMP"), + ("var_char", "VARCHAR"), + ("varchar", "VARCHAR"), +]; +/// The stack a standalone dynamic pivot proof parses on, as the analysis workers have. +pub(crate) const PIVOT_WORKER_STACK_BYTES: usize = 16 * 1024 * 1024; diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/main/mod.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/main/mod.rs index 5893bfd5e8..6c5394355c 100644 --- a/crates/sqlbuild-analysis/src/assembly/analysis_session/main/mod.rs +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/main/mod.rs @@ -1,5 +1,6 @@ pub mod expression_shapes; pub mod finish_analysis_session; +pub mod prove_dynamic_contract; pub mod provide_deferred_analyses; pub mod run_analysis_session; pub mod start_analysis_session; diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/main/prove_dynamic_contract.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/main/prove_dynamic_contract.rs new file mode 100644 index 0000000000..96759642d4 --- /dev/null +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/main/prove_dynamic_contract.rs @@ -0,0 +1,26 @@ +//! Python's `analyze_dynamic_column_contract` for one model outside the analysis session. + +use crate::assembly::analysis_session::_helpers::dynamic_pivot::{PivotFacts, pivot_outcome}; +use crate::assembly::analysis_session::constants::PIVOT_WORKER_STACK_BYTES; +use crate::assembly::analysis_session::models::{PivotOutcome, PivotRequest}; + +/// The model's proof, or a deferral where Python's result is not reproduced exactly. +pub fn prove_dynamic_contract(request: &PivotRequest) -> Result { + let facts = PivotFacts { + dialect: &request.dialect, + column_types: &request.column_types, + authoritative_types: &request.authoritative_types, + column_nullability: &request.column_nullability, + families_by_table: &request.families_by_table, + }; + std::thread::scope(|scope| { + std::thread::Builder::new() + .stack_size(PIVOT_WORKER_STACK_BYTES) + .spawn_scoped(scope, || { + pivot_outcome(&request.sql, &request.families, &facts) + }) + .map_err(|error| error.to_string())? + .join() + .map_err(|_| "the native dynamic pivot proof panicked".to_owned()) + }) +} diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/models.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/models.rs index a52f3e38ce..d10d56c0ad 100644 --- a/crates/sqlbuild-analysis/src/assembly/analysis_session/models.rs +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/models.rs @@ -33,6 +33,52 @@ pub struct ModelRequest { pub has_set_operation: bool, /// The snapshot validity columns a published shape appends. pub snapshot_columns: Option<(String, String)>, + /// The SQL Python's dynamic pivot proof parses: placeholders replaced by their defaults. + pub pivot_sql: String, + pub dynamic_families: Vec, +} + +/// One declared family of columns a runtime-dynamic pivot generates. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct DynamicFamily { + pub name: String, + pub pivot_column: String, + pub value_column: String, + pub aggregate: String, + pub data_type: String, + pub name_pattern: Option, +} + +/// Python's `DynamicColumnContractProof`. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ContractProof { + pub output_proven: bool, + pub fixed_columns: Vec, + /// Each family's name and inferred type. + pub families: Vec<(String, Option)>, + pub input_relations: Vec, + pub failure_reason: Option, + pub bare_dynamic_pivot: bool, +} + +/// One model's dynamic pivot proof inputs, for a model the session does not analyse. +#[derive(Debug, Clone, Default)] +pub struct PivotRequest { + pub dialect: String, + pub column_types: Shapes, + pub authoritative_types: Shapes, + pub column_nullability: Shapes, + pub families_by_table: Vec<(String, Vec)>, + pub sql: String, + pub families: Vec, +} + +/// A model's dynamic pivot proof: none declared, left to Python, or proven natively. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum PivotOutcome { + Absent, + Deferred, + Proof(ContractProof), } /// Everything one compile's model analysis reads. @@ -49,6 +95,8 @@ pub struct SessionRequest { pub complete_schemas: Shapes, /// The Python binding catalog's `schemas` before analysis. pub catalog_schemas: Shapes, + /// Python's `dynamic_families_by_table`, in dict order. + pub dynamic_families_by_table: Vec<(String, Vec)>, pub models: Vec, } @@ -78,6 +126,8 @@ pub enum LineageFacts { PythonAnalysis, /// The lineage of the analysis Python returned for this model's enrichment deferral. PythonEnrichment, + /// Plain lineage facts of the native re-analysis with known inputs. + NativeEnrichment(Vec), } /// Python's `PolyglotAnalysisResult` for one model. @@ -147,6 +197,8 @@ pub struct SessionOutcome { pub schema_additions: Shapes, /// Relations whose analysis shapes the Python binding catalog records. pub analysis_names: Vec, + /// Each model's dynamic pivot proof, in request order. + pub dynamic_contracts: Vec, } /// One expression source's shape: inferred, absent, or left to Python. diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/tests/helpers.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/tests/helpers.rs index fccee1d92b..1f274fdfc1 100644 --- a/crates/sqlbuild-analysis/src/assembly/analysis_session/tests/helpers.rs +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/tests/helpers.rs @@ -4,8 +4,9 @@ use crate::assembly::analysis_session::main::provide_deferred_analyses::provide_ use crate::assembly::analysis_session::main::run_analysis_session::run_analysis_session; use crate::assembly::analysis_session::main::start_analysis_session::start_analysis_session; use crate::assembly::analysis_session::models::{ - AnalysisSession, Deferral, DeferredAnalysis, ModelOutcome, ModelReference, ModelRequest, - SessionRequest, SessionStep, + AnalysisSession, ColumnFact, ContractProof, Deferral, DeferredAnalysis, DynamicFamily, + ModelOutcome, ModelReference, ModelRequest, PivotOutcome, PivotRequest, SessionRequest, + SessionStep, }; use crate::assembly::analysis_session::tests::test_types::ModelSpec; use crate::assembly::analysis_session::types::{Pairs, Shapes}; @@ -79,6 +80,8 @@ pub(crate) fn model(name: &str, sql: &str, sources: &[&str], refs: &[&str]) -> M recover_cte_facts: false, has_set_operation: sql.contains("UNION"), snapshot_columns: None, + pivot_sql: sql.to_owned(), + dynamic_families: Vec::new(), } } @@ -100,6 +103,7 @@ pub(crate) fn orders_request(models: Vec) -> SessionRequest { )]), complete_schemas: raw_orders.clone(), catalog_schemas: raw_orders, + dynamic_families_by_table: Vec::new(), models, } } @@ -210,3 +214,86 @@ pub(crate) fn session_lines(models: &[ModelSpec]) -> (Vec>, Vec, +) -> PivotRequest { + let raw_orders: Shapes = shapes(&[( + "raw_orders", + &[ + ("order_id", "INTEGER"), + ("customer_id", "INTEGER"), + ("status", "VARCHAR"), + ("amount", "DOUBLE"), + ], + )]); + let upstream: Shapes = shapes(&[("status_amounts", &[("customer_id", "INTEGER")])]); + let families: Vec = family + .map(|(pivot_column, value_column, aggregate)| { + amounts(pivot_column, value_column, aggregate) + }) + .into_iter() + .collect(); + PivotRequest { + dialect: dialect.to_owned(), + column_types: [raw_orders.clone(), upstream].concat(), + authoritative_types: raw_orders, + column_nullability: shapes(&[("raw_orders", &[("customer_id", "non_null")])]), + families_by_table: vec![( + "status_amounts".to_owned(), + vec![amounts("status", "amount", "MAX")], + )], + sql: sql.to_owned(), + families, + } +} + +fn amounts(pivot_column: &str, value_column: &str, aggregate: &str) -> DynamicFamily { + DynamicFamily { + name: "amounts".to_owned(), + pivot_column: pivot_column.to_owned(), + value_column: value_column.to_owned(), + aggregate: aggregate.to_owned(), + data_type: "DOUBLE".to_owned(), + name_pattern: None, + } +} + +/// A failed proof with Python's reason. +pub(crate) fn failed(reason: &str) -> PivotOutcome { + PivotOutcome::Proof(ContractProof { + output_proven: false, + fixed_columns: Vec::new(), + families: Vec::new(), + input_relations: Vec::new(), + failure_reason: Some(reason.to_owned()), + bare_dynamic_pivot: false, + }) +} + +/// A proven output: `(name, type, nullability)` fixed columns and the `amounts` family type. +pub(crate) fn proven( + fixed: &[(&str, &str, &str)], + family_type: Option<&str>, + input: &str, + bare: bool, +) -> PivotOutcome { + PivotOutcome::Proof(ContractProof { + output_proven: true, + fixed_columns: fixed + .iter() + .map(|(name, data_type, nullability)| ColumnFact { + name: (*name).to_owned(), + data_type: Some((*data_type).to_owned()), + nullability: (*nullability).to_owned(), + }) + .collect(), + families: vec![("amounts".to_owned(), family_type.map(str::to_owned))], + input_relations: vec![input.to_owned()], + failure_reason: None, + bare_dynamic_pivot: bare, + }) +} diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/tests/test_analysis_session.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/tests/test_analysis_session.rs index 4f53c3093f..c56625f012 100644 --- a/crates/sqlbuild-analysis/src/assembly/analysis_session/tests/test_analysis_session.rs +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/tests/test_analysis_session.rs @@ -1,8 +1,12 @@ +use crate::assembly::analysis_session::main::prove_dynamic_contract::prove_dynamic_contract; use crate::assembly::analysis_session::main::start_analysis_session::start_analysis_session; +use crate::assembly::analysis_session::models::PivotOutcome; use crate::assembly::analysis_session::tests::helpers::{ - catalog, model_requests, orders_request, session_lines, + catalog, failed, model_requests, orders_request, pivot_request, proven, session_lines, +}; +use crate::assembly::analysis_session::tests::test_types::{ + PivotTestCase, SessionTestCase, UnscheduledTestCase, }; -use crate::assembly::analysis_session::tests::test_types::{SessionTestCase, UnscheduledTestCase}; #[test] fn given_models_when_running_the_session_then_defers_and_publishes_as_python_does() { @@ -53,14 +57,14 @@ fn given_models_when_running_the_session_then_defers_and_publishes_as_python_doe ], }, SessionTestCase { - description: "an unknown column is a binding diagnostic and an untyped output defers", + description: "an unknown column is a binding diagnostic and enriches natively", models: &[( "orders_bad", "SELECT missing_column FROM __source(\"raw_orders\")", &["raw_orders"], &[], )], - expected_steps: &[&["defer enrichment 0"], &[]], + expected_steps: &[&[]], expected_outcomes: &[&[ "missing_column - unknown", "succeeded=true star=false/false diagnostics=[\"B002\"]", @@ -83,7 +87,6 @@ fn given_models_when_running_the_session_then_defers_and_publishes_as_python_doe ), ], expected_steps: &[ - &["defer enrichment 0"], &["publish events event_id:UNKNOWN", "defer analysis 1"], &[], ], @@ -164,3 +167,127 @@ fn given_unschedulable_models_when_starting_then_python_analyses() { ); } } + +#[test] +fn given_dynamic_pivots_when_proving_then_matches_python_or_defers() { + let test_cases = [ + PivotTestCase { + description: "a bare DuckDB pivot keeps its GROUP BY columns", + dialect: "duckdb", + sql: "PIVOT __source(\"raw_orders\") ON status USING MAX(amount) GROUP BY customer_id", + family: Some(("status", "amount", "MAX")), + expected_outcome: proven( + &[("customer_id", "INTEGER", "non_null")], + Some("DOUBLE"), + "raw_orders", + true, + ), + }, + PivotTestCase { + description: "a cast aggregate renders its target type", + dialect: "duckdb", + sql: "PIVOT __source(\"raw_orders\") ON status \ + USING MIN(CAST(amount AS DECIMAL(12, 2))) GROUP BY customer_id", + family: Some(("status", "amount", "MIN")), + expected_outcome: proven( + &[("customer_id", "INTEGER", "non_null")], + Some("DECIMAL(12, 2)"), + "raw_orders", + true, + ), + }, + PivotTestCase { + description: "a CTE input without GROUP BY keeps the remaining input columns", + dialect: "snowflake", + sql: "WITH base AS (SELECT customer_id, status, amount FROM __source(\"raw_orders\")) \ + SELECT * FROM base PIVOT (SUM(amount) FOR status IN (ANY ORDER BY status))", + family: Some(("status", "amount", "SUM")), + expected_outcome: proven( + &[("customer_id", "INTEGER", "non_null")], + None, + "raw_orders", + false, + ), + }, + PivotTestCase { + description: "a passthrough redeclaring the upstream families", + dialect: "duckdb", + sql: "SELECT * FROM __ref(\"status_amounts\")", + family: Some(("status", "amount", "MAX")), + expected_outcome: proven( + &[("customer_id", "INTEGER", "unknown")], + Some("DOUBLE"), + "status_amounts", + false, + ), + }, + PivotTestCase { + description: "static pivot values", + dialect: "snowflake", + sql: "SELECT * FROM __source(\"raw_orders\") PIVOT (MAX(amount) FOR status IN ('placed'))", + family: Some(("status", "amount", "MAX")), + expected_outcome: failed( + "static pivot values must use ordinary exact column declarations", + ), + }, + PivotTestCase { + description: "DuckDB values on the ON clause", + dialect: "duckdb", + sql: "PIVOT __source(\"raw_orders\") ON status IN ('placed') USING MAX(amount)", + family: Some(("status", "amount", "MAX")), + expected_outcome: failed( + "dynamic column contracts currently require exactly one pivot column", + ), + }, + PivotTestCase { + description: "a dialect without dynamic pivots", + dialect: "postgres", + sql: "SELECT 1", + family: Some(("status", "amount", "MAX")), + expected_outcome: failed( + "adapter dialect 'postgres' does not support compiler-proven dynamic pivots", + ), + }, + PivotTestCase { + description: "a declared pivot column the pivot does not use", + dialect: "duckdb", + sql: "PIVOT __source(\"raw_orders\") ON status USING MAX(amount) GROUP BY customer_id", + family: Some(("order_id", "amount", "MAX")), + expected_outcome: failed( + "dynamic family 'amounts' declares pivot_column 'order_id' but the output pivot uses 'status'", + ), + }, + PivotTestCase { + description: "SQL Python reports as a parse error", + dialect: "duckdb", + sql: "PIVOT FROM WHERE", + family: Some(("status", "amount", "MAX")), + expected_outcome: PivotOutcome::Deferred, + }, + PivotTestCase { + description: "non-ASCII names Python casefolds differently", + dialect: "duckdb", + sql: "SELECT * FROM __ref(\"stra\u{df}e\")", + family: Some(("status", "amount", "MAX")), + expected_outcome: PivotOutcome::Deferred, + }, + PivotTestCase { + description: "no declared families", + dialect: "duckdb", + sql: "SELECT 1", + family: None, + expected_outcome: PivotOutcome::Absent, + }, + ]; + for test_case in test_cases { + let request = pivot_request(test_case.dialect, test_case.sql, test_case.family); + + let outcome = prove_dynamic_contract(&request).expect("the proof runs"); + + assert_eq!( + outcome, test_case.expected_outcome, + "{}", + test_case.description + ); + } +} diff --git a/crates/sqlbuild-analysis/src/assembly/analysis_session/tests/test_types.rs b/crates/sqlbuild-analysis/src/assembly/analysis_session/tests/test_types.rs index f94d532933..c09b783048 100644 --- a/crates/sqlbuild-analysis/src/assembly/analysis_session/tests/test_types.rs +++ b/crates/sqlbuild-analysis/src/assembly/analysis_session/tests/test_types.rs @@ -1,3 +1,5 @@ +use crate::assembly::analysis_session::models::PivotOutcome; + /// A model as `(name, sql, source names, model refs)`. pub(crate) type ModelSpec = ( &'static str, @@ -40,3 +42,12 @@ pub(crate) struct WavesTestCase { pub(crate) producers: &'static [&'static [usize]], pub(crate) expected_waves: Option<&'static [&'static [usize]]>, } + +pub(crate) struct PivotTestCase { + pub(crate) description: &'static str, + pub(crate) dialect: &'static str, + pub(crate) sql: &'static str, + /// `(pivot column, value column, aggregate)` of one `amounts` family, or none declared. + pub(crate) family: Option<(&'static str, &'static str, &'static str)>, + pub(crate) expected_outcome: PivotOutcome, +} diff --git a/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/session.rs b/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/session.rs index 2e3c48a906..87af8f9b33 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/session.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/session.rs @@ -4,13 +4,14 @@ use pyo3::prelude::{Bound, PyModule, PyModuleMethods, PyRef, PyResult, Python}; use pyo3::{pyclass, pyfunction, pymethods, wrap_pyfunction}; use sqlbuild_analysis::assembly::analysis_session::main::expression_shapes::expression_shapes; use sqlbuild_analysis::assembly::analysis_session::main::finish_analysis_session::finish_analysis_session; +use sqlbuild_analysis::assembly::analysis_session::main::prove_dynamic_contract::prove_dynamic_contract; use sqlbuild_analysis::assembly::analysis_session::main::provide_deferred_analyses::provide_deferred_analyses; use sqlbuild_analysis::assembly::analysis_session::main::run_analysis_session::run_analysis_session; use sqlbuild_analysis::assembly::analysis_session::main::start_analysis_session::start_analysis_session; use sqlbuild_analysis::assembly::analysis_session::models::{ - AnalysisSession, ColumnFact, Deferral, DeferredAnalysis, ExpressionShape, + AnalysisSession, ColumnFact, Deferral, DeferredAnalysis, DynamicFamily, ExpressionShape, ExpressionShapeRequest, LineageFacts, LineageRow, ModelOutcome, ModelReference, ModelRequest, - SessionOutcome, SessionRequest, SessionStep, + PivotOutcome, PivotRequest, SessionOutcome, SessionRequest, SessionStep, }; use sqlbuild_analysis::assembly::analysis_session::types::{Pairs, Shapes}; use sqlbuild_analysis::semantic_validation::types::DiagnosticRow; @@ -25,7 +26,9 @@ use crate::bindings::types::CompilerDetach; type ColumnRow = (String, Option, String); /// `(output column, transform code, confidence code, [(resource type, resource, column)])`. type LineageItem = (String, u8, u8, Vec<(String, String, String)>); -/// A model's name, SQL, placeholders, references, lineage, required names and flags. +/// `(name, pivot column, value column, aggregate, type, name pattern)`. +type FamilyRow = (String, String, String, String, String, Option); +/// A model's name, SQL, placeholders, references, lineage, required names, flags and pivot. type ModelRow = ( String, String, @@ -36,6 +39,8 @@ type ModelRow = ( bool, bool, Option<(String, String)>, + String, + Vec, ); /// The profile settings, relation shapes and models one session analyses. type RequestRow = ( @@ -47,6 +52,7 @@ type RequestRow = ( Shapes, Shapes, Shapes, + Vec<(String, Vec)>, Vec, ); /// `(kind, model, cleaned_sql, schemas, binding diagnostics, lineage)`. @@ -82,8 +88,29 @@ type OutcomeRow = ( bool, String, ); -/// `(models, catalog schema additions, analysis-shape names)`. -type FinishRow = (Vec, Shapes, Vec); +/// `(output proven, fixed columns, (family, inferred type), inputs, failure, bare pivot)`. +type ProofRow = ( + bool, + Vec, + Vec<(String, Option)>, + Vec, + Option, + bool, +); +/// The dialect, relation facts, families by table, SQL and families of one pivot proof. +type PivotRequestRow = ( + String, + Shapes, + Shapes, + Shapes, + Vec<(String, Vec)>, + String, + Vec, +); +/// `("absent" | "deferred" | "proof", proof)`. +type ContractRow = (&'static str, Option); +/// `(models, catalog schema additions, analysis-shape names, dynamic pivot proofs)`. +type FinishRow = (Vec, Shapes, Vec, Vec); /// One compile's native model analysis; any failure hands the whole analysis back to Python. #[pyclass(module = "sqlbuild._native")] @@ -204,6 +231,7 @@ fn session_request(request: RequestRow) -> SessionRequest { column_nullability, complete_schemas, catalog_schemas, + dynamic_families_by_table, models, ) = request; SessionRequest { @@ -215,6 +243,7 @@ fn session_request(request: RequestRow) -> SessionRequest { column_nullability, complete_schemas, catalog_schemas, + dynamic_families_by_table: families_by_table(dynamic_families_by_table), models: models.into_iter().map(model_request).collect(), } } @@ -230,6 +259,8 @@ fn model_request(row: ModelRow) -> ModelRequest { recover_cte_facts, has_set_operation, snapshot_columns, + pivot_sql, + dynamic_families, ) = row; ModelRequest { name, @@ -247,6 +278,8 @@ fn model_request(row: ModelRow) -> ModelRequest { recover_cte_facts, has_set_operation, snapshot_columns, + pivot_sql, + dynamic_families: dynamic_families.into_iter().map(dynamic_family).collect(), } } @@ -325,15 +358,57 @@ fn finish_row(outcome: SessionOutcome) -> FinishRow { outcome.models.into_iter().map(outcome_row).collect(), outcome.schema_additions, outcome.analysis_names, + outcome + .dynamic_contracts + .into_iter() + .map(contract_row) + .collect(), ) } +fn families_by_table(rows: Vec<(String, Vec)>) -> Vec<(String, Vec)> { + rows.into_iter() + .map(|(name, families)| (name, families.into_iter().map(dynamic_family).collect())) + .collect() +} + +fn dynamic_family(row: FamilyRow) -> DynamicFamily { + let (name, pivot_column, value_column, aggregate, data_type, name_pattern) = row; + DynamicFamily { + name, + pivot_column, + value_column, + aggregate, + data_type, + name_pattern, + } +} + +fn contract_row(outcome: PivotOutcome) -> ContractRow { + match outcome { + PivotOutcome::Absent => ("absent", None), + PivotOutcome::Deferred => ("deferred", None), + PivotOutcome::Proof(proof) => ( + "proof", + Some(( + proof.output_proven, + proof.fixed_columns.into_iter().map(column_row).collect(), + proof.families, + proof.input_relations, + proof.failure_reason, + proof.bare_dynamic_pivot, + )), + ), + } +} + fn outcome_row(outcome: ModelOutcome) -> OutcomeRow { let analysis = outcome.analysis; let (source, lineage) = match analysis.lineage { LineageFacts::Native(rows) => ("native", lineage_items(rows)), LineageFacts::PythonAnalysis => ("analysis", Vec::new()), LineageFacts::PythonEnrichment => ("enrichment", Vec::new()), + LineageFacts::NativeEnrichment(rows) => ("facts", lineage_items(rows)), }; ( analysis.analysis_succeeded, @@ -350,9 +425,28 @@ fn outcome_row(outcome: ModelOutcome) -> OutcomeRow { ) } +/// One model's dynamic pivot proof, or `("deferred", None)` for Python's proof. +#[pyfunction] +fn prove_dynamic_column_contract(py: Python<'_>, request: PivotRequestRow) -> ContractRow { + let (dialect, column_types, authoritative_types, column_nullability, by_table, sql, families) = + request; + let request = PivotRequest { + dialect, + column_types, + authoritative_types, + column_nullability, + families_by_table: families_by_table(by_table), + sql, + families: families.into_iter().map(dynamic_family).collect(), + }; + py.compiler_detach(|| prove_dynamic_contract(&request)) + .map_or(("deferred", None), contract_row) +} + pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { module.add_class::()?; module.add_function(wrap_pyfunction!(start_model_analysis_session, module)?)?; module.add_function(wrap_pyfunction!(infer_expression_source_shapes, module)?)?; + module.add_function(wrap_pyfunction!(prove_dynamic_column_contract, module)?)?; Ok(()) } diff --git a/scripts/compiler_differential/_helpers/analysis_corpus/analysis_session.py b/scripts/compiler_differential/_helpers/analysis_corpus/analysis_session.py index 2c00e8a6d6..4a7c34a23e 100644 --- a/scripts/compiler_differential/_helpers/analysis_corpus/analysis_session.py +++ b/scripts/compiler_differential/_helpers/analysis_corpus/analysis_session.py @@ -13,6 +13,18 @@ f"SELECT s.*, s.amount * 2 AS doubled_amount\nFROM {_ORDERS} s\n" ) +_PIVOT_PATH: str = "models/marts/status_amounts.sql" + + +def _pivot_model(*, pivot_column: str, body: str) -> str: + return ( + 'MODEL (\n description "Order amounts pivoted by status",\n contract enforced,\n' + " materialized table,\n columns (customer_id (type INTEGER)),\n" + f" dynamic_columns (\n status_amounts (\n pivot_column {pivot_column},\n" + " value_column amount,\n aggregate MAX,\n type DOUBLE\n )\n ),\n);\n\n" + f"{body}" + ) + def analysis_session_failure_cases() -> tuple[FailureCase, ...]: """Return this lane's failure cases; the lane appends here without editing shared lists.""" @@ -47,4 +59,39 @@ def analysis_session_failure_cases() -> tuple[FailureCase, ...]: "SELECT customer_id, total_amount, order_count FROM totals\n", }, ), + failure_case( + name="analysis-session-enriched-subquery-type-mismatch", + expected_code="K002", + files={ + FAILURE_MART_PATH: 'MODEL (\n description "Order totals per customer",\n' + " contract enforced,\n" + " columns (\n customer_id (type INTEGER),\n amount (type VARCHAR),\n ),\n" + ");\n\n" + f"SELECT customer_id, amount\nFROM (SELECT * FROM {_ORDERS}) staged\n", + }, + ), + failure_case( + name="analysis-session-pivot-values-on-clause", + expected_code="K011", + expected_message="exactly one pivot column", + files={ + _PIVOT_PATH: _pivot_model( + pivot_column="status", + body="PIVOT __source(\"raw_orders\")\nON status IN ('placed')\n" + "USING MAX(amount)\nGROUP BY customer_id\n", + ), + }, + ), + failure_case( + name="analysis-session-pivot-family-column-mismatch", + expected_code="K011", + expected_message="declares pivot_column 'region'", + files={ + _PIVOT_PATH: _pivot_model( + pivot_column="region", + body='PIVOT __source("raw_orders")\nON status\nUSING MAX(amount)\n' + "GROUP BY customer_id\n", + ), + }, + ), ) diff --git a/src/sqlbuild/_native.pyi b/src/sqlbuild/_native.pyi index 1cd39de451..dfda5a5faf 100644 --- a/src/sqlbuild/_native.pyi +++ b/src/sqlbuild/_native.pyi @@ -397,6 +397,20 @@ class NativeModelAnalysisSession: ], list[tuple[str, list[tuple[str, str]]]], list[str], + list[ + tuple[ + str, + tuple[ + bool, + list[tuple[str, str | None, str]], + list[tuple[str, str | None]], + list[str], + str | None, + bool, + ] + | None, + ] + ], ] | None ): ... @@ -406,6 +420,29 @@ class NativeModelAnalysisSession: def start_model_analysis_session( catalog: object, request: tuple[object, ...], / ) -> NativeModelAnalysisSession | None: ... +def prove_dynamic_column_contract( + request: tuple[ + str, + list[tuple[str, list[tuple[str, str]]]], + list[tuple[str, list[tuple[str, str]]]], + list[tuple[str, list[tuple[str, str]]]], + list[tuple[str, list[tuple[str, str, str, str, str, str | None]]]], + str, + list[tuple[str, str, str, str, str, str | None]], + ], + /, +) -> tuple[ + str, + tuple[ + bool, + list[tuple[str, str | None, str]], + list[tuple[str, str | None]], + list[str], + str | None, + bool, + ] + | None, +]: ... def infer_expression_source_shapes( catalog: object, request: tuple[str, bool, list[tuple[str, str]], list[str]], / ) -> tuple[list[tuple[bool, list[tuple[str, str]] | None]], str | None]: ... diff --git a/src/sqlbuild/compiler/analysis_session/_helpers/dynamic_contract.py b/src/sqlbuild/compiler/analysis_session/_helpers/dynamic_contract.py new file mode 100644 index 0000000000..803935326f --- /dev/null +++ b/src/sqlbuild/compiler/analysis_session/_helpers/dynamic_contract.py @@ -0,0 +1,52 @@ +"""Prove one model's dynamic column contract natively, outside the analysis session.""" + +from __future__ import annotations + +import sqlbuild._native as _native +from sqlbuild.compiler.analysis_session._helpers.deferral_records import record_analysis_deferral +from sqlbuild.compiler.analysis_session._helpers.session_rows import ( + contract_proof, + family_rows, + shape_rows, +) +from sqlbuild.compiler.analysis_session.constants import CONTRACT_DEFERRED, DEFERRAL_DYNAMIC_PIVOT +from sqlbuild.compiler.analysis_session.types import ContractRow, ShapeRows +from sqlbuild.compiler.compile.classes.python_model_analysis import PythonModelAnalysis +from sqlbuild.compiler.compile.models import DynamicColumnContractProof +from sqlbuild.compiler.lineage.types import InferredNullability +from sqlbuild.spec.contracts.models import SchemaDynamicColumnFamily + + +def native_dynamic_contract( + *, + query_sql: str, + dialect: str | None, + families: tuple[SchemaDynamicColumnFamily, ...], + column_types_by_table: dict[str, dict[str, str]], + authoritative_column_types_by_table: dict[str, dict[str, str]], + column_nullability_by_table: dict[str, dict[str, InferredNullability]], + dynamic_families_by_table: dict[str, tuple[SchemaDynamicColumnFamily, ...]], +) -> DynamicColumnContractProof | None: + """Python's proof for declared families, or None (recorded) where Python must prove it.""" + + shapes: tuple[ShapeRows | None, ...] = ( + shape_rows(column_types_by_table), + shape_rows(authoritative_column_types_by_table), + shape_rows(column_nullability_by_table), + ) + row: ContractRow = (CONTRACT_DEFERRED, None) + if all(rows is not None for rows in shapes): + row = _native.prove_dynamic_column_contract( + ( + dialect or "generic", + shapes[0] or [], + shapes[1] or [], + shapes[2] or [], + family_rows(dynamic_families_by_table), + query_sql, + [PythonModelAnalysis.family_row(family) for family in families], + ) + ) + if row[0] == CONTRACT_DEFERRED: + record_analysis_deferral(kind=DEFERRAL_DYNAMIC_PIVOT) + return contract_proof(row[1]) diff --git a/src/sqlbuild/compiler/analysis_session/_helpers/session_rows.py b/src/sqlbuild/compiler/analysis_session/_helpers/session_rows.py index 88b2dab191..093c361e6e 100644 --- a/src/sqlbuild/compiler/analysis_session/_helpers/session_rows.py +++ b/src/sqlbuild/compiler/analysis_session/_helpers/session_rows.py @@ -10,18 +10,25 @@ ColumnRow, DeferredRow, DiagnosticRow, + FamilyRow, LineageItem, + ProofRow, ShapeRows, ) from sqlbuild.compiler.compile.classes.python_model_analysis import PythonModelAnalysis from sqlbuild.compiler.compile.models import ( CompactLineageFacts, + CompiledLineageColumnFact, + CompiledLineageSourceFact, + DynamicColumnContractProof, + DynamicColumnFamilyProof, InferredColumn, PolyglotAnalysisResult, ) from sqlbuild.compiler.lineage.types import InferredNullability from sqlbuild.compiler.sql_analysis.constants import CASE_SENSITIVE_BINDING_DIALECTS from sqlbuild.compiler.sql_analysis.models import SqlBindingDiagnostic +from sqlbuild.spec.contracts.models import SchemaDynamicColumnFamily def shape_rows(shapes: Mapping[str, Mapping[str, object]]) -> ShapeRows | None: @@ -61,10 +68,30 @@ def session_request( list(profile.function_return_types.items()), request.rich_type_inference, *shapes, + family_rows(request.dynamic_families_by_table), python.model_rows(), ) +def contract_proof(row: ProofRow | None) -> DynamicColumnContractProof | None: + """A native dynamic pivot proof as Python's, None when absent or left to Python.""" + + if row is None: + return None + output_proven, columns, families, input_relations, failure_reason, bare = row + return DynamicColumnContractProof( + output_proven=output_proven, + fixed_columns=inferred_columns(columns) or (), + families=tuple( + DynamicColumnFamilyProof(name=name, inferred_type=inferred_type) + for name, inferred_type in families + ), + input_relations=tuple(input_relations), + failure_reason=failure_reason, + bare_dynamic_pivot=bare, + ) + + def case_sensitive_shapes(*, profile: ExpressionInferenceProfile, dialect: str | None) -> bool: """Python's `inferred_binding_shape` test for keeping authored identifier quoting.""" @@ -143,3 +170,38 @@ def compact_lineage(rows: list[LineageItem]) -> CompactLineageFacts: ) ) return CompactLineageFacts(string_pool=tuple(pool), rows=tuple(compact_rows)) + + +def lineage_facts(rows: list[LineageItem]) -> tuple[CompiledLineageColumnFact, ...]: + """Native lineage rows as the plain facts Python's re-analysis returns.""" + + return tuple( + CompiledLineageColumnFact( + output_column=output_column, + upstream_columns=_source_facts(sources), + transform_kind=CompactLineageFacts.transform_kind(transform_code), + confidence=CompactLineageFacts.confidence(confidence_code), + ) + for output_column, transform_code, confidence_code, sources in rows + ) + + +def _source_facts(sources: list[tuple[str, str, str]]) -> tuple[CompiledLineageSourceFact, ...]: + return tuple( + CompiledLineageSourceFact( + resource_type=resource_type, resource_name=resource_name, column_name=column_name + ) + for resource_type, resource_name, column_name in sources + ) + + +def family_rows( + families_by_table: dict[str, tuple[SchemaDynamicColumnFamily, ...]], +) -> list[tuple[str, list[FamilyRow]]]: + """Declared dynamic families by relation, as native request rows.""" + + return [(name, _family_rows(families)) for name, families in families_by_table.items()] + + +def _family_rows(families: tuple[SchemaDynamicColumnFamily, ...]) -> list[FamilyRow]: + return [PythonModelAnalysis.family_row(family) for family in families] diff --git a/src/sqlbuild/compiler/analysis_session/classes/native_model_analysis.py b/src/sqlbuild/compiler/analysis_session/classes/native_model_analysis.py index 7674467755..8fa6049f45 100644 --- a/src/sqlbuild/compiler/analysis_session/classes/native_model_analysis.py +++ b/src/sqlbuild/compiler/analysis_session/classes/native_model_analysis.py @@ -12,14 +12,19 @@ from sqlbuild.compiler.analysis_session._helpers.session_rows import ( binding_diagnostics, compact_lineage, + contract_proof, deferred_row, inferred_columns, + lineage_facts, ) from sqlbuild.compiler.analysis_session.constants import ( + CONTRACT_DEFERRED, DEFERRAL_ANALYSIS, + DEFERRAL_DYNAMIC_PIVOT, DEFERRAL_ENRICHMENT, DEFERRAL_LEGACY_ANALYSIS, DEFERRAL_SESSION, + LINEAGE_FACTS, LINEAGE_NATIVE, NATIVE_ANALYSIS_FAILURE_MESSAGE, NATIVE_SESSION_FAILURE_MESSAGE, @@ -30,6 +35,7 @@ DeferredRow, FinishRow, OutcomeRow, + ProofRow, ShapeRows, StepRow, ) @@ -79,19 +85,22 @@ def analyses( sqlbuild_error=session.failure, ) return None - outcomes, schema_additions, analysis_names = finished + outcomes, schema_additions, analysis_names, contracts = finished self._record_catalog_changes( schema_additions=schema_additions, analysis_names=analysis_names ) + self._deferrals[DEFERRAL_DYNAMIC_PIVOT] += sum( + kind == CONTRACT_DEFERRED for kind, _ in contracts + ) for kind, count in self._deferrals.items(): record_analysis_deferral(kind=kind, count=count) self._python.record_uncached() return { request.model_input.model_file.file_path.stem: self._model_analysis( - index=index, request=request, outcome=outcome + index=index, request=request, outcome=outcome, contract=contract ) - for index, (request, outcome) in enumerate( - zip(self._python.requests, outcomes, strict=True) + for index, (request, outcome, (_, contract)) in enumerate( + zip(self._python.requests, outcomes, contracts, strict=True) ) } @@ -168,7 +177,12 @@ def _record_catalog_changes( ) def _model_analysis( - self, *, index: int, request: ModelSqlAnalysisRequest, outcome: OutcomeRow + self, + *, + index: int, + request: ModelSqlAnalysisRequest, + outcome: OutcomeRow, + contract: ProofRow | None, ) -> ModelSqlAnalysis: ( succeeded, @@ -184,6 +198,8 @@ def _model_analysis( lineage: Sequence[CompiledLineageColumnFact] = ( (compact_lineage(lineage_rows) if succeeded else ()) if lineage_source == LINEAGE_NATIVE + else lineage_facts(lineage_rows) + if lineage_source == LINEAGE_FACTS else self._kept[(index, lineage_source)].lineage_columns ) return ModelSqlAnalysis( @@ -200,4 +216,5 @@ def _model_analysis( fused_binding_validated=True, cleaned_sql=cleaned_sql, validated_schema=request.binding_schema, + dynamic_column_contract=contract_proof(contract), ) diff --git a/src/sqlbuild/compiler/analysis_session/constants.py b/src/sqlbuild/compiler/analysis_session/constants.py index 51a32d88cf..9d819d795d 100644 --- a/src/sqlbuild/compiler/analysis_session/constants.py +++ b/src/sqlbuild/compiler/analysis_session/constants.py @@ -6,10 +6,13 @@ DEFERRAL_SESSION: str = "session" DEFERRAL_LEGACY_ANALYSIS: str = "legacy_analysis" DEFERRAL_ENRICHMENT: str = "input_enrichment" +DEFERRAL_DYNAMIC_PIVOT: str = "dynamic_pivot" DEFERRAL_EXPRESSION_SHAPES: str = "expression_shapes" DEFERRAL_ANALYSIS: str = "analysis" DEFERRAL_ENRICH: str = "enrichment" LINEAGE_NATIVE: str = "native" +LINEAGE_FACTS: str = "facts" +CONTRACT_DEFERRED: str = "deferred" NATIVE_ANALYSIS_FAILURE_MESSAGE: str = "native compact query analysis failed" NATIVE_SESSION_FAILURE_MESSAGE: str = "native model analysis failed; falling back" NATIVE_SHAPES_FAILURE_MESSAGE: str = "native expression-source shape inference failed; falling back" diff --git a/src/sqlbuild/compiler/analysis_session/main/_prove_native_dynamic_contract.py b/src/sqlbuild/compiler/analysis_session/main/_prove_native_dynamic_contract.py new file mode 100644 index 0000000000..9d47358f42 --- /dev/null +++ b/src/sqlbuild/compiler/analysis_session/main/_prove_native_dynamic_contract.py @@ -0,0 +1,31 @@ +"""Prove a dynamic column contract natively for the preview compiler engine.""" + +from __future__ import annotations + +from sqlbuild.compiler.analysis_session._helpers.dynamic_contract import native_dynamic_contract +from sqlbuild.compiler.compile.models import DynamicColumnContractProof +from sqlbuild.compiler.lineage.types import InferredNullability +from sqlbuild.spec.contracts.models import SchemaDynamicColumnFamily + + +def prove_native_dynamic_contract( + *, + query_sql: str, + dialect: str | None, + families: tuple[SchemaDynamicColumnFamily, ...], + column_types_by_table: dict[str, dict[str, str]], + authoritative_column_types_by_table: dict[str, dict[str, str]], + column_nullability_by_table: dict[str, dict[str, InferredNullability]], + dynamic_families_by_table: dict[str, tuple[SchemaDynamicColumnFamily, ...]], +) -> DynamicColumnContractProof | None: + """Return Python's proof for declared families, or None to prove it in Python.""" + + return native_dynamic_contract( + query_sql=query_sql, + dialect=dialect, + families=families, + column_types_by_table=column_types_by_table, + authoritative_column_types_by_table=authoritative_column_types_by_table, + column_nullability_by_table=column_nullability_by_table, + dynamic_families_by_table=dynamic_families_by_table, + ) diff --git a/src/sqlbuild/compiler/analysis_session/types.py b/src/sqlbuild/compiler/analysis_session/types.py index 746ce29cc0..8fcf150865 100644 --- a/src/sqlbuild/compiler/analysis_session/types.py +++ b/src/sqlbuild/compiler/analysis_session/types.py @@ -22,4 +22,9 @@ bool, str, ] -type FinishRow = tuple[list[OutcomeRow], ShapeRows, list[str]] +type FamilyRow = tuple[str, str, str, str, str, str | None] +type ProofRow = tuple[ + bool, list[ColumnRow], list[tuple[str, str | None]], list[str], str | None, bool +] +type ContractRow = tuple[str, ProofRow | None] +type FinishRow = tuple[list[OutcomeRow], ShapeRows, list[str], list[ContractRow]] diff --git a/src/sqlbuild/compiler/compile/_helpers/native_stages/assembly.py b/src/sqlbuild/compiler/compile/_helpers/native_stages/assembly.py index fc4fe25baa..145745939d 100644 --- a/src/sqlbuild/compiler/compile/_helpers/native_stages/assembly.py +++ b/src/sqlbuild/compiler/compile/_helpers/native_stages/assembly.py @@ -12,6 +12,9 @@ from sqlbuild.compiler.analysis_session.main._infer_native_expression_source_shapes import ( infer_native_expression_source_shapes, ) +from sqlbuild.compiler.analysis_session.main._prove_native_dynamic_contract import ( + prove_native_dynamic_contract, +) from sqlbuild.compiler.analysis_session.models import NativeModelAnalysisRequest from sqlbuild.compiler.compile._helpers.assembly.semantic_shapes import ( get_expression_source_shapes, @@ -91,8 +94,18 @@ def dynamic_column_contract_by_engine( sql_analysis: ModelSqlAnalysis | None, python_proof: partial[DynamicColumnContractProof | None], ) -> DynamicColumnContractProof | None: - """Return the proof the model's native analysis carries, or prove the contract in Python.""" + """Return the proof native analysis carries or proves alone, or prove it in Python.""" if sql_analysis is not None and sql_analysis.dynamic_column_contract is not None: return sql_analysis.dynamic_column_contract + if ( + sql_analysis is None + and python_proof.keywords["families"] + and native_stage_enabled(NativeStage.MODEL_ANALYSIS) + ): + native: DynamicColumnContractProof | None = prove_native_dynamic_contract( + **python_proof.keywords + ) + if native is not None: + return native return python_proof() diff --git a/src/sqlbuild/compiler/compile/classes/python_model_analysis.py b/src/sqlbuild/compiler/compile/classes/python_model_analysis.py index 20a7c1b100..025f8187d0 100644 --- a/src/sqlbuild/compiler/compile/classes/python_model_analysis.py +++ b/src/sqlbuild/compiler/compile/classes/python_model_analysis.py @@ -6,7 +6,10 @@ from sqlbuild.adapter.contract.models import ExpressionInferenceProfile from sqlbuild.compiler.compile._helpers.analysis.cache import record_analysis_cache_metrics -from sqlbuild.compiler.compile._helpers.analysis.columns import table_function_analysis_name +from sqlbuild.compiler.compile._helpers.analysis.columns import ( + substitute_placeholder_defaults, + table_function_analysis_name, +) from sqlbuild.compiler.compile._helpers.analysis.compact import ( analyze_columns_and_lineage_with_polyglot, ) @@ -33,6 +36,7 @@ from sqlbuild.compiler.planner.types import ContractPolicy, MaterializationType from sqlbuild.compiler.references.types import SqlReferenceKind from sqlbuild.compiler.sql_analysis.models import SqlBindingDiagnostic +from sqlbuild.spec.contracts.models import SchemaDynamicColumnFamily _SET_OPERATION_PATTERN: re.Pattern[str] = re.compile( r"\b(?:UNION|INTERSECT|EXCEPT)\b", re.IGNORECASE @@ -116,6 +120,19 @@ def enrich( recover_cte_facts=_recovers_cte_facts(request.model_input), ) + @staticmethod + def family_row(family: SchemaDynamicColumnFamily) -> tuple[str, str, str, str, str, str | None]: + """One declared dynamic column family as the native request row.""" + + return ( + family.name, + family.pivot_column, + family.value_column, + family.aggregate, + family.type, + family.name_pattern, + ) + @staticmethod def legacy_precomputed( *, @@ -189,6 +206,21 @@ def _model_row(request: ModelSqlAnalysisRequest) -> tuple[object, ...]: if values.get("materialized") == MaterializationType.SNAPSHOT else None ), + ( + substitute_placeholder_defaults( + query_sql=request.query_sql, placeholders=request.placeholders + ) + if request.placeholders + else request.query_sql + ), + [ + PythonModelAnalysis.family_row(family) + for family in ( + model_input.schema_entry.dynamic_columns + if model_input.schema_entry is not None + else () + ) + ], ) diff --git a/tests/e2e/scripts/compiler_differential/test_compiler_differential.py b/tests/e2e/scripts/compiler_differential/test_compiler_differential.py index f0184ff6c8..d21f2bd436 100644 --- a/tests/e2e/scripts/compiler_differential/test_compiler_differential.py +++ b/tests/e2e/scripts/compiler_differential/test_compiler_differential.py @@ -444,7 +444,6 @@ def test_given_failure_expectation_when_code_is_not_the_first_error_then_harness "Analysis deferrals (python): none recorded", "Analysis deferrals (native-preview):", " metadata_checks metadata_validation.py (project ", - " input_enrichment analysis_session (project ", "Compiler differential passed: 1 projects identical", ), expected_sites=frozenset( @@ -453,9 +452,7 @@ def test_given_failure_expectation_when_code_is_not_the_first_error_then_harness "_validate_sql_syntax_with_message parse_one", } ), - expected_deferrals=frozenset( - {"metadata_checks metadata_validation.py", "input_enrichment analysis_session"} - ), + expected_deferrals=frozenset({"metadata_checks metadata_validation.py"}), ), WheelSiteReportTestCase( description="preview_deferrals_are_counted_per_kind_and_site", @@ -467,11 +464,7 @@ def test_given_failure_expectation_when_code_is_not_the_first_error_then_harness ), expected_sites=frozenset(), expected_deferrals=frozenset( - { - "legacy_fallback orders.sql", - "metadata_checks metadata_validation.py", - "input_enrichment analysis_session", - } + {"legacy_fallback orders.sql", "metadata_checks metadata_validation.py"} ), ), ], diff --git a/tests/integration/src/sqlbuild/compiler/analysis_session/_test_types.py b/tests/integration/src/sqlbuild/compiler/analysis_session/_test_types.py index 8c14cdd29c..28ee16c825 100644 --- a/tests/integration/src/sqlbuild/compiler/analysis_session/_test_types.py +++ b/tests/integration/src/sqlbuild/compiler/analysis_session/_test_types.py @@ -14,6 +14,8 @@ class GeneratedAnalysisParityTestCase: dialects: tuple[str | None, ...] expected_minimum_native: int expected_minimum_expression_shapes: int + expected_minimum_pivot_proofs: int + expected_minimum_proven_pivots: int @dataclass(frozen=True) diff --git a/tests/integration/src/sqlbuild/compiler/analysis_session/helpers.py b/tests/integration/src/sqlbuild/compiler/analysis_session/helpers.py index 080c39fe00..dfc7e9b2b8 100644 --- a/tests/integration/src/sqlbuild/compiler/analysis_session/helpers.py +++ b/tests/integration/src/sqlbuild/compiler/analysis_session/helpers.py @@ -23,6 +23,9 @@ from sqlbuild.compiler.analysis_session.main._infer_native_expression_source_shapes import ( infer_native_expression_source_shapes, ) +from sqlbuild.compiler.analysis_session.main._prove_native_dynamic_contract import ( + prove_native_dynamic_contract, +) from sqlbuild.compiler.analysis_session.models import NativeModelAnalysisRequest from sqlbuild.compiler.compile._helpers.assembly.semantic_shapes import ( get_expression_source_shapes, @@ -33,6 +36,7 @@ CompactLineageFacts, CompileAdapterContext, CompileProjectInputs, + DynamicColumnContractProof, ModelSqlAnalysis, ) from sqlbuild.compiler.discovery.main.discover import discover_project_inputs @@ -78,26 +82,63 @@ python_functions_inherit_default_namespace=True, sql_lexical_syntax=DuckDbAdapter.sql_lexical_syntax, ) -_PIVOT_MODEL: str = """MODEL ( +_PIVOT_HEADER: str = """MODEL ( description "Order amounts pivoted by status", contract enforced, materialized table, columns (customer_id (type INTEGER)), dynamic_columns ( status_amounts ( - pivot_column status, - value_column amount, - aggregate MAX, + pivot_column {pivot}, + value_column {value}, + aggregate {aggregate}, type DOUBLE ) ), ); -PIVOT __source("raw_orders") -ON status -USING MAX(amount) -GROUP BY customer_id """ +_PIVOT_BODIES: tuple[str, ...] = ( + 'PIVOT __source("raw_orders")\nON status\nUSING {aggregate}(amount)\nGROUP BY customer_id', + 'PIVOT __source("raw_orders")\nON status\nUSING {aggregate}(CAST(amount AS DECIMAL(12, 2)))', + "WITH base AS (\n SELECT customer_id, status, CAST(amount AS DOUBLE) AS amount\n" + ' FROM __source("raw_orders")\n)\nSELECT * FROM base\n' + "PIVOT ({aggregate}(amount) FOR status IN (ANY ORDER BY status))", + 'WITH base AS (\n SELECT * FROM __source("raw_orders")\n),\npivoted AS (\n' + " PIVOT base ON status USING {aggregate}(amount) GROUP BY customer_id\n)\n" + "SELECT * FROM pivoted", + "PIVOT __source(\"raw_orders\")\nON status IN ('placed', 'shipped')\n" + "USING {aggregate}(amount)\nGROUP BY customer_id", + 'PIVOT __source("raw_events")\nON kind\nUSING {aggregate}(order_id)\nGROUP BY event_id', + 'SELECT * FROM __source("raw_orders") o\nJOIN __source("raw_customers") c USING (customer_id)', +) +_PIVOT_AGGREGATES: tuple[str, ...] = ("MAX", "MIN", "ANY_VALUE", "SUM") + + +def _pivot_models(rng: random.Random) -> dict[str, str]: + """A typed pivot every passthrough reads, and seeded pivot shapes proven or refused.""" + + files: dict[str, str] = { + "models/marts/status_amounts.sql": _PIVOT_HEADER.format( + pivot="status", value="amount", aggregate="MAX" + ) + + _PIVOT_BODIES[0].format(aggregate="MAX") + + "\n", + "models/marts/status_passthrough.sql": _PIVOT_HEADER.format( + pivot="status", value="amount", aggregate="MAX" + ) + + 'SELECT * FROM __ref("status_amounts")\n', + } + for index, body in enumerate(rng.sample(_PIVOT_BODIES, k=4)): + aggregate: str = rng.choice(_PIVOT_AGGREGATES) + header: str = _PIVOT_HEADER.format( + pivot=rng.choice(("status", "status", "kind")), + value=rng.choice(("amount", "amount", "order_id")), + aggregate=aggregate, + ) + files[f"models/marts/pivot_{index}.sql"] = header + body.format(aggregate=aggregate) + "\n" + return files + type _Relation = tuple[str, tuple[str, ...]] type _Template = Callable[[random.Random, list[_Relation]], tuple[str, str, tuple[str, ...]]] @@ -221,7 +262,7 @@ def generated_analysis_files(*, rng: random.Random, model_count: int) -> dict[st files: dict[str, str] = { "sqlbuild_project.toml": _PROJECT_TOML, "sources/raw.yml": _SOURCES, - "models/marts/status_amounts.sql": _PIVOT_MODEL, + **_pivot_models(rng), } relations: list[_Relation] = list(_INITIAL_RELATIONS) for index in range(model_count): @@ -258,6 +299,9 @@ class AnalysisParity: native: list[object] = field(default_factory=list) analysed_models: int = 0 expression_shapes: int = 0 + pivot_proofs: int = 0 + standalone_proofs: int = 0 + proven_pivots: int = 0 def compare_analyses( @@ -295,6 +339,25 @@ def models_by_both( _catalog_view(profile.binding_catalog), _catalog_view(native_catalog), ) + return native or python + + def proofs_by_both( + *, + sql_analysis: ModelSqlAnalysis | None, + python_proof: partial[DynamicColumnContractProof | None], + ) -> DynamicColumnContractProof | None: + native: DynamicColumnContractProof | None = getattr( + sql_analysis, "dynamic_column_contract", None + ) + python: DynamicColumnContractProof | None = python_proof() + standalone: DynamicColumnContractProof | None = prove_native_dynamic_contract( + **python_proof.keywords + ) + parity.pivot_proofs += native is not None + parity.proven_pivots += native is not None and native.output_proven + parity.standalone_proofs += standalone is not None + _append(parity, "dynamic pivot proof", python, native or python) + _append(parity, "standalone dynamic pivot proof", python, standalone or python) return python def shapes_by_both( @@ -320,6 +383,7 @@ def shapes_by_both( with monkeypatch.context() as patch: patch.setattr(project_assembly, "analyze_model_sql_by_engine", models_by_both) patch.setattr(project_assembly, "expression_source_shapes_by_engine", shapes_by_both) + patch.setattr(project_assembly, "dynamic_column_contract_by_engine", proofs_by_both) with suppress(CompileInputError): _ = project_assembly.assemble_compiled_project( inputs=inputs, @@ -387,7 +451,6 @@ def _analysis_views(analyses: dict[str, ModelSqlAnalysis] | None) -> object: analysis.polyglot_analysis.binding_validated, analysis.cleaned_sql, analysis.placeholders, - analysis.dynamic_column_contract, ) for name, analysis in analyses.items() } diff --git a/tests/integration/src/sqlbuild/compiler/analysis_session/test_native_model_analysis.py b/tests/integration/src/sqlbuild/compiler/analysis_session/test_native_model_analysis.py index ee5c88edd7..f21fe53fc9 100644 --- a/tests/integration/src/sqlbuild/compiler/analysis_session/test_native_model_analysis.py +++ b/tests/integration/src/sqlbuild/compiler/analysis_session/test_native_model_analysis.py @@ -44,13 +44,15 @@ "test_case", [ GeneratedAnalysisParityTestCase( - description="stars, CTEs, set operations, untyped inputs, contracts and a pivot", + description="stars, CTEs, set operations, untyped inputs, contracts and pivots", seed=20261008, count=6, model_count=24, dialects=("duckdb", "postgres", "snowflake", "bigquery"), expected_minimum_native=500, expected_minimum_expression_shapes=40, + expected_minimum_pivot_proofs=80, + expected_minimum_proven_pivots=12, ) ], ids=lambda case: case.description, @@ -75,11 +77,15 @@ def test_given_generated_projects_when_analysing_natively_then_matches_python( assert mismatches(inputs=parity.names, expected=parity.python, actual=parity.native) == [] assert kinds["analysis_session:session"] == kinds["analysis_session:expression_shapes"] == 0 + assert kinds["analysis_session:dynamic_pivot"] == 0 assert ( parity.analysed_models - kinds["analysis_session:legacy_analysis"] >= test_case.expected_minimum_native ) assert parity.expression_shapes >= test_case.expected_minimum_expression_shapes + assert parity.pivot_proofs >= test_case.expected_minimum_pivot_proofs + assert parity.standalone_proofs >= test_case.expected_minimum_pivot_proofs + assert parity.proven_pivots >= test_case.expected_minimum_proven_pivots @pytest.mark.parametrize( From a86b0084967e01bac6fcc7655842c549d2c06338 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Fri, 9 Oct 2026 05:38:37 +0100 Subject: [PATCH 2/2] test(native): gate enrichment deferrals, CTE recovery and standalone pivot proofs --- .../compiler/analysis_session/_test_types.py | 12 +++ .../compiler/analysis_session/helpers.py | 101 ++++++++++++++++-- .../test_native_model_analysis.py | 62 ++++++++++- 3 files changed, 164 insertions(+), 11 deletions(-) diff --git a/tests/integration/src/sqlbuild/compiler/analysis_session/_test_types.py b/tests/integration/src/sqlbuild/compiler/analysis_session/_test_types.py index 28ee16c825..f3ea3a0aac 100644 --- a/tests/integration/src/sqlbuild/compiler/analysis_session/_test_types.py +++ b/tests/integration/src/sqlbuild/compiler/analysis_session/_test_types.py @@ -16,6 +16,8 @@ class GeneratedAnalysisParityTestCase: expected_minimum_expression_shapes: int expected_minimum_pivot_proofs: int expected_minimum_proven_pivots: int + expected_maximum_enrichment_deferrals: int + expected_minimum_native_enrichments: int @dataclass(frozen=True) @@ -37,3 +39,13 @@ class SessionFailureTestCase: seed: int model_count: int expected_kinds: dict[str, int] + + +@dataclass(frozen=True) +class StandalonePivotProofTestCase: + """A pivot model left out of model analysis, whose proof assembly takes natively.""" + + description: str + analysed_models: frozenset[str] + expected_native_proofs: int + expected_proven_by_model: dict[str, bool | None] diff --git a/tests/integration/src/sqlbuild/compiler/analysis_session/helpers.py b/tests/integration/src/sqlbuild/compiler/analysis_session/helpers.py index dfc7e9b2b8..47095774c0 100644 --- a/tests/integration/src/sqlbuild/compiler/analysis_session/helpers.py +++ b/tests/integration/src/sqlbuild/compiler/analysis_session/helpers.py @@ -14,7 +14,9 @@ import pytest import sqlbuild._native as native_module +import sqlbuild.compiler.analysis_session.classes.native_model_analysis as native_model_analysis import sqlbuild.compiler.compile._helpers.assembly.project as project_assembly +import sqlbuild.compiler.compile._helpers.native_stages.assembly as native_stage_assembly from sqlbuild.adapter.contract.models import ExpressionInferenceProfile from sqlbuild.adapters.duckdb.classes.duckdb_adapter import DuckDbAdapter from sqlbuild.compiler.analysis_session.main._analyze_native_model_sql import ( @@ -35,6 +37,7 @@ from sqlbuild.compiler.compile.models import ( CompactLineageFacts, CompileAdapterContext, + CompiledLineageColumnFact, CompileProjectInputs, DynamicColumnContractProof, ModelSqlAnalysis, @@ -64,6 +67,17 @@ description: Customers feed. expression: >- (SELECT 10 AS customer_id, 'Ada' AS Customer_Name, 'east' AS region) + - name: raw_payments + description: Typed payments feed. + expression: >- + (SELECT 10 AS customer_id, 'east' AS region, CAST(5 AS DECIMAL(10, 2)) AS amount) + columns: + - name: customer_id + type: INTEGER + - name: region + type: VARCHAR + - name: amount + type: DECIMAL(10, 2) - name: raw_events description: Untyped events feed. columns: @@ -113,21 +127,38 @@ 'SELECT * FROM __source("raw_orders") o\nJOIN __source("raw_customers") c USING (customer_id)', ) _PIVOT_AGGREGATES: tuple[str, ...] = ("MAX", "MIN", "ANY_VALUE", "SUM") +_TYPED_PIVOT: str = ( + _PIVOT_HEADER.format(pivot="status", value="amount", aggregate="MAX") + + _PIVOT_BODIES[0].format(aggregate="MAX") + + "\n" +) +_TYPED_PIVOT_PASSTHROUGH: str = ( + _PIVOT_HEADER.format(pivot="status", value="amount", aggregate="MAX") + + 'SELECT * FROM __ref("status_amounts")\n' +) +_CONTRACT_CTE_MODELS: dict[str, str] = { + "models/marts/payments_by_customer.sql": ( + 'MODEL (\n description "Payments passed through a CTE",\n contract enforced,\n' + " columns (customer_id (type INTEGER), region (type VARCHAR)),\n);\n\n" + 'WITH base AS (SELECT customer_id, region FROM __source("raw_payments"))\n' + "SELECT customer_id, region FROM base\n" + ), + "models/marts/payment_totals.sql": ( + 'MODEL (\n description "Payment totals aggregated in a CTE",\n contract enforced,\n' + " columns (total_amount (type DECIMAL(38, 2)), region (type VARCHAR)),\n);\n\n" + "WITH base AS (\n SELECT SUM(amount) AS total_amount, region\n" + ' FROM __source("raw_payments")\n GROUP BY region\n)\n' + "SELECT total_amount, region FROM base\n" + ), +} def _pivot_models(rng: random.Random) -> dict[str, str]: """A typed pivot every passthrough reads, and seeded pivot shapes proven or refused.""" files: dict[str, str] = { - "models/marts/status_amounts.sql": _PIVOT_HEADER.format( - pivot="status", value="amount", aggregate="MAX" - ) - + _PIVOT_BODIES[0].format(aggregate="MAX") - + "\n", - "models/marts/status_passthrough.sql": _PIVOT_HEADER.format( - pivot="status", value="amount", aggregate="MAX" - ) - + 'SELECT * FROM __ref("status_amounts")\n', + "models/marts/status_amounts.sql": _TYPED_PIVOT, + "models/marts/status_passthrough.sql": _TYPED_PIVOT_PASSTHROUGH, } for index, body in enumerate(rng.sample(_PIVOT_BODIES, k=4)): aggregate: str = rng.choice(_PIVOT_AGGREGATES) @@ -256,6 +287,47 @@ def _subquery(rng: random.Random, inputs: list[_Relation]) -> tuple[str, str, tu ) +def pivot_project_files() -> dict[str, str]: + """A typed pivot with its passthrough and one plain model, for selective assembly.""" + + return { + "sqlbuild_project.toml": _PROJECT_TOML, + "sources/raw.yml": _SOURCES, + "models/marts/status_amounts.sql": _TYPED_PIVOT, + "models/marts/status_passthrough.sql": _TYPED_PIVOT_PASSTHROUGH, + "models/staging/orders_list.sql": ( + 'MODEL (description "Order ids");\n\nSELECT order_id FROM __source("raw_orders")\n' + ), + } + + +def native_pivot_proofs( + *, monkeypatch: pytest.MonkeyPatch +) -> list[DynamicColumnContractProof | None]: + """Record every standalone native pivot proof assembly takes, failing any wheel pivot proof.""" + + proofs: list[DynamicColumnContractProof | None] = [] + native_proof: Callable[..., DynamicColumnContractProof | None] = ( + native_stage_assembly.prove_native_dynamic_contract + ) + + def recorded(**keywords: Any) -> DynamicColumnContractProof | None: + proofs.append(native_proof(**keywords)) + return proofs[-1] + + python_proof: Callable[..., DynamicColumnContractProof | None] = ( + project_assembly.analyze_dynamic_column_contract + ) + + def wheel_not_expected(**keywords: Any) -> DynamicColumnContractProof | None: + assert not keywords["families"], "the Python wheel proved a pivot native should prove" + return python_proof(**keywords) + + monkeypatch.setattr(native_stage_assembly, "prove_native_dynamic_contract", recorded) + monkeypatch.setattr(project_assembly, "analyze_dynamic_column_contract", wheel_not_expected) + return proofs + + def generated_analysis_files(*, rng: random.Random, model_count: int) -> dict[str, str]: """Sources (typed, inferred and untyped), a pivot and models drawing on every template.""" @@ -263,6 +335,7 @@ def generated_analysis_files(*, rng: random.Random, model_count: int) -> dict[st "sqlbuild_project.toml": _PROJECT_TOML, "sources/raw.yml": _SOURCES, **_pivot_models(rng), + **_CONTRACT_CTE_MODELS, } relations: list[_Relation] = list(_INITIAL_RELATIONS) for index in range(model_count): @@ -302,6 +375,7 @@ class AnalysisParity: pivot_proofs: int = 0 standalone_proofs: int = 0 proven_pivots: int = 0 + native_enrichments: int = 0 def compare_analyses( @@ -380,7 +454,16 @@ def shapes_by_both( ) return python + native_lineage_facts: Callable[..., tuple[CompiledLineageColumnFact, ...]] = ( + native_model_analysis.lineage_facts + ) + + def counted_native_enrichment(rows: list[Any]) -> tuple[CompiledLineageColumnFact, ...]: + parity.native_enrichments += 1 + return native_lineage_facts(rows) + with monkeypatch.context() as patch: + patch.setattr(native_model_analysis, "lineage_facts", counted_native_enrichment) patch.setattr(project_assembly, "analyze_model_sql_by_engine", models_by_both) patch.setattr(project_assembly, "expression_source_shapes_by_engine", shapes_by_both) patch.setattr(project_assembly, "dynamic_column_contract_by_engine", proofs_by_both) diff --git a/tests/integration/src/sqlbuild/compiler/analysis_session/test_native_model_analysis.py b/tests/integration/src/sqlbuild/compiler/analysis_session/test_native_model_analysis.py index f21fe53fc9..4e9527ea64 100644 --- a/tests/integration/src/sqlbuild/compiler/analysis_session/test_native_model_analysis.py +++ b/tests/integration/src/sqlbuild/compiler/analysis_session/test_native_model_analysis.py @@ -14,12 +14,20 @@ analyze_native_model_sql, ) from sqlbuild.compiler.analysis_session.models import NativeModelAnalysisRequest -from sqlbuild.compiler.compile.models import CompileProjectInputs +from sqlbuild.compiler.compile._helpers.assembly.project import assemble_compiled_project +from sqlbuild.compiler.compile.models import ( + CompiledProject, + CompileProjectInputs, + DynamicColumnContractProof, +) +from sqlbuild.compiler.frontier.constants import COMPILER_ENGINE_ENV_VAR +from sqlbuild.compiler.frontier.types import CompilerEngine from sqlbuild.compiler.sql_analysis.constants import ANALYSIS_RECORD_DIR_ENV_VAR from tests.integration.src.sqlbuild.compiler.analysis_session._test_types import ( AnalysisFallbackTestCase, GeneratedAnalysisParityTestCase, SessionFailureTestCase, + StandalonePivotProofTestCase, ) from tests.integration.src.sqlbuild.compiler.analysis_session.helpers import ( AnalysisParity, @@ -30,6 +38,8 @@ deferral_kinds, failing_provide_sessions, generated_analysis_files, + native_pivot_proofs, + pivot_project_files, ) from tests.integration.src.sqlbuild.compiler.helpers import mismatches @@ -44,7 +54,7 @@ "test_case", [ GeneratedAnalysisParityTestCase( - description="stars, CTEs, set operations, untyped inputs, contracts and pivots", + description="stars, CTEs, set operations, untyped inputs, contracts, contract CTEs, pivots", seed=20261008, count=6, model_count=24, @@ -53,6 +63,8 @@ expected_minimum_expression_shapes=40, expected_minimum_pivot_proofs=80, expected_minimum_proven_pivots=12, + expected_maximum_enrichment_deferrals=120, + expected_minimum_native_enrichments=90, ) ], ids=lambda case: case.description, @@ -86,6 +98,11 @@ def test_given_generated_projects_when_analysing_natively_then_matches_python( assert parity.pivot_proofs >= test_case.expected_minimum_pivot_proofs assert parity.standalone_proofs >= test_case.expected_minimum_pivot_proofs assert parity.proven_pivots >= test_case.expected_minimum_proven_pivots + assert ( + kinds["analysis_session:input_enrichment"] + <= test_case.expected_maximum_enrichment_deferrals + ) + assert parity.native_enrichments >= test_case.expected_minimum_native_enrichments @pytest.mark.parametrize( @@ -175,5 +192,46 @@ def test_given_session_failure_after_deferrals_when_analysing_then_records_only_ assert deferral_kinds(record_dir) == Counter(test_case.expected_kinds) +@pytest.mark.parametrize( + "test_case", + [ + StandalonePivotProofTestCase( + description="a selection that leaves the pivots out of model analysis", + analysed_models=frozenset({"orders_list"}), + expected_native_proofs=2, + expected_proven_by_model={ + "orders_list": None, + "status_amounts": True, + "status_passthrough": True, + }, + ) + ], + ids=lambda case: case.description, +) +def test_given_unanalysed_pivot_model_when_assembling_then_native_proves_it_without_the_wheel( + test_case: StandalonePivotProofTestCase, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + inputs: CompileProjectInputs = compile_inputs( + project_dir=tmp_path / "project", files=pivot_project_files() + ) + monkeypatch.setenv(COMPILER_ENGINE_ENV_VAR, CompilerEngine.NATIVE_PREVIEW.value) + proofs: list[DynamicColumnContractProof | None] = native_pivot_proofs(monkeypatch=monkeypatch) + + project: CompiledProject = assemble_compiled_project( + inputs=inputs, + inference_profile=ExpressionInferenceProfile(sql_analysis_dialect="duckdb"), + analysis_model_names=test_case.analysed_models, + ) + + assert ( + all(proof is not None for proof in proofs), + len(proofs), + { + model.name: getattr(model.dynamic_column_contract, "output_proven", None) + for model in project.models + }, + ) == (True, test_case.expected_native_proofs, test_case.expected_proven_by_model) + + if __name__ == "__main__": pytest.main([__file__, "-vv"])