[DataGridPremium] Fix aggregation with sorting

packages/x-data-grid-premium/src/DataGridPremium/useDataGridPremiumComponent.tsx

@@ -4,7 +4,6 @@ import { RefObject } from '@mui/x-internals/types';
 import {
   useGridInitialization,
   useGridInitializeState,
-  useGridVirtualizer,
   useGridClipboard,
   useGridColumnMenu,
   useGridColumns,
@@ -24,7 +23,6 @@ import {
   useGridRows,
   useGridRowsPreProcessors,
   rowsStateInitializer,
-  useGridRowsMeta,
   useGridParamsApi,
   useGridRowSelection,
   useGridSorting,
@@ -79,6 +77,7 @@ import {
   propsStateInitializer,
   rowReorderStateInitializer,
   type GridConfiguration,
+  useFirstRender,
 } from '@mui/x-data-grid-pro/internals';
 import { useGridSelector } from '@mui/x-data-grid-pro';
 import { GridPrivateApiPremium } from '../models/gridApiPremium';
@@ -202,7 +201,6 @@ export const useDataGridPremiumComponent = (
   useGridInitializeState(aiAssistantStateInitializer, apiRef, props);
   useGridInitializeState(chartsIntegrationStateInitializer, apiRef, props);
 
-  useGridVirtualizer(apiRef, props);
   useGridSidebar(apiRef, props);
   useGridPivoting(apiRef, props, inProps.columns, inProps.rows);
   useGridRowGrouping(apiRef, props);
@@ -231,7 +229,6 @@ export const useDataGridPremiumComponent = (
   useGridColumnReorder(apiRef, props);
   useGridColumnResize(apiRef, props);
   useGridPagination(apiRef, props);
-  useGridRowsMeta(apiRef, props);
   useGridRowReorder(apiRef, props);
   useGridScroll(apiRef, props);
   useGridInfiniteLoader(apiRef, props);
@@ -254,6 +251,9 @@ export const useDataGridPremiumComponent = (
   useGridPivotingExportState(apiRef);
 
   // Should be the last thing to run, because all pre-processors should have been registered by now.
+  useFirstRender(() => {
+    apiRef.current.runAppliersForPendingProcessors();
+  });
   React.useEffect(() => {
     apiRef.current.runAppliersForPendingProcessors();
   });

packages/x-data-grid-premium/src/hooks/features/aggregation/useGridAggregation.ts

@@ -10,6 +10,7 @@ import {
   gridRenderContextSelector,
   gridVisibleColumnFieldsSelector,
   gridSortModelSelector,
+  gridRowMaximumTreeDepthSelector,
 } from '@mui/x-data-grid-pro';
 import {
   useGridRegisterPipeProcessor,
@@ -89,126 +90,128 @@ export const useGridAggregation = (
   );
 
   const abortControllerRef = React.useRef<AbortController | null>(null);
-  const applyAggregation = React.useCallback(
-    (reason?: 'filter' | 'sort') => {
-      // Abort previous if any
-      if (abortControllerRef.current) {
-        abortControllerRef.current.abort();
-      }
-      const abortController = new AbortController();
-      abortControllerRef.current = abortController;
-
-      const aggregationRules = getAggregationRules(
-        gridColumnLookupSelector(apiRef),
-        gridAggregationModelSelector(apiRef),
-        props.aggregationFunctions,
-        !!props.dataSource,
-      );
-      const aggregatedFields = Object.keys(aggregationRules);
-      const currentAggregationLookup = gridAggregationLookupSelector(apiRef);
-      const needsSorting = shouldApplySorting(aggregationRules, aggregatedFields);
-      if (reason === 'sort' && !needsSorting) {
-        // no need to re-apply aggregation on `sortedRowsSet` if sorting is not needed
-        return;
-      }
+  const applyAggregation = React.useCallback(() => {
+    // Abort previous if we're proceeding
+    if (abortControllerRef.current) {
+      abortControllerRef.current.abort();
+    }
+    const abortController = new AbortController();
+    abortControllerRef.current = abortController;
 
-      const renderContext = gridRenderContextSelector(apiRef);
-      const visibleColumns = gridVisibleColumnFieldsSelector(apiRef);
+    const aggregationRules = getAggregationRules(
+      gridColumnLookupSelector(apiRef),
+      gridAggregationModelSelector(apiRef),
+      props.aggregationFunctions,
+      !!props.dataSource,
+    );
+    const aggregatedFields = Object.keys(aggregationRules);
+    const currentAggregationLookup = gridAggregationLookupSelector(apiRef);
+
+    const renderContext = gridRenderContextSelector(apiRef);
+    const visibleColumns = gridVisibleColumnFieldsSelector(apiRef);
+
+    const chunks: string[][] = [];
+    const sortedAggregatedFields = gridSortModelSelector(apiRef)
+      .map((s) => s.field)
+      .filter((field) => aggregatedFields.includes(field));
+    const visibleAggregatedFields = visibleColumns
+      .slice(renderContext.firstColumnIndex, renderContext.lastColumnIndex + 1)
+      .filter((field) => aggregatedFields.includes(field));
+    const visibleAggregatedFieldsWithSort = [
+      ...visibleAggregatedFields,
+      ...sortedAggregatedFields.filter((field) => !visibleAggregatedFields.includes(field)),
+    ];
+    const hasAggregatedSortedField =
+      gridRowMaximumTreeDepthSelector(apiRef) > 1 && sortedAggregatedFields.length > 0;
+
+    if (visibleAggregatedFields.length > 0) {
+      chunks.push(visibleAggregatedFieldsWithSort);
+    }
+    const otherAggregatedFields = aggregatedFields.filter(
+      (field) => !visibleAggregatedFieldsWithSort.includes(field),
+    );
 
-      const chunks: string[][] = [];
-      const visibleAggregatedFields = visibleColumns
-        .slice(renderContext.firstColumnIndex, renderContext.lastColumnIndex + 1)
-        .filter((field) => aggregatedFields.includes(field));
-      if (visibleAggregatedFields.length > 0) {
-        chunks.push(visibleAggregatedFields);
-      }
-      const otherAggregatedFields = aggregatedFields.filter(
-        (field) => !visibleAggregatedFields.includes(field),
-      );
+    const chunkSize = 20; // columns per chunk
+    for (let i = 0; i < otherAggregatedFields.length; i += chunkSize) {
+      chunks.push(otherAggregatedFields.slice(i, i + chunkSize));
+    }
 
-      const chunkSize = 20; // columns per chunk
-      for (let i = 0; i < otherAggregatedFields.length; i += chunkSize) {
-        chunks.push(otherAggregatedFields.slice(i, i + chunkSize));
-      }
+    let chunkIndex = 0;
+    const aggregationLookup: GridAggregationLookup = {};
+    let chunkStartTime = performance.now();
+    const timeLimit = 1000 / 120;
 
-      let chunkIndex = 0;
-      const aggregationLookup: GridAggregationLookup = {};
-      let chunkStartTime = performance.now();
-      const timeLimit = 1000 / 120;
+    const processChunk = () => {
+      if (abortController.signal.aborted) {
+        return;
+      }
 
-      const processChunk = () => {
-        if (abortController.signal.aborted) {
-          return;
-        }
+      const currentChunk = chunks[chunkIndex];
+      if (!currentChunk) {
+        apiRef.current.publishEvent('aggregationLookupSet');
+        abortControllerRef.current = null;
+        return;
+      }
 
-        const currentChunk = chunks[chunkIndex];
-        if (!currentChunk) {
-          const sortModel = gridSortModelSelector(apiRef).map((s) => s.field);
-          const hasAggregatedSorting = sortModel.some((field) => aggregationRules[field]);
-          if (hasAggregatedSorting) {
-            apiRef.current.applySorting();
-          }
-          abortControllerRef.current = null;
-          return;
+      const applySorting = shouldApplySorting(aggregationRules, currentChunk);
+
+      // createAggregationLookup now RETURNS new partial lookup
+      const partialLookup = createAggregationLookup({
+        apiRef,
+        getAggregationPosition: props.getAggregationPosition,
+        aggregatedFields: currentChunk,
+        aggregationRules,
+        aggregationRowsScope: props.aggregationRowsScope,
+        isDataSource: !!props.dataSource,
+        applySorting,
+      });
+
+      for (const key of Object.keys(partialLookup)) {
+        for (const field of Object.keys(partialLookup[key])) {
+          aggregationLookup[key] ??= {};
+          aggregationLookup[key][field] = partialLookup[key][field];
         }
+      }
 
-        const applySorting = shouldApplySorting(aggregationRules, currentChunk);
-
-        // createAggregationLookup now RETURNS new partial lookup
-        const partialLookup = createAggregationLookup({
-          apiRef,
-          getAggregationPosition: props.getAggregationPosition,
-          aggregatedFields: currentChunk,
-          aggregationRules,
-          aggregationRowsScope: props.aggregationRowsScope,
-          isDataSource: !!props.dataSource,
-          applySorting,
-        });
-
-        for (const key of Object.keys(partialLookup)) {
-          for (const field of Object.keys(partialLookup[key])) {
-            aggregationLookup[key] ??= {};
-            aggregationLookup[key][field] = partialLookup[key][field];
-          }
-        }
+      apiRef.current.setState((state) => ({
+        ...state,
+        aggregation: { ...state.aggregation, lookup: { ...aggregationLookup } },
+      }));
 
-        apiRef.current.setState((state) => ({
-          ...state,
-          aggregation: { ...state.aggregation, lookup: { ...aggregationLookup } },
-        }));
+      if (chunkIndex === 0 && hasAggregatedSortedField) {
+        apiRef.current.applySorting();
+      }
 
-        chunkIndex += 1;
+      chunkIndex += 1;
 
-        if (performance.now() - chunkStartTime < timeLimit) {
-          processChunk();
-          return;
-        }
+      if (performance.now() - chunkStartTime < timeLimit) {
+        processChunk();
+        return;
+      }
 
-        setTimeout(() => {
-          chunkStartTime = performance.now();
-          processChunk();
-        }, 0);
-      };
+      setTimeout(() => {
+        chunkStartTime = performance.now();
+        processChunk();
+      }, 0);
+    };
 
-      processChunk();
+    processChunk();
 
-      // processChunk() does nothing if there are no aggregated fields
-      // make sure that the lookup is empty in this case
-      if (aggregatedFields.length === 0 && !isObjectEmpty(currentAggregationLookup)) {
-        apiRef.current.setState((state) => ({
-          ...state,
-          aggregation: { ...state.aggregation, lookup: {} },
-        }));
-      }
-    },
-    [
-      apiRef,
-      props.getAggregationPosition,
-      props.aggregationFunctions,
-      props.aggregationRowsScope,
-      props.dataSource,
-    ],
-  );
+    // processChunk() does nothing if there are no aggregated fields
+    // make sure that the lookup is empty in this case
+    if (aggregatedFields.length === 0 && !isObjectEmpty(currentAggregationLookup)) {
+      apiRef.current.setState((state) => ({
+        ...state,
+        aggregation: { ...state.aggregation, lookup: {} },
+      }));
+    }
+  }, [
+    apiRef,
+    props.getAggregationPosition,
+    props.aggregationFunctions,
+    props.aggregationRowsScope,
+    props.dataSource,
+  ]);
 
   React.useEffect(() => {
     return () => {
@@ -219,7 +222,7 @@ export const useGridAggregation = (
     };
   }, []);
 
-  const deferredApplyAggregation = useRunOncePerLoop(applyAggregation);
+  const { schedule: deferredApplyAggregation } = useRunOncePerLoop(applyAggregation);
 
   const aggregationApi: GridAggregationApi = {
     setAggregationModel,
@@ -285,7 +288,32 @@ export const useGridAggregation = (
   useGridEvent(apiRef, 'aggregationModelChange', checkAggregationRulesDiff);
   useGridEvent(apiRef, 'columnsChange', checkAggregationRulesDiff);
   useGridEvent(apiRef, 'filteredRowsSet', deferredApplyAggregation);
-  useGridEvent(apiRef, 'sortedRowsSet', () => deferredApplyAggregation('sort'));
+
+  const lastSortModel = React.useRef(gridSortModelSelector(apiRef));
+  useGridEvent(apiRef, 'sortedRowsSet', () => {
+    const sortModel = gridSortModelSelector(apiRef);
+    if (lastSortModel.current === sortModel) {
+      return;
+    }
+    lastSortModel.current = sortModel;
+
+    const aggregationRules = getAggregationRules(
+      gridColumnLookupSelector(apiRef),
+      gridAggregationModelSelector(apiRef),
+      props.aggregationFunctions,
+      !!props.dataSource,
+    );
+    const aggregatedFields = Object.keys(aggregationRules);
+    if (!aggregatedFields.length) {
+      return;
+    }
+    const needsSorting = shouldApplySorting(aggregationRules, aggregatedFields);
+    if (!needsSorting) {
+      return;
+    }
+
+    deferredApplyAggregation();
+  });
 
   /**
    * EFFECTS
@@ -295,4 +323,10 @@ export const useGridAggregation = (
       apiRef.current.setAggregationModel(props.aggregationModel);
     }
   }, [apiRef, props.aggregationModel]);
+
+  React.useEffect(() => {
+    if (props.getAggregationPosition !== undefined) {
+      deferredApplyAggregation();
+    }
+  }, [deferredApplyAggregation, props.getAggregationPosition]);
 };

packages/x-data-grid-premium/src/hooks/features/chartsIntegration/useGridChartsIntegration.tsx

@@ -879,6 +879,11 @@ export const useGridChartsIntegration = (
     'sortedRowsSet',
     runIf(isChartsIntegrationAvailable, () => debouncedHandleRowDataUpdate(syncedChartIds)),
   );
+  useGridEvent(
+    apiRef,
+    'aggregationLookupSet',
+    runIf(isChartsIntegrationAvailable, () => debouncedHandleRowDataUpdate(syncedChartIds)),
+  );
 
   React.useEffect(() => {
     if (!activeChartId && availableChartIds.length > 0) {