feat: Add ground truth dividers to Annotations table and correct all …

…annotation counts (#949)

#433

Demo:
https://dev-annotation-dividers.cryoet.dev.si.czi.technology/runs/241

@codemonkey800's original PR:
3e17d89

- Displays a divider before each ground truth/non ground truth sequence
of rows in the table.
   - Also displays dividers with 0 count if that section is empty.
   - See the issue for discussions re: various edge cases.
- Changes all counts (total, filtered, ground truth, and non ground
truth) to count the unique `shape_type` and `annotation_id` pairs
represented in `annotation_files` using `distinct_on` to match the
number of rows displayed in the tables.
   - Removes the need to sum counts on the frontend.

I'm not super satisfied with having to have 2 callbacks in
AnnotationsTable.tsx for displaying dividers, but couldn't really think
of something better that still worked well with react-table and
satisfied the UX requirements for displaying dividers when there are no
rows for that section. Making the data type a union between `Annotation`
and something else that represents a divider messes up the types of the
column helper, and allowing `<Table>` to take a wrapper component would
be the same problem but probably even uglier because it would have to
take 2, one to wrap `<table>` and one to wrap the `<tr>`s. If I had to
I'd probably pick the union type option and just send down the actual
divider `ReactNode`s in the `data` array as rows, but that just feels
like abusing react-table...

frontend/packages/data-portal/app/components/Run/AnnotationTable.tsx

@@ -1,9 +1,15 @@
 /* eslint-disable react/no-unstable-nested-components */
 
 import { Button, Icon } from '@czi-sds/components'
-import { ColumnDef, createColumnHelper } from '@tanstack/react-table'
-import { range } from 'lodash-es'
-import { ComponentProps, useCallback, useMemo } from 'react'
+import { useSearchParams } from '@remix-run/react'
+import {
+  ColumnDef,
+  createColumnHelper,
+  Row,
+  Table,
+} from '@tanstack/react-table'
+import { range, toNumber } from 'lodash-es'
+import { ComponentProps, ReactNode, useCallback, useMemo } from 'react'
 
 import { AuthorList } from 'app/components/AuthorList'
 import { I18n } from 'app/components/I18n'
@@ -16,6 +22,7 @@ import {
   MethodType,
 } from 'app/constants/methodTypes'
 import { MAX_PER_PAGE } from 'app/constants/pagination'
+import { QueryParams } from 'app/constants/query'
 import { AnnotationTableWidths } from 'app/constants/table'
 import { TestIds } from 'app/constants/testIds'
 import { useDownloadModalQueryParamState } from 'app/hooks/useDownloadModalQueryParamState'
@@ -65,7 +72,8 @@ function ConfidenceValue({ value }: { value: number }) {
 
 export function AnnotationTable() {
   const { isLoadingDebounced } = useIsLoading()
-  const { run } = useRunById()
+  const [searchParams] = useSearchParams()
+  const { run, annotationFilesAggregates } = useRunById()
   const { toggleDrawer } = useMetadataDrawer()
   const { setActiveAnnotation } = useAnnotation()
   const { t } = useI18n()
@@ -382,11 +390,99 @@ export function AnnotationTable() {
     [run.annotation_table],
   )
 
+  const currentPage = toNumber(
+    searchParams.get(QueryParams.AnnotationsPage) ?? 1,
+  )
+
+  /**
+   * Attaches divider(s) before a row.
+   *  - The ground truth divider can only be attached to the first row.
+   *    - The first page always shows the divider, even if there are 0 ground truth rows.
+   *    - Subsequent pages only show the divider if the first row is ground truth.
+   *  - The non ground truth divider is attached to the first non ground truth row.
+   */
+  const getGroundTruthDividersForRow = (
+    table: Table<Annotation>,
+    row: Row<Annotation>,
+  ): ReactNode => {
+    return (
+      <>
+        {row.index === 0 &&
+          (currentPage === 1 || row.original.ground_truth_status) && (
+            <RowDivider
+              groundTruth
+              count={annotationFilesAggregates.groundTruthCount}
+            />
+          )}
+
+        {row.id ===
+          table.getRowModel().rows.find((r) => !r.original.ground_truth_status)
+            ?.id && (
+          <RowDivider
+            groundTruth={false}
+            count={annotationFilesAggregates.otherCount}
+          />
+        )}
+      </>
+    )
+  }
+
+  /**
+   * Adds divider(s) to the end of the table when there are no rows to attach to.
+   *  - The ground truth divider won't have a row to attach to only if there are 0 rows in the
+   *    table.
+   *  - The non ground truth divider won't have a row to attach to only if there are 0 non ground
+   *    truth rows in the table and this is the last page.
+   */
+  const getGroundTruthDividersWhenNoRows = (): ReactNode => {
+    return (
+      <>
+        {annotationFilesAggregates.filteredCount === 0 && (
+          <RowDivider groundTruth count={0} />
+        )}
+
+        {annotationFilesAggregates.otherCount === 0 &&
+          annotationFilesAggregates.filteredCount <=
+            currentPage * MAX_PER_PAGE && (
+            <RowDivider groundTruth={false} count={0} />
+          )}
+      </>
+    )
+  }
+
   return (
     <PageTable
       data={isLoadingDebounced ? LOADING_ANNOTATIONS : annotations}
       columns={columns}
+      getBeforeRowElement={getGroundTruthDividersForRow}
+      getAfterTableElement={getGroundTruthDividersWhenNoRows}
       hoverType="none"
     />
   )
 }
+
+function RowDivider({
+  groundTruth,
+  count,
+}: {
+  groundTruth: boolean
+  count: number
+}) {
+  const { t } = useI18n()
+
+  return (
+    <tr
+      className="bg-sds-gray-100 border-t border-sds-gray-300"
+      data-testid={TestIds.AnnotationTableDivider}
+    >
+      <td
+        className="text-sds-header-xxs text-sds-gray-500 p-sds-s leading-sds-header-xs"
+        colSpan={1000}
+      >
+        {t(groundTruth ? 'groundTruthAnnotations' : 'otherAnnotations', {
+          count,
+        })}
+      </td>
+    </tr>
+  )
+}

frontend/packages/data-portal/app/components/Run/RunHeader.tsx

@@ -42,7 +42,7 @@ function FileSummary({ data }: { data: FileSummaryData[] }) {
 
 export function RunHeader() {
   const multipleTomogramsEnabled = useFeatureFlag('multipleTomograms')
-  const { run } = useRunById()
+  const { run, annotationFilesAggregates } = useRunById()
   const { toggleDrawer } = useMetadataDrawer()
   const { t } = useI18n()
 
@@ -61,11 +61,7 @@ export function RunHeader() {
       (stats) => stats.tomograms_aggregate.aggregate?.count ?? 0,
     ),
   )
-  const annotationsCount = sum(
-    run.tomogram_stats.flatMap(
-      (stats) => stats.annotations_aggregate.aggregate?.count ?? 0,
-    ),
-  )
+  const annotationsCount = annotationFilesAggregates.totalCount
 
   return (
     <PageHeader

frontend/packages/data-portal/app/components/Table/PageTable.tsx

@@ -1,19 +1,13 @@
-import { ComponentProps } from 'react'
-
 import { cns } from 'app/utils/cns'
 
-import { Table } from './Table'
+import { Table, TableProps } from './Table'
 
-// FIXME: refactor this to be useMemo-based instead?
-export function PageTable<T>({
-  hoverType,
-  ...props
-}: Pick<
-  ComponentProps<typeof Table<T>>,
-  'data' | 'columns' | 'onTableRowClick'
-> & {
+export interface PageTableProps<T> extends TableProps<T> {
   hoverType?: 'group' | 'none'
-}) {
+}
+
+// FIXME: refactor this to be useMemo-based instead?
+export function PageTable<T>({ hoverType, ...props }: PageTableProps<T>) {
   return (
     <Table
       classes={{

frontend/packages/data-portal/app/components/Table/Table.tsx

@@ -2,7 +2,7 @@ import {
   CellHeader,
   Table as SDSTable,
   TableHeader,
-  TableProps,
+  TableProps as SDSTableProps,
   TableRow,
 } from '@czi-sds/components'
 import TableContainer from '@mui/material/TableContainer'
@@ -11,20 +11,16 @@ import {
   flexRender,
   getCoreRowModel,
   Row,
+  Table as ReactTable,
   useReactTable,
 } from '@tanstack/react-table'
+import { Fragment, ReactNode } from 'react'
 
 import { ErrorBoundary } from 'app/components/ErrorBoundary'
 import { useLayout } from 'app/context/Layout.context'
 import { cns } from 'app/utils/cns'
 
-export function Table<T>({
-  classes,
-  columns,
-  data,
-  tableProps,
-  onTableRowClick,
-}: {
+export interface TableProps<T> {
   classes?: {
     body?: string
     cell?: string
@@ -35,9 +31,21 @@ export function Table<T>({
   }
   columns: ColumnDef<T>[]
   data: T[]
-  tableProps?: TableProps
+  tableProps?: SDSTableProps
+  getBeforeRowElement?: (table: ReactTable<T>, row: Row<T>) => ReactNode
+  getAfterTableElement?: (table: ReactTable<T>) => ReactNode
   onTableRowClick?(row: Row<T>): void
-}) {
+}
+
+export function Table<T>({
+  classes,
+  columns,
+  data,
+  tableProps,
+  getBeforeRowElement,
+  getAfterTableElement,
+  onTableRowClick,
+}: TableProps<T>) {
   const { hasFilters } = useLayout()
 
   const table = useReactTable<T>({
@@ -85,21 +93,26 @@ export function Table<T>({
 
         <tbody className={classes?.body}>
           {table.getRowModel().rows.map((row) => (
-            <TableRow
-              className={classes?.row}
-              key={row.id}
-              onClick={() => onTableRowClick?.(row)}
-            >
-              {row.getVisibleCells().map((cell) => (
-                <ErrorBoundary
-                  key={cell.id}
-                  logId={getLogId(`cell-${cell.id}`)}
-                >
-                  {flexRender(cell.column.columnDef.cell, cell.getContext())}
-                </ErrorBoundary>
-              ))}
-            </TableRow>
+            <Fragment key={row.id}>
+              {getBeforeRowElement?.(table, row)}
+
+              <TableRow
+                className={classes?.row}
+                onClick={() => onTableRowClick?.(row)}
+              >
+                {row.getVisibleCells().map((cell) => (
+                  <ErrorBoundary
+                    key={cell.id}
+                    logId={getLogId(`cell-${cell.id}`)}
+                  >
+                    {flexRender(cell.column.columnDef.cell, cell.getContext())}
+                  </ErrorBoundary>
+                ))}
+              </TableRow>
+            </Fragment>
           ))}
+
+          {getAfterTableElement?.(table)}
         </tbody>
       </SDSTable>
     </TableContainer>

frontend/packages/data-portal/app/constants/testIds.ts

@@ -1,5 +1,6 @@
 export enum TestIds {
   AnnotationId = 'annotation-id',
+  AnnotationTableDivider = 'annotation-table-divider',
   EnvelopeIcon = 'envelope-icon',
   MetadataDrawer = 'metadata-drawer',
   MetadataDrawerCloseButton = 'metadata-drawer-close-button',