From 2fe5ddf76aef144b7a81cb1fb2054d0d11916466 Mon Sep 17 00:00:00 2001
From: c-bata <contact@c-bata.link>
Date: Sun, 25 Oct 2020 03:39:03 +0900
Subject: [PATCH] Filter trials by state

---
 .../static/components/DataGrid.tsx            | 194 ++++++++++++------
 .../static/components/StudyDetail.tsx         |   4 +-
 2 files changed, 133 insertions(+), 65 deletions(-)

diff --git a/optuna_dashboard/static/components/DataGrid.tsx b/optuna_dashboard/static/components/DataGrid.tsx
index 528b72f3..1f98b4c3 100644
--- a/optuna_dashboard/static/components/DataGrid.tsx
+++ b/optuna_dashboard/static/components/DataGrid.tsx
@@ -14,17 +14,17 @@ import {
 } from "@material-ui/core"
 import KeyboardArrowDownIcon from "@material-ui/icons/KeyboardArrowDown"
 import KeyboardArrowUpIcon from "@material-ui/icons/KeyboardArrowUp"
+import { Clear } from "@material-ui/icons"
 
 type Order = "asc" | "desc"
 
+const defaultRowsPerPageOption = [10, 50, 100, { label: "All", value: -1 }]
+
 const useStyles = makeStyles((theme: Theme) =>
   createStyles({
     root: {
       width: "100%",
     },
-    table: {
-      minWidth: 750,
-    },
     visuallyHidden: {
       border: 0,
       clip: "rect(0 0 0 0)",
@@ -36,27 +36,25 @@ const useStyles = makeStyles((theme: Theme) =>
       top: 20,
       width: 1,
     },
+    filterable: {
+      color: theme.palette.primary.main,
+      textDecoration: "underline",
+      cursor: "pointer",
+    },
   })
 )
 
-const defaultInitialRowsPerPage = 10
-const defaultRowsPerPageOption = [10, 50, 100, { label: "All", value: -1 }]
-
-function descendingComparator<T>(a: T, b: T, orderBy: keyof T) {
-  if (b[orderBy] < a[orderBy]) {
-    return -1
-  }
-  if (b[orderBy] > a[orderBy]) {
-    return 1
-  }
-  return 0
-}
-
 interface DataGridColumn<T> {
   field: keyof T
   label: string
-  sortable: boolean
-  toCellValue?: (dataIndex: number) => string | React.ReactNode
+  sortable?: boolean
+  filterable?: boolean
+  toCellValue?: (rowIndex: number) => string | React.ReactNode
+}
+
+interface RowFilter<T> {
+  field: keyof T
+  value: any
 }
 
 function DataGrid<T>(props: {
@@ -64,41 +62,31 @@ function DataGrid<T>(props: {
   rows: T[]
   keyField: keyof T
   dense?: boolean
-  collapseBody?: (dataIndex: number) => React.ReactNode
+  collapseBody?: (rowIndex: number) => React.ReactNode
   initialRowsPerPage?: number
   rowsPerPageOption?: Array<number | { value: number; label: string }>
 }) {
   const classes = useStyles()
-  const {
-    columns,
-    rows,
-    keyField,
-    dense,
-    collapseBody,
-    initialRowsPerPage,
-    rowsPerPageOption,
-  } = props
+  const { columns, rows, keyField, dense, collapseBody } = props
+  let { initialRowsPerPage, rowsPerPageOption } = props
   const [order, setOrder] = React.useState<Order>("asc")
   const [orderBy, setOrderBy] = React.useState<keyof T>(keyField)
   const [page, setPage] = React.useState(0)
-  const [rowsPerPage, setRowsPerPage] = React.useState(
-    initialRowsPerPage || defaultInitialRowsPerPage
-  )
+  const [filters, setFilters] = React.useState<RowFilter<T>[]>([])
 
-  const handleRequestSort = (
-    event: React.MouseEvent<unknown>,
-    property: keyof T
-  ) => {
-    const isAsc = orderBy === property && order === "asc"
-    setOrder(isAsc ? "desc" : "asc")
-    setOrderBy(property)
-  }
-  const createSortHandler = (property: keyof T) => (
-    event: React.MouseEvent<unknown>
-  ) => {
-    handleRequestSort(event, property)
+  const getRowIndex = (row: T): number => {
+    return rows.findIndex((row2) => row[keyField] === row2[keyField])
   }
 
+  // Pagination
+  rowsPerPageOption = rowsPerPageOption || defaultRowsPerPageOption
+  initialRowsPerPage = initialRowsPerPage // use first element as default
+    ? initialRowsPerPage
+    : isNumber(rowsPerPageOption[0])
+    ? rowsPerPageOption[0]
+    : rowsPerPageOption[0].value
+  const [rowsPerPage, setRowsPerPage] = React.useState(initialRowsPerPage)
+
   const handleChangePage = (event: unknown, newPage: number) => {
     setPage(newPage)
   }
@@ -110,23 +98,53 @@ function DataGrid<T>(props: {
     setPage(0)
   }
 
-  const emptyRows =
-    rowsPerPage - Math.min(rowsPerPage, rows.length - page * rowsPerPage)
+  // Filtering
+  const fieldAlreadyFiltered = (field: keyof T): boolean =>
+    filters.some((f) => f.field === field)
 
-  const sortedRows = stableSort<T>(rows, getComparator(order, orderBy))
-  const paginateRows =
+  const handleClickFilterCell = (field: keyof T, value: any) => {
+    if (fieldAlreadyFiltered(field)) {
+      return
+    }
+    const newFilters = [...filters, { field: field, value: value }]
+    setFilters(newFilters)
+  }
+
+  const clearFilter = (field: keyof T): void => {
+    setFilters(filters.filter((f) => f.field !== field))
+  }
+
+  const filteredRows = rows.filter((row) =>
+    filters.length === 0
+      ? true
+      : filters.some((f) => {
+          return row[f.field] === f.value
+        })
+  )
+
+  // Sorting
+  const createSortHandler = (property: keyof T) => (
+    event: React.MouseEvent<unknown>
+  ) => {
+    const isAsc = orderBy === property && order === "asc"
+    setOrder(isAsc ? "desc" : "asc")
+    setOrderBy(property)
+  }
+  const sortedRows = stableSort<T>(filteredRows, getComparator(order, orderBy))
+  const currentPageRows =
     rowsPerPage > 0
       ? sortedRows.slice(page * rowsPerPage, page * rowsPerPage + rowsPerPage)
       : sortedRows
+  const emptyRows =
+    rowsPerPage - Math.min(rowsPerPage, sortedRows.length - page * rowsPerPage)
 
   return (
     <div className={classes.root}>
       <TableContainer>
         <Table
-          className={classes.table}
           aria-labelledby="tableTitle"
           size={dense ? "small" : "medium"}
-          aria-label="enhanced table"
+          aria-label="data grid"
         >
           <TableHead>
             <TableRow>
@@ -154,19 +172,31 @@ function DataGrid<T>(props: {
                   ) : (
                     column.label
                   )}
+                  {fieldAlreadyFiltered(column.field) ? (
+                    <IconButton
+                      size={dense ? "small" : "medium"}
+                      color="inherit"
+                      onClick={(e) => {
+                        clearFilter(column.field)
+                      }}
+                    >
+                      <Clear />
+                    </IconButton>
+                  ) : null}
                 </TableCell>
               ))}
             </TableRow>
           </TableHead>
           <TableBody>
-            {paginateRows.map((row, index) => (
+            {currentPageRows.map((row, index) => (
               <DataGridRow<T>
                 columns={columns}
-                rowIndex={page * rowsPerPage + index}
+                rowIndex={getRowIndex(row)}
                 row={row}
                 keyField={keyField}
                 collapseBody={collapseBody}
-                key={`data-grid-row-${row[keyField]}`}
+                key={`${row[keyField]}`}
+                handleClickFilterCell={handleClickFilterCell}
               />
             ))}
             {emptyRows > 0 && (
@@ -178,9 +208,9 @@ function DataGrid<T>(props: {
         </Table>
       </TableContainer>
       <TablePagination
-        rowsPerPageOptions={rowsPerPageOption || defaultRowsPerPageOption}
+        rowsPerPageOptions={rowsPerPageOption}
         component="div"
-        count={rows.length}
+        count={filteredRows.length}
         rowsPerPage={rowsPerPage}
         page={page}
         onChangePage={handleChangePage}
@@ -195,9 +225,18 @@ function DataGridRow<T>(props: {
   rowIndex: number
   row: T
   keyField: keyof T
-  collapseBody?: (dataIndex: number) => React.ReactNode
+  collapseBody?: (rowIndex: number) => React.ReactNode
+  handleClickFilterCell: (field: keyof T, value: any) => void
 }) {
-  const { columns, rowIndex, row, keyField, collapseBody } = props
+  const classes = useStyles()
+  const {
+    columns,
+    rowIndex,
+    row,
+    keyField,
+    collapseBody,
+    handleClickFilterCell,
+  } = props
   const [open, setOpen] = React.useState(false)
 
   return (
@@ -214,13 +253,26 @@ function DataGridRow<T>(props: {
             </IconButton>
           </TableCell>
         ) : null}
-        {columns.map((column) => (
-          <TableCell key={`${row[keyField]}:${column.field}`}>
-            {column.toCellValue
-              ? column.toCellValue(rowIndex)
-              : row[column.field]}
-          </TableCell>
-        ))}
+        {columns.map((column) => {
+          const cellItem = column.toCellValue
+            ? column.toCellValue(rowIndex)
+            : row[column.field]
+
+          return column.filterable ? (
+            <TableCell
+              key={`${row[keyField]}:${column.field}`}
+              onClick={(e) => {
+                handleClickFilterCell(column.field, row[column.field])
+              }}
+            >
+              <div className={classes.filterable}>{cellItem}</div>
+            </TableCell>
+          ) : (
+            <TableCell key={`${row[keyField]}:${column.field}`}>
+              {cellItem}
+            </TableCell>
+          )
+        })}
       </TableRow>
       {collapseBody ? (
         <TableRow>
@@ -244,6 +296,16 @@ function getComparator<T>(
     : (a, b) => -descendingComparator<T>(a, b, orderBy)
 }
 
+function descendingComparator<T>(a: T, b: T, orderBy: keyof T) {
+  if (b[orderBy] < a[orderBy]) {
+    return -1
+  }
+  if (b[orderBy] > a[orderBy]) {
+    return 1
+  }
+  return 0
+}
+
 function stableSort<T>(array: T[], comparator: (a: T, b: T) => number) {
   const stabilizedThis = array.map((el, index) => [el, index] as [T, number])
   stabilizedThis.sort((a, b) => {
@@ -254,4 +316,10 @@ function stableSort<T>(array: T[], comparator: (a: T, b: T) => number) {
   return stabilizedThis.map((el) => el[0])
 }
 
+const isNumber = (
+  rowsPerPage: number | { value: number; label: string }
+): rowsPerPage is number => {
+  return typeof rowsPerPage === "number"
+}
+
 export { DataGrid, DataGridColumn }
diff --git a/optuna_dashboard/static/components/StudyDetail.tsx b/optuna_dashboard/static/components/StudyDetail.tsx
index 11985d59..dec45506 100644
--- a/optuna_dashboard/static/components/StudyDetail.tsx
+++ b/optuna_dashboard/static/components/StudyDetail.tsx
@@ -129,14 +129,14 @@ const TrialTable: FC<{ trials: Trial[] }> = ({ trials = [] }) => {
     {
       field: "state",
       label: "State",
-      sortable: false,
+      sortable: true,
+      filterable: true,
       toCellValue: (i) => trials[i].state.toString(),
     },
     { field: "value", label: "Value", sortable: true },
     {
       field: "params",
       label: "Params",
-      sortable: false,
       toCellValue: (i) =>
         trials[i].params.map((p) => p.name + ": " + p.value).join(", "),
     },