Flatten table on sort

extension/src/experiments/model/index.ts

@@ -438,16 +438,25 @@ 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) {
         continue
       }
+
       if (commit.subRows) {
         commit.subRows = commit.subRows.map(experiment => ({
           ...experiment,
@@ -816,6 +825,7 @@ export class ExperimentsModel extends ModelWithPersistence {
       const experiments = this.getExperimentsByCommit(
         commitWithSelectedAndStarred
       )
+
       const unfilteredCommit = collectUnfiltered(
         commitWithSelectedAndStarred,
         experiments,
@@ -829,4 +839,40 @@ 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 getFlattenedRowData(workspaceRow: Commit): Commit[] {
+    const commitsBySha: { [sha: string]: Commit[] } =
+      this.applyFiltersToFlattenedCommits()
+    const rows = []
+
+    for (const { branch, sha } of this.rowOrder) {
+      const commits = commitsBySha[sha]
+      if (!commits) {
+        continue
+      }
+
+      rows.push(...commits.map(commit => ({ ...commit, branch })))
+    }
+
+    return [workspaceRow, ...sortExperiments(this.getSorts(), rows)]
+  }
 }

extension/src/test/suite/experiments/index.test.ts

@@ -1720,60 +1720,63 @@ suite('Experiments Test Suite', () => {
   })
 
   describe('Sorting', () => {
-    it('should be able to sort', async () => {
-      const { experiments, messageSpy } = await buildExperimentsWebview({
-        disposer: disposable,
-        availableNbCommits: { main: 20 },
-        expShow: generateTestExpShowOutput(
-          {},
+    const mockExpShowOutput = generateTestExpShowOutput(
+      {},
+      {
+        rev: '2d879497587b80b2d9e61f072d9dbe9c07a65357',
+        experiments: [
           {
-            rev: '2d879497587b80b2d9e61f072d9dbe9c07a65357',
-            experiments: [
-              {
-                params: {
-                  'params.yaml': {
-                    data: {
-                      test: 2
-                    }
-                  }
+            params: {
+              'params.yaml': {
+                data: {
+                  test: 2
                 }
-              },
-              {
-                params: {
-                  'params.yaml': {
-                    data: {
-                      test: 1
-                    }
-                  }
+              }
+            }
+          },
+          {
+            params: {
+              'params.yaml': {
+                data: {
+                  test: 1
                 }
-              },
-              {
-                params: {
-                  'params.yaml': {
-                    data: {
-                      test: 3
-                    }
-                  }
+              }
+            }
+          },
+          {
+            params: {
+              'params.yaml': {
+                data: {
+                  test: 3
                 }
               }
-            ]
+            }
           }
-        ),
+        ],
+        data: { params: { 'params.yaml': { data: { test: 5 } } } }
+      }
+    )
+
+    const getIds = (rows: Commit[]) =>
+      rows.map(({ id, subRows }) => {
+        const data: { id: string; subRows?: string[] } = { id }
+
+        if (subRows) {
+          data.subRows = subRows.map(({ id }) => id)
+        }
+        return data
+      })
+
+    it('should be able to flatten the table rows and sort', async () => {
+      const { experiments, messageSpy } = await buildExperimentsWebview({
+        disposer: disposable,
+        availableNbCommits: { main: 20 },
+        expShow: mockExpShowOutput,
         rowOrder: [
           { sha: '2d879497587b80b2d9e61f072d9dbe9c07a65357', branch: 'main' }
         ]
       })
 
-      const getIds = (rows: Commit[]) =>
-        rows.map(({ id, subRows }) => {
-          const data: { id: string; subRows?: string[] } = { id }
-
-          if (subRows) {
-            data.subRows = subRows.map(({ id }) => id)
-          }
-          return data
-        })
-
       const { rows, sorts: noSorts } = messageSpy.lastCall.args[0]
 
       expect(getIds(rows)).to.deep.equal([
@@ -1816,13 +1819,85 @@ suite('Experiments Test Suite', () => {
       expect(getIds(sortedRows)).to.deep.equal([
         { id: EXPERIMENT_WORKSPACE_ID },
         {
-          id: '2d879497587b80b2d9e61f072d9dbe9c07a65357',
-          subRows: ['exp-2', 'exp-1', 'exp-3']
+          id: 'exp-2'
+        },
+        {
+          id: 'exp-1'
+        },
+        {
+          id: 'exp-3'
+        },
+        {
+          id: '2d879497587b80b2d9e61f072d9dbe9c07a65357'
         }
       ])
 
       expect(sorts).to.deep.equal([{ descending: false, path: sortPath }])
     }).timeout(WEBVIEW_TEST_TIMEOUT)
+
+    it('should be able to filter out parent commit rows when sorted', async () => {
+      const { experiments, experimentsModel, messageSpy } =
+        await buildExperimentsWebview({
+          disposer: disposable,
+          availableNbCommits: { main: 20 },
+          expShow: mockExpShowOutput,
+          rowOrder: [
+            { sha: '2d879497587b80b2d9e61f072d9dbe9c07a65357', branch: 'main' }
+          ]
+        })
+
+      const { rows, sorts: noSorts } = messageSpy.lastCall.args[0]
+
+      expect(getIds(rows)).to.deep.equal([
+        { id: EXPERIMENT_WORKSPACE_ID },
+        {
+          id: '2d879497587b80b2d9e61f072d9dbe9c07a65357',
+          subRows: ['exp-1', 'exp-2', 'exp-3']
+        }
+      ])
+
+      expect(noSorts).to.deep.equal([])
+
+      const paramPath = buildMetricOrParamPath(
+        ColumnType.PARAMS,
+        'params.yaml',
+        'test'
+      )
+
+      stub(experimentsModel, 'getFilters').returns([
+        {
+          operator: Operator.LESS_THAN,
+          path: paramPath,
+          value: 4
+        }
+      ])
+      stub(SortQuickPicks, 'pickSortToAdd')
+        .onFirstCall()
+        .resolves({ descending: true, path: paramPath })
+
+      messageSpy.resetHistory()
+      const messageSent = waitForSpyCall(messageSpy, messageSpy.callCount)
+
+      await experiments.addSort()
+      await messageSent
+
+      const { rows: sortedRows, sorts } = messageSpy.lastCall.args[0]
+
+      expect(getIds(sortedRows)).to.deep.equal([
+        { id: EXPERIMENT_WORKSPACE_ID },
+        {
+          id: 'exp-3'
+        },
+        {
+          id: 'exp-1'
+        },
+        {
+          id: 'exp-2'
+        }
+      ])
+
+      expect(sorts).to.deep.equal([{ descending: true, path: paramPath }])
+    }).timeout(WEBVIEW_TEST_TIMEOUT)
   })
 
   describe('persisted state', () => {

extension/src/test/suite/experiments/model/sortBy/tree.test.ts

@@ -95,6 +95,7 @@ suite('Experiments Sort By Tree Test Suite', () => {
     return closeAllEditors()
   })
 
+  // eslint-disable-next-line sonarjs/cognitive-complexity
   describe('ExperimentsSortByTree', () => {
     it('should appear in the UI', async () => {
       await expect(
@@ -161,6 +162,21 @@ suite('Experiments Sort By Tree Test Suite', () => {
             get(exp, selector)
           )
 
+      const getSortedParamsArray = (selector = testParamPathArray) => {
+        const rows = messageSpy.getCall(-1).firstArg.rows
+        const params = []
+
+        for (const row of rows) {
+          const param = get(row, selector)
+
+          if (param) {
+            params.push(param)
+          }
+        }
+
+        return params
+      }
+
       stub(WorkspaceExperiments.prototype, 'getDvcRoots').returns([dvcDemoPath])
       stub(WorkspaceExperiments.prototype, 'getOnlyOrPickProject').resolves(
         dvcDemoPath
@@ -176,9 +192,10 @@ suite('Experiments Sort By Tree Test Suite', () => {
       )
       await tableChangedPromise
       mockShowQuickPick.reset()
-      expect(getParamsArray(), 'single sort with table command').to.deep.equal([
-        1, 2, 3, 4
-      ])
+      expect(
+        getSortedParamsArray(),
+        'single sort with table command'
+      ).to.deep.equal([1, 2, 3, 4])
 
       const tableSortRemoved = experimentsUpdatedEvent(experiments)
 
@@ -193,7 +210,7 @@ suite('Experiments Sort By Tree Test Suite', () => {
 
       await addSortWithMocks(otherTestParamPath, false)
       expect(
-        getParamsArray(),
+        getSortedParamsArray(),
         `row order is maintained after applying a sort on ${otherTestParamPath}`
       ).to.deep.equal([1, 3, 2, 4])
 
@@ -212,7 +229,7 @@ suite('Experiments Sort By Tree Test Suite', () => {
         }
       ])
       expect(
-        getParamsArray(),
+        getSortedParamsArray(),
         'the result of both sorts is sent to the webview'
       ).to.deep.equal([3, 1, 4, 2])
 
@@ -231,7 +248,7 @@ suite('Experiments Sort By Tree Test Suite', () => {
         }
       ])
       expect(
-        getParamsArray(),
+        getSortedParamsArray(),
         'the result of the switched sort is sent to the webview'
       ).to.deep.equal([4, 2, 3, 1])
 
@@ -243,7 +260,7 @@ suite('Experiments Sort By Tree Test Suite', () => {
         }
       )
       expect(
-        getParamsArray(),
+        getSortedParamsArray(),
         'when removing a sort that changes the order of ties, those ties should reflect their original order'
       ).to.deep.equal([2, 4, 1, 3])
 

webview/src/experiments/components/table/body/SortedTableContent.tsx

@@ -0,0 +1,30 @@
+import React, { Fragment, RefObject } from 'react'
+import { Row } from '@tanstack/react-table'
+import { Experiment } from 'dvc/src/experiments/webview/contract'
+import { TableBody } from './TableBody'
+
+interface SortedTableContentProps {
+  rows: Row<Experiment>[]
+  tableRef: RefObject<HTMLTableElement>
+  tableHeadHeight: number
+}
+
+export const SortedTableContent: React.FC<SortedTableContentProps> = ({
+  rows,
+  tableHeadHeight,
+  tableRef
+}) => {
+  return (
+    <>
+      {rows.map((row, i) => (
+        <TableBody
+          key={row.id}
+          tableHeaderHeight={tableHeadHeight}
+          root={tableRef.current}
+          row={row}
+          isLast={i === rows.length - 1}
+        />
+      ))}
+    </>
+  )
+}