feat: refactor datatable logic

frontend/components/dataset/dataset.tsx

@@ -7,11 +7,11 @@ import { useCallback, useEffect, useMemo, useState } from "react";
 
 import AddToLabelingQueuePopover from "@/components/traces/add-to-labeling-queue-popover";
 import { Button } from "@/components/ui/button.tsx";
+import { ColumnsMenu } from "@/components/ui/columns-menu";
 import DeleteSelectedRows from "@/components/ui/delete-selected-rows.tsx";
 import { InfiniteDataTable } from "@/components/ui/infinite-datatable";
 import { useInfiniteScroll } from "@/components/ui/infinite-datatable/hooks";
 import { DataTableStateProvider } from "@/components/ui/infinite-datatable/model/datatable-store";
-import ColumnsMenu from "@/components/ui/infinite-datatable/ui/columns-menu.tsx";
 import { type Datapoint, type Dataset as DatasetType } from "@/lib/dataset/types";
 import { useToast } from "@/lib/hooks/use-toast";
 import { cn } from "@/lib/utils";
@@ -71,8 +71,6 @@ const columns: ColumnDef<Datapoint>[] = [
   },
 ];
 
-const defaultDatasetColumnOrder = ["__row_selection", "index", "createdAt", "data", "target", "metadata"];
-
 const DatasetContent = ({ dataset, enableDownloadParquet, publicApiBaseUrl }: DatasetProps) => {
   const router = useRouter();
   const searchParams = useSearchParams();
@@ -289,7 +287,6 @@ const DatasetContent = ({ dataset, enableDownloadParquet, publicApiBaseUrl }: Da
         </div>
         <div className="flex overflow-hidden flex-1">
           <InfiniteDataTable
-            columns={columns}
             data={datapoints}
             hasMore={hasMore}
             isFetching={isFetching}
@@ -305,7 +302,6 @@ const DatasetContent = ({ dataset, enableDownloadParquet, publicApiBaseUrl }: Da
             }}
             onRowSelectionChange={setRowSelection}
             className="flex-1"
-            lockedColumns={["__row_selection"]}
             selectionPanel={(selectedRowIds) => (
               <div className="flex flex-col space-y-2">
                 <DeleteSelectedRows
@@ -316,13 +312,7 @@ const DatasetContent = ({ dataset, enableDownloadParquet, publicApiBaseUrl }: Da
               </div>
             )}
           >
-            <ColumnsMenu
-              lockedColumns={["__row_selection"]}
-              columnLabels={columns.map((column: ColumnDef<Datapoint>) => ({
-                id: column.id!,
-                label: typeof column.header === "string" ? column.header : column.id!,
-              }))}
-            />
+            <ColumnsMenu />
           </InfiniteDataTable>
         </div>
         <div className="flex text-secondary-foreground text-sm">{totalCount} datapoints</div>
@@ -356,7 +346,12 @@ const DatasetContent = ({ dataset, enableDownloadParquet, publicApiBaseUrl }: Da
 
 export default function Dataset(props: DatasetProps) {
   return (
-    <DataTableStateProvider storageKey="dataset-table" defaultColumnOrder={defaultDatasetColumnOrder}>
+    <DataTableStateProvider
+      storageKey="dataset-table"
+      columnDefs={columns}
+      enableRowSelection
+      lockedColumns={["__row_selection"]}
+    >
       <DatasetContent {...props} />
     </DataTableStateProvider>
   );

frontend/components/datasets/datasets.tsx

@@ -6,10 +6,10 @@ import { useParams, useRouter, useSearchParams } from "next/navigation";
 import { useCallback, useState } from "react";
 
 import { Button } from "@/components/ui/button";
+import { ColumnsMenu } from "@/components/ui/columns-menu";
 import DeleteSelectedRows from "@/components/ui/delete-selected-rows.tsx";
 import { useInfiniteScroll } from "@/components/ui/infinite-datatable/hooks";
 import { DataTableStateProvider } from "@/components/ui/infinite-datatable/model/datatable-store";
-import ColumnsMenu from "@/components/ui/infinite-datatable/ui/columns-menu.tsx";
 import DataTableFilter, { DataTableFilterList } from "@/components/ui/infinite-datatable/ui/datatable-filter";
 import { type ColumnFilter } from "@/components/ui/infinite-datatable/ui/datatable-filter/utils";
 import { DataTableSearch } from "@/components/ui/infinite-datatable/ui/datatable-search";
@@ -50,8 +50,6 @@ const columns: ColumnDef<DatasetInfo>[] = [
   },
 ];
 
-const defaultDatasetsColumnOrder = ["__row_selection", "id", "name", "datapointsCount", "createdAt"];
-
 const datasetsTableFilters: ColumnFilter[] = [
   {
     name: "ID",
@@ -209,7 +207,6 @@ function DatasetsContent() {
             enableRowSelection={true}
             getRowHref={(row) => `/project/${projectId}/datasets/${row.original.id}`}
             getRowId={(row: DatasetInfo) => row.id}
-            columns={columns}
             data={datasets}
             hasMore={hasMore}
             isFetching={isFetching}
@@ -219,7 +216,6 @@ function DatasetsContent() {
               rowSelection,
             }}
             onRowSelectionChange={setRowSelection}
-            lockedColumns={["__row_selection"]}
             emptyRow={filter.length === 0 && !search ? EmptyRow : undefined}
             selectionPanel={(selectedRowIds) => (
               <div className="flex flex-col space-y-2">
@@ -233,13 +229,7 @@ function DatasetsContent() {
           >
             <div className="flex flex-1 w-full space-x-2 pt-1">
               <DataTableFilter columns={datasetsTableFilters} />
-              <ColumnsMenu
-                lockedColumns={["__row_selection"]}
-                columnLabels={columns.map((column) => ({
-                  id: column.id!,
-                  label: typeof column.header === "string" ? column.header : column.id!,
-                }))}
-              />
+              <ColumnsMenu />
               <DataTableSearch className="mr-0.5" placeholder="Search by dataset name..." />
             </div>
             <DataTableFilterList />
@@ -252,7 +242,12 @@ function DatasetsContent() {
 
 export default function Datasets() {
   return (
-    <DataTableStateProvider storageKey="datasets-table" defaultColumnOrder={defaultDatasetsColumnOrder}>
+    <DataTableStateProvider
+      storageKey="datasets-table"
+      columnDefs={columns}
+      enableRowSelection
+      lockedColumns={["__row_selection"]}
+    >
       <DatasetsContent />
     </DataTableStateProvider>
   );

frontend/components/debugger-sessions/debugger-session-view/sidebar/runs-tab.tsx

@@ -5,7 +5,7 @@ import { useParams } from "next/navigation";
 import { useCallback, useEffect } from "react";
 import { useShallow } from "zustand/react/shallow";
 
-import { FETCH_SIZE, tracePickerColumnOrder, tracePickerColumns } from "@/components/traces/trace-picker/columns";
+import { FETCH_SIZE, tracePickerColumns } from "@/components/traces/trace-picker/columns";
 import { InfiniteDataTable } from "@/components/ui/infinite-datatable";
 import { DataTableStateProvider } from "@/components/ui/infinite-datatable/model/datatable-store";
 import RefreshButton from "@/components/ui/infinite-datatable/ui/refresh-button";
@@ -73,7 +73,6 @@ const RunsContent = () => {
   return (
     <InfiniteDataTable<TraceRow>
       className="w-full px-4 py-2"
-      columns={tracePickerColumns}
       data={historyRuns}
       getRowId={(t) => t.id}
       onRowClick={handleRowClick}
@@ -83,7 +82,6 @@ const RunsContent = () => {
       isLoading={isHistoryLoading}
       fetchNextPage={noop}
       estimatedRowHeight={36}
-      lockedColumns={["status"]}
     >
       <div className="flex gap-2 w-full items-center">
         <RefreshButton onClick={fetchHistory} variant="outline" />
@@ -94,7 +92,7 @@ const RunsContent = () => {
 
 export default function RunsTab() {
   return (
-    <DataTableStateProvider defaultColumnOrder={tracePickerColumnOrder} pageSize={FETCH_SIZE}>
+    <DataTableStateProvider columnDefs={tracePickerColumns} pageSize={FETCH_SIZE} lockedColumns={["status"]}>
       <RunsContent />
     </DataTableStateProvider>
   );

frontend/components/debugger-sessions/debugger-sessions.tsx

@@ -7,11 +7,11 @@ import { type ReactNode, useCallback, useState } from "react";
 
 import ClientTimestampFormatter from "@/components/client-timestamp-formatter";
 import { Badge } from "@/components/ui/badge.tsx";
+import { ColumnsMenu } from "@/components/ui/columns-menu";
 import Header from "@/components/ui/header";
 import { InfiniteDataTable } from "@/components/ui/infinite-datatable";
 import { useInfiniteScroll } from "@/components/ui/infinite-datatable/hooks/use-infinite-scroll";
 import { DataTableStateProvider } from "@/components/ui/infinite-datatable/model/datatable-store";
-import ColumnsMenu from "@/components/ui/infinite-datatable/ui/columns-menu.tsx";
 import Mono from "@/components/ui/mono";
 import { TableCell, TableRow } from "@/components/ui/table.tsx";
 import { type DebuggerSession, type DebuggerSessionStatus } from "@/lib/actions/debugger-sessions";
@@ -82,8 +82,6 @@ const columns: ColumnDef<DebuggerSession>[] = [
   },
 ];
 
-const defaultDebuggerSessionsColumnOrder = ["id", "name", "status", "createdAt"];
-
 const EmptyRow = (
   <TableRow className="flex">
     <TableCell className="text-center p-4 rounded-b w-full h-auto">
@@ -167,7 +165,6 @@ function DebuggerSessionsContent() {
         <div className="flex overflow-hidden flex-1">
           <InfiniteDataTable
             getRowId={(row: DebuggerSession) => row.id}
-            columns={columns}
             data={debuggerSessions ?? []}
             hasMore={hasMore}
             getRowHref={(row) => `debugger-sessions/${row.id}`}
@@ -181,12 +178,7 @@ function DebuggerSessionsContent() {
             emptyRow={EmptyRow}
           >
             <div className="flex flex-1 w-full space-x-2 pt-1">
-              <ColumnsMenu
-                columnLabels={columns.map((column) => ({
-                  id: column.id!,
-                  label: typeof column.header === "string" ? column.header : column.id!,
-                }))}
-              />
+              <ColumnsMenu />
             </div>
           </InfiniteDataTable>
         </div>
@@ -197,10 +189,7 @@ function DebuggerSessionsContent() {
 
 export default function DebuggerSessions() {
   return (
-    <DataTableStateProvider
-      storageKey="debugger-sessions-table"
-      defaultColumnOrder={defaultDebuggerSessionsColumnOrder}
-    >
+    <DataTableStateProvider storageKey="debugger-sessions-table" columnDefs={columns}>
       <DebuggerSessionsContent />
     </DataTableStateProvider>
   );

frontend/components/evaluation/eval-columns-menu/eval-columns-menu.tsx

@@ -1,21 +1,16 @@
 "use client";
 
 import { useParams } from "next/navigation";
-import { useMemo } from "react";
+import { useCallback, useMemo } from "react";
 
 import { useEvalStore } from "@/components/evaluation/store";
-import { type ColumnActions, ColumnsMenu, type CustomColumnPanelConfig } from "@/components/ui/columns-menu";
+import { ColumnsMenu, type CustomColumnPanelConfig } from "@/components/ui/columns-menu";
+import { selectAllColumnDefs, useDataTableStore } from "@/components/ui/infinite-datatable/model/datatable-store";
 
-interface EvalColumnsMenuProps {
-  lockedColumns?: string[];
-  columnLabels?: { id: string; label: string; onDelete?: () => void }[];
-}
-
-export default function EvalColumnsMenu({ lockedColumns = [], columnLabels = [] }: EvalColumnsMenuProps) {
+export default function EvalColumnsMenu() {
   const { evaluationId } = useParams();
   const isShared = useEvalStore((s) => s.isShared);
-  const addCustomColumn = useEvalStore((s) => s.addCustomColumn);
-  const updateCustomColumn = useEvalStore((s) => s.updateCustomColumn);
+  const store = useDataTableStore();
 
   const panelConfig = useMemo<CustomColumnPanelConfig>(
     () => ({
@@ -24,31 +19,16 @@ export default function EvalColumnsMenu({ lockedColumns = [], columnLabels = []
       buildTestQuery: (sql) =>
         `SELECT ${sql} as \`test\` FROM evaluation_datapoints WHERE evaluation_id = {evaluationId:UUID} LIMIT 1`,
       testQueryParameters: { evaluationId: evaluationId as string },
-      getColumnDefs: () => useEvalStore.getState().columnDefs,
+      getColumnDefs: () => selectAllColumnDefs(store.getState()),
       namePlaceholder: "e.g. Span Count",
       sqlPlaceholder: "e.g. arrayCount(x -> 1, trace_spans)",
       aiInputPlaceholder: "e.g. Count the number of spans in trace_spans",
       sqlHint: "Expression is added as a column: SELECT <expr> FROM evaluation_datapoints",
     }),
-    [evaluationId]
+    [evaluationId, store]
   );
 
-  const columnActions = useMemo<ColumnActions>(
-    () => ({
-      addCustomColumn,
-      updateCustomColumn,
-      getColumnDef: (columnId) => useEvalStore.getState().columnDefs.find((c) => c.id === columnId),
-    }),
-    [addCustomColumn, updateCustomColumn]
-  );
+  const getColumnDefs = useCallback(() => selectAllColumnDefs(store.getState()), [store]);
 
-  return (
-    <ColumnsMenu
-      lockedColumns={lockedColumns}
-      columnLabels={columnLabels}
-      panelConfig={panelConfig}
-      columnActions={columnActions}
-      showCreateButton={!isShared}
-    />
-  );
+  return <ColumnsMenu panelConfig={panelConfig} getColumnDefs={getColumnDefs} showCreateButton={!isShared} />;
 }