Avoid more unnecessary experiments table re-renders (#4503)

Author: mattseddon

webview/src/experiments/components/table/Table.tsx

@@ -1,11 +1,14 @@
-import React, { useRef, useState, CSSProperties } from 'react'
+import React, { useRef, useState, CSSProperties, useEffect } from 'react'
 import cx from 'classnames'
 import { useDispatch } from 'react-redux'
 import styles from './styles.module.scss'
 import { TableHead } from './header/TableHead'
 import { TableContent } from './body/TableContent'
 import { InstanceProp } from '../../util/interfaces'
-import { clearSelectedRows } from '../../state/rowSelectionSlice'
+import {
+  clearSelectedRows,
+  updateRowOrder
+} from '../../state/rowSelectionSlice'
 
 export const Table: React.FC<InstanceProp> = ({ instance }) => {
   const dispatch = useDispatch()
@@ -15,6 +18,27 @@ export const Table: React.FC<InstanceProp> = ({ instance }) => {
 
   const tableRef = useRef<HTMLTableElement>(null)
 
+  const { setColumnOrder, getHeaderGroups, getAllLeafColumns, getRowModel } =
+    instance
+
+  const { rows, flatRows } = getRowModel()
+
+  useEffect(() => {
+    dispatch(
+      updateRowOrder(
+        flatRows.map(
+          ({ depth, original: { branch, id, executorStatus, starred } }) => ({
+            branch,
+            depth,
+            executorStatus,
+            id,
+            starred
+          })
+        )
+      )
+    )
+  }, [dispatch, flatRows])
+
   return (
     <div
       className={styles.tableContainer}
@@ -34,15 +58,17 @@ export const Table: React.FC<InstanceProp> = ({ instance }) => {
         }}
       >
         <TableHead
-          instance={instance}
+          columnOrder={getAllLeafColumns().map(({ id }) => id)}
+          headerGroups={getHeaderGroups()}
+          setColumnOrder={setColumnOrder}
           root={tableRef.current}
           setExpColumnNeedsShadow={setExpColumnNeedsShadow}
           setTableHeadHeight={setTableHeadHeight}
         />
         <TableContent
-          instance={instance}
-          tableRef={tableRef}
+          rows={rows}
           tableHeadHeight={tableHeadHeight}
+          tableRef={tableRef}
         />
       </table>
     </div>

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

@@ -1,8 +1,13 @@
 import React from 'react'
-import { RowContent } from './Row'
+import { TableRow } from './Row'
 import styles from '../styles.module.scss'
 import { RowProp } from '../../../util/interfaces'
 
-export const NestedRow: React.FC<RowProp> = ({ row }) => {
-  return <RowContent row={row} className={styles.nestedRow} />
+export const NestedRow: React.FC<RowProp & { isExpanded: boolean }> = ({
+  row,
+  isExpanded
+}) => {
+  return (
+    <TableRow row={row} isExpanded={isExpanded} className={styles.nestedRow} />
+  )
 }

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

@@ -1,5 +1,5 @@
 import cx from 'classnames'
-import React, { useCallback, useMemo } from 'react'
+import React, { memo, useCallback, useMemo } from 'react'
 import { useDispatch, useSelector } from 'react-redux'
 import { EXPERIMENT_WORKSPACE_ID } from 'dvc/src/cli/dvc/contract'
 import { StubCell, CellWrapper } from './Cell'
@@ -16,8 +16,9 @@ import {
   toggleRowSelected
 } from '../../../state/rowSelectionSlice'
 
-export const RowContent: React.FC<RowProp & { className?: string }> = ({
+const Row: React.FC<RowProp & { className?: string; isExpanded: boolean }> = ({
   row,
+  isExpanded,
   className
 }): JSX.Element => {
   const changes = useSelector(
@@ -27,10 +28,11 @@ export const RowContent: React.FC<RowProp & { className?: string }> = ({
   const { selectedRows } = useSelector(
     (state: ExperimentsState) => state.rowSelection
   )
-  const { getVisibleCells, original, getIsExpanded, subRows } = row
+  const { getVisibleCells, original, subRows } = row
   const { branch, displayColor, error, starred, id } = original
   const [stubCell, ...cells] = getVisibleCells()
   const isWorkspace = id === EXPERIMENT_WORKSPACE_ID
+
   const changesIfWorkspace = isWorkspace ? changes : undefined
 
   const isRowSelected = !!selectedRows[getCompositeId(id, branch)]
@@ -95,7 +97,7 @@ export const RowContent: React.FC<RowProp & { className?: string }> = ({
           plotColor={displayColor}
           starred={starred}
           isRowSelected={isRowSelected}
-          showSubRowStates={!getIsExpanded() && !isWorkspace}
+          showSubRowStates={!isExpanded && !isWorkspace}
           subRowStates={subRowStates}
           toggleExperiment={() => toggleExperiment(id)}
           toggleRowSelection={toggleRowSelection}
@@ -117,3 +119,5 @@ export const RowContent: React.FC<RowProp & { className?: string }> = ({
     </ContextMenu>
   )
 }
+
+export const TableRow = memo(Row)

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

@@ -1,42 +1,38 @@
 import React from 'react'
 import cx from 'classnames'
 import { EXPERIMENT_WORKSPACE_ID } from 'dvc/src/cli/dvc/contract'
-import { ExperimentGroup } from './ExperimentGroup'
-import { RowContent } from './Row'
+import { NestedRow } from './NestedRow'
+import { TableRow } from './Row'
 import { WorkspaceRowGroup } from './WorkspaceRowGroup'
 import styles from '../styles.module.scss'
-import { InstanceProp, RowProp } from '../../../util/interfaces'
+import { RowProp } from '../../../util/interfaces'
 
-interface TableBodyProps extends RowProp, InstanceProp {
+interface TableBodyProps extends RowProp {
   root: HTMLElement | null
   tableHeaderHeight: number
   isLast?: boolean
 }
 
 export const TableBody: React.FC<TableBodyProps> = ({
   row,
-  instance,
   root,
   tableHeaderHeight,
   isLast
 }) => {
   const contentProps = {
+    isExpanded: row.getIsExpanded(),
     key: row.id,
     row
   }
   const content =
     row.depth > 0 ? (
-      <ExperimentGroup {...contentProps} />
+      <NestedRow {...contentProps} />
     ) : (
-      <RowContent {...contentProps} />
+      <TableRow {...contentProps} />
     )
 
   return row.original.id === EXPERIMENT_WORKSPACE_ID ? (
-    <WorkspaceRowGroup
-      tableHeaderHeight={tableHeaderHeight}
-      root={root}
-      instance={instance}
-    >
+    <WorkspaceRowGroup tableHeaderHeight={tableHeaderHeight} root={root}>
       {content}
     </WorkspaceRowGroup>
   ) : (

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

@@ -16,7 +16,7 @@ const tableStateData = {
 }
 
 jest.mock('../../../../shared/api')
-jest.mock('./ExperimentGroup')
+jest.mock('./NestedRow')
 jest.mock('./Row')
 
 describe('TableContent', () => {
@@ -815,18 +815,36 @@ describe('TableContent', () => {
   } as unknown as Table<Experiment>
 
   const renderTableContent = (rowsInstance = instance) => {
+    const { rows, flatRows } = rowsInstance.getRowModel()
+
     return render(
       <Provider
         store={configureStore({
           preloadedState: {
+            rowSelection: {
+              lastSelectedRowId: undefined,
+              rowOrder: flatRows.map(
+                ({
+                  depth,
+                  original: { branch, id, executorStatus, starred }
+                }) => ({
+                  branch,
+                  depth,
+                  executorStatus,
+                  id,
+                  starred
+                })
+              ),
+              selectedRows: {}
+            },
             tableData: tableStateData
           },
           reducer: experimentsReducers
         })}
       >
         <table>
           <TableContent
-            instance={rowsInstance}
+            rows={rows}
             tableHeadHeight={50}
             tableRef={createRef()}
           />

webview/src/experiments/components/table/body/TableContent.test.tsx

@@ -1,40 +1,21 @@
-import React, { Fragment, RefObject, useEffect } from 'react'
-import { useDispatch } from 'react-redux'
+import React, { Fragment, RefObject } from 'react'
+import { Row } from '@tanstack/react-table'
+import { Experiment } from 'dvc/src/experiments/webview/contract'
 import { TableBody } from './TableBody'
 import { collectBranchWithRows } from './util'
 import { BranchDivider } from './branchDivider/BranchDivider'
-import { InstanceProp } from '../../../util/interfaces'
-import { updateRowOrder } from '../../../state/rowSelectionSlice'
 
-interface TableContentProps extends InstanceProp {
+interface TableContentProps {
+  rows: Row<Experiment>[]
   tableRef: RefObject<HTMLTableElement>
   tableHeadHeight: number
 }
 
 export const TableContent: React.FC<TableContentProps> = ({
-  instance,
-  tableRef,
-  tableHeadHeight
+  rows,
+  tableHeadHeight,
+  tableRef
 }) => {
-  const { rows, flatRows } = instance.getRowModel()
-  const dispatch = useDispatch()
-
-  useEffect(() => {
-    dispatch(
-      updateRowOrder(
-        flatRows.map(
-          ({ depth, original: { branch, id, executorStatus, starred } }) => ({
-            branch,
-            depth,
-            executorStatus,
-            id,
-            starred
-          })
-        )
-      )
-    )
-  }, [dispatch, flatRows])
-
   return (
     <>
       {collectBranchWithRows(rows).map(([branch, branchRows]) => {
@@ -51,7 +32,6 @@ export const TableContent: React.FC<TableContentProps> = ({
                     tableHeaderHeight={tableHeadHeight}
                     root={tableRef.current}
                     row={row}
-                    instance={instance}
                     isLast={i === branchRows.length - 1}
                   />
                 </Fragment>

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

@@ -1,12 +1,12 @@
 import cx from 'classnames'
-import React, { PropsWithChildren } from 'react'
+import React, { PropsWithChildren, ReactNode } from 'react'
 import { useInView } from 'react-intersection-observer'
-import { InstanceProp } from '../../../util/interfaces'
 import styles from '../styles.module.scss'
 
-interface WorkspaceRowGroupProps extends InstanceProp {
+interface WorkspaceRowGroupProps {
   root: HTMLElement | null
   tableHeaderHeight: number
+  children: ReactNode
 }
 
 export const WorkspaceRowGroup: React.FC<

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

@@ -1,7 +1,12 @@
 import { Experiment } from 'dvc/src/experiments/webview/contract'
-import React, { DragEvent, useRef, useEffect } from 'react'
+import React, { DragEvent, useRef, useEffect, memo } from 'react'
 import { useSelector, useDispatch } from 'react-redux'
-import { Table, Header } from '@tanstack/react-table'
+import {
+  Header,
+  HeaderGroup,
+  Updater,
+  ColumnOrderState
+} from '@tanstack/react-table'
 import { MergedHeaderGroups } from './MergeHeaderGroups'
 import { setDropTarget } from '../../../state/headerDropTargetSlice'
 import { ExperimentsState } from '../../../store'
@@ -15,22 +20,22 @@ import styles from '../styles.module.scss'
 import { reorderColumns } from '../../../util/messages'
 
 interface TableHeadProps {
-  instance: Table<Experiment>
+  headerGroups: HeaderGroup<Experiment>[]
+  columnOrder: string[]
+  setColumnOrder: (updater: Updater<ColumnOrderState>) => void
   root: HTMLElement | null
   setExpColumnNeedsShadow: (needsShadow: boolean) => void
   setTableHeadHeight: (height: number) => void
 }
 
-export const TableHead = ({
-  instance,
+const THead = ({
+  columnOrder,
+  headerGroups,
+  setColumnOrder,
   root,
   setExpColumnNeedsShadow,
   setTableHeadHeight
 }: TableHeadProps) => {
-  const { setColumnOrder, getHeaderGroups, getAllLeafColumns } = instance
-  const headerGroups = getHeaderGroups()
-  const allColumns = getAllLeafColumns()
-
   const headerDropTargetId = useSelector(
     (state: ExperimentsState) => state.headerDropTarget
   )
@@ -58,7 +63,7 @@ export const TableHead = ({
     )
     if (displacerHeader) {
       draggingIds.current = leafColumnIds(displacerHeader)
-      fullColumnOrder.current = allColumns.map(({ id }) => id)
+      fullColumnOrder.current = columnOrder
     }
   }
 
@@ -130,3 +135,5 @@ export const TableHead = ({
     </thead>
   )
 }
+
+export const TableHead = memo(THead)