@finos/legend-query-builder 4.14.25 → 4.14.27

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/components/shared/BasicValueSpecificationEditor.tsx

@@ -70,7 +70,13 @@ import {
 } from '@finos/legend-shared';
 import { flowResult } from 'mobx';
 import { observer } from 'mobx-react-lite';
-import { useEffect, useRef, useState } from 'react';
+import {
+  forwardRef,
+  useEffect,
+  useImperativeHandle,
+  useRef,
+  useState,
+} from 'react';
 import {
   instanceValue_setValue,
   instanceValue_setValues,
@@ -204,23 +210,26 @@ const VariableExpressionParameterEditor = observer(
 );
 
 const StringPrimitiveInstanceValueEditor = observer(
-  (props: {
-    valueSpecification: PrimitiveInstanceValue;
-    className?: string | undefined;
-    setValueSpecification: (val: ValueSpecification) => void;
-    resetValue: () => void;
-    selectorConfig?:
-      | {
-          values: string[] | undefined;
-          isLoading: boolean;
-          reloadValues:
-            | DebouncedFunc<(inputValue: string) => GeneratorFn<void>>
-            | undefined;
-          cleanUpReloadValues?: () => void;
-        }
-      | undefined;
-    obseverContext: ObserverContext;
-  }) => {
+  forwardRef<
+    HTMLInputElement,
+    {
+      valueSpecification: PrimitiveInstanceValue;
+      className?: string | undefined;
+      setValueSpecification: (val: ValueSpecification) => void;
+      resetValue: () => void;
+      selectorConfig?:
+        | {
+            values: string[] | undefined;
+            isLoading: boolean;
+            reloadValues:
+              | DebouncedFunc<(inputValue: string) => GeneratorFn<void>>
+              | undefined;
+            cleanUpReloadValues?: () => void;
+          }
+        | undefined;
+      obseverContext: ObserverContext;
+    }
+  >(function StringPrimitiveInstanceValueEditor(props, ref) {
     const {
       valueSpecification,
       className,
@@ -307,6 +316,7 @@ const StringPrimitiveInstanceValueEditor = observer(
             value={value}
             placeholder={value === '' ? '(empty)' : undefined}
             onChange={changeInputValue}
+            ref={ref}
           />
         )}
         <button
@@ -319,7 +329,7 @@ const StringPrimitiveInstanceValueEditor = observer(
         </button>
       </div>
     );
-  },
+  }),
 );
 
 const BooleanPrimitiveInstanceValueEditor = observer(
@@ -367,14 +377,17 @@ const BooleanPrimitiveInstanceValueEditor = observer(
 );
 
 const NumberPrimitiveInstanceValueEditor = observer(
-  (props: {
-    valueSpecification: PrimitiveInstanceValue;
-    isInteger: boolean;
-    className?: string | undefined;
-    resetValue: () => void;
-    setValueSpecification: (val: ValueSpecification) => void;
-    obseverContext: ObserverContext;
-  }) => {
+  forwardRef<
+    HTMLInputElement,
+    {
+      valueSpecification: PrimitiveInstanceValue;
+      isInteger: boolean;
+      className?: string | undefined;
+      resetValue: () => void;
+      setValueSpecification: (val: ValueSpecification) => void;
+      obseverContext: ObserverContext;
+    }
+  >(function NumberPrimitiveInstanceValueEditor(props, ref) {
     const {
       valueSpecification,
       isInteger,
@@ -387,6 +400,7 @@ const NumberPrimitiveInstanceValueEditor = observer(
       (valueSpecification.values[0] as number).toString(),
     );
     const inputRef = useRef<HTMLInputElement>(null);
+    useImperativeHandle(ref, () => inputRef.current as HTMLInputElement, []);
     const numericValue = isInteger
       ? Number.parseInt(Number(value).toString(), 10)
       : Number(value);
@@ -484,7 +498,7 @@ const NumberPrimitiveInstanceValueEditor = observer(
         </button>
       </div>
     );
-  },
+  }),
 );
 
 const EnumValueInstanceValueEditor = observer(
@@ -947,26 +961,29 @@ const DateInstanceValueEditor = observer(
  *
  * See https://github.com/finos/legend-studio/pull/1021
  */
-export const BasicValueSpecificationEditor: React.FC<{
-  valueSpecification: ValueSpecification;
-  graph: PureModel;
-  obseverContext: ObserverContext;
-  typeCheckOption: TypeCheckOption;
-  className?: string | undefined;
-  setValueSpecification: (val: ValueSpecification) => void;
-  resetValue: () => void;
-  isConstant?: boolean;
-  selectorConfig?:
-    | {
-        values: string[] | undefined;
-        isLoading: boolean;
-        reloadValues:
-          | DebouncedFunc<(inputValue: string) => GeneratorFn<void>>
-          | undefined;
-        cleanUpReloadValues?: () => void;
-      }
-    | undefined;
-}> = (props) => {
+export const BasicValueSpecificationEditor = forwardRef<
+  HTMLInputElement,
+  {
+    valueSpecification: ValueSpecification;
+    graph: PureModel;
+    obseverContext: ObserverContext;
+    typeCheckOption: TypeCheckOption;
+    className?: string | undefined;
+    setValueSpecification: (val: ValueSpecification) => void;
+    resetValue: () => void;
+    isConstant?: boolean;
+    selectorConfig?:
+      | {
+          values: string[] | undefined;
+          isLoading: boolean;
+          reloadValues:
+            | DebouncedFunc<(inputValue: string) => GeneratorFn<void>>
+            | undefined;
+          cleanUpReloadValues?: () => void;
+        }
+      | undefined;
+  }
+>(function BasicValueSpecificationEditor(props, ref) {
   const {
     className,
     valueSpecification,
@@ -990,6 +1007,7 @@ export const BasicValueSpecificationEditor: React.FC<{
             resetValue={resetValue}
             selectorConfig={selectorConfig}
             obseverContext={obseverContext}
+            ref={ref}
           />
         );
       case PRIMITIVE_TYPE.BOOLEAN:
@@ -1016,6 +1034,7 @@ export const BasicValueSpecificationEditor: React.FC<{
             className={className}
             resetValue={resetValue}
             obseverContext={obseverContext}
+            ref={ref}
           />
         );
       case PRIMITIVE_TYPE.DATE:
@@ -1134,6 +1153,7 @@ export const BasicValueSpecificationEditor: React.FC<{
             className={className}
             resetValue={resetValue}
             obseverContext={obseverContext}
+            ref={ref}
           />
         );
       }
@@ -1141,4 +1161,4 @@ export const BasicValueSpecificationEditor: React.FC<{
   }
 
   return <UnsupportedValueSpecificationEditor />;
-};
+});

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);
     }