Improve collection of summary columns efficiency (#4391)

Author: mattseddon

extension/src/experiments/columns/model.ts

@@ -6,7 +6,12 @@ import {
   collectParamsFiles
 } from './collect'
 import { EXPERIMENT_COLUMN_ID, timestampColumn } from './constants'
-import { SummaryAcc, collectFromColumnOrder } from './util'
+import {
+  MAX_SUMMARY_ORDER_LENGTH,
+  SummaryAcc,
+  collectFromColumnOrder as collectSummaryColumnOrder,
+  limitSummaryOrder
+} from './util'
 import { Column, ColumnType } from '../webview/contract'
 import { ExpShowOutput } from '../../cli/dvc/contract'
 import { PersistenceKey } from '../../persistence/constants'
@@ -49,17 +54,18 @@ export class ColumnsModel extends PathSelectionModel<Column> {
     const acc: SummaryAcc = { metrics: [], params: [] }
     for (const path of this.columnOrderState) {
       const reachedMaxSummaryOrderLength =
-        acc.metrics.length >= 3 && acc.params.length >= 3
-      if (
-        this.status[path] !== Status.SELECTED ||
-        reachedMaxSummaryOrderLength
-      ) {
+        acc.metrics.length >= MAX_SUMMARY_ORDER_LENGTH &&
+        acc.params.length >= MAX_SUMMARY_ORDER_LENGTH
+      if (reachedMaxSummaryOrderLength) {
+        return limitSummaryOrder(acc)
+      }
+      if (this.status[path] !== Status.SELECTED) {
         continue
       }
-      collectFromColumnOrder(path, acc)
+      collectSummaryColumnOrder(path, acc)
     }
 
-    return [...acc.params.slice(0, 3), ...acc.metrics.slice(0, 3)]
+    return limitSummaryOrder(acc)
   }
 
   public getColumnWidths(): Record<string, number> {

extension/src/experiments/columns/util.ts

@@ -1,14 +1,29 @@
 import { ColumnType } from '../webview/contract'
 
+export const MAX_SUMMARY_ORDER_LENGTH = 3
+
 export type SummaryAcc = {
   metrics: string[]
   params: string[]
 }
 
 export const collectFromColumnOrder = (path: string, acc: SummaryAcc) => {
-  if (path.startsWith(ColumnType.METRICS)) {
+  if (
+    path.startsWith(ColumnType.METRICS) &&
+    acc.metrics.length < MAX_SUMMARY_ORDER_LENGTH
+  ) {
     acc.metrics.push(path)
-  } else if (path.startsWith(ColumnType.PARAMS)) {
+    return
+  }
+  if (
+    path.startsWith(ColumnType.PARAMS) &&
+    acc.params.length < MAX_SUMMARY_ORDER_LENGTH
+  ) {
     acc.params.push(path)
   }
 }
+
+export const limitSummaryOrder = (acc: SummaryAcc): string[] => [
+  ...acc.params.slice(0, MAX_SUMMARY_ORDER_LENGTH),
+  ...acc.metrics.slice(0, MAX_SUMMARY_ORDER_LENGTH)
+]

extension/src/experiments/model/tree.ts

@@ -218,7 +218,11 @@ export class ExperimentsTree
     return []
   }
 
-  private formatExperiment(experiment: ExperimentAugmented, dvcRoot: string) {
+  private formatExperiment(
+    experiment: ExperimentAugmented,
+    dvcRoot: string,
+    summaryColumns: string[]
+  ) {
     return {
       collapsibleState: experiment.hasChildren
         ? TreeItemCollapsibleState.Expanded
@@ -233,11 +237,7 @@ export class ExperimentsTree
       iconPath: this.getExperimentIcon(experiment),
       id: experiment.id,
       label: experiment.label,
-      tooltip: this.getTooltip(
-        experiment.error,
-        experiment,
-        this.experiments.getRepository(dvcRoot).getSummaryColumnOrder()
-      ),
+      tooltip: this.getTooltip(experiment.error, experiment, summaryColumns),
       type: experiment.type
     }
   }
@@ -252,21 +252,29 @@ export class ExperimentsTree
       return [{ error: cliError }]
     }
 
+    const summaryColumns = repository.getSummaryColumnOrder()
     return repository
       .getWorkspaceAndCommits()
-      .map(experiment => this.formatExperiment(experiment, dvcRoot))
+      .map(experiment =>
+        this.formatExperiment(experiment, dvcRoot, summaryColumns)
+      )
   }
 
   private getExperimentsByCommit(
     dvcRoot: string,
     commit: Experiment
   ): ExperimentItem[] {
+    const repository = this.experiments.getRepository(dvcRoot)
+    const summaryColumns = repository.getSummaryColumnOrder()
     return (
-      this.experiments
-        .getRepository(dvcRoot)
+      repository
         .getCommitExperiments(commit)
         ?.map(experiment =>
-          this.formatExperiment(experiment as ExperimentAugmented, dvcRoot)
+          this.formatExperiment(
+            experiment as ExperimentAugmented,
+            dvcRoot,
+            summaryColumns
+          )
         ) || []
     )
   }

extension/src/plots/model/index.ts

@@ -214,6 +214,7 @@ export class PlotsModel extends ModelWithPersistence {
 
   public getSelectedRevisionDetails() {
     const selectedRevisions: Revision[] = []
+    const summaryColumns = this.experiments.getSummaryColumnOrder()
     for (const experiment of this.experiments.getSelectedRevisions()) {
       const { commit, description, label, displayColor, id } = experiment
       const revision: Revision = {
@@ -223,10 +224,7 @@ export class PlotsModel extends ModelWithPersistence {
         fetched: this.fetchedRevs.has(id),
         id,
         label,
-        summaryColumns: getRevisionSummaryColumns(
-          this.experiments.getSummaryColumnOrder(),
-          experiment
-        )
+        summaryColumns: getRevisionSummaryColumns(summaryColumns, experiment)
       }
 
       if (commit) {