[DataGridPremium] Use the aggregated value on tree data real groups

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

@@ -47,6 +47,19 @@ const getAggregationCellValue = ({
     if (aggregationRowsScope === 'filtered' && filteredRowsLookup[rowId] === false) {
       return;
     }
+
+    // If the row is a group, we want to aggregate based on its children
+    // For instance in the following tree, we want the aggregated values of A to be based on A.A, A.B.A and A.B.B but not A.B
+    // A
+    //   A.A
+    //   A.B
+    //     A.B.A
+    //     A.B.B
+    const rowNode = apiRef.current.getRowNode(rowId)!;
+    if (rowNode.children?.length) {
+      return;
+    }
+
     values.push(apiRef.current.getCellValue(rowId, field));
   });
 

packages/grid/x-data-grid-premium/src/hooks/features/aggregation/wrapColumnWithAggregation.tsx

@@ -188,8 +188,9 @@ export const wrapColumnWithAggregationValue = ({
     field: string,
   ): GridAggregationLookup[GridRowId][string] | null => {
     let cellAggregationPosition: GridAggregationPosition | null = null;
+    const rowNode = apiRef.current.getRowNode(id)!;
 
-    if (id.toString().startsWith('auto-generated-row-')) {
+    if (rowNode.children?.length) {
       cellAggregationPosition = 'inline';
     } else if (id.toString().startsWith('auto-generated-group-footer-')) {
       cellAggregationPosition = 'footer';
@@ -200,8 +201,7 @@ export const wrapColumnWithAggregationValue = ({
     }
 
     // TODO: Add custom root id
-    const groupId =
-      cellAggregationPosition === 'inline' ? id : apiRef.current.getRowNode(id)!.parent ?? '';
+    const groupId = cellAggregationPosition === 'inline' ? id : rowNode.parent ?? '';
 
     const aggregationResult = gridAggregationLookupSelector(apiRef)[groupId]?.[field];
     if (!aggregationResult || aggregationResult.position !== cellAggregationPosition) {

packages/grid/x-data-grid-premium/src/tests/aggregation.DataGridPremium.test.tsx

@@ -278,6 +278,86 @@ describe('<DataGridPremium /> - Aggregation', () => {
     });
   });
 
+  describe('Tree Data', () => {
+    const TreeDataTest = (props: Omit<DataGridPremiumProps, 'columns'>) => {
+      return (
+        <Test
+          treeData
+          defaultGroupingExpansionDepth={-1}
+          columns={[
+            {
+              field: 'value',
+              headerName: 'Value',
+              type: 'number',
+            },
+          ]}
+          getTreeDataPath={(row) => row.hierarchy}
+          getRowId={(row) => row.hierarchy.join('/')}
+          groupingColDef={{ headerName: 'Files', width: 350 }}
+          getAggregationPosition={(rowNode) => (rowNode != null ? 'inline' : null)}
+          initialState={{
+            aggregation: {
+              model: {
+                value: 'sum',
+              },
+            },
+          }}
+          {...props}
+        />
+      );
+    };
+
+    it('should use aggregated values instead of provided values on data groups', () => {
+      render(
+        <TreeDataTest
+          rows={[
+            {
+              hierarchy: ['A'],
+              value: 10,
+            },
+            {
+              hierarchy: ['A', 'A'],
+              value: 1,
+            },
+            {
+              hierarchy: ['A', 'B'],
+              value: 2,
+            },
+          ]}
+        />,
+      );
+
+      expect(getColumnValues(1)).to.deep.equal(['3' /* Agg "A" */, '1', '2']);
+    });
+
+    it('should only aggregate based on leaves', () => {
+      render(
+        <TreeDataTest
+          rows={[
+            {
+              hierarchy: ['A'],
+              value: 2,
+            },
+            {
+              hierarchy: ['A', 'A'],
+              value: 2,
+            },
+            {
+              hierarchy: ['A', 'A', 'A'],
+              value: 1,
+            },
+            {
+              hierarchy: ['A', 'A', 'B'],
+              value: 1,
+            },
+          ]}
+        />,
+      );
+
+      expect(getColumnValues(1)).to.deep.equal(['2' /* Agg "A" */, '2' /* Agg "A.A" */, '1', '1']);
+    });
+  });
+
   describe('Column Menu', () => {
     it('should render select on aggregable column', () => {
       render(<Test />);