feat(tesseract): Support separate pre-aggregations for different multi-stage subqueries (cube-js#10684)

Author: waralexrom

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

@@ -297,14 +297,57 @@ export class PreAggregationPartitionRangeLoader {
         .map(targetTableName => `SELECT * FROM ${targetTableName}${emptyResult ? ' WHERE 1 = 0' : ''}`)
         .join(' UNION ALL ');
 
+      const baseTargetTableName = allTableTargetNames.length === 1 && !emptyResult ? allTableTargetNames[0] : `(${unionTargetTableName})`;
+
+      // Build per-usage target table names if usageMapping is present
+      let usageTargetTableNames: Record<string, string> | undefined;
+      if (this.preAggregation.usageMapping) {
+        usageTargetTableNames = {};
+        for (const [suffix, usageInfo] of Object.entries(this.preAggregation.usageMapping)) {
+          if (usageInfo.dateRange && this.preAggregation.partitionGranularity) {
+            // Load partition ranges specific to this usage's dateRange.
+            // Use partitionRange (generated locally via timeSeries, always in DEFAULT_TS_FORMAT)
+            // instead of buildRangeEnd (from DB, may include Z suffix depending on driver timestampFormat).
+            const usageDateRange = PreAggregationPartitionRangeLoader.intersectDateRanges(
+              [loadResults[0]?.partitionRange?.[0] || null, loadResults[loadResults.length - 1]?.partitionRange?.[1] || null] as QueryDateRange,
+              usageInfo.dateRange as QueryDateRange,
+            );
+            if (usageDateRange) {
+              const usagePartitions = loadResults.filter(r => {
+                if (!r.partitionRange) return true;
+                const [pStart, pEnd] = r.partitionRange;
+                const [uStart, uEnd] = usageDateRange;
+                return pEnd >= uStart && pStart <= uEnd;
+              });
+              const usageTableNames = usagePartitions.map(r => r.targetTableName);
+              if (usageTableNames.length === 1) {
+                [usageTargetTableNames[suffix]] = usageTableNames;
+              } else if (usageTableNames.length > 0) {
+                const usageUnion = usageTableNames
+                  .map(t => `SELECT * FROM ${t}`)
+                  .join(' UNION ALL ');
+                usageTargetTableNames[suffix] = `(${usageUnion})`;
+              } else {
+                usageTargetTableNames[suffix] = baseTargetTableName;
+              }
+            } else {
+              usageTargetTableNames[suffix] = baseTargetTableName;
+            }
+          } else {
+            usageTargetTableNames[suffix] = baseTargetTableName;
+          }
+        }
+      }
+
       return {
-        targetTableName: allTableTargetNames.length === 1 && !emptyResult ? allTableTargetNames[0] : `(${unionTargetTableName})`,
+        targetTableName: baseTargetTableName,
         refreshKeyValues: loadResults.map(t => t.refreshKeyValues),
         lastUpdatedAt,
         buildRangeEnd: !emptyResult && loadResults.length && loadResults[loadResults.length - 1].buildRangeEnd,
         lambdaTable,
         rollupLambdaId: this.preAggregation.rollupLambdaId,
         isMultiTableUnion: allTableTargetNames.length > 1,
+        usageTargetTableNames,
       };
     } else {
       return new PreAggregationLoader(

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

@@ -149,6 +149,7 @@ export type LoadPreAggregationResult = {
   rollupLambdaId?: string;
   partitionRange?: QueryDateRange;
   isMultiTableUnion?: boolean;
+  usageTargetTableNames?: Record<string, string>;
 };
 
 export type PreAggregationTableToTempTable = [string, LoadPreAggregationResult];
@@ -192,6 +193,7 @@ export type PreAggregationDescription = {
   sealAt?: string;
   rollupLambdaId?: string;
   lastRollupLambda?: boolean;
+  usageMapping?: Record<string, { dateRange?: QueryDateRange }>;
 };
 
 export const tablesToVersionEntries = (schema, tables: TableCacheEntry[]): VersionEntry[] => R.sortBy(
@@ -570,10 +572,22 @@ export class PreAggregations {
             );
           }
 
-          return [p.tableName, usedPreAggregation];
+          return [p.tableName, usedPreAggregation] as PreAggregationTableToTempTable;
         };
 
-        return preAggregationPromise().then(res => preAggregationsTablesToTempTables.concat([res]));
+        return preAggregationPromise().then(([tableName, result]) => {
+          const { usageTargetTableNames } = result;
+          if (usageTargetTableNames && Object.keys(usageTargetTableNames).length > 0) {
+            const entries: PreAggregationTableToTempTable[] = Object.entries(usageTargetTableNames).map(
+              ([suffix, usageTarget]) => [
+                `${tableName}${suffix}`,
+                { ...result, targetTableName: usageTarget, usageTargetTableNames: undefined },
+              ]
+            );
+            return preAggregationsTablesToTempTables.concat(entries);
+          }
+          return preAggregationsTablesToTempTables.concat([[tableName, result]]);
+        });
       }).reduce((promise, fn) => promise.then(fn), Promise.resolve([]));
 
     return preAggregationsTablesToTempTablesPromise.then(preAggregationsTablesToTempTables => ({

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

@@ -429,9 +429,7 @@ export class QueryCache {
       ? queryAndParams
       : [queryAndParams, []];
     const replacedKeyQuery: string = preAggregationsTablesToTempTables.reduce(
-      (query, [tableName, { targetTableName }]) => (
-        QueryCache.replaceAll(tableName, targetTableName, query)
-      ),
+      (query, [tableName, { targetTableName }]) => QueryCache.replaceAll(tableName, targetTableName, query),
       keyQuery
     );
     return Array.isArray(queryAndParams)

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

@@ -959,11 +959,9 @@ export class BaseQuery {
     try {
       const buildResult = nativeBuildSqlAndParams(queryParams);
 
-      const [query, params, preAggregation] = buildResult;
+      const [query, params, preAggResult] = buildResult;
       const paramsArray = [...params];
-      if (preAggregation) {
-        this.preAggregations.preAggregationForQuery = preAggregation;
-      }
+      this.applyNativePreAggResult(preAggResult);
       return [query, paramsArray];
     } catch (e) {
       if (e.name === 'TesseractUserError') {
@@ -1009,8 +1007,21 @@ export class BaseQuery {
 
     const buildResult = nativeBuildSqlAndParams(queryParams);
 
-    const [, , preAggregation] = buildResult;
-    return preAggregation;
+    const [, , preAggResult] = buildResult;
+    this.applyNativePreAggResult(preAggResult);
+    return this.preAggregations.preAggregationForQuery;
+  }
+
+  applyNativePreAggResult(preAggResult) {
+    if (!preAggResult) return;
+    if (Array.isArray(preAggResult)) {
+      this.preAggregations.preAggregationUsageInfos = preAggResult;
+      const first = preAggResult[0];
+      this.preAggregations.preAggregationForQuery =
+        this.getPreAggregationByName(first.cubeName, first.preAggregationName);
+    } else {
+      this.preAggregations.preAggregationForQuery = preAggResult;
+    }
   }
 
   allCubeMembers(path) {
@@ -3496,7 +3507,7 @@ export class BaseQuery {
   }
 
   escapeStringLiteral(str) {
-    return `'${str.replace(/'/g, "''")}'`;
+    return `'${str.replace(/'/g, '\'\'')}'`;
   }
 
   autoPrefixAndEvaluateSql(cubeName, sql, isMemberExpr = false) {

packages/cubejs-schema-compiler/src/adapter/PreAggregations.ts

@@ -94,6 +94,17 @@ export type FullPreAggregationDescription = any;
  */
 export type TransformedQuery = any;
 
+export type UsageDateRangeInfo = {
+  dateRange?: [string, string];
+};
+
+export type PreAggregationUsageInfo = {
+  cubeName: string;
+  preAggregationName: string;
+  external: boolean;
+  usages: Record<string, UsageDateRangeInfo>;
+};
+
 export class PreAggregations {
   private readonly query: BaseQuery;
 
@@ -107,6 +118,8 @@ export class PreAggregations {
 
   public preAggregationForQuery: PreAggregationForQuery | undefined = undefined;
 
+  public preAggregationUsageInfos: PreAggregationUsageInfo[] | undefined = undefined;
+
   public constructor(query: BaseQuery, historyQueries, cubeLatticeCache) {
     this.query = query;
     this.historyQueries = historyQueries;
@@ -137,6 +150,10 @@ export class PreAggregations {
     const isInPreAggregationQuery = this.query.options.preAggregationQuery;
     if (!isInPreAggregationQuery) {
       const preAggregationForQuery = this.findPreAggregationForQuery();
+      // Check usageInfos after findPreAggregationForQuery (which may populate them)
+      if (this.preAggregationUsageInfos && this.preAggregationUsageInfos.length > 0) {
+        return this.preAggregationDescriptionsForUsageInfos(this.preAggregationUsageInfos);
+      }
       if (preAggregationForQuery) {
         return this.preAggregationDescriptionsFor(preAggregationForQuery);
       }
@@ -165,6 +182,48 @@ export class PreAggregations {
     return join.joins.map(j => j.originalTo).concat([join.root]);
   }
 
+  private preAggregationDescriptionsForUsageInfos(usageInfos: PreAggregationUsageInfo[]): FullPreAggregationDescription[] {
+    return usageInfos.flatMap(usageInfo => {
+      const preAggObj = this.getRollupPreAggregationByName(usageInfo.cubeName, usageInfo.preAggregationName);
+      if (!preAggObj || !('preAggregationName' in preAggObj)) {
+        return [];
+      }
+      const foundPreAgg = preAggObj as PreAggregationForQuery;
+
+      // One description per physical pre-aggregation, with usageMapping attached
+      const descriptions = this.preAggregationDescriptionsFor(foundPreAgg);
+
+      // Compute the union of all usage date ranges so that partitions cover
+      // every usage (e.g. time_shift may require earlier partitions).
+      const mergedDateRange = PreAggregations.mergeUsageDateRanges(usageInfo.usages);
+
+      return descriptions.map(desc => ({
+        ...desc,
+        usageMapping: usageInfo.usages,
+        ...(mergedDateRange && desc.matchedTimeDimensionDateRange ? { matchedTimeDimensionDateRange: mergedDateRange } : {}),
+      }));
+    });
+  }
+
+  private static mergeUsageDateRanges(usages: Record<string, UsageDateRangeInfo>): [string, string] | null {
+    let minDate: string | null = null;
+    let maxDate: string | null = null;
+
+    for (const usage of Object.values(usages)) {
+      if (usage.dateRange) {
+        const [from, to] = usage.dateRange;
+        if (!minDate || from < minDate) {
+          minDate = from;
+        }
+        if (!maxDate || to > maxDate) {
+          maxDate = to;
+        }
+      }
+    }
+
+    return minDate && maxDate ? [minDate, maxDate] : null;
+  }
+
   private preAggregationDescriptionsFor(foundPreAggregation: PreAggregationForQuery): FullPreAggregationDescription[] {
     let preAggregations: PreAggregationForQuery[] = [foundPreAggregation];
     if (foundPreAggregation.preAggregation.type === 'rollupJoin') {