@finos/legend-query-builder 4.14.26 → 4.14.28

package/src/components/result/tds/QueryBuilderTDSGridResult.tsx

@@ -14,11 +14,10 @@
  * limitations under the License.
  */
 
-import { clsx } from '@finos/legend-art';
 import { observer } from 'mobx-react-lite';
 import type { QueryBuilderState } from '../../../stores/QueryBuilderState.js';
 import { PRIMITIVE_TYPE, type TDSExecutionResult } from '@finos/legend-graph';
-import { useState, useCallback } from 'react';
+import { useState, useCallback, useEffect } from 'react';
 import {
   DataGrid,
   type DataGridApi,
@@ -28,23 +27,45 @@ import {
   type DataGridGetContextMenuItemsParams,
   type DataGridIRowNode,
   type DataGridMenuItemDef,
+  type DataGridIAggFuncParams,
 } from '@finos/legend-lego/data-grid';
 import {
   getRowDataFromExecutionResult,
   type IQueryRendererParamsWithGridType,
   filterByOrOutValues,
 } from './QueryBuilderTDSResultShared.js';
-import type {
-  QueryBuilderResultState,
-  QueryBuilderTDSResultCellData,
-  QueryBuilderTDSResultCellDataType,
-  QueryBuilderTDSRowDataType,
+import {
+  type QueryBuilderResultState,
+  QueryBuilderResultWavgAggregationState,
+  type QueryBuilderTDSResultCellData,
+  type QueryBuilderTDSResultCellDataType,
+  type QueryBuilderTDSRowDataType,
 } from '../../../stores/QueryBuilderResultState.js';
 import { QueryBuilderTDSState } from '../../../stores/fetch-structure/tds/QueryBuilderTDSState.js';
 import { DEFAULT_LOCALE } from '../../../graph-manager/QueryBuilderConst.js';
-import { isNumber, isString, isValidURL } from '@finos/legend-shared';
+import {
+  assertErrorThrown,
+  isNumber,
+  isString,
+  isValidURL,
+} from '@finos/legend-shared';
 import { useApplicationStore } from '@finos/legend-application';
 import { QUERY_BUILDER_TEST_ID } from '../../../__lib__/QueryBuilderTesting.js';
+import {
+  clsx,
+  Modal,
+  ModalBody,
+  ModalHeader,
+  ModalFooter,
+  ModalFooterButton,
+  Dialog,
+  CustomSelectorInput,
+} from '@finos/legend-art';
+
+export const enum QueryBuilderDataGridCustomAggregationFunction {
+  wavg = 'wavg',
+  WAVG = 'WAVG',
+}
 
 const getAggregationTDSColumnCustomizations = (
   result: TDSExecutionResult,
@@ -71,7 +92,7 @@ const getAggregationTDSColumnCustomizations = (
     case PRIMITIVE_TYPE.FLOAT:
       return {
         filter: 'agNumberColumnFilter',
-        allowedAggFuncs: ['count', 'sum', 'max', 'min', 'avg'],
+        allowedAggFuncs: ['count', 'sum', 'max', 'min', 'avg', 'wavg'],
       };
     default:
       return {
@@ -222,16 +243,19 @@ export const QueryBuilderTDSGridResult = observer(
   }) => {
     const { executionResult, queryBuilderState } = props;
     const applicationStore = useApplicationStore();
+    const darkMode =
+      !applicationStore.layoutService.TEMPORARY__isLightColorThemeEnabled;
     const [columnAPi, setColumnApi] = useState<DataGridColumnApi | undefined>(
       undefined,
     );
+    const [aggFuncParams, setAggFuncParams] = useState<
+      DataGridIAggFuncParams | undefined
+    >(undefined);
     const resultState = queryBuilderState.resultState;
     const isLocalModeEnabled = queryBuilderState.isLocalModeEnabled;
     const colDefs = isLocalModeEnabled
       ? getLocalColDefs(executionResult, resultState)
       : getColDefs(executionResult, resultState);
-    const darkMode =
-      !applicationStore.layoutService.TEMPORARY__isLightColorThemeEnabled;
 
     const onSaveGridColumnState = (): void => {
       if (!columnAPi) {
@@ -242,6 +266,10 @@ export const QueryBuilderTDSGridResult = observer(
         isPivotModeEnabled: columnAPi.isPivotMode(),
         isLocalModeEnabled: true,
         previewLimit: resultState.previewLimit,
+        ...(resultState.wavgAggregationState?.weightedColumnIdPairs && {
+          weightedColumnPairs:
+            resultState.wavgAggregationState.weightedColumnIdPairs,
+        }),
       });
     };
 
@@ -336,6 +364,103 @@ export const QueryBuilderTDSGridResult = observer(
       ],
     );
 
+    const weightedColumnOptions = columnAPi
+      ?.getColumns()
+      ?.filter((c) => c.getColDef().cellDataType === 'number')
+      .map((col) => ({
+        label: col.getColId(),
+        value: col.getColId(),
+      }));
+
+    const selectedWeightedColumn =
+      aggFuncParams?.colDef.field &&
+      resultState.wavgAggregationState?.weightedColumnIdPairs.get(
+        aggFuncParams.colDef.field,
+      )
+        ? {
+            label: resultState.wavgAggregationState.weightedColumnIdPairs.get(
+              aggFuncParams.colDef.field,
+            ),
+            value: resultState.wavgAggregationState.weightedColumnIdPairs.get(
+              aggFuncParams.colDef.field,
+            ),
+          }
+        : null;
+
+    const onWeightedColumnOptionChange = async (
+      option: { label: string; value: string } | null,
+    ): Promise<void> => {
+      if (aggFuncParams?.colDef.field && option?.value) {
+        resultState.wavgAggregationState?.addWeightedColumnIdPair(
+          aggFuncParams.colDef.field,
+          option.value,
+        );
+      }
+    };
+
+    const weightedAverage = (param: DataGridIAggFuncParams): void => {
+      if (param.colDef.field) {
+        if (!resultState.wavgAggregationState) {
+          resultState.setWavgAggregationState(
+            new QueryBuilderResultWavgAggregationState(),
+          );
+        }
+        resultState.wavgAggregationState?.addWeightedColumnIdPair(
+          param.colDef.field,
+          param.colDef.field,
+        );
+        resultState.wavgAggregationState?.setIsApplyingWavg(true);
+        setAggFuncParams(param);
+      } else {
+        applicationStore.notificationService.notifyError(
+          'The id of this column can`t be retrieved to perform weighted average',
+        );
+      }
+    };
+
+    const weightedAverageHelper = (param: DataGridIAggFuncParams): number => {
+      try {
+        const column = param.colDef.field;
+        if (column) {
+          const weightedColumnId =
+            resultState.wavgAggregationState?.weightedColumnIdPairs.get(column);
+          if (weightedColumnId) {
+            const weightedColumnSum = param.rowNode.allLeafChildren
+              .map((node) => node.data[weightedColumnId])
+              .reduce((a, b) => a + b) as number;
+            const weightedColumnMultiply = param.rowNode.allLeafChildren
+              .map((node) => node.data[weightedColumnId] * node.data[column])
+              .reduce((a, b) => a + b);
+            if (weightedColumnSum !== 0) {
+              onSaveGridColumnState();
+              return weightedColumnMultiply / weightedColumnSum;
+            } else {
+              applicationStore.notificationService.notifyError(
+                'The weighted column sum is 0',
+              );
+            }
+          } else {
+            applicationStore.notificationService.notifyError(
+              'The weighted column Id is not defined',
+            );
+          }
+        }
+      } catch (error) {
+        assertErrorThrown(error);
+        applicationStore.notificationService.notifyError(error);
+      }
+      return -1;
+    };
+
+    useEffect(() => {
+      if (aggFuncParams) {
+        aggFuncParams.columnApi.setColumnAggFunc(
+          aggFuncParams.colDef.field!,
+          QueryBuilderDataGridCustomAggregationFunction.WAVG,
+        );
+      }
+    }, [resultState.wavgAggregationState, aggFuncParams]);
+
     return (
       <div
         data-testid={QUERY_BUILDER_TEST_ID.QUERY_BUILDER_RESULT_VALUES_TDS}
@@ -372,6 +497,10 @@ export const QueryBuilderTDSGridResult = observer(
               suppressFieldDotNotation={true}
               suppressContextMenu={false}
               columnDefs={colDefs}
+              aggFuncs={{
+                wavg: weightedAverage,
+                WAVG: weightedAverageHelper,
+              }}
               sideBar={['columns', 'filters']}
               onColumnVisible={onSaveGridColumnState}
               onColumnPinned={onSaveGridColumnState}
@@ -411,6 +540,54 @@ export const QueryBuilderTDSGridResult = observer(
               }
             />
           )}
+          {resultState.wavgAggregationState?.isApplyingWavg && (
+            <Dialog
+              open={resultState.wavgAggregationState.isApplyingWavg}
+              onClose={() =>
+                resultState.wavgAggregationState?.setIsApplyingWavg(false)
+              }
+              classes={{
+                root: 'editor-modal__root-container',
+                container: 'editor-modal__container',
+                paper: 'editor-modal__content',
+              }}
+            >
+              <Modal
+                darkMode={
+                  !applicationStore.layoutService
+                    .TEMPORARY__isLightColorThemeEnabled
+                }
+                className="query-editor__blocking-alert"
+              >
+                <ModalHeader title="Applying Weighted Average" />
+                <ModalBody>
+                  <div className="query-builder__result__tds-grid__text">
+                    choose a weighted column from dropdown
+                  </div>
+                  <CustomSelectorInput
+                    options={weightedColumnOptions}
+                    onChange={onWeightedColumnOptionChange}
+                    value={selectedWeightedColumn}
+                    placeholder={'Choose a weighted column'}
+                    darkMode={
+                      !applicationStore.layoutService
+                        .TEMPORARY__isLightColorThemeEnabled
+                    }
+                  />
+                </ModalBody>
+                <ModalFooter>
+                  <ModalFooterButton
+                    onClick={() => {
+                      resultState.wavgAggregationState?.setIsApplyingWavg(
+                        false,
+                      );
+                    }}
+                    text="Apply"
+                  />
+                </ModalFooter>
+              </Modal>
+            </Dialog>
+          )}
         </div>
       </div>
     );

package/src/stores/QueryBuilderResultState.ts

@@ -48,6 +48,7 @@ import { QUERY_BUILDER_EVENT } from '../__lib__/QueryBuilderEvent.js';
 import { ExecutionPlanState } from './execution-plan/ExecutionPlanState.js';
 import type { DataGridColumnState } from '@finos/legend-lego/data-grid';
 import { downloadStream } from '@finos/legend-application';
+import { QueryBuilderDataGridCustomAggregationFunction } from '../components/result/tds/QueryBuilderTDSGridResult.js';
 
 export const DEFAULT_LIMIT = 1000;
 
@@ -83,8 +84,37 @@ type QueryBuilderDataGridConfig = {
   isPivotModeEnabled: boolean | undefined;
   isLocalModeEnabled: boolean | undefined;
   previewLimit?: number | undefined;
+  weightedColumnPairs?: Map<string, string> | undefined;
 };
 
+export class QueryBuilderResultWavgAggregationState {
+  isApplyingWavg = false;
+  weightedColumnIdPairs: Map<string, string>;
+
+  constructor() {
+    makeObservable(this, {
+      isApplyingWavg: observable,
+      weightedColumnIdPairs: observable,
+      setIsApplyingWavg: action,
+      addWeightedColumnIdPair: action,
+      removeWeightedColumnIdPair: action,
+    });
+    this.weightedColumnIdPairs = new Map<string, string>();
+  }
+
+  setIsApplyingWavg(val: boolean): void {
+    this.isApplyingWavg = val;
+  }
+
+  addWeightedColumnIdPair(col: string, weightedColumnId: string): void {
+    this.weightedColumnIdPairs.set(col, weightedColumnId);
+  }
+
+  removeWeightedColumnIdPair(col: string): void {
+    this.weightedColumnIdPairs.delete(col);
+  }
+}
+
 export class QueryBuilderResultState {
   readonly queryBuilderState: QueryBuilderState;
   readonly executionPlanState: ExecutionPlanState;
@@ -105,6 +135,7 @@ export class QueryBuilderResultState {
   isSelectingCells: boolean;
 
   gridConfig: QueryBuilderDataGridConfig | undefined;
+  wavgAggregationState: QueryBuilderResultWavgAggregationState | undefined;
 
   constructor(queryBuilderState: QueryBuilderState) {
     makeObservable(this, {
@@ -120,7 +151,9 @@ export class QueryBuilderResultState {
       isSelectingCells: observable,
       isQueryUsageViewerOpened: observable,
       gridConfig: observable,
+      wavgAggregationState: observable,
       setGridConfig: action,
+      setWavgAggregationState: action,
       setIsSelectingCells: action,
       setIsRunningQuery: action,
       setExecutionResult: action,
@@ -153,6 +186,12 @@ export class QueryBuilderResultState {
     this.gridConfig = val;
   }
 
+  setWavgAggregationState(
+    val: QueryBuilderResultWavgAggregationState | undefined,
+  ): void {
+    this.wavgAggregationState = val;
+  }
+
   setIsSelectingCells(val: boolean): void {
     this.isSelectingCells = val;
   }
@@ -199,11 +238,50 @@ export class QueryBuilderResultState {
     }
   }
 
+  processWeightedColumnPairsMap(
+    config: QueryGridConfig,
+  ): Map<string, string> | undefined {
+    if (config.weightedColumnPairs) {
+      const wavgColumns = config.columns
+        .filter(
+          (col) =>
+            (col as DataGridColumnState).aggFunc ===
+            QueryBuilderDataGridCustomAggregationFunction.WAVG,
+        )
+        .map((col) => (col as DataGridColumnState).colId);
+      const weightedColumnPairsMap = new Map<string, string>();
+      config.weightedColumnPairs.forEach((wc) => {
+        if (wc[0] && wc[1]) {
+          weightedColumnPairsMap.set(wc[0], wc[1]);
+        }
+      });
+      for (const wavgCol of weightedColumnPairsMap.keys()) {
+        if (!wavgColumns.includes(wavgCol)) {
+          weightedColumnPairsMap.delete(wavgCol);
+        }
+      }
+      return weightedColumnPairsMap;
+    }
+    return undefined;
+  }
+
   handlePreConfiguredGridConfig(config: QueryGridConfig): void {
-    const newConfig = {
-      ...config,
-      columns: config.columns as DataGridColumnState[],
-    };
+    let newConfig;
+    const weightedColumnPairsMap = this.processWeightedColumnPairsMap(config);
+    if (weightedColumnPairsMap) {
+      this.wavgAggregationState = new QueryBuilderResultWavgAggregationState();
+      this.wavgAggregationState.weightedColumnIdPairs = weightedColumnPairsMap;
+      newConfig = {
+        ...config,
+        weightedColumnPairs: weightedColumnPairsMap,
+        columns: config.columns as DataGridColumnState[],
+      };
+    } else {
+      newConfig = {
+        ...config,
+        columns: config.columns as DataGridColumnState[],
+      };
+    }
     if (config.previewLimit) {
       this.setPreviewLimit(config.previewLimit);
     }

package/src/stores/QueryBuilderState.ts

@@ -32,6 +32,7 @@ import {
   ActionState,
   hashArray,
   assertTrue,
+  assertNonNullable,
 } from '@finos/legend-shared';
 import { QueryBuilderFilterState } from './filter/QueryBuilderFilterState.js';
 import { QueryBuilderFetchStructureState } from './fetch-structure/QueryBuilderFetchStructureState.js';
@@ -54,6 +55,7 @@ import {
   type ValueSpecification,
   type Type,
   type QueryGridConfig,
+  type QueryExecutionContext,
   GRAPH_MANAGER_EVENT,
   CompilationError,
   extractSourceInformationCoordinates,
@@ -73,6 +75,7 @@ import {
   InstanceValue,
   Multiplicity,
   RuntimePointer,
+  QueryExplicitExecutionContext,
 } from '@finos/legend-graph';
 import { buildLambdaFunction } from './QueryBuilderValueSpecificationBuilder.js';
 import type {
@@ -287,6 +290,24 @@ export abstract class QueryBuilderState implements CommandRegistrar {
     return this.allVariables.map((e) => e.name);
   }
 
+  getQueryExecutionContext(): QueryExecutionContext {
+    const queryExeContext = new QueryExplicitExecutionContext();
+    const runtimeValue = guaranteeType(
+      this.executionContextState.runtimeValue,
+      RuntimePointer,
+      'Query runtime must be of type runtime pointer',
+    );
+    assertNonNullable(
+      this.executionContextState.mapping,
+      'Query required mapping to update',
+    );
+    queryExeContext.mapping = PackageableElementExplicitReference.create(
+      this.executionContextState.mapping,
+    );
+    queryExeContext.runtime = runtimeValue.packageableRuntime;
+    return queryExeContext;
+  }
+
   /**
    * Gets information about the current queryBuilderState.
    * This information can be used as a part of analytics