Add filter options to experiment table header context menu (#4415)

Author: mattseddon

extension/src/experiments/index.ts

@@ -7,7 +7,7 @@ import {
   workspace
 } from 'vscode'
 import omit from 'lodash.omit'
-import { addStarredToColumns } from './columns/like'
+import { ColumnLike, addStarredToColumns } from './columns/like'
 import { setContextForEditorTitleIcons } from './context'
 import { ExperimentsModel } from './model'
 import {
@@ -17,6 +17,7 @@ import {
 } from './model/quickPick'
 import { pickAndModifyParams } from './model/modify/quickPick'
 import {
+  pickColumnToFilter,
   pickFilterToAdd,
   pickFiltersToRemove
 } from './model/filterBy/quickPick'
@@ -302,11 +303,13 @@ export class Experiments extends BaseRepository<TableData> {
     return this.experiments.getFilters()
   }
 
-  public async addFilter() {
-    const columns = this.columns.getTerminalNodes()
-    const columnLikes = addStarredToColumns(columns)
+  public async addFilter(overrideColumn?: ColumnLike) {
+    const column = await this.pickColumnToFilter(overrideColumn)
+    if (!column) {
+      return
+    }
 
-    const filterToAdd = await pickFilterToAdd(columnLikes)
+    const filterToAdd = await pickFilterToAdd(column)
     if (!filterToAdd) {
       return
     }
@@ -624,6 +627,7 @@ export class Experiments extends BaseRepository<TableData> {
       () => this.selectColumns(),
       () => this.selectFirstColumns(),
       (branchesSelected: string[]) => this.selectBranches(branchesSelected),
+      (column: ColumnLike) => this.addFilter(column),
       () => this.data.update()
     )
 
@@ -687,4 +691,13 @@ export class Experiments extends BaseRepository<TableData> {
       return setUserConfigValue(ConfigKey.DO_NOT_INFORM_MAX_PLOTTED, true)
     }
   }
+
+  private pickColumnToFilter(overrideColumn?: ColumnLike) {
+    if (overrideColumn) {
+      return overrideColumn
+    }
+    const columns = this.columns.getTerminalNodes()
+    const columnLikes = addStarredToColumns(columns)
+    return pickColumnToFilter(columnLikes)
+  }
 }

extension/src/experiments/model/filterBy/quickPick.test.ts

@@ -47,26 +47,15 @@ const mixedParam = {
 }
 
 describe('pickFilterToAdd', () => {
-  it('should return early if no param or metric is picked', async () => {
-    const params = [epochsParam]
-    mockedQuickPickValue.mockResolvedValueOnce(undefined)
-    const filter = await pickFilterToAdd(params)
-    expect(filter).toBeUndefined()
-  })
-
   it('should return early if no operator is picked', async () => {
-    const params = [epochsParam]
-    mockedQuickPickValue.mockResolvedValueOnce(epochsParam)
     mockedQuickPickValue.mockResolvedValueOnce(undefined)
-    const filter = await pickFilterToAdd(params)
+    const filter = await pickFilterToAdd(epochsParam)
     expect(filter).toBeUndefined()
   })
 
   it('should call showQuickPick with the correct operators for a mixed type param', async () => {
-    const params = [mixedParam]
-    mockedQuickPickValue.mockResolvedValueOnce(mixedParam)
     mockedQuickPickValue.mockResolvedValueOnce(undefined)
-    await pickFilterToAdd(params)
+    await pickFilterToAdd(mixedParam)
     expect(mockedQuickPickValue).toHaveBeenCalledWith(
       OPERATORS.filter(({ types }) => types.includes('number')),
       {
@@ -76,19 +65,15 @@ describe('pickFilterToAdd', () => {
   })
 
   it('should return early if no value is provided', async () => {
-    const params = [epochsParam]
-    mockedQuickPickValue.mockResolvedValueOnce(epochsParam)
     mockedQuickPickValue.mockResolvedValueOnce('==')
     mockedGetInput.mockResolvedValueOnce(undefined)
-    const filter = await pickFilterToAdd(params)
+    const filter = await pickFilterToAdd(epochsParam)
     expect(filter).toBeUndefined()
   })
 
   it('should return without asking for a value when a boolean param is selected', async () => {
-    const params = [boolParam]
-    mockedQuickPickValue.mockResolvedValueOnce(boolParam)
     mockedQuickPickValue.mockResolvedValueOnce(Operator.IS_TRUE)
-    const filter = await pickFilterToAdd(params)
+    const filter = await pickFilterToAdd(boolParam)
     expect(filter).toStrictEqual({
       operator: Operator.IS_TRUE,
       path: boolParam.path,
@@ -104,11 +89,9 @@ describe('pickFilterToAdd', () => {
   })
 
   it('should return a filter definition if all of the steps are completed', async () => {
-    const params = [epochsParam]
-    mockedQuickPickValue.mockResolvedValueOnce(epochsParam)
     mockedQuickPickValue.mockResolvedValueOnce('==')
     mockedGetInput.mockResolvedValueOnce('5')
-    const filter = await pickFilterToAdd(params)
+    const filter = await pickFilterToAdd(epochsParam)
     expect(filter).toStrictEqual({
       operator: '==',
       path: epochsParam.path,

extension/src/experiments/model/filterBy/quickPick.ts

@@ -115,19 +115,19 @@ const addFilterValue = async (path: string, operator: Operator) => {
   }
 }
 
-export const pickFilterToAdd = async (
+export const pickColumnToFilter = (
   columns: ColumnLike[] | undefined
-): Promise<FilterDefinition | undefined> => {
-  const picked = await pickFromColumnLikes(columns, {
+): Thenable<ColumnLike | undefined> =>
+  pickFromColumnLikes(columns, {
     title: Title.SELECT_PARAM_OR_METRIC_FILTER
   })
-  if (!picked) {
-    return
-  }
 
+export const pickFilterToAdd = async ({
+  firstValueType,
+  path
+}: ColumnLike): Promise<FilterDefinition | undefined> => {
   const typedOperators = OPERATORS.filter(
-    operator =>
-      picked.firstValueType && operator.types.includes(picked.firstValueType)
+    operator => firstValueType && operator.types.includes(firstValueType)
   )
 
   const operator = await quickPickValue<Operator>(typedOperators, {
@@ -144,12 +144,12 @@ export const pickFilterToAdd = async (
   ) {
     return {
       operator,
-      path: picked.path,
+      path,
       value: undefined
     }
   }
 
-  return addFilterValue(picked.path, operator)
+  return addFilterValue(path, operator)
 }
 
 export const pickFiltersToRemove = (

extension/src/experiments/webview/messages.ts

@@ -23,6 +23,8 @@ import { ConfigKey, setConfigValue } from '../../vscode/config'
 import { NUM_OF_COMMITS_TO_INCREASE } from '../../cli/dvc/constants'
 import { Pipeline } from '../../pipeline'
 import { collectColumnsWithChangedValues } from '../columns/collect'
+import { ColumnLike } from '../columns/like'
+import { getFilterId } from '../model/filterBy'
 
 export class WebviewMessages {
   private readonly dvcRoot: string
@@ -35,7 +37,7 @@ export class WebviewMessages {
   private readonly notifyChanged: () => void
   private readonly selectColumns: () => Promise<void>
   private readonly selectFirstColumns: () => Promise<void>
-
+  private readonly addFilter: (column: ColumnLike) => Promise<void>
   private readonly selectBranches: (
     branchesSelected: string[]
   ) => Promise<string[] | undefined>
@@ -54,6 +56,7 @@ export class WebviewMessages {
     selectBranches: (
       branchesSelected: string[]
     ) => Promise<string[] | undefined>,
+    addFilter: (column: ColumnLike) => Promise<void>,
     update: () => Promise<void>
   ) {
     this.dvcRoot = dvcRoot
@@ -65,6 +68,7 @@ export class WebviewMessages {
     this.selectColumns = selectColumns
     this.selectFirstColumns = selectFirstColumns
     this.selectBranches = selectBranches
+    this.addFilter = addFilter
     this.update = update
   }
 
@@ -101,6 +105,11 @@ export class WebviewMessages {
         return this.addColumnSort(message.payload)
       case MessageFromWebviewType.REMOVE_COLUMN_SORT:
         return this.removeColumnSort(message.payload)
+      case MessageFromWebviewType.FILTER_COLUMN:
+        return this.addColumnFilter(message.payload)
+      case MessageFromWebviewType.REMOVE_COLUMN_FILTERS:
+        return this.removeColumnFilter(message.payload)
+
       case MessageFromWebviewType.APPLY_EXPERIMENT_TO_WORKSPACE:
         return commands.executeCommand(
           RegisteredCliCommands.EXPERIMENT_VIEW_APPLY,
@@ -434,6 +443,40 @@ export class WebviewMessages {
     return this.notifyChanged()
   }
 
+  private addColumnFilter(selectedPath: string) {
+    const column = this.columns
+      .getTerminalNodes()
+      .find(({ path }) => path === selectedPath)
+
+    if (!column) {
+      return
+    }
+
+    void this.addFilter(column)
+
+    sendTelemetryEvent(
+      EventName.VIEWS_EXPERIMENTS_TABLE_FILTER_COLUMN,
+      undefined,
+      undefined
+    )
+  }
+
+  private removeColumnFilter(selectedPath: string) {
+    for (const filter of this.experiments.getFilters()) {
+      if (filter.path === selectedPath) {
+        const id = getFilterId(filter)
+        this.experiments.removeFilter(id)
+      }
+    }
+    this.notifyChanged()
+
+    sendTelemetryEvent(
+      EventName.VIEWS_EXPERIMENTS_TABLE_REMOVE_COLUMN_FILTER,
+      undefined,
+      undefined
+    )
+  }
+
   private focusSortsTree() {
     const commandPromise = commands.executeCommand(
       'dvc.views.experimentsSortByTree.focus'

extension/src/telemetry/constants.ts

@@ -37,6 +37,7 @@ export const EventName = Object.assign(
       'views.experimentTable.toggleStars',
     VIEWS_EXPERIMENTS_TABLE_EXPERIMENT_TOGGLE:
       'views.experimentTable.toggleStatus',
+    VIEWS_EXPERIMENTS_TABLE_FILTER_COLUMN: 'views.experimentTable.filterColumn',
     VIEWS_EXPERIMENTS_TABLE_FOCUS_CHANGED:
       'views.experimentsTable.focusChanged',
     VIEWS_EXPERIMENTS_TABLE_FOCUS_FILTERS_TREE:
@@ -49,6 +50,8 @@ export const EventName = Object.assign(
     VIEWS_EXPERIMENTS_TABLE_OPEN_PARAMS_FILE:
       'views.experimentsTable.paramsFileOpened',
     VIEWS_EXPERIMENTS_TABLE_REFRESH: 'views.experimentsTable.refresh',
+    VIEWS_EXPERIMENTS_TABLE_REMOVE_COLUMN_FILTER:
+      'views.experimentsTable.removeColumnFilter',
     VIEWS_EXPERIMENTS_TABLE_REMOVE_COLUMN_SORT:
       'views.experimentsTable.columnSortRemoved',
     VIEWS_EXPERIMENTS_TABLE_RESET_COMMITS:
@@ -252,6 +255,8 @@ export interface IEventNamePropertyMapping {
   }
   [EventName.VIEWS_EXPERIMENTS_TABLE_RESET_COMMITS]: undefined
   [EventName.VIEWS_EXPERIMENTS_TABLE_CREATED]: undefined
+  [EventName.VIEWS_EXPERIMENTS_TABLE_FILTER_COLUMN]: undefined
+  [EventName.VIEWS_EXPERIMENTS_TABLE_REMOVE_COLUMN_FILTER]: undefined
   [EventName.VIEWS_EXPERIMENTS_TABLE_FOCUS_CHANGED]: WebviewFocusChangedProperties
   [EventName.VIEWS_EXPERIMENTS_TABLE_HIDE_COLUMN_PATH]: {
     path: string