fix(tesseract): Per branch aliases in rollupLambda unions (cube-js#11184)

Author: ovr

rust/cube/cubesqlplanner/cubesqlplanner/src/logical_plan/optimizers/pre_aggregation/compiled_pre_aggregation.rs

@@ -20,7 +20,21 @@ pub struct PreAggregationJoin {
 
 #[derive(Clone, Debug)]
 pub struct PreAggregationUnion {
-    pub items: Vec<Rc<PreAggregationTable>>,
+    pub items: Vec<PreAggregationUnionItem>,
+}
+
+/// A single member rollup of a `rollupLambda` union, paired with the
+/// member symbols of *that* rollup. The lambda exposes the first member
+/// rollup's symbols, but each branch stores its columns under its own
+/// cube aliases (e.g. `requests_stream__tenant_id` vs `requests__tenant_id`),
+/// so the physical builder needs each branch's own symbols to read the
+/// right column while projecting the lambda's unified alias.
+#[derive(Clone, Debug)]
+pub struct PreAggregationUnionItem {
+    pub table: Rc<PreAggregationTable>,
+    pub measures: Vec<Rc<MemberSymbol>>,
+    pub dimensions: Vec<Rc<MemberSymbol>>,
+    pub time_dimensions: Vec<Rc<MemberSymbol>>,
 }
 
 #[derive(Clone, Debug)]

rust/cube/cubesqlplanner/cubesqlplanner/src/logical_plan/optimizers/pre_aggregation/optimizer.rs

@@ -414,11 +414,14 @@ impl PreAggregationOptimizer {
                 let items = union
                     .items
                     .iter()
-                    .map(|t| {
-                        Rc::new(PreAggregationTable {
+                    .map(|item| PreAggregationUnionItem {
+                        table: Rc::new(PreAggregationTable {
                             usage_index: Some(usage_index),
-                            ..t.as_ref().clone()
-                        })
+                            ..item.table.as_ref().clone()
+                        }),
+                        measures: item.measures.clone(),
+                        dimensions: item.dimensions.clone(),
+                        time_dimensions: item.time_dimensions.clone(),
                     })
                     .collect();
                 Rc::new(PreAggregationSource::Union(PreAggregationUnion { items }))

rust/cube/cubesqlplanner/cubesqlplanner/src/logical_plan/optimizers/pre_aggregation/pre_aggregations_compiler.rs

@@ -7,6 +7,7 @@ use crate::logical_plan::PreAggregationJoin;
 use crate::logical_plan::PreAggregationJoinItem;
 use crate::logical_plan::PreAggregationTable;
 use crate::logical_plan::PreAggregationUnion;
+use crate::logical_plan::PreAggregationUnionItem;
 use crate::planner::join_hints::JoinHints;
 use crate::planner::multi_fact_join_groups::{MeasuresJoinHints, MultiFactJoinGroups};
 use crate::planner::planners::JoinPlanner;
@@ -302,13 +303,22 @@ impl PreAggregationsCompiler {
         for (i, rollup) in pre_aggrs_for_lambda.clone().iter().enumerate() {
             match rollup.source.as_ref() {
                 PreAggregationSource::Single(table) => {
-                    sources.push(Rc::new(table.clone()));
+                    // Carry this branch's own symbols: each member rollup stores its
+                    // columns under its own cube aliases, while the lambda exposes the
+                    // first rollup's symbols. The physical builder maps lambda members
+                    // to each branch's stored column through these.
+                    sources.push(PreAggregationUnionItem {
+                        table: Rc::new(table.clone()),
+                        measures: rollup.measures.clone(),
+                        dimensions: rollup.dimensions.clone(),
+                        time_dimensions: rollup.time_dimensions.clone(),
+                    });
                 }
                 _ => {
                     return Err(CubeError::user(format!("Rollup lambda can't be nested")));
                 }
             }
-            if i > 1 {
+            if i >= 1 {
                 Self::match_symbols(&rollup.measures, &pre_aggrs_for_lambda[0].measures)?;
                 Self::match_symbols(&rollup.dimensions, &pre_aggrs_for_lambda[0].dimensions)?;
                 Self::match_time_dimensions(

rust/cube/cubesqlplanner/cubesqlplanner/src/logical_plan/pre_aggregation.rs

@@ -178,7 +178,10 @@ impl PrettyPrint for PreAggregation {
                 result.println("Union:", &state);
                 let state = state.new_level();
                 for item in union.items.iter() {
-                    result.println(&format!("-{}.{}", item.cube_name, item.name), &state);
+                    result.println(
+                        &format!("-{}.{}", item.table.cube_name, item.table.name),
+                        &state,
+                    );
                 }
             }
         }

rust/cube/cubesqlplanner/cubesqlplanner/src/physical_plan_builder/processors/pre_aggregation.rs

@@ -11,6 +11,7 @@ use crate::physical_plan::{
 };
 use crate::physical_plan_builder::PhysicalPlanBuilder;
 use crate::planner::sql_templates::PlanSqlTemplates;
+use crate::planner::MemberSymbol;
 use crate::planner::SqlJoinCondition;
 use cubenativeutils::CubeError;
 use std::rc::Rc;
@@ -73,68 +74,110 @@ impl PreAggregationProcessor<'_> {
         Ok(from)
     }
 
+    /// Resolve the column a union branch stores a lambda member under. The
+    /// lambda exposes the first member rollup's symbols, but each branch keeps
+    /// its own cube aliases (e.g. `requests_stream__tenant_id` vs the lambda's
+    /// `requests__tenant_id`).
+    fn find_branch_member(
+        lambda_member: &Rc<MemberSymbol>,
+        branch_members: &[Rc<MemberSymbol>],
+        branch_cube_name: &str,
+    ) -> Result<Rc<MemberSymbol>, CubeError> {
+        if let Some(member) = branch_members
+            .iter()
+            .find(|m| m.full_name() == lambda_member.full_name())
+        {
+            return Ok(member.clone());
+        }
+        let short_name = lambda_member.name();
+        if let Some(member) = branch_members
+            .iter()
+            .find(|m| m.name() == short_name && m.cube_name() == branch_cube_name)
+        {
+            return Ok(member.clone());
+        }
+        Err(CubeError::internal(format!(
+            "Lambda pre-aggregation member '{}' has no match in union branch '{}'",
+            lambda_member.full_name(),
+            branch_cube_name
+        )))
+    }
+
     fn make_pre_aggregation_union_source(
         &self,
         pre_aggregation: &PreAggregation,
         union: &PreAggregationUnion,
     ) -> Result<Rc<From>, CubeError> {
         if union.items.len() == 1 {
-            let table_source = self.make_pre_aggregation_table_source(&union.items[0])?;
+            let table_source = self.make_pre_aggregation_table_source(&union.items[0].table)?;
             return Ok(From::new(FromSource::Single(table_source)));
         }
         let query_tools = self.builder.query_tools();
 
         let mut union_sources = Vec::new();
         for item in union.items.iter() {
-            let table_source = self.make_pre_aggregation_table_source(&item)?;
+            let branch_cube_name = &item.table.cube_name;
+            let table_source = self.make_pre_aggregation_table_source(&item.table)?;
             let from = From::new(FromSource::Single(table_source));
             let mut select_builder = SelectBuilder::new(from);
             for dim in pre_aggregation.dimensions().iter() {
-                let name_in_table =
-                    PlanSqlTemplates::member_alias_name(&item.cube_alias, &dim.name(), &None);
-                let alias = dim.alias();
+                // Read this branch's stored column for the lambda dimension and
+                // project it under the lambda's unified alias.
+                let branch_dim = Self::find_branch_member(dim, &item.dimensions, branch_cube_name)?;
                 select_builder.add_projection_reference_member(
                     &dim,
-                    QualifiedColumnName::new(None, name_in_table),
-                    Some(alias),
+                    QualifiedColumnName::new(None, branch_dim.alias()),
+                    Some(dim.alias()),
                 );
             }
+            // Match time dimensions on their base member so the granularity
+            // suffix is applied consistently on both the read and output sides.
+            let branch_time_bases = item
+                .time_dimensions
+                .iter()
+                .map(|td| {
+                    if let Ok(t) = td.as_time_dimension() {
+                        t.base_symbol().clone()
+                    } else {
+                        td.clone()
+                    }
+                })
+                .collect::<Vec<_>>();
             for dim in pre_aggregation.time_dimensions().iter() {
-                let (alias, granularity) = if let Ok(td) = dim.as_time_dimension() {
-                    (td.base_symbol().alias(), td.granularity().clone())
+                let (lambda_base, granularity) = if let Ok(td) = dim.as_time_dimension() {
+                    (td.base_symbol().clone(), td.granularity().clone())
                 } else {
-                    (dim.alias(), None)
+                    (dim.clone(), None)
                 };
 
-                let name_in_table = PlanSqlTemplates::member_alias_name(
-                    &item.cube_alias,
-                    &dim.name(),
-                    &granularity,
-                );
+                let branch_base =
+                    Self::find_branch_member(&lambda_base, &branch_time_bases, branch_cube_name)?;
 
-                let suffix = if let Some(granularity) = granularity {
-                    format!("_{}", granularity.clone())
+                let read_suffix = if let Some(granularity) = &granularity {
+                    format!("_{}", granularity)
+                } else {
+                    String::new()
+                };
+                let name_in_table = format!("{}{}", branch_base.alias(), read_suffix);
+
+                let out_suffix = if let Some(granularity) = &granularity {
+                    format!("_{}", granularity)
                 } else {
                     "_day".to_string()
                 };
-                let alias = format!("{}{}", alias, suffix);
+                let alias = format!("{}{}", lambda_base.alias(), out_suffix);
                 select_builder.add_projection_reference_member(
                     &dim,
-                    QualifiedColumnName::new(None, name_in_table.clone()),
+                    QualifiedColumnName::new(None, name_in_table),
                     Some(alias),
                 );
             }
             for meas in pre_aggregation.measures().iter() {
-                let name_in_table = PlanSqlTemplates::member_alias_name(
-                    &item.cube_alias,
-                    &meas.name(),
-                    &meas.alias_suffix(),
-                );
-                let alias = meas.alias();
+                let branch_meas = Self::find_branch_member(meas, &item.measures, branch_cube_name)?;
                 select_builder.add_projection_reference_member(
                     &meas,
-                    QualifiedColumnName::new(None, name_in_table.clone()),
-                    Some(alias),
+                    QualifiedColumnName::new(None, branch_meas.alias()),
+                    Some(meas.alias()),
                 );
             }
             let context = SqlNodesFactory::new();