feat: Initial support for grouped join pushdown (#9032)

* Load request extended with list of grouped subqueries to join with. Each subquery come with SQL string, alias, condition and join type. SQL is treated as static query, fully porcessed by SQL API. Condition - as member expression, because it can contain references to dimensions
* `BaseQuery` will add grouped subqueries to SQL after subquery dimensions and join tree parts
* Wrapper rewrite rules handle `wrapper_pullup_replacer` in `join` position of `wrapped_select`, but still mostly with empty list inside. Notable exception is aggregation flattening - it is allowed to try to flatten aggregation into WrappedSelect with non-empty joins.
* Wrapper pull up rules no handle trivial pulling up over single `wrapped_select_join` and over list `WrappedSelectJoins` list
* Wrapper replacers now have `grouped_subqueries` in a context, it is used to track qualifiers of grouped subqueries, to allow rewriting of column expressions referencing grouped subquery with push-to-cube active
* Converter from egraph to logical plan reconstructs schema for joins
* Penalize joins harder than wrappers. For queries ungrouped queries with grouped join, like `SELECT * FROM (ungrouped) JOIN (grouped);`, this would allow to prefer representation without join, but with wrapper. Without this cost change join representation would have higher `ast_size_outside_wrapper` component, but it is less important then `wrapper_nodes`, so wrapper representation would be more expensive.
* SQL generation for `WrappedSelectNode` will generate each grouped subquery separately, generate member expressions for join conditions and push all that to Cube
* Rewrite rules are too relaxed for now: they allow any `push_to_cube=false` query as grouped, which is incorrect, as we may drop push-to-Cube when wrapping `WrappedSelect` with another layer, even if inner `WrappedSelect` is projection. This is covered by check in SQL generation: only `CubeScan(grouped=true)` are allowed. This require new flag in wrapper replacers to fix.

Supporting changes:
* Extract `CubeScanWrappedSqlNode` from `CubeScanWrapperNode`. Both a logical plan nodes for now, `Wrapper` is generated during rewrites, and `WrappedSql` stores SQL, generated generated during `evaluate_wrapped_sql`
* Extracted `__cubeJoinField` join check to function, and make it stricter, only `left.__cubeJoinField = right.__cubeJoinField` is allowed
* Avoid clone calls in `find_column_by_alias` and `__cubeJoinField` join check
* Add `PushToCubeContext` in SQL generation code, used to track everything necessary to generate SQL for push-to-Cube case. For now it contains ungrouped `CubeScanNode` that whould be pushed into, and qualifiers and generated SQL for all grouped subqueries. Qualifiers are necessary to generate column expressions on top of ungrouped-grouped join.
* New `copy_value!` macro, to extend `copy_flag!`, allows to copy single value between different unary language variants
* Prettify extracted best plan and cost in logs

Author: mcheshkov

packages/cubejs-api-gateway/openspec.yml

@@ -329,6 +329,24 @@ components:
         - $ref: "#/components/schemas/V1LoadRequestQueryFilterBase"
         - $ref: "#/components/schemas/V1LoadRequestQueryFilterLogicalOr"
         - $ref: "#/components/schemas/V1LoadRequestQueryFilterLogicalAnd"
+    V1LoadRequestQueryJoinSubquery:
+      type: "object"
+      properties:
+        sql:
+          type: "string"
+        # TODO This is _always_ a member expression, maybe pass as parsed, without intermediate string?
+        "on":
+          type: "string"
+        # TODO why string? it's enum
+        joinType:
+          type: "string"
+        alias:
+          type: "string"
+      required:
+        - sql
+        - "on"
+        - joinType
+        - alias
     V1LoadRequestQuery:
       type: "object"
       properties:
@@ -366,6 +384,12 @@ components:
             $ref: "#/components/schemas/V1LoadRequestQueryFilterItem"
         ungrouped:
           type: "boolean"
+        # vector of (subquery sql: string, join condition: member expression, join type: enum)
+        # they will be added to end of joinQuery in BaseQuery, in same order as here
+        subqueryJoins:
+          type: "array"
+          items:
+            $ref: "#/components/schemas/V1LoadRequestQueryJoinSubquery"
     V1LoadRequest:
       type: "object"
       properties:

packages/cubejs-api-gateway/src/gateway.ts

@@ -1312,7 +1312,12 @@ class ApiGateway {
   }
 
   private hasExpressionsInQuery(query: Query): boolean {
-    const arraysToCheck = [query.measures, query.dimensions, query.segments];
+    const arraysToCheck = [
+      query.measures,
+      query.dimensions,
+      query.segments,
+      (query.subqueryJoins ?? []).map(join => join.on),
+    ];
 
     return arraysToCheck.some(array => array?.some(item => typeof item === 'string' && item.startsWith('{')));
   }
@@ -1323,6 +1328,10 @@ class ApiGateway {
       measures: (query.measures || []).map(m => (typeof m === 'string' ? this.parseMemberExpression(m) : m)),
       dimensions: (query.dimensions || []).map(m => (typeof m === 'string' ? this.parseMemberExpression(m) : m)),
       segments: (query.segments || []).map(m => (typeof m === 'string' ? this.parseMemberExpression(m) : m)),
+      subqueryJoins: (query.subqueryJoins ?? []).map(join => (typeof join.on === 'string' ? {
+        ...join,
+        on: this.parseMemberExpression(join.on),
+      } : join)),
     };
   }
 
@@ -1361,6 +1370,10 @@ class ApiGateway {
       measures: (query.measures || []).map(m => (typeof m !== 'string' ? this.evalMemberExpression(m as ParsedMemberExpression) : m)),
       dimensions: (query.dimensions || []).map(m => (typeof m !== 'string' ? this.evalMemberExpression(m as ParsedMemberExpression) : m)),
       segments: (query.segments || []).map(m => (typeof m !== 'string' ? this.evalMemberExpression(m as ParsedMemberExpression) : m)),
+      subqueryJoins: (query.subqueryJoins ?? []).map(join => (typeof join.on !== 'string' ? {
+        ...join,
+        on: this.evalMemberExpression(join.on as ParsedMemberExpression)
+      } : join)),
     };
   }
 

packages/cubejs-api-gateway/src/query.js

@@ -96,6 +96,15 @@ const oneCondition = Joi.object().keys({
   and: Joi.array().items(oneFilter, Joi.link('...').description('oneCondition schema')),
 }).xor('or', 'and');
 
+const subqueryJoin = Joi.object().keys({
+  sql: Joi.string(),
+  // TODO This is _always_ a member expression, maybe pass as parsed, without intermediate string?
+  // TODO there are three different types instead of alternatives for this actually
+  on: Joi.alternatives(Joi.string(), memberExpression, parsedMemberExpression),
+  joinType: Joi.string().valid('LEFT', 'INNER'),
+  alias: Joi.string(),
+});
+
 const querySchema = Joi.object().keys({
   // TODO add member expression alternatives only for SQL API queries?
   measures: Joi.array().items(Joi.alternatives(id, memberExpression, parsedMemberExpression)),
@@ -122,6 +131,7 @@ const querySchema = Joi.object().keys({
   renewQuery: Joi.boolean(),
   ungrouped: Joi.boolean(),
   responseFormat: Joi.valid('default', 'compact'),
+  subqueryJoins: Joi.array().items(subqueryJoin),
 });
 
 const normalizeQueryOrder = order => {

packages/cubejs-api-gateway/src/types/query.ts

@@ -67,6 +67,15 @@ interface QueryTimeDimension {
   granularity?: QueryTimeDimensionGranularity;
 }
 
+type SubqueryJoins = {
+  sql: string,
+  // TODO This is _always_ a member expression, maybe pass as parsed, without intermediate string?
+  // TODO there are three different types instead of alternatives for this actually
+  on: string | ParsedMemberExpression | MemberExpression,
+  joinType: 'LEFT' | 'INNER',
+  alias: string,
+};
+
 /**
  * Incoming network query data type.
  */
@@ -85,6 +94,9 @@ interface Query {
   renewQuery?: boolean;
   ungrouped?: boolean;
   responseFormat?: ResultType;
+
+  // TODO incoming query, query with parsed exprs and query with evaluated exprs are all different types
+  subqueryJoins?: Array<SubqueryJoins>,
 }
 
 /**

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

@@ -60,6 +60,23 @@ const SecondsDurations = {
  * @property {*} headCommitId
  */
 
+/**
+ * @typedef {Object} JoinRoot
+ * @property {string} sql
+ * @property {string} alias
+ */
+
+/**
+ * @typedef {Object} JoinItem
+ * @property {string} sql
+ * @property {string} alias
+ * @property {string} on
+ */
+
+/**
+ * @typedef {[JoinRoot, ...JoinItem]} JoinChain
+ */
+
 /**
  * BaseQuery class. BaseQuery object encapsulates the logic of
  * transforming an incoming to a specific cube request to the
@@ -224,6 +241,7 @@ export class BaseQuery {
       multiStageQuery: this.options.multiStageQuery,
       multiStageDimensions: this.options.multiStageDimensions,
       multiStageTimeDimensions: this.options.multiStageTimeDimensions,
+      subqueryJoins: this.options.subqueryJoins,
     });
     this.from = this.options.from;
     this.multiStageQuery = this.options.multiStageQuery;
@@ -269,6 +287,11 @@ export class BaseQuery {
     this.preAggregationsSchemaOption = this.options.preAggregationsSchema ?? DEFAULT_PREAGGREGATIONS_SCHEMA;
     this.externalQueryClass = this.options.externalQueryClass;
 
+    /**
+     * @type {Array<{sql: string, on: {expression: Function}, joinType: 'LEFT' | 'INNER', alias: string}>}
+     */
+    this.customSubQueryJoins = this.options.subqueryJoins ?? [];
+
     // Set the default order only when options.order is not provided at all
     // if options.order is set (empty array [] or with data) - use it as is
     this.order = this.options.order ?? this.defaultOrder();
@@ -1604,19 +1627,44 @@ export class BaseQuery {
     return this.joinSql([
       { sql: cubeSql, alias: cubeAlias },
       ...(subQueryDimensionsByCube[join.root] || []).map(d => this.subQueryJoin(d)),
-      ...joins
+      ...joins,
+      ...this.customSubQueryJoins.map((customJoin) => this.customSubQueryJoin(customJoin)),
     ]);
   }
 
   joinSql(toJoin) {
     const [root, ...rest] = toJoin;
     const joins = rest.map(
-      j => `LEFT JOIN ${j.sql} ${this.asSyntaxJoin} ${j.alias} ON ${j.on}`
+      j => {
+        const joinType = j.joinType ?? 'LEFT';
+        return `${joinType} JOIN ${j.sql} ${this.asSyntaxJoin} ${j.alias} ON ${j.on}`;
+      }
     );
 
     return [`${root.sql} ${this.asSyntaxJoin} ${root.alias}`, ...joins].join('\n');
   }
 
+  /**
+   *
+   * @param {{sql: string, on: {cubeName: string, expression: Function}, joinType: 'LEFT' | 'INNER', alias: string}} customJoin
+   * @returns {JoinItem}
+   */
+  customSubQueryJoin(customJoin) {
+    const on = this.evaluateSql(customJoin.on.cubeName, customJoin.on.expression);
+
+    return {
+      sql: `(${customJoin.sql})`,
+      alias: customJoin.alias,
+      on,
+      joinType: customJoin.joinType,
+    };
+  }
+
+  /**
+   *
+   * @param {string} dimension
+   * @returns {JoinItem}
+   */
   subQueryJoin(dimension) {
     const { prefix, subQuery, cubeName } = this.subQueryDescription(dimension);
     const primaryKeys = this.cubeEvaluator.primaryKeys[cubeName];

packages/cubejs-testing/test/__snapshots__/smoke-cubesql.test.ts.snap

@@ -57,6 +57,32 @@ Array [
 ]
 `;
 
+exports[`SQL API Postgres (Data) join with filtered grouped query: join grouped with filter 1`] = `
+Array [
+  Object {
+    "count": "2",
+    "status": "processed",
+  },
+  Object {
+    "count": "2",
+    "status": "new",
+  },
+]
+`;
+
+exports[`SQL API Postgres (Data) join with grouped query: join grouped 1`] = `
+Array [
+  Object {
+    "count": "2",
+    "status": "processed",
+  },
+  Object {
+    "count": "1",
+    "status": "shipped",
+  },
+]
+`;
+
 exports[`SQL API Postgres (Data) metabase max number: metabase max number 1`] = `
 Array [
   Object {

packages/cubejs-testing/test/smoke-cubesql.test.ts

@@ -476,6 +476,65 @@ filter_subq AS (
       expect(res.rows).toMatchSnapshot('select __user and literal in wrapper');
     });
 
+    test('join with grouped query', async () => {
+      const query = `
+        SELECT
+          "Orders".status AS status,
+          COUNT(*) AS count
+        FROM
+          "Orders"
+          INNER JOIN
+          (
+            SELECT
+              status,
+              SUM(totalAmount)
+            FROM
+              "Orders"
+            GROUP BY 1
+            ORDER BY 2 DESC
+            LIMIT 2
+          ) top_orders
+        ON
+          "Orders".status = top_orders.status
+        GROUP BY 1
+        ORDER BY 1
+        `;
+
+      const res = await connection.query(query);
+      // Expect only top statuses 2 by total amount: processed and shipped
+      expect(res.rows).toMatchSnapshot('join grouped');
+    });
+
+    test('join with filtered grouped query', async () => {
+      const query = `
+        SELECT
+          "Orders".status AS status,
+          COUNT(*) AS count
+        FROM
+          "Orders"
+          INNER JOIN
+          (
+            SELECT
+              status,
+              SUM(totalAmount)
+            FROM
+              "Orders"
+            WHERE
+              status NOT IN ('shipped')
+            GROUP BY 1
+            ORDER BY 2 DESC
+            LIMIT 2
+          ) top_orders
+        ON
+          "Orders".status = top_orders.status
+        GROUP BY 1
+        `;
+
+      const res = await connection.query(query);
+      // Expect only top statuses 2 by total amount, with shipped filtered out: processed and new
+      expect(res.rows).toMatchSnapshot('join grouped with filter');
+    });
+
     test('where segment is false', async () => {
       const query =
         'SELECT value AS val, * FROM "SegmentTest" WHERE segment_eq_1 IS FALSE ORDER BY value;';

rust/cubesql/cubeclient/.openapi-generator/FILES

@@ -16,6 +16,7 @@ src/models/v1_load_request_query_filter_base.rs
 src/models/v1_load_request_query_filter_item.rs
 src/models/v1_load_request_query_filter_logical_and.rs
 src/models/v1_load_request_query_filter_logical_or.rs
+src/models/v1_load_request_query_join_subquery.rs
 src/models/v1_load_request_query_time_dimension.rs
 src/models/v1_load_response.rs
 src/models/v1_load_result.rs

rust/cubesql/cubeclient/src/models/mod.rs

@@ -28,6 +28,8 @@ pub mod v1_load_request_query_filter_logical_and;
 pub use self::v1_load_request_query_filter_logical_and::V1LoadRequestQueryFilterLogicalAnd;
 pub mod v1_load_request_query_filter_logical_or;
 pub use self::v1_load_request_query_filter_logical_or::V1LoadRequestQueryFilterLogicalOr;
+pub mod v1_load_request_query_join_subquery;
+pub use self::v1_load_request_query_join_subquery::V1LoadRequestQueryJoinSubquery;
 pub mod v1_load_request_query_time_dimension;
 pub use self::v1_load_request_query_time_dimension::V1LoadRequestQueryTimeDimension;
 pub mod v1_load_response;

rust/cubesql/cubeclient/src/models/v1_load_request_query.rs

@@ -28,6 +28,8 @@ pub struct V1LoadRequestQuery {
     pub filters: Option<Vec<crate::models::V1LoadRequestQueryFilterItem>>,
     #[serde(rename = "ungrouped", skip_serializing_if = "Option::is_none")]
     pub ungrouped: Option<bool>,
+    #[serde(rename = "subqueryJoins", skip_serializing_if = "Option::is_none")]
+    pub subquery_joins: Option<Vec<crate::models::V1LoadRequestQueryJoinSubquery>>,
 }
 
 impl V1LoadRequestQuery {
@@ -42,6 +44,7 @@ impl V1LoadRequestQuery {
             offset: None,
             filters: None,
             ungrouped: None,
+            subquery_joins: None,
         }
     }
 }