refactor(core): avoid throwaway ModelPlanNode build for aliased model scans

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,15 @@ 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 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<usize> =
+        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 +67,65 @@ 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,
+    /// 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,
+}
+
+impl TreeNodeRewriter for ModelRewriter<'_> {
+    type Node = LogicalPlan;
+
+    fn f_down(&mut self, plan: LogicalPlan) -> Result<Transformed<LogicalPlan>> {
+        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<Transformed<LogicalPlan>> {
+        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 +420,13 @@ impl ModelAnalyzeRule {
         scope_manager: &mut ScopeManager,
         current_scope_id: ScopeId,
     ) -> Result<Transformed<LogicalPlan>> {
-        plan.transform_up(&mut |plan| -> Result<Transformed<LogicalPlan>> {
-            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 +543,8 @@ impl ModelAnalyzeRule {
         required_fields: Vec<Expr>,
         original_table_scan: Option<LogicalPlan>,
     ) -> Result<ModelPlanNode> {
+        #[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 +649,75 @@ 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<Vec<Expr>> {
+        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.
+    fn shortcut_aliased_table_scan(
+        &self,
+        plan: LogicalPlan,
+        scope_manager: &mut ScopeManager,
+        current_scope_id: ScopeId,
+    ) -> Result<Transformed<LogicalPlan>> {
+        let LogicalPlan::SubqueryAlias(SubqueryAlias { input, alias, .. }) = &plan else {
+            return Ok(Transformed::no(plan));
+        };
+        let LogicalPlan::TableScan(scan) = input.as_ref() else {
+            return Ok(Transformed::no(plan));
+        };
+        if !belong_to_mdl(
+            &self.analyzed_wren_mdl.wren_mdl(),
+            scan.table_name.clone(),
+            Arc::clone(&self.session_state),
+        ) {
+            return Ok(Transformed::no(plan));
+        }
+        let Some(model) = self
+            .analyzed_wren_mdl
+            .wren_mdl
+            .get_model(scan.table_name.table())
+        else {
+            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),
+        });
+        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<AnalyzedWrenMDL>,
@@ -612,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<Expr> = 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,
@@ -671,23 +781,11 @@ impl ModelAnalyzeRule {
                     .wren_mdl()
                     .get_model(model_node.plan_name())
                 {
-                    let field: Vec<Expr> = 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 {

core/wren-core/core/src/mdl/mod.rs

@@ -677,6 +677,94 @@ 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::<Manifest>(&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,
+        )?);
+        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(
+            &ctx,
+            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?;
+        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(())
+    }
+
     #[cfg(feature = "multi-thread")]
     #[test]
     fn test_sync_transform() -> Result<()> {
@@ -3260,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")