fix(schema-compiler): Fix pre-agg matching for 'rollupJoin' / 'rollupLambda' pre-aggregations

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

@@ -507,7 +507,18 @@ export class PreAggregations {
     filterDimensionsSingleValueEqual =
       allValuesEq1(filterDimensionsSingleValueEqual) ? new Set(filterDimensionsSingleValueEqual?.keys()) : null;
 
+    // Build reverse query joins map, which is used for
+    // rollupLambda and rollupJoin pre-aggs matching later
+    const joinsMap: Record<string, string> = {};
+    if (query.join) {
+      for (const j of query.join.joins) {
+        joinsMap[j.to] = j.from;
+      }
+    }
+
     return {
+      joinGraphRoot: query.join?.root,
+      joinsMap,
       sortedDimensions,
       sortedTimeDimensions,
       timeDimensions,
@@ -726,22 +737,51 @@ export class PreAggregations {
         }
       }
 
-      // In 'rollupJoin' / 'rollupLambda' pre-aggregations fullName members will be empty, because there are
-      // no connections in the joinTree between cubes from different datasources
-      const dimsToMatch = references.rollups.length > 0 ? references.dimensions : references.fullNameDimensions;
-
-      const dimensionsMatch = (dimensions, doBackAlias) => R.all(
-        d => (
-          doBackAlias ?
-            backAlias(dimsToMatch) :
-            (dimsToMatch)
-        ).indexOf(d) !== -1,
-        dimensions
-      );
+      let dimsToMatch: string[];
+      let timeDimsToMatch: PreAggregationTimeDimensionReference[];
+      let dimensionsMatch: (dimensions: string[], doBackAlias: boolean) => boolean;
+
+      if (references.rollups.length > 0) {
+        // In 'rollupJoin' / 'rollupLambda' pre-aggregations fullName members will be empty, because there are
+        // no connections in the joinTree between cubes from different datasources
+        // but joinGraph of the query has all the connections, necessary for serving the query,
+        // so we use this information to complete the full paths of members from the root of the query
+        // up to the pre-agg cube.
+        dimsToMatch = references.dimensions;
+        timeDimsToMatch = references.timeDimensions;
+
+        const buildPath = (cube: string): string[] => {
+          const path = [cube];
+          const parentMap = transformedQuery.joinsMap;
+          while (parentMap[cube]) {
+            cube = parentMap[cube];
+            path.push(cube);
+          }
+          return path.reverse();
+        };
+
+        dimensionsMatch = (dimensions, doBackAlias) => {
+          let target = doBackAlias ? backAlias(dimsToMatch) : dimsToMatch;
+          target = target.map(dim => {
+            const [cube, field] = dim.split('.');
+            if (cube === transformedQuery.joinGraphRoot) {
+              return dim;
+            }
+            const path = buildPath(cube);
+            return `${path.join('.')}.${field}`;
+          });
 
-      // In 'rollupJoin' / 'rollupLambda' pre-aggregations fullName members will be empty, because there are
-      // no connections in the joinTree between cubes from different datasources
-      const timeDimsToMatch = references.rollups.length > 0 ? references.timeDimensions : references.fullNameTimeDimensions;
+          return dimensions.every(d => target.includes(d));
+        };
+      } else {
+        dimsToMatch = references.fullNameDimensions;
+        timeDimsToMatch = references.fullNameTimeDimensions;
+
+        dimensionsMatch = (dimensions, doBackAlias) => {
+          const target = doBackAlias ? backAlias(dimsToMatch) : dimsToMatch;
+          return dimensions.every(d => target.includes(d));
+        };
+      }
 
       const timeDimensionsMatch = (timeDimensionsList, doBackAlias) => R.allPass(
         timeDimensionsList.map(

packages/cubejs-schema-compiler/test/integration/postgres/pre-aggregations.test.ts

@@ -599,7 +599,83 @@ describe('PreAggregations', () => {
         }
       ]
     });
-    `);
+
+    cube('cube_1', {
+      sql: \`SELECT 1 as id, 'dim_1' as dim_1\`,
+
+      joins: {
+        cube_2: {
+          relationship: 'many_to_one',
+          sql: \`\${CUBE.dim_1} = \${cube_2.dim_1}\`
+        }
+      },
+
+      dimensions: {
+        id: {
+          sql: 'id',
+          type: 'string',
+          primary_key: true
+        },
+
+        dim_1: {
+          sql: 'dim_1',
+          type: 'string'
+        },
+      },
+
+      pre_aggregations: {
+        aaa: {
+          dimensions: [
+            dim_1
+          ]
+        },
+        rollupJoin: {
+          type: 'rollupJoin',
+          dimensions: [
+            dim_1,
+            cube_2.dim_1,
+            cube_2.dim_2  // XXX
+          ],
+          rollups: [
+            aaa,
+            cube_2.bbb
+          ]
+        }
+      }
+    });
+
+    cube('cube_2', {
+      sql: \`SELECT 2 as id, 'dim_1' as dim_1, 'dim_2' as dim_2\`,
+
+      dimensions: {
+        id: {
+          sql: 'id',
+          type: 'string',
+          primary_key: true
+        },
+
+        dim_1: {
+          sql: 'dim_1',
+          type: 'string'
+        },
+
+        dim_2: {
+          sql: 'dim_2',
+          type: 'string'
+        },
+      },
+
+      pre_aggregations: {
+        bbb: {
+          dimensions: [
+            dim_1,
+            dim_2,
+          ]
+        }
+      }
+    });
+
+  `);
 
   it('simple pre-aggregation', async () => {
     await compiler.compile();
@@ -2773,4 +2849,36 @@ describe('PreAggregations', () => {
     expect(loadSql[0]).not.toMatch(/GROUP BY/);
     expect(loadSql[0]).toMatch(/THEN 1 END `real_time_lambda_visitors__count`/);
   });
+
+  it('rollupJoin pre-aggregation', async () => {
+    await compiler.compile();
+
+    const query = new PostgresQuery({ joinGraph, cubeEvaluator, compiler }, {
+      dimensions: ['cube_1.dim_1', 'cube_2.dim_2'],
+      timezone: 'America/Los_Angeles',
+      preAggregationsSchema: ''
+    });
+
+    const queryAndParams = query.buildSqlAndParams();
+    console.log(queryAndParams);
+    const preAggregationsDescription: any = query.preAggregations?.preAggregationsDescription();
+    console.log(preAggregationsDescription);
+    expect(preAggregationsDescription.length).toBe(2);
+    const aaa = preAggregationsDescription.find(p => p.preAggregationId === 'cube_1.aaa');
+    const bbb = preAggregationsDescription.find(p => p.preAggregationId === 'cube_2.bbb');
+    expect(aaa).toBeDefined();
+    expect(bbb).toBeDefined();
+
+    expect(query.preAggregations?.preAggregationForQuery?.canUsePreAggregation).toEqual(true);
+    expect(query.preAggregations?.preAggregationForQuery?.preAggregationName).toEqual('rollupJoin');
+
+    return dbRunner.evaluateQueryWithPreAggregations(query).then(res => {
+      expect(res).toEqual(
+        [{
+          cube_1__dim_1: 'dim_1',
+          cube_2__dim_2: 'dim_2',
+        }]
+      );
+    });
+  });
 });