[pull] master from cube-js:master

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 () => {

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'/
+      );
+    });
+  }
+});

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<CompiledPreAggregation>,
         date_range: Option<(String, String)>,
         is_user_query: bool,
+        time_shifts: &TimeShiftState,
     ) -> Result<Option<Rc<Query>>, 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<String> {
+        let mut symbols: Vec<Rc<MemberSymbol>> = 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<String>,
+        read_members: &HashSet<String>,
+        time_shifts: &TimeShiftState,
+    ) -> bool {
+        if time_shifts.is_empty() {
+            return true;
+        }
+        let is_read = |member: &Rc<MemberSymbol>| {
+            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<State>,
@@ -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();