feat: Multiple rollups in rollupLambda support (#6008)

Author: paveltiunov

packages/cubejs-backend-shared/src/time.ts

@@ -106,3 +106,16 @@ export const addSecondsToLocalTimestamp = (timestamp: string, timezone: string,
     .add(seconds, 'second')
     .toDate();
 };
+
+export const reformatInIsoLocal = (timestamp: string): string => {
+  if (!timestamp) {
+    return timestamp;
+  }
+  if (timestamp.length === 23) {
+    return timestamp;
+  }
+  if (timestamp.length === 24) {
+    return timestamp.replace('Z', '');
+  }
+  return moment.tz(timestamp, 'UTC').utc().format(moment.HTML5_FMT.DATETIME_LOCAL_MS);
+};

packages/cubejs-query-orchestrator/src/orchestrator/PreAggregations.ts

@@ -8,7 +8,7 @@ import {
   FROM_PARTITION_RANGE,
   getEnv,
   inDbTimeZone,
-  MAX_SOURCE_ROW_LIMIT,
+  MAX_SOURCE_ROW_LIMIT, reformatInIsoLocal,
   timeSeries,
   TO_PARTITION_RANGE,
   utcToLocalTimeZone,
@@ -192,6 +192,8 @@ export type PreAggregationDescription = {
   buildRangeEnd?: string;
   updateWindowSeconds?: number;
   sealAt?: string;
+  rollupLambdaId?: string;
+  lastRollupLambda?: boolean;
 };
 
 const tablesToVersionEntries = (schema, tables: TableCacheEntry[]): VersionEntry[] => R.sortBy(
@@ -456,6 +458,8 @@ type LoadPreAggregationResult = {
   buildRangeEnd?: string;
   lambdaTable?: InlineTable;
   queryKey?: any[];
+  rollupLambdaId?: string;
+  partitionRange?: QueryDateRange;
 };
 
 export class PreAggregationLoader {
@@ -1440,7 +1444,7 @@ export class PreAggregationPartitionRangeLoader {
     // eslint-disable-next-line no-use-before-define
     private readonly preAggregations: PreAggregations,
     private readonly preAggregation: PreAggregationDescription,
-    private readonly preAggregationsTablesToTempTables: any,
+    private readonly preAggregationsTablesToTempTables: [string, LoadPreAggregationResult][],
     private readonly loadCache: any,
     private readonly options: PreAggsPartitionRangeLoaderOpts = {
       maxPartitions: 10000,
@@ -1562,7 +1566,7 @@ export class PreAggregationPartitionRangeLoader {
     );
     const [_, buildRangeEnd] = buildRange;
     const loadRange: [string, string] = [...range];
-    if (this.preAggregation.unionWithSourceData && buildRangeEnd < range[1]) {
+    if (buildRangeEnd < range[1]) {
       loadRange[1] = buildRangeEnd;
     }
     const sealAt = addSecondsToLocalTimestamp(
@@ -1581,7 +1585,7 @@ export class PreAggregationPartitionRangeLoader {
         .map(q => ({ ...q, sql: this.replacePartitionSqlAndParams(q.sql, range, partitionTableName) })),
       previewSql: this.preAggregation.previewSql &&
         this.replacePartitionSqlAndParams(this.preAggregation.previewSql, range, partitionTableName),
-      buildRangeEnd,
+      buildRangeEnd: loadRange[1],
       sealAt, // Used only for kSql pre aggregations
     };
   }
@@ -1600,34 +1604,61 @@ export class PreAggregationPartitionRangeLoader {
         this.loadCache,
         this.options,
       ));
-      const resolveResults = await Promise.all(partitionLoaders.map(l => l.loadPreAggregation(false)));
-      const loadResults = resolveResults.filter(res => res !== null);
+      const resolveResults = await Promise.all(partitionLoaders.map(async (l, i) => {
+        const result = await l.loadPreAggregation(false);
+        return result && {
+          ...result,
+          partitionRange: partitionRanges[i]
+        };
+      }));
+      let loadResults = resolveResults.filter(res => res !== null);
       if (this.options.externalRefresh && loadResults.length === 0) {
         throw new Error(
           // eslint-disable-next-line no-use-before-define
           PreAggregations.noPreAggregationPartitionsBuiltMessage(partitionLoaders.map(p => p.preAggregation))
         );
       }
+
+      let lambdaTable: InlineTable;
+      let emptyResult = false;
+
+      if (this.preAggregation.rollupLambdaId) {
+        if (this.lambdaQuery && loadResults.length > 0) {
+          const { buildRangeEnd } = loadResults[loadResults.length - 1];
+          lambdaTable = await this.downloadLambdaTable(buildRangeEnd);
+        }
+        const rollupLambdaResults = this.preAggregationsTablesToTempTables.filter(tempTableResult => tempTableResult[1].rollupLambdaId === this.preAggregation.rollupLambdaId);
+        const lastResult = rollupLambdaResults[rollupLambdaResults.length - 1];
+        const filteredResults = loadResults.filter(
+          r => (this.preAggregation.lastRollupLambda || reformatInIsoLocal(r.buildRangeEnd) === reformatInIsoLocal(r.partitionRange[1])) &&
+            (!lastResult || !lastResult[1].buildRangeEnd || reformatInIsoLocal(lastResult[1].buildRangeEnd) < reformatInIsoLocal(r.partitionRange[0]))
+        );
+        if (filteredResults.length === 0) {
+          emptyResult = true;
+          loadResults = [loadResults[loadResults.length - 1]];
+        } else {
+          loadResults = filteredResults;
+        }
+      }
+
       const allTableTargetNames = loadResults.map(targetTableName => targetTableName.targetTableName);
       let lastUpdatedAt = getLastUpdatedAtTimestamp(loadResults.map(r => r.lastUpdatedAt));
-      let lambdaTable: InlineTable;
 
-      if (this.lambdaQuery && loadResults.length > 0) {
-        const { buildRangeEnd } = loadResults[loadResults.length - 1];
-        lambdaTable = await this.downloadLambdaTable(buildRangeEnd);
+      if (lambdaTable) {
         allTableTargetNames.push(lambdaTable.name);
         lastUpdatedAt = Date.now();
       }
 
       const unionTargetTableName = allTableTargetNames
-        .map(targetTableName => `SELECT * FROM ${targetTableName}`)
+        .map(targetTableName => `SELECT * FROM ${targetTableName}${emptyResult ? ' WHERE 1 = 0' : ''}`)
         .join(' UNION ALL ');
       return {
-        targetTableName: allTableTargetNames.length === 1 ? allTableTargetNames[0] : `(${unionTargetTableName})`,
+        targetTableName: allTableTargetNames.length === 1 && !emptyResult ? allTableTargetNames[0] : `(${unionTargetTableName})`,
         refreshKeyValues: loadResults.map(t => t.refreshKeyValues),
         lastUpdatedAt,
-        buildRangeEnd: undefined,
+        buildRangeEnd: !emptyResult && loadResults.length && loadResults[loadResults.length - 1].buildRangeEnd,
         lambdaTable,
+        rollupLambdaId: this.preAggregation.rollupLambdaId,
       };
     } else {
       return new PreAggregationLoader(
@@ -1713,9 +1744,7 @@ export class PreAggregationPartitionRangeLoader {
     ));
     if (partitionRanges.length > this.options.maxPartitions) {
       throw new Error(
-        `The maximum number of partitions (${
-          this.options.maxPartitions
-        }) was reached for the pre-aggregation`
+        `Pre-aggregation '${this.preAggregation.tableName}' requested to build ${partitionRanges.length} partitions which exceeds the maximum number of partitions per pre-aggregation of ${this.options.maxPartitions}`
       );
     }
     return { buildRange: dateRange, partitionRanges };

packages/cubejs-query-orchestrator/test/unit/QueryOrchestrator.test.js

@@ -3,14 +3,22 @@ import { QueryOrchestrator } from '../../src/orchestrator/QueryOrchestrator';
 
 class MockDriver {
   constructor({ csvImport } = {}) {
-    this.tables = [];
+    this.tablesObj = [];
     this.tablesReady = [];
     this.executedQueries = [];
     this.cancelledQueries = [];
     this.csvImport = csvImport;
     this.now = new Date().getTime();
   }
 
+  get tables() {
+    return this.tablesObj.map(t => t.tableName || t);
+  }
+
+  resetTables() {
+    this.tablesObj = [];
+  }
+
   query(query) {
     this.executedQueries.push(query);
     let promise = Promise.resolve([query]);
@@ -26,6 +34,10 @@ class MockDriver {
       promise = promise.then(() => [{ now: new Date().toJSON() }]);
     }
 
+    if (query.match(/^SELECT '(\d+-\d+-\d+)'/)) {
+      promise = promise.then(() => [{ date: new Date(`${query.match(/^SELECT '(\d+-\d+-\d+)'/)[1]}T00:00:00.000Z`).toJSON() }]);
+    }
+
     if (query.match(/^SELECT MAX\(timestamp\)/)) {
       promise = promise.then(() => [{ max: new Date('2021-06-01T00:00:00.000Z').toJSON() }]);
     }
@@ -56,7 +68,11 @@ class MockDriver {
     if (this.tablesQueryDelay) {
       await this.delay(this.tablesQueryDelay);
     }
-    return this.tables.filter(t => t.split('.')[0] === schema).map(t => ({ table_name: t.replace(`${schema}.`, '') }));
+    return this.tablesObj.filter(t => (t.tableName || t).split('.')[0] === schema)
+      .map(t => ({
+        table_name: (t.tableName || t).replace(`${schema}.`, ''),
+        build_range_end: t.buildRangeEnd
+      }));
   }
 
   delay(timeout) {
@@ -69,15 +85,15 @@ class MockDriver {
   }
 
   loadPreAggregationIntoTable(preAggregationTableName, loadSql) {
-    this.tables.push(preAggregationTableName.substring(0, 100));
+    this.tablesObj.push({ tableName: preAggregationTableName.substring(0, 100) });
     const promise = this.query(loadSql);
     const resPromise = promise.then(() => this.tablesReady.push(preAggregationTableName.substring(0, 100)));
     resPromise.cancel = promise.cancel;
     return resPromise;
   }
 
   async dropTable(tableName) {
-    this.tables = this.tables.filter(t => t !== tableName);
+    this.tablesObj = this.tablesObj.filter(t => (t.tableName || t) !== tableName);
     return this.query(`DROP TABLE ${tableName}`);
   }
 
@@ -109,12 +125,12 @@ class ExternalMockDriver extends MockDriver {
   }
 
   async uploadTable(table) {
-    this.tables.push(table.substring(0, 100));
+    this.tablesObj.push({ tableName: table.substring(0, 100) });
     throw new Error('uploadTable has been called instead of uploadTableWithIndexes');
   }
 
-  async uploadTableWithIndexes(table, columns, tableData, indexesSql) {
-    this.tables.push(table.substring(0, 100));
+  async uploadTableWithIndexes(table, columns, tableData, indexesSql, uniqueKeyColumns, queryTracingObj, externalOptions) {
+    this.tablesObj.push({ tableName: table.substring(0, 100), buildRangeEnd: queryTracingObj?.buildRangeEnd });
     if (tableData.csvFile) {
       this.csvFiles.push(tableData.csvFile);
     }
@@ -209,7 +225,7 @@ describe('QueryOrchestrator', () => {
         }),
       },
       preAggregationsOptions: {
-        maxPartitions: 32,
+        maxPartitions: 100,
         queueOptions: () => ({
           executionTimeout: 2,
           concurrency: 2,
@@ -535,7 +551,7 @@ describe('QueryOrchestrator', () => {
   });
 
   test('save structure versions', async () => {
-    mockDriver.tables = [];
+    mockDriver.resetTables();
     await queryOrchestrator.fetchQuery({
       query: 'SELECT * FROM stb_pre_aggregations.orders',
       values: [],
@@ -1017,7 +1033,7 @@ describe('QueryOrchestrator', () => {
     await expect(async () => {
       await queryOrchestrator.fetchQuery(query);
     }).rejects.toThrow(
-      'The maximum number of partitions (32) was reached for the pre-aggregation'
+      'Pre-aggregation \'stb_pre_aggregations.orders_d\' requested to build 745 partitions which exceeds the maximum number of partitions per pre-aggregation of 100'
     );
   });
 
@@ -1137,6 +1153,73 @@ describe('QueryOrchestrator', () => {
     expect(result.data[0]).toMatch(/orders_d20210601/);
   });
 
+  test('lambda partitions', async () => {
+    const query = (matchedTimeDimensionDateRange) => ({
+      query: 'SELECT * FROM stb_pre_aggregations.orders_d UNION ALL SELECT * FROM stb_pre_aggregations.orders_h',
+      values: [],
+      cacheKeyQueries: {
+        queries: []
+      },
+      preAggregations: [{
+        preAggregationsSchema: 'stb_pre_aggregations',
+        tableName: 'stb_pre_aggregations.orders_d',
+        loadSql: [
+          'CREATE TABLE stb_pre_aggregations.orders_d AS SELECT * FROM public.orders WHERE timestamp >= ? AND timestamp <= ?',
+          ['__FROM_PARTITION_RANGE', '__TO_PARTITION_RANGE']
+        ],
+        invalidateKeyQueries: [['SELECT CASE WHEN NOW() > ? THEN NOW() END as now', ['__TO_PARTITION_RANGE'], {
+          renewalThreshold: 1,
+          updateWindowSeconds: 86400,
+          renewalThresholdOutsideUpdateWindow: 86400,
+          incremental: true
+        }]],
+        preAggregationStartEndQueries: [
+          ['SELECT MIN(timestamp) FROM orders', []],
+          ['SELECT \'2021-05-31\'', []],
+        ],
+        external: true,
+        partitionGranularity: 'day',
+        timezone: 'UTC',
+        rollupLambdaId: 'orders.d_lambda',
+        matchedTimeDimensionDateRange
+      }, {
+        preAggregationsSchema: 'stb_pre_aggregations',
+        tableName: 'stb_pre_aggregations.orders_h',
+        loadSql: [
+          'CREATE TABLE stb_pre_aggregations.orders_h AS SELECT * FROM public.orders WHERE timestamp >= ? AND timestamp <= ?',
+          ['__FROM_PARTITION_RANGE', '__TO_PARTITION_RANGE']
+        ],
+        invalidateKeyQueries: [['SELECT CASE WHEN NOW() > ? THEN NOW() END as now', ['__TO_PARTITION_RANGE'], {
+          renewalThreshold: 1,
+          updateWindowSeconds: 86400,
+          renewalThresholdOutsideUpdateWindow: 86400,
+          incremental: true
+        }]],
+        preAggregationStartEndQueries: [
+          ['SELECT \'2021-05-30\'', []],
+          ['SELECT MAX(timestamp) FROM orders', []],
+        ],
+        external: true,
+        partitionGranularity: 'hour',
+        timezone: 'UTC',
+        rollupLambdaId: 'orders.d_lambda',
+        lastRollupLambda: true,
+        matchedTimeDimensionDateRange
+      }],
+      requestId: 'lambda partitions',
+      external: true,
+    });
+    let result = await queryOrchestrator.fetchQuery(query());
+    console.log(JSON.stringify(result, null, 2));
+    expect(result.data[0]).toMatch(/orders_d20210501/);
+    expect(result.data[0]).not.toMatch(/orders_h2021053000/);
+    expect(result.data[0]).toMatch(/orders_h2021053100/);
+
+    result = await queryOrchestrator.fetchQuery(query(['2021-05-31T00:00:00.000', '2021-05-31T23:59:59.999']));
+    console.log(JSON.stringify(result, null, 2));
+    expect(result.data[0]).toMatch(/orders_h2021053100/);
+  });
+
   test('loadRefreshKeys', async () => {
     const preAggregationsLoadCacheByDataSource = {};
     const preAggregationExternalRefreshKey = {

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

@@ -576,7 +576,7 @@ class BaseQuery {
     const preAggForQuery = this.preAggregations.findPreAggregationForQuery();
     const result = {};
     if (preAggForQuery && preAggForQuery.preAggregation.unionWithSourceData) {
-      const lambdaPreAgg = preAggForQuery.referencedPreAggregations[0];
+      const lambdaPreAgg = preAggForQuery.referencedPreAggregations[preAggForQuery.referencedPreAggregations.length - 1];
       // TODO(cristipp) Use source query instead of preaggregation references.
       const references = this.cubeEvaluator.evaluatePreAggregationReferences(lambdaPreAgg.cube, lambdaPreAgg.preAggregation);
       const lambdaQuery = this.newSubQuery(