Add support for TDS v2 (relation) groupBy in QueryBuilder (#3892)

* feat: Add roundtrip test for new groupBy * feat: Ensure ColSpec function2 is serialized * fix: Fix typo in QueryBuilderRelationProjectValueSpecBuilder * feat: Handle converting query builder state to typed groupBy protocol * feat: Start adding support to convert from relation protocol to aggregation state * feat: Support converting TDSv2 groupBy protocol to state and fix handling nested properties * fix: Fix test expected value * fix: Remove console.log statements * fix: Remove commented out code * fix: Fix copyright year * fix: Fix missing match function name * fix: Add changeset file * fix: Fix lint and copyright issues * feat: Add additional test cases to QueryRoundtripGrammar tests * refactor: Clean up V1_buildTypedGroupByFunctionExpression * feat: Remove unused set variable * fix: Tmp fix of expression genericType in V1_buildTypedGroupByFunctionExpression * feat: Add getNumericAggregateOperatorReturnType helper function * feat: Fix setting return type in V1_buildTypedGroupByFunctionExpression * feat: Use getNumericAggregateOperatorReturnType function in aggregate operators * feat: Clean up code in V1_QueryValueSpecificationBuilderHelper * feat: Clean up and add comments to buildRelationAggregation function * feat: Remove options param from buildRelationAggregation function * refactor: Update some variable names in V1_buildTypedGroupByFunctionExpression * refactor: Clean up and add comments to QueryBuilderTypedAggregationStateBuilder functions * feat: Add QueryBuilderLambdaRoundtrip test cases * fix: Remove extra param in function call * fix: Update changeset files * feat: Don't use getNumericAggregateOperatorReturnType function in operator classes * refactor: Update error message * refactor: Check for fully qualified function name in V1_buildGroupByFunctionExpression * feat: Make isTypedProjectionExpression check for fully qualified function name * feat: Make isTypedGroupByExpression check for fully qualified function name * refactor: Set groupBy() expression return type in processTypedAggregationColSpec function * refactor: Update comment * refactor: Update comment
finos · Feb 18, 2025 · db7ebfb · db7ebfb
1 parent 2fa1cbb
commit db7ebfb
Show file tree

Hide file tree

Showing 18 changed files with 1,177 additions and 28 deletions.
diff --git a/.changeset/five-tomatoes-guess.md b/.changeset/five-tomatoes-guess.md
@@ -0,0 +1,5 @@
+---
+'@finos/legend-graph': patch
+---
+
+Ensure function2 is handled in ColSpec ValueSpecification
diff --git a/.changeset/light-icons-sparkle.md b/.changeset/light-icons-sparkle.md
@@ -0,0 +1,5 @@
+---
+'@finos/legend-query-builder': patch
+---
+
+Add support for TDS v2 (relation) groupBy in QueryBuilder
diff --git a/packages/legend-graph/src/graph-manager/action/changeDetection/ValueSpecificationObserver.ts b/packages/legend-graph/src/graph-manager/action/changeDetection/ValueSpecificationObserver.ts
@@ -547,5 +547,8 @@ function _observe_ColSpec(
   if (metamodel.function1) {
     observe_ValueSpecification(metamodel.function1, context);
   }
+  if (metamodel.function2) {
+    observe_ValueSpecification(metamodel.function2, context);
+  }
   return metamodel;
 }
diff --git a/...anager/protocol/pure/v1/transformation/pureGraph/from/V1_ValueSpecificationTransformer.ts b/...anager/protocol/pure/v1/transformation/pureGraph/from/V1_ValueSpecificationTransformer.ts
@@ -419,9 +419,20 @@ class V1_ValueSpecificationTransformer
           this.useAppliedFunction,
         ),
       );
+      const fun2 = col.function2?.accept_ValueSpecificationVisitor(
+        new V1_ValueSpecificationTransformer(
+          this.inScope,
+          this.open,
+          this.isParameter,
+          this.useAppliedFunction,
+        ),
+      );
       if (fun1) {
         colProtocol.function1 = guaranteeType(fun1, V1_Lambda);
       }
+      if (fun2) {
+        colProtocol.function2 = guaranteeType(fun2, V1_Lambda);
+      }
       return colProtocol;
     });
     classInstance.value = colSpecArray;
@@ -444,9 +455,20 @@ class V1_ValueSpecificationTransformer
         this.useAppliedFunction,
       ),
     );
+    const fun2 = val.function2?.accept_ValueSpecificationVisitor(
+      new V1_ValueSpecificationTransformer(
+        this.inScope,
+        this.open,
+        this.isParameter,
+        this.useAppliedFunction,
+      ),
+    );
     if (fun1) {
       colProtocol.function1 = guaranteeType(fun1, V1_Lambda);
     }
+    if (fun2) {
+      colProtocol.function2 = guaranteeType(fun1, V1_Lambda);
+    }
     classInstance.value = colProtocol;
     return classInstance;
   }

diff --git a/...es/legend-graph/src/graph/metamodel/pure/valueSpecification/RelationValueSpecification.ts b/...es/legend-graph/src/graph/metamodel/pure/valueSpecification/RelationValueSpecification.ts
@@ -38,6 +38,7 @@ export class ColSpec implements Hashable {
       CORE_HASH_STRUCTURE.RELATION_COL_SPEC,
       this.name,
       this.function1 ?? '',
+      this.function2 ?? '',
     ]);
   }
 }

diff --git a/...d-manual-tests/src/__tests__/query-builder/QueryRoundtripGrammar.engine-roundtrip-test.ts b/...d-manual-tests/src/__tests__/query-builder/QueryRoundtripGrammar.engine-roundtrip-test.ts
@@ -124,6 +124,7 @@ const TEST_CASES: QueryTestCase[] = [
     queryGrammar:
       "|showcase::northwind::model::crm::Customer.all()->filter(x|$x.companyTitle == 'company title')->project(~['Company Name':x|$x.companyName, 'Company Title':x|$x.companyTitle])->filter(row|$row.'Company Name' == 'company name')",
   },
+  // result modifier
   {
     testName: '[LEGACY] Result Modifier: Sort',
     model: 'Northwind',
@@ -153,13 +154,42 @@ const TEST_CASES: QueryTestCase[] = [
     convertedRelation:
       "|showcase::northwind::model::crm::Customer.all()->filter(x|$x.companyTitle == 'company title')->project(~['Company Name':x|$x.companyName, 'Company Title':x|$x.companyTitle])->filter(row|$row.'Company Name' == 'company name')",
   },
-  //aggregation
+  // aggregation
   {
-    testName: '[AGGREGATION] Simple wavg query',
+    testName: '[LEGACY AGGREGATION] Simple wavg query',
     model: 'Northwind',
     queryGrammar:
       "|showcase::northwind::model::OrderLineItem.all()->groupBy([], [agg(x|$x.quantity->meta::pure::functions::math::wavgUtility::wavgRowMapper($x.unitPrice), y|$y->wavg())], ['Quantity (wavg)'])",
   },
+  {
+    testName: '[LEGACY AGGREGATION] Simple group by count query',
+    model: 'Northwind',
+    queryGrammar:
+      "|showcase::northwind::model::Order.all()->groupBy([x|$x.shipToName], [agg(x|$x.id,x|$x->count())], ['Ship To Name','Id (count)'])",
+    convertedRelation:
+      "|showcase::northwind::model::Order.all()->project(~['Ship To Name':x|$x.shipToName, 'Id (count)':x|$x.id])->groupBy(~['Ship To Name'], ~['Id (count)':x|$x.'Id (count)':x|$x->count()])",
+  },
+  {
+    testName: '[AGGREGATION] Group by count query with nested property',
+    model: 'Northwind',
+    queryGrammar:
+      "|showcase::northwind::model::Order.all()->project(~['Ship To Name':x|$x.shipToName, 'Customer/Id (count)':x|$x.customer.id])->groupBy(~['Ship To Name'], ~['Customer/Id (count)':x|$x.'Customer/Id (count)':x|$x->count()])",
+  },
+  {
+    testName:
+      '[AGGREGATION] Group by count query with pre-filter and post-filter',
+    model: 'Northwind',
+    queryGrammar:
+      "|showcase::northwind::model::Order.all()->filter(x|$x.shipToName == 'test')->project(~['Ship To Name':x|$x.shipToName, 'Id (count)':x|$x.id])->groupBy(~['Ship To Name'], ~['Id (count)':x|$x.'Id (count)':x|$x->count()])->filter(row|$row.'Id (count)' >= 5)",
+  },
+  {
+    testName:
+      '[AGGREGATION] Group by count query with post-filter before groupBy',
+    model: 'Northwind',
+    queryGrammar:
+      "|showcase::northwind::model::Order.all()->project(~['Ship To Name':x|$x.shipToName, 'Id (count)':x|$x.id])->filter(row|$row.'Id (count)' >= 5)->groupBy(~['Ship To Name'], ~['Id (count)':x|$x.'Id (count)':x|$x->count()])",
+    isUnsupported: true,
+  },
 ];
 
 const globalGraphManagerStates = new Map<string, GraphManagerState>();

diff --git a/...query-builder/src/graph-manager/protocol/pure/QueryBuilder_PureProtocolProcessorPlugin.ts b/...query-builder/src/graph-manager/protocol/pure/QueryBuilder_PureProtocolProcessorPlugin.ts
@@ -21,11 +21,11 @@ import {
   V1_buildGetAllFunctionExpression,
   V1_buildGetAllVersionsFunctionExpression,
   V1_buildGetAllVersionsInRangeFunctionExpression,
-  V1_buildGroupByFunctionExpression,
   V1_buildOLAPGroupByFunctionExpression,
   V1_buildProjectFunctionExpression,
   V1_buildSubTypePropertyExpressionTypeInference,
   V1_buildWatermarkFunctionExpression,
+  V1_buildGroupByFunctionExpression,
 } from './v1/V1_QueryValueSpecificationBuilderHelper.js';
 import {
   type V1_GraphBuilderContext,
@@ -159,10 +159,10 @@ export class QueryBuilder_PureProtocolProcessorPlugin extends PureProtocolProces
             processingContext,
           );
         } else if (
-          matchFunctionName(
-            functionName,
+          matchFunctionName(functionName, [
             QUERY_BUILDER_SUPPORTED_FUNCTIONS.TDS_GROUP_BY,
-          )
+            QUERY_BUILDER_SUPPORTED_FUNCTIONS.RELATION_GROUP_BY,
+          ])
         ) {
           return V1_buildGroupByFunctionExpression(
             functionName,

diff --git a/...ery-builder/src/graph-manager/protocol/pure/v1/V1_QueryValueSpecificationBuilderHelper.ts b/...ery-builder/src/graph-manager/protocol/pure/v1/V1_QueryValueSpecificationBuilderHelper.ts
@@ -30,7 +30,7 @@ import {
   type V1_ProcessingContext,
   type ValueSpecification,
   type Type,
-  type SimpleFunctionExpression,
+  SimpleFunctionExpression,
   V1_ValueSpecification,
   V1_buildBaseSimpleFunctionExpression,
   V1_buildGenericFunctionExpression,
@@ -997,7 +997,7 @@ export const V1_buildProjectFunctionExpression = (
   return expression;
 };
 
-export const V1_buildGroupByFunctionExpression = (
+export const V1_buildTDSGroupByFunctionExpression = (
   functionName: string,
   parameters: V1_ValueSpecification[],
   openVariables: string[],
@@ -1140,6 +1140,227 @@ export const V1_buildGroupByFunctionExpression = (
   return expression;
 };
 
+export const V1_buildTypedGroupByFunctionExpression = (
+  functionName: string,
+  parameters: V1_ValueSpecification[],
+  openVariables: string[],
+  compileContext: V1_GraphBuilderContext,
+  processingContext: V1_ProcessingContext,
+): SimpleFunctionExpression => {
+  assertTrue(
+    parameters.length === 3,
+    `Can't build relation groupBy() expression: groupBy() expects 2 arguments`,
+  );
+
+  const precedingExpression = (
+    parameters[0] as V1_ValueSpecification
+  ).accept_ValueSpecificationVisitor(
+    new V1_ValueSpecificationBuilder(
+      compileContext,
+      processingContext,
+      openVariables,
+    ),
+  );
+  assertNonNullable(
+    precedingExpression.genericType,
+    `Can't build relation groupBy() expression: preceding expression return type is missing`,
+  );
+
+  // Assert that preceding function is a project() function
+  const projectFunction = guaranteeType(
+    precedingExpression,
+    SimpleFunctionExpression,
+  );
+  if (
+    projectFunction.functionName !==
+    extractElementNameFromPath(
+      QUERY_BUILDER_SUPPORTED_FUNCTIONS.RELATION_PROJECT,
+    )
+  ) {
+    throw new UnsupportedOperationError(
+      `Can't build relation groupBy() expression: preceding expression must be project() column expression`,
+    );
+  }
+
+  // Get normal (grouped) columns
+  const columnsClassInstance = parameters[1];
+  assertType(
+    columnsClassInstance,
+    V1_ClassInstance,
+    `Can't build relation groupBy() expression: groupBy() expects argument #1 to be a ClassInstance`,
+  );
+  const columnExpressions = guaranteeType(
+    columnsClassInstance.value,
+    V1_ColSpecArray,
+    `Can't build relation groupBy() expression: groupBy() expects argument #1 to hold col spec array instance value`,
+  );
+
+  // Get aggregation columns
+  const aggregationColumnsClassInstance = parameters[2];
+  assertType(
+    aggregationColumnsClassInstance,
+    V1_ClassInstance,
+    `Can't build groupBy() expression: groupBy() expects argument #2 to be a ClassInstance`,
+  );
+  const aggregationExpressions = guaranteeType(
+    aggregationColumnsClassInstance.value,
+    V1_ColSpecArray,
+    `Can't build relation groupBy() expression: groupBy() expects argument #2 to hold col spec array instance value`,
+  );
+
+  // Make sure top-level lambdas have their lambda parameter types set properly
+  const topLevelLambdaParameters: V1_Variable[] =
+    aggregationExpressions.colSpecs
+      .map((colSpec) => [colSpec.function1, colSpec.function2])
+      .flat()
+      .filter(isNonNullable)
+      .map((value) => value.parameters)
+      .flat();
+  topLevelLambdaParameters.forEach((variable) => {
+    if (!variable.genericType) {
+      const variableExpression = new VariableExpression(
+        variable.name,
+        precedingExpression.multiplicity,
+      );
+      variableExpression.genericType = precedingExpression.genericType;
+      processingContext.addInferredVariables(variable.name, variableExpression);
+    }
+  });
+
+  const projectRelationReturnType = guaranteeType(
+    projectFunction.genericType?.value.typeArguments?.[0]?.value.rawType,
+    RelationType,
+    `Can't build relation groupBy() expression: project() function does not return a relation`,
+  );
+
+  // build column expressions
+  const processedColumnExpressions = new ColSpecArrayInstance(Multiplicity.ONE);
+  const processedColSpecArray = new ColSpecArray();
+  processedColumnExpressions.values = [processedColSpecArray];
+  const relationType = new RelationType(RelationType.ID);
+  processedColSpecArray.colSpecs = columnExpressions.colSpecs.map((colSpec) => {
+    const pColSpec = new ColSpec();
+    pColSpec.name = colSpec.name;
+
+    // Add the column using the return type of the preceding project
+    const column = projectRelationReturnType.columns.find(
+      (_column) => _column.name === colSpec.name,
+    );
+    if (column) {
+      relationType.columns.push(column);
+    }
+    return pColSpec;
+  });
+
+  // build aggregation column expressions
+  const processedAggregationExpressions = new ColSpecArrayInstance(
+    Multiplicity.ONE,
+  );
+  const processedAggregationColSpecArray = new ColSpecArray();
+  processedAggregationExpressions.values = [processedAggregationColSpecArray];
+  processedAggregationColSpecArray.colSpecs =
+    aggregationExpressions.colSpecs.map((colSpec) => {
+      const pColSpec = new ColSpec();
+      pColSpec.name = colSpec.name;
+
+      // Build the map lambda
+      const mapFunction = guaranteeType(
+        colSpec.function1,
+        V1_Lambda,
+        `Can't build relation col spec() expression: expects function1 to be a lambda`,
+      );
+      const mapLambda: ValueSpecification = buildProjectionColumnLambda(
+        mapFunction,
+        openVariables,
+        compileContext,
+        processingContext,
+      );
+      pColSpec.function1 = mapLambda;
+
+      // Build the reduce lambda
+      const reduceFunction = guaranteeType(
+        colSpec.function2,
+        V1_Lambda,
+        `Can't build relation col spec() expression: expects function2 to be a lambda`,
+      );
+      const reduceLambda = guaranteeType(
+        reduceFunction.accept_ValueSpecificationVisitor(
+          new V1_ValueSpecificationBuilder(
+            compileContext,
+            processingContext,
+            openVariables,
+          ),
+        ),
+        LambdaFunctionInstanceValue,
+        `Can't build relation col spec() expression: expected aggregation function to be a lambda`,
+      );
+      pColSpec.function2 = reduceLambda;
+
+      // For now, we temporarily set the return type of the column in the groupBy() expression to be
+      // the same as the return type of the column in the preceding project() expression.
+      // The actual return type for the groupBy() expression will be determined when we process/build the graph.
+      const returnType = projectRelationReturnType.columns.find(
+        (_column) => _column.name === colSpec.name,
+      )?.type;
+      if (returnType) {
+        relationType.columns.push(new RelationColumn(colSpec.name, returnType));
+      } else {
+        throw new UnsupportedOperationError(
+          `Unable to find projected column with name ${colSpec.name}`,
+        );
+      }
+
+      return pColSpec;
+    });
+
+  const expression = V1_buildBaseSimpleFunctionExpression(
+    [
+      precedingExpression,
+      processedColumnExpressions,
+      processedAggregationExpressions,
+    ],
+    functionName,
+    compileContext,
+  );
+  const relationGenericType = new GenericType(Relation.INSTANCE);
+  const relationTypeGenericType = new GenericType(relationType);
+  relationGenericType.typeArguments = [
+    GenericTypeExplicitReference.create(relationTypeGenericType),
+  ];
+  expression.genericType =
+    GenericTypeExplicitReference.create(relationGenericType);
+
+  return expression;
+};
+
+export const V1_buildGroupByFunctionExpression = (
+  functionName: string,
+  parameters: V1_ValueSpecification[],
+  openVariables: string[],
+  compileContext: V1_GraphBuilderContext,
+  processingContext: V1_ProcessingContext,
+): SimpleFunctionExpression => {
+  if (
+    functionName === QUERY_BUILDER_SUPPORTED_FUNCTIONS.RELATION_GROUP_BY ||
+    parameters.length === 3
+  ) {
+    return V1_buildTypedGroupByFunctionExpression(
+      functionName,
+      parameters,
+      openVariables,
+      compileContext,
+      processingContext,
+    );
+  }
+  return V1_buildTDSGroupByFunctionExpression(
+    functionName,
+    parameters,
+    openVariables,
+    compileContext,
+    processingContext,
+  );
+};
+
 export const V1_buildWatermarkFunctionExpression = (
   functionName: string,
   parameters: V1_ValueSpecification[],