feat: add metadata column type badge

determined-ai · Nov 1, 2024 · 9ddcc0e · 9ddcc0e
1 parent 21b0256
commit 9ddcc0e
Show file tree

Hide file tree

Showing 6 changed files with 92 additions and 63 deletions.
diff --git a/webui/react/src/components/ColumnPickerMenu.tsx b/webui/react/src/components/ColumnPickerMenu.tsx
@@ -1,3 +1,4 @@
+import Badge from 'hew/Badge';
 import Button from 'hew/Button';
 import Checkbox, { CheckboxChangeEvent } from 'hew/Checkbox';
 import Dropdown from 'hew/Dropdown';
@@ -10,9 +11,11 @@ import { Loadable } from 'hew/utils/loadable';
 import React, { ChangeEvent, useCallback, useMemo, useState } from 'react';
 import { FixedSizeList as List } from 'react-window';
 
-import { V1LocationType } from 'services/api-ts-sdk';
+import { runColumns } from 'pages/FlatRuns/columns';
+import { V1ColumnType, V1LocationType } from 'services/api-ts-sdk';
 import { ProjectColumn } from 'types';
 import { ensureArray } from 'utils/data';
+import { formatColumnKey, removeColumnTypePrefix } from 'utils/flatRun';
 
 import css from './ColumnPickerMenu.module.scss';
 
@@ -75,7 +78,7 @@ const ColumnPickerTab: React.FC<ColumnTabProps> = ({
   onVisibleColumnChange,
   onHeatmapSelectionRemove,
 }) => {
-  const checkedColumns = useMemo(
+  const checkedColumnNames = useMemo(
     () => (compare ? new Set(columnState.slice(0, pinnedColumnsCount)) : new Set(columnState)),
     [columnState, compare, pinnedColumnsCount],
   );
@@ -95,18 +98,21 @@ const ColumnPickerTab: React.FC<ColumnTabProps> = ({
   }, [searchString, totalColumns, tab]);
 
   const allFilteredColumnsChecked = useMemo(() => {
-    return filteredColumns.every((col) => columnState.includes(col.column));
+    return filteredColumns.every((col) => columnState.includes(formatColumnKey(col)));
   }, [columnState, filteredColumns]);
 
   const handleShowHideAll = useCallback(() => {
     const filteredColumnMap: Record<string, boolean> = filteredColumns.reduce(
-      (acc, col) => ({ ...acc, [col.column]: columnState.includes(col.column) }),
+      (acc, col) => ({
+        ...acc,
+        [formatColumnKey(col)]: columnState.includes(formatColumnKey(col)),
+      }),
       {},
     );
 
     const newColumns = allFilteredColumnsChecked
       ? columnState.filter((col) => !filteredColumnMap[col])
-      : [...new Set([...columnState, ...filteredColumns.map((col) => col.column)])];
+      : [...new Set([...columnState, ...filteredColumns.map((col) => formatColumnKey(col))])];
     const pinnedCount = allFilteredColumnsChecked
       ? // If uncheck something pinned, reduce the pinnedColumnsCount
         newColumns.filter((col) => columnState.indexOf(col) < pinnedColumnsCount).length
@@ -123,32 +129,34 @@ const ColumnPickerTab: React.FC<ColumnTabProps> = ({
 
   const handleColumnChange = useCallback(
     (event: CheckboxChangeEvent) => {
-      const { id, checked } = event.target;
-      if (id === undefined) return;
+      const { id: targetCol, checked } = event.target;
+
+      if (targetCol === undefined) return;
+
       if (compare) {
         // pin or unpin column
-        const newColumns = columnState.filter((c) => c !== id);
+        const newColumns = columnState.filter((c) => c !== targetCol);
         let pinnedCount = pinnedColumnsCount;
         if (checked) {
-          newColumns.splice(pinnedColumnsCount, 0, id);
+          newColumns.splice(pinnedColumnsCount, 0, targetCol);
           pinnedCount = Math.max(pinnedColumnsCount + 1, 0);
         } else {
-          newColumns.splice(pinnedColumnsCount - 1, 0, id);
+          newColumns.splice(pinnedColumnsCount - 1, 0, targetCol);
           pinnedCount = Math.max(pinnedColumnsCount - 1, 0);
         }
         onVisibleColumnChange?.(newColumns, pinnedCount);
       } else {
         let pinnedCount = pinnedColumnsCount;
         // If uncheck something pinned, reduce the pinnedColumnsCount
-        if (!checked && columnState.indexOf(id) < pinnedColumnsCount) {
+        if (!checked && columnState.indexOf(targetCol) < pinnedColumnsCount) {
           pinnedCount = Math.max(pinnedColumnsCount - 1, 0);
         }
         // If uncheck something had heatmap skipped, reset to heatmap visible
         if (!checked) {
-          onHeatmapSelectionRemove?.(id);
+          onHeatmapSelectionRemove?.(targetCol);
         }
         const newColumnSet = new Set(columnState);
-        checked ? newColumnSet.add(id) : newColumnSet.delete(id);
+        checked ? newColumnSet.add(targetCol) : newColumnSet.delete(targetCol);
         onVisibleColumnChange?.([...newColumnSet], pinnedCount);
       }
     },
@@ -165,24 +173,48 @@ const ColumnPickerTab: React.FC<ColumnTabProps> = ({
   const rows = useCallback(
     ({ index, style }: { index: number; style: React.CSSProperties }) => {
       const col = filteredColumns[index];
+      const colType =
+        (runColumns as readonly string[]).includes(col.column) &&
+        col.type === V1ColumnType.UNSPECIFIED
+          ? 'BOOLEAN'
+          : removeColumnTypePrefix(col.type);
+      const getColDisplayName = (col: ProjectColumn) => {
+        return (
+          <>
+            {col.displayName || col.column} <Badge text={colType} />
+          </>
+        );
+      };
+      const getId = () => {
+        if (col.location === V1LocationType.RUNMETADATA) return formatColumnKey(col);
+
+        return col.column;
+      };
+      const getChecked = () => {
+        if (col.location === V1LocationType.RUNMETADATA)
+          return checkedColumnNames.has(formatColumnKey(col));
+
+        return checkedColumnNames.has(col.column);
+      };
+
       return (
         <div
           className={css.rows}
           data-test="row"
-          data-test-id={col.column}
-          key={col.column}
+          data-test-id={getId()}
+          key={getId()}
           style={style}>
           <Checkbox
-            checked={checkedColumns.has(col.column)}
+            checked={getChecked()}
             data-test="checkbox"
-            id={col.column}
+            id={getId()}
             onChange={handleColumnChange}>
-            {col.displayName || col.column}
+            {getColDisplayName(col)}
           </Checkbox>
         </div>
       );
     },
-    [filteredColumns, checkedColumns, handleColumnChange],
+    [filteredColumns, checkedColumnNames, handleColumnChange],
   );
 
   return (

diff --git a/webui/react/src/components/Searches/columns.ts b/webui/react/src/components/Searches/columns.ts
@@ -130,7 +130,6 @@ export const getColumnDefs = ({
   },
   checkpointCount: {
     id: 'checkpointCount',
-    isNumerical: true,
     renderer: (record: ExperimentWithTrial) => ({
       allowOverlay: false,
       data: Number(record.experiment.checkpoints),
@@ -143,7 +142,6 @@ export const getColumnDefs = ({
   },
   checkpointSize: {
     id: 'checkpointSize',
-    isNumerical: true,
     renderer: (record: ExperimentWithTrial) =>
       handleEmptyCell(record.experiment.checkpointSize, (data) => ({
         allowOverlay: false,
@@ -174,7 +172,6 @@ export const getColumnDefs = ({
   },
   duration: {
     id: 'duration',
-    isNumerical: true,
     renderer: (record: ExperimentWithTrial) =>
       handleEmptyCell(record.experiment.duration, () => ({
         allowOverlay: false,
@@ -297,7 +294,6 @@ export const getColumnDefs = ({
   },
   numTrials: {
     id: 'numTrials',
-    isNumerical: true,
     renderer: (record: ExperimentWithTrial) => ({
       allowOverlay: false,
       data: record.experiment.numTrials,
@@ -335,7 +331,6 @@ export const getColumnDefs = ({
   },
   searcherMetric: {
     id: 'searcherMetric',
-    isNumerical: false,
     renderer: (record: ExperimentWithTrial) =>
       handleEmptyCell(record.experiment.searcherMetric, (data) => ({
         allowOverlay: false,
@@ -361,7 +356,6 @@ export const getColumnDefs = ({
   },
   startTime: {
     id: 'startTime',
-    isNumerical: true,
     renderer: (record: ExperimentWithTrial) => ({
       allowOverlay: false,
       copyData: getTimeInEnglish(new Date(record.experiment.startTime)),
@@ -440,7 +434,6 @@ export const getColumnDefs = ({
 export const searcherMetricsValColumn = (columnWidth?: number): ColumnDef<ExperimentWithTrial> => {
   return {
     id: 'searcherMetricsVal',
-    isNumerical: true,
     renderer: (record: ExperimentWithTrial) =>
       handleEmptyCell(record.bestTrial?.searcherMetricsVal, (data) => ({
         allowOverlay: false,

diff --git a/webui/react/src/pages/F_ExpList/expListColumns.ts b/webui/react/src/pages/F_ExpList/expListColumns.ts
@@ -136,7 +136,6 @@ export const getColumnDefs = ({
   },
   checkpointCount: {
     id: 'checkpointCount',
-    isNumerical: true,
     renderer: (record: ExperimentWithTrial) => ({
       allowOverlay: false,
       data: Number(record.experiment.checkpoints),
@@ -149,7 +148,6 @@ export const getColumnDefs = ({
   },
   checkpointSize: {
     id: 'checkpointSize',
-    isNumerical: true,
     renderer: (record: ExperimentWithTrial) => ({
       allowOverlay: false,
       copyData: record.experiment.checkpointSize
@@ -176,7 +174,6 @@ export const getColumnDefs = ({
   },
   duration: {
     id: 'duration',
-    isNumerical: true,
     renderer: (record: ExperimentWithTrial) => ({
       allowOverlay: false,
       copyData: getDurationInEnglish(record.experiment),
@@ -305,7 +302,6 @@ export const getColumnDefs = ({
   },
   numTrials: {
     id: 'numTrials',
-    isNumerical: true,
     renderer: (record: ExperimentWithTrial) => ({
       allowOverlay: false,
       data: record.experiment.numTrials,
@@ -346,7 +342,6 @@ export const getColumnDefs = ({
   },
   searcherMetric: {
     id: 'searcherMetric',
-    isNumerical: false,
     renderer: (record: ExperimentWithTrial) => {
       const sMetric = record.experiment.searcherMetric ?? '';
       return {
@@ -374,7 +369,6 @@ export const getColumnDefs = ({
   },
   startTime: {
     id: 'startTime',
-    isNumerical: true,
     renderer: (record: ExperimentWithTrial) => ({
       allowOverlay: false,
       copyData: getTimeInEnglish(new Date(record.experiment.startTime)),
@@ -456,7 +450,6 @@ export const searcherMetricsValColumn = (
 ): ColumnDef<ExperimentWithTrial> => {
   return {
     id: 'searcherMetricsVal',
-    isNumerical: true,
     renderer: (record: ExperimentWithTrial) => {
       const sMetricValue = record.bestTrial?.searcherMetricsVal;
 

diff --git a/webui/react/src/pages/FlatRuns/FlatRuns.tsx b/webui/react/src/pages/FlatRuns/FlatRuns.tsx
@@ -75,6 +75,7 @@ import userSettings from 'stores/userSettings';
 import { DetailedUser, FlatRun, FlatRunAction, ProjectColumn, RunState } from 'types';
 import handleError from 'utils/error';
 import { combine } from 'utils/filterFormSet';
+import { formatColumnKey } from 'utils/flatRun';
 import { eagerSubscribe } from 'utils/observable';
 import { pluralizer } from 'utils/string';
 
@@ -295,7 +296,7 @@ const FlatRuns: React.FC<Props> = ({ projectId, workspaceId, searchId }) => {
     const projectColumnsMap: Loadable<Record<string, ProjectColumn>> = Loadable.map(
       projectColumns,
       (columns) => {
-        return columns.reduce((acc, col) => ({ ...acc, [col.column]: col }), {});
+        return columns.reduce((acc, col) => ({ ...acc, [formatColumnKey(col)]: col }), {});
       },
     );
     const columnDefs = getColumnDefs({
@@ -333,6 +334,7 @@ const FlatRuns: React.FC<Props> = ({ projectId, workspaceId, searchId }) => {
           };
 
         let dataPath: string | undefined = undefined;
+        const columnDefKey = formatColumnKey(currentColumn);
         switch (currentColumn.location) {
           case V1LocationType.EXPERIMENT:
             dataPath = `experiment.${currentColumn.column}`;
@@ -373,11 +375,11 @@ const FlatRuns: React.FC<Props> = ({ projectId, workspaceId, searchId }) => {
               settings.heatmapOn &&
               !settings.heatmapSkipped.includes(currentColumn.column)
             ) {
-              columnDefs[currentColumn.column] = defaultNumberColumn(
-                currentColumn.column,
+              columnDefs[columnDefKey] = defaultNumberColumn(
+                columnDefKey,
                 currentColumn.displayName || currentColumn.column,
-                settings.columnWidths[currentColumn.column] ??
-                  defaultColumnWidths[currentColumn.column as RunColumn] ??
+                settings.columnWidths[columnDefKey] ??
+                  defaultColumnWidths[columnDefKey as RunColumn] ??
                   MIN_COLUMN_WIDTH,
                 dataPath,
                 {
@@ -386,45 +388,46 @@ const FlatRuns: React.FC<Props> = ({ projectId, workspaceId, searchId }) => {
                 },
               );
             } else {
-              columnDefs[currentColumn.column] = defaultNumberColumn(
-                currentColumn.column,
+              columnDefs[columnDefKey] = defaultNumberColumn(
+                columnDefKey,
                 currentColumn.displayName || currentColumn.column,
-                settings.columnWidths[currentColumn.column] ??
-                  defaultColumnWidths[currentColumn.column as RunColumn] ??
+                settings.columnWidths[columnDefKey] ??
+                  defaultColumnWidths[columnDefKey as RunColumn] ??
                   MIN_COLUMN_WIDTH,
                 dataPath,
+                undefined,
               );
             }
             break;
           }
           case V1ColumnType.DATE:
-            columnDefs[currentColumn.column] = defaultDateColumn(
-              currentColumn.column,
+            columnDefs[columnDefKey] = defaultDateColumn(
+              columnDefKey,
               currentColumn.displayName || currentColumn.column,
-              settings.columnWidths[currentColumn.column] ??
-                defaultColumnWidths[currentColumn.column as RunColumn] ??
+              settings.columnWidths[columnDefKey] ??
+                defaultColumnWidths[columnDefKey as RunColumn] ??
                 MIN_COLUMN_WIDTH,
               dataPath,
             );
             break;
           case V1ColumnType.ARRAY:
-            columnDefs[currentColumn.column] = defaultArrayColumn(
-              currentColumn.column,
+            columnDefs[columnDefKey] = defaultArrayColumn(
+              columnDefKey,
               currentColumn.displayName || currentColumn.column,
-              settings.columnWidths[currentColumn.column] ??
-                defaultColumnWidths[currentColumn.column as RunColumn] ??
+              settings.columnWidths[columnDefKey] ??
+                defaultColumnWidths[columnDefKey as RunColumn] ??
                 MIN_COLUMN_WIDTH,
               dataPath,
             );
             break;
           case V1ColumnType.TEXT:
           case V1ColumnType.UNSPECIFIED:
           default:
-            columnDefs[currentColumn.column] = defaultTextColumn(
-              currentColumn.column,
+            columnDefs[columnDefKey] = defaultTextColumn(
+              columnDefKey,
               currentColumn.displayName || currentColumn.column,
-              settings.columnWidths[currentColumn.column] ??
-                defaultColumnWidths[currentColumn.column as RunColumn] ??
+              settings.columnWidths[columnDefKey] ??
+                defaultColumnWidths[columnDefKey as RunColumn] ??
                 MIN_COLUMN_WIDTH,
               dataPath,
             );
@@ -434,7 +437,7 @@ const FlatRuns: React.FC<Props> = ({ projectId, workspaceId, searchId }) => {
             .getOrElse([])
             .find((h) => h.metricsName === currentColumn.column);
 
-          columnDefs[currentColumn.column] = searcherMetricsValColumn(
+          columnDefs[columnDefKey] = searcherMetricsValColumn(
             settings.columnWidths[currentColumn.column],
             heatmap && settings.heatmapOn && !settings.heatmapSkipped.includes(currentColumn.column)
               ? {
@@ -444,7 +447,7 @@ const FlatRuns: React.FC<Props> = ({ projectId, workspaceId, searchId }) => {
               : undefined,
           );
         }
-        return columnDefs[currentColumn.column];
+        return columnDefs[columnDefKey];
       })
       .flatMap((col) => (col ? [col] : []));
     return gridColumns;