refactor: read analysis natively, plan SQL tests on a native thread and trim hot spots

Cargo.toml

@@ -84,6 +84,8 @@ ruff_python_ast = "=0.0.4"
 ruff_python_parser = "=0.0.4"
 ruff_text_size = "=0.0.4"
 # Unix only: `sqb` shares the compute log prune lock with Python's `filelock` through `flock`.
+# Fast deterministic hasher for hot in-memory maps whose iteration order never reaches output.
+rustc-hash = "2"
 rustix = { version = "1.1.4", default-features = false, features = ["fs", "std"] }
 saphyr-parser = "0.1"
 regex = "1"

crates/sqlbuild-analysis/Cargo.toml

@@ -17,6 +17,7 @@ polyglot-sql.workspace = true
 polyglot-sql-function-catalogs.workspace = true
 rayon.workspace = true
 regex.workspace = true
+rustc-hash.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 sha2.workspace = true

crates/sqlbuild-analysis/src/compiler/_helpers/sql_tests/extraction.rs

@@ -39,6 +39,19 @@ const SCENARIO_CONTEXT: &str = "SQL scenario";
 const DECLARATION_CALLS: [&str; 4] = ["enum", "const", "var", "param"];
 /// Longest projection the alias help repeats verbatim.
 const ALIAS_HELP_EXPRESSION_LIMIT: usize = 80;
+const FROM_KEYWORD: &str = "FROM";
+const SELECT_LIST_END_KEYWORDS: [&str; 10] = [
+    FROM_KEYWORD,
+    "WHERE",
+    "GROUP",
+    "HAVING",
+    "QUALIFY",
+    "WINDOW",
+    "ORDER",
+    "LIMIT",
+    "OFFSET",
+    "FETCH",
+];
 const RESERVED_NAMES: [&str; 5] = [
     "__actual",
     "__expected__typed",
@@ -777,15 +790,7 @@ fn projection_names(
             }
         }
     }
-    let mut end = select_list_from(branch, select_end, syntax)?.unwrap_or(branch.len());
-    for keyword in [
-        "WHERE", "GROUP", "HAVING", "QUALIFY", "WINDOW", "ORDER", "LIMIT", "OFFSET", "FETCH",
-    ] {
-        if let Some(position) = find_top_level_clause_keyword(branch, select_end, keyword, syntax)?
-        {
-            end = end.min(position);
-        }
-    }
+    let end = select_list_end(branch, select_end, syntax)?.unwrap_or(branch.len());
     let expressions = split_top_level(&branch[select_end..end], b',', syntax)?;
     if expressions.is_empty() {
         return Err(
@@ -1379,20 +1384,32 @@ pub(crate) fn find_top_level_keyword(
     })
 }
 
-/// The top-level `FROM` that ends a select list, skipping the `FROM` of `IS [NOT] DISTINCT FROM`.
-fn select_list_from(
+/// The first top-level clause keyword ending a select list, skipping names and `IS DISTINCT FROM`.
+fn select_list_end(
     sql: &str,
     start: usize,
     syntax: &LexicalSyntax,
 ) -> Result<Option<usize>, String> {
-    let mut search_start = start;
-    while let Some(position) = find_top_level_clause_keyword(sql, search_start, "FROM", syntax)? {
-        if !trailing_distinct_operator(&sql[start..position], syntax)? {
-            return Ok(Some(position));
+    scan_code(sql, start, syntax, |index, depth| {
+        let Some(first) = byte_at(sql, index).filter(|_| depth == 0) else {
+            return Ok(CodeStep::Advance);
+        };
+        let first = first.to_ascii_uppercase();
+        let Some(keyword) = SELECT_LIST_END_KEYWORDS
+            .iter()
+            .filter(|keyword| keyword.as_bytes()[0] == first)
+            .find(|keyword| consume_keyword(sql, index, keyword).is_some())
+        else {
+            return Ok(CodeStep::Advance);
+        };
+        if previous_code_byte(sql, index, syntax)? == Some(b'.') {
+            return Ok(CodeStep::Advance);
         }
-        search_start = position + "FROM".len();
-    }
-    Ok(None)
+        if *keyword == FROM_KEYWORD && trailing_distinct_operator(&sql[start..index], syntax)? {
+            return Ok(CodeStep::Advance);
+        }
+        Ok(CodeStep::Stop(index))
+    })
 }
 
 /// Whether `prefix` ends with `IS DISTINCT` or `IS NOT DISTINCT`, ignoring comments.
@@ -1437,22 +1454,6 @@ fn trailing_word<'a>(
     Ok(word_start.map(|start| (&code[..start], &code[start..])))
 }
 
-fn find_top_level_clause_keyword(
-    sql: &str,
-    start: usize,
-    keyword: &str,
-    syntax: &LexicalSyntax,
-) -> Result<Option<usize>, String> {
-    let mut search_start = start;
-    while let Some(position) = find_top_level_keyword(sql, search_start, keyword, syntax)? {
-        if previous_code_byte(sql, position, syntax)? != Some(b'.') {
-            return Ok(Some(position));
-        }
-        search_start = position + keyword.len();
-    }
-    Ok(None)
-}
-
 fn previous_code_byte(sql: &str, end: usize, syntax: &LexicalSyntax) -> Result<Option<u8>, String> {
     let mut previous = None;
     let mut index = 0;

crates/sqlbuild-analysis/src/compiler/_helpers/sql_tests/helper_scope.rs

@@ -5,6 +5,7 @@ use std::collections::{BTreeMap, HashMap, HashSet};
 use sqlbuild_core::graph_walk::main::depth_first_postorder::depth_first_postorder;
 
 use polyglot_sql::Dialect;
+use rustc_hash::FxHashSet;
 
 use crate::compiler::_helpers::sql_tests::cte_rename::defined_cte_keys;
 use crate::compiler::_helpers::sql_tests::cte_slices::SliceDialect;
@@ -549,7 +550,7 @@ fn is_generated_cte_name(name: &str) -> bool {
 /// Distinct lowercase identifiers outside strings and comments, in first-appearance order.
 fn identifier_tokens(sql: &str, patterns: &SqlTestPatterns) -> Vec<String> {
     let protected = protected_ranges(&patterns.lexical, sql);
-    let mut seen: HashSet<String> = HashSet::new();
+    let mut seen: FxHashSet<String> = FxHashSet::default();
     let mut tokens: Vec<String> = Vec::new();
     let mut push = |token: &str| {
         let lowered = token.to_ascii_lowercase();

crates/sqlbuild-analysis/src/compiler/tests/test_sql_test_extraction_rules.rs

@@ -169,6 +169,22 @@ fn given_authored_sql_tests_when_extracting_then_each_rule_holds() {
             sql: "WITH __ref__orders AS (SELECT 1 AS id), __expected__orders AS (SELECT 'completed' IS DISTINCT FROM 'x' AS status, a IS NOT /* c */ DISTINCT FROM b AS same FROM __ref__orders) SELECT 1",
             expected_outcome: accepted_outcome(false),
         },
+        ExtractionRuleTestCase {
+            description: "clause keywords after a dot or inside parentheses do not end the select list",
+            mode: "model",
+            raw: false,
+            syntax: "generic",
+            sql: "WITH __ref__orders AS (SELECT 1 AS id), __expected__orders AS (SELECT t.order, t.limit AS l, (SELECT max(id) FROM __ref__orders WHERE id > 0 ORDER BY 1) AS m FROM __ref__orders t WHERE t.id > 0 ORDER BY 1) SELECT 1",
+            expected_outcome: accepted_outcome(false),
+        },
+        ExtractionRuleTestCase {
+            description: "the first top-level clause keyword ends a select list without FROM",
+            mode: "model",
+            raw: false,
+            syntax: "generic",
+            sql: "WITH __ref__orders AS (SELECT 1 AS id), __expected__orders AS (SELECT 1 AS id WHERE TRUE LIMIT 1) SELECT 1",
+            expected_outcome: accepted_outcome(false),
+        },
         ExtractionRuleTestCase {
             description: "IS DISTINCT FROM without an alias is rejected",
             mode: "model",

crates/sqlbuild-analysis/src/query_analysis/_helpers/borrowed_facts.rs

@@ -1,13 +1,15 @@
 //! Immutable CTE output graph over one borrowed syntax tree.
 
 use std::cell::Cell;
-use std::collections::{BTreeSet, HashMap, HashSet};
+use std::collections::BTreeSet;
 use std::rc::Rc;
 
 use polyglot_sql::expressions::{Column, Join, JoinKind, Select, Star, Values, With};
 use polyglot_sql::{
     DialectType, Expression, ExpressionWalk, ProjectionNullability, TransformKind, ValidationSchema,
 };
+use rustc_hash::{FxHashMap, FxHashSet};
+
 const SQL_WILDCARD: &str = "*";
 
 #[derive(Clone, Debug)]
@@ -22,14 +24,14 @@ pub(super) struct OutputFact {
 #[derive(Clone)]
 struct Source {
     outputs: Rc<Vec<OutputFact>>,
-    columns: Rc<HashMap<String, Option<usize>>>,
+    columns: Rc<FxHashMap<String, Option<usize>>>,
     physical: Option<String>,
     null_extended: bool,
 }
 
 impl Source {
     fn new(outputs: Vec<OutputFact>, physical: Option<String>) -> Self {
-        let mut columns: HashMap<String, Option<usize>> = HashMap::new();
+        let mut columns: FxHashMap<String, Option<usize>> = FxHashMap::default();
         for (index, output) in outputs.iter().enumerate() {
             let key = output.name.to_lowercase();
             let value = (!columns.contains_key(&key)).then_some(index);
@@ -44,13 +46,13 @@ impl Source {
     }
 }
 
-type Relations = HashMap<String, Source>;
+type Relations = FxHashMap<String, Source>;
 #[derive(Default)]
 struct Sources {
     items: Vec<(String, Source)>,
-    using_columns: HashMap<String, OutputFact>,
+    using_columns: FxHashMap<String, OutputFact>,
 }
-type Aliases = HashMap<String, Option<OutputFact>>;
+type Aliases = FxHashMap<String, Option<OutputFact>>;
 
 struct BindingProof {
     enabled: bool,
@@ -110,7 +112,7 @@ fn infer_with_binding_proof(
                 | DialectType::Databricks
         ),
     };
-    let mut relations = Relations::new();
+    let mut relations = Relations::default();
     if let Some(schema) = schema {
         for table in &schema.tables {
             let outputs: Vec<OutputFact> = table
@@ -147,7 +149,7 @@ fn query(expression: &Expression, inherited: &Relations, proof: &BindingProof) -
     ));
     let mut relations = inherited.clone();
     if let Some(with) = query_with(expression) {
-        let mut names: HashSet<String> = HashSet::new();
+        let mut names: FxHashSet<String> = FxHashSet::default();
         for cte in &with.ctes {
             if proof.enabled {
                 proof.require(!with.recursive && names.insert(cte.alias.name.to_lowercase()));
@@ -195,11 +197,19 @@ fn set_outputs(
     right: Vec<OutputFact>,
     by_name: bool,
 ) -> Vec<OutputFact> {
+    let mut right_by_name: FxHashMap<String, usize> = FxHashMap::default();
+    if by_name {
+        for (index, output) in right.iter().enumerate() {
+            right_by_name
+                .entry(output.name.to_ascii_lowercase())
+                .or_insert(index);
+        }
+    }
     for (index, output) in left.iter_mut().enumerate() {
         let other = if by_name {
-            right
-                .iter()
-                .find(|other| other.name.eq_ignore_ascii_case(&output.name))
+            right_by_name
+                .get(&output.name.to_ascii_lowercase())
+                .map(|position| &right[*position])
         } else {
             right.get(index)
         };
@@ -214,11 +224,12 @@ fn set_outputs(
         }
     }
     if by_name {
+        let mut left_names: FxHashSet<String> = left
+            .iter()
+            .map(|output| output.name.to_ascii_lowercase())
+            .collect();
         for mut output in right {
-            if !left
-                .iter()
-                .any(|other| other.name.eq_ignore_ascii_case(&output.name))
-            {
+            if left_names.insert(output.name.to_ascii_lowercase()) {
                 output.nullability = ProjectionNullability::Nullable;
                 left.push(output);
             }
@@ -273,7 +284,7 @@ fn select_outputs(select: &Select, relations: &Relations, proof: &BindingProof)
         sources = joined_sources(sources, join, relations, proof);
     }
     let mut outputs: Vec<OutputFact> = Vec::new();
-    let mut aliases = Aliases::new();
+    let mut aliases = Aliases::default();
     for projection in &select.expressions {
         let projection = unwrapped(projection);
         if let Some(star) = projection_star(projection) {
@@ -402,7 +413,7 @@ fn source(
         Expression::Lateral(lateral) if matches!(unwrapped(&lateral.this), Expression::Function(function) if function.name.eq_ignore_ascii_case("FLATTEN")) =>
         {
             proof.require(false);
-            let fact = expression_fact(&lateral.this, preceding, &Aliases::new());
+            let fact = expression_fact(&lateral.this, preceding, &Aliases::default());
             let names: Vec<String> = if lateral.column_aliases.is_empty() {
                 ["SEQ", "KEY", "PATH", "INDEX", "VALUE", "THIS"]
                     .into_iter()
@@ -438,7 +449,7 @@ fn values_outputs(values: &Values) -> Vec<OutputFact> {
     for row in &values.expressions {
         let mut facts: Vec<OutputFact> = Vec::new();
         for (index, expression) in row.expressions.iter().enumerate() {
-            let mut fact = expression_fact(expression, &Sources::default(), &Aliases::new());
+            let mut fact = expression_fact(expression, &Sources::default(), &Aliases::default());
             fact.name = values.column_aliases.get(index).map_or_else(
                 || format!("column{}", index + 1),
                 |alias| alias.name.clone(),
@@ -483,7 +494,7 @@ fn joined_sources(
         items: vec![(name.clone(), right.clone())],
         ..Default::default()
     };
-    let mut joined: HashMap<String, OutputFact> = HashMap::new();
+    let mut joined: FxHashMap<String, OutputFact> = FxHashMap::default();
     for identifier in &join.using {
         let column = Column {
             name: identifier.clone(),
@@ -493,8 +504,8 @@ fn joined_sources(
             span: None,
             inferred_type: None,
         };
-        let left = column_fact(&column, &sources, &Aliases::new());
-        let right = column_fact(&column, &right_scope, &Aliases::new());
+        let left = column_fact(&column, &sources, &Aliases::default());
+        let right = column_fact(&column, &right_scope, &Aliases::default());
         proof.require(
             left.as_ref().is_some_and(|fact| fact.resolved)
                 && right.as_ref().is_some_and(|fact| fact.resolved),
@@ -535,7 +546,7 @@ fn joined_sources(
     {
         proof.require(
             scalar_binding_expression(on)
-                && expression_fact(on, &sources, &Aliases::new()).resolved,
+                && expression_fact(on, &sources, &Aliases::default()).resolved,
         );
     }
     sources
@@ -708,7 +719,7 @@ fn nullability(
 
 fn star_outputs(star: &Star, sources: &Sources) -> Vec<OutputFact> {
     let mut outputs: Vec<OutputFact> = Vec::new();
-    let mut merged: HashSet<String> = HashSet::new();
+    let mut merged: FxHashSet<String> = FxHashSet::default();
     for (name, source) in &sources.items {
         if star
             .table
@@ -756,7 +767,7 @@ fn star_outputs(star: &Star, sources: &Sources) -> Vec<OutputFact> {
                 .find(|alias| alias.alias.name.eq_ignore_ascii_case(&output.name))
             {
                 let name = output.name.clone();
-                output = expression_fact(&alias.this, sources, &Aliases::new());
+                output = expression_fact(&alias.this, sources, &Aliases::default());
                 output.name = name;
             }
             if let Some((_, alias)) = star

crates/sqlbuild-analysis/src/query_analysis/tests/helpers.rs

@@ -542,6 +542,25 @@ pub(crate) fn borrowed_facts_preserve_union_dependencies_and_join_nullability()
         facts[0].nullability,
         polyglot_sql::ProjectionNullability::NonNull
     );
+    let expression = polyglot_sql::parse_one(
+        "SELECT order_id FROM orders UNION ALL BY NAME SELECT customer_id AS extra_id, customer_id AS ORDER_ID FROM customers",
+        polyglot_sql::DialectType::Snowflake,
+    ).expect("valid SQL");
+    let facts = crate::query_analysis::borrowed_facts::infer(
+        &expression,
+        Some(&schema),
+        polyglot_sql::DialectType::Snowflake,
+    );
+    assert_eq!(
+        facts
+            .iter()
+            .map(|fact| (fact.name.as_str(), fact.nullability, fact.upstream.len()))
+            .collect::<Vec<_>>(),
+        vec![
+            ("order_id", polyglot_sql::ProjectionNullability::NonNull, 2),
+            ("extra_id", polyglot_sql::ProjectionNullability::Nullable, 1),
+        ]
+    );
     let expression = polyglot_sql::parse_one(
         "SELECT o.order_id, c.customer_id FROM orders o LEFT JOIN customers c ON o.order_id = c.customer_id",
         polyglot_sql::DialectType::Snowflake,