diff --git a/packages/cubejs-schema-compiler/src/adapter/BaseQuery.js b/packages/cubejs-schema-compiler/src/adapter/BaseQuery.js index 6e768435a7426..21c6c24160d6c 100644 --- a/packages/cubejs-schema-compiler/src/adapter/BaseQuery.js +++ b/packages/cubejs-schema-compiler/src/adapter/BaseQuery.js @@ -1000,6 +1000,7 @@ export class BaseQuery { ungrouped: this.options.ungrouped, exportAnnotatedSql: false, preAggregationQuery: this.options.preAggregationQuery, + preAggregationId: this.options.preAggregationId || null, securityContext: this.contextSymbols.securityContext, cubestoreSupportMultistage: this.options.cubestoreSupportMultistage ?? getEnv('cubeStoreRollingWindowJoin'), disableExternalPreAggregations: !!this.options.disableExternalPreAggregations, diff --git a/rust/cubesqlplanner/cubesqlplanner/src/cube_bridge/base_query_options.rs b/rust/cubesqlplanner/cubesqlplanner/src/cube_bridge/base_query_options.rs index 05effb43fb170..4c1cac9cc792f 100644 --- a/rust/cubesqlplanner/cubesqlplanner/src/cube_bridge/base_query_options.rs +++ b/rust/cubesqlplanner/cubesqlplanner/src/cube_bridge/base_query_options.rs @@ -71,6 +71,8 @@ pub struct BaseQueryOptionsStatic { pub cubestore_support_multistage: Option, #[serde(rename = "disableExternalPreAggregations")] pub disable_external_pre_aggregations: bool, + #[serde(rename = "preAggregationId")] + pub pre_aggregation_id: Option, } #[nativebridge::native_bridge(BaseQueryOptionsStatic)] diff --git a/rust/cubesqlplanner/cubesqlplanner/src/logical_plan/optimizers/pre_aggregation/dimension_matcher.rs b/rust/cubesqlplanner/cubesqlplanner/src/logical_plan/optimizers/pre_aggregation/dimension_matcher.rs index 54feb41f93c4c..158fd8ba7e0b3 100644 --- a/rust/cubesqlplanner/cubesqlplanner/src/logical_plan/optimizers/pre_aggregation/dimension_matcher.rs +++ b/rust/cubesqlplanner/cubesqlplanner/src/logical_plan/optimizers/pre_aggregation/dimension_matcher.rs @@ -227,11 +227,7 @@ impl<'a> DimensionMatcher<'a> { add_to_matched_dimension: bool, ) -> Result { let granularity = if self.pre_aggregation.allow_non_strict_date_range_match { - if let Some(granularity) = time_dimension.granularity_obj() { - granularity.min_granularity()? - } else { - time_dimension.granularity().clone() - } + time_dimension.granularity().clone() } else { time_dimension.rollup_granularity(self.query_tools.clone())? }; diff --git a/rust/cubesqlplanner/cubesqlplanner/src/logical_plan/optimizers/pre_aggregation/optimizer.rs b/rust/cubesqlplanner/cubesqlplanner/src/logical_plan/optimizers/pre_aggregation/optimizer.rs index db52c997367c5..7fcb9477b654c 100644 --- a/rust/cubesqlplanner/cubesqlplanner/src/logical_plan/optimizers/pre_aggregation/optimizer.rs +++ b/rust/cubesqlplanner/cubesqlplanner/src/logical_plan/optimizers/pre_aggregation/optimizer.rs @@ -28,6 +28,7 @@ impl PreAggregationOptimizer { &mut self, plan: Rc, disable_external_pre_aggregations: bool, + pre_aggregation_id: Option<&str>, ) -> Result>, CubeError> { let cube_names = collect_cube_names_from_node(&plan)?; let mut compiler = PreAggregationsCompiler::try_new(self.query_tools.clone(), &cube_names)?; @@ -36,6 +37,12 @@ impl PreAggregationOptimizer { compiler.compile_all_pre_aggregations(disable_external_pre_aggregations)?; for pre_aggregation in compiled_pre_aggregations.iter() { + if let Some(id) = pre_aggregation_id { + let full_name = format!("{}.{}", pre_aggregation.cube_name, pre_aggregation.name); + if full_name != id { + continue; + } + } let new_query = self.try_rewrite_query(plan.clone(), pre_aggregation)?; if new_query.is_some() { return Ok(new_query); diff --git a/rust/cubesqlplanner/cubesqlplanner/src/planner/query_properties.rs b/rust/cubesqlplanner/cubesqlplanner/src/planner/query_properties.rs index d46b10e715a23..85d8280d67322 100644 --- a/rust/cubesqlplanner/cubesqlplanner/src/planner/query_properties.rs +++ b/rust/cubesqlplanner/cubesqlplanner/src/planner/query_properties.rs @@ -108,6 +108,7 @@ pub struct QueryProperties { query_join_hints: Rc>, allow_multi_stage: bool, disable_external_pre_aggregations: bool, + pre_aggregation_id: Option, } impl QueryProperties { @@ -410,6 +411,7 @@ impl QueryProperties { let total_query = options.static_data().total_query.unwrap_or(false); let disable_external_pre_aggregations = options.static_data().disable_external_pre_aggregations; + let pre_aggregation_id = options.static_data().pre_aggregation_id.clone(); let mut res = Self { measures, @@ -431,6 +433,7 @@ impl QueryProperties { query_join_hints, allow_multi_stage: true, disable_external_pre_aggregations, + pre_aggregation_id, }; res.apply_static_filters()?; Ok(Rc::new(res)) @@ -482,6 +485,7 @@ impl QueryProperties { query_join_hints, allow_multi_stage, disable_external_pre_aggregations, + pre_aggregation_id: None, }; res.apply_static_filters()?; @@ -753,6 +757,10 @@ impl QueryProperties { self.disable_external_pre_aggregations } + pub fn pre_aggregation_id(&self) -> Option<&str> { + self.pre_aggregation_id.as_deref() + } + pub fn all_filters(&self) -> Option { let items = self .time_dimensions_filters diff --git a/rust/cubesqlplanner/cubesqlplanner/src/planner/top_level_planner.rs b/rust/cubesqlplanner/cubesqlplanner/src/planner/top_level_planner.rs index 690a7aeec8d51..3b69cd453491e 100644 --- a/rust/cubesqlplanner/cubesqlplanner/src/planner/top_level_planner.rs +++ b/rust/cubesqlplanner/cubesqlplanner/src/planner/top_level_planner.rs @@ -76,9 +76,12 @@ impl TopLevelPlanner { ); let disable_external_pre_aggregations = self.request.disable_external_pre_aggregations(); - if let Some(result) = pre_aggregation_optimizer - .try_optimize(plan.clone(), disable_external_pre_aggregations)? - { + let pre_aggregation_id = self.request.pre_aggregation_id(); + if let Some(result) = pre_aggregation_optimizer.try_optimize( + plan.clone(), + disable_external_pre_aggregations, + pre_aggregation_id, + )? { if pre_aggregation_optimizer.get_used_pre_aggregations().len() == 1 { ( result, diff --git a/rust/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/base_query_options.rs b/rust/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/base_query_options.rs index b992ea7d4c59e..07a6b1931a009 100644 --- a/rust/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/base_query_options.rs +++ b/rust/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/base_query_options.rs @@ -65,6 +65,8 @@ pub struct MockBaseQueryOptions { cubestore_support_multistage: Option, #[builder(default = false)] disable_external_pre_aggregations: bool, + #[builder(default)] + pre_aggregation_id: Option, } impl_static_data!( @@ -82,7 +84,8 @@ impl_static_data!( pre_aggregation_query, total_query, cubestore_support_multistage, - disable_external_pre_aggregations + disable_external_pre_aggregations, + pre_aggregation_id ); pub fn members_from_strings(strings: Vec) -> Vec { diff --git a/rust/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/mock_dimension_definition.rs b/rust/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/mock_dimension_definition.rs index 8d790a4f708e0..96abec097c30f 100644 --- a/rust/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/mock_dimension_definition.rs +++ b/rust/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/mock_dimension_definition.rs @@ -59,7 +59,7 @@ impl MockDimensionDefinition { pub fn from_yaml(yaml: &str) -> Result, CubeError> { let yaml_def: YamlDimensionDefinition = serde_yaml::from_str(yaml) .map_err(|e| CubeError::user(format!("Failed to parse YAML: {}", e)))?; - Ok(yaml_def.build()) + Ok(Rc::new(yaml_def.build().definition)) } } diff --git a/rust/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/mock_evaluator.rs b/rust/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/mock_evaluator.rs index 30eb738187320..6629edc45ad8e 100644 --- a/rust/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/mock_evaluator.rs +++ b/rust/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/mock_evaluator.rs @@ -233,23 +233,29 @@ impl CubeEvaluator for MockCubeEvaluator { let granularity = &path[3]; - let valid_granularities = [ + // Check custom granularities in schema first + if let Some(custom) = self.schema.get_granularity(&path[0], &path[1], granularity) { + return Ok(custom as Rc); + } + + // Fall back to predefined granularities + let predefined = [ "second", "minute", "hour", "day", "week", "month", "quarter", "year", ]; - if !valid_granularities.contains(&granularity.as_str()) { - return Err(CubeError::user(format!( - "Unsupported granularity: '{}'. Supported: second, minute, hour, day, week, month, quarter, year", + if predefined.contains(&granularity.as_str()) { + use crate::test_fixtures::cube_bridge::MockGranularityDefinition; + Ok(Rc::new( + MockGranularityDefinition::builder() + .interval(format!("1 {}", granularity)) + .build(), + ) as Rc) + } else { + Err(CubeError::user(format!( + "Granularity '{}' not found", granularity - ))); + ))) } - - use crate::test_fixtures::cube_bridge::MockGranularityDefinition; - Ok(Rc::new( - MockGranularityDefinition::builder() - .interval(granularity.clone()) - .build(), - ) as Rc) } fn pre_aggregations_for_cube_as_array( @@ -308,3 +314,72 @@ impl CubeEvaluator for MockCubeEvaluator { self } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::test_fixtures::cube_bridge::MockSchema; + + fn create_custom_granularity_schema() -> MockSchema { + MockSchema::from_yaml_file("common/custom_granularity_test.yaml") + } + + fn resolve( + evaluator: &MockCubeEvaluator, + granularity: &str, + ) -> Result, CubeError> { + evaluator.resolve_granularity(vec![ + "orders".to_string(), + "created_at".to_string(), + "granularities".to_string(), + granularity.to_string(), + ]) + } + + #[test] + fn test_resolve_predefined_granularity() { + let schema = create_custom_granularity_schema(); + let evaluator = schema.create_evaluator(); + + let result = resolve(&evaluator, "day").expect("should resolve predefined granularity"); + assert_eq!(result.static_data().interval, "1 day"); + assert_eq!(result.static_data().origin, None); + assert_eq!(result.static_data().offset, None); + } + + #[test] + fn test_resolve_custom_granularity() { + let schema = create_custom_granularity_schema(); + let evaluator = schema.create_evaluator(); + + let result = resolve(&evaluator, "half_year").expect("should resolve custom granularity"); + assert_eq!(result.static_data().interval, "6 months"); + assert_eq!(result.static_data().origin, Some("2024-01-01".to_string())); + assert_eq!(result.static_data().offset, None); + } + + #[test] + fn test_resolve_custom_granularity_with_offset() { + let schema = create_custom_granularity_schema(); + let evaluator = schema.create_evaluator(); + + let result = resolve(&evaluator, "fiscal_year").expect("should resolve custom granularity"); + assert_eq!(result.static_data().interval, "1 year"); + assert_eq!(result.static_data().offset, Some("1 month".to_string())); + } + + #[test] + fn test_resolve_unknown_granularity_error() { + let schema = create_custom_granularity_schema(); + let evaluator = schema.create_evaluator(); + + let result = resolve(&evaluator, "nonexistent"); + assert!(result.is_err()); + let err = result.err().unwrap(); + assert!( + err.message.contains("Granularity 'nonexistent' not found"), + "unexpected error: {}", + err.message + ); + } +} diff --git a/rust/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/mock_granularity_definition.rs b/rust/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/mock_granularity_definition.rs index db762c779661a..43763e0d573c6 100644 --- a/rust/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/mock_granularity_definition.rs +++ b/rust/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/mock_granularity_definition.rs @@ -12,9 +12,9 @@ use typed_builder::TypedBuilder; pub struct MockGranularityDefinition { #[builder(setter(into))] interval: String, - #[builder(default, setter(strip_option, into))] + #[builder(default, setter(strip_option(fallback = origin_opt), into))] origin: Option, - #[builder(default, setter(strip_option, into))] + #[builder(default, setter(strip_option(fallback = offset_opt), into))] offset: Option, #[builder(default, setter(strip_option))] sql: Option>, diff --git a/rust/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/mock_schema.rs b/rust/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/mock_schema.rs index 9f9f37c3fa24c..9702c8213ac4f 100644 --- a/rust/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/mock_schema.rs +++ b/rust/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/mock_schema.rs @@ -1,8 +1,8 @@ use crate::test_fixtures::cube_bridge::yaml::YamlSchema; use crate::test_fixtures::cube_bridge::{ MockBaseTools, MockCubeDefinition, MockCubeEvaluator, MockDimensionDefinition, MockDriverTools, - MockJoinGraph, MockJoinItemDefinition, MockMeasureDefinition, MockPreAggregationDescription, - MockSegmentDefinition, + MockGranularityDefinition, MockJoinGraph, MockJoinItemDefinition, MockMeasureDefinition, + MockPreAggregationDescription, MockSegmentDefinition, }; use cubenativeutils::CubeError; use std::collections::HashMap; @@ -20,6 +20,8 @@ pub struct MockCube { pub dimensions: HashMap>, pub segments: HashMap>, pub pre_aggregations: Vec<(String, Rc)>, + /// Outer key = dimension_name, inner key = granularity_name + pub granularities: HashMap>>, } impl MockSchema { @@ -108,6 +110,18 @@ impl MockSchema { .and_then(|cube| cube.segments.get(segment_name).cloned()) } + pub fn get_granularity( + &self, + cube_name: &str, + dimension_name: &str, + granularity_name: &str, + ) -> Option> { + self.cubes + .get(cube_name) + .and_then(|cube| cube.granularities.get(dimension_name)) + .and_then(|grans| grans.get(granularity_name).cloned()) + } + pub fn get_pre_aggregation( &self, cube_name: &str, @@ -275,6 +289,7 @@ impl MockSchemaBuilder { dimensions: HashMap::new(), segments: HashMap::new(), pre_aggregations: Vec::new(), + granularities: HashMap::new(), joins: HashMap::new(), } } @@ -309,6 +324,7 @@ pub struct MockCubeBuilder { dimensions: HashMap>, segments: HashMap>, pre_aggregations: Vec<(String, Rc)>, + granularities: HashMap>>, #[allow(dead_code)] joins: HashMap, } @@ -356,6 +372,19 @@ impl MockCubeBuilder { self } + pub fn add_granularity( + mut self, + dimension_name: &str, + granularity_name: &str, + definition: MockGranularityDefinition, + ) -> Self { + self.granularities + .entry(dimension_name.to_string()) + .or_default() + .insert(granularity_name.to_string(), Rc::new(definition)); + self + } + #[allow(dead_code)] pub fn add_join(mut self, name: impl Into, definition: MockJoinItemDefinition) -> Self { self.joins.insert(name.into(), definition); @@ -376,6 +405,7 @@ impl MockCubeBuilder { dimensions: self.dimensions, segments: self.segments, pre_aggregations: self.pre_aggregations, + granularities: self.granularities, }; self.schema_builder.cubes.insert(self.cube_name, cube); @@ -538,6 +568,7 @@ impl MockViewBuilder { dimensions: all_dimensions, segments: all_segments, pre_aggregations: Vec::new(), + granularities: HashMap::new(), }; self.schema_builder.cubes.insert(self.view_name, view_cube); @@ -1304,4 +1335,75 @@ mod tests { fn test_from_yaml_file_not_found() { MockSchema::from_yaml_file("nonexistent.yaml"); } + + #[test] + fn test_schema_with_granularities() { + use crate::test_fixtures::cube_bridge::MockGranularityDefinition; + + let schema = MockSchemaBuilder::new() + .add_cube("orders") + .add_dimension( + "id", + MockDimensionDefinition::builder() + .dimension_type("number".to_string()) + .sql("id".to_string()) + .primary_key(Some(true)) + .build(), + ) + .add_dimension( + "created_at", + MockDimensionDefinition::builder() + .dimension_type("time".to_string()) + .sql("created_at".to_string()) + .build(), + ) + .add_granularity( + "created_at", + "half_year", + MockGranularityDefinition::builder() + .interval("6 months") + .origin("2024-01-01") + .build(), + ) + .add_granularity( + "created_at", + "fiscal_year", + MockGranularityDefinition::builder() + .interval("1 year") + .offset("1 month") + .build(), + ) + .finish_cube() + .build(); + + // Verify granularity accessor + let half_year = schema + .get_granularity("orders", "created_at", "half_year") + .expect("half_year should exist"); + assert_eq!(half_year.static_data().interval, "6 months"); + assert_eq!( + half_year.static_data().origin, + Some("2024-01-01".to_string()) + ); + + let fiscal_year = schema + .get_granularity("orders", "created_at", "fiscal_year") + .expect("fiscal_year should exist"); + assert_eq!(fiscal_year.static_data().interval, "1 year"); + assert_eq!( + fiscal_year.static_data().offset, + Some("1 month".to_string()) + ); + + // Missing granularity returns None + assert!(schema + .get_granularity("orders", "created_at", "nonexistent") + .is_none()); + assert!(schema + .get_granularity("orders", "id", "half_year") + .is_none()); + assert!(schema + .get_granularity("nonexistent", "created_at", "half_year") + .is_none()); + } } diff --git a/rust/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/yaml/base_query_options.rs b/rust/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/yaml/base_query_options.rs index de51b83d809c3..843813f6a820f 100644 --- a/rust/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/yaml/base_query_options.rs +++ b/rust/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/yaml/base_query_options.rs @@ -34,6 +34,8 @@ pub struct YamlBaseQueryOptions { pub cubestore_support_multistage: Option, #[serde(default)] pub disable_external_pre_aggregations: Option, + #[serde(default)] + pub pre_aggregation_id: Option, } #[derive(Debug, Deserialize)] diff --git a/rust/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/yaml/dimension.rs b/rust/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/yaml/dimension.rs index 4330a71ee4306..0a38329b2b3bc 100644 --- a/rust/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/yaml/dimension.rs +++ b/rust/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/yaml/dimension.rs @@ -1,10 +1,25 @@ use crate::cube_bridge::case_variant::CaseVariant; use crate::test_fixtures::cube_bridge::yaml::case::YamlCaseVariant; use crate::test_fixtures::cube_bridge::yaml::timeshift::YamlTimeShiftDefinition; -use crate::test_fixtures::cube_bridge::MockDimensionDefinition; +use crate::test_fixtures::cube_bridge::{MockDimensionDefinition, MockGranularityDefinition}; use serde::Deserialize; use std::rc::Rc; +#[derive(Debug, Deserialize)] +pub struct YamlGranularityEntry { + pub name: String, + pub interval: String, + #[serde(default)] + pub origin: Option, + #[serde(default)] + pub offset: Option, +} + +pub struct YamlDimensionBuildResult { + pub definition: MockDimensionDefinition, + pub granularities: Vec<(String, MockGranularityDefinition)>, +} + #[derive(Debug, Deserialize)] pub struct YamlDimensionDefinition { #[serde(rename = "type")] @@ -31,10 +46,12 @@ pub struct YamlDimensionDefinition { longitude: Option, #[serde(default)] time_shift: Vec, + #[serde(default)] + granularities: Vec, } impl YamlDimensionDefinition { - pub fn build(self) -> Rc { + pub fn build(self) -> YamlDimensionBuildResult { let time_shift = if !self.time_shift.is_empty() { Some(self.time_shift.into_iter().map(|ts| ts.build()).collect()) } else { @@ -48,21 +65,37 @@ impl YamlDimensionDefinition { } }); - Rc::new( - MockDimensionDefinition::builder() - .dimension_type(self.dimension_type) - .multi_stage(self.multi_stage) - .add_group_by_references(self.add_group_by_references) - .sub_query(self.sub_query) - .propagate_filters_to_sub_query(self.propagate_filters_to_sub_query) - .values(self.values) - .primary_key(self.primary_key) - .sql_opt(self.sql) - .case(case) - .latitude_opt(self.latitude) - .longitude_opt(self.longitude) - .time_shift(time_shift) - .build(), - ) + let granularities: Vec<(String, MockGranularityDefinition)> = self + .granularities + .into_iter() + .map(|entry| { + let def = MockGranularityDefinition::builder() + .interval(entry.interval) + .origin_opt(entry.origin) + .offset_opt(entry.offset) + .build(); + (entry.name, def) + }) + .collect(); + + let definition = MockDimensionDefinition::builder() + .dimension_type(self.dimension_type) + .multi_stage(self.multi_stage) + .add_group_by_references(self.add_group_by_references) + .sub_query(self.sub_query) + .propagate_filters_to_sub_query(self.propagate_filters_to_sub_query) + .values(self.values) + .primary_key(self.primary_key) + .sql_opt(self.sql) + .case(case) + .latitude_opt(self.latitude) + .longitude_opt(self.longitude) + .time_shift(time_shift) + .build(); + + YamlDimensionBuildResult { + definition, + granularities, + } } } diff --git a/rust/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/yaml/schema.rs b/rust/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/yaml/schema.rs index 236e90a7cd989..047a3ae6f732d 100644 --- a/rust/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/yaml/schema.rs +++ b/rust/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/yaml/schema.rs @@ -111,11 +111,13 @@ impl YamlSchema { let mut cube_builder = builder.add_cube(cube.name).cube_def(cube_def); for dim_entry in cube.dimensions { - let dim_rc = dim_entry.definition.build(); - let dim_def = Rc::try_unwrap(dim_rc) - .ok() - .expect("Rc should have single owner"); - cube_builder = cube_builder.add_dimension(dim_entry.name, dim_def); + let result = dim_entry.definition.build(); + cube_builder = + cube_builder.add_dimension(dim_entry.name.clone(), result.definition); + for (gran_name, gran_def) in result.granularities { + cube_builder = + cube_builder.add_granularity(&dim_entry.name, &gran_name, gran_def); + } } for meas_entry in cube.measures { diff --git a/rust/cubesqlplanner/cubesqlplanner/src/test_fixtures/schemas/yaml_files/common/custom_granularity_test.yaml b/rust/cubesqlplanner/cubesqlplanner/src/test_fixtures/schemas/yaml_files/common/custom_granularity_test.yaml new file mode 100644 index 0000000000000..7e81575d887df --- /dev/null +++ b/rust/cubesqlplanner/cubesqlplanner/src/test_fixtures/schemas/yaml_files/common/custom_granularity_test.yaml @@ -0,0 +1,63 @@ +cubes: + - name: orders + sql: "SELECT * FROM orders" + dimensions: + - name: id + type: number + sql: id + primary_key: true + - name: status + type: string + sql: status + - name: created_at + type: time + sql: created_at + granularities: + - name: half_year + interval: "6 months" + origin: "2024-01-01" + - name: bi_weekly + interval: "2 weeks" + - name: fiscal_year + interval: "1 year" + offset: "1 month" + measures: + - name: count + type: count + sql: "COUNT(*)" + - name: total_amount + type: sum + sql: amount + - name: avg_amount + type: avg + sql: amount + pre_aggregations: + - name: custom_half_year_rollup + type: rollup + measures: + - count + - total_amount + - avg_amount + dimensions: + - status + time_dimension: created_at + granularity: half_year + + - name: daily_rollup + type: rollup + measures: + - count + - total_amount + - avg_amount + dimensions: + - status + time_dimension: created_at + granularity: day + + - name: custom_half_year_non_strict + type: rollup + measures: + - count + time_dimension: created_at + granularity: half_year + allow_non_strict_date_range_match: true diff --git a/rust/cubesqlplanner/cubesqlplanner/src/test_fixtures/test_utils/test_context.rs b/rust/cubesqlplanner/cubesqlplanner/src/test_fixtures/test_utils/test_context.rs index 33a3303ea768a..763bd66944b73 100644 --- a/rust/cubesqlplanner/cubesqlplanner/src/test_fixtures/test_utils/test_context.rs +++ b/rust/cubesqlplanner/cubesqlplanner/src/test_fixtures/test_utils/test_context.rs @@ -246,6 +246,7 @@ impl TestContext { .disable_external_pre_aggregations .unwrap_or(false), ) + .pre_aggregation_id(yaml_options.pre_aggregation_id) .build(), ) } diff --git a/rust/cubesqlplanner/cubesqlplanner/src/tests/pre_aggregation_sql_generation.rs b/rust/cubesqlplanner/cubesqlplanner/src/tests/pre_aggregation_sql_generation.rs index 2e0ac2b749c2a..f0d2e327dfe74 100644 --- a/rust/cubesqlplanner/cubesqlplanner/src/tests/pre_aggregation_sql_generation.rs +++ b/rust/cubesqlplanner/cubesqlplanner/src/tests/pre_aggregation_sql_generation.rs @@ -467,6 +467,165 @@ fn test_segment_no_match_missing_in_pre_agg() { assert!(pre_aggrs.is_empty()); } +// --- Custom granularity pre-aggregation tests --- + +#[test] +fn test_custom_granularity_full_match() { + let schema = MockSchema::from_yaml_file("common/custom_granularity_test.yaml") + .only_pre_aggregations(&["custom_half_year_rollup"]); + let ctx = TestContext::new(schema).unwrap(); + + let (sql, pre_aggrs) = ctx + .build_sql_with_used_pre_aggregations(indoc! {" + measures: + - orders.count + dimensions: + - orders.status + time_dimensions: + - dimension: orders.created_at + granularity: half_year + "}) + .unwrap(); + + assert_eq!(pre_aggrs.len(), 1); + assert_eq!(pre_aggrs[0].name(), "custom_half_year_rollup"); + + insta::assert_snapshot!(sql); +} + +#[test] +fn test_standard_pre_agg_coarser_custom_query() { + let schema = MockSchema::from_yaml_file("common/custom_granularity_test.yaml") + .only_pre_aggregations(&["daily_rollup"]); + let ctx = TestContext::new(schema).unwrap(); + + let (sql, pre_aggrs) = ctx + .build_sql_with_used_pre_aggregations(indoc! {" + measures: + - orders.count + dimensions: + - orders.status + time_dimensions: + - dimension: orders.created_at + granularity: half_year + "}) + .unwrap(); + + assert_eq!(pre_aggrs.len(), 1); + assert_eq!(pre_aggrs[0].name(), "daily_rollup"); + + insta::assert_snapshot!(sql); +} + +#[test] +fn test_custom_pre_agg_finer_query_no_match() { + let schema = MockSchema::from_yaml_file("common/custom_granularity_test.yaml") + .only_pre_aggregations(&["custom_half_year_rollup"]); + let ctx = TestContext::new(schema).unwrap(); + + let (_sql, pre_aggrs) = ctx + .build_sql_with_used_pre_aggregations(indoc! {" + measures: + - orders.count + dimensions: + - orders.status + time_dimensions: + - dimension: orders.created_at + granularity: day + "}) + .unwrap(); + + assert!(pre_aggrs.is_empty()); +} + +#[test] +fn test_custom_pre_agg_finer_standard_query_no_match() { + let schema = MockSchema::from_yaml_file("common/custom_granularity_test.yaml") + .only_pre_aggregations(&["custom_half_year_rollup"]); + let ctx = TestContext::new(schema).unwrap(); + + let (_sql, pre_aggrs) = ctx + .build_sql_with_used_pre_aggregations(indoc! {" + measures: + - orders.count + dimensions: + - orders.status + time_dimensions: + - dimension: orders.created_at + granularity: month + "}) + .unwrap(); + + assert!(pre_aggrs.is_empty()); +} + +#[test] +fn test_custom_granularity_non_additive_full_match() { + let schema = MockSchema::from_yaml_file("common/custom_granularity_test.yaml") + .only_pre_aggregations(&["custom_half_year_rollup"]); + let ctx = TestContext::new(schema).unwrap(); + + let (sql, pre_aggrs) = ctx + .build_sql_with_used_pre_aggregations(indoc! {" + measures: + - orders.avg_amount + dimensions: + - orders.status + time_dimensions: + - dimension: orders.created_at + granularity: half_year + "}) + .unwrap(); + + assert_eq!(pre_aggrs.len(), 1); + assert_eq!(pre_aggrs[0].name(), "custom_half_year_rollup"); + + insta::assert_snapshot!(sql); +} + +#[test] +fn test_custom_granularity_non_additive_coarser_no_match() { + let schema = MockSchema::from_yaml_file("common/custom_granularity_test.yaml") + .only_pre_aggregations(&["daily_rollup"]); + let ctx = TestContext::new(schema).unwrap(); + + let (_sql, pre_aggrs) = ctx + .build_sql_with_used_pre_aggregations(indoc! {" + measures: + - orders.avg_amount + dimensions: + - orders.status + time_dimensions: + - dimension: orders.created_at + granularity: half_year + "}) + .unwrap(); + + assert!(pre_aggrs.is_empty()); +} + +#[test] +fn test_custom_granularity_non_strict_self_match() { + let schema = MockSchema::from_yaml_file("common/custom_granularity_test.yaml") + .only_pre_aggregations(&["custom_half_year_non_strict"]); + let ctx = TestContext::new(schema).unwrap(); + + let (sql, pre_aggrs) = ctx + .build_sql_with_used_pre_aggregations(indoc! {" + measures: + - orders.count + time_dimensions: + - dimension: orders.created_at + granularity: half_year + "}) + .unwrap(); + + assert_eq!(pre_aggrs.len(), 1); + assert_eq!(pre_aggrs[0].name(), "custom_half_year_non_strict"); + + insta::assert_snapshot!(sql); +} + #[test] fn test_segment_with_coarser_granularity() { let schema = MockSchema::from_yaml_file("common/pre_aggregation_matching_test.yaml") diff --git a/rust/cubesqlplanner/cubesqlplanner/src/tests/snapshots/cubesqlplanner__tests__pre_aggregation_sql_generation__custom_granularity_full_match.snap b/rust/cubesqlplanner/cubesqlplanner/src/tests/snapshots/cubesqlplanner__tests__pre_aggregation_sql_generation__custom_granularity_full_match.snap new file mode 100644 index 0000000000000..cb15946ac3039 --- /dev/null +++ b/rust/cubesqlplanner/cubesqlplanner/src/tests/snapshots/cubesqlplanner__tests__pre_aggregation_sql_generation__custom_granularity_full_match.snap @@ -0,0 +1,8 @@ +--- +source: cubesqlplanner/src/tests/pre_aggregation_sql_generation.rs +expression: sql +--- +SELECT "orders__status" "orders__status", ('2024-01-01T00:00:00.000' ::timestamp + INTERVAL '6 month' * FLOOR(EXTRACT(EPOCH FROM ("orders__created_at_half_year" - '2024-01-01T00:00:00.000'::timestamp)) / EXTRACT(EPOCH FROM INTERVAL '6 month'))) "orders__created_at_half_year", sum("orders__count") "orders__count" +FROM orders__custom_half_year_rollup AS "orders__custom_half_year_rollup" +GROUP BY 1, 2 +ORDER BY 2 ASC diff --git a/rust/cubesqlplanner/cubesqlplanner/src/tests/snapshots/cubesqlplanner__tests__pre_aggregation_sql_generation__custom_granularity_non_additive_full_match.snap b/rust/cubesqlplanner/cubesqlplanner/src/tests/snapshots/cubesqlplanner__tests__pre_aggregation_sql_generation__custom_granularity_non_additive_full_match.snap new file mode 100644 index 0000000000000..e13670ff0edea --- /dev/null +++ b/rust/cubesqlplanner/cubesqlplanner/src/tests/snapshots/cubesqlplanner__tests__pre_aggregation_sql_generation__custom_granularity_non_additive_full_match.snap @@ -0,0 +1,8 @@ +--- +source: cubesqlplanner/src/tests/pre_aggregation_sql_generation.rs +expression: sql +--- +SELECT "orders__status" "orders__status", ('2024-01-01T00:00:00.000' ::timestamp + INTERVAL '6 month' * FLOOR(EXTRACT(EPOCH FROM ("orders__created_at_half_year" - '2024-01-01T00:00:00.000'::timestamp)) / EXTRACT(EPOCH FROM INTERVAL '6 month'))) "orders__created_at_half_year", sum("orders__avg_amount") "orders__avg_amount" +FROM orders__custom_half_year_rollup AS "orders__custom_half_year_rollup" +GROUP BY 1, 2 +ORDER BY 2 ASC diff --git a/rust/cubesqlplanner/cubesqlplanner/src/tests/snapshots/cubesqlplanner__tests__pre_aggregation_sql_generation__custom_granularity_non_strict_self_match.snap b/rust/cubesqlplanner/cubesqlplanner/src/tests/snapshots/cubesqlplanner__tests__pre_aggregation_sql_generation__custom_granularity_non_strict_self_match.snap new file mode 100644 index 0000000000000..1671948991514 --- /dev/null +++ b/rust/cubesqlplanner/cubesqlplanner/src/tests/snapshots/cubesqlplanner__tests__pre_aggregation_sql_generation__custom_granularity_non_strict_self_match.snap @@ -0,0 +1,8 @@ +--- +source: cubesqlplanner/src/tests/pre_aggregation_sql_generation.rs +expression: sql +--- +SELECT ('2024-01-01T00:00:00.000' ::timestamp + INTERVAL '6 month' * FLOOR(EXTRACT(EPOCH FROM ("orders__created_at_half_year" - '2024-01-01T00:00:00.000'::timestamp)) / EXTRACT(EPOCH FROM INTERVAL '6 month'))) "orders__created_at_half_year", sum("orders__count") "orders__count" +FROM orders__custom_half_year_non_strict AS "orders__custom_half_year_non_strict" +GROUP BY 1 +ORDER BY 1 ASC diff --git a/rust/cubesqlplanner/cubesqlplanner/src/tests/snapshots/cubesqlplanner__tests__pre_aggregation_sql_generation__standard_pre_agg_coarser_custom_query.snap b/rust/cubesqlplanner/cubesqlplanner/src/tests/snapshots/cubesqlplanner__tests__pre_aggregation_sql_generation__standard_pre_agg_coarser_custom_query.snap new file mode 100644 index 0000000000000..da109d805b174 --- /dev/null +++ b/rust/cubesqlplanner/cubesqlplanner/src/tests/snapshots/cubesqlplanner__tests__pre_aggregation_sql_generation__standard_pre_agg_coarser_custom_query.snap @@ -0,0 +1,8 @@ +--- +source: cubesqlplanner/src/tests/pre_aggregation_sql_generation.rs +expression: sql +--- +SELECT "orders__status" "orders__status", ('2024-01-01T00:00:00.000' ::timestamp + INTERVAL '6 month' * FLOOR(EXTRACT(EPOCH FROM ("orders__created_at_day" - '2024-01-01T00:00:00.000'::timestamp)) / EXTRACT(EPOCH FROM INTERVAL '6 month'))) "orders__created_at_half_year", sum("orders__count") "orders__count" +FROM orders__daily_rollup AS "orders__daily_rollup" +GROUP BY 1, 2 +ORDER BY 2 ASC