Collapse duplicate rows when table is flattened

extension/src/experiments/columns/collect/order.ts

@@ -1,5 +1,9 @@
 import { Column, ColumnType } from '../../webview/contract'
-import { EXPERIMENT_COLUMN_ID } from '../constants'
+import {
+  EXPERIMENT_COLUMN_ID,
+  BRANCH_COLUMN_ID,
+  COMMIT_COLUMN_ID
+} from '../constants'
 
 export const collectColumnOrder = async (
   existingColumnOrder: string[],
@@ -25,6 +29,14 @@ export const collectColumnOrder = async (
     existingColumnOrder.unshift(EXPERIMENT_COLUMN_ID)
   }
 
+  if (!existingColumnOrder.includes(BRANCH_COLUMN_ID)) {
+    existingColumnOrder.splice(1, 0, BRANCH_COLUMN_ID)
+  }
+
+  if (!existingColumnOrder.includes(COMMIT_COLUMN_ID)) {
+    existingColumnOrder.splice(2, 0, COMMIT_COLUMN_ID)
+  }
+
   return [
     ...existingColumnOrder,
     ...acc.timestamp,

extension/src/experiments/columns/constants.ts

@@ -11,3 +11,12 @@ export const timestampColumn: Column = {
 }
 
 export const EXPERIMENT_COLUMN_ID = 'id'
+
+export const COMMIT_COLUMN_ID = 'commit'
+export const BRANCH_COLUMN_ID = 'branch'
+
+export const DEFAULT_COLUMN_IDS = [
+  EXPERIMENT_COLUMN_ID,
+  BRANCH_COLUMN_ID,
+  COMMIT_COLUMN_ID
+]

extension/src/experiments/columns/model.ts

@@ -12,6 +12,11 @@ import {
   limitSummaryOrder
 } from './util'
 import { collectColumnOrder } from './collect/order'
+import {
+  BRANCH_COLUMN_ID,
+  COMMIT_COLUMN_ID,
+  DEFAULT_COLUMN_IDS
+} from './constants'
 import { Column, ColumnType } from '../webview/contract'
 import { ExpShowOutput } from '../../cli/dvc/contract'
 import { PersistenceKey } from '../../persistence/constants'
@@ -104,10 +109,10 @@ export class ColumnsModel extends PathSelectionModel<Column> {
 
   public selectFirst(firstColumns: string[]) {
     const columnOrder = [
-      'id',
+      ...DEFAULT_COLUMN_IDS,
       ...firstColumns,
       ...this.getColumnOrder().filter(
-        column => !['id', ...firstColumns].includes(column)
+        column => ![...DEFAULT_COLUMN_IDS, ...firstColumns].includes(column)
       )
     ]
     this.setColumnOrder(columnOrder)
@@ -201,6 +206,13 @@ export class ColumnsModel extends PathSelectionModel<Column> {
         return this.setColumnOrderFromData(selectedColumns)
       }
     }
+
+    const maybeMissingDefaultColumns = [COMMIT_COLUMN_ID, BRANCH_COLUMN_ID]
+    for (const id of maybeMissingDefaultColumns) {
+      if (!this.columnOrderState.includes(id)) {
+        return this.setColumnOrderFromData(selectedColumns)
+      }
+    }
   }
 
   private transformAndSetChanges(data: ExpShowOutput) {

extension/src/experiments/model/index.ts

@@ -438,11 +438,19 @@ export class ExperimentsModel extends ModelWithPersistence {
   }
 
   public getRowData() {
-    const commitsBySha = this.applyFiltersToCommits()
+    const workspaceRow = {
+      branch: WORKSPACE_BRANCH,
+      ...this.addDetails(this.workspace)
+    }
+    const sorts = this.getSorts()
+    const flattenRowData = sorts.length > 0
+    if (flattenRowData) {
+      return this.getFlattenedRowData(workspaceRow)
+    }
+
+    const commitsBySha: { [sha: string]: Commit } = this.applyFiltersToCommits()
+    const rows: Commit[] = [workspaceRow]
 
-    const rows: Commit[] = [
-      { branch: WORKSPACE_BRANCH, ...this.addDetails(this.workspace) }
-    ]
     for (const { branch, sha } of this.rowOrder) {
       const commit = commitsBySha[sha]
       if (!commit) {
@@ -829,4 +837,54 @@ export class ExperimentsModel extends ModelWithPersistence {
     }
     return commitsBySha
   }
+
+  private applyFiltersToFlattenedCommits() {
+    const commitsBySha: { [sha: string]: Commit[] } = {}
+    const filters = this.getFilters()
+
+    for (const commit of this.commits) {
+      const commitWithSelectedAndStarred = this.addDetails(commit)
+      const experiments = this.getExperimentsByCommit(
+        commitWithSelectedAndStarred
+      )
+
+      commitsBySha[commit.sha as string] = [
+        commitWithSelectedAndStarred,
+        ...(experiments || [])
+      ].filter(exp => !!filterExperiment(filters, exp))
+    }
+
+    return commitsBySha
+  }
+
+  private addBranchToFlattenedCommit(commitAndExps: Commit[], branch: string) {
+    return commitAndExps.map(commitOrExp => ({
+      ...commitOrExp,
+      branch: commitOrExp.branch || branch,
+      otherBranches: commitOrExp?.otherBranches
+        ? [...commitOrExp.otherBranches, branch]
+        : []
+    }))
+  }
+
+  private getFlattenedRowData(workspaceRow: Commit): Commit[] {
+    const commitsBySha: { [sha: string]: Commit[] } =
+      this.applyFiltersToFlattenedCommits()
+    const rowsBySha: { [sha: string]: Commit[] } = {}
+
+    for (const { branch, sha } of this.rowOrder) {
+      if (!commitsBySha[sha]) {
+        continue
+      }
+
+      const commitAndExps = rowsBySha[sha] || commitsBySha[sha]
+
+      rowsBySha[sha] = this.addBranchToFlattenedCommit(commitAndExps, branch)
+    }
+
+    return [
+      workspaceRow,
+      ...sortExperiments(this.getSorts(), Object.values(rowsBySha).flat())
+    ]
+  }
 }

extension/src/experiments/webview/contract.ts

@@ -58,6 +58,7 @@ export type Experiment = {
   executorStatus?: ExecutorStatus
   timestamp?: string | null
   branch?: string | typeof WORKSPACE_BRANCH
+  otherBranches?: string[]
   baselineSha?: string | undefined
   studioLinkType?: StudioLinkType
 }

extension/src/test/fixtures/expShow/base/columns.ts

@@ -10,6 +10,8 @@ const nestedParamsFile = join('nested', 'params.yaml')
 
 export const dataColumnOrder: string[] = [
   'id',
+  'branch',
+  'commit',
   'Created',
   'metrics:summary.json:accuracy',
   'metrics:summary.json:loss',

extension/src/test/fixtures/expShow/sorted/columns.ts

@@ -0,0 +1,199 @@
+import { join } from 'path'
+import { Column, ColumnType } from '../../../../experiments/webview/contract'
+import { buildMetricOrParamPath } from '../../../../experiments/columns/paths'
+import { timestampColumn } from '../../../../experiments/columns/constants'
+
+const nestedParamsFile = join('nested', 'params.yaml')
+
+const data: Column[] = [
+  timestampColumn,
+  {
+    type: ColumnType.METRICS,
+    hasChildren: true,
+    label: 'summary.json',
+    parentPath: buildMetricOrParamPath(ColumnType.METRICS),
+    path: buildMetricOrParamPath(ColumnType.METRICS, 'summary.json')
+  },
+  {
+    type: ColumnType.METRICS,
+    hasChildren: false,
+    label: 'loss',
+    parentPath: buildMetricOrParamPath(ColumnType.METRICS, 'summary.json'),
+    path: buildMetricOrParamPath(ColumnType.METRICS, 'summary.json', 'loss'),
+    pathArray: [ColumnType.METRICS, 'summary.json', 'loss'],
+    firstValueType: 'number'
+  },
+  {
+    type: ColumnType.METRICS,
+    hasChildren: false,
+    label: 'accuracy',
+    parentPath: buildMetricOrParamPath(ColumnType.METRICS, 'summary.json'),
+    path: buildMetricOrParamPath(
+      ColumnType.METRICS,
+      'summary.json',
+      'accuracy'
+    ),
+    pathArray: [ColumnType.METRICS, 'summary.json', 'accuracy'],
+    firstValueType: 'number'
+  },
+  {
+    type: ColumnType.METRICS,
+    hasChildren: false,
+    label: 'val_loss',
+    parentPath: buildMetricOrParamPath(ColumnType.METRICS, 'summary.json'),
+    path: buildMetricOrParamPath(
+      ColumnType.METRICS,
+      'summary.json',
+      'val_loss'
+    ),
+    pathArray: [ColumnType.METRICS, 'summary.json', 'val_loss'],
+    firstValueType: 'number'
+  },
+  {
+    type: ColumnType.METRICS,
+    hasChildren: false,
+    label: 'val_accuracy',
+    parentPath: buildMetricOrParamPath(ColumnType.METRICS, 'summary.json'),
+    path: buildMetricOrParamPath(
+      ColumnType.METRICS,
+      'summary.json',
+      'val_accuracy'
+    ),
+    pathArray: [ColumnType.METRICS, 'summary.json', 'val_accuracy'],
+    firstValueType: 'number'
+  },
+  {
+    type: ColumnType.PARAMS,
+    hasChildren: true,
+    label: 'params.yaml',
+    parentPath: ColumnType.PARAMS,
+    path: buildMetricOrParamPath(ColumnType.PARAMS, 'params.yaml')
+  },
+  {
+    type: ColumnType.PARAMS,
+    hasChildren: false,
+    label: 'code_names',
+    parentPath: buildMetricOrParamPath(ColumnType.PARAMS, 'params.yaml'),
+    path: buildMetricOrParamPath(
+      ColumnType.PARAMS,
+      'params.yaml',
+      'code_names'
+    ),
+    pathArray: [ColumnType.PARAMS, 'params.yaml', 'code_names'],
+    firstValueType: 'array'
+  },
+  {
+    type: ColumnType.PARAMS,
+    hasChildren: false,
+    label: 'epochs',
+    parentPath: buildMetricOrParamPath(ColumnType.PARAMS, 'params.yaml'),
+    path: buildMetricOrParamPath(ColumnType.PARAMS, 'params.yaml', 'epochs'),
+    pathArray: [ColumnType.PARAMS, 'params.yaml', 'epochs'],
+    firstValueType: 'number'
+  },
+  {
+    type: ColumnType.PARAMS,
+    hasChildren: false,
+    label: 'learning_rate',
+    parentPath: buildMetricOrParamPath(ColumnType.PARAMS, 'params.yaml'),
+    path: buildMetricOrParamPath(
+      ColumnType.PARAMS,
+      'params.yaml',
+      'learning_rate'
+    ),
+    pathArray: [ColumnType.PARAMS, 'params.yaml', 'learning_rate'],
+    firstValueType: 'number'
+  },
+  {
+    type: ColumnType.PARAMS,
+    hasChildren: false,
+    label: 'dvc_logs_dir',
+    parentPath: buildMetricOrParamPath(ColumnType.PARAMS, 'params.yaml'),
+    path: buildMetricOrParamPath(
+      ColumnType.PARAMS,
+      'params.yaml',
+      'dvc_logs_dir'
+    ),
+    pathArray: [ColumnType.PARAMS, 'params.yaml', 'dvc_logs_dir'],
+    firstValueType: 'string'
+  },
+  {
+    type: ColumnType.PARAMS,
+    hasChildren: false,
+    label: 'log_file',
+    parentPath: buildMetricOrParamPath(ColumnType.PARAMS, 'params.yaml'),
+    path: buildMetricOrParamPath(ColumnType.PARAMS, 'params.yaml', 'log_file'),
+    pathArray: [ColumnType.PARAMS, 'params.yaml', 'log_file'],
+    firstValueType: 'string'
+  },
+  {
+    type: ColumnType.PARAMS,
+    hasChildren: false,
+    label: 'dropout',
+    parentPath: buildMetricOrParamPath(ColumnType.PARAMS, 'params.yaml'),
+    path: buildMetricOrParamPath(ColumnType.PARAMS, 'params.yaml', 'dropout'),
+    pathArray: [ColumnType.PARAMS, 'params.yaml', 'dropout'],
+    firstValueType: 'number'
+  },
+  {
+    type: ColumnType.PARAMS,
+    hasChildren: true,
+    label: 'process',
+    parentPath: buildMetricOrParamPath(ColumnType.PARAMS, 'params.yaml'),
+    path: buildMetricOrParamPath(ColumnType.PARAMS, 'params.yaml', 'process')
+  },
+  {
+    type: ColumnType.PARAMS,
+    hasChildren: false,
+    label: 'threshold',
+    parentPath: buildMetricOrParamPath(
+      ColumnType.PARAMS,
+      'params.yaml',
+      'process'
+    ),
+    path: buildMetricOrParamPath(
+      ColumnType.PARAMS,
+      'params.yaml',
+      'process',
+      'threshold'
+    ),
+    pathArray: [ColumnType.PARAMS, 'params.yaml', 'process', 'threshold'],
+    firstValueType: 'number'
+  },
+  {
+    type: ColumnType.PARAMS,
+    hasChildren: true,
+    label: nestedParamsFile,
+    parentPath: ColumnType.PARAMS,
+    path: buildMetricOrParamPath(ColumnType.PARAMS, nestedParamsFile)
+  },
+  {
+    type: ColumnType.PARAMS,
+    hasChildren: false,
+    label: 'test',
+    parentPath: buildMetricOrParamPath(ColumnType.PARAMS, nestedParamsFile),
+    path: buildMetricOrParamPath(ColumnType.PARAMS, nestedParamsFile, 'test'),
+    pathArray: [ColumnType.PARAMS, nestedParamsFile, 'test'],
+    firstValueType: 'boolean'
+  },
+  {
+    type: ColumnType.PARAMS,
+    hasChildren: false,
+    label: 'test_arg',
+    parentPath: buildMetricOrParamPath(
+      ColumnType.PARAMS,
+      'params.yaml',
+      'process'
+    ),
+    path: buildMetricOrParamPath(
+      ColumnType.PARAMS,
+      'params.yaml',
+      'process',
+      'test_arg'
+    ),
+    pathArray: [ColumnType.PARAMS, 'params.yaml', 'process', 'test_arg'],
+    firstValueType: 'string'
+  }
+]
+
+export default data