fix: DH-17831: Fix partition table losing quick filters (#2682)

- partition tables were losing there quick filters
- needed to rebuild the filters on table change for partition tables

---------

Co-authored-by: Mike Bender <mikebender@deephaven.io>

Author: dgodinez-dh

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

@@ -11,12 +11,18 @@ import {
 import IrisGrid from './IrisGrid';
 import IrisGridTestUtils from './IrisGridTestUtils';
 import type IrisGridProxyModel from './IrisGridProxyModel';
+import { isPartitionedGridModel } from './PartitionedGridModel';
 
 jest.mock('@deephaven/grid', () => ({
   ...jest.requireActual('@deephaven/grid'),
   isExpandableColumnGridModel: jest.fn(),
 }));
 
+jest.mock('./PartitionedGridModel', () => ({
+  ...jest.requireActual('./PartitionedGridModel'),
+  isPartitionedGridModel: jest.fn(() => false),
+}));
+
 const { asMock } = TestUtils;
 
 const VIEW_SIZE = 500;
@@ -592,6 +598,29 @@ describe('handleResizeAllColumns', () => {
   });
 });
 
+describe('handleTableChanged', () => {
+  it('calls rebuildFilters for a partitioned model with partition mode', () => {
+    const component = makeComponent();
+    component.rebuildFilters = jest.fn();
+    // Set partitionConfig while isPartitionedGridModel is still false so
+    // IrisGridModelUpdater does not attempt to apply it to the regular model
+    act(() => {
+      component.setState({
+        partitionConfig: { mode: 'partition', partitions: [] },
+      });
+    });
+    // Mock true only for the handleTableChanged call; deps are unchanged so
+    // IrisGridModelUpdater's updatePartitionConfig won't re-fire during the
+    // movedColumns setState triggered inside handleTableChanged
+    asMock(isPartitionedGridModel).mockReturnValue(true);
+    act(() => {
+      component.handleTableChanged();
+    });
+    asMock(isPartitionedGridModel).mockReturnValue(false);
+    expect(component.rebuildFilters).toHaveBeenCalled();
+  });
+});
+
 describe('Advanced Filter', () => {
   it.each([
     { columnIndex: -1, expectedVisibility: false },

packages/iris-grid/src/IrisGrid.tsx

@@ -3389,6 +3389,11 @@ class IrisGrid extends Component<IrisGridProps, IrisGridState> {
     const { model } = this.props;
     // movedColumns reset triggers metricCalculator update in the Grid component
     this.setState({ movedColumns: model.initialMovedColumns });
+    // For partitioned tables, we want to rebuild filters on table change to ensure filters are applied to the new partition
+    const { partitionConfig } = this.state;
+    if (isPartitionedGridModel(model) && partitionConfig?.mode !== 'keys') {
+      this.rebuildFilters();
+    }
   }
 
   handleViewChanged(metrics?: GridMetrics): void {

packages/iris-grid/src/IrisGridProxyModel.test.ts

@@ -1,7 +1,8 @@
 import dh from '@deephaven/jsapi-shim';
 import type IrisGridModel from './IrisGridModel';
-import type IrisGridProxyModel from './IrisGridProxyModel';
+import IrisGridProxyModel from './IrisGridProxyModel';
 import IrisGridTestUtils from './IrisGridTestUtils';
+import { type PartitionConfig } from './PartitionedGridModel';
 
 const irisGridTestUtils = new IrisGridTestUtils(dh);
 
@@ -136,6 +137,77 @@ describe('IrisGridProxyModel', () => {
     expect(testUnderlyingSetter).toHaveBeenCalledTimes(1);
   });
 
+  describe('set partitionConfig', () => {
+    const mockPartitionedTable = {
+      getMergedTable: jest.fn(() =>
+        Promise.resolve(irisGridTestUtils.makeTable())
+      ),
+      getKeyTable: jest.fn(() =>
+        Promise.resolve(irisGridTestUtils.makeTable())
+      ),
+      getKeys: jest.fn(() => Promise.resolve([])),
+      getTable: jest.fn(() => Promise.resolve(irisGridTestUtils.makeTable())),
+      keyColumns: [],
+      columns: [],
+      close: jest.fn(),
+      addEventListener: jest.fn(),
+      removeEventListener: jest.fn(),
+    };
+
+    let setNextModelSpy: jest.SpyInstance;
+
+    beforeEach(() => {
+      setNextModelSpy = jest
+        .spyOn(IrisGridProxyModel.prototype, 'setNextModel')
+        .mockImplementation(() => null);
+    });
+
+    afterEach(() => {
+      setNextModelSpy.mockRestore();
+    });
+
+    function makePartitionedModel() {
+      return irisGridTestUtils.makeModel(
+        mockPartitionedTable as unknown as Parameters<
+          typeof irisGridTestUtils.makeModel
+        >[0]
+      );
+    }
+
+    it('does not swap model when the same config reference is set again', () => {
+      const model = makePartitionedModel();
+      const config: PartitionConfig = { mode: 'partition', partitions: ['a'] };
+      model.partitionConfig = config;
+      setNextModelSpy.mockClear();
+      model.partitionConfig = config;
+      expect(setNextModelSpy).not.toHaveBeenCalled();
+    });
+
+    it('does not swap model when a structurally equal config is set', () => {
+      const model = makePartitionedModel();
+      model.partitionConfig = { mode: 'partition', partitions: ['a'] };
+      setNextModelSpy.mockClear();
+      model.partitionConfig = { mode: 'partition', partitions: ['a'] };
+      expect(setNextModelSpy).not.toHaveBeenCalled();
+    });
+
+    it('swaps model when mode changes', () => {
+      const model = makePartitionedModel();
+      model.partitionConfig = { mode: 'partition', partitions: ['a'] };
+      setNextModelSpy.mockClear();
+      model.partitionConfig = { mode: 'merged', partitions: ['a'] };
+      expect(setNextModelSpy).toHaveBeenCalled();
+    });
+
+    it('swaps model when partition values change', () => {
+      const model = makePartitionedModel();
+      model.partitionConfig = { mode: 'partition', partitions: ['a'] };
+      setNextModelSpy.mockClear();
+      model.partitionConfig = { mode: 'partition', partitions: ['b'] };
+      expect(setNextModelSpy).toHaveBeenCalled();
+    });
+  });
+
   // Functions must be set on prototype
   test('Proxies functions when necessary', () => {
     const testFn = jest.fn();

packages/iris-grid/src/IrisGridProxyModel.ts

@@ -314,6 +314,20 @@ class IrisGridProxyModel extends IrisGridModel implements PartitionedGridModel {
     if (!this.isPartitionRequired) {
       throw new Error('Partitions are not available');
     }
+    // Skip if config hasn't changed to avoid unnecessary model swaps (e.g. on remount)
+    if (
+      this.partition === partitionConfig ||
+      (this.partition != null &&
+        partitionConfig != null &&
+        this.partition.mode === partitionConfig.mode &&
+        this.partition.partitions.length ===
+          partitionConfig.partitions.length &&
+        this.partition.partitions.every(
+          (p, i) => p === partitionConfig.partitions[i]
+        ))
+    ) {
+      return;
+    }
     log.debug('set partitionConfig', partitionConfig);
     this.partition = partitionConfig;