Column Customization: Persist Column Settings (#6253)

* Motivation for features / changes
We now have a way for users to customize which columns they would like
to appear in their data table and in what order. However, we currently
revert back to default on every page load. This change stores their
selection in localStorage so that it persists a crossed sessions.

* Technical description of changes
This just follows the same structure as all persisted settings in OSS.

Author: JamesHollyer

tensorboard/webapp/metrics/metrics_module.ts

@@ -32,6 +32,8 @@ import {
   getMetricsScalarSmoothing,
   getMetricsStepSelectorEnabled,
   getMetricsTooltipSort,
+  getRangeSelectionHeaders,
+  getSingleSelectionHeaders,
   isMetricsSettingsPaneOpen,
   METRICS_FEATURE_KEY,
   METRICS_SETTINGS_DEFAULT,
@@ -115,6 +117,18 @@ export function getMetricsTimeSeriesLinkedTimeEnabled() {
   });
 }
 
+export function getSingleSelectionHeadersFactory() {
+  return createSelector(getSingleSelectionHeaders, (singleSelectionHeaders) => {
+    return {singleSelectionHeaders};
+  });
+}
+
+export function getRangeSelectionHeadersFactory() {
+  return createSelector(getRangeSelectionHeaders, (rangeSelectionHeaders) => {
+    return {rangeSelectionHeaders};
+  });
+}
+
 @NgModule({
   imports: [
     CommonModule,
@@ -157,6 +171,12 @@ export function getMetricsTimeSeriesLinkedTimeEnabled() {
     PersistentSettingsConfigModule.defineGlobalSetting(
       getMetricsTimeSeriesLinkedTimeEnabled
     ),
+    PersistentSettingsConfigModule.defineGlobalSetting(
+      getSingleSelectionHeadersFactory
+    ),
+    PersistentSettingsConfigModule.defineGlobalSetting(
+      getRangeSelectionHeadersFactory
+    ),
   ],
   providers: [
     {

tensorboard/webapp/metrics/store/metrics_reducers.ts

@@ -461,13 +461,19 @@ const reducer = createReducer(
       partialSettings.rangeSelectionEnabled ?? state.rangeSelectionEnabled;
     const linkedTimeEnabled =
       partialSettings.linkedTimeEnabled ?? state.linkedTimeEnabled;
+    const singleSelectionHeaders =
+      partialSettings.singleSelectionHeaders ?? state.singleSelectionHeaders;
+    const rangeSelectionHeaders =
+      partialSettings.rangeSelectionHeaders ?? state.rangeSelectionHeaders;
 
     return {
       ...state,
       isSettingsPaneOpen,
       stepSelectorEnabled,
       rangeSelectionEnabled,
       linkedTimeEnabled,
+      singleSelectionHeaders,
+      rangeSelectionHeaders,
       settings: {
         ...state.settings,
         ...metricsSettings,

tensorboard/webapp/metrics/store/metrics_reducers_test.ts

@@ -2804,6 +2804,67 @@ describe('metrics reducers', () => {
 
       expect(nextState.isSettingsPaneOpen).toBe(false);
     });
+
+    it('loads singleSelectionHeaders setting into the next state', () => {
+      const beforeState = buildMetricsState({
+        singleSelectionHeaders: [
+          {type: ColumnHeaderType.RUN, enabled: true},
+          {type: ColumnHeaderType.SMOOTHED, enabled: true},
+          {type: ColumnHeaderType.VALUE, enabled: true},
+        ],
+      });
+
+      const nextState = reducers(
+        beforeState,
+        globalSettingsLoaded({
+          partialSettings: {
+            singleSelectionHeaders: [
+              {type: ColumnHeaderType.SMOOTHED, enabled: true},
+              {type: ColumnHeaderType.RUN, enabled: true},
+              {type: ColumnHeaderType.VALUE, enabled: false},
+            ],
+          },
+        })
+      );
+
+      expect(nextState.singleSelectionHeaders).toEqual([
+        {type: ColumnHeaderType.SMOOTHED, enabled: true},
+        {type: ColumnHeaderType.RUN, enabled: true},
+        {type: ColumnHeaderType.VALUE, enabled: false},
+      ]);
+    });
+
+    it('loads rangeSelectionHeaders setting into the next state', () => {
+      const beforeState = buildMetricsState({
+        rangeSelectionHeaders: [
+          {type: ColumnHeaderType.RUN, enabled: true},
+          {type: ColumnHeaderType.MIN_VALUE, enabled: true},
+          {type: ColumnHeaderType.MAX_VALUE, enabled: true},
+          {type: ColumnHeaderType.MEAN, enabled: false},
+        ],
+      });
+
+      const nextState = reducers(
+        beforeState,
+        globalSettingsLoaded({
+          partialSettings: {
+            rangeSelectionHeaders: [
+              {type: ColumnHeaderType.RUN, enabled: true},
+              {type: ColumnHeaderType.MEAN, enabled: true},
+              {type: ColumnHeaderType.MAX_VALUE, enabled: true},
+              {type: ColumnHeaderType.MIN_VALUE, enabled: false},
+            ],
+          },
+        })
+      );
+
+      expect(nextState.rangeSelectionHeaders).toEqual([
+        {type: ColumnHeaderType.RUN, enabled: true},
+        {type: ColumnHeaderType.MEAN, enabled: true},
+        {type: ColumnHeaderType.MAX_VALUE, enabled: true},
+        {type: ColumnHeaderType.MIN_VALUE, enabled: false},
+      ]);
+    });
   });
 
   it('loads Step Selector Setting into the next state', () => {

tensorboard/webapp/persistent_settings/_data_source/BUILD

@@ -24,6 +24,7 @@ tf_ng_module(
     ],
     deps = [
         "//tensorboard/webapp/metrics:types",
+        "//tensorboard/webapp/metrics/views/card_renderer:scalar_card_types",
     ],
 )
 
@@ -53,6 +54,7 @@ tf_ng_module(
         ":types",
         "//tensorboard/webapp/angular:expect_angular_core_testing",
         "//tensorboard/webapp/metrics:types",
+        "//tensorboard/webapp/metrics/views/card_renderer:scalar_card_types",
         "@npm//@angular/core",
         "@npm//@types/jasmine",
         "@npm//rxjs",

tensorboard/webapp/persistent_settings/_data_source/persistent_settings_data_source.ts

@@ -94,6 +94,14 @@ export class OSSSettingsConverter extends SettingsConverter<
     if (settings.linkedTimeEnabled !== undefined) {
       serializableSettings.linkedTimeEnabled = settings.linkedTimeEnabled;
     }
+    if (settings.singleSelectionHeaders !== undefined) {
+      serializableSettings.singleSelectionHeaders =
+        settings.singleSelectionHeaders;
+    }
+    if (settings.rangeSelectionHeaders !== undefined) {
+      serializableSettings.rangeSelectionHeaders =
+        settings.rangeSelectionHeaders;
+    }
     return serializableSettings;
   }
 
@@ -200,6 +208,20 @@ export class OSSSettingsConverter extends SettingsConverter<
       settings.linkedTimeEnabled = backendSettings.linkedTimeEnabled;
     }
 
+    if (
+      backendSettings.hasOwnProperty('singleSelectionHeaders') &&
+      typeof backendSettings.singleSelectionHeaders === 'object'
+    ) {
+      settings.singleSelectionHeaders = backendSettings.singleSelectionHeaders;
+    }
+
+    if (
+      backendSettings.hasOwnProperty('rangeSelectionHeaders') &&
+      typeof backendSettings.rangeSelectionHeaders === 'object'
+    ) {
+      settings.rangeSelectionHeaders = backendSettings.rangeSelectionHeaders;
+    }
+
     return settings;
   }
 }
@@ -263,6 +285,7 @@ export class PersistentSettingsDataSourceImpl<UiSettings, StorageSettings>
         localStorage.getItem(LEGACY_METRICS_LOCAL_STORAGE_KEY) ?? '{}'
       )
     );
+
     const settings = this.converter.backendToUi(
       this.deserialize(localStorage.getItem(GLOBAL_LOCAL_STORAGE_KEY) ?? '{}')
     );

tensorboard/webapp/persistent_settings/_data_source/persistent_settings_data_source_test.ts

@@ -16,6 +16,7 @@ import {Injectable} from '@angular/core';
 import {TestBed} from '@angular/core/testing';
 import {firstValueFrom} from 'rxjs';
 import {TooltipSort} from '../../metrics/types';
+import {ColumnHeaderType} from '../../metrics/views/card_renderer/scalar_card_types';
 import {
   OSSSettingsConverter,
   PersistentSettingsDataSourceImpl,
@@ -182,6 +183,44 @@ describe('persistent_settings data_source test', () => {
           linkedTimeEnabled: true,
         });
       });
+      it('properly converts singleSelectionEnabled', async () => {
+        getItemSpy.withArgs(TEST_ONLY.GLOBAL_LOCAL_STORAGE_KEY).and.returnValue(
+          JSON.stringify({
+            singleSelectionHeaders: [
+              {type: ColumnHeaderType.RUN, enabled: true},
+              {type: ColumnHeaderType.VALUE, enabled: false},
+            ],
+          })
+        );
+
+        const actual = await firstValueFrom(dataSource.getSettings());
+
+        expect(actual).toEqual({
+          singleSelectionHeaders: [
+            {type: ColumnHeaderType.RUN, enabled: true},
+            {type: ColumnHeaderType.VALUE, enabled: false},
+          ],
+        });
+      });
+      it('properly converts rangeSelectionEnabled', async () => {
+        getItemSpy.withArgs(TEST_ONLY.GLOBAL_LOCAL_STORAGE_KEY).and.returnValue(
+          JSON.stringify({
+            rangeSelectionHeaders: [
+              {type: ColumnHeaderType.RUN, enabled: true},
+              {type: ColumnHeaderType.MIN_VALUE, enabled: true},
+            ],
+          })
+        );
+
+        const actual = await firstValueFrom(dataSource.getSettings());
+
+        expect(actual).toEqual({
+          rangeSelectionHeaders: [
+            {type: ColumnHeaderType.RUN, enabled: true},
+            {type: ColumnHeaderType.MIN_VALUE, enabled: true},
+          ],
+        });
+      });
     });
 
     describe('#setSettings', () => {

tensorboard/webapp/persistent_settings/_data_source/types.ts

@@ -14,6 +14,7 @@ limitations under the License.
 ==============================================================================*/
 
 import {TooltipSort} from '../../metrics/types';
+import {ColumnHeader} from '../../metrics/views/card_renderer/scalar_card_types';
 
 export enum ThemeValue {
   BROWSER_DEFAULT = 'browser_default',
@@ -43,6 +44,8 @@ export declare interface BackendSettings {
   stepSelectorEnabled?: boolean;
   rangeSelectionEnabled?: boolean;
   linkedTimeEnabled?: boolean;
+  singleSelectionHeaders?: ColumnHeader[];
+  rangeSelectionHeaders?: ColumnHeader[];
 }
 
 /**
@@ -65,4 +68,6 @@ export interface PersistableSettings {
   stepSelectorEnabled?: boolean;
   rangeSelectionEnabled?: boolean;
   linkedTimeEnabled?: boolean;
+  singleSelectionHeaders?: ColumnHeader[];
+  rangeSelectionHeaders?: ColumnHeader[];
 }