feat(tesseract): Support use_original_sql_pre_aggregations in rollup (cube-js#11146)

Author: ovr

packages/cubejs-backend-native/test/bridge/object-bridges-coverage.test.ts

@@ -62,6 +62,7 @@ const BRIDGES: BridgeSpec[] = [
       'timezone',
       'total_query',
       'ungrouped',
+      'use_original_sql_pre_aggregations_in_pre_aggregation',
     ],
   },
   {

packages/cubejs-schema-compiler/src/adapter/BaseQuery.js

@@ -954,6 +954,7 @@ export class BaseQuery {
       ungrouped: this.options.ungrouped,
       exportAnnotatedSql: exportAnnotatedSql === true,
       preAggregationQuery: this.options.preAggregationQuery,
+      useOriginalSqlPreAggregationsInPreAggregation: this.options.useOriginalSqlPreAggregationsInPreAggregation,
       preAggregationId: this.options.preAggregationId || null,
       totalQuery: this.options.totalQuery,
       joinHints: this.options.joinHints,

rust/cube/cubesqlplanner/cubesqlplanner/src/cube_bridge/base_query_options.rs

@@ -206,6 +206,8 @@ pub struct BaseQueryOptionsStatic {
     pub export_annotated_sql: bool,
     #[serde(rename = "preAggregationQuery")]
     pub pre_aggregation_query: Option<bool>,
+    #[serde(rename = "useOriginalSqlPreAggregationsInPreAggregation")]
+    pub use_original_sql_pre_aggregations_in_pre_aggregation: Option<bool>,
     #[serde(rename = "totalQuery")]
     pub total_query: Option<bool>,
     #[serde(rename = "cubestoreSupportMultistage")]

rust/cube/cubesqlplanner/cubesqlplanner/src/planner/query_properties.rs

@@ -167,6 +167,10 @@ pub struct QueryProperties {
     ungrouped: bool,
     #[builder(default)]
     pre_aggregation_query: bool,
+    /// When building a rollup pre-aggregation, source it from the cube's
+    /// `originalSql` pre-aggregation table instead of the raw cube SQL.
+    #[builder(default)]
+    use_original_sql_pre_aggregations_in_pre_aggregation: bool,
     #[builder(default)]
     total_query: bool,
     #[builder(default = Rc::new(JoinHints::new()))]
@@ -354,6 +358,10 @@ impl QueryProperties {
         self.pre_aggregation_query
     }
 
+    pub fn use_original_sql_pre_aggregations_in_pre_aggregation(&self) -> bool {
+        self.use_original_sql_pre_aggregations_in_pre_aggregation
+    }
+
     pub fn disable_external_pre_aggregations(&self) -> bool {
         self.disable_external_pre_aggregations
     }
@@ -1144,6 +1152,7 @@ impl PartialEq for QueryProperties {
             ungrouped,
             ignore_cumulative,
             pre_aggregation_query,
+            use_original_sql_pre_aggregations_in_pre_aggregation,
             total_query,
             allow_multi_stage,
             disable_external_pre_aggregations,
@@ -1169,6 +1178,8 @@ impl PartialEq for QueryProperties {
             && *ungrouped == other.ungrouped
             && *ignore_cumulative == other.ignore_cumulative
             && *pre_aggregation_query == other.pre_aggregation_query
+            && *use_original_sql_pre_aggregations_in_pre_aggregation
+                == other.use_original_sql_pre_aggregations_in_pre_aggregation
             && *total_query == other.total_query
             && *allow_multi_stage == other.allow_multi_stage
             && *disable_external_pre_aggregations == other.disable_external_pre_aggregations

rust/cube/cubesqlplanner/cubesqlplanner/src/planner/query_properties_compiler.rs

@@ -84,6 +84,10 @@ impl QueryPropertiesCompiler {
             .and_then(|v| v.parse::<usize>().ok());
         let ungrouped = options.static_data().ungrouped.unwrap_or(false);
         let pre_aggregation_query = options.static_data().pre_aggregation_query.unwrap_or(false);
+        let use_original_sql_pre_aggregations_in_pre_aggregation = options
+            .static_data()
+            .use_original_sql_pre_aggregations_in_pre_aggregation
+            .unwrap_or(false);
         let total_query = options.static_data().total_query.unwrap_or(false);
         let disable_external_pre_aggregations =
             options.static_data().disable_external_pre_aggregations;
@@ -109,6 +113,9 @@ impl QueryPropertiesCompiler {
             .offset(offset)
             .ungrouped(ungrouped)
             .pre_aggregation_query(pre_aggregation_query)
+            .use_original_sql_pre_aggregations_in_pre_aggregation(
+                use_original_sql_pre_aggregations_in_pre_aggregation,
+            )
             .total_query(total_query)
             .query_join_hints(query_join_hints)
             .disable_external_pre_aggregations(disable_external_pre_aggregations)

rust/cube/cubesqlplanner/cubesqlplanner/src/planner/top_level_planner.rs

@@ -50,10 +50,17 @@ impl TopLevelPlanner {
         };
 
         let templates = self.query_tools.plan_sql_templates(is_external)?;
-
         let physical_plan_builder =
             PhysicalPlanBuilder::new(self.query_tools.query_tools().clone(), templates.clone());
-        let original_sql_pre_aggregations = if !self.request.is_pre_aggregation_query() {
+
+        // Substitute a cube's base SQL with its `originalSql` pre-aggregation table when:
+        // reading (regular query), or building a rollup that opted in via
+        // `useOriginalSqlPreAggregationsInPreAggregation`.
+        let original_sql_pre_aggregations = if !self.request.is_pre_aggregation_query()
+            || self
+                .request
+                .use_original_sql_pre_aggregations_in_pre_aggregation()
+        {
             OriginalSqlCollector::new(self.query_tools.query_tools().clone())
                 .collect(&optimized_plan)?
         } else {

rust/cube/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/base_query_options.rs

@@ -65,6 +65,8 @@ pub struct MockBaseQueryOptions {
     #[builder(default)]
     pre_aggregation_query: Option<bool>,
     #[builder(default)]
+    use_original_sql_pre_aggregations_in_pre_aggregation: Option<bool>,
+    #[builder(default)]
     total_query: Option<bool>,
     #[builder(default)]
     cubestore_support_multistage: Option<bool>,
@@ -93,6 +95,7 @@ impl_static_data!(
     ungrouped,
     export_annotated_sql,
     pre_aggregation_query,
+    use_original_sql_pre_aggregations_in_pre_aggregation,
     total_query,
     cubestore_support_multistage,
     disable_external_pre_aggregations,

rust/cube/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/yaml/base_query_options.rs

@@ -30,6 +30,8 @@ pub struct YamlBaseQueryOptions {
     #[serde(default)]
     pub pre_aggregation_query: Option<bool>,
     #[serde(default)]
+    pub use_original_sql_pre_aggregations_in_pre_aggregation: Option<bool>,
+    #[serde(default)]
     pub total_query: Option<bool>,
     #[serde(default)]
     pub cubestore_support_multistage: Option<bool>,

rust/cube/cubesqlplanner/cubesqlplanner/src/test_fixtures/test_utils/test_context.rs

@@ -443,6 +443,9 @@ impl TestContext {
                 .ungrouped(yaml_options.ungrouped)
                 .export_annotated_sql(yaml_options.export_annotated_sql.unwrap_or(false))
                 .pre_aggregation_query(yaml_options.pre_aggregation_query)
+                .use_original_sql_pre_aggregations_in_pre_aggregation(
+                    yaml_options.use_original_sql_pre_aggregations_in_pre_aggregation,
+                )
                 .total_query(yaml_options.total_query)
                 .cubestore_support_multistage(yaml_options.cubestore_support_multistage)
                 .disable_external_pre_aggregations(

rust/cube/cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/sql_generation.rs

@@ -10,6 +10,7 @@ use crate::test_fixtures::cube_bridge::{
     MockMemberExpressionDefinition, MockMemberSql, MockSchema,
 };
 use crate::test_fixtures::test_utils::TestContext;
+use cubenativeutils::CubeError;
 use indoc::indoc;
 use std::rc::Rc;
 
@@ -1195,3 +1196,94 @@ fn test_count_distinct_approx_multistage_pre_agg_reads_cardinality() {
         sql
     );
 }
+
+// A cube `foo` whose `originalSql` pre-aggregation (`main`) materializes its base
+// SQL, plus a `second` rollup. The mock renders the originalSql pre-agg table as
+// `foo__main` and the raw cube SQL references `foo_table`.
+fn use_original_sql_pre_aggregations_in_pre_aggregation_schema() -> Result<MockSchema, CubeError> {
+    MockSchema::from_yaml(indoc! {"
+        cubes:
+          - name: foo
+            sql: SELECT * FROM foo_table
+            dimensions:
+              - name: time
+                type: time
+                sql: timestamp
+            measures:
+              - name: total
+                type: sum
+                sql: amount
+            pre_aggregations:
+              - name: main
+                type: originalSql
+              - name: second
+                type: rollup
+                measures:
+                  - total
+                time_dimension: time
+                granularity: day
+    "})
+}
+
+// Building a rollup with `useOriginalSqlPreAggregations` must source it from the cube's
+// `originalSql` pre-aggregation table instead of the raw cube SQL
+#[test]
+fn test_rollup_build_with_use_original_sql_pre_aggregations_in_pre_aggregation_reads_original_sql_pre_agg(
+) -> Result<(), CubeError> {
+    let ctx = TestContext::new(use_original_sql_pre_aggregations_in_pre_aggregation_schema()?)?;
+
+    let query_yaml = indoc! {"
+        measures:
+          - foo.total
+        time_dimensions:
+          - dimension: foo.time
+            granularity: day
+        pre_aggregation_query: true
+        use_original_sql_pre_aggregations_in_pre_aggregation: true
+    "};
+
+    let sql = ctx.build_sql(query_yaml)?;
+
+    assert!(
+        sql.contains("foo__main"),
+        "Build SQL should source from the originalSql pre-agg table, got:\n{}",
+        sql
+    );
+    assert!(
+        !sql.contains("foo_table"),
+        "Build SQL should not read the raw cube table when useOriginalSqlPreAggregations is set, got:\n{}",
+        sql
+    );
+    Ok(())
+}
+
+// Without the flag, a rollup build reads the raw cube SQL — the originalSql
+// pre-agg table must not be substituted in.
+#[test]
+fn test_rollup_build_without_use_original_sql_pre_aggregations_in_pre_aggregation_reads_raw_table(
+) -> Result<(), CubeError> {
+    let ctx = TestContext::new(use_original_sql_pre_aggregations_in_pre_aggregation_schema()?)?;
+
+    let query_yaml = indoc! {"
+        measures:
+          - foo.total
+        time_dimensions:
+          - dimension: foo.time
+            granularity: day
+        pre_aggregation_query: true
+    "};
+
+    let sql = ctx.build_sql(query_yaml)?;
+
+    assert!(
+        sql.contains("foo_table"),
+        "Build SQL should read the raw cube table without useOriginalSqlPreAggregations, got:\n{}",
+        sql
+    );
+    assert!(
+        !sql.contains("foo__main"),
+        "Build SQL should not source from the originalSql pre-agg table without the flag, got:\n{}",
+        sql
+    );
+    Ok(())
+}