fix: DH-22326: Un-hide group-by columns, keep other hidden columns hidden when applying rollup

packages/iris-grid/src/IrisGrid.test.tsx

@@ -265,6 +265,91 @@ describe('handleResizeColumn', () => {
   });
 });
 
+describe('handleRollupChange', () => {
+  it('un-hides group-by columns by name', () => {
+    const columns = irisGridTestUtils.makeColumns(3);
+    const irisGrid = makeComponent(
+      irisGridTestUtils.makeModel(irisGridTestUtils.makeTable({ columns }))
+    );
+    const { metricCalculator } = irisGrid.state;
+    const resetColumnWidthByName = jest.spyOn(
+      metricCalculator,
+      'resetColumnWidthByName'
+    );
+
+    const groupByNames = [columns[1].name, columns[2].name];
+
+    act(() => {
+      irisGrid.handleRollupChange({
+        columns: groupByNames,
+        showConstituents: true,
+        showNonAggregatedColumns: true,
+      });
+    });
+
+    expect(resetColumnWidthByName).toHaveBeenCalledWith(groupByNames[0]);
+    expect(resetColumnWidthByName).toHaveBeenCalledWith(groupByNames[1]);
+    expect(irisGrid.state.rollupConfig?.columns).toEqual(groupByNames);
+  });
+
+  it('does not call resetColumnWidthByName when there are no group-by columns', () => {
+    const irisGrid = makeComponent(
+      irisGridTestUtils.makeModel(
+        irisGridTestUtils.makeTable({
+          columns: irisGridTestUtils.makeColumns(3),
+        })
+      )
+    );
+    const { metricCalculator } = irisGrid.state;
+    const resetColumnWidthByName = jest.spyOn(
+      metricCalculator,
+      'resetColumnWidthByName'
+    );
+
+    act(() => {
+      irisGrid.handleRollupChange({
+        columns: [],
+        showConstituents: true,
+        showNonAggregatedColumns: true,
+      });
+    });
+
+    expect(resetColumnWidthByName).not.toHaveBeenCalled();
+  });
+
+  it('un-hides a group-by column that is absent from the current (already rolled-up) model', () => {
+    // Simulates editing an existing rollup where a newly-added group-by
+    // column is not present in the current model (e.g. non-aggregated columns
+    // are hidden), so a model-index lookup against this.props.model would
+    // miss it. Resetting by name must still clear its hidden width.
+    const columns = irisGridTestUtils.makeColumns(3);
+    const model = irisGridTestUtils.makeModel(
+      irisGridTestUtils.makeTable({ columns })
+    );
+    const irisGrid = makeComponent(model);
+    const { metricCalculator } = irisGrid.state;
+
+    // Simulate the column having been hidden previously (width 0 stored by name).
+    const newGroupByName = 'NotInCurrentModel';
+    metricCalculator.userColumnWidthsByName.set(newGroupByName, 0);
+
+    // Spy AFTER seeding so the spy still calls through.
+    jest.spyOn(model, 'getColumnIndexByName').mockReturnValue(undefined);
+
+    act(() => {
+      irisGrid.handleRollupChange({
+        columns: [newGroupByName],
+        showConstituents: false,
+        showNonAggregatedColumns: false,
+      });
+    });
+
+    expect(metricCalculator.userColumnWidthsByName.has(newGroupByName)).toBe(
+      false
+    );
+  });
+});
+
 // auto resize -> reset user width and set calculated width to content width
 // manual resize -> set user width to content width
 describe('handleResizeAllColumns', () => {

packages/iris-grid/src/IrisGrid.tsx

@@ -3797,8 +3797,18 @@ class IrisGrid extends Component<IrisGridProps, IrisGridState> {
   handleRollupChange(rollupConfig: UIRollupConfig): void {
     log.info('Rollup change', rollupConfig);
 
+    // Un-hide group-by columns for the rollup. Reset by name (rather than by
+    // model index) because when editing an existing rollup the current model
+    // is the rolled-up one and may not contain a newly-added group-by column,
+    // so an index lookup against this.props.model would miss it.
+    if (rollupConfig?.columns != null && rollupConfig.columns.length > 0) {
+      const { metricCalculator } = this.state;
+      rollupConfig.columns.forEach(name => {
+        metricCalculator.resetColumnWidthByName(name);
+      });
+    }
+
     this.resetGridViewState();
-    this.showAllColumns();
     this.clearAllFilters();
 
     this.startLoading(

packages/iris-grid/src/IrisGridMetricCalculator.ts

@@ -306,6 +306,22 @@ export class IrisGridMetricCalculator extends GridMetricCalculator {
     }
   }
 
+  /**
+   * Resets the user width for a column by its name. Works regardless of
+   * whether the column is currently present in the model, so it is safe to
+   * call when the model is about to be swapped (e.g. when applying a rollup).
+   * @param name The name of the column to reset
+   */
+  resetColumnWidthByName(name: ColumnName): void {
+    this.userColumnWidthsByName.delete(name);
+    if (this.cachedModelColumnNames != null) {
+      const index = this.cachedModelColumnNames.indexOf(name);
+      if (index >= 0) {
+        super.resetColumnWidth(index);
+      }
+    }
+  }
+
   /**
    * Resets all user column widths
    */