datacube: fix aggregation count issue (2) (#3877)

finos · Feb 7, 2025 · ddcd23b · ddcd23b
1 parent 672ccc5
commit ddcd23b
Show file tree

Hide file tree

Showing 5 changed files with 62 additions and 61 deletions.
diff --git a/.changeset/soft-dingos-happen.md b/.changeset/soft-dingos-happen.md
@@ -0,0 +1,4 @@
+---
+'@finos/legend-application-data-cube': patch
+'@finos/legend-data-cube': patch
+---
diff --git a/packages/legend-application-data-cube/src/stores/LegendDataCubeCacheManager.ts b/packages/legend-application-data-cube/src/stores/LegendDataCubeCacheManager.ts
@@ -96,21 +96,15 @@ export class LegendDataCubeDataCubeCacheManager {
       let colType: string;
       switch (col.type as string) {
         case PRIMITIVE_TYPE.BOOLEAN: {
-          colType = 'BOOLEAN';
-          break;
-        }
-        case PRIMITIVE_TYPE.NUMBER: {
-          colType = 'DOUBLE';
+          colType = 'BIT';
           break;
         }
         case PRIMITIVE_TYPE.INTEGER: {
-          colType = 'FLOAT';
-          break;
-        }
-        case PRIMITIVE_TYPE.DECIMAL: {
-          colType = 'DECIMAL';
+          colType = 'INTEGER';
           break;
         }
+        case PRIMITIVE_TYPE.NUMBER:
+        case PRIMITIVE_TYPE.DECIMAL:
         case PRIMITIVE_TYPE.FLOAT: {
           colType = 'FLOAT';
           break;
@@ -119,7 +113,7 @@ export class LegendDataCubeDataCubeCacheManager {
         case PRIMITIVE_TYPE.STRICTDATE:
         case PRIMITIVE_TYPE.DATETIME:
         case PRIMITIVE_TYPE.DATE: {
-          colType = 'STRING';
+          colType = 'VARCHAR';
           break;
         }
         case PRIMITIVE_TYPE.STRING: {
@@ -170,7 +164,12 @@ export class LegendDataCubeDataCubeCacheManager {
     const rows = data.map((row) => {
       const tdsRow = new TDSRow();
       tdsRow.values = columnNames.map(
-        (column) => row[column] as string | number | boolean | null,
+        (column) =>
+          // NOTE: DuckDB WASM returns ArrayBuffer for numeric value, such as for count(*)
+          // so we need to convert it to number
+          (ArrayBuffer.isView(row[column])
+            ? row[column].valueOf()
+            : row[column]) as string | number | boolean | null,
       );
       return tdsRow;
     });

diff --git a/packages/legend-application-data-cube/src/stores/LegendDataCubeDataCubeEngine.ts b/packages/legend-application-data-cube/src/stores/LegendDataCubeDataCubeEngine.ts
@@ -53,7 +53,6 @@ import {
   V1_Column,
   V1_Database,
   V1_Date,
-  V1_Double,
   V1_DuckDBDatasourceSpecification,
   V1_EngineRuntime,
   V1_IdentifiedConnection,
@@ -69,7 +68,6 @@ import {
   V1_TestAuthenticationStrategy,
   V1_VarChar,
   V1_Bit,
-  V1_Decimal,
   V1_Float,
   PackageableElementPointerType,
   DatabaseType,
@@ -543,20 +541,14 @@ export class LegendDataCubeDataCubeEngine extends DataCubeEngine {
           column.type = new V1_Bit();
           break;
         }
-        case PRIMITIVE_TYPE.NUMBER: {
-          column.type = new V1_Double();
-          break;
-        }
         case PRIMITIVE_TYPE.INTEGER: {
           column.type = new V1_Integer();
           break;
         }
-        case PRIMITIVE_TYPE.FLOAT: {
-          column.type = new V1_Float();
-          break;
-        }
+        case PRIMITIVE_TYPE.NUMBER:
+        case PRIMITIVE_TYPE.FLOAT:
         case PRIMITIVE_TYPE.DECIMAL: {
-          column.type = new V1_Decimal();
+          column.type = new V1_Float();
           break;
         }
         case PRIMITIVE_TYPE.DATE:

diff --git a/packages/legend-data-cube/src/stores/view/grid/DataCubeGridClientEngine.ts b/packages/legend-data-cube/src/stores/view/grid/DataCubeGridClientEngine.ts
@@ -31,7 +31,11 @@ import {
   pruneObject,
 } from '@finos/legend-shared';
 import { buildExecutableQuery } from '../../core/DataCubeQueryBuilder.js';
-import { PureClientVersion, type TabularDataSet } from '@finos/legend-graph';
+import {
+  PRIMITIVE_TYPE,
+  PureClientVersion,
+  type TabularDataSet,
+} from '@finos/legend-graph';
 import { makeObservable, observable, runInAction } from 'mobx';
 import type {
   DataCubeConfiguration,
@@ -43,15 +47,12 @@ import { isPivotResultColumnName } from '../../core/DataCubeQueryEngine.js';
 import { buildQuerySnapshot } from './DataCubeGridQuerySnapshotBuilder.js';
 import { AlertType } from '../../services/DataCubeAlertService.js';
 import type { DataCubeViewState } from '../DataCubeViewState.js';
-import {
-  _aggCountCol,
-  buildGridDataFetchExecutableQuery,
-} from './DataCubeGridQueryBuilder.js';
-import { _colSpecArrayParam } from '../../core/DataCubeQuerySnapshotBuilderUtils.js';
+import { buildGridDataFetchExecutableQuery } from './DataCubeGridQueryBuilder.js';
 import { DataCubeSettingKey } from '../../../__lib__/DataCubeSetting.js';
 import { DEFAULT_ALERT_WINDOW_CONFIG } from '../../services/DataCubeLayoutService.js';
 import type { DataCubeExecutionResult } from '../../core/DataCubeEngine.js';
 import { _lambda } from '../../core/DataCubeQueryBuilderUtils.js';
+import { sum } from 'mathjs';
 
 type DataCubeGridClientCellValue = string | number | boolean | null | undefined;
 type DataCubeGridClientRowData = {
@@ -136,8 +137,7 @@ export const INTERNAL__GRID_CLIENT_ROOT_AGGREGATION_COLUMN_ID =
   'INTERNAL__rootAggregation';
 export const INTERNAL__GRID_CLIENT_DATA_FETCH_MANUAL_TRIGGER_COLUMN_ID =
   'INTERNAL__dataFetchManualTrigger';
-export const INTERNAL__GRID_CLIENT_ROW_GROUPING_COUNT_AGG_COLUMN_ID =
-  'INTERNAL__count';
+export const INTERNAL__GRID_CLIENT_ROW_GROUPING_COUNT_AGG_COLUMN_ID = 'count';
 
 export enum DataCubeGridClientPinnedAlignement {
   LEFT = 'left',
@@ -228,17 +228,19 @@ function buildRowData(
           ? value
           : INTERNAL__GRID_CLIENT_MISSING_VALUE;
       if (snapshot.data.pivot && snapshot.data.groupBy) {
-        row[INTERNAL__GRID_CLIENT_ROW_GROUPING_COUNT_AGG_COLUMN_ID] =
-          result.columns
-            .filter(
-              (col) =>
-                isPivotResultColumnName(col) &&
-                col.endsWith(
-                  INTERNAL__GRID_CLIENT_ROW_GROUPING_COUNT_AGG_COLUMN_ID,
-                ),
-            )
-            .map((col) => (row[col] as number | undefined) ?? 0)
-            .reduce((a, b) => a + b, 0);
+        row[INTERNAL__GRID_CLIENT_ROW_GROUPING_COUNT_AGG_COLUMN_ID] = Number(
+          sum(
+            result.columns
+              .filter(
+                (col) =>
+                  isPivotResultColumnName(col) &&
+                  col.endsWith(
+                    INTERNAL__GRID_CLIENT_ROW_GROUPING_COUNT_AGG_COLUMN_ID,
+                  ),
+              )
+              .map((col) => (row[col] as number | undefined) ?? 0),
+          ).toString(),
+        );
       }
     });
     return row;
@@ -260,15 +262,7 @@ async function getCastColumns(
   snapshot.data.groupExtendedColumns = [];
   snapshot.data.sortColumns = [];
   snapshot.data.limit = 0;
-  const query = buildExecutableQuery(snapshot, view.source, view.engine, {
-    postProcessor: (_snapshot, sequence, funcMap, configuration, engine) => {
-      // when both pivot and groupBy present, we need to account for the count column
-      // in the cast expression
-      if (funcMap.pivot && currentSnapshot.data.groupBy) {
-        _colSpecArrayParam(funcMap.pivot, 1).colSpecs.push(_aggCountCol());
-      }
-    },
-  });
+  const query = buildExecutableQuery(snapshot, view.source, view.engine);
 
   const result = await view.engine.executeQuery(
     _lambda([], [query]),
@@ -285,10 +279,15 @@ async function getCastColumns(
     },
   );
 
-  return result.result.builder.columns.map((column) => ({
-    name: column.name,
-    type: column.type as string,
-  }));
+  return result.result.builder.columns.map((column) => {
+    const type = column.type as string;
+    return {
+      name: column.name,
+      // FIXME: this is a workaround to help plan generation does not handle well decimal type
+      // Remove this once https://github.com/finos/legend-engine/pull/3400 is productionized
+      type: type === PRIMITIVE_TYPE.DECIMAL ? PRIMITIVE_TYPE.FLOAT : type,
+    };
+  });
 }
 
 export type DataCubeGridClientDataFetchRequest = {
@@ -399,6 +398,7 @@ export class DataCubeGridClientServerSideDataSource
       this._grid.isPaginationEnabled,
     );
     let result: DataCubeExecutionResult;
+    let rowData: DataCubeGridClientRowData[];
 
     try {
       result = await this._view.engine.executeQuery(
@@ -415,6 +415,7 @@ export class DataCubeGridClientServerSideDataSource
             : undefined,
         },
       );
+      rowData = buildRowData(result.result.result, newSnapshot);
     } catch (error) {
       assertErrorThrown(error);
       this._view.alertService.alertError(error, {
@@ -425,7 +426,6 @@ export class DataCubeGridClientServerSideDataSource
       return;
     }
 
-    const rowData = buildRowData(result.result.result, newSnapshot);
     if (
       this._view.settingService.getBooleanValue(
         DataCubeSettingKey.DEBUGGER__ENABLE_DEBUG_MODE,

diff --git a/packages/legend-data-cube/src/stores/view/grid/DataCubeGridQueryBuilder.ts b/packages/legend-data-cube/src/stores/view/grid/DataCubeGridQueryBuilder.ts
@@ -22,6 +22,7 @@ import {
   V1_GenericType,
   V1_RelationType,
   V1_getGenericTypeFullPath,
+  V1_createRelationTypeColumn,
 } from '@finos/legend-graph';
 import { type DataCubeQuerySnapshot } from '../../core/DataCubeQuerySnapshot.js';
 import {
@@ -102,14 +103,15 @@ function _addCountAggColumnToPivot(
 ): void {
   if (funcMap.pivot && funcMap.pivotCast) {
     _colSpecArrayParam(funcMap.pivot, 1).colSpecs.push(_aggCountCol());
-    const castColumns = guaranteeType(
+    const relationType = guaranteeType(
       guaranteeType(
         guaranteeType(funcMap.pivotCast.parameters[0], V1_GenericTypeInstance)
           .genericType.typeArguments[0],
         V1_GenericType,
       ).rawType,
       V1_RelationType,
-    ).columns.map((column) =>
+    );
+    const castColumns = relationType.columns.map((column) =>
       _colSpec(
         column.name,
         undefined,
@@ -133,10 +135,14 @@ function _addCountAggColumnToPivot(
           PIVOT_COLUMN_NAME_VALUE_SEPARATOR +
           INTERNAL__GRID_CLIENT_ROW_GROUPING_COUNT_AGG_COLUMN_ID,
       )
-      .map((col) => _colSpec(col, undefined, undefined, PRIMITIVE_TYPE.INTEGER))
       .forEach((col) => {
-        castColumns.push(col);
-        countAggColumns.push(col);
+        // Here we directly modify the list of cast columns in the relation type
+        relationType.columns.push(
+          V1_createRelationTypeColumn(col, PRIMITIVE_TYPE.INTEGER),
+        );
+        countAggColumns.push(
+          _colSpec(col, undefined, undefined, PRIMITIVE_TYPE.INTEGER),
+        );
       });
   }
 }