Improve performance of experiment column collection (#4353)

Author: mattseddon

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

@@ -18,6 +18,12 @@ export const collectDeps = (acc: ColumnAccumulator, data: ExpData) => {
     return
   }
   for (const [file, { hash }] of Object.entries(deps)) {
+    const path = buildDepPath(file)
+    if (acc.collected.has(path)) {
+      return
+    }
+    acc.collected.add(path)
+
     const pathArray = getPathArray(file)
     const label = pathArray.pop() as string
 
@@ -27,7 +33,6 @@ export const collectDeps = (acc: ColumnAccumulator, data: ExpData) => {
       '/',
       label
     )
-    const path = buildDepPath(file)
 
     mergeAncestors(
       acc,

extension/src/experiments/columns/collect/index.test.ts

@@ -22,13 +22,13 @@ const mockedGetConfigValue = jest.mocked(getConfigValue)
 mockedGetConfigValue.mockImplementation(() => 5)
 
 describe('collectColumns', () => {
-  it('should return a value equal to the columns fixture when given the output fixture', () => {
-    const columns = collectColumns(outputFixture)
+  it('should return a value equal to the columns fixture when given the output fixture', async () => {
+    const columns = await collectColumns(outputFixture)
     expect(columns).toStrictEqual(columnsFixture)
   })
 
-  it('should output both params and metrics when both are present', () => {
-    const columns = collectColumns(
+  it('should output both params and metrics when both are present', async () => {
+    const columns = await collectColumns(
       generateTestExpShowOutput({
         metrics: {
           1: {
@@ -50,8 +50,8 @@ describe('collectColumns', () => {
     expect(metrics).toBeDefined()
   })
 
-  it('should omit params when none exist in the source data', () => {
-    const columns = collectColumns(
+  it('should omit params when none exist in the source data', async () => {
+    const columns = await collectColumns(
       generateTestExpShowOutput({
         metrics: {
           1: {
@@ -66,13 +66,13 @@ describe('collectColumns', () => {
     expect(metrics).toBeDefined()
   })
 
-  it('should return an empty array if no params and metrics are provided', () => {
-    const columns = collectColumns(generateTestExpShowOutput({}))
+  it('should return an empty array if no params and metrics are provided', async () => {
+    const columns = await collectColumns(generateTestExpShowOutput({}))
     expect(columns).toStrictEqual([])
   })
 
-  it('should aggregate multiple different field names', () => {
-    const columns = collectColumns(
+  it('should aggregate multiple different field names', async () => {
+    const columns = await collectColumns(
       generateTestExpShowOutput(
         {
           params: {
@@ -127,8 +127,8 @@ describe('collectColumns', () => {
     ])
   })
 
-  it('should create concatenated columns for nesting deeper than 5', () => {
-    const columns = collectColumns(
+  it('should create concatenated columns for nesting deeper than 5', async () => {
+    const columns = await collectColumns(
       generateTestExpShowOutput({
         params: {
           'params.yaml': {

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

@@ -28,38 +28,50 @@ const collectFromExperiment = (
 
   const { data } = experiment
   if (data) {
-    collectMetricsAndParams(acc, data)
-    collectDeps(acc, data)
+    return Promise.all([
+      collectMetricsAndParams(acc, data),
+      collectDeps(acc, data)
+    ])
   }
 }
 
-const collectFromExperiments = (
+const collectFromExperiments = async (
   acc: ColumnAccumulator,
   experiments?: ExpRange[] | null
 ) => {
   if (!experiments) {
     return
   }
+
+  const promises = []
   for (const { revs } of experiments) {
-    collectFromExperiment(acc, revs[0])
+    promises.push(collectFromExperiment(acc, revs[0]))
   }
+
+  await Promise.all(promises)
 }
 
-export const collectColumns = (output: ExpShowOutput): Column[] => {
-  const acc: ColumnAccumulator = {}
+export const collectColumns = async (
+  output: ExpShowOutput
+): Promise<Column[]> => {
+  const acc: ColumnAccumulator = { collected: new Set(), columns: {} }
 
-  acc.timestamp = timestampColumn
+  acc.columns.timestamp = timestampColumn
 
+  const promises = []
   for (const expState of output) {
     if (experimentHasError(expState)) {
       continue
     }
 
-    collectFromExperiment(acc, expState)
-    collectFromExperiments(acc, expState.experiments)
+    promises.push(
+      collectFromExperiment(acc, expState),
+      collectFromExperiments(acc, expState.experiments)
+    )
   }
+  await Promise.all(promises)
 
-  const columns = Object.values(acc)
+  const columns = Object.values(acc.columns)
   const hasNoData = isEqual(columns, [timestampColumn])
 
   return hasNoData ? [] : columns

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

@@ -29,28 +29,35 @@ const collectMetricOrParam = (
   label: string,
   value: Value
 ) => {
+  const path = buildMetricOrParamPath(type, ...pathArray, label)
+  if (acc.collected.has(path)) {
+    return
+  }
+  acc.collected.add(path)
+
   const { limitedDepthAncestors, limitedDepthLabel } = limitAncestorDepth(
     pathArray,
     METRIC_PARAM_SEPARATOR,
     FILE_SEPARATOR,
     label
   )
-  const path = buildMetricOrParamPath(type, ...limitedDepthAncestors, label)
+  const limitedPath = buildMetricOrParamPath(
+    type,
+    ...limitedDepthAncestors,
+    limitedDepthLabel
+  )
+
   mergeAncestors(
     acc,
     type,
-    path,
+    limitedPath,
     limitedDepthAncestors,
     (...pathArray: string[]) => buildMetricOrParamPath(type, ...pathArray)
   )
 
   collectColumn(
     acc,
-    buildMetricOrParamPath(
-      type,
-      ...(limitedDepthAncestors || pathArray),
-      label
-    ),
+    limitedPath,
     buildMetricOrParamPath(type, ...limitedDepthAncestors),
     [type, ...pathArray, label],
     limitedDepthLabel,
@@ -89,17 +96,28 @@ const walkMetricsOrParamsFile = (
   }
 }
 
+const collectMetricsOrParams = (
+  acc: ColumnAccumulator,
+  type: ColumnType,
+  metricsOrParams: MetricsOrParams | null
+): void => {
+  if (!metricsOrParams) {
+    return
+  }
+  walkMetricsOrParamsFile(acc, type, metricsOrParams)
+}
+
 export const collectMetricsAndParams = (
   acc: ColumnAccumulator,
   data: ExpData
 ) => {
   const { metrics, params } = data
-  if (metrics) {
-    walkMetricsOrParamsFile(acc, ColumnType.METRICS, metrics)
-  }
-  if (params) {
-    walkMetricsOrParamsFile(acc, ColumnType.PARAMS, params)
-  }
+  return Promise.all([
+    // eslint-disable-next-line sonarjs/no-use-of-empty-return-value
+    collectMetricsOrParams(acc, ColumnType.METRICS, metrics),
+    // eslint-disable-next-line sonarjs/no-use-of-empty-return-value
+    collectMetricsOrParams(acc, ColumnType.PARAMS, params)
+  ])
 }
 
 const collectChange = (

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

@@ -2,7 +2,10 @@ import { Column, ColumnType } from '../../webview/contract'
 import { Value } from '../../../cli/dvc/contract'
 import { ConfigKey, getConfigValue } from '../../../vscode/config'
 
-export type ColumnAccumulator = Record<string, Column>
+export type ColumnAccumulator = {
+  columns: Record<string, Column>
+  collected: Set<string>
+}
 
 const joinPathArray = (
   pathSegments: string[],
@@ -76,10 +79,10 @@ const mergeParentColumnByPath = (
   parentPath: string,
   label: string
 ) => {
-  if (acc[path]) {
-    acc[path].hasChildren = true
+  if (acc.columns[path]) {
+    acc.columns[path].hasChildren = true
   } else {
-    acc[path] = {
+    acc.columns[path] = {
       hasChildren: true,
       label,
       parentPath,
@@ -106,17 +109,15 @@ export const mergeAncestors = (
   limitedDepthAncestors: string[],
   join: (...pathArray: string[]) => string
 ) => {
-  if (!acc[path]) {
-    for (let i = 1; i <= limitedDepthAncestors.length; i++) {
-      const pathArray = limitedDepthAncestors.slice(0, i)
-      mergeParentColumnByPath(
-        acc,
-        type,
-        join(...pathArray),
-        join(...pathArray.slice(0, -1)),
-        pathArray[pathArray.length - 1]
-      )
-    }
+  for (let i = 1; i <= limitedDepthAncestors.length; i++) {
+    const pathArray = limitedDepthAncestors.slice(0, i)
+    mergeParentColumnByPath(
+      acc,
+      type,
+      join(...pathArray),
+      join(...pathArray.slice(0, -1)),
+      pathArray[pathArray.length - 1]
+    )
   }
 }
 
@@ -146,8 +147,11 @@ export const collectColumn = (
   label: string,
   value: Value
 ) => {
-  if (acc[path]) {
-    return
-  }
-  acc[path] = buildValueColumn(path, parentPath, pathArray, label, value)
+  acc.columns[path] = buildValueColumn(
+    path,
+    parentPath,
+    pathArray,
+    label,
+    value
+  )
 }

extension/src/experiments/columns/model.test.ts

@@ -133,6 +133,7 @@ describe('ColumnsModel', () => {
     )
     await model.transformAndSet(deeplyNestedOutputFixture)
     expect(mockedGetConfigValue).toHaveBeenCalled()
+
     expect(model.getSelected()).toStrictEqual(deeplyNestedColumnsWithHeightOf1)
   })
 

extension/src/test/fixtures/expShow/deeplyNested/maxHeight.ts

@@ -352,20 +352,21 @@ export const deeplyNestedColumnsWithHeightOf2: Column[] = [
 export const deeplyNestedColumnsWithHeightOf1: Column[] = [
   timestampColumn,
   {
+    firstValueType: 'string',
     hasChildren: false,
     label: 'params.yaml:nested1.doubled',
     parentPath: 'params',
-    path: 'params:doubled',
+    path: 'params:params.yaml:nested1.doubled',
     pathArray: ['params', 'params.yaml', 'nested1', 'doubled'],
-    type: ColumnType.PARAMS,
-    firstValueType: 'string'
+    type: ColumnType.PARAMS
   },
   {
+    firstValueType: 'string',
     hasChildren: false,
     label:
       'params.yaml:nested1.nested2.nested3.nested4.nested5.nested6.nested7',
     parentPath: 'params',
-    path: 'params:nested7',
+    path: 'params:params.yaml:nested1.nested2.nested3.nested4.nested5.nested6.nested7',
     pathArray: [
       'params',
       'params.yaml',
@@ -377,14 +378,14 @@ export const deeplyNestedColumnsWithHeightOf1: Column[] = [
       'nested6',
       'nested7'
     ],
-    type: ColumnType.PARAMS,
-    firstValueType: 'string'
+    type: ColumnType.PARAMS
   },
   {
+    firstValueType: 'string',
     hasChildren: false,
     label: 'params.yaml:nested1.nested2.nested3.nested4.nested5b.nested6',
     parentPath: 'params',
-    path: 'params:nested6',
+    path: 'params:params.yaml:nested1.nested2.nested3.nested4.nested5b.nested6',
     pathArray: [
       'params',
       'params.yaml',
@@ -395,16 +396,33 @@ export const deeplyNestedColumnsWithHeightOf1: Column[] = [
       'nested5b',
       'nested6'
     ],
-    type: ColumnType.PARAMS,
-    firstValueType: 'string'
+    type: ColumnType.PARAMS
+  },
+  {
+    firstValueType: 'string',
+    hasChildren: false,
+    label: 'params.yaml:nested1.nested2.nested3.nested4.nested5b.doubled',
+    parentPath: 'params',
+    path: 'params:params.yaml:nested1.nested2.nested3.nested4.nested5b.doubled',
+    pathArray: [
+      'params',
+      'params.yaml',
+      'nested1',
+      'nested2',
+      'nested3',
+      'nested4',
+      'nested5b',
+      'doubled'
+    ],
+    type: ColumnType.PARAMS
   },
   {
+    firstValueType: 'number',
     hasChildren: false,
     label: 'params.yaml:outlier',
     parentPath: 'params',
-    path: 'params:outlier',
+    path: 'params:params.yaml:outlier',
     pathArray: ['params', 'params.yaml', 'outlier'],
-    type: ColumnType.PARAMS,
-    firstValueType: 'number'
+    type: ColumnType.PARAMS
   }
 ]