From 3b33796032126fe9cb631b523205f41791eca5b0 Mon Sep 17 00:00:00 2001 From: Amir Fathi Date: Tue, 28 Jul 2026 18:34:18 +0000 Subject: [PATCH 1/2] refactor(core): avoid throwaway ModelPlanNode build for aliased model scans For a model scanned through a table alias (FROM customer AS e), the bottom-up ModelAnalyzeRule walk visits the bare TableScan before its SubqueryAlias parent. Required columns are keyed by the alias, so this first pass finds none and builds a ModelPlanNode via the wildcard- expansion path; analyze_subquery_alias_model then discards it and rebuilds the node correctly once the alias is known. The first build runs the full RLAC/CLAC parsing for nothing. PR #2449 already fixed the correctness bug this caused (the throwaway build was denying access to columns under column-level security instead of pruning them); the wasted build itself remained. This adds a pre-order TreeNodeRewriter pass that recognizes SubqueryAlias -> TableScan for a model directly and builds the ModelPlanNode once, with the alias's real required columns, skipping the discarded intermediate build. Everything else keeps using the original bottom-up analyze_model_internal pass. Fixes #2451 --- .../src/logical_plan/analyze/model_anlayze.rs | 172 +++++++++++++++--- core/wren-core/core/src/mdl/mod.rs | 41 +++++ 2 files changed, 187 insertions(+), 26 deletions(-) diff --git a/core/wren-core/core/src/logical_plan/analyze/model_anlayze.rs b/core/wren-core/core/src/logical_plan/analyze/model_anlayze.rs index 790df3995a..2f636a1393 100644 --- a/core/wren-core/core/src/logical_plan/analyze/model_anlayze.rs +++ b/core/wren-core/core/src/logical_plan/analyze/model_anlayze.rs @@ -4,7 +4,9 @@ use crate::logical_plan::utils::{belong_to_mdl, expr_to_columns}; use crate::mdl::context::SessionPropertiesRef; use crate::mdl::utils::quoted; use crate::mdl::{AnalyzedWrenMDL, Dataset, SessionStateRef}; -use datafusion::common::tree_node::{Transformed, TransformedResult, TreeNode}; +use datafusion::common::tree_node::{ + Transformed, TransformedResult, TreeNode, TreeNodeRecursion, TreeNodeRewriter, +}; use datafusion::common::{internal_err, plan_err, Column, DFSchemaRef, Result, Spans}; use datafusion::config::ConfigOptions; use datafusion::error::DataFusionError; @@ -21,6 +23,14 @@ use std::collections::HashSet; use std::fmt::Debug; use std::sync::Arc; +#[cfg(test)] +thread_local! { + /// Counts [`ModelAnalyzeRule::build_model_plan_node`] calls; thread-local since tests run + /// each on its own thread. + pub(crate) static BUILD_MODEL_PLAN_NODE_CALLS: std::cell::Cell = + const { std::cell::Cell::new(0) }; +} + /// [ModelAnalyzeRule] responsible for analyzing the model plan node. Turn TableScan from a model to a ModelPlanNode. /// We collect the required fields from the projection, filter, aggregation, and join, /// and pass them to the ModelPlanNode. @@ -56,6 +66,62 @@ impl Drop for ModelStackGuard { } } +/// Drives [`ModelAnalyzeRule::analyze_model`], trying +/// [`ModelAnalyzeRule::shortcut_aliased_table_scan`] before the original +/// [`ModelAnalyzeRule::analyze_model_internal`] pass. +struct ModelRewriter<'a> { + rule: &'a ModelAnalyzeRule, + scope_manager: &'a mut ScopeManager, + current_scope_id: ScopeId, + shortcut_taken: bool, +} + +impl TreeNodeRewriter for ModelRewriter<'_> { + type Node = LogicalPlan; + + fn f_down(&mut self, plan: LogicalPlan) -> Result> { + let result = self.rule.shortcut_aliased_table_scan( + plan, + self.scope_manager, + self.current_scope_id, + )?; + self.shortcut_taken = result.tnr == TreeNodeRecursion::Jump; + Ok(result) + } + + fn f_up(&mut self, plan: LogicalPlan) -> Result> { + let plan = if std::mem::take(&mut self.shortcut_taken) { + plan + } else { + self.rule + .analyze_model_internal(plan, self.scope_manager, self.current_scope_id)? + .data + }; + // If the plan contains subquery, we should analyze the subquery recursively + plan.map_subqueries(|plan| { + if let LogicalPlan::Subquery(subquery) = &plan { + let root_scope = + self.scope_manager.get_scope_mut(self.current_scope_id)?; + let Some(child_scope_id) = root_scope.pop_child_scope() else { + return internal_err!("No child scope found for subquery"); + }; + let transformed = self + .rule + .analyze_model( + Arc::unwrap_or_clone(Arc::clone(&subquery.subquery)), + self.scope_manager, + child_scope_id, + )? + .data; + return Ok(Transformed::yes(LogicalPlan::Subquery( + subquery.with_plan(Arc::new(transformed)), + ))); + } + Ok(Transformed::no(plan)) + }) + } +} + impl Debug for ModelAnalyzeRule { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("ModelAnalyzeRule").finish() @@ -350,31 +416,13 @@ impl ModelAnalyzeRule { scope_manager: &mut ScopeManager, current_scope_id: ScopeId, ) -> Result> { - plan.transform_up(&mut |plan| -> Result> { - let plan = self - .analyze_model_internal(plan, scope_manager, current_scope_id)? - .data; - // If the plan contains subquery, we should analyze the subquery recursively - plan.map_subqueries(|plan| { - if let LogicalPlan::Subquery(subquery) = &plan { - let root_scope = scope_manager.get_scope_mut(current_scope_id)?; - let Some(child_scope_id) = root_scope.pop_child_scope() else { - return internal_err!("No child scope found for subquery"); - }; - let transformed = self - .analyze_model( - Arc::unwrap_or_clone(Arc::clone(&subquery.subquery)), - scope_manager, - child_scope_id, - )? - .data; - return Ok(Transformed::yes(LogicalPlan::Subquery( - subquery.with_plan(Arc::new(transformed)), - ))); - } - Ok(Transformed::no(plan)) - }) - }) + let mut rewriter = ModelRewriter { + rule: self, + scope_manager, + current_scope_id, + shortcut_taken: false, + }; + plan.rewrite(&mut rewriter) } /// Analyze the model and generate the ModelPlanNode @@ -491,6 +539,8 @@ impl ModelAnalyzeRule { required_fields: Vec, original_table_scan: Option, ) -> Result { + #[cfg(test)] + BUILD_MODEL_PLAN_NODE_CALLS.with(|c| c.set(c.get() + 1)); let model_name = model.name().to_string(); { let mut stack = self.building_models.lock(); @@ -595,6 +645,76 @@ impl ModelAnalyzeRule { .data() } + /// Pre-order shortcut for `SubqueryAlias -> TableScan` (e.g. `FROM a_model AS alias`). + /// Builds the `ModelPlanNode` directly with the alias's own required columns, + /// mirroring [`Self::analyze_subquery_alias_model`], instead of the bottom-up pass. + fn shortcut_aliased_table_scan( + &self, + plan: LogicalPlan, + scope_manager: &mut ScopeManager, + current_scope_id: ScopeId, + ) -> Result> { + let LogicalPlan::SubqueryAlias(SubqueryAlias { input, alias, .. }) = plan else { + return Ok(Transformed::no(plan)); + }; + if !matches!(input.as_ref(), LogicalPlan::TableScan(_)) { + return Ok(Transformed::no(LogicalPlan::SubqueryAlias( + SubqueryAlias::try_new(input, alias)?, + ))); + } + let LogicalPlan::TableScan(table_scan) = Arc::unwrap_or_clone(input) else { + unreachable!("just matched TableScan above") + }; + if !belong_to_mdl( + &self.analyzed_wren_mdl.wren_mdl(), + table_scan.table_name.clone(), + Arc::clone(&self.session_state), + ) { + return Ok(Transformed::no(LogicalPlan::SubqueryAlias( + SubqueryAlias::try_new( + Arc::new(LogicalPlan::TableScan(table_scan)), + alias, + )?, + ))); + } + let Some(model) = self + .analyzed_wren_mdl + .wren_mdl + .get_model(table_scan.table_name.table()) + else { + return Ok(Transformed::no(LogicalPlan::SubqueryAlias( + SubqueryAlias::try_new( + Arc::new(LogicalPlan::TableScan(table_scan)), + alias, + )?, + ))); + }; + let field: Vec = if let Some(used_columns) = + scope_manager.try_get_required_columns(current_scope_id, &alias) + { + used_columns.iter().cloned().collect() + } else { + // If the required columns are not found in the current scope but the table is visited, + // it could be a count(*) query + if scope_manager + .try_get_visited_dataset(current_scope_id, &alias) + .is_none() + { + return internal_err!( + "Table {} not found in the visited dataset and required columns map", + alias + ); + }; + vec![] + }; + let model_plan_node = self.build_model_plan_node(model, field, None)?; + let model_plan = LogicalPlan::Extension(Extension { + node: Arc::new(model_plan_node), + }); + let subquery = LogicalPlanBuilder::from(model_plan).alias(alias)?.build()?; + Ok(Transformed::new(subquery, true, TreeNodeRecursion::Jump)) + } + fn analyze_table_scan( &self, analyzed_wren_mdl: Arc, diff --git a/core/wren-core/core/src/mdl/mod.rs b/core/wren-core/core/src/mdl/mod.rs index c1c07e501b..7f33ebe2e0 100644 --- a/core/wren-core/core/src/mdl/mod.rs +++ b/core/wren-core/core/src/mdl/mod.rs @@ -677,6 +677,47 @@ mod test { ColumnLevelOperator, DataSource, JoinType, RelationshipBuilder, SessionProperty, }; + /// An aliased model scan should build its `ModelPlanNode` once, not once as a + /// discarded wildcard build and once for real. + #[tokio::test] + async fn test_aliased_model_scan_builds_model_plan_node_once() -> Result<()> { + let test_data: PathBuf = + [env!("CARGO_MANIFEST_DIR"), "tests", "data", "mdl.json"] + .iter() + .collect(); + let mdl_json = fs::read_to_string(test_data.as_path())?; + let mdl = match serde_json::from_str::(&mdl_json) { + Ok(mdl) => mdl, + Err(e) => return not_impl_err!("Failed to parse mdl json: {}", e), + }; + let analyzed_mdl = Arc::new(AnalyzedWrenMDL::analyze( + mdl, + Arc::new(HashMap::default()), + Mode::Unparse, + )?); + + crate::logical_plan::analyze::model_anlayze::BUILD_MODEL_PLAN_NODE_CALLS + .with(|c| c.set(0)); + let actual = mdl::transform_sql_with_ctx( + &create_wren_ctx(None, analyzed_mdl.wren_mdl().data_source().as_ref()), + Arc::clone(&analyzed_mdl), + &[], + Arc::new(HashMap::new()), + "select e.c_custkey from test.test.customer e", + ) + .await?; + assert_eq!( + crate::logical_plan::analyze::model_anlayze::BUILD_MODEL_PLAN_NODE_CALLS + .with(|c| c.get()), + 1, + "an aliased model scan should build its ModelPlanNode once, not build a \ + throwaway wildcard node first" + ); + assert_sql_valid_executable(&actual).await?; + + Ok(()) + } + #[cfg(feature = "multi-thread")] #[test] fn test_sync_transform() -> Result<()> { From fd2665d2f71c3df384e4644b5c1fa4e5d83167b1 Mon Sep 17 00:00:00 2001 From: Amir Fathi Date: Mon, 3 Aug 2026 12:16:46 +0000 Subject: [PATCH 2/2] refactor(core): address review feedback on aliased model scan shortcut - Update the stale comment on test_clac_unreferenced_column_pruned_not_denied: it described the throwaway-plan mechanism this PR deletes. - Add snapshot assertions (assert_snapshot!) to the new test so it asserts the no-behavior-change claim, not just the build-call count. - Cover the two branches the new test missed: count(*) on an aliased model (the visited_dataset fallback) and a join of two aliased models (shortcut_taken across siblings). - Drop the needless SubqueryAlias rebuild on non-matching paths in shortcut_aliased_table_scan; match on &plan first and return the original Transformed::no(plan) on every early exit. - Comment the Jump/f_up invariant shortcut_taken depends on. - Extract the required-columns/visited_dataset-fallback lookup duplicated across shortcut_aliased_table_scan, analyze_table_scan and analyze_subquery_alias_model into resolve_required_fields. - Fix the counter's doc comment: the real dependency is #[tokio::test]'s current-thread runtime, not "one thread per test". Addresses goldmedal's review on PR #2612. --- .../src/logical_plan/analyze/model_anlayze.rs | 132 ++++++++---------- core/wren-core/core/src/mdl/mod.rs | 58 +++++++- 2 files changed, 106 insertions(+), 84 deletions(-) diff --git a/core/wren-core/core/src/logical_plan/analyze/model_anlayze.rs b/core/wren-core/core/src/logical_plan/analyze/model_anlayze.rs index 2f636a1393..9045534113 100644 --- a/core/wren-core/core/src/logical_plan/analyze/model_anlayze.rs +++ b/core/wren-core/core/src/logical_plan/analyze/model_anlayze.rs @@ -25,8 +25,9 @@ use std::sync::Arc; #[cfg(test)] thread_local! { - /// Counts [`ModelAnalyzeRule::build_model_plan_node`] calls; thread-local since tests run - /// each on its own thread. + /// Counts [`ModelAnalyzeRule::build_model_plan_node`] calls; thread-local because + /// `#[tokio::test]`'s default current-thread runtime keeps caller and counter on + /// the same thread (a `flavor = "multi_thread"` test would need something else). pub(crate) static BUILD_MODEL_PLAN_NODE_CALLS: std::cell::Cell = const { std::cell::Cell::new(0) }; } @@ -73,6 +74,9 @@ struct ModelRewriter<'a> { rule: &'a ModelAnalyzeRule, scope_manager: &'a mut ScopeManager, current_scope_id: ScopeId, + /// Set by `f_down` when the pre-order shortcut fires, so `f_up` skips the + /// bottom-up pass for that node. Correct only because a `Jump` from `f_down` + /// still runs `f_up` for the same node (`transform_children` turns `Jump` into `Continue`). shortcut_taken: bool, } @@ -645,6 +649,32 @@ impl ModelAnalyzeRule { .data() } + /// Resolve the required fields for `table_ref` in `current_scope_id`: the + /// recorded required columns if any, `vec![]` if the table was only visited + /// (e.g. `count(*)`), or an error if neither is present. + fn resolve_required_fields( + &self, + scope_manager: &ScopeManager, + current_scope_id: ScopeId, + table_ref: &TableReference, + ) -> Result> { + if let Some(used_columns) = + scope_manager.try_get_required_columns(current_scope_id, table_ref) + { + Ok(used_columns.iter().cloned().collect()) + } else if scope_manager + .try_get_visited_dataset(current_scope_id, table_ref) + .is_some() + { + Ok(vec![]) + } else { + internal_err!( + "Table {} not found in the visited dataset and required columns map", + table_ref + ) + } + } + /// Pre-order shortcut for `SubqueryAlias -> TableScan` (e.g. `FROM a_model AS alias`). /// Builds the `ModelPlanNode` directly with the alias's own required columns, /// mirroring [`Self::analyze_subquery_alias_model`], instead of the bottom-up pass. @@ -654,59 +684,32 @@ impl ModelAnalyzeRule { scope_manager: &mut ScopeManager, current_scope_id: ScopeId, ) -> Result> { - let LogicalPlan::SubqueryAlias(SubqueryAlias { input, alias, .. }) = plan else { + let LogicalPlan::SubqueryAlias(SubqueryAlias { input, alias, .. }) = &plan else { return Ok(Transformed::no(plan)); }; - if !matches!(input.as_ref(), LogicalPlan::TableScan(_)) { - return Ok(Transformed::no(LogicalPlan::SubqueryAlias( - SubqueryAlias::try_new(input, alias)?, - ))); - } - let LogicalPlan::TableScan(table_scan) = Arc::unwrap_or_clone(input) else { - unreachable!("just matched TableScan above") + let LogicalPlan::TableScan(scan) = input.as_ref() else { + return Ok(Transformed::no(plan)); }; if !belong_to_mdl( &self.analyzed_wren_mdl.wren_mdl(), - table_scan.table_name.clone(), + scan.table_name.clone(), Arc::clone(&self.session_state), ) { - return Ok(Transformed::no(LogicalPlan::SubqueryAlias( - SubqueryAlias::try_new( - Arc::new(LogicalPlan::TableScan(table_scan)), - alias, - )?, - ))); + return Ok(Transformed::no(plan)); } let Some(model) = self .analyzed_wren_mdl .wren_mdl - .get_model(table_scan.table_name.table()) + .get_model(scan.table_name.table()) else { - return Ok(Transformed::no(LogicalPlan::SubqueryAlias( - SubqueryAlias::try_new( - Arc::new(LogicalPlan::TableScan(table_scan)), - alias, - )?, - ))); - }; - let field: Vec = if let Some(used_columns) = - scope_manager.try_get_required_columns(current_scope_id, &alias) - { - used_columns.iter().cloned().collect() - } else { - // If the required columns are not found in the current scope but the table is visited, - // it could be a count(*) query - if scope_manager - .try_get_visited_dataset(current_scope_id, &alias) - .is_none() - { - return internal_err!( - "Table {} not found in the visited dataset and required columns map", - alias - ); - }; - vec![] + return Ok(Transformed::no(plan)); }; + // original_table_scan is None below, so `input` is never consumed: clone just + // the alias instead of re-matching `plan` to take ownership of it. + let alias = alias.clone(); + + let field = + self.resolve_required_fields(scope_manager, current_scope_id, &alias)?; let model_plan_node = self.build_model_plan_node(model, field, None)?; let model_plan = LogicalPlan::Extension(Extension { node: Arc::new(model_plan_node), @@ -732,24 +735,11 @@ impl ModelAnalyzeRule { let table_name = table_scan.table_name.table(); if let Some(model) = analyzed_wren_mdl.wren_mdl.get_model(table_name) { let table_ref = alias.unwrap_or(table_scan.table_name.clone()); - let field: Vec = if let Some(used_columns) = - scope_manager.try_get_required_columns(current_scope_id, &table_ref) - { - used_columns.iter().cloned().collect() - } else { - // If the required columns are not found in the current scope but the table is visited, - // it could be a count(*) query - if scope_manager - .try_get_visited_dataset(current_scope_id, &table_ref) - .is_none() - { - return internal_err!( - "Table {} not found in the visited dataset and required columns map", - table_ref - ); - }; - vec![] - }; + let field = self.resolve_required_fields( + scope_manager, + current_scope_id, + &table_ref, + )?; let model_plan_node = self.build_model_plan_node( Arc::clone(&model), field, @@ -791,23 +781,11 @@ impl ModelAnalyzeRule { .wren_mdl() .get_model(model_node.plan_name()) { - let field: Vec = if let Some(used_columns) = - scope_manager.try_get_required_columns(current_scope_id, &alias) - { - used_columns.iter().cloned().collect() - } else { - // If the required columns are not found in the current scope but the table is visited, - // it could be a count(*) query - if scope_manager - .try_get_visited_dataset(current_scope_id, &alias) - .is_none() - { - return internal_err!( - "Table {} not found in the visited dataset and required columns map", - alias); - }; - vec![] - }; + let field = self.resolve_required_fields( + scope_manager, + current_scope_id, + &alias, + )?; let model_plan_node = self.build_model_plan_node(Arc::clone(&model), field, None)?; let model_plan = LogicalPlan::Extension(Extension { diff --git a/core/wren-core/core/src/mdl/mod.rs b/core/wren-core/core/src/mdl/mod.rs index 7f33ebe2e0..5536a25d7c 100644 --- a/core/wren-core/core/src/mdl/mod.rs +++ b/core/wren-core/core/src/mdl/mod.rs @@ -695,11 +695,12 @@ mod test { Arc::new(HashMap::default()), Mode::Unparse, )?); + let ctx = create_wren_ctx(None, analyzed_mdl.wren_mdl().data_source().as_ref()); crate::logical_plan::analyze::model_anlayze::BUILD_MODEL_PLAN_NODE_CALLS .with(|c| c.set(0)); let actual = mdl::transform_sql_with_ctx( - &create_wren_ctx(None, analyzed_mdl.wren_mdl().data_source().as_ref()), + &ctx, Arc::clone(&analyzed_mdl), &[], Arc::new(HashMap::new()), @@ -714,6 +715,52 @@ mod test { throwaway wildcard node first" ); assert_sql_valid_executable(&actual).await?; + assert_snapshot!(actual, @"SELECT e.c_custkey FROM (SELECT customer.c_custkey FROM (SELECT __source.c_custkey AS c_custkey FROM customer AS __source) AS customer) AS e"); + + // count(*) on an aliased model has no required columns, so it's the only + // remaining path that reaches the `visited_dataset` fallback (`field = + // vec![]`) for an aliased scan; it must still build once, not twice. + crate::logical_plan::analyze::model_anlayze::BUILD_MODEL_PLAN_NODE_CALLS + .with(|c| c.set(0)); + let actual = mdl::transform_sql_with_ctx( + &ctx, + Arc::clone(&analyzed_mdl), + &[], + Arc::new(HashMap::new()), + "select count(*) from test.test.customer e", + ) + .await?; + assert_eq!( + crate::logical_plan::analyze::model_anlayze::BUILD_MODEL_PLAN_NODE_CALLS + .with(|c| c.get()), + 1, + "count(*) on an aliased model should still build its ModelPlanNode once" + ); + assert_sql_valid_executable(&actual).await?; + assert_snapshot!(actual, @r#"SELECT count(1) AS "count(*)" FROM (SELECT __source.c_custkey AS c_custkey, __source.c_name AS c_name FROM customer AS __source) AS e"#); + + // A join of two aliased models is the only thing that exercises + // `shortcut_taken` across siblings: `ModelRewriter` carries one `bool` field + // for the whole walk, and it must reflect only the child just visited, never + // leak from the left join input into the right's `f_up`. + crate::logical_plan::analyze::model_anlayze::BUILD_MODEL_PLAN_NODE_CALLS + .with(|c| c.set(0)); + let actual = mdl::transform_sql_with_ctx( + &ctx, + Arc::clone(&analyzed_mdl), + &[], + Arc::new(HashMap::new()), + "select e.c_custkey, o.o_orderkey from test.test.customer e join test.test.orders o on e.c_custkey = o.o_custkey", + ) + .await?; + assert_eq!( + crate::logical_plan::analyze::model_anlayze::BUILD_MODEL_PLAN_NODE_CALLS + .with(|c| c.get()), + 2, + "a join of two aliased models should build each ModelPlanNode once" + ); + assert_sql_valid_executable(&actual).await?; + assert_snapshot!(actual, @"SELECT e.c_custkey, o.o_orderkey FROM (SELECT customer.c_custkey FROM (SELECT __source.c_custkey AS c_custkey FROM customer AS __source) AS customer) AS e INNER JOIN (SELECT orders.o_custkey, orders.o_orderkey FROM (SELECT __source.o_custkey AS o_custkey, __source.o_orderkey AS o_orderkey FROM orders AS __source) AS orders) AS o ON e.c_custkey = o.o_custkey"); Ok(()) } @@ -3301,12 +3348,9 @@ mod test { #[tokio::test] async fn test_clac_unreferenced_column_pruned_not_denied() -> Result<()> { - // A CLS-protected column the query does NOT reference must be pruned, - // not denied — even when the model is referenced with a table alias - // (which builds a throwaway inner plan with empty required fields, so the - // model gets wildcard-expanded), and for `count(*)`. Before the fix, an - // aliased scan of `customer` — or `count(*)` — that never selected - // `c_name` still denied `customer.c_name`. + // A CLS-protected column the query does NOT reference must be pruned, not + // denied, in every scan shape (bare, aliased, count(*), wildcard). Before + // the fix, an aliased or count(*) scan of `customer` still denied `c_name`. let ctx = create_wren_ctx(None, None); let manifest = ManifestBuilder::new() .catalog("wren")