diff --git a/packages/cubejs-query-orchestrator/test/unit/QueryQueue.abstract.ts b/packages/cubejs-query-orchestrator/test/unit/QueryQueue.abstract.ts index 8f3244eb98400..dc1fff74cfb93 100644 --- a/packages/cubejs-query-orchestrator/test/unit/QueryQueue.abstract.ts +++ b/packages/cubejs-query-orchestrator/test/unit/QueryQueue.abstract.ts @@ -250,29 +250,40 @@ export const QueryQueueTest = (name: string, options: QueryQueueTestOptions) => const onlyLocalTest = options.cacheAndQueueDriver !== 'cubestore' ? test : xtest; test('orphaned', async () => { - // recover if previous test broken something - for (let i = 1; i <= 4; i++) { - await queue.executeInQueue('delay', `11${i}`, { delay: 50, result: `${i}` }, 0); - } - cancelledQuery = null; - delayCount = 0; - let result = queue.executeInQueue('delay', '111', { delay: 800, result: '1' }, 0); - delayFn(null, 50).then(() => queue.executeInQueue('delay', '112', { delay: 800, result: '2' }, 0)).catch(e => e); - delayFn(null, 75).then(() => queue.executeInQueue('delay', '113', { delay: 800, result: '3' }, 0)).catch(e => e); - // orphaned timeout should be applied - delayFn(null, 100).then(() => queue.executeInQueue('delay', '114', { delay: 900, result: '4' }, 0)).catch(e => e); + // Two queries hold the single worker slot. orphanedTimeout keeps them out of the + // orphaned set themselves: the memory driver reports active queries as orphaned once + // their timeout passes, Cube Store does not. + const pending = [ + queue.executeInQueue('delay', '121', { delay: 1200, result: '1', orphanedTimeout: 60 }, 0).catch(e => e), + ]; + await delayFn(null, 50); + pending.push(queue.executeInQueue('delay', '122', { delay: 1200, result: '2', orphanedTimeout: 60 }, 0).catch(e => e)); + await delayFn(null, 50); + // 121 and 122 keep the worker busy for ~2.4s, so this one is still queued when its + // 1s orphaned timeout expires + pending.push(queue.executeInQueue('delay', '123', { delay: 50, result: '3', orphanedTimeout: 1 }, 0).catch(e => e)); + + // Reconciliation is what cancels orphaned queries and nothing else triggers it while + // the worker is busy. + const deadline = Date.now() + 2000; + while (cancelledQuery !== '123' && Date.now() < deadline) { + await queue.reconcileQueue(); + await delayFn(null, 100); + } - expect(await result).toBe('10'); - await queue.executeInQueue('delay', '112', { delay: 800, result: '2' }, 0); + expect(cancelledQuery).toBe('123'); - result = await queue.executeInQueue('delay', '113', { delay: 900, result: '3' }, 0); - expect(result).toBe('32'); + // every client gave up on ContinueWaitError long before this point + const outcomes = await Promise.all(pending); + outcomes.forEach((e) => expect(e).toBeInstanceOf(ContinueWaitError)); + await awaitProcessing(); - await delayFn(null, 500); - expect(cancelledQuery).toBe('114'); - await queue.executeInQueue('delay', '114', { delay: 50, result: '4' }, 0); + // 123 was cancelled before the worker could pick it up + expect(delayCount).toBe(2); + // cancellation removed it from the queue, so the same key can be queued again + expect(await queue.executeInQueue('delay', '123', { delay: 50, result: '3' }, 0)).toBe('32'); }); test('orphaned with custom ttl', async () => { diff --git a/packages/cubejs-schema-compiler/test/unit/pre-agg-interpolated-cube-refs.test.ts b/packages/cubejs-schema-compiler/test/unit/pre-agg-interpolated-cube-refs.test.ts new file mode 100644 index 0000000000000..fdd1a8e7a333f --- /dev/null +++ b/packages/cubejs-schema-compiler/test/unit/pre-agg-interpolated-cube-refs.test.ts @@ -0,0 +1,122 @@ +import { prepareJsCompiler } from './PrepareCompiler'; +import { PostgresQuery } from '../../src/adapter/PostgresQuery'; + +/** + * Pre-aggregation references built by interpolating the cube itself + * (`` (CUBE) => `${CUBE}.issued_date` ``) instead of a member. The cube + * stringifies to its name, so such a reference names the member as text; both + * planners have to resolve it to the same member. + */ +describe('pre-aggregation references interpolating the cube', () => { + const model = (preAggregations: string) => ` + const getCubeFields = (cube, names) => names.map((name) => cube[name]); + + cube('invoices', { + sql: 'SELECT * FROM invoices', + + measures: { + count: { type: 'count' }, + total: { sql: 'amount', type: 'sum' }, + }, + + dimensions: { + id: { sql: 'id', type: 'number', primaryKey: true }, + org_id: { sql: 'org_id', type: 'string' }, + issued_date: { sql: 'issued_date', type: 'time' }, + payment_received_date: { sql: 'payment_received_date', type: 'time' }, + }, + + preAggregations: ${preAggregations}, + }); + `; + + const preAggregationsWithInterpolatedTimeDimension = `{ + by_org_and_issued_date: { + type: 'rollup', + measures: (CUBE) => getCubeFields(CUBE, ['count', 'total']), + dimensions: (CUBE) => getCubeFields(CUBE, ['org_id']), + timeDimension: (CUBE) => \`\${CUBE}.issued_date\`, + granularity: 'day', + partitionGranularity: 'month', + }, + by_org_and_payment_received_date: { + type: 'rollup', + measures: (CUBE) => getCubeFields(CUBE, ['count', 'total']), + dimensions: (CUBE) => getCubeFields(CUBE, ['org_id']), + timeDimension: (CUBE) => \`\${CUBE}.payment_received_date\`, + granularity: 'day', + partitionGranularity: 'month', + }, + by_org_all_time: { + type: 'rollup', + measures: (CUBE) => getCubeFields(CUBE, ['count', 'total']), + dimensions: (CUBE) => getCubeFields(CUBE, ['org_id']), + }, + }`; + + const preAggregationsWithInterpolatedMembers = `{ + by_org_and_issued_date: { + type: 'rollup', + measures: (CUBE) => [\`\${CUBE}.count\`, \`\${CUBE}.total\`], + dimensions: (CUBE) => [\`\${CUBE}.org_id\`], + timeDimension: (CUBE) => \`\${CUBE}.issued_date\`, + granularity: 'day', + partitionGranularity: 'month', + }, + }`; + + const preAggregationWithGranularitySuffix = `{ + by_org_and_issued_date: { + type: 'rollup', + measures: (CUBE) => getCubeFields(CUBE, ['count']), + timeDimension: (CUBE) => \`\${CUBE}.issued_date_day\`, + granularity: 'day', + partitionGranularity: 'month', + }, + }`; + + async function buildQuery(preAggregations: string, useNativeSqlPlanner: boolean) { + const { compiler, joinGraph, cubeEvaluator } = prepareJsCompiler(model(preAggregations)); + await compiler.compile(); + + return new PostgresQuery({ joinGraph, cubeEvaluator, compiler }, { + measures: ['invoices.count'], + dimensions: ['invoices.org_id'], + timeDimensions: [{ + dimension: 'invoices.issued_date', + granularity: 'day', + dateRange: ['2020-01-01', '2020-03-31'], + }], + timezone: 'UTC', + useNativeSqlPlanner, + }); + } + + for (const useNativeSqlPlanner of [false, true]) { + const planner = useNativeSqlPlanner ? 'tesseract' : 'legacy'; + + it(`resolves an interpolated time dimension (${planner})`, async () => { + const query = await buildQuery(preAggregationsWithInterpolatedTimeDimension, useNativeSqlPlanner); + query.buildSqlAndParams(); + + const descriptions: any = query.preAggregations?.preAggregationsDescription(); + expect(descriptions.map(d => d.preAggregationId)).toEqual(['invoices.by_org_and_issued_date']); + }); + + it(`resolves interpolated measure and dimension references (${planner})`, async () => { + const query = await buildQuery(preAggregationsWithInterpolatedMembers, useNativeSqlPlanner); + query.buildSqlAndParams(); + + const descriptions: any = query.preAggregations?.preAggregationsDescription(); + expect(descriptions.map(d => d.preAggregationId)).toEqual(['invoices.by_org_and_issued_date']); + }); + + it(`reports the member an interpolated reference names when it does not exist (${planner})`, async () => { + const query = await buildQuery(preAggregationWithGranularitySuffix, useNativeSqlPlanner); + + expect(() => query.buildSqlAndParams()).toThrow( + /'issued_date_day' not found for path 'invoices.issued_date_day'/ + ); + }); + } +}); diff --git a/rust/cube/cubesqlplanner/cubesqlplanner/src/logical_plan/optimizers/pre_aggregation/optimizer.rs b/rust/cube/cubesqlplanner/cubesqlplanner/src/logical_plan/optimizers/pre_aggregation/optimizer.rs index 25754b62f082a..1a582b7af4b99 100644 --- a/rust/cube/cubesqlplanner/cubesqlplanner/src/logical_plan/optimizers/pre_aggregation/optimizer.rs +++ b/rust/cube/cubesqlplanner/cubesqlplanner/src/logical_plan/optimizers/pre_aggregation/optimizer.rs @@ -3,6 +3,7 @@ use super::*; use crate::logical_plan::visitor::{LogicalPlanRewriter, NodeRewriteResult}; use crate::logical_plan::*; use crate::planner::collectors::{collect_cube_names_from_symbols, has_multi_stage_members}; +use crate::planner::filter::typed_filter::resolve_base_symbol; use crate::planner::filter::FilterItem; use crate::planner::filter::FilterOp; use crate::planner::join_hints::JoinHints; @@ -140,9 +141,13 @@ impl PreAggregationOptimizer { let external = pre_aggregation.external.unwrap_or(false); let date_range = Self::extract_date_range(&query.filter(), &self.query_tools, time_shifts, external); - if let Some(rewritten) = - self.try_rewrite_simple_query(query, pre_aggregation, date_range, is_user_query)? - { + if let Some(rewritten) = self.try_rewrite_simple_query( + query, + pre_aggregation, + date_range, + is_user_query, + time_shifts, + )? { return Ok(Some(rewritten)); } } @@ -156,6 +161,7 @@ impl PreAggregationOptimizer { pre_aggregation: &Rc, date_range: Option<(String, String)>, is_user_query: bool, + time_shifts: &TimeShiftState, ) -> Result>, CubeError> { // Row identity for an ungrouped read is judged against the join this // very node will render, taken from the node itself rather than @@ -174,6 +180,14 @@ impl PreAggregationOptimizer { pre_aggregation, row_grain, )? { + if !Self::can_carry_time_shifts( + pre_aggregation, + &matched_measures, + &Self::read_member_names(&query.schema(), &query.filter()), + time_shifts, + ) { + return Ok(None); + } let source = self.make_pre_aggregation_source(pre_aggregation, &matched_measures, date_range)?; let new_query = Query::builder() @@ -477,6 +491,88 @@ impl PreAggregationOptimizer { } } + // A stored member is shifted by offsetting its column as a whole, which + // only reproduces the shifted values when the shift can be attributed to + // that column. A column built from several members of which just some are + // shifted has no such offset — moving it would carry along rows the shift + // must leave in place — and the lookup cannot attribute a shift to it + // either, so the two agree: whenever a shift is involved but cannot be + // attributed, the pre-aggregation cannot serve the shifted leaf. + // + // Grouping members — time dimensions, dimensions and segments — are + // substituted by column, so a pre-aggregation can only serve a shifted + // leaf when every stored member a shift reaches is one whose column can + // carry that shift. `shift_for_substituted_column` decides that, and the + // rendering node asks it too, so a member admitted here is one that will + // actually be offset. + // + // A measure column holds an aggregate, and a shift changes which rows + // feed it rather than the value itself, so no offset applies at all. A + // stored measure reading a shifted member is therefore always unusable, + // however cleanly the shift could be attributed to it. Only the measures + // matching consumed are examined, since the rest are never read. + // Resolved names of every member the query reads, so a stored member no + // one reads cannot decide anything. + fn read_member_names(schema: &LogicalSchema, filter: &LogicalFilter) -> HashSet { + let mut symbols: Vec> = schema + .dimensions + .iter() + .chain(schema.time_dimensions.iter()) + .chain(schema.measures.iter()) + .cloned() + .collect(); + for item in filter + .dimensions_filters + .iter() + .chain(filter.time_dimensions_filters.iter()) + .chain(filter.segments.iter()) + { + item.find_all_member_evaluators(&mut symbols); + } + symbols + .into_iter() + .map(|symbol| { + resolve_base_symbol(&symbol) + .resolve_reference_chain() + .full_name() + }) + .collect() + } + + fn can_carry_time_shifts( + pre_aggregation: &CompiledPreAggregation, + matched_measures: &HashSet, + read_members: &HashSet, + time_shifts: &TimeShiftState, + ) -> bool { + if time_shifts.is_empty() { + return true; + } + let is_read = |member: &Rc| { + read_members.contains( + &resolve_base_symbol(member) + .resolve_reference_chain() + .full_name(), + ) + }; + let grouping_members_carry_shift = pre_aggregation + .time_dimensions + .iter() + .chain(pre_aggregation.dimensions.iter()) + .chain(pre_aggregation.segments.iter()) + .filter(|member| is_read(member)) + .all(|member| { + !time_shifts.has_shift_under(member) + || time_shifts.shift_for_substituted_column(member).is_some() + }); + grouping_members_carry_shift + && pre_aggregation + .measures + .iter() + .filter(|measure| matched_measures.contains(&measure.full_name())) + .all(|measure| !time_shifts.has_shift_under(measure)) + } + fn extract_date_range( filter: &LogicalFilter, query_tools: &Rc, @@ -496,8 +592,7 @@ impl PreAggregationOptimizer { // Apply time shift for this dimension if present. // SQL renders `column + interval`, so actual data range is `date - interval`. if let Some(interval) = time_shifts - .dimensions_shifts - .get(&base_filter.member_name()) + .get_for_symbol(base_filter.raw_member_evaluator_ref()) .and_then(|s| s.interval.as_ref()) { let tz = query_tools.timezone(); diff --git a/rust/cube/cubesqlplanner/cubesqlplanner/src/logical_plan/optimizers/pre_aggregation/pre_aggregations_compiler.rs b/rust/cube/cubesqlplanner/cubesqlplanner/src/logical_plan/optimizers/pre_aggregation/pre_aggregations_compiler.rs index ed31b21fc2b07..df68099ec8e2e 100644 --- a/rust/cube/cubesqlplanner/cubesqlplanner/src/logical_plan/optimizers/pre_aggregation/pre_aggregations_compiler.rs +++ b/rust/cube/cubesqlplanner/cubesqlplanner/src/logical_plan/optimizers/pre_aggregation/pre_aggregations_compiler.rs @@ -13,11 +13,17 @@ use crate::planner::multi_fact_join_groups::{MeasuresJoinHints, MultiFactJoinGro use crate::planner::planners::JoinPlanner; use crate::planner::planners::ResolvedJoinItem; use crate::planner::state::State; +use crate::planner::Compiler; use crate::planner::GranularityHelper; use crate::planner::MemberSymbol; +use crate::planner::SqlCall; +use crate::planner::SqlCallReference; +use crate::planner::SymbolPath; +use crate::planner::SymbolPathType; use crate::planner::TimeDimensionSymbol; use crate::utils::debug::DebugSql; use cubenativeutils::CubeError; +use cubenativeutils::CubeErrorCauseType; use itertools::Itertools; use std::collections::HashMap; use std::fmt::Debug; @@ -112,19 +118,14 @@ impl PreAggregationsCompiler { } let measures = if let Some(refs) = description.measure_references()? { - Self::symbols_from_ref( - self.query_tools.clone(), - &name.cube_name, - refs, - Self::check_is_measure, - )? + Self::symbols_from_ref(self.query_tools.clone(), name, refs, Self::check_is_measure)? } else { Vec::new() }; let dimensions = if let Some(refs) = description.dimension_references()? { Self::symbols_from_ref( self.query_tools.clone(), - &name.cube_name, + name, refs, Self::check_is_dimension, )? @@ -192,12 +193,7 @@ impl PreAggregationsCompiler { 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, - )? + Self::symbols_from_ref(self.query_tools.clone(), name, refs, Self::check_is_segment)? } else { Vec::new() }; @@ -553,19 +549,19 @@ impl PreAggregationsCompiler { fn symbols_from_ref Result<(), CubeError>>( query_tools: Rc, - cube_name: &String, + name: &PreAggregationFullName, ref_func: Rc, check_type_fn: F, ) -> Result>, CubeError> { let evaluator_compiler_cell = query_tools.compiler().clone(); let mut evaluator_compiler = evaluator_compiler_cell.borrow_mut(); - let sql_call = evaluator_compiler.compile_sql_call(cube_name, ref_func)?; - let mut res = Vec::new(); - for symbol in sql_call.get_dependencies().iter() { - check_type_fn(&symbol)?; - res.push(symbol.clone()); + let sql_call = evaluator_compiler.compile_sql_call(&name.cube_name, ref_func)?; + let symbols = + Self::reference_symbols(&query_tools, &mut evaluator_compiler, name, &sql_call)?; + for symbol in symbols.iter() { + check_type_fn(symbol)?; } - Ok(res) + Ok(symbols) } fn time_dimension_symbol_from_ref( @@ -577,24 +573,75 @@ impl PreAggregationsCompiler { let mut evaluator_compiler = evaluator_compiler_cell.borrow_mut(); let sql_call = evaluator_compiler.compile_sql_call(&name.cube_name, ref_func)?; - let mut symbols = Vec::new(); - - for symbol in sql_call.get_dependencies().into_iter() { - Self::check_is_time_dimension(&symbol)?; - symbols.push(symbol); + let symbols = + Self::reference_symbols(&query_tools, &mut evaluator_compiler, name, &sql_call)?; + for symbol in symbols.iter() { + Self::check_is_time_dimension(symbol)?; } symbols.into_iter().next().ok_or_else(|| { let path = sql_call.debug_sql(true); - let member_name = path.rsplit('.').next().unwrap_or(&path); - - CubeError::user(format!( - "'{}' not found for path '{}' in pre-aggregation '{}.{}'", - member_name, path, name.cube_name, name.name - )) + Self::reference_not_found_error(&path, name) }) } + /// Members a pre-aggregation reference declaration names, in declaration + /// order. An element that interpolated the cube instead of the member is + /// resolved here; an element naming no member at all is an error, so a + /// reference is never silently dropped. + fn reference_symbols( + query_tools: &Rc, + evaluator_compiler: &mut Compiler, + name: &PreAggregationFullName, + sql_call: &SqlCall, + ) -> Result>, CubeError> { + let mut result = Vec::new(); + for item in sql_call.reference_items() { + let symbol = match item { + SqlCallReference::Symbol(symbol) => symbol, + SqlCallReference::Path(path) => { + let full_name = path.join("."); + let symbol_path = + SymbolPath::parse(query_tools.cube_evaluator().clone(), &full_name) + // A path the data model doesn't know is reported with the + // pre-aggregation it came from; anything else (a failure + // reaching the model at all) is passed through as it is. + .map_err(|e| match e.cause { + CubeErrorCauseType::User => { + Self::reference_not_found_error(&full_name, name) + } + _ => e, + })?; + match symbol_path.path_type() { + SymbolPathType::Dimension => { + evaluator_compiler.add_dimension_evaluator_by_path(symbol_path)? + } + SymbolPathType::Measure => { + evaluator_compiler.add_measure_evaluator_by_path(symbol_path)? + } + SymbolPathType::Segment => { + evaluator_compiler.add_segment_evaluator_by_path(symbol_path)? + } + _ => return Err(Self::reference_not_found_error(&full_name, name)), + } + } + SqlCallReference::Unresolved(rendered) => { + return Err(Self::reference_not_found_error(&rendered, name)) + } + }; + result.push(symbol); + } + Ok(result) + } + + fn reference_not_found_error(path: &str, name: &PreAggregationFullName) -> CubeError { + let member_name = path.rsplit('.').next().unwrap_or(path); + CubeError::user(format!( + "'{}' not found for path '{}' in pre-aggregation '{}.{}'", + member_name, path, name.cube_name, name.name + )) + } + fn check_is_measure(symbol: &MemberSymbol) -> Result<(), CubeError> { symbol .as_measure() @@ -640,13 +687,11 @@ mod tests { use crate::test_fixtures::test_utils::TestContext; use indoc::indoc; - fn create_time_dimension_context() -> TestContext { - // `time_dimension: \"{CUBE}.created_at\"` models a JS reference built - // via string interpolation — `(CUBE) => `${CUBE}.created_at``. The JS - // planner resolves it (reference evaluation stringifies `${CUBE}` to - // the cube name), but here `{CUBE}` compiles to a cube reference and - // `.created_at` stays literal text, so the compiled reference has no - // member symbol dependencies. + fn create_reference_context() -> TestContext { + // Template syntax like `\"{CUBE}.created_at\"` models a reference built by + // interpolating the cube itself — `(CUBE) => `${CUBE}.created_at``, where + // the member name arrives as literal text next to a cube reference + // instead of as a member symbol. let schema = MockSchema::from_yaml(indoc! {" cubes: - name: orders @@ -656,36 +701,88 @@ mod tests { type: number sql: id primary_key: true + - name: status + type: string + sql: status - name: created_at type: time sql: created_at + - name: city + type: string + sql: city measures: - name: count type: count + - name: total + type: sum + sql: amount + segments: + - name: completed + sql: \"{CUBE}.status = 'completed'\" pre_aggregations: - - name: working_rollup + - name: symbol_rollup type: rollup measures: - count time_dimension: created_at granularity: day - - name: broken_rollup_unsupported + - name: interpolated_rollup type: rollup measures: - - count + - '{CUBE}.count' + dimensions: + - '{CUBE}.status' + segments: + - '{CUBE}.completed' time_dimension: '{CUBE}.created_at' granularity: day - - name: broken_rollup_no_granularity + - name: interpolated_rollup_no_granularity type: rollup measures: - count time_dimension: '{CUBE}.created_at' + - name: interpolated_rollup_mixed_list + type: rollup + measures: + - '{CUBE}.total' + - count + dimensions: + - status + - '{CUBE}.city' + time_dimension: '{CUBE}.created_at' + granularity: day - name: broken_rollup_granularity_suffix type: rollup measures: - count time_dimension: '{CUBE}.created_at_day' granularity: day + - name: broken_rollup_unknown_measure + type: rollup + measures: + - '{CUBE}.unknown_total' + time_dimension: created_at + granularity: day + - name: interpolated_rollup_granularity_segment + type: rollup + measures: + - count + time_dimension: '{CUBE}.created_at.day' + granularity: day + - name: symbol_rollup_granularity_segment + type: rollup + measures: + - count + time_dimension: created_at.day + granularity: day + - name: broken_rollup_expression_dimension + type: rollup + measures: + - count + dimensions: + - \"{CUBE}.status = 'completed'\" + time_dimension: created_at + granularity: day "}) .unwrap(); TestContext::new(schema).unwrap() @@ -704,8 +801,8 @@ mod tests { #[test] fn test_time_dimension_resolves_to_member_symbol() { - let ctx = create_time_dimension_context(); - let compiled = compile_pre_agg(&ctx, "working_rollup").unwrap(); + let ctx = create_reference_context(); + let compiled = compile_pre_agg(&ctx, "symbol_rollup").unwrap(); assert_eq!(compiled.time_dimensions.len(), 1); assert_eq!( @@ -716,32 +813,86 @@ mod tests { } #[test] - fn test_time_dimension_resolved_to_cube_ref_returns_error() { - let ctx = create_time_dimension_context(); - let err = compile_pre_agg(&ctx, "broken_rollup_unsupported") - .expect_err("Pre-aggregation with unresolvable time dimension should fail to compile"); + fn test_interpolated_references_resolve_to_member_symbols() { + let ctx = create_reference_context(); + let compiled = compile_pre_agg(&ctx, "interpolated_rollup").unwrap(); + assert_eq!( - err.message, - "'created_at' not found for path 'orders.created_at' in pre-aggregation 'orders.broken_rollup_unsupported'" + compiled + .measures + .iter() + .map(|m| m.full_name()) + .collect_vec(), + vec!["orders.count".to_string()] + ); + assert_eq!( + compiled + .dimensions + .iter() + .map(|d| d.full_name()) + .collect_vec(), + vec!["orders.status".to_string()] + ); + assert_eq!( + compiled + .segments + .iter() + .map(|sg| sg.full_name()) + .collect_vec(), + vec!["expr:orders.completed".to_string()] + ); + assert_eq!(compiled.time_dimensions.len(), 1); + assert_eq!( + compiled.time_dimensions[0].full_name(), + "orders.created_at_day" ); + assert_eq!(compiled.granularity, Some("day".to_string())); } #[test] - fn test_time_dimension_resolved_to_cube_ref_without_granularity_returns_error() { - let ctx = create_time_dimension_context(); - let err = compile_pre_agg(&ctx, "broken_rollup_no_granularity") - .expect_err("Pre-aggregation with unresolvable time dimension should fail to compile"); + fn test_interpolated_time_dimension_without_granularity_resolves() { + let ctx = create_reference_context(); + let compiled = compile_pre_agg(&ctx, "interpolated_rollup_no_granularity").unwrap(); + + assert_eq!(compiled.time_dimensions.len(), 1); + assert_eq!(compiled.time_dimensions[0].full_name(), "orders.created_at"); + assert_eq!(compiled.granularity, None); + } + + // One list mixing both forms keeps every member it names, in declaration + // order — join hints and lambda member matching read the list positionally. + #[test] + fn test_interpolated_and_symbol_references_in_one_list() { + let ctx = create_reference_context(); + let compiled = compile_pre_agg(&ctx, "interpolated_rollup_mixed_list").unwrap(); + assert_eq!( - err.message, - "'created_at' not found for path 'orders.created_at' in pre-aggregation 'orders.broken_rollup_no_granularity'" + compiled + .measures + .iter() + .map(|m| m.full_name()) + .collect_vec(), + vec!["orders.total".to_string(), "orders.count".to_string()] + ); + assert_eq!( + compiled + .dimensions + .iter() + .map(|d| d.full_name()) + .collect_vec(), + vec!["orders.status".to_string(), "orders.city".to_string()] + ); + assert_eq!( + compiled.time_dimensions[0].full_name(), + "orders.created_at_day" ); } - // Interpolated reference with a granularity-suffixed member name, - // e.g. `(CUBE) => `${CUBE}.created_at_day``. + // An interpolated reference naming the granularity-suffixed member instead + // of the member itself, e.g. `(CUBE) => `${CUBE}.created_at_day``. #[test] - fn test_time_dimension_with_granularity_suffix_returns_error() { - let ctx = create_time_dimension_context(); + fn test_interpolated_time_dimension_with_granularity_suffix_returns_error() { + let ctx = create_reference_context(); let err = compile_pre_agg(&ctx, "broken_rollup_granularity_suffix") .expect_err("Pre-aggregation with unresolvable time dimension should fail to compile"); assert_eq!( @@ -750,6 +901,48 @@ mod tests { ); } + // Naming the granularity inside the reference, `${CUBE}.created_at.day`, + // instead of through `granularity:`. Rejected — and rejected the same way as + // the equivalent symbol reference `CUBE.created_at.day`. + #[test] + fn test_granularity_inside_the_reference_is_rejected_like_the_symbol_form() { + let ctx = create_reference_context(); + let interpolated = compile_pre_agg(&ctx, "interpolated_rollup_granularity_segment") + .expect_err("Granularity inside a time dimension reference should fail to compile"); + let symbol = compile_pre_agg(&ctx, "symbol_rollup_granularity_segment") + .expect_err("Granularity inside a time dimension reference should fail to compile"); + + assert_eq!(interpolated.message, symbol.message); + assert_eq!( + interpolated.message, + "Pre-aggregation time dimension must be a dimension" + ); + } + + #[test] + fn test_interpolated_measure_that_does_not_exist_returns_error() { + let ctx = create_reference_context(); + let err = compile_pre_agg(&ctx, "broken_rollup_unknown_measure") + .expect_err("Pre-aggregation with unresolvable measure should fail to compile"); + assert_eq!( + err.message, + "'unknown_total' not found for path 'orders.unknown_total' in pre-aggregation 'orders.broken_rollup_unknown_measure'" + ); + } + + // An element built as an expression rather than a member reference names no + // member, so it is reported instead of dropped from the reference list. + #[test] + fn test_reference_that_names_no_member_returns_error() { + let ctx = create_reference_context(); + let err = compile_pre_agg(&ctx, "broken_rollup_expression_dimension") + .expect_err("Pre-aggregation with an expression reference should fail to compile"); + assert_eq!( + err.message, + "'status = 'completed'' not found for path 'orders.status = 'completed'' in pre-aggregation 'orders.broken_rollup_expression_dimension'" + ); + } + #[test] fn test_compile_simple_rollup() { let schema = MockSchema::from_yaml_file("common/pre_aggregations_test.yaml"); diff --git a/rust/cube/cubesqlplanner/cubesqlplanner/src/physical_plan/filter/base_filter.rs b/rust/cube/cubesqlplanner/cubesqlplanner/src/physical_plan/filter/base_filter.rs index c72348f2b5dd5..6a7f307b78446 100644 --- a/rust/cube/cubesqlplanner/cubesqlplanner/src/physical_plan/filter/base_filter.rs +++ b/rust/cube/cubesqlplanner/cubesqlplanner/src/physical_plan/filter/base_filter.rs @@ -27,8 +27,7 @@ impl ToSql for BaseFilter { { let time_shift = visitor .time_shifts() - .dimensions_shifts - .get(&symbol_to_match.full_name()) + .get_for_symbol(&symbol_to_match) .and_then(|shift| shift.interval.as_ref()); return self.typed_filter().to_sql_for_filter_params( filter_params_item, diff --git a/rust/cube/cubesqlplanner/cubesqlplanner/src/physical_plan/sql_nodes/factory.rs b/rust/cube/cubesqlplanner/cubesqlplanner/src/physical_plan/sql_nodes/factory.rs index e6a28fe17e3f6..92d22405403cd 100644 --- a/rust/cube/cubesqlplanner/cubesqlplanner/src/physical_plan/sql_nodes/factory.rs +++ b/rust/cube/cubesqlplanner/cubesqlplanner/src/physical_plan/sql_nodes/factory.rs @@ -287,7 +287,11 @@ impl SqlNodesFactory { }; let input = if !self.time_shifts.is_empty() { - TimeShiftSqlNode::new(self.time_shifts.clone(), input) + TimeShiftSqlNode::new( + self.time_shifts.clone(), + self.pre_aggregation_dimensions_references.clone(), + input, + ) } else { input }; diff --git a/rust/cube/cubesqlplanner/cubesqlplanner/src/physical_plan/sql_nodes/time_shift.rs b/rust/cube/cubesqlplanner/cubesqlplanner/src/physical_plan/sql_nodes/time_shift.rs index 6516be32ad44b..e572a1876fac5 100644 --- a/rust/cube/cubesqlplanner/cubesqlplanner/src/physical_plan/sql_nodes/time_shift.rs +++ b/rust/cube/cubesqlplanner/cubesqlplanner/src/physical_plan/sql_nodes/time_shift.rs @@ -1,4 +1,5 @@ use super::SqlNode; +use crate::physical_plan::sql_nodes::render_references::RenderReferences; use crate::physical_plan::SqlEvaluatorVisitor; use crate::planner::planners::multi_stage::TimeShiftState; use crate::planner::query_tools::QueryTools; @@ -11,14 +12,27 @@ use std::rc::Rc; /// Applies a per-dimension time shift to time dimensions whose /// full name is in `shifts`, by rendering the dimension expression /// shifted by the configured interval. +/// +/// `substituted` names the dimensions rendered as a stored column instead +/// of being evaluated. Their SQL is never expanded, so the shift cannot be +/// picked up further down and has to be applied to the column itself. pub struct TimeShiftSqlNode { shifts: TimeShiftState, + substituted: RenderReferences, input: Rc, } impl TimeShiftSqlNode { - pub fn new(shifts: TimeShiftState, input: Rc) -> Rc { - Rc::new(Self { shifts, input }) + pub fn new( + shifts: TimeShiftState, + substituted: RenderReferences, + input: Rc, + ) -> Rc { + Rc::new(Self { + shifts, + substituted, + input, + }) } pub fn input(&self) -> &Rc { @@ -38,8 +52,34 @@ impl SqlNode for TimeShiftSqlNode { let res = match node.as_ref() { MemberSymbol::Dimension(ev) => { if !ev.is_reference() && ev.is_time() { - if let Some(shift) = self.shifts.dimensions_shifts.get(&ev.full_name()) { - let shift = shift.interval.clone().unwrap().to_sql(); + // The first probe is by exact name on purpose: a dimension + // that gets evaluated has its shift applied when the + // recursion reaches the owned member it wraps, and matching + // it here as well would add the interval twice. Only a + // substituted dimension, which is never expanded, resolves + // through the chain. + let shift = self + .shifts + .dimensions_shifts + .get(&ev.full_name()) + .or_else(|| { + if self.substituted.contains_key(&ev.full_name()) { + self.shifts.shift_for_substituted_column(node) + } else { + None + } + }); + if let Some(shift) = shift { + let shift = shift + .interval + .as_ref() + .ok_or_else(|| { + CubeError::internal(format!( + "Time shift for dimension {} has no interval", + ev.full_name() + )) + })? + .to_sql(); let inner_visitor = visitor.with_arg_needs_paren_safe(false); let input = self.input.to_sql( &inner_visitor, diff --git a/rust/cube/cubesqlplanner/cubesqlplanner/src/planner/planners/multi_stage/time_shift_state.rs b/rust/cube/cubesqlplanner/cubesqlplanner/src/planner/planners/multi_stage/time_shift_state.rs index e603f0e21860d..25a6f97353f8e 100644 --- a/rust/cube/cubesqlplanner/cubesqlplanner/src/planner/planners/multi_stage/time_shift_state.rs +++ b/rust/cube/cubesqlplanner/cubesqlplanner/src/planner/planners/multi_stage/time_shift_state.rs @@ -1,7 +1,11 @@ +use crate::planner::collectors::find_owned_by_cube_child; +use crate::planner::filter::typed_filter::resolve_base_symbol; use crate::planner::symbols::CalendarDimensionTimeShift; +use crate::planner::symbols::MemberSymbol; use crate::planner::DimensionTimeShift; use cubenativeutils::CubeError; use std::collections::HashMap; +use std::rc::Rc; /// Per-dimension time-shift accumulator used during multi-stage /// planning. Keyed by dimension full name; aggregates the shifts @@ -16,6 +20,53 @@ impl TimeShiftState { self.dimensions_shifts.is_empty() } + /// Looks up the shift for a symbol that may still be wrapped in a + /// `TimeDimension`, be a reference to the shifted member, or wrap it in + /// its own SQL. Keys are built either from the chain-resolved dimension + /// or, for dimension-specific shifts, from the owned member the declared + /// dimension wraps, so both forms are probed. + pub fn get_for_symbol(&self, symbol: &Rc) -> Option<&DimensionTimeShift> { + let resolved = resolve_base_symbol(symbol).resolve_reference_chain(); + if let Some(shift) = self.dimensions_shifts.get(&resolved.full_name()) { + return Some(shift); + } + let owned = find_owned_by_cube_child(&resolved).ok()?; + self.dimensions_shifts.get(&owned.full_name()) + } + + /// The shift a stored column standing for this member can carry. + /// + /// A column is shifted by offsetting it, which only stands in for the + /// shifted member when the member is a time dimension evaluated in place: + /// a reference is rendered through to what it points at, and a non-time + /// member has no meaning under an interval. Both the gate that admits a + /// pre-aggregation and the node that renders from one ask this, so the two + /// cannot come to different conclusions. + pub fn shift_for_substituted_column( + &self, + symbol: &Rc, + ) -> Option<&DimensionTimeShift> { + let dimension = resolve_base_symbol(symbol).as_dimension().ok()?; + if dimension.is_reference() || !dimension.is_time() { + return None; + } + self.get_for_symbol(symbol) + } + + /// True when the symbol itself, or any member it is built from, is + /// shifted. Unlike `get_for_symbol` this answers whether a shift is + /// involved at all, not whether one can be attributed to the symbol. + pub fn has_shift_under(&self, symbol: &Rc) -> bool { + let symbol = resolve_base_symbol(symbol); + if self.dimensions_shifts.contains_key(&symbol.full_name()) { + return true; + } + symbol + .get_dependencies() + .iter() + .any(|dep| self.has_shift_under(dep)) + } + /// Splits the accumulated shifts into two maps: regular /// `DimensionTimeShift`s applied at render time, and /// `CalendarDimensionTimeShift`s that come from a calendar diff --git a/rust/cube/cubesqlplanner/cubesqlplanner/src/planner/sql_call.rs b/rust/cube/cubesqlplanner/cubesqlplanner/src/planner/sql_call.rs index b7aee3eb3417f..af0478eca5500 100644 --- a/rust/cube/cubesqlplanner/cubesqlplanner/src/planner/sql_call.rs +++ b/rust/cube/cubesqlplanner/cubesqlplanner/src/planner/sql_call.rs @@ -9,9 +9,18 @@ use crate::planner::{CubeNameSymbol, CubeTableSymbol}; use crate::utils::sql_expression_scanner::analyze_template_arg_contexts; use cubenativeutils::CubeError; use itertools::Itertools; +use lazy_static::lazy_static; +use regex::Regex; use std::collections::HashMap; use std::rc::Rc; +lazy_static! { + /// A whole template element made of a cube-name placeholder followed only by + /// dotted identifiers. + static ref INTERPOLATED_REFERENCE_RE: Regex = + Regex::new(r"^\s*\{arg:(\d+)\}((?:\.[A-Za-z_$][A-Za-z0-9_$]*)+)\s*$").unwrap(); +} + /// Reference to a cube from a SQL template. /// /// - `Name` — the cube as an identifier (rendered from `{CUBE}` or @@ -88,6 +97,18 @@ impl SqlDependency { } } +/// What one element of a reference declaration names, as read off the compiled +/// template by `SqlCall::reference_items`. +#[derive(Clone, Debug)] +pub enum SqlCallReference { + Symbol(Rc), + /// Member path of an element that interpolated the cube name, still to be + /// resolved against the data model. + Path(Vec), + /// Element naming no member, rendered for diagnostics. + Unresolved(String), +} + /// Namespace for the placeholder prefixes recognised inside a /// `SqlCall` template: /// @@ -547,6 +568,117 @@ impl SqlCall { } } + /// What each element of a reference declaration (a pre-aggregation + /// `measures:` / `dimensions:` / `segments:` / `time_dimension:`) names, + /// in declaration order. + /// + /// An element referencing a member yields that member's symbol. An element + /// that interpolated the cube itself — ``(CUBE) => `${CUBE}.created_at` `` — + /// depends on the cube name and keeps the member as literal text, so it + /// yields the path to resolve: the cube reference's path followed by the + /// literal segments. Anything else names no member and is reported as + /// unresolved, rendered for diagnostics. + /// + /// An element wrapping a member reference in an expression still yields that + /// member, since the member is a dependency of its own — only the cube-name + /// form has nothing to fall back on. + pub fn reference_items(&self) -> Vec { + let elements = match &self.template { + SqlTemplate::String(s) => std::slice::from_ref(s), + SqlTemplate::StringVec(strings) => strings.as_slice(), + }; + let mut taken = vec![false; self.deps.len()]; + let mut result = Vec::new(); + for element in elements { + let arg_indices = Self::template_arg_indices(element); + let names_symbol = arg_indices + .iter() + .any(|index| self.deps.get(*index).is_some_and(|dep| dep.is_symbol())); + if names_symbol { + for index in arg_indices { + // An index the recorded dependencies don't cover names + // nothing; the rest of the element is still read. + let Some(symbol) = self.deps.get(index).and_then(|dep| dep.as_symbol()) else { + continue; + }; + if taken[index] { + continue; + } + taken[index] = true; + result.push(SqlCallReference::Symbol(symbol.clone())); + } + continue; + } + match self.interpolated_reference_path(element) { + Some(path) => result.push(SqlCallReference::Path(path)), + None => result.push(SqlCallReference::Unresolved( + self.render_for_diagnostics(element), + )), + } + } + // A symbol no element referenced: keep it rather than lose a member the + // declaration depends on. + for (index, dep) in self.deps.iter().enumerate() { + if !taken[index] { + if let Some(symbol) = dep.as_symbol() { + result.push(SqlCallReference::Symbol(symbol.clone())); + } + } + } + result + } + + // Path of an element made of a cube-name placeholder followed only by dotted + // identifiers; `None` when the element has any other shape. + fn interpolated_reference_path(&self, element: &str) -> Option> { + let captures = INTERPOLATED_REFERENCE_RE.captures(element)?; + let index = captures.get(1)?.as_str().parse::().ok()?; + let cube_ref = self.deps.get(index)?.as_cube_ref()?.as_name()?; + let mut path = cube_ref.path().clone(); + path.extend( + captures + .get(2)? + .as_str() + .split('.') + .skip(1) + .map(String::from), + ); + Some(path) + } + + // `{arg:N}` indices in the order they appear in the element. + fn template_arg_indices(element: &str) -> Vec { + let needle = format!("{{{}:", SqlCallArg::ARG_PREFIX); + let mut result = Vec::new(); + let mut rest = element; + while let Some(start) = rest.find(&needle) { + rest = &rest[start + needle.len()..]; + let Some(end) = rest.find('}') else { + break; + }; + if let Ok(index) = rest[..end].parse::() { + result.push(index); + } + rest = &rest[end + 1..]; + } + result + } + + // Element with its dependencies replaced by the members and cubes they name, + // for error messages about an element that names no member. + fn render_for_diagnostics(&self, element: &str) -> String { + let deps = self + .deps + .iter() + .map(|dep| match dep { + SqlDependency::Symbol(symbol) => symbol.full_name(), + SqlDependency::CubeRef(cube_ref) => cube_ref.cube_name().clone(), + }) + .collect_vec(); + Self::substitute_template(element, &deps, &[], &[], &[], &[]) + .unwrap_or_else(|_| element.to_string()) + } + /// Number of member-symbol dependencies. Cube refs are not /// counted. pub fn dependencies_count(&self) -> usize { @@ -703,3 +835,206 @@ impl crate::utils::debug::DebugSql for SqlCall { .unwrap() } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::test_fixtures::cube_bridge::MockSchema; + use crate::test_fixtures::test_utils::TestContext; + use indoc::indoc; + + fn test_context() -> TestContext { + let schema = MockSchema::from_yaml(indoc! {" + 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 + measures: + - name: count + type: count + "}) + .unwrap(); + TestContext::new(schema).unwrap() + } + + fn cube_name_dep(cube_name: &str, path: &[&str]) -> SqlDependency { + SqlDependency::CubeRef(CubeRef::Name(CubeNameSymbol::new( + cube_name.to_string(), + path.iter().map(|p| p.to_string()).collect(), + ))) + } + + fn call(template: SqlTemplate, deps: Vec) -> SqlCall { + SqlCall::new( + template, + deps, + vec![], + vec![], + SecutityContextProps::default(), + ) + } + + fn single(template: &str, deps: Vec) -> SqlCall { + call(SqlTemplate::String(template.to_string()), deps) + } + + /// Reference items as comparable strings: `symbol:`, + /// `path:`, `unresolved:`. + fn described(sql_call: &SqlCall) -> Vec { + sql_call + .reference_items() + .iter() + .map(|item| match item { + SqlCallReference::Symbol(symbol) => format!("symbol:{}", symbol.full_name()), + SqlCallReference::Path(path) => format!("path:{}", path.join(".")), + SqlCallReference::Unresolved(rendered) => format!("unresolved:{}", rendered), + }) + .collect() + } + + #[test] + fn test_cube_name_followed_by_member_yields_path() { + let sql_call = single("{arg:0}.created_at", vec![cube_name_dep("orders", &[])]); + + assert_eq!(described(&sql_call), vec!["path:orders.created_at"]); + } + + #[test] + fn test_surrounding_whitespace_is_ignored() { + let sql_call = single(" {arg:0}.created_at\n", vec![cube_name_dep("orders", &[])]); + + assert_eq!(described(&sql_call), vec!["path:orders.created_at"]); + } + + #[test] + fn test_join_path_is_kept() { + // `${CUBE.users}.name` — the cube reference carries the cubes traversed, + // and the member name follows as literal text. + let sql_call = single("{arg:0}.name", vec![cube_name_dep("users", &["orders"])]); + + assert_eq!(described(&sql_call), vec!["path:orders.users.name"]); + } + + #[test] + fn test_literal_segments_after_the_cube_are_all_kept() { + let sql_call = single("{arg:0}.users.name", vec![cube_name_dep("orders", &[])]); + + assert_eq!(described(&sql_call), vec!["path:orders.users.name"]); + } + + #[test] + fn test_every_element_of_a_reference_list_is_resolved() { + let sql_call = call( + SqlTemplate::StringVec(vec![ + "{arg:0}.count".to_string(), + "{arg:1}.name".to_string(), + ]), + vec![cube_name_dep("orders", &[]), cube_name_dep("users", &[])], + ); + + assert_eq!( + described(&sql_call), + vec!["path:orders.count", "path:users.name"] + ); + } + + // A member reference produces a symbol dependency of its own. + #[test] + fn test_member_dependency_is_reported_as_a_symbol() { + let ctx = test_context(); + let symbol = ctx.create_dimension("orders.created_at").unwrap(); + let sql_call = single("{arg:0}", vec![SqlDependency::Symbol(symbol)]); + + assert_eq!(described(&sql_call), vec!["symbol:orders.created_at"]); + } + + // Declaration order survives a list mixing both forms — join hints and + // lambda member matching read the compiled list positionally. + #[test] + fn test_declaration_order_is_kept_for_a_mixed_list() { + let ctx = test_context(); + let symbol = ctx.create_dimension("orders.status").unwrap(); + let sql_call = call( + SqlTemplate::StringVec(vec![ + "{arg:0}.created_at".to_string(), + "{arg:1}".to_string(), + ]), + vec![cube_name_dep("orders", &[]), SqlDependency::Symbol(symbol)], + ); + + assert_eq!( + described(&sql_call), + vec!["path:orders.created_at", "symbol:orders.status"] + ); + } + + // An element naming a member through an expression names no single member. + #[test] + fn test_element_carrying_an_expression_is_unresolved() { + for (template, expected) in [ + ("{arg:0}.created_at + 1", "unresolved:orders.created_at + 1"), + ( + "date_trunc('day', {arg:0}.created_at)", + "unresolved:date_trunc('day', orders.created_at)", + ), + ("{arg:0}", "unresolved:orders"), + ("{arg:0}.", "unresolved:orders."), + ("{arg:0}.2days", "unresolved:orders.2days"), + ] { + let sql_call = single(template, vec![cube_name_dep("orders", &[])]); + + assert_eq!( + described(&sql_call), + vec![expected], + "unexpected reference items for `{}`", + template + ); + } + } + + // `${CUBE.sql()}` renders the cube's table expression, not its name, so it + // cannot start a member path. + #[test] + fn test_cube_table_reference_is_unresolved() { + let ctx = test_context(); + let cube_table = ctx + .query_tools() + .compiler() + .borrow_mut() + .add_cube_table_evaluator("orders".to_string(), vec![]) + .unwrap(); + let sql_call = single( + "{arg:0}.created_at", + vec![SqlDependency::CubeRef(CubeRef::Table(cube_table))], + ); + + assert_eq!(described(&sql_call), vec!["unresolved:orders.created_at"]); + } + + // An index the recorded dependencies don't cover must not be read as one. + #[test] + fn test_placeholder_out_of_bounds_next_to_a_member_is_skipped() { + let ctx = test_context(); + let symbol = ctx.create_dimension("orders.status").unwrap(); + let sql_call = single("{arg:0} || {arg:7}", vec![SqlDependency::Symbol(symbol)]); + + assert_eq!(described(&sql_call), vec!["symbol:orders.status"]); + } + + #[test] + fn test_placeholder_out_of_bounds_is_unresolved() { + let sql_call = single("{arg:3}.created_at", vec![cube_name_dep("orders", &[])]); + + assert_eq!(described(&sql_call), vec!["unresolved:{arg:3}.created_at"]); + } +} diff --git a/rust/cube/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/mock_member_sql.rs b/rust/cube/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/mock_member_sql.rs index 0e9546017538a..f7c734fe72255 100644 --- a/rust/cube/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/mock_member_sql.rs +++ b/rust/cube/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/mock_member_sql.rs @@ -104,12 +104,63 @@ impl MockMemberSql { })) } + /// Pre-aggregation array references where an element may interpolate the + /// cube itself: `["{CUBE}.count", "{CUBE.status}", "city"]`. A brace-free + /// element is a plain member path, recorded the way + /// `pre_agg_array_refs` records it, so both forms can be mixed in one list. + pub fn pre_agg_array_templates(members: Vec) -> Result, CubeError> { + let mut args = SqlTemplateArgs::default(); + let mut args_names = Vec::new(); + let mut template_elements = Vec::new(); + + for member in &members { + if member.contains('{') { + template_elements.push(Self::parse_template_into( + member, + &mut args, + &mut args_names, + )?); + } else { + let path_parts: Vec = member.split('.').map(|s| s.to_string()).collect(); + if path_parts.iter().any(|p| p.is_empty()) { + return Err(CubeError::user(format!( + "Invalid path in pre-aggregation: {}", + member + ))); + } + let arg_name = path_parts[0].clone(); + if !args_names.contains(&arg_name) { + args_names.push(arg_name); + } + let index = args.insert_symbol_path(path_parts); + template_elements.push(format!("{{arg:{}}}", index)); + } + } + + Ok(Rc::new(Self { + template: SqlTemplate::StringVec(template_elements), + args, + args_names, + })) + } + /// Parse the template string and extract symbol paths /// Converts "{path.to.symbol}" to "{arg:N}" and collects paths fn parse_template(template: &str) -> Result<(String, SqlTemplateArgs, Vec), CubeError> { - let mut result = String::new(); let mut args = SqlTemplateArgs::default(); let mut args_names = Vec::new(); + let result = Self::parse_template_into(template, &mut args, &mut args_names)?; + Ok((result, args, args_names)) + } + + // Parses one template, recording its dependencies into the given args so + // several templates can share one dependency list. + fn parse_template_into( + template: &str, + args: &mut SqlTemplateArgs, + args_names: &mut Vec, + ) -> Result { + let mut result = String::new(); let mut chars = template.chars().peekable(); @@ -172,8 +223,7 @@ impl MockMemberSql { // planner passes at render time. if let Some(body) = path.strip_prefix("FILTER_PARAMS:") { let (cube_name, name, column) = Self::parse_filter_params_body(body)?; - let column = - Self::parse_column_references(&column, &mut args, &mut args_names)?; + let column = Self::parse_column_references(&column, args, args_names)?; let index = args.insert_filter_params(FilterParamsItem { cube_name, name, @@ -216,7 +266,7 @@ impl MockMemberSql { } } - Ok((result, args, args_names)) + Ok(result) } // Splits a `.:` FILTER_PARAMS body. diff --git a/rust/cube/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/yaml/pre_aggregation.rs b/rust/cube/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/yaml/pre_aggregation.rs index 4a7323e46b0e3..4e6f93e90f1b7 100644 --- a/rust/cube/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/yaml/pre_aggregation.rs +++ b/rust/cube/cubesqlplanner/cubesqlplanner/src/test_fixtures/cube_bridge/yaml/pre_aggregation.rs @@ -161,7 +161,11 @@ impl YamlPreAggregationDefinition { } fn build_array_references(members: Vec) -> Result, CubeError> { - MockMemberSql::pre_agg_array_refs(members).map(|m| m as Rc) + if members.iter().any(|m| m.contains('{')) { + MockMemberSql::pre_agg_array_templates(members).map(|m| m as Rc) + } else { + MockMemberSql::pre_agg_array_refs(members).map(|m| m as Rc) + } } fn build_single_reference(member: String) -> Result, CubeError> { diff --git a/rust/cube/cubesqlplanner/cubesqlplanner/src/test_fixtures/schemas/yaml_files/common/integration_derived_time_dim_shift_pre_agg.yaml b/rust/cube/cubesqlplanner/cubesqlplanner/src/test_fixtures/schemas/yaml_files/common/integration_derived_time_dim_shift_pre_agg.yaml new file mode 100644 index 0000000000000..50a75a61c0549 --- /dev/null +++ b/rust/cube/cubesqlplanner/cubesqlplanner/src/test_fixtures/schemas/yaml_files/common/integration_derived_time_dim_shift_pre_agg.yaml @@ -0,0 +1,57 @@ +cubes: + - name: pa_customers + sql: "SELECT * FROM pa_customers" + joins: + - name: pa_returns + relationship: one_to_many + sql: "{pa_customers}.id = {pa_returns.customer_id}" + dimensions: + - name: id + type: number + sql: id + primary_key: true + + # Time dimension derived from another cube's time dimension: + # not owned by its cube and not a plain reference, so its shift + # is keyed by the owned member it wraps. + - name: return_day + type: time + sql: "DATE_TRUNC('day', {pa_returns.created_at})" + measures: + - name: total_value + type: sum + sql: lifetime_value + + - name: total_value_prev_month + type: number + sql: "{CUBE.total_value}" + multi_stage: true + time_shift: + - interval: "1 month" + type: prior + timeDimension: pa_customers.return_day + + pre_aggregations: + - name: value_by_return_day_month + type: rollup + measures: + - total_value + time_dimension: pa_customers.return_day + granularity: month + + - name: pa_returns + sql: "SELECT * FROM pa_returns" + dimensions: + - name: id + type: number + sql: id + primary_key: true + - name: customer_id + type: number + sql: customer_id + - name: created_at + type: time + sql: created_at + measures: + - name: count + type: count diff --git a/rust/cube/cubesqlplanner/cubesqlplanner/src/test_fixtures/schemas/yaml_files/common/integration_multi_dep_time_dim_shift_pre_agg.yaml b/rust/cube/cubesqlplanner/cubesqlplanner/src/test_fixtures/schemas/yaml_files/common/integration_multi_dep_time_dim_shift_pre_agg.yaml new file mode 100644 index 0000000000000..90b7e5993c439 --- /dev/null +++ b/rust/cube/cubesqlplanner/cubesqlplanner/src/test_fixtures/schemas/yaml_files/common/integration_multi_dep_time_dim_shift_pre_agg.yaml @@ -0,0 +1,109 @@ +cubes: + - name: mdd_events + sql: "SELECT * FROM mdd_events" + dimensions: + - name: id + type: number + sql: id + primary_key: true + - name: happened_at + type: time + sql: happened_at + - name: recorded_at + type: time + sql: recorded_at + - name: batch_at + type: time + sql: batch_at + + # Time dimension built from two owned time dimensions. The shift + # below covers only one of them, so the stored rollup column + # cannot represent the shifted values. + # Plain reference to the shifted dimension. Rendered through to what + # it points at, so a stored copy of it is never offset. + - name: happened_at_ref + type: time + sql: "{CUBE.happened_at}" + + - name: effective_at + type: time + sql: "COALESCE({CUBE.happened_at}, {CUBE.recorded_at})" + segments: + # Built from the partially shifted dimension, so a stored copy of it + # is computed from unshifted values just like a stored dimension. + - name: recent_effective + sql: "{CUBE.effective_at} >= '2024-02-01'" + + measures: + - name: total + type: sum + sql: val + + # Reads the shifted time dimension inside its own SQL, so a stored + # copy is aggregated from unshifted values. + - name: late_total + type: sum + sql: "CASE WHEN {CUBE.happened_at} >= '2024-02-01' THEN val END" + + - name: total_prev_month + type: number + sql: "{CUBE.total}" + multi_stage: true + time_shift: + - interval: "1 month" + type: prior + timeDimension: mdd_events.happened_at + + - name: late_prev_month + type: number + sql: "{CUBE.late_total}" + multi_stage: true + time_shift: + - interval: "1 month" + type: prior + timeDimension: mdd_events.happened_at + + pre_aggregations: + # Stores the partially shifted dimension under `dimensions:` rather + # than as its time dimension, so it is substituted by column just the + # same while the time dimension itself is unshifted. + - name: total_by_batch_month_with_effective + type: rollup + measures: + - total + dimensions: + - effective_at + time_dimension: mdd_events.batch_at + granularity: month + + - name: total_by_batch_month_with_ref + type: rollup + measures: + - total + dimensions: + - happened_at_ref + time_dimension: mdd_events.batch_at + granularity: month + + - name: late_total_by_batch_month + type: rollup + measures: + - late_total + time_dimension: mdd_events.batch_at + granularity: month + + - name: total_by_batch_month_with_segment + type: rollup + measures: + - total + segments: + - recent_effective + time_dimension: mdd_events.batch_at + granularity: month + + - name: total_by_effective_month + type: rollup + measures: + - total + time_dimension: mdd_events.effective_at + granularity: month diff --git a/rust/cube/cubesqlplanner/cubesqlplanner/src/test_fixtures/schemas/yaml_files/common/integration_multi_stage_multiplied_pre_agg.yaml b/rust/cube/cubesqlplanner/cubesqlplanner/src/test_fixtures/schemas/yaml_files/common/integration_multi_stage_multiplied_pre_agg.yaml index b107f014db1fe..a744407de8920 100644 --- a/rust/cube/cubesqlplanner/cubesqlplanner/src/test_fixtures/schemas/yaml_files/common/integration_multi_stage_multiplied_pre_agg.yaml +++ b/rust/cube/cubesqlplanner/cubesqlplanner/src/test_fixtures/schemas/yaml_files/common/integration_multi_stage_multiplied_pre_agg.yaml @@ -142,3 +142,14 @@ cubes: - count time_dimension: created_at granularity: month + +views: + - name: customers_view + cubes: + - join_path: customers + includes: + - total_lifetime_value + - total_lifetime_value_prev_month_by_returns + - join_path: customers.returns + includes: + - created_at diff --git a/rust/cube/cubesqlplanner/cubesqlplanner/src/test_fixtures/schemas/yaml_files/common/pre_aggregation_matching_test.yaml b/rust/cube/cubesqlplanner/cubesqlplanner/src/test_fixtures/schemas/yaml_files/common/pre_aggregation_matching_test.yaml index a5f9ea15659f9..c2c4b569915f4 100644 --- a/rust/cube/cubesqlplanner/cubesqlplanner/src/test_fixtures/schemas/yaml_files/common/pre_aggregation_matching_test.yaml +++ b/rust/cube/cubesqlplanner/cubesqlplanner/src/test_fixtures/schemas/yaml_files/common/pre_aggregation_matching_test.yaml @@ -181,3 +181,17 @@ cubes: - id - status - city + + # Same members as `segment_rollup`, but every reference is written as a + # cube-name interpolation — `(CUBE) => `${CUBE}.count``. + - name: interpolated_refs_rollup + type: rollup + measures: + - '{CUBE}.count' + - '{CUBE}.total_amount' + dimensions: + - '{CUBE}.status' + segments: + - '{CUBE}.high_priority' + time_dimension: '{CUBE}.created_at' + granularity: day diff --git a/rust/cube/cubesqlplanner/cubesqlplanner/src/test_fixtures/schemas/yaml_files/seeds/integration_derived_time_dim_shift_tables.sql b/rust/cube/cubesqlplanner/cubesqlplanner/src/test_fixtures/schemas/yaml_files/seeds/integration_derived_time_dim_shift_tables.sql new file mode 100644 index 0000000000000..c4b7e391e0345 --- /dev/null +++ b/rust/cube/cubesqlplanner/cubesqlplanner/src/test_fixtures/schemas/yaml_files/seeds/integration_derived_time_dim_shift_tables.sql @@ -0,0 +1,27 @@ +DROP TABLE IF EXISTS pa_returns CASCADE; +DROP TABLE IF EXISTS pa_customers CASCADE; + +CREATE TABLE pa_customers ( + id INTEGER PRIMARY KEY, + lifetime_value NUMERIC(10, 2) NOT NULL +); + +INSERT INTO pa_customers (id, lifetime_value) VALUES + (1, 1000.00), + (2, 2000.00), + (3, 500.00); + +CREATE TABLE pa_returns ( + id INTEGER PRIMARY KEY, + customer_id INTEGER NOT NULL REFERENCES pa_customers(id), + created_at TIMESTAMP NOT NULL +); + +-- One customer per month, so each month's total_value is that customer's +-- lifetime_value. December 2023 sits before the queried range and is only +-- reachable when the shifted leaf widens its pre-aggregation date range. +INSERT INTO pa_returns (id, customer_id, created_at) VALUES + (1, 1, '2023-12-12 10:00:00'), + (2, 2, '2024-01-15 10:00:00'), + (3, 3, '2024-02-14 10:00:00'), + (4, 1, '2024-03-05 10:00:00'); diff --git a/rust/cube/cubesqlplanner/cubesqlplanner/src/test_fixtures/schemas/yaml_files/seeds/integration_multi_dep_time_dim_shift_tables.sql b/rust/cube/cubesqlplanner/cubesqlplanner/src/test_fixtures/schemas/yaml_files/seeds/integration_multi_dep_time_dim_shift_tables.sql new file mode 100644 index 0000000000000..79d5608b19aea --- /dev/null +++ b/rust/cube/cubesqlplanner/cubesqlplanner/src/test_fixtures/schemas/yaml_files/seeds/integration_multi_dep_time_dim_shift_tables.sql @@ -0,0 +1,18 @@ +DROP TABLE IF EXISTS mdd_events CASCADE; + +CREATE TABLE mdd_events ( + id INTEGER PRIMARY KEY, + happened_at TIMESTAMP, + recorded_at TIMESTAMP, + batch_at TIMESTAMP NOT NULL, + val NUMERIC(10, 2) NOT NULL +); + +-- Row 3 has no happened_at, so the shift does not move it: it stays in +-- February while every other row shifts by a month. That is what makes the +-- shifted values impossible to reproduce by offsetting the stored column. +INSERT INTO mdd_events (id, happened_at, recorded_at, batch_at, val) VALUES + (1, '2023-12-12 10:00:00', NULL, '2024-01-05 10:00:00', 1000.00), + (2, '2024-01-15 10:00:00', NULL, '2024-02-05 10:00:00', 2000.00), + (3, NULL, '2024-02-14 10:00:00', '2024-02-05 10:00:00', 500.00), + (4, '2024-03-05 10:00:00', NULL, '2024-03-05 10:00:00', 1000.00); diff --git a/rust/cube/cubesqlplanner/cubesqlplanner/src/tests/filter_params_time_shift.rs b/rust/cube/cubesqlplanner/cubesqlplanner/src/tests/filter_params_time_shift.rs new file mode 100644 index 0000000000000..20003884a6d3f --- /dev/null +++ b/rust/cube/cubesqlplanner/cubesqlplanner/src/tests/filter_params_time_shift.rs @@ -0,0 +1,91 @@ +use crate::test_fixtures::cube_bridge::MockSchema; +use crate::test_fixtures::test_utils::TestContext; +use indoc::indoc; + +// A cube whose `sql` pushes a time dimension down through FILTER_PARAMS, where +// that dimension is derived from the time dimension of another cube and a +// multi-stage measure shifts it. +fn schema() -> MockSchema { + MockSchema::from_yaml(indoc! {" + cubes: + - name: fps_returns + sql: \"SELECT * FROM fps_returns WHERE {FILTER_PARAMS_COLUMN:fps_orders.return_day:DATE_TRUNC('day', created_at)}\" + dimensions: + - name: id + type: number + sql: id + primary_key: true + - name: order_id + type: number + sql: order_id + - name: created_at + type: time + sql: created_at + measures: + - name: count + type: count + + - name: fps_orders + sql: \"SELECT * FROM fps_orders\" + joins: + - name: fps_returns + relationship: one_to_many + sql: \"{fps_orders}.id = {fps_returns.order_id}\" + dimensions: + - name: id + type: number + sql: id + primary_key: true + - name: return_day + type: time + sql: \"DATE_TRUNC('day', {fps_returns.created_at})\" + measures: + - name: total + type: sum + sql: amount + + - name: total_prev_month + type: number + sql: \"{CUBE.total}\" + multi_stage: true + time_shift: + - interval: \"1 month\" + type: prior + timeDimension: fps_orders.return_day + "}) + .unwrap() +} + +// The shifted stage reads the dimension offset by the interval, so the +// predicate pushed into the cube's sql has to be offset the same way. Leaving +// it bare would filter the source rows by unshifted bounds while the stage +// groups by shifted values. +#[test] +fn pushed_down_column_is_offset_in_the_shifted_stage() { + let ctx = TestContext::new(schema()).unwrap(); + + let (sql, _) = ctx + .build_sql_and_params(indoc! {" + measures: + - fps_orders.total + - fps_orders.total_prev_month + time_dimensions: + - dimension: fps_orders.return_day + granularity: month + dateRange: + - \"2024-01-01\" + - \"2024-03-31\" + "}) + .unwrap(); + + assert!( + sql.contains("(DATE_TRUNC('day', created_at) + interval '1 month')"), + "the shifted stage must push down the offset column\nsql: {}", + sql + ); + assert!( + sql.contains("(DATE_TRUNC('day', created_at) >="), + "the unshifted stage must push down the bare column\nsql: {}", + sql + ); +} diff --git a/rust/cube/cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/multi_stage.rs b/rust/cube/cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/multi_stage.rs index f7cdf37665385..3214f4ac8bbf5 100644 --- a/rust/cube/cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/multi_stage.rs +++ b/rust/cube/cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/multi_stage.rs @@ -1,6 +1,7 @@ use crate::test_fixtures::cube_bridge::MockSchema; use crate::test_fixtures::test_utils::TestContext; use indoc::indoc; +use itertools::Itertools; const SEED: &str = "integration_multi_stage_tables.sql"; const YAML: &str = "common/integration_multi_stage_multiplied_pre_agg.yaml"; @@ -73,6 +74,77 @@ async fn test_multi_stage_time_shift_pre_agg_with_leaf_measure() { } } +#[tokio::test(flavor = "multi_thread")] +async fn test_multi_stage_time_shift_pre_agg_with_leaf_measure_via_view() { + // Same shape as test_multi_stage_time_shift_pre_agg_with_leaf_measure, but + // through a view: the time dimension filter is view-qualified while the + // accumulated time shifts are keyed by the underlying cube member. The + // shifted leaf must still get its own widened date_range. + // + // The widened range only decides which rollup partitions are loaded, so + // the assertions on the usages are what guard it; the executed rows are a + // parity check against the cube-level query. + let schema = MockSchema::from_yaml_file(YAML) + .only_pre_aggregations(&["customers_lifetime_by_returns_month"]); + let ctx = TestContext::new(schema).unwrap(); + + let query = indoc! {r#" + measures: + - customers_view.total_lifetime_value + - customers_view.total_lifetime_value_prev_month_by_returns + time_dimensions: + - dimension: customers_view.created_at + granularity: month + dateRange: + - "2024-01-01" + - "2024-03-31" + order: + - id: customers_view.created_at + "#}; + + let (_sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(query).unwrap(); + let names: Vec<&str> = pre_aggrs.iter().map(|u| u.name().as_str()).collect(); + + assert_eq!( + pre_aggrs.len(), + 2, + "Expected 2 usages (shifted + unshifted leaf); got {:?}", + names + ); + assert!( + names + .iter() + .all(|n| *n == "customers_lifetime_by_returns_month"), + "Both usages must be customers_lifetime_by_returns_month; got {:?}", + names + ); + + let shifted_range = Some(( + "2023-12-01T00:00:00.000".to_string(), + "2024-02-29T23:59:59.999".to_string(), + )); + let original_range = Some(( + "2024-01-01T00:00:00.000".to_string(), + "2024-03-31T23:59:59.999".to_string(), + )); + let shifted = pre_aggrs + .iter() + .find(|u| u.date_range == shifted_range) + .expect("Expected a usage with shifted date_range"); + let unshifted = pre_aggrs + .iter() + .find(|u| u.date_range == original_range) + .expect("Expected a usage with original date_range"); + assert_ne!( + shifted.index, unshifted.index, + "Shifted and unshifted usages must have different usage indexes" + ); + + if let Some(result) = ctx.try_execute(query, SEED).await { + insta::assert_snapshot!(result); + } +} + #[tokio::test(flavor = "multi_thread")] async fn test_multi_stage_time_shift_pre_agg_with_multi_stage_measure() { let schema = MockSchema::from_yaml_file(YAML) @@ -244,3 +316,315 @@ async fn test_multi_stage_pre_agg_covering_multiplying_filter() { insta::assert_snapshot!(result); } } + +const DERIVED_YAML: &str = "common/integration_derived_time_dim_shift_pre_agg.yaml"; +const DERIVED_SEED: &str = "integration_derived_time_dim_shift_tables.sql"; + +#[tokio::test(flavor = "multi_thread")] +async fn test_multi_stage_time_shift_pre_agg_on_derived_time_dimension() { + // The shift is declared on a time dimension derived from another cube's + // time dimension, and the pre-aggregation materializes that derived + // dimension. The shift must reach both the widened date_range and the + // rollup column the shifted leaf reads. + let schema = MockSchema::from_yaml_file(DERIVED_YAML) + .only_pre_aggregations(&["value_by_return_day_month"]); + let ctx = TestContext::new(schema).unwrap(); + + let query = indoc! {r#" + measures: + - pa_customers.total_value + - pa_customers.total_value_prev_month + time_dimensions: + - dimension: pa_customers.return_day + granularity: month + dateRange: + - "2024-01-01" + - "2024-03-31" + order: + - id: pa_customers.return_day + "#}; + + let (sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(query).unwrap(); + let names: Vec<&str> = pre_aggrs.iter().map(|u| u.name().as_str()).collect(); + + assert_eq!( + pre_aggrs.len(), + 2, + "Expected 2 usages (shifted + unshifted leaf); got {:?}", + names + ); + + let shifted_range = Some(( + "2023-12-01T00:00:00.000".to_string(), + "2024-02-29T23:59:59.999".to_string(), + )); + let original_range = Some(( + "2024-01-01T00:00:00.000".to_string(), + "2024-03-31T23:59:59.999".to_string(), + )); + assert!( + pre_aggrs.iter().any(|u| u.date_range == shifted_range), + "Expected a usage with shifted date_range; got {:?}", + pre_aggrs.iter().map(|u| u.date_range.clone()).collect_vec() + ); + assert!( + pre_aggrs.iter().any(|u| u.date_range == original_range), + "Expected a usage with original date_range; got {:?}", + pre_aggrs.iter().map(|u| u.date_range.clone()).collect_vec() + ); + + // The shifted leaf reads the rollup column, so the interval has to be + // applied to that column — the derived member's own SQL is never + // expanded here and cannot carry the shift. Every occurrence must sit + // directly on the column: one landing on an already shifted expression + // would mean the interval was added twice. + let shifts = sql.matches("interval '1 month'").count(); + let shifts_on_column = sql + .matches(r#""pa_customers__return_day_month" + interval '1 month'"#) + .count(); + assert!( + shifts > 0, + "Shifted leaf must offset the pre-aggregation column:\n{}", + sql + ); + assert_eq!( + shifts_on_column, shifts, + "Every shift must be applied to the rollup column itself:\n{}", + sql + ); + + if let Some(result) = ctx.try_execute(query, DERIVED_SEED).await { + insta::assert_snapshot!(result); + } +} + +const MULTI_DEP_YAML: &str = "common/integration_multi_dep_time_dim_shift_pre_agg.yaml"; +const MULTI_DEP_SEED: &str = "integration_multi_dep_time_dim_shift_tables.sql"; + +// A rollup stores one column for a time dimension built from two owned time +// dimensions, while the shift covers only one of them. Offsetting the stored +// column would move every row, including those the shift must leave in place, +// so no offset of that column can reproduce the shifted values and the rollup +// cannot serve the shifted leaf at all. +#[tokio::test(flavor = "multi_thread")] +async fn test_multi_stage_time_shift_pre_agg_on_partially_shifted_time_dimension() { + let schema = MockSchema::from_yaml_file(MULTI_DEP_YAML) + .only_pre_aggregations(&["total_by_effective_month"]); + let ctx = TestContext::new(schema).unwrap(); + + let query = indoc! {r#" + measures: + - mdd_events.total + - mdd_events.total_prev_month + time_dimensions: + - dimension: mdd_events.effective_at + granularity: month + dateRange: + - "2024-01-01" + - "2024-03-31" + order: + - id: mdd_events.effective_at + "#}; + + let (_sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(query).unwrap(); + assert!( + pre_aggrs.is_empty(), + "Rollup cannot represent the shifted leaf, so the query must fall back to base SQL; got {:?}", + pre_aggrs.iter().map(|u| u.name().clone()).collect_vec() + ); + + // Snapshot holds the values the same query produces from base SQL. + if let Some(result) = ctx.try_execute(query, MULTI_DEP_SEED).await { + insta::assert_snapshot!(result); + } +} + +// Same partially shifted dimension, but stored under the rollup's `dimensions:` +// while its time dimension carries no shift. A stored dimension is substituted +// by column exactly like a stored time dimension, so reading it unshifted is +// just as wrong and the rollup must be rejected here too. +#[tokio::test(flavor = "multi_thread")] +async fn test_multi_stage_time_shift_pre_agg_on_partially_shifted_stored_dimension() { + let schema = MockSchema::from_yaml_file(MULTI_DEP_YAML) + .only_pre_aggregations(&["total_by_batch_month_with_effective"]); + let ctx = TestContext::new(schema).unwrap(); + + let query = indoc! {r#" + measures: + - mdd_events.total + - mdd_events.total_prev_month + dimensions: + - mdd_events.effective_at + time_dimensions: + - dimension: mdd_events.batch_at + granularity: month + dateRange: + - "2024-01-01" + - "2024-03-31" + order: + - id: mdd_events.batch_at + - id: mdd_events.effective_at + "#}; + + let (_sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(query).unwrap(); + assert!( + pre_aggrs.is_empty(), + "Rollup stores the partially shifted dimension as a column, so it cannot serve the shifted leaf; got {:?}", + pre_aggrs.iter().map(|u| u.name().clone()).collect_vec() + ); + + // Base SQL keeps the row with no happened_at where it is while shifting + // the rest, which is what no offset of a stored column can reproduce. + if let Some(result) = ctx.try_execute(query, MULTI_DEP_SEED).await { + insta::assert_snapshot!(result); + } +} + +// The partially shifted dimension reaches the rollup through a segment this +// time. A stored segment is substituted by column as well, so the same +// rejection applies. +#[tokio::test(flavor = "multi_thread")] +async fn test_multi_stage_time_shift_pre_agg_on_partially_shifted_stored_segment() { + let schema = MockSchema::from_yaml_file(MULTI_DEP_YAML) + .only_pre_aggregations(&["total_by_batch_month_with_segment"]); + let ctx = TestContext::new(schema).unwrap(); + + let query = indoc! {r#" + measures: + - mdd_events.total + - mdd_events.total_prev_month + segments: + - mdd_events.recent_effective + time_dimensions: + - dimension: mdd_events.batch_at + granularity: month + dateRange: + - "2024-01-01" + - "2024-03-31" + order: + - id: mdd_events.batch_at + "#}; + + let (_sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(query).unwrap(); + assert!( + pre_aggrs.is_empty(), + "Rollup stores a segment built from the partially shifted dimension; got {:?}", + pre_aggrs.iter().map(|u| u.name().clone()).collect_vec() + ); + + if let Some(result) = ctx.try_execute(query, MULTI_DEP_SEED).await { + insta::assert_snapshot!(result); + } +} + +// The shifted time dimension is read inside a measure's own SQL this time. +// The rollup stores that measure aggregated from unshifted values, which no +// offset recovers, so it cannot serve the shifted leaf either. +#[tokio::test(flavor = "multi_thread")] +async fn test_multi_stage_time_shift_pre_agg_on_measure_reading_shifted_dimension() { + let schema = MockSchema::from_yaml_file(MULTI_DEP_YAML) + .only_pre_aggregations(&["late_total_by_batch_month"]); + let ctx = TestContext::new(schema).unwrap(); + + let query = indoc! {r#" + measures: + - mdd_events.late_total + - mdd_events.late_prev_month + time_dimensions: + - dimension: mdd_events.batch_at + granularity: month + dateRange: + - "2024-01-01" + - "2024-03-31" + order: + - id: mdd_events.batch_at + "#}; + + let (_sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(query).unwrap(); + assert!( + pre_aggrs.is_empty(), + "Rollup stores a measure computed from the shifted dimension; got {:?}", + pre_aggrs.iter().map(|u| u.name().clone()).collect_vec() + ); + + if let Some(result) = ctx.try_execute(query, MULTI_DEP_SEED).await { + insta::assert_snapshot!(result); + } +} + +// The stored dimension is a plain reference to the shifted one. Its column +// holds unshifted values and the renderer resolves the reference through to +// what it points at rather than offsetting the column, so the rollup cannot +// serve the shifted leaf. +#[tokio::test(flavor = "multi_thread")] +async fn test_multi_stage_time_shift_pre_agg_on_stored_reference_dimension() { + let schema = MockSchema::from_yaml_file(MULTI_DEP_YAML) + .only_pre_aggregations(&["total_by_batch_month_with_ref"]); + let ctx = TestContext::new(schema).unwrap(); + + let query = indoc! {r#" + measures: + - mdd_events.total + - mdd_events.total_prev_month + dimensions: + - mdd_events.happened_at_ref + time_dimensions: + - dimension: mdd_events.batch_at + granularity: month + dateRange: + - "2024-01-01" + - "2024-03-31" + order: + - id: mdd_events.batch_at + - id: mdd_events.happened_at_ref + "#}; + + let (_sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(query).unwrap(); + assert!( + pre_aggrs.is_empty(), + "Rollup stores a reference to the shifted dimension; got {:?}", + pre_aggrs.iter().map(|u| u.name().clone()).collect_vec() + ); + + if let Some(result) = ctx.try_execute(query, MULTI_DEP_SEED).await { + insta::assert_snapshot!(result); + } +} + +// The rollup stores a member the shift reaches, but the query never selects +// it, so nothing ever reads that column and the stored measures are still +// usable. Rejecting on a member the query does not read costs a rollup for +// no correctness gain. +#[tokio::test(flavor = "multi_thread")] +async fn test_multi_stage_time_shift_pre_agg_keeps_rollup_when_unshiftable_member_unused() { + let schema = MockSchema::from_yaml_file(MULTI_DEP_YAML) + .only_pre_aggregations(&["total_by_batch_month_with_ref"]); + let ctx = TestContext::new(schema).unwrap(); + + let query = indoc! {r#" + measures: + - mdd_events.total + - mdd_events.total_prev_month + time_dimensions: + - dimension: mdd_events.batch_at + granularity: month + dateRange: + - "2024-01-01" + - "2024-03-31" + order: + - id: mdd_events.batch_at + "#}; + + let (_sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(query).unwrap(); + assert_eq!( + pre_aggrs.len(), + 2, + "Rollup stores happened_at_ref but the query does not select it, so it stays usable" + ); + + // Same values the query produces from base SQL, so keeping the rollup + // costs nothing in correctness. + if let Some(result) = ctx.try_execute(query, MULTI_DEP_SEED).await { + insta::assert_snapshot!(result); + } +} diff --git a/rust/cube/cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/snapshots/cubesqlplanner__tests__integration__pre_aggregations__multi_stage__multi_stage_time_shift_pre_agg_keeps_rollup_when_unshiftable_member_unused.snap b/rust/cube/cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/snapshots/cubesqlplanner__tests__integration__pre_aggregations__multi_stage__multi_stage_time_shift_pre_agg_keeps_rollup_when_unshiftable_member_unused.snap new file mode 100644 index 0000000000000..5b8954d1c3ea9 --- /dev/null +++ b/rust/cube/cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/snapshots/cubesqlplanner__tests__integration__pre_aggregations__multi_stage__multi_stage_time_shift_pre_agg_keeps_rollup_when_unshiftable_member_unused.snap @@ -0,0 +1,9 @@ +--- +source: cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/multi_stage.rs +expression: result +--- +mdd_events__batch_at_month | mdd_events__total | mdd_events__total_prev_month +---------------------------+-------------------+----------------------------- +2024-01-01T00:00:00.000Z | 1000 | 1000 +2024-02-01T00:00:00.000Z | 2500 | 2500 +2024-03-01T00:00:00.000Z | 1000 | 1000 diff --git a/rust/cube/cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/snapshots/cubesqlplanner__tests__integration__pre_aggregations__multi_stage__multi_stage_time_shift_pre_agg_on_derived_time_dimension.snap b/rust/cube/cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/snapshots/cubesqlplanner__tests__integration__pre_aggregations__multi_stage__multi_stage_time_shift_pre_agg_on_derived_time_dimension.snap new file mode 100644 index 0000000000000..7853eddcaff49 --- /dev/null +++ b/rust/cube/cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/snapshots/cubesqlplanner__tests__integration__pre_aggregations__multi_stage__multi_stage_time_shift_pre_agg_on_derived_time_dimension.snap @@ -0,0 +1,9 @@ +--- +source: cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/multi_stage.rs +expression: result +--- +pa_customers__return_day_month | pa_customers__total_value | pa_customers__total_value_prev_month +-------------------------------+---------------------------+------------------------------------- +2024-01-01T00:00:00.000Z | 2000 | 1000 +2024-02-01T00:00:00.000Z | 500 | 2000 +2024-03-01T00:00:00.000Z | 1000 | 500 diff --git a/rust/cube/cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/snapshots/cubesqlplanner__tests__integration__pre_aggregations__multi_stage__multi_stage_time_shift_pre_agg_on_measure_reading_shifted_dimension.snap b/rust/cube/cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/snapshots/cubesqlplanner__tests__integration__pre_aggregations__multi_stage__multi_stage_time_shift_pre_agg_on_measure_reading_shifted_dimension.snap new file mode 100644 index 0000000000000..bfafbb868d1ac --- /dev/null +++ b/rust/cube/cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/snapshots/cubesqlplanner__tests__integration__pre_aggregations__multi_stage__multi_stage_time_shift_pre_agg_on_measure_reading_shifted_dimension.snap @@ -0,0 +1,9 @@ +--- +source: cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/multi_stage.rs +expression: result +--- +mdd_events__batch_at_month | mdd_events__late_total | mdd_events__late_prev_month +---------------------------+------------------------+---------------------------- +2024-01-01 00:00:00 | NULL | NULL +2024-02-01 00:00:00 | NULL | 2000.00 +2024-03-01 00:00:00 | 1000.00 | 1000.00 diff --git a/rust/cube/cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/snapshots/cubesqlplanner__tests__integration__pre_aggregations__multi_stage__multi_stage_time_shift_pre_agg_on_partially_shifted_stored_dimension.snap b/rust/cube/cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/snapshots/cubesqlplanner__tests__integration__pre_aggregations__multi_stage__multi_stage_time_shift_pre_agg_on_partially_shifted_stored_dimension.snap new file mode 100644 index 0000000000000..a4290ef2d3163 --- /dev/null +++ b/rust/cube/cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/snapshots/cubesqlplanner__tests__integration__pre_aggregations__multi_stage__multi_stage_time_shift_pre_agg_on_partially_shifted_stored_dimension.snap @@ -0,0 +1,13 @@ +--- +source: cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/multi_stage.rs +expression: result +--- +mdd_events__effective_at | mdd_events__batch_at_month | mdd_events__total | mdd_events__total_prev_month +-------------------------+----------------------------+-------------------+----------------------------- +2023-12-12 10:00:00 | 2024-01-01 00:00:00 | 1000.00 | NULL +2024-01-12 10:00:00 | 2024-01-01 00:00:00 | NULL | 1000.00 +2024-01-15 10:00:00 | 2024-02-01 00:00:00 | 2000.00 | NULL +2024-02-14 10:00:00 | 2024-02-01 00:00:00 | 500.00 | 500.00 +2024-02-15 10:00:00 | 2024-02-01 00:00:00 | NULL | 2000.00 +2024-03-05 10:00:00 | 2024-03-01 00:00:00 | 1000.00 | NULL +2024-04-05 10:00:00 | 2024-03-01 00:00:00 | NULL | 1000.00 diff --git a/rust/cube/cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/snapshots/cubesqlplanner__tests__integration__pre_aggregations__multi_stage__multi_stage_time_shift_pre_agg_on_partially_shifted_stored_segment.snap b/rust/cube/cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/snapshots/cubesqlplanner__tests__integration__pre_aggregations__multi_stage__multi_stage_time_shift_pre_agg_on_partially_shifted_stored_segment.snap new file mode 100644 index 0000000000000..5155164223869 --- /dev/null +++ b/rust/cube/cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/snapshots/cubesqlplanner__tests__integration__pre_aggregations__multi_stage__multi_stage_time_shift_pre_agg_on_partially_shifted_stored_segment.snap @@ -0,0 +1,8 @@ +--- +source: cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/multi_stage.rs +expression: result +--- +mdd_events__batch_at_month | mdd_events__total | mdd_events__total_prev_month +---------------------------+-------------------+----------------------------- +2024-02-01 00:00:00 | 500.00 | 2500.00 +2024-03-01 00:00:00 | 1000.00 | 1000.00 diff --git a/rust/cube/cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/snapshots/cubesqlplanner__tests__integration__pre_aggregations__multi_stage__multi_stage_time_shift_pre_agg_on_partially_shifted_time_dimension.snap b/rust/cube/cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/snapshots/cubesqlplanner__tests__integration__pre_aggregations__multi_stage__multi_stage_time_shift_pre_agg_on_partially_shifted_time_dimension.snap new file mode 100644 index 0000000000000..e213e988f214c --- /dev/null +++ b/rust/cube/cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/snapshots/cubesqlplanner__tests__integration__pre_aggregations__multi_stage__multi_stage_time_shift_pre_agg_on_partially_shifted_time_dimension.snap @@ -0,0 +1,9 @@ +--- +source: cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/multi_stage.rs +expression: result +--- +mdd_events__effective_at_month | mdd_events__total | mdd_events__total_prev_month +-------------------------------+-------------------+----------------------------- +2024-01-01 00:00:00 | 2000.00 | 1000.00 +2024-02-01 00:00:00 | 500.00 | 2500.00 +2024-03-01 00:00:00 | 1000.00 | NULL diff --git a/rust/cube/cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/snapshots/cubesqlplanner__tests__integration__pre_aggregations__multi_stage__multi_stage_time_shift_pre_agg_on_stored_reference_dimension.snap b/rust/cube/cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/snapshots/cubesqlplanner__tests__integration__pre_aggregations__multi_stage__multi_stage_time_shift_pre_agg_on_stored_reference_dimension.snap new file mode 100644 index 0000000000000..2c44a73637c57 --- /dev/null +++ b/rust/cube/cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/snapshots/cubesqlplanner__tests__integration__pre_aggregations__multi_stage__multi_stage_time_shift_pre_agg_on_stored_reference_dimension.snap @@ -0,0 +1,13 @@ +--- +source: cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/multi_stage.rs +expression: result +--- +mdd_events__happened_at_ref | mdd_events__batch_at_month | mdd_events__total | mdd_events__total_prev_month +----------------------------+----------------------------+-------------------+----------------------------- +2023-12-12 10:00:00 | 2024-01-01 00:00:00 | 1000.00 | NULL +2024-01-12 10:00:00 | 2024-01-01 00:00:00 | NULL | 1000.00 +2024-01-15 10:00:00 | 2024-02-01 00:00:00 | 2000.00 | NULL +2024-02-15 10:00:00 | 2024-02-01 00:00:00 | NULL | 2000.00 +NULL | 2024-02-01 00:00:00 | 500.00 | 500.00 +2024-03-05 10:00:00 | 2024-03-01 00:00:00 | 1000.00 | NULL +2024-04-05 10:00:00 | 2024-03-01 00:00:00 | NULL | 1000.00 diff --git a/rust/cube/cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/snapshots/cubesqlplanner__tests__integration__pre_aggregations__multi_stage__multi_stage_time_shift_pre_agg_with_leaf_measure_via_view.snap b/rust/cube/cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/snapshots/cubesqlplanner__tests__integration__pre_aggregations__multi_stage__multi_stage_time_shift_pre_agg_with_leaf_measure_via_view.snap new file mode 100644 index 0000000000000..4a75563dd53d6 --- /dev/null +++ b/rust/cube/cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/snapshots/cubesqlplanner__tests__integration__pre_aggregations__multi_stage__multi_stage_time_shift_pre_agg_with_leaf_measure_via_view.snap @@ -0,0 +1,9 @@ +--- +source: cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/multi_stage.rs +expression: result +--- +customers_view__created_at_month | customers_view__total_lifetime_value | customers_view__total_lifetime_value_prev_month_by_returns +---------------------------------+--------------------------------------+----------------------------------------------------------- +2024-01-01T00:00:00.000Z | 1000 | NULL +2024-02-01T00:00:00.000Z | 2000 | 1000 +2024-03-01T00:00:00.000Z | 500 | 2000 diff --git a/rust/cube/cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/snapshots/cubesqlplanner__tests__integration__pre_aggregations__sql_generation__interpolated_refs_full_match_result.snap b/rust/cube/cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/snapshots/cubesqlplanner__tests__integration__pre_aggregations__sql_generation__interpolated_refs_full_match_result.snap new file mode 100644 index 0000000000000..b0e7963d1d1d1 --- /dev/null +++ b/rust/cube/cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/snapshots/cubesqlplanner__tests__integration__pre_aggregations__sql_generation__interpolated_refs_full_match_result.snap @@ -0,0 +1,11 @@ +--- +source: cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/sql_generation.rs +expression: result +--- +orders__status | orders__created_at_day | orders__count | orders__total_amount +---------------+------------------------+---------------+--------------------- +cancelled | 2025-02-15 00:00:00 | 1 | 25.00 +completed | 2025-01-10 00:00:00 | 1 | 100.00 +completed | 2025-01-31 00:00:00 | 1 | 300.00 +pending | 2025-01-10 00:00:00 | 1 | 200.00 +pending | 2025-03-01 00:00:00 | 1 | 175.00 diff --git a/rust/cube/cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/snapshots/cubesqlplanner__tests__integration__pre_aggregations__sql_generation__interpolated_refs_with_coarser_granularity_result.snap b/rust/cube/cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/snapshots/cubesqlplanner__tests__integration__pre_aggregations__sql_generation__interpolated_refs_with_coarser_granularity_result.snap new file mode 100644 index 0000000000000..a24338aecc9f1 --- /dev/null +++ b/rust/cube/cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/snapshots/cubesqlplanner__tests__integration__pre_aggregations__sql_generation__interpolated_refs_with_coarser_granularity_result.snap @@ -0,0 +1,10 @@ +--- +source: cubesqlplanner/cubesqlplanner/src/tests/integration/pre_aggregations/sql_generation.rs +expression: result +--- +orders__status | orders__created_at_month | orders__count +---------------+--------------------------+-------------- +cancelled | 2025-02-01 00:00:00 | 1 +completed | 2025-01-01 00:00:00 | 2 +pending | 2025-01-01 00:00:00 | 1 +pending | 2025-03-01 00:00:00 | 1 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 8994412eda9a1..54faf53106fa5 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 @@ -1692,3 +1692,101 @@ async fn test_ungrouped_cross_cube_view_query_matches_rollup_covering_both_prima Ok(()) } + +// --- References written as a cube-name interpolation --- +// +// `interpolated_refs_rollup` names the same members as `segment_rollup`, but +// through `` `${CUBE}.count` ``-style references. The two must be picked for the +// same queries and read back the same rows. The rows themselves are pinned by +// the snapshot, which is only compared when Postgres execution is enabled. + +async fn interpolated_and_symbol_refs_agree( + query_yaml: &str, + snapshot_name: &str, +) -> Result<(), CubeError> { + let interpolated_ctx = TestContext::new( + MockSchema::from_yaml_file("common/pre_aggregation_matching_test.yaml") + .only_pre_aggregations(&["interpolated_refs_rollup"]), + )?; + let symbol_ctx = TestContext::new( + MockSchema::from_yaml_file("common/pre_aggregation_matching_test.yaml") + .only_pre_aggregations(&["segment_rollup"]), + )?; + + let (_sql, interpolated_pre_aggrs) = + interpolated_ctx.build_sql_with_used_pre_aggregations(query_yaml)?; + assert_eq!(interpolated_pre_aggrs.len(), 1); + assert_eq!(interpolated_pre_aggrs[0].name(), "interpolated_refs_rollup"); + + let (_sql, symbol_pre_aggrs) = symbol_ctx.build_sql_with_used_pre_aggregations(query_yaml)?; + assert_eq!(symbol_pre_aggrs.len(), 1); + assert_eq!(symbol_pre_aggrs[0].name(), "segment_rollup"); + + let interpolated_result = interpolated_ctx + .try_execute_pg(query_yaml, "pre_aggregation_matching_tables.sql") + .await; + let symbol_result = symbol_ctx + .try_execute_pg(query_yaml, "pre_aggregation_matching_tables.sql") + .await; + + assert_eq!(interpolated_result, symbol_result); + + // Without Postgres execution there are no rows to compare, so make it + // explicit that the row assertions below do run when it is enabled. + #[cfg(feature = "integration-postgres")] + assert!( + interpolated_result.is_some(), + "Postgres execution is enabled but returned no result" + ); + + if let Some(result) = interpolated_result { + insta::assert_snapshot!(snapshot_name, result); + } + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread")] +async fn test_interpolated_refs_full_match() -> Result<(), CubeError> { + interpolated_and_symbol_refs_agree( + indoc! {" + measures: + - orders.count + - orders.total_amount + dimensions: + - orders.status + segments: + - orders.high_priority + time_dimensions: + - dimension: orders.created_at + granularity: day + order: + - id: orders.status + - id: orders.created_at + "}, + "interpolated_refs_full_match_result", + ) + .await +} + +#[tokio::test(flavor = "multi_thread")] +async fn test_interpolated_refs_with_coarser_granularity() -> Result<(), CubeError> { + interpolated_and_symbol_refs_agree( + indoc! {" + measures: + - orders.count + dimensions: + - orders.status + segments: + - orders.high_priority + time_dimensions: + - dimension: orders.created_at + granularity: month + order: + - id: orders.status + - id: orders.created_at + "}, + "interpolated_refs_with_coarser_granularity_result", + ) + .await +} diff --git a/rust/cube/cubesqlplanner/cubesqlplanner/src/tests/mod.rs b/rust/cube/cubesqlplanner/cubesqlplanner/src/tests/mod.rs index 2788c6a90d863..5e5ef25d6ad5d 100644 --- a/rust/cube/cubesqlplanner/cubesqlplanner/src/tests/mod.rs +++ b/rust/cube/cubesqlplanner/cubesqlplanner/src/tests/mod.rs @@ -8,6 +8,7 @@ mod dimension_symbol; mod filter; mod filter_params_callback_column; mod filter_params_segment; +mod filter_params_time_shift; mod join_hints_collector; mod measure_symbol; mod member_expressions_on_views; diff --git a/rust/cubestore/cubestore/src/cluster/mod.rs b/rust/cubestore/cubestore/src/cluster/mod.rs index 446f0a34c2445..52e8e92c09fab 100644 --- a/rust/cubestore/cubestore/src/cluster/mod.rs +++ b/rust/cubestore/cubestore/src/cluster/mod.rs @@ -26,7 +26,7 @@ use crate::cluster::transport::{ClusterTransport, MetaStoreTransport, WorkerConn use crate::config::injection::{DIService, Injector}; use crate::config::is_router; #[allow(unused_imports)] -use crate::config::{Config, ConfigObj, RepartitionStrategy}; +use crate::config::{Config, ConfigObj, RepartitionStrategy, TopKAggregateStrategy}; use crate::metastore::chunks::chunk_file_name; use crate::metastore::job::{Job, JobRunnerPool, JobStatus, JobType}; use crate::metastore::{ @@ -231,12 +231,36 @@ pub struct ClusterImpl { crate::di_service!(ClusterImpl, [Cluster]); +/// Planning decisions the sending node takes from its own configuration and that the receiving node +/// must reproduce. Each of these shapes both the worker subtree and the node that combines it, so a +/// receiver planning its half from a different value returns wrong rows instead of failing. They +/// travel with the query so that the sender's configuration decides for the whole plan. +#[derive(Copy, Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct PlanningFlags { + pub group_by_limit_factor: usize, + pub topk_strategy: TopKAggregateStrategy, +} + +impl PlanningFlags { + pub fn from_config(config: &dyn ConfigObj) -> PlanningFlags { + PlanningFlags { + group_by_limit_factor: config.group_by_limit_factor(), + topk_strategy: config.topk_aggregate_strategy(), + } + } +} + /// Parameters that the worker node uses to plan queries. Generally, it needs to construct the same /// query plans as the router node (or if there are multiple levels of cluster send, the node from /// which it received the query). We include the necessary information here. #[derive(Copy, Clone, Debug, Serialize, Deserialize)] pub struct WorkerPlanningParams { pub worker_partition_count: usize, + /// Absent from a sender that predates the flags. Such a sender planned its half from its own + /// configuration, so the receiver falls back to its own -- the same value, as long as the two + /// binaries still agree on the defaults. + #[serde(default)] + pub flags: Option, } impl WorkerPlanningParams { @@ -244,6 +268,7 @@ impl WorkerPlanningParams { pub fn no_worker() -> WorkerPlanningParams { WorkerPlanningParams { worker_partition_count: 1, + flags: None, } } } @@ -2316,6 +2341,106 @@ mod tests { use crate::config::Config; use std::fs; + /// The wire shape a node that predates the planning flags writes. + #[derive(Serialize, Deserialize)] + struct WorkerPlanningParamsWithoutFlags { + worker_partition_count: usize, + } + + fn to_flexbuffers(value: &T) -> Vec { + let mut ser = flexbuffers::FlexbufferSerializer::new(); + value.serialize(&mut ser).unwrap(); + ser.take_buffer() + } + + fn from_flexbuffers<'de, T: Deserialize<'de>>(buffer: &'de [u8]) -> T { + T::deserialize(flexbuffers::Reader::get_root(buffer).unwrap()).unwrap() + } + + #[test] + fn planning_flags_absent_from_a_sender_that_omits_them() { + // A message from a node that predates the flags must still be read, and must leave the + // flags absent so that the receiver falls back to its own configuration -- which is what + // that sender planned its half from. + let buffer = to_flexbuffers(&WorkerPlanningParamsWithoutFlags { + worker_partition_count: 3, + }); + let params: WorkerPlanningParams = from_flexbuffers(&buffer); + assert_eq!(params.worker_partition_count, 3); + assert!(params.flags.is_none()); + } + + #[test] + fn planning_flags_are_ignored_by_a_receiver_that_does_not_know_them() { + // The reverse direction: a peer that predates the flags must still read the rest of the + // message rather than fail on the field it does not know. + let buffer = to_flexbuffers(&WorkerPlanningParams { + worker_partition_count: 3, + flags: Some(PlanningFlags { + group_by_limit_factor: 2, + topk_strategy: TopKAggregateStrategy::FullMerge, + }), + }); + let params: WorkerPlanningParamsWithoutFlags = from_flexbuffers(&buffer); + assert_eq!(params.worker_partition_count, 3); + } + + #[test] + fn planning_flags_round_trip() { + let flags = PlanningFlags { + group_by_limit_factor: 5, + topk_strategy: TopKAggregateStrategy::VectorizedStreaming, + }; + let params: WorkerPlanningParams = + from_flexbuffers(&to_flexbuffers(&WorkerPlanningParams { + worker_partition_count: 7, + flags: Some(flags), + })); + assert_eq!(params.worker_partition_count, 7); + assert_eq!(params.flags, Some(flags)); + } + + #[test] + fn planning_flags_strategy_wire_names() { + // The names are the cross-node contract, so a renamed variant must not change them. + #[derive(Serialize)] + struct FlagsWithStrategyAsString { + group_by_limit_factor: usize, + topk_strategy: &'static str, + } + + for (name, strategy) in [ + ("streaming", TopKAggregateStrategy::Streaming), + ( + "vectorized_streaming", + TopKAggregateStrategy::VectorizedStreaming, + ), + ("full_merge", TopKAggregateStrategy::FullMerge), + ] { + let buffer = to_flexbuffers(&FlagsWithStrategyAsString { + group_by_limit_factor: 1, + topk_strategy: name, + }); + let flags: PlanningFlags = from_flexbuffers(&buffer); + assert_eq!(flags.topk_strategy, strategy, "wire name {}", name); + } + } + + /// The values a receiver falls back to when the sender sends no flags. + #[test] + fn planning_flags_from_config() { + let config = Config::test("planning_flags_from_config") + .update_config(|mut c| { + c.group_by_limit_factor = 4; + c.topk_aggregate_strategy = TopKAggregateStrategy::FullMerge; + c + }) + .config_obj(); + let flags = PlanningFlags::from_config(config.as_ref()); + assert_eq!(flags.group_by_limit_factor, 4); + assert_eq!(flags.topk_strategy, TopKAggregateStrategy::FullMerge); + } + fn config_with_workers(name: &str, workers: Vec) -> Arc { Config::test(name) .update_config(|mut c| { diff --git a/rust/cubestore/cubestore/src/config/mod.rs b/rust/cubestore/cubestore/src/config/mod.rs index c33f66b2666b2..ffdb34b160ffa 100644 --- a/rust/cubestore/cubestore/src/config/mod.rs +++ b/rust/cubestore/cubestore/src/config/mod.rs @@ -55,6 +55,7 @@ use futures::future::join_all; use log::Level; use log::{debug, error}; use mockall::automock; +use serde::{Deserialize, Serialize}; use simple_logger::SimpleLogger; use std::fmt::Display; use std::future::Future; @@ -585,20 +586,27 @@ pub trait ConfigObj: DIService { fn repartition_check_overlapping_children(&self) -> bool; /// Factor `f` controlling when the worker-side partial hash aggregate trims its output to the /// top-k groups. Trimming happens only when the number of local groups exceeds `f * k`, where - /// `k = limit + offset`. `0` disables the optimization. + /// `k = limit + offset`. `0` disables the optimization. Whether the trim runs decides the shape + /// of both halves of a split plan, so it rides in [`PlanningFlags`]; a worker uses its own value + /// only when the sender sent no flags. fn group_by_limit_factor(&self) -> usize; /// When the worker group-by-limit hash trim is active, controls where the worker's hash table /// lives: `false` (default) coalesces the partial aggregate's input to one partition (one hash /// table per worker, "over merge"); `true` keeps the raw multi-partition input so it runs per /// partition ("under merge"). + /// + /// Unlike [`ConfigObj::group_by_limit_factor`], this one stays node-local and is not part of + /// [`PlanningFlags`]: it only moves a partition coalesce below the partial aggregate, leaving + /// the subtree's schema and the partition count the router sees the same either way. fn group_by_limit_per_partition(&self) -> bool; /// Replace the sort-preserving merge feeding a grouped Linear (hash) aggregate with a plain /// partition coalesce (the hash aggregate ignores input order, so the per-row merge is wasted). fn coalesce_under_hash_aggregate(&self) -> bool; - /// Router-side merge strategy for distributed value-ordered top-k. + /// Router-side merge strategy for distributed value-ordered top-k. The router's value governs + /// the whole query; see [`TopKAggregateStrategy`]. fn topk_aggregate_strategy(&self) -> TopKAggregateStrategy; fn allow_decimal128(&self) -> bool; @@ -1304,14 +1312,25 @@ pub async fn init_test_logger() { /// Router-side merge strategy for distributed value-ordered top-k (`SELECT ... GROUP BY x ORDER BY /// agg(...) LIMIT k`). Selected by `CUBESTORE_TOPK_STRATEGY`. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +/// +/// Both halves of a split plan must be planned from the same value, or the router combines a worker +/// stream whose ordering it does not have and returns wrong rows instead of failing. So it rides in +/// [`PlanningFlags`], and a worker uses its own value only when the sender sent no flags. +/// +/// The wire names are the env names and must survive a variant rename. No catch-all variant: an +/// unknown strategy fails the message deserialize, which the receiver logs and the sender sees as a +/// dropped connection. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] pub enum TopKAggregateStrategy { /// Original streaming NRA merge with per-row state (default). + #[serde(rename = "streaming")] Streaming, /// Same streaming NRA merge (keeps early termination, bounded router memory), but vectorized. + #[serde(rename = "vectorized_streaming")] VectorizedStreaming, /// ClickHouse-style full re-aggregation on the router + fetch-limited sort. Drops early /// termination, so the router materializes every distinct group. + #[serde(rename = "full_merge")] FullMerge, } diff --git a/rust/cubestore/cubestore/src/queryplanner/optimizations/distributed_partial_aggregate.rs b/rust/cubestore/cubestore/src/queryplanner/optimizations/distributed_partial_aggregate.rs index e936671706bce..4baf4d47559b6 100644 --- a/rust/cubestore/cubestore/src/queryplanner/optimizations/distributed_partial_aggregate.rs +++ b/rust/cubestore/cubestore/src/queryplanner/optimizations/distributed_partial_aggregate.rs @@ -1,4 +1,3 @@ -use crate::cluster::WorkerPlanningParams; use crate::queryplanner::check_memory::CheckMemoryExec; use crate::queryplanner::group_by_limit_aggregate::GroupByLimitAggregateExec; use crate::queryplanner::inline_aggregate::{InlineAggregateExec, InlineAggregateMode}; @@ -106,9 +105,7 @@ pub fn push_aggregate_to_workers( .next() .unwrap(), w.worker_sort_and_limit.clone(), - WorkerPlanningParams { - worker_partition_count: w.properties().output_partitioning().partition_count(), - }, + w.properties().output_partitioning().partition_count(), )) } else { return Ok(p_final); @@ -314,9 +311,7 @@ pub fn push_worker_sort_and_limit( w.limit_and_reverse.clone(), w.required_input_ordering.clone(), w.worker_sort_and_limit.clone(), - WorkerPlanningParams { - worker_partition_count: w.properties().output_partitioning().partition_count(), - }, + w.properties().output_partitioning().partition_count(), ))); } @@ -1260,9 +1255,7 @@ mod tests { limit_and_reverse, None, None, - WorkerPlanningParams { - worker_partition_count: 1, - }, + 1, )) } diff --git a/rust/cubestore/cubestore/src/queryplanner/optimizations/mod.rs b/rust/cubestore/cubestore/src/queryplanner/optimizations/mod.rs index bd0f1b782962f..2175a32372be6 100644 --- a/rust/cubestore/cubestore/src/queryplanner/optimizations/mod.rs +++ b/rust/cubestore/cubestore/src/queryplanner/optimizations/mod.rs @@ -7,8 +7,7 @@ pub mod rolling_optimizer; mod trace_data_loaded; use super::serialized_plan::PreSerializedPlan; -use crate::cluster::{Cluster, WorkerPlanningParams}; -use crate::config::TopKAggregateStrategy; +use crate::cluster::{Cluster, PlanningFlags}; use crate::queryplanner::optimizations::distributed_partial_aggregate::{ add_limit_to_workers, drop_sort_merge_under_global_aggregate, ensure_partition_merge, push_aggregate_to_workers, push_sorted_partial_aggregate_below_merge, @@ -40,13 +39,13 @@ pub struct CubeQueryPlanner { /// Set on the router cluster: Option>, /// Set on the worker - worker_partition_count: Option, + worker_partition_count: Option, serialized_plan: Arc, memory_handler: Arc, data_loaded_size: Option>, - group_by_limit_factor: usize, + /// On the router, this node's own configuration; on a worker, what the router sent. + planning_flags: PlanningFlags, group_by_limit_per_partition: bool, - topk_strategy: TopKAggregateStrategy, } impl CubeQueryPlanner { @@ -54,9 +53,8 @@ impl CubeQueryPlanner { cluster: Arc, serialized_plan: Arc, memory_handler: Arc, - group_by_limit_factor: usize, + planning_flags: PlanningFlags, group_by_limit_per_partition: bool, - topk_strategy: TopKAggregateStrategy, ) -> CubeQueryPlanner { CubeQueryPlanner { cluster: Some(cluster), @@ -64,30 +62,29 @@ impl CubeQueryPlanner { serialized_plan, memory_handler, data_loaded_size: None, - group_by_limit_factor, + planning_flags, group_by_limit_per_partition, - topk_strategy, } } + /// The worker plans from the flags the router sent, not from its own configuration: the two + /// halves of a split plan only fit together when both were planned from the same values. pub fn new_on_worker( serialized_plan: Arc, - worker_planning_params: WorkerPlanningParams, + worker_partition_count: usize, memory_handler: Arc, data_loaded_size: Option>, - group_by_limit_factor: usize, + planning_flags: PlanningFlags, group_by_limit_per_partition: bool, - topk_strategy: TopKAggregateStrategy, ) -> CubeQueryPlanner { CubeQueryPlanner { serialized_plan, cluster: None, - worker_partition_count: Some(worker_planning_params), + worker_partition_count: Some(worker_partition_count), memory_handler, data_loaded_size, - group_by_limit_factor, + planning_flags, group_by_limit_per_partition, - topk_strategy, } } } @@ -108,9 +105,9 @@ impl QueryPlanner for CubeQueryPlanner { let p = DefaultPhysicalPlanner::with_extension_planners(vec![ Arc::new(CubeExtensionPlanner { cluster: self.cluster.clone(), - worker_planning_params: self.worker_partition_count, + worker_partition_count: self.worker_partition_count, serialized_plan: self.serialized_plan.clone(), - topk_strategy: self.topk_strategy, + planning_flags: self.planning_flags, }), Arc::new(RollingWindowPlanner {}), ]) @@ -121,7 +118,7 @@ impl QueryPlanner for CubeQueryPlanner { self.memory_handler.clone(), self.data_loaded_size.clone(), ctx_state.config().options(), - self.group_by_limit_factor, + self.planning_flags.group_by_limit_factor, self.group_by_limit_per_partition, ); result diff --git a/rust/cubestore/cubestore/src/queryplanner/panic.rs b/rust/cubestore/cubestore/src/queryplanner/panic.rs index 4e47faf2a6e61..11b4e03db4dbf 100644 --- a/rust/cubestore/cubestore/src/queryplanner/panic.rs +++ b/rust/cubestore/cubestore/src/queryplanner/panic.rs @@ -1,4 +1,3 @@ -use crate::cluster::WorkerPlanningParams; use crate::queryplanner::planning::WorkerExec; use async_trait::async_trait; use datafusion::arrow::datatypes::Schema; @@ -184,8 +183,7 @@ pub fn plan_panic_worker() -> Result, DataFusionError> { // a WorkerExec for some reason. (Also, it's important that DF optimizations run identically // when it comes to aggregates pushed down through ClusterSend and the like -- it's actually // NOT important for panic worker planning.) - WorkerPlanningParams { - worker_partition_count: 1, - }, + /* worker_partition_count */ + 1, ))) } diff --git a/rust/cubestore/cubestore/src/queryplanner/planning.rs b/rust/cubestore/cubestore/src/queryplanner/planning.rs index 97c87789c96bb..ab985f7eb717a 100644 --- a/rust/cubestore/cubestore/src/queryplanner/planning.rs +++ b/rust/cubestore/cubestore/src/queryplanner/planning.rs @@ -32,8 +32,7 @@ use itertools::{EitherOrBoth, Itertools}; use std::any::Any; use std::fmt::Formatter; -use crate::cluster::{Cluster, WorkerPlanningParams}; -use crate::config::TopKAggregateStrategy; +use crate::cluster::{Cluster, PlanningFlags}; use crate::metastore::multi_index::MultiPartition; use crate::metastore::table::{Table, TablePath}; use crate::metastore::{ @@ -1969,9 +1968,10 @@ fn pull_up_cluster_send(mut p: LogicalPlan) -> Result>, // Set on the workers. - pub worker_planning_params: Option, + pub worker_partition_count: Option, pub serialized_plan: Arc, - pub topk_strategy: TopKAggregateStrategy, + /// On the router, this node's own configuration; on a worker, what the router sent. + pub planning_flags: PlanningFlags, } #[async_trait] @@ -2116,16 +2116,19 @@ impl CubeExtensionPlanner { limit_and_reverse, worker_sort_and_limit, required_input_ordering, + self.planning_flags, )?)) } else { - let worker_planning_params = self.worker_planning_params.expect("cluster_send_partition_count must be set when CubeExtensionPlanner::cluster is None"); + let worker_partition_count = self.worker_partition_count.expect( + "worker_partition_count must be set when CubeExtensionPlanner::cluster is None", + ); Ok(Arc::new(WorkerExec::new( input, max_batch_rows, limit_and_reverse, required_input_ordering, worker_sort_and_limit, - worker_planning_params, + worker_partition_count, ))) } } @@ -2150,13 +2153,11 @@ impl WorkerExec { limit_and_reverse: Option<(usize, bool)>, required_input_ordering: Option, worker_sort_and_limit: Option, - worker_planning_params: WorkerPlanningParams, + worker_partition_count: usize, ) -> WorkerExec { // This, importantly, gives us the same PlanProperties as ClusterSendExec. - let properties = ClusterSendExec::compute_properties( - input.properties(), - worker_planning_params.worker_partition_count, - ); + let properties = + ClusterSendExec::compute_properties(input.properties(), worker_partition_count); WorkerExec { input, max_batch_rows, diff --git a/rust/cubestore/cubestore/src/queryplanner/query_executor.rs b/rust/cubestore/cubestore/src/queryplanner/query_executor.rs index ff71268ac5555..312597d66dd5c 100644 --- a/rust/cubestore/cubestore/src/queryplanner/query_executor.rs +++ b/rust/cubestore/cubestore/src/queryplanner/query_executor.rs @@ -1,5 +1,5 @@ use crate::cluster::{ - pick_worker_by_ids, pick_worker_by_partitions, Cluster, WorkerPlanningParams, + pick_worker_by_ids, pick_worker_by_partitions, Cluster, PlanningFlags, WorkerPlanningParams, }; use crate::config::injection::DIService; use crate::config::ConfigObj; @@ -567,9 +567,8 @@ impl QueryExecutorImpl { cluster, serialized_plan, self.memory_handler.clone(), - self.config.group_by_limit_factor(), + PlanningFlags::from_config(self.config.as_ref()), self.config.group_by_limit_per_partition(), - self.config.topk_aggregate_strategy(), )) } @@ -580,14 +579,20 @@ impl QueryExecutorImpl { worker_planning_params: WorkerPlanningParams, data_loaded_size: Option>, ) -> Result, CubeError> { + // A sender that does not send the flags planned from its own configuration, and this + // node's configuration reproduces it as long as the value is unset (both binaries default + // to the same one) or set on every node. A value set on the router alone is the one case + // it cannot reproduce, so keep such a value set cluster-wide until every node sends flags. + let planning_flags = worker_planning_params + .flags + .unwrap_or_else(|| PlanningFlags::from_config(self.config.as_ref())); self.make_context(CubeQueryPlanner::new_on_worker( serialized_plan, - worker_planning_params, + worker_planning_params.worker_partition_count, self.memory_handler.clone(), data_loaded_size.clone(), - self.config.group_by_limit_factor(), + planning_flags, self.config.group_by_limit_per_partition(), - self.config.topk_aggregate_strategy(), )) } @@ -1482,6 +1487,8 @@ pub struct ClusterSendExec { pub required_input_ordering: Option, /// Not used in execution, only stored to allow consistent optimization on router and worker. pub worker_sort_and_limit: Option<(Vec<(usize, bool, bool)>, usize)>, + /// The flags this node planned with, sent to the worker so it plans its half the same way. + pub planning_flags: PlanningFlags, } pub type PartitionWithFilters = (u64, RowRange); @@ -1506,6 +1513,7 @@ impl ClusterSendExec { limit_and_reverse: Option<(usize, bool)>, worker_sort_and_limit: Option<(Vec<(usize, bool, bool)>, usize)>, required_input_ordering: Option, + planning_flags: PlanningFlags, ) -> Result { let partitions = Self::distribute_to_workers( cluster.config().as_ref(), @@ -1525,6 +1533,7 @@ impl ClusterSendExec { limit_and_reverse, required_input_ordering, worker_sort_and_limit, + planning_flags, }) } @@ -1551,6 +1560,7 @@ impl ClusterSendExec { WorkerPlanningParams { // Or, self.partitions.len(). worker_partition_count: self.properties().output_partitioning().partition_count(), + flags: Some(self.planning_flags), } } @@ -1847,6 +1857,7 @@ impl ClusterSendExec { limit_and_reverse: self.limit_and_reverse, worker_sort_and_limit: self.worker_sort_and_limit.clone(), required_input_ordering: new_required_input_ordering, + planning_flags: self.planning_flags, } } @@ -1915,6 +1926,7 @@ impl ExecutionPlan for ClusterSendExec { limit_and_reverse: self.limit_and_reverse, worker_sort_and_limit: self.worker_sort_and_limit.clone(), required_input_ordering: self.required_input_ordering.clone(), + planning_flags: self.planning_flags, })) } @@ -2530,7 +2542,53 @@ fn slice_copy(a: &dyn Array, start: usize, len: usize) -> ArrayRef { #[cfg(test)] mod tests { use super::*; + use crate::cluster::MockCluster; + use crate::config::TopKAggregateStrategy; + use crate::queryplanner::planning::PlanningMeta; use datafusion::arrow::datatypes::Field; + use datafusion::common::DFSchema; + use datafusion::logical_expr::EmptyRelation; + use std::collections::HashMap; + + /// The flags the router stamped must reach the worker, not silently decay to the absent case + /// that makes the worker plan from its own configuration. + #[test] + fn cluster_send_exec_sends_its_planning_flags() -> Result<(), CubeError> { + let flags = PlanningFlags { + group_by_limit_factor: 3, + topk_strategy: TopKAggregateStrategy::FullMerge, + }; + let input: Arc = Arc::new(EmptyExec::new(Arc::new(Schema::empty()))); + let plan = PreSerializedPlan::try_new( + LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: Arc::new(DFSchema::empty()), + }), + PlanningMeta { + indices: Vec::new(), + multi_part_subtree: HashMap::new(), + pushable_chunk_filters: Vec::new(), + }, + None, + )?; + let exec = ClusterSendExec { + properties: ClusterSendExec::compute_properties(input.properties(), 2), + partitions: Vec::new(), + cluster: Arc::new(MockCluster::new()), + serialized_plan: Arc::new(plan), + input_for_optimizations: input, + use_streaming: false, + limit_and_reverse: None, + required_input_ordering: None, + worker_sort_and_limit: None, + planning_flags: flags, + }; + + let params = exec.worker_planning_params(); + assert_eq!(params.worker_partition_count, 2); + assert_eq!(params.flags, Some(flags)); + Ok(()) + } #[test] fn test_batch_to_dataframe() -> Result<(), CubeError> { diff --git a/rust/cubestore/cubestore/src/queryplanner/topk/plan.rs b/rust/cubestore/cubestore/src/queryplanner/topk/plan.rs index 4ad1a190e8afa..5fce7903eee1f 100644 --- a/rust/cubestore/cubestore/src/queryplanner/topk/plan.rs +++ b/rust/cubestore/cubestore/src/queryplanner/topk/plan.rs @@ -616,7 +616,7 @@ pub fn plan_topk( // Full-merge strategy: workers send their groups unsorted, the router re-aggregates with one // vectorized hash aggregate and takes the top-k with a fetch-limited sort. Drops early // termination (so the router materializes every distinct group), so it is opt-in and skips HLL. - if ext_planner.topk_strategy == TopKAggregateStrategy::FullMerge + if ext_planner.planning_flags.topk_strategy == TopKAggregateStrategy::FullMerge && agg_fun .iter() .all(|(f, _)| *f != TopKAggregateFunction::Merge) @@ -679,7 +679,8 @@ pub fn plan_topk( None }; - let merge_version = if ext_planner.topk_strategy == TopKAggregateStrategy::VectorizedStreaming + let merge_version = if ext_planner.planning_flags.topk_strategy + == TopKAggregateStrategy::VectorizedStreaming && topk_v2_supported(lower_node, group_expr_len, agg_fun.as_slice()) { TopKMergeVersion::V2 diff --git a/rust/cubestore/cubestore/src/sql/mod.rs b/rust/cubestore/cubestore/src/sql/mod.rs index 1c80d9796caeb..0ce3abad3b156 100644 --- a/rust/cubestore/cubestore/src/sql/mod.rs +++ b/rust/cubestore/cubestore/src/sql/mod.rs @@ -856,13 +856,7 @@ impl SqlService for SqlServiceImpl { } else { let worker = &workers[0]; cluster - .run_select( - worker, - plan, - WorkerPlanningParams { - worker_partition_count: 1, - }, - ) + .run_select(worker, plan, WorkerPlanningParams::no_worker()) .await?; } panic!("worker did not panic")