Persist preSwitchCategorizedUncheckedKeys and `postSwitchCategorize…

…dUncheckedKeys` to preserve the diff swtich state when switching experiments (#4927) * Persist pre switch state Signed-off-by: harupy <17039389+harupy@users.noreply.github.com> * clean up Signed-off-by: harupy <17039389+harupy@users.noreply.github.com> * Add persistedState Signed-off-by: harupy <17039389+harupy@users.noreply.github.com> * Add comments Signed-off-by: harupy <17039389+harupy@users.noreply.github.com>
mlflow · Oct 22, 2021 · bec12da · bec12da
1 parent 68067c2
commit bec12da
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 32 deletions.
diff --git a/mlflow/server/js/src/experiment-tracking/components/ExperimentView.js b/mlflow/server/js/src/experiment-tracking/components/ExperimentView.js
@@ -57,7 +57,6 @@ import {
   DEFAULT_ORDER_BY_KEY,
   DEFAULT_ORDER_BY_ASC,
   DEFAULT_START_TIME,
-  DEFAULT_CATEGORIZED_UNCHECKED_KEYS,
   ATTRIBUTE_COLUMN_SORT_LABEL,
   ATTRIBUTE_COLUMN_SORT_KEY,
   COLUMN_SORT_BY_ASC,
@@ -170,10 +169,6 @@ export class ExperimentView extends Component {
       showDeleteRunModal: false,
       // True if a model for restoring one or more runs should be displayed
       showRestoreRunModal: false,
-      // Columns unselected before turning on the diff-view switch
-      preSwitchCategorizedUncheckedKeys: DEFAULT_CATEGORIZED_UNCHECKED_KEYS,
-      // Columns unselected as the result of turning on the diff-view switch
-      postSwitchCategorizedUncheckedKeys: DEFAULT_CATEGORIZED_UNCHECKED_KEYS,
     };
   }
 
@@ -1168,32 +1163,44 @@ export class ExperimentView extends Component {
   }
 
   handleDiffSwitchChange() {
+    let newCategorizedUncheckedKeys;
+    let switchPersistedState;
+    if (!this.state.persistedState.diffSwitchSelected) {
+      // When turning on the diff switch
+      const { categorizedUncheckedKeys } = this.state.persistedState;
+      newCategorizedUncheckedKeys = ExperimentViewUtil.getCategorizedUncheckedKeysDiffView({
+        ...this.props,
+        categorizedUncheckedKeys,
+      });
+      switchPersistedState = {
+        preSwitchCategorizedUncheckedKeys: categorizedUncheckedKeys,
+        postSwitchCategorizedUncheckedKeys: newCategorizedUncheckedKeys,
+      };
+    } else {
+      // When turning off the diff switch
+      const {
+        preSwitchCategorizedUncheckedKeys,
+        postSwitchCategorizedUncheckedKeys,
+        categorizedUncheckedKeys: currCategorizedUncheckedKeys,
+      } = this.state.persistedState;
+      newCategorizedUncheckedKeys = ExperimentViewUtil.getRestoredCategorizedUncheckedKeys({
+        preSwitchCategorizedUncheckedKeys,
+        postSwitchCategorizedUncheckedKeys,
+        currCategorizedUncheckedKeys,
+      });
+      switchPersistedState = {};
+    }
+
     this.setState(
       {
         persistedState: new ExperimentViewPersistedState({
           ...this.state.persistedState,
           diffSwitchSelected: !this.state.persistedState.diffSwitchSelected,
+          ...switchPersistedState,
         }).toJSON(),
       },
       () => {
-        let categorizedUncheckedKeys;
-        if (this.state.persistedState.diffSwitchSelected) {
-          categorizedUncheckedKeys = ExperimentViewUtil.getCategorizedUncheckedKeysDiffView({
-            ...this.props,
-            categorizedUncheckedKeys: this.state.persistedState.categorizedUncheckedKeys,
-          });
-          this.setState({
-            preSwitchCategorizedUncheckedKeys: this.state.persistedState.categorizedUncheckedKeys,
-            postSwitchCategorizedUncheckedKeys: categorizedUncheckedKeys,
-          });
-        } else {
-          categorizedUncheckedKeys = ExperimentViewUtil.getRestoredCategorizedUncheckedKeys({
-            preSwitchCategorizedUncheckedKeys: this.state.preSwitchCategorizedUncheckedKeys,
-            postSwitchCategorizedUncheckedKeys: this.state.postSwitchCategorizedUncheckedKeys,
-            currCategorizedUncheckedKeys: this.state.persistedState.categorizedUncheckedKeys,
-          });
-        }
-        this.handleColumnSelectionCheck(categorizedUncheckedKeys);
+        this.handleColumnSelectionCheck(newCategorizedUncheckedKeys);
       },
     );
   }

diff --git a/mlflow/server/js/src/experiment-tracking/components/ExperimentView.test.js b/mlflow/server/js/src/experiment-tracking/components/ExperimentView.test.js
@@ -498,13 +498,13 @@ describe('Diff Switch', () => {
 
     // Switch turned on
     instance.handleDiffSwitchChange();
-    expect(wrapper.state().preSwitchCategorizedUncheckedKeys).toEqual({
+    expect(wrapper.state().persistedState.preSwitchCategorizedUncheckedKeys).toEqual({
       [COLUMN_TYPES.ATTRIBUTES]: ['a2'],
       [COLUMN_TYPES.PARAMS]: ['p2'],
       [COLUMN_TYPES.METRICS]: ['m2'],
       [COLUMN_TYPES.TAGS]: ['t2'],
     });
-    expect(wrapper.state().postSwitchCategorizedUncheckedKeys).toEqual({
+    expect(wrapper.state().persistedState.postSwitchCategorizedUncheckedKeys).toEqual({
       [COLUMN_TYPES.ATTRIBUTES]: ['a1'],
       [COLUMN_TYPES.PARAMS]: ['p1'],
       [COLUMN_TYPES.METRICS]: ['m1'],

diff --git a/mlflow/server/js/src/experiment-tracking/sdk/MlflowLocalStorageMessages.js b/mlflow/server/js/src/experiment-tracking/sdk/MlflowLocalStorageMessages.js
@@ -15,7 +15,7 @@
  *    local storage.
  */
 import Immutable from 'immutable';
-import { COLUMN_TYPES } from '../constants';
+import { DEFAULT_CATEGORIZED_UNCHECKED_KEYS } from '../constants';
 
 /**
  * This class wraps attributes of the ExperimentPage component's state that should be
@@ -55,14 +55,13 @@ export const ExperimentViewPersistedState = Immutable.Record(
     unbaggedMetrics: [],
     unbaggedParams: [],
     // Unchecked keys in the columns dropdown
-    categorizedUncheckedKeys: {
-      [COLUMN_TYPES.ATTRIBUTES]: [],
-      [COLUMN_TYPES.PARAMS]: [],
-      [COLUMN_TYPES.METRICS]: [],
-      [COLUMN_TYPES.TAGS]: [],
-    },
+    categorizedUncheckedKeys: DEFAULT_CATEGORIZED_UNCHECKED_KEYS,
     // Switch to select only columns with differences
     diffSwitchSelected: false,
+    // Columns unselected before turning on the diff-view switch
+    preSwitchCategorizedUncheckedKeys: DEFAULT_CATEGORIZED_UNCHECKED_KEYS,
+    // Columns unselected as the result of turning on the diff-view switch
+    postSwitchCategorizedUncheckedKeys: DEFAULT_CATEGORIZED_UNCHECKED_KEYS,
   },
   'ExperimentViewPersistedState',
 );