fix(tesseract): Segment support in pre-aggregations (cube-js#10422)

Author: waralexrom

packages/cubejs-schema-compiler/test/integration/postgres/sql-generation.test.ts

@@ -3058,6 +3058,7 @@ SELECT 1 AS revenue,  cast('2024-01-01' AS timestamp) as time UNION ALL
       preAggregationsSchema: ''
     });
 
+    query.buildSqlAndParams();
     const preAggregationsDescription: any = query.preAggregations?.preAggregationsDescription()[0];
 
     const res = await testWithPreAggregation(preAggregationsDescription, query);

rust/cubesqlplanner/cubesqlplanner/src/cube_bridge/pre_aggregation_description.rs

@@ -33,6 +33,9 @@ pub trait PreAggregationDescription {
     #[nbridge(field, optional)]
     fn time_dimension_reference(&self) -> Result<Option<Rc<dyn MemberSql>>, CubeError>;
 
+    #[nbridge(field, optional)]
+    fn segment_references(&self) -> Result<Option<Rc<dyn MemberSql>>, CubeError>;
+
     #[nbridge(field, optional)]
     fn rollup_references(&self) -> Result<Option<Rc<dyn MemberSql>>, CubeError>;
 }

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

@@ -47,6 +47,7 @@ pub struct CompiledPreAggregation {
     pub measures: Vec<Rc<MemberSymbol>>,
     pub dimensions: Vec<Rc<MemberSymbol>>,
     pub time_dimensions: Vec<Rc<MemberSymbol>>,
+    pub segments: Vec<Rc<MemberSymbol>>,
     pub allow_non_strict_date_range_match: bool,
 }
 
@@ -60,6 +61,7 @@ impl Debug for CompiledPreAggregation {
             .field("measures", &self.measures)
             .field("dimensions", &self.dimensions)
             .field("time_dimensions", &self.time_dimensions)
+            .field("segments", &self.segments)
             .field(
                 "allow_non_strict_date_range_match",
                 &self.allow_non_strict_date_range_match,

rust/cubesqlplanner/cubesqlplanner/src/logical_plan/optimizers/pre_aggregation/dimension_matcher.rs

@@ -35,6 +35,7 @@ pub struct DimensionMatcher<'a> {
     pre_aggregation: &'a CompiledPreAggregation,
     pre_aggregation_dimensions: HashMap<String, bool>,
     pre_aggregation_time_dimensions: HashMap<String, (Option<Rc<TimeDimensionSymbol>>, bool)>,
+    pre_aggregation_segments: HashMap<String, bool>,
     result: MatchState,
 }
 
@@ -56,11 +57,17 @@ impl<'a> DimensionMatcher<'a> {
                 }
             })
             .collect::<HashMap<_, _>>();
+        let pre_aggregation_segments = pre_aggregation
+            .segments
+            .iter()
+            .map(|s| (s.full_name(), false))
+            .collect();
         Self {
             query_tools,
             pre_aggregation,
             pre_aggregation_dimensions,
             pre_aggregation_time_dimensions,
+            pre_aggregation_segments,
             result: MatchState::Full,
         }
     }
@@ -128,6 +135,12 @@ impl<'a> DimensionMatcher<'a> {
                 MatchState::Partial
             };
         self.result = self.result.combine(&time_dimension_coverage_result);
+        let segment_coverage_result = if self.pre_aggregation_segments.values().all(|v| *v) {
+            MatchState::Full
+        } else {
+            MatchState::Partial
+        };
+        self.result = self.result.combine(&segment_coverage_result);
         self.result
     }
 
@@ -143,7 +156,30 @@ impl<'a> DimensionMatcher<'a> {
             MemberSymbol::TimeDimension(time_dimension) => {
                 self.try_match_time_dimension(time_dimension, add_to_matched_dimension)
             }
-            MemberSymbol::MemberExpression(_member_expression) => Ok(MatchState::NotMatched), //TODO We don't allow to use pre-aggregations with member expressions before SQL API is ready for it
+            MemberSymbol::MemberExpression(me) => {
+                if let Some(found) = self.pre_aggregation_segments.get_mut(&me.full_name()) {
+                    if add_to_matched_dimension {
+                        *found = true;
+                    }
+                    Ok(MatchState::Full)
+                } else {
+                    let dependencies = me.get_dependencies();
+                    if dependencies.is_empty() {
+                        Ok(MatchState::NotMatched)
+                    } else {
+                        let mut result = MatchState::Partial;
+                        for dep in dependencies {
+                            let dep_match =
+                                self.try_match_symbol(&dep, add_to_matched_dimension)?;
+                            if dep_match == MatchState::NotMatched {
+                                return Ok(MatchState::NotMatched);
+                            }
+                            result = result.combine(&dep_match);
+                        }
+                        Ok(result)
+                    }
+                }
+            }
             _ => Ok(MatchState::NotMatched),
         }
     }
@@ -655,4 +691,139 @@ mod tests {
             MatchState::Full,
         );
     }
+
+    #[test]
+    fn test_segment_full_match() {
+        let ctx = create_test_context();
+        assert_eq!(
+            match_pre_agg(
+                &ctx,
+                "segment_rollup",
+                indoc! {"
+                    measures:
+                      - orders.count
+                    dimensions:
+                      - orders.status
+                    segments:
+                      - orders.high_priority
+                    time_dimensions:
+                      - dimension: orders.created_at
+                        granularity: day
+                "},
+            ),
+            MatchState::Full,
+        );
+    }
+
+    #[test]
+    fn test_segment_partial_match_unused_segment() {
+        let ctx = create_test_context();
+        // No segment in query, but pre-agg has segment => partial (unused segment coverage)
+        assert_eq!(
+            match_pre_agg(
+                &ctx,
+                "segment_rollup",
+                indoc! {"
+                    measures:
+                      - orders.count
+                    dimensions:
+                      - orders.status
+                    time_dimensions:
+                      - dimension: orders.created_at
+                        granularity: day
+                "},
+            ),
+            MatchState::Partial,
+        );
+    }
+
+    #[test]
+    fn test_segment_not_matched_missing_in_pre_agg() {
+        let ctx = create_test_context();
+        // Segment in query, but pre-agg without segments => NotMatched
+        assert_eq!(
+            match_pre_agg(
+                &ctx,
+                "main_rollup",
+                indoc! {"
+                    measures:
+                      - orders.count
+                    dimensions:
+                      - orders.status
+                      - orders.city
+                    segments:
+                      - orders.high_priority
+                "},
+            ),
+            MatchState::NotMatched,
+        );
+    }
+
+    #[test]
+    fn test_expression_segment_matched_via_dimensions() {
+        let ctx = create_test_context();
+        // status_filter depends on {CUBE.status}, status is in main_rollup dimensions
+        assert_eq!(
+            match_pre_agg(
+                &ctx,
+                "main_rollup",
+                indoc! {"
+                    measures:
+                      - orders.count
+                    dimensions:
+                      - orders.status
+                      - orders.city
+                    segments:
+                      - orders.status_filter
+                "},
+            ),
+            MatchState::Partial,
+        );
+    }
+
+    #[test]
+    fn test_expression_segment_not_matched_missing_dependency() {
+        let ctx = create_test_context();
+        // status_filter depends on {CUBE.status}, but daily_countries_rollup has no status dimension
+        assert_eq!(
+            match_pre_agg(
+                &ctx,
+                "daily_countries_rollup",
+                indoc! {"
+                    measures:
+                      - orders.count
+                    dimensions:
+                      - orders.country
+                    segments:
+                      - orders.status_filter
+                    time_dimensions:
+                      - dimension: orders.created_at
+                        granularity: day
+                "},
+            ),
+            MatchState::NotMatched,
+        );
+    }
+
+    #[test]
+    fn test_plain_segment_not_matched_via_dimensions() {
+        let ctx = create_test_context();
+        // high_priority has no {CUBE.x} deps, main_rollup has no segments => NotMatched
+        assert_eq!(
+            match_pre_agg(
+                &ctx,
+                "main_rollup",
+                indoc! {"
+                    measures:
+                      - orders.count
+                    dimensions:
+                      - orders.status
+                      - orders.city
+                    segments:
+                      - orders.high_priority
+                "},
+            ),
+            MatchState::NotMatched,
+        );
+    }
 }

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

@@ -399,7 +399,8 @@ impl PreAggregationOptimizer {
                 .dimensions
                 .iter()
                 .cloned()
-                .chain(pre_aggregation.time_dimensions.iter().map(|d| d.clone()))
+                .chain(pre_aggregation.time_dimensions.iter().cloned())
+                .chain(pre_aggregation.segments.iter().cloned())
                 .collect(),
             measures: filtered_measures.clone(),
             multiplied_measures: HashSet::new(),
@@ -414,6 +415,7 @@ impl PreAggregationOptimizer {
             .name(pre_aggregation.name.clone())
             .time_dimensions(pre_aggregation.time_dimensions.clone())
             .dimensions(pre_aggregation.dimensions.clone())
+            .segments(pre_aggregation.segments.clone())
             .measures(filtered_measures)
             .schema(Rc::new(schema))
             .external(pre_aggregation.external.unwrap_or_default())

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

@@ -165,6 +165,16 @@ impl PreAggregationsCompiler {
         } else {
             Vec::new()
         };
+        let segments = if let Some(refs) = description.segment_references()? {
+            Self::symbols_from_ref(
+                self.query_tools.clone(),
+                &name.cube_name,
+                refs,
+                Self::check_is_segment,
+            )?
+        } else {
+            Vec::new()
+        };
         let allow_non_strict_date_range_match = description
             .static_data()
             .allow_non_strict_date_range_match
@@ -208,6 +218,7 @@ impl PreAggregationsCompiler {
             measures,
             dimensions,
             time_dimensions,
+            segments,
             allow_non_strict_date_range_match,
         });
         self.compiled_cache.insert(name.clone(), res.clone());
@@ -265,6 +276,7 @@ impl PreAggregationsCompiler {
         let measures = pre_aggrs_for_lambda[0].measures.clone();
         let dimensions = pre_aggrs_for_lambda[0].dimensions.clone();
         let time_dimensions = pre_aggrs_for_lambda[0].time_dimensions.clone();
+        let segments = pre_aggrs_for_lambda[0].segments.clone();
         let allow_non_strict_date_range_match = description
             .static_data()
             .allow_non_strict_date_range_match
@@ -282,6 +294,7 @@ impl PreAggregationsCompiler {
             measures,
             dimensions,
             time_dimensions,
+            segments,
             allow_non_strict_date_range_match,
         });
         self.compiled_cache.insert(name.clone(), res.clone());
@@ -504,6 +517,15 @@ impl PreAggregationsCompiler {
         }
         Ok(())
     }
+
+    fn check_is_segment(symbol: &MemberSymbol) -> Result<(), CubeError> {
+        symbol.as_member_expression().map_err(|_| {
+            CubeError::user(
+                "Pre-aggregation segment reference must be a member expression".to_string(),
+            )
+        })?;
+        Ok(())
+    }
 }
 
 #[cfg(test)]
@@ -759,4 +781,39 @@ mod tests {
         let result2 = PreAggregationFullName::from_string("too.many.parts");
         assert!(result2.is_err());
     }
+
+    #[test]
+    fn test_compile_rollup_with_segments() {
+        let schema = MockSchema::from_yaml_file("common/pre_aggregation_matching_test.yaml");
+        let test_context = TestContext::new(schema).unwrap();
+        let query_tools = test_context.query_tools().clone();
+
+        let cube_names = vec!["orders".to_string()];
+        let mut compiler = PreAggregationsCompiler::try_new(query_tools, &cube_names).unwrap();
+
+        let pre_agg_name =
+            PreAggregationFullName::new("orders".to_string(), "segment_rollup".to_string());
+        let compiled = compiler.compile_pre_aggregation(&pre_agg_name).unwrap();
+
+        assert_eq!(compiled.name, "segment_rollup");
+        assert_eq!(compiled.cube_name, "orders");
+        assert_eq!(compiled.granularity, Some("day".to_string()));
+
+        // Check segments
+        assert_eq!(compiled.segments.len(), 1);
+        assert_eq!(
+            compiled.segments[0].full_name(),
+            "expr:orders.high_priority"
+        );
+
+        // Check measures
+        assert_eq!(compiled.measures.len(), 2);
+        let measure_names: Vec<String> = compiled.measures.iter().map(|m| m.full_name()).collect();
+        assert!(measure_names.contains(&"orders.count".to_string()));
+        assert!(measure_names.contains(&"orders.total_amount".to_string()));
+
+        // Check dimensions
+        assert_eq!(compiled.dimensions.len(), 1);
+        assert_eq!(compiled.dimensions[0].full_name(), "orders.status");
+    }
 }

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

@@ -15,6 +15,8 @@ pub struct PreAggregation {
     dimensions: Vec<Rc<MemberSymbol>>,
     #[builder(default)]
     time_dimensions: Vec<Rc<MemberSymbol>>,
+    #[builder(default)]
+    segments: Vec<Rc<MemberSymbol>>,
     external: bool,
     #[builder(default)]
     granularity: Option<String>,
@@ -43,6 +45,10 @@ impl PreAggregation {
         &self.time_dimensions
     }
 
+    pub fn segments(&self) -> &Vec<Rc<MemberSymbol>> {
+        &self.segments
+    }
+
     pub fn external(&self) -> bool {
         self.external
     }
@@ -115,6 +121,11 @@ impl PreAggregation {
             );
         }
 
+        for segment in self.segments().iter() {
+            let alias = segment.alias();
+            res.insert(segment.full_name(), QualifiedColumnName::new(None, alias));
+        }
+
         if let PreAggregationSource::Join(join) = self.source().as_ref() {
             for item in join.items.iter() {
                 for member in item.from_members.iter().chain(item.to_members.iter()) {

rust/cubesqlplanner/cubesqlplanner/src/planner/filter/base_segment.rs

@@ -34,6 +34,7 @@ impl BaseSegment {
             name.clone(),
             MemberExpressionExpression::SqlCall(expression),
             None,
+            None,
             query_tools.base_tools().clone(),
         )?;
         let full_name = full_name.unwrap_or(member_expression_symbol.full_name());