refactor(core): avoid throwaway ModelPlanNode build for aliased model scans - #2612
refactor(core): avoid throwaway ModelPlanNode build for aliased model scans#2612AmirF194 wants to merge 2 commits into
Conversation
… 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 Canner#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 Canner#2451
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (2)
WalkthroughModel analysis now uses a pre-order tree rewriter to shortcut aliased model scans, resolve required fields through one helper, construct ChangesModel scan rewrite
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant LogicalPlan
participant ModelRewriter
participant RequiredFieldResolver
participant ModelPlanNode
participant RegressionTests
LogicalPlan->>ModelRewriter: rewrite model plan
ModelRewriter->>RequiredFieldResolver: resolve aliased required fields
RequiredFieldResolver-->>ModelRewriter: return fields or empty visited-only set
ModelRewriter->>ModelPlanNode: build aliased model node
ModelPlanNode-->>ModelRewriter: return rewritten plan
RegressionTests->>ModelRewriter: analyze column, count(*), or join query
ModelRewriter-->>RegressionTests: return executable transformed SQL
Possibly related PRs
Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
core/wren-core/core/src/logical_plan/analyze/model_anlayze.rs (1)
648-717: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRequired-fields resolution logic duplicated a third time.
The
try_get_required_columns/try_get_visited_datasetfallback /internal_err!block here (Line 692-709) is now duplicated almost verbatim inanalyze_table_scan(Line 735-752) andanalyze_subquery_alias_model(Line 794-810). Consider extracting a shared helper (e.g.fn resolve_required_fields(&self, scope_manager, current_scope_id, table_ref) -> Result<Vec<Expr>>) to avoid a third copy of this logic drifting out of sync.♻️ Sketch of a shared helper
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 ) } }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@core/wren-core/core/src/logical_plan/analyze/model_anlayze.rs` around lines 648 - 717, Extract the duplicated required-fields resolution from shortcut_aliased_table_scan, analyze_table_scan, and analyze_subquery_alias_model into a shared resolve_required_fields helper using the existing ScopeManager lookups and internal_err! fallback. Replace each inline block with calls to the helper, passing the appropriate scope, scope ID, and table reference while preserving the current Vec<Expr> results and error behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@core/wren-core/core/src/logical_plan/analyze/model_anlayze.rs`:
- Around line 648-717: Extract the duplicated required-fields resolution from
shortcut_aliased_table_scan, analyze_table_scan, and
analyze_subquery_alias_model into a shared resolve_required_fields helper using
the existing ScopeManager lookups and internal_err! fallback. Replace each
inline block with calls to the helper, passing the appropriate scope, scope ID,
and table reference while preserving the current Vec<Expr> results and error
behavior.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 1b1c6b52-1854-4f02-8358-ea07c25ff9d2
📒 Files selected for processing (2)
core/wren-core/core/src/logical_plan/analyze/model_anlayze.rscore/wren-core/core/src/mdl/mod.rs
goldmedal
left a comment
goldmedal
left a comment
There was a problem hiding this comment.
Sound change — verified no behavior change and the build-count win holds. I diffed generated SQL across 19 plan shapes against b590b650 (identical), instrumented build counts (halved for aliased scans, never increased), and ran all CI gates locally including core-py-ci. Comments below are cleanups, none blocking.
Please address in this PR
1. test_clac_unreferenced_column_pruned_not_denied — update the stale comment. It documents the exact mechanism this PR deletes:
"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)"
After this change the aliased-select half no longer wildcard-expands — only the count(*) half does. This is the PR that invalidates the comment, and it's the #2449 guard, so it shouldn't be left describing gone behavior.
2. Add a snapshot assertion to the new test. assert_sql_valid_executable only proves the SQL parses and runs — it would pass on a plan that silently pruned or added a column, which is the actual risk being claimed against. The file uses assert_snapshot! throughout; one here would make the test assert the no-behavior-change claim rather than just the call count.
3. Cover the two branches the new test misses. Both are one-liners with the existing helpers:
count(*)on an aliased model — the only path reaching thevisited_datasetfallback (field = vec![]), and the only remaining wildcard-expansion path for aliased scans.- A join of two aliased models — the only thing exercising
shortcut_takenacross siblings.
Optional
4. Drop the needless SubqueryAlias rebuild on non-matching paths. The three early returns do Transformed::no(SubqueryAlias::try_new(input, alias)?), so every non-model SubqueryAlias in every plan pays a schema recompute in f_down to hand back an unchanged node. Match on &plan first:
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, so input never needs consuming — this also removes the unreachable!() and the Arc::unwrap_or_clone.
5. Comment the invariant shortcut_taken depends on. It's correct only because Jump from f_down still runs f_up for that node (transform_children rewrites Jump→Continue, then transform_parent calls f_up). I confirmed this holds in the pinned DataFusion, but if a bump changed it the failure is silent — a parent's analyze_model_internal gets skipped. One line on the field would make the coupling explicit.
6. Extract the block duplicated from analyze_subquery_alias_model. The required-columns lookup + visited_dataset fallback + build + alias-wrap is copied verbatim, including the internal_err! text. A helper over (model, alias, scope_manager, current_scope_id) would stop the two drifting.
7. Nit. The counter's doc comment explains per-test threads, but the real dependency is #[tokio::test]'s current-thread runtime — under flavor = "multi_thread" the increments land elsewhere and the assert reads 0.
- 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 Canner#2612.
|
Thanks for the thorough review, especially catching the stale comment and the missing branches. Pushed 1. Stale comment. Updated 2. Snapshot assertion. Added 3. Missing branches. Added the two cases to the same test: 4. SubqueryAlias rebuild. Applied your replacement in 5. 6. Extracted the duplicated block. Pulled the required-columns/ 7. Counter doc comment. Fixed to name Re-verified after all changes: |
2 participants
Summary
For
FROM customer AS e, the bottom-upModelAnalyzeRulewalk visits the bareTableScanbefore itsSubqueryAliasparent, so it looks up required columnsunder the alias, finds none, and builds a full wildcard-expanded
ModelPlanNodethat
analyze_subquery_alias_modelimmediately discards and rebuilds correctlyonce the alias is known. This skips that first, discarded build for the common
SubqueryAlias -> TableScanshape, via a pre-order rewriter pass that builds thenode once with the alias's real required columns.
No behavior change: same plan, same SQL output, one fewer RLAC/CLAC parse per
aliased model scan.
refactor:, notfix:.What failure does this repair?
N/A, no failure. This is a
refactor:: it removes a wasted build, it does notchange the query result. PR #2449 already fixed the correctness bug the
throwaway build used to cause (wildcard-expansion access denial on unreferenced
CLS columns); that fix is untouched by this change, it lives in
ModelSourceNode::new(plan.rs), which this PR does not modify.How is it tested?
test_aliased_model_scan_builds_model_plan_node_once(
core/src/mdl/mod.rs) asserts, via a#[cfg(test)]call counter onbuild_model_plan_node, that an aliased model scan builds it once. Confirmedit fails on current
main(counts 2) and passes with this change.wren-semantic-coresuite (148 tests, includingsqllogictest'smodel.slt/type.slt/view.slt/tpch.slt) passes unchanged.cargo check --all-targets,cargo clippy --all-targets --all-features -- -D warnings, andcargo fmt --checkare all clean.build call, not a measured latency number.
Duplicate check
Searched open PRs (
is:pr is:open) forModelPlanNode,analyze_model, andthis issue number in the body: none touch
model_anlayze.rsor reference#2451.
git log -Gon the touched functions shows no prior attempt at thisspecific change.
Fixes #2451
Summary by CodeRabbit
Bug Fixes
count(*)queries.Tests