diff --git a/packages/cubejs-backend-native/test/bridge/object-bridges-coverage.test.ts b/packages/cubejs-backend-native/test/bridge/object-bridges-coverage.test.ts index c7ac0e632d105..a49039c536502 100644 --- a/packages/cubejs-backend-native/test/bridge/object-bridges-coverage.test.ts +++ b/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', ], }, { diff --git a/packages/cubejs-schema-compiler/src/adapter/BaseQuery.js b/packages/cubejs-schema-compiler/src/adapter/BaseQuery.js index 552916937d301..8a515e8e81906 100644 --- a/packages/cubejs-schema-compiler/src/adapter/BaseQuery.js +++ b/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, diff --git a/rust/cube/cubesqlplanner/cubesqlplanner/src/cube_bridge/base_query_options.rs b/rust/cube/cubesqlplanner/cubesqlplanner/src/cube_bridge/base_query_options.rs index f8deae6722bc8..5208860ba436d 100644 --- a/rust/cube/cubesqlplanner/cubesqlplanner/src/cube_bridge/base_query_options.rs +++ b/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, + #[serde(rename = "useOriginalSqlPreAggregationsInPreAggregation")] + pub use_original_sql_pre_aggregations_in_pre_aggregation: Option, #[serde(rename = "totalQuery")] pub total_query: Option, #[serde(rename = "cubestoreSupportMultistage")] diff --git a/rust/cube/cubesqlplanner/cubesqlplanner/src/planner/query_properties.rs b/rust/cube/cubesqlplanner/cubesqlplanner/src/planner/query_properties.rs index 08e88deb156e7..c8c50894e4fb3 100644 --- a/rust/cube/cubesqlplanner/cubesqlplanner/src/planner/query_properties.rs +++ b/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 diff --git a/rust/cube/cubesqlplanner/cubesqlplanner/src/planner/query_properties_compiler.rs b/rust/cube/cubesqlplanner/cubesqlplanner/src/planner/query_properties_compiler.rs index 9e5d781d68007..5979a16cd2352 100644 --- a/rust/cube/cubesqlplanner/cubesqlplanner/src/planner/query_properties_compiler.rs +++ b/rust/cube/cubesqlplanner/cubesqlplanner/src/planner/query_properties_compiler.rs @@ -84,6 +84,10 @@ impl QueryPropertiesCompiler { .and_then(|v| v.parse::().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) diff --git a/rust/cube/cubesqlplanner/cubesqlplanner/src/planner/top_level_planner.rs b/rust/cube/cubesqlplanner/cubesqlplanner/src/planner/top_level_planner.rs index f7e3ff0d091fd..427eb8c2b46ad 100644 --- a/rust/cube/cubesqlplanner/cubesqlplanner/src/planner/top_level_planner.rs +++ b/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 { diff --git a/rust/cube/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/base_query_options.rs b/rust/cube/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/base_query_options.rs index bca4b3d7d78ca..092de5bff1982 100644 --- a/rust/cube/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/base_query_options.rs +++ b/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, #[builder(default)] + use_original_sql_pre_aggregations_in_pre_aggregation: Option, + #[builder(default)] total_query: Option, #[builder(default)] cubestore_support_multistage: Option, @@ -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, diff --git a/rust/cube/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/yaml/base_query_options.rs b/rust/cube/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/yaml/base_query_options.rs index b9bfa737531ec..b63ae658a6ed8 100644 --- a/rust/cube/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/yaml/base_query_options.rs +++ b/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, #[serde(default)] + pub use_original_sql_pre_aggregations_in_pre_aggregation: Option, + #[serde(default)] pub total_query: Option, #[serde(default)] pub cubestore_support_multistage: Option, diff --git a/rust/cube/cubesqlplanner/cubesqlplanner/src/test_fixtures/test_utils/test_context.rs b/rust/cube/cubesqlplanner/cubesqlplanner/src/test_fixtures/test_utils/test_context.rs index a6c1dc99eab64..1b90aaa93717a 100644 --- a/rust/cube/cubesqlplanner/cubesqlplanner/src/test_fixtures/test_utils/test_context.rs +++ b/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( diff --git a/rust/cube/cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/sql_generation.rs b/rust/cube/cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/sql_generation.rs index 83f26924ed5e2..34dbb391c0cff 100644 --- a/rust/cube/cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/sql_generation.rs +++ b/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::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(()) +}