From 7aaa2929b8fa146e4f9f4289af67f2a61d6a8a3c Mon Sep 17 00:00:00 2001
From: Vlad Babich <vladbabich@deephaven.io>
Date: Mon, 4 May 2026 13:09:46 -0600
Subject: [PATCH] DH-22326: Un-hide group-by columns, keep other hidden columns
 hidden when applying rollup

---
 packages/iris-grid/src/IrisGrid.test.tsx | 47 ++++++++++++++++++++++++
 packages/iris-grid/src/IrisGrid.tsx      | 13 ++++++-
 2 files changed, 59 insertions(+), 1 deletion(-)

diff --git a/packages/iris-grid/src/IrisGrid.test.tsx b/packages/iris-grid/src/IrisGrid.test.tsx
index 3606e4bfca..0389205107 100644
--- a/packages/iris-grid/src/IrisGrid.test.tsx
+++ b/packages/iris-grid/src/IrisGrid.test.tsx
@@ -265,6 +265,53 @@ describe('handleResizeColumn', () => {
   });
 });
 
+describe('handleRollupChange', () => {
+  it('un-hides group-by columns', () => {
+    const columns = irisGridTestUtils.makeColumns(3);
+    const irisGrid = makeComponent(
+      irisGridTestUtils.makeModel(irisGridTestUtils.makeTable({ columns }))
+    );
+    const { metricCalculator } = irisGrid.state;
+    const resetColumnWidth = jest.spyOn(metricCalculator, 'resetColumnWidth');
+
+    const groupByNames = [columns[1].name, columns[2].name];
+
+    act(() => {
+      irisGrid.handleRollupChange({
+        columns: groupByNames,
+        showConstituents: true,
+        showNonAggregatedColumns: true,
+      });
+    });
+
+    expect(resetColumnWidth).toHaveBeenCalledWith(1);
+    expect(resetColumnWidth).toHaveBeenCalledWith(2);
+    expect(irisGrid.state.rollupConfig?.columns).toEqual(groupByNames);
+  });
+
+  it('does not call resetColumnWidth when there are no group-by columns', () => {
+    const irisGrid = makeComponent(
+      irisGridTestUtils.makeModel(
+        irisGridTestUtils.makeTable({
+          columns: irisGridTestUtils.makeColumns(3),
+        })
+      )
+    );
+    const { metricCalculator } = irisGrid.state;
+    const resetColumnWidth = jest.spyOn(metricCalculator, 'resetColumnWidth');
+
+    act(() => {
+      irisGrid.handleRollupChange({
+        columns: [],
+        showConstituents: true,
+        showNonAggregatedColumns: true,
+      });
+    });
+
+    expect(resetColumnWidth).not.toHaveBeenCalled();
+  });
+});
+
 // auto resize -> reset user width and set calculated width to content width
 // manual resize -> set user width to content width
 describe('handleResizeAllColumns', () => {
diff --git a/packages/iris-grid/src/IrisGrid.tsx b/packages/iris-grid/src/IrisGrid.tsx
index 412c5d3941..c8c12447af 100644
--- a/packages/iris-grid/src/IrisGrid.tsx
+++ b/packages/iris-grid/src/IrisGrid.tsx
@@ -3797,8 +3797,19 @@ class IrisGrid extends Component<IrisGridProps, IrisGridState> {
   handleRollupChange(rollupConfig: UIRollupConfig): void {
     log.info('Rollup change', rollupConfig);
 
+    // Un-hide group-by columns for the rollup
+    if (rollupConfig?.columns != null && rollupConfig.columns.length > 0) {
+      const { model } = this.props;
+      const { metricCalculator } = this.state;
+      rollupConfig.columns.forEach(name => {
+        const modelIndex = model.getColumnIndexByName(name);
+        if (modelIndex != null) {
+          metricCalculator.resetColumnWidth(modelIndex);
+        }
+      });
+    }
+
     this.resetGridViewState();
-    this.showAllColumns();
     this.clearAllFilters();
 
     this.startLoading(