diff --git a/core/wren-core/core/src/logical_plan/analyze/plan.rs b/core/wren-core/core/src/logical_plan/analyze/plan.rs index d84b72603b..8e6ef0e0d1 100644 --- a/core/wren-core/core/src/logical_plan/analyze/plan.rs +++ b/core/wren-core/core/src/logical_plan/analyze/plan.rs @@ -334,8 +334,19 @@ impl ModelPlanNodeBuilder { )); } - self.directed_graph - .add_node(Dataset::Model(Arc::clone(&model))); + // Ensure the base model is present as a node. If any calculated-column + // sub-graph already merged the base model in (the common case when the + // model has calculated columns), reuse it instead of adding a + // duplicate — duplicate base-model nodes break the join-chain walk in + // `RelationChain::with_chain`. + let base_dataset = Dataset::Model(Arc::clone(&model)); + if !self + .directed_graph + .node_indices() + .any(|idx| self.directed_graph[idx] == base_dataset) + { + self.directed_graph.add_node(base_dataset); + } if !is_dag(&self.directed_graph) { return plan_err!("cyclic dependency detected: {}", model.name()); } @@ -901,10 +912,29 @@ fn merge_graph( graph: &mut Graph, new_graph: &Graph, ) -> Result<()> { + // Build a lookup of datasets already present in the destination graph so + // that nodes representing the same `Dataset` (typically the shared base + // model when merging multiple calculated-column sub-graphs) are reused + // instead of duplicated. Without this, `RelationChain::with_chain` walks + // the merged graph via `find_edge(start, next)` and silently truncates the + // chain at the first duplicate base-model node, dropping every relationship + // after the first one. + let mut existing: HashMap = graph + .node_indices() + .map(|idx| (graph[idx].clone(), idx)) + .collect(); + let mut node_map = HashMap::new(); for node in new_graph.node_indices() { - let new_node = graph.add_node(new_graph[node].clone()); - node_map.insert(node, new_node); + let dataset = new_graph[node].clone(); + let dest_index = if let Some(&existing_idx) = existing.get(&dataset) { + existing_idx + } else { + let new_idx = graph.add_node(dataset.clone()); + existing.insert(dataset, new_idx); + new_idx + }; + node_map.insert(node, dest_index); } for edge in new_graph.edge_indices() { @@ -913,7 +943,13 @@ fn merge_graph( }; let source = node_map.get(&source).unwrap(); let target = node_map.get(&target).unwrap(); - graph.add_edge(*source, *target, new_graph[edge].clone()); + // Skip duplicate edges between the same pair of nodes: the same + // relationship may appear in multiple calc-col sub-graphs (e.g. two + // calc cols traversing the same relationship), and adding parallel + // edges would produce redundant joins downstream. + if graph.find_edge(*source, *target).is_none() { + graph.add_edge(*source, *target, new_graph[edge].clone()); + } } Ok(()) } diff --git a/core/wren-core/core/src/logical_plan/analyze/relation_chain.rs b/core/wren-core/core/src/logical_plan/analyze/relation_chain.rs index 122710ea8b..1a12968d7d 100644 --- a/core/wren-core/core/src/logical_plan/analyze/relation_chain.rs +++ b/core/wren-core/core/src/logical_plan/analyze/relation_chain.rs @@ -61,7 +61,7 @@ impl RelationChain { #[allow(clippy::too_many_arguments)] pub fn with_chain( source: Self, - mut start: NodeIndex, + start: NodeIndex, iter: impl Iterator, directed_graph: Graph, model_required_fields: &HashMap>, @@ -70,10 +70,22 @@ impl RelationChain { properties: SessionPropertiesRef, ) -> Result { let mut relation_chain = source; + // Track the most recently visited node so that we can walk both linear + // chains (A -> B -> C, where each link starts from the previously + // visited node) and star fan-outs (A -> B and A -> C, both starting + // from the base model). At each step we prefer an edge from the + // previous node; if none exists we fall back to an edge from the + // original starting node. This lets a single model project calculated + // columns that span multiple distinct relationships without truncating + // the join chain. + let mut prev = start; for next in iter { let target = directed_graph.node_weight(next).unwrap(); - let Some(link_index) = directed_graph.find_edge(start, next) else { + let link_index = directed_graph + .find_edge(prev, next) + .or_else(|| directed_graph.find_edge(start, next)); + let Some(link_index) = link_index else { break; }; let link = directed_graph.edge_weight(link_index).unwrap(); @@ -119,7 +131,7 @@ impl RelationChain { link.condition.clone(), Box::new(relation_chain), ); - start = next; + prev = next; } Ok(relation_chain) } diff --git a/core/wren-core/core/src/mdl/mod.rs b/core/wren-core/core/src/mdl/mod.rs index 418aa509d6..80f404dbf2 100644 --- a/core/wren-core/core/src/mdl/mod.rs +++ b/core/wren-core/core/src/mdl/mod.rs @@ -2610,6 +2610,171 @@ mod test { Ok(()) } + /// Regression test for the multi-relationship calculated-column bug. + /// + /// A model with multiple calculated columns that traverse *different* + /// relationships used to fail at planning time because the merged dataset + /// graph contained duplicate copies of the base model (one per merged + /// sub-graph). `RelationChain::with_chain` located the next link via + /// `directed_graph.find_edge(start, next)`; when the second sub-graph + /// started with a duplicate base-model node, no edge was found from the + /// first relationship's terminal node and the chain was silently + /// truncated. The result was an error like + /// + /// No field named __relation__1. + /// + /// because only the first relationship's columns ended up in the + /// projection. `merge_graph` now dedupes nodes/edges by `Dataset`, the base + /// model is added conditionally, and `with_chain` prefers an edge from the + /// previously visited node (falling back to the start node) so star + /// fan-outs across multiple relationships plan correctly. + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] + async fn test_multiple_relationship_calculated_columns() -> Result<()> { + // Run the transform on a dedicated thread with an enlarged stack: + // wren-core's plan analyzer uses deep recursion and can exceed the + // default 2 MiB test-thread stack when joining several relationships. + // Production tokio worker threads have a larger stack and are + // unaffected; this is not infinite recursion. + std::thread::Builder::new() + .stack_size(8 * 1024 * 1024) + .spawn(|| { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + rt.block_on(test_multiple_relationship_calculated_columns_inner()) + }) + .unwrap() + .join() + .unwrap() + } + + async fn test_multiple_relationship_calculated_columns_inner() -> Result<()> { + let ctx = create_wren_ctx(None, None); + + // dc_inventory has two distinct relationships: one to `product`, one to + // `market_retailer`. Three calculated columns reference those + // relationships: `market` and `retailer` go through `market_retailer`, + // `item_num` goes through `product`. + let manifest = ManifestBuilder::new() + .catalog("wren") + .schema("test") + .model( + ModelBuilder::new("dc_inventory") + .table_reference("dc_inventory") + .column(ColumnBuilder::new("keypm", "int").build()) + .column(ColumnBuilder::new("market_retailer_key", "int").build()) + .column(ColumnBuilder::new("oh_units", "int").build()) + .column( + ColumnBuilder::new_relationship( + "product", + "product", + "dc_inventory_to_product", + ) + .build(), + ) + .column( + ColumnBuilder::new_relationship( + "market_retailer", + "market_retailer", + "dc_inventory_to_market_retailer", + ) + .build(), + ) + .column( + ColumnBuilder::new_calculated("market", "string") + .expression("market_retailer.market") + .build(), + ) + .column( + ColumnBuilder::new_calculated("retailer", "string") + .expression("market_retailer.retailer") + .build(), + ) + .column( + ColumnBuilder::new_calculated("item_num", "string") + .expression("product.item_num") + .build(), + ) + .primary_key("keypm") + .build(), + ) + .model( + ModelBuilder::new("product") + .table_reference("product") + .column(ColumnBuilder::new("keypm", "int").build()) + .column(ColumnBuilder::new("item_num", "string").build()) + .primary_key("keypm") + .build(), + ) + .model( + ModelBuilder::new("market_retailer") + .table_reference("market_retailer") + .column(ColumnBuilder::new("market_retailer_key", "int").build()) + .column(ColumnBuilder::new("market", "string").build()) + .column(ColumnBuilder::new("retailer", "string").build()) + .primary_key("market_retailer_key") + .build(), + ) + .relationship( + RelationshipBuilder::new("dc_inventory_to_product") + .model("dc_inventory") + .model("product") + .join_type(JoinType::ManyToOne) + .condition("dc_inventory.keypm = product.keypm") + .build(), + ) + .relationship( + RelationshipBuilder::new("dc_inventory_to_market_retailer") + .model("dc_inventory") + .model("market_retailer") + .join_type(JoinType::ManyToOne) + .condition( + "dc_inventory.market_retailer_key = market_retailer.market_retailer_key", + ) + .build(), + ) + .build(); + + let analyzed_mdl = Arc::new(AnalyzedWrenMDL::analyze( + manifest, + Arc::new(HashMap::default()), + Mode::Unparse, + )?); + + // Reproduces the production failure: SELECT * forces all calculated + // columns to be projected, exercising both relationships. + let sql = "SELECT * FROM dc_inventory"; + let result = transform_sql_with_ctx( + &ctx, + Arc::clone(&analyzed_mdl), + &[], + Arc::new(HashMap::new()), + sql, + ) + .await; + + // Before the fix this errors with: + // Schema error: No field named __relation__N.item_num + // (or .market / .retailer depending on traversal order). + let transformed = + result.expect("multi-relationship calc cols should plan successfully"); + // All three calc cols must be projected for SELECT *. + assert!( + transformed.contains("market"), + "transformed SQL missing 'market': {transformed}" + ); + assert!( + transformed.contains("retailer"), + "transformed SQL missing 'retailer': {transformed}" + ); + assert!( + transformed.contains("item_num"), + "transformed SQL missing 'item_num': {transformed}" + ); + Ok(()) + } + #[tokio::test] async fn test_rlac_on_to_many_calculated_field() -> Result<()> { let ctx = create_wren_ctx(None, None);