From 627d64dffe1c015b3f50af27ee5ef9cc17501a27 Mon Sep 17 00:00:00 2001
From: Hoonji <736199+hoonji@users.noreply.github.com>
Date: Sun, 21 Apr 2024 18:08:34 +0900
Subject: [PATCH 1/2] Memoizes column header functions

---
 .../webapp/metrics/views/card_renderer/BUILD  |  1 +
 .../scalar_card_data_table.ng.html            |  4 +-
 .../card_renderer/scalar_card_data_table.ts   | 27 +++++++----
 .../webapp/runs/views/runs_table/BUILD        |  1 +
 .../views/runs_table/runs_data_table.ng.html  |  4 +-
 .../runs/views/runs_table/runs_data_table.ts  | 47 +++++++++----------
 6 files changed, 47 insertions(+), 37 deletions(-)

diff --git a/tensorboard/webapp/metrics/views/card_renderer/BUILD b/tensorboard/webapp/metrics/views/card_renderer/BUILD
index 0790e5e8f3..288f052416 100644
--- a/tensorboard/webapp/metrics/views/card_renderer/BUILD
+++ b/tensorboard/webapp/metrics/views/card_renderer/BUILD
@@ -286,6 +286,7 @@ tf_ng_module(
         ":utils",
         "//tensorboard/webapp/metrics:types",
         "//tensorboard/webapp/runs:types",
+        "//tensorboard/webapp/util:memoize",
         "//tensorboard/webapp/widgets/card_fob:types",
         "//tensorboard/webapp/widgets/data_table",
         "//tensorboard/webapp/widgets/data_table:types",
diff --git a/tensorboard/webapp/metrics/views/card_renderer/scalar_card_data_table.ng.html b/tensorboard/webapp/metrics/views/card_renderer/scalar_card_data_table.ng.html
index 880ce02d9c..f787fb2a85 100644
--- a/tensorboard/webapp/metrics/views/card_renderer/scalar_card_data_table.ng.html
+++ b/tensorboard/webapp/metrics/views/card_renderer/scalar_card_data_table.ng.html
@@ -30,7 +30,7 @@
   (loadAllColumns)="loadAllColumns.emit()"
 >
   <ng-container header>
-    <ng-container *ngFor="let header of getHeaders()">
+    <ng-container *ngFor="let header of extendHeaders(columnHeaders)">
       <tb-data-table-header-cell
         *ngIf="header.enabled && (header.type !== ColumnHeaderType.SMOOTHED || smoothingEnabled)"
         [header]="header"
@@ -42,7 +42,7 @@
   <ng-container content>
     <ng-container *ngFor="let dataRow of getTimeSelectionTableData()">
       <tb-data-table-content-row>
-        <ng-container *ngFor="let header of getHeaders()">
+        <ng-container *ngFor="let header of extendHeaders(columnHeaders)">
           <tb-data-table-content-cell
             *ngIf="header.enabled && (header.type !== ColumnHeaderType.SMOOTHED || smoothingEnabled)"
             [header]="header"
diff --git a/tensorboard/webapp/metrics/views/card_renderer/scalar_card_data_table.ts b/tensorboard/webapp/metrics/views/card_renderer/scalar_card_data_table.ts
index f8d9b0a36b..6764b0a3b5 100644
--- a/tensorboard/webapp/metrics/views/card_renderer/scalar_card_data_table.ts
+++ b/tensorboard/webapp/metrics/views/card_renderer/scalar_card_data_table.ts
@@ -43,6 +43,7 @@ import {
   AddColumnEvent,
 } from '../../../widgets/data_table/types';
 import {isDatumVisible} from './utils';
+import {memoize} from '../../../util/memoize';
 
 @Component({
   selector: 'scalar-card-data-table',
@@ -74,15 +75,23 @@ export class ScalarCardDataTable {
 
   ColumnHeaderType = ColumnHeaderType;
 
-  getHeaders(): ColumnHeader[] {
-    return [
-      {
-        name: 'color',
-        displayName: '',
-        type: ColumnHeaderType.COLOR,
-        enabled: true,
-      },
-    ].concat(this.columnHeaders);
+  // Columns must be memoized to stop needless re-rendering of the content and headers in these
+  // columns. This has been known to cause problems with the controls in these columns,
+  // specifically the add button.
+  extendHeaders = memoize(this.internalExtendHeaders);
+
+  private internalExtendHeaders(headers: ColumnHeader[]) {
+    return ([] as Array<ColumnHeader>).concat(
+      [
+        {
+          name: 'color',
+          displayName: '',
+          type: ColumnHeaderType.COLOR,
+          enabled: true,
+        },
+      ],
+      headers
+    );
   }
 
   getMinPointInRange(
diff --git a/tensorboard/webapp/runs/views/runs_table/BUILD b/tensorboard/webapp/runs/views/runs_table/BUILD
index 719cb0bd5b..f425ef4b20 100644
--- a/tensorboard/webapp/runs/views/runs_table/BUILD
+++ b/tensorboard/webapp/runs/views/runs_table/BUILD
@@ -107,6 +107,7 @@ tf_ng_module(
         "//tensorboard/webapp/types:ui",
         "//tensorboard/webapp/util:colors",
         "//tensorboard/webapp/util:matcher",
+        "//tensorboard/webapp/util:memoize",
         "//tensorboard/webapp/widgets/custom_modal",
         "//tensorboard/webapp/widgets/data_table",
         "//tensorboard/webapp/widgets/data_table:filter_dialog",
diff --git a/tensorboard/webapp/runs/views/runs_table/runs_data_table.ng.html b/tensorboard/webapp/runs/views/runs_table/runs_data_table.ng.html
index 83688127bd..41cf95f2a0 100644
--- a/tensorboard/webapp/runs/views/runs_table/runs_data_table.ng.html
+++ b/tensorboard/webapp/runs/views/runs_table/runs_data_table.ng.html
@@ -41,7 +41,7 @@
     (loadAllColumns)="loadAllColumns.emit()"
   >
     <ng-container header>
-      <ng-container *ngFor="let header of getHeaders()">
+      <ng-container *ngFor="let header of extendHeaders(headers)">
         <tb-data-table-header-cell
           [header]="header"
           [sortingInfo]="sortingInfo"
@@ -67,7 +67,7 @@
     <ng-container content>
       <ng-container *ngFor="let dataRow of data; trackBy: trackByRuns">
         <tb-data-table-content-row [attr.data-id]="dataRow.id">
-          <ng-container *ngFor="let header of getHeaders()">
+          <ng-container *ngFor="let header of extendHeaders(headers)">
             <tb-data-table-content-cell
               [header]="header"
               [datum]="dataRow[header.name]"
diff --git a/tensorboard/webapp/runs/views/runs_table/runs_data_table.ts b/tensorboard/webapp/runs/views/runs_table/runs_data_table.ts
index d4cf74814f..95a1a2f4aa 100644
--- a/tensorboard/webapp/runs/views/runs_table/runs_data_table.ts
+++ b/tensorboard/webapp/runs/views/runs_table/runs_data_table.ts
@@ -30,6 +30,7 @@ import {
   ReorderColumnEvent,
   AddColumnEvent,
 } from '../../../widgets/data_table/types';
+import {memoize} from '../../../util/memoize';
 @Component({
   selector: 'runs-data-table',
   templateUrl: 'runs_data_table.ng.html',
@@ -67,32 +68,30 @@ export class RunsDataTable {
   @Output() addFilter = new EventEmitter<FilterAddedEvent>();
   @Output() loadAllColumns = new EventEmitter<null>();
 
-  // These columns must be stored and reused to stop needless re-rendering of
-  // the content and headers in these columns. This has been known to cause
-  // problems with the controls in these columns, specifically the color picker.
-  beforeColumns = [
-    {
-      name: 'selected',
-      displayName: '',
-      type: ColumnHeaderType.CUSTOM,
-      enabled: true,
-    },
-  ];
+  // Columns must be memoized to stop needless re-rendering of the content and headers in these
+  // columns. This has been known to cause problems with the controls in these columns,
+  // specifically the add button.
+  extendHeaders = memoize(this.internalExtendHeaders);
 
-  afterColumns = [
-    {
-      name: 'color',
-      displayName: '',
-      type: ColumnHeaderType.COLOR,
-      enabled: true,
-    },
-  ];
-
-  getHeaders() {
+  private internalExtendHeaders(headers: ColumnHeader[]) {
     return ([] as Array<ColumnHeader>).concat(
-      this.beforeColumns,
-      this.headers,
-      this.afterColumns
+      [
+        {
+          name: 'selected',
+          displayName: '',
+          type: ColumnHeaderType.CUSTOM,
+          enabled: true,
+        },
+      ],
+      headers,
+      [
+        {
+          name: 'color',
+          displayName: '',
+          type: ColumnHeaderType.COLOR,
+          enabled: true,
+        },
+      ]
     );
   }
 

From 24d55042f337276c3ec9039ac5600f9d7619e364 Mon Sep 17 00:00:00 2001
From: Hoonji <736199+hoonji@users.noreply.github.com>
Date: Sun, 21 Apr 2024 18:41:05 +0900
Subject: [PATCH 2/2] Fixes test

---
 .../runs/views/runs_table/runs_data_table_test.ts      | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/tensorboard/webapp/runs/views/runs_table/runs_data_table_test.ts b/tensorboard/webapp/runs/views/runs_table/runs_data_table_test.ts
index f2b5432481..6f5d9c411b 100644
--- a/tensorboard/webapp/runs/views/runs_table/runs_data_table_test.ts
+++ b/tensorboard/webapp/runs/views/runs_table/runs_data_table_test.ts
@@ -220,25 +220,25 @@ describe('runs_data_table', () => {
       });
     });
 
-    // If the call to getHeaders produces a new color header on each call, then
+    // If the call to extendHeaders produces a new color header on each call, then
     // the content gets re-rendered when the color picker opens. This causes the
     // color picker modal to close immediately. This test ensures that we keep
     // the same reference to the color header and continue to use that reference
     // instead of creating a new one on each call.
-    it('getHeaders does not produce new color header on each call', () => {
+    it('extendHeaders does not produce new color header on each call', () => {
       const fixture = createComponent({});
-
       const runsDataTable = fixture.debugElement.query(
         By.directive(RunsDataTable)
       );
+      const {headers} = runsDataTable.componentInstance;
 
       expect(
         runsDataTable.componentInstance
-          .getHeaders()
+          .extendHeaders(headers)
           .find((header: ColumnHeader) => header.name === 'color')
       ).toBe(
         runsDataTable.componentInstance
-          .getHeaders()
+          .extendHeaders(headers)
           .find((header: ColumnHeader) => header.name === 'color')
       );
     });