@finos/legend-query-builder 4.16.6 → 4.16.8

package/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[],

package/src/stores/QueryBuilderStateBuilder.ts

@@ -101,6 +101,11 @@ import {
   isTypedProjectionExpression,
   processTypedTDSProjectExpression,
 } from './fetch-structure/tds/projection/QueryBuilderTypedProjectionStateBuilder.js';
+import {
+  isTypedGroupByExpression,
+  processTypedAggregationColSpec,
+  processTypedGroupByExpression,
+} from './fetch-structure/tds/aggregation/QueryBuilderTypedAggregationStateBuilder.js';
 
 const processGetAllExpression = (
   expression: SimpleFunctionExpression,
@@ -740,16 +745,24 @@ export class QueryBuilderValueSpecificationProcessor
       );
       return;
     } else if (
-      matchFunctionName(
-        functionName,
+      matchFunctionName(functionName, [
         QUERY_BUILDER_SUPPORTED_FUNCTIONS.TDS_GROUP_BY,
-      )
+        QUERY_BUILDER_SUPPORTED_FUNCTIONS.RELATION_GROUP_BY,
+      ])
     ) {
-      processTDSGroupByExpression(
-        valueSpecification,
-        this.queryBuilderState,
-        this.parentLambda,
-      );
+      if (isTypedGroupByExpression(valueSpecification)) {
+        processTypedGroupByExpression(
+          valueSpecification,
+          this.queryBuilderState,
+          this.parentLambda,
+        );
+      } else {
+        processTDSGroupByExpression(
+          valueSpecification,
+          this.queryBuilderState,
+          this.parentLambda,
+        );
+      }
       return;
     } else if (
       matchFunctionName(functionName, QUERY_BUILDER_SUPPORTED_FUNCTIONS.TDS_AGG)
@@ -1002,6 +1015,25 @@ export class QueryBuilderValueSpecificationProcessor
         }
       });
 
+      return;
+    } else if (
+      matchFunctionName(this.parentExpression.functionName, [
+        QUERY_BUILDER_SUPPORTED_FUNCTIONS.RELATION_GROUP_BY,
+      ])
+    ) {
+      const spec = valueSpecification.values;
+      assertTrue(
+        spec.length === 1,
+        `Can't process col spec array instance: value expected to be of size 1`,
+      );
+      guaranteeNonNullable(spec[0]).colSpecs.forEach((col) => {
+        processTypedAggregationColSpec(
+          col,
+          this.parentExpression,
+          this.queryBuilderState,
+        );
+      });
+
       return;
     }
     throw new UnsupportedOperationError(

package/src/stores/fetch-structure/tds/aggregation/QueryBuilderRelationAggregationValueSpecBuilder.ts

@@ -0,0 +1,196 @@
+/**
+ * Copyright (c) 2020-present, Goldman Sachs
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import {
+  ColSpecArray,
+  type ValueSpecification,
+  extractElementNameFromPath,
+  SimpleFunctionExpression,
+  ColSpecArrayInstance,
+  Multiplicity,
+  ColSpec,
+  GenericType,
+  GenericTypeExplicitReference,
+  Relation,
+  RelationColumn,
+  RelationType,
+  LambdaFunctionInstanceValue,
+  AbstractPropertyExpression,
+  PropertyExplicitReference,
+  Property,
+  VariableExpression,
+} from '@finos/legend-graph';
+import { QUERY_BUILDER_SUPPORTED_FUNCTIONS } from '../../../../graph/QueryBuilderMetaModelConst.js';
+import type { QueryBuilderTDSState } from '../QueryBuilderTDSState.js';
+import { buildGenericLambdaFunctionInstanceValue } from '../../../QueryBuilderValueSpecificationHelper.js';
+import {
+  guaranteeNonNullable,
+  guaranteeType,
+  UnsupportedOperationError,
+} from '@finos/legend-shared';
+import { DEFAULT_LAMBDA_VARIABLE_NAME } from '@finos/legend-data-cube';
+
+export const buildRelationAggregation = (
+  precedingExpression: ValueSpecification,
+  tdsState: QueryBuilderTDSState,
+): SimpleFunctionExpression => {
+  // Verify that preceding expression is a relation project
+  const projectFunction = guaranteeType(
+    precedingExpression,
+    SimpleFunctionExpression,
+  );
+  if (
+    projectFunction.functionName !==
+    extractElementNameFromPath(
+      QUERY_BUILDER_SUPPORTED_FUNCTIONS.RELATION_PROJECT,
+    )
+  ) {
+    throw new UnsupportedOperationError(
+      `Can't build relation groupBy() expression: previous expression must be project() column expression`,
+    );
+  }
+
+  // Build groupBy() expression
+  const groupByFunction = new SimpleFunctionExpression(
+    extractElementNameFromPath(
+      QUERY_BUILDER_SUPPORTED_FUNCTIONS.RELATION_GROUP_BY,
+    ),
+  );
+  const queryBuilderState = tdsState.queryBuilderState;
+
+  // Build normal (grouped) columns
+  const groupByCols = new ColSpecArrayInstance(Multiplicity.ONE, undefined);
+  const groupByColSpecArray = new ColSpecArray();
+  groupByCols.values = [groupByColSpecArray];
+
+  // Build aggregation columns
+  const aggregationCols = new ColSpecArrayInstance(Multiplicity.ONE, undefined);
+  const aggregationColSpecArray = new ColSpecArray();
+  aggregationCols.values = [aggregationColSpecArray];
+
+  // Build relation return type
+  const relationType = new RelationType(RelationType.ID);
+
+  // Add normal (non-aggregated) columns to groupByCols
+  tdsState.projectionColumns
+    .filter(
+      (projectionColumnState) =>
+        !tdsState.aggregationState.columns.some(
+          (aggregationColumnState) =>
+            aggregationColumnState.projectionColumnState ===
+            projectionColumnState,
+        ),
+    )
+    .forEach((projectionColumnState) => {
+      // Create and add column to groupByColSpecArray
+      const colSpec = new ColSpec();
+      colSpec.name = projectionColumnState.columnName;
+      groupByColSpecArray.colSpecs.push(colSpec);
+
+      // Add column return type to relationType
+      const returnType = guaranteeNonNullable(
+        projectionColumnState.getColumnType(),
+        `Can't create value spec for projection column ${projectionColumnState.columnName}. Missing type.`,
+      );
+      relationType.columns.push(
+        new RelationColumn(projectionColumnState.columnName, returnType),
+      );
+    });
+
+  // Add aggregation columns to aggregationCols
+  tdsState.aggregationState.columns.forEach((aggregationColumnState) => {
+    // Create and add column to aggregationColSpecArray
+    const colSpec = new ColSpec();
+    colSpec.name = aggregationColumnState.columnName;
+    aggregationColSpecArray.colSpecs.push(colSpec);
+
+    // Build map function (function1)
+    // First, get the ColSpec from the preceding projection so we can get the name of the projected column.
+    const projectionColSpec = guaranteeNonNullable(
+      guaranteeNonNullable(
+        (projectFunction.parametersValues[1] as ColSpecArrayInstance).values[0],
+        'Could not find ColSpec array in project() function first parameter',
+      ).colSpecs.find(
+        (_colSpec) => _colSpec.name === aggregationColumnState.columnName,
+      ),
+      `Could not find projected column matching aggregation column '${aggregationColumnState.columnName}'`,
+    );
+    const projectedPropertyExpression = guaranteeType(
+      guaranteeType(projectionColSpec.function1, LambdaFunctionInstanceValue)
+        .values[0]?.expressionSequence[0],
+      AbstractPropertyExpression,
+    );
+    const projectedProperty = guaranteeType(
+      projectedPropertyExpression.func.value,
+      Property,
+    );
+    // Second, build a new AbstractPropertyExpression for our map function
+    const newPropertyExpression = new AbstractPropertyExpression('');
+    newPropertyExpression.func = PropertyExplicitReference.create(
+      new Property(
+        projectionColSpec.name,
+        projectedProperty.multiplicity,
+        projectedProperty.genericType,
+        projectedProperty._OWNER,
+      ),
+    );
+    newPropertyExpression.parametersValues = [
+      new VariableExpression(DEFAULT_LAMBDA_VARIABLE_NAME, Multiplicity.ONE),
+    ];
+    const mapLambda = buildGenericLambdaFunctionInstanceValue(
+      DEFAULT_LAMBDA_VARIABLE_NAME,
+      [newPropertyExpression],
+      queryBuilderState.graphManagerState.graph,
+    );
+    colSpec.function1 = mapLambda;
+
+    // Reduce function (function2)
+    const reduceLambda = buildGenericLambdaFunctionInstanceValue(
+      aggregationColumnState.lambdaParameterName,
+      [
+        aggregationColumnState.operator.buildAggregateExpressionFromState(
+          aggregationColumnState,
+        ),
+      ],
+      aggregationColumnState.aggregationState.tdsState.queryBuilderState
+        .graphManagerState.graph,
+    );
+    colSpec.function2 = reduceLambda;
+
+    // Add column return type to relationType
+    const returnType = guaranteeNonNullable(
+      aggregationColumnState.getColumnType(),
+      `Can't create value spec for aggregation column ${aggregationColumnState.columnName}. Missing type.`,
+    );
+    relationType.columns.push(
+      new RelationColumn(aggregationColumnState.columnName, returnType),
+    );
+  });
+
+  groupByFunction.parametersValues = [
+    precedingExpression,
+    groupByCols,
+    aggregationCols,
+  ];
+  const relationGenericType = new GenericType(Relation.INSTANCE);
+  const relationTypeGenericType = new GenericType(relationType);
+  relationGenericType.typeArguments = [
+    GenericTypeExplicitReference.create(relationTypeGenericType),
+  ];
+  groupByFunction.genericType =
+    GenericTypeExplicitReference.create(relationGenericType);
+  return groupByFunction;
+};