From df8c0f136eb932bef569ccd6672bbc50a65f62eb Mon Sep 17 00:00:00 2001
From: Olzhas Nurpeisov <olzhik1123@gmail.com>
Date: Mon, 23 Mar 2026 18:13:17 +0000
Subject: [PATCH 1/2] feat: refactor datatable

---
 frontend/components/dataset/dataset.tsx       |  20 ++--
 frontend/components/datasets/datasets.tsx     |  20 ++--
 .../sidebar/runs-tab.tsx                      |   5 +-
 .../debugger-sessions/debugger-sessions.tsx   |  16 +--
 .../eval-columns-menu/eval-columns-menu.tsx   |  21 +---
 .../evaluation-datapoints-table-content.tsx   |  36 +-----
 .../evaluation-datapoints-table/index.tsx     |  19 +---
 .../evaluations/evaluations-groups-bar.tsx    |  37 +++---
 .../components/evaluations/evaluations.tsx    |  27 ++---
 .../evaluators/evaluators-table.tsx           |  11 +-
 frontend/components/evaluators/evaluators.tsx |   9 +-
 frontend/components/evaluators/lib/consts.tsx |   2 -
 .../playground/playground-history-table.tsx   |  27 +----
 .../components/playgrounds/playgrounds.tsx    |  20 ++--
 frontend/components/queues/queues.tsx         |  20 ++--
 .../signal/create-signal-job/index.tsx        |  19 +---
 .../components/signal/events-table/index.tsx  |  13 +--
 .../components/signal/jobs-table/index.tsx    |  13 +--
 .../components/signal/runs-table/columns.tsx  |   2 -
 .../components/signal/runs-table/index.tsx    |  13 +--
 .../signal/triggers-table/columns.tsx         |   2 -
 .../signal/triggers-table/index.tsx           |  14 +--
 frontend/components/signals/index.tsx         |   2 +-
 frontend/components/sql/editor-panel.tsx      |   2 +-
 .../traces/sessions-table/columns.tsx         |  16 ---
 .../traces/sessions-table/index.tsx           |  17 +--
 .../components/traces/spans-table/columns.tsx |  16 ---
 .../components/traces/spans-table/index.tsx   |  15 +--
 .../traces/trace-picker/columns.tsx           |   2 -
 .../components/traces/trace-picker/index.tsx  |   4 +-
 .../trace-picker/trace-picker-content.tsx     |   1 -
 .../traces/traces-table/columns.tsx           |  16 ---
 .../components/traces/traces-table/index.tsx  |  56 +--------
 .../traces-columns-menu/index.tsx             |  20 +---
 .../ui/columns-menu/columns-list-panel.tsx    |  21 ++--
 .../ui/columns-menu/columns-menu.tsx          |  78 +++++++------
 frontend/components/ui/columns-menu/types.ts  |   1 +
 .../ui/infinite-datatable/index.tsx           |   2 -
 .../model/datatable-store.tsx                 |  87 ++++++++++----
 .../ui/infinite-datatable/model/types.ts      |   9 +-
 .../ui/columns-menu-item.tsx                  |  49 --------
 .../ui/infinite-datatable/ui/columns-menu.tsx | 106 ------------------
 .../ui/infinite-datatable/ui/header.tsx       |   7 +-
 43 files changed, 258 insertions(+), 635 deletions(-)
 delete mode 100644 frontend/components/ui/infinite-datatable/ui/columns-menu-item.tsx
 delete mode 100644 frontend/components/ui/infinite-datatable/ui/columns-menu.tsx

diff --git a/frontend/components/dataset/dataset.tsx b/frontend/components/dataset/dataset.tsx
index fad7fe87c..9e35579e1 100644
--- a/frontend/components/dataset/dataset.tsx
+++ b/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();
@@ -305,7 +303,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 +313,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 +347,12 @@ const DatasetContent = ({ dataset, enableDownloadParquet, publicApiBaseUrl }: Da
 
 export default function Dataset(props: DatasetProps) {
   return (
-    <DataTableStateProvider storageKey="dataset-table" defaultColumnOrder={defaultDatasetColumnOrder}>
+    <DataTableStateProvider
+      storageKey="dataset-table"
+      columns={columns}
+      enableRowSelection
+      lockedColumns={["__row_selection"]}
+    >
       <DatasetContent {...props} />
     </DataTableStateProvider>
   );
diff --git a/frontend/components/datasets/datasets.tsx b/frontend/components/datasets/datasets.tsx
index f9470ca33..fefb43fd6 100644
--- a/frontend/components/datasets/datasets.tsx
+++ b/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",
@@ -219,7 +217,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 +230,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 +243,12 @@ function DatasetsContent() {
 
 export default function Datasets() {
   return (
-    <DataTableStateProvider storageKey="datasets-table" defaultColumnOrder={defaultDatasetsColumnOrder}>
+    <DataTableStateProvider
+      storageKey="datasets-table"
+      columns={columns}
+      enableRowSelection
+      lockedColumns={["__row_selection"]}
+    >
       <DatasetsContent />
     </DataTableStateProvider>
   );
diff --git a/frontend/components/debugger-sessions/debugger-session-view/sidebar/runs-tab.tsx b/frontend/components/debugger-sessions/debugger-session-view/sidebar/runs-tab.tsx
index 7b77e132b..4439ab42d 100644
--- a/frontend/components/debugger-sessions/debugger-session-view/sidebar/runs-tab.tsx
+++ b/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";
@@ -83,7 +83,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 +93,7 @@ const RunsContent = () => {
 
 export default function RunsTab() {
   return (
-    <DataTableStateProvider defaultColumnOrder={tracePickerColumnOrder} pageSize={FETCH_SIZE}>
+    <DataTableStateProvider columns={tracePickerColumns} pageSize={FETCH_SIZE} lockedColumns={["status"]}>
       <RunsContent />
     </DataTableStateProvider>
   );
diff --git a/frontend/components/debugger-sessions/debugger-sessions.tsx b/frontend/components/debugger-sessions/debugger-sessions.tsx
index 4deb510b8..2854a9363 100644
--- a/frontend/components/debugger-sessions/debugger-sessions.tsx
+++ b/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">
@@ -181,12 +179,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 +190,7 @@ function DebuggerSessionsContent() {
 
 export default function DebuggerSessions() {
   return (
-    <DataTableStateProvider
-      storageKey="debugger-sessions-table"
-      defaultColumnOrder={defaultDebuggerSessionsColumnOrder}
-    >
+    <DataTableStateProvider storageKey="debugger-sessions-table" columns={columns}>
       <DebuggerSessionsContent />
     </DataTableStateProvider>
   );
diff --git a/frontend/components/evaluation/eval-columns-menu/eval-columns-menu.tsx b/frontend/components/evaluation/eval-columns-menu/eval-columns-menu.tsx
index f38f66de0..d7114b76f 100644
--- a/frontend/components/evaluation/eval-columns-menu/eval-columns-menu.tsx
+++ b/frontend/components/evaluation/eval-columns-menu/eval-columns-menu.tsx
@@ -6,16 +6,12 @@ import { useMemo } from "react";
 import { useEvalStore } from "@/components/evaluation/store";
 import { type ColumnActions, ColumnsMenu, type CustomColumnPanelConfig } from "@/components/ui/columns-menu";
 
-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 removeCustomColumn = useEvalStore((s) => s.removeCustomColumn);
 
   const panelConfig = useMemo<CustomColumnPanelConfig>(
     () => ({
@@ -37,18 +33,11 @@ export default function EvalColumnsMenu({ lockedColumns = [], columnLabels = []
     () => ({
       addCustomColumn,
       updateCustomColumn,
+      removeCustomColumn,
       getColumnDef: (columnId) => useEvalStore.getState().columnDefs.find((c) => c.id === columnId),
     }),
-    [addCustomColumn, updateCustomColumn]
+    [addCustomColumn, updateCustomColumn, removeCustomColumn]
   );
 
-  return (
-    <ColumnsMenu
-      lockedColumns={lockedColumns}
-      columnLabels={columnLabels}
-      panelConfig={panelConfig}
-      columnActions={columnActions}
-      showCreateButton={!isShared}
-    />
-  );
+  return <ColumnsMenu panelConfig={panelConfig} columnActions={columnActions} showCreateButton={!isShared} />;
 }
diff --git a/frontend/components/evaluation/evaluation-datapoints-table/evaluation-datapoints-table-content.tsx b/frontend/components/evaluation/evaluation-datapoints-table/evaluation-datapoints-table-content.tsx
index 1fd988454..070593ef3 100644
--- a/frontend/components/evaluation/evaluation-datapoints-table/evaluation-datapoints-table-content.tsx
+++ b/frontend/components/evaluation/evaluation-datapoints-table/evaluation-datapoints-table-content.tsx
@@ -1,7 +1,6 @@
 import { Settings as SettingsIcon } from "lucide-react";
 import { usePathname, useRouter, useSearchParams } from "next/navigation";
 import { useCallback, useEffect, useMemo } from "react";
-import { useStore } from "zustand";
 import { useShallow } from "zustand/react/shallow";
 
 import EvalColumnsMenu from "@/components/evaluation/eval-columns-menu";
@@ -16,7 +15,6 @@ import {
   DropdownMenuTrigger,
 } from "@/components/ui/dropdown-menu";
 import { InfiniteDataTable } from "@/components/ui/infinite-datatable";
-import { useDataTableStore } from "@/components/ui/infinite-datatable/model/datatable-store";
 import DataTableFilter, { DataTableFilterList } from "@/components/ui/infinite-datatable/ui/datatable-filter";
 import { Switch } from "@/components/ui/switch";
 import { type EvalRow } from "@/lib/evaluation/types";
@@ -47,30 +45,6 @@ const EvaluationDatapointsTableContent = ({
   const heatmapEnabled = useEvalStore((s) => s.heatmapEnabled);
   const setHeatmapEnabled = useEvalStore((s) => s.setHeatmapEnabled);
   const setScoreRanges = useEvalStore((s) => s.setScoreRanges);
-  const removeCustomColumn = useEvalStore((s) => s.removeCustomColumn);
-
-  // Datatable store for column sync
-  const datatableStore = useDataTableStore();
-  const { columnOrder, setColumnOrder } = useStore(datatableStore, (s) => ({
-    columnOrder: s.columnOrder,
-    setColumnOrder: s.setColumnOrder,
-  }));
-
-  // Sync datatable columnOrder with eval store columnDefs
-  useEffect(() => {
-    const visibleIds = columns.filter((c) => !c.meta?.hidden).map((c) => c.id!);
-    const currentSet = new Set(columnOrder);
-    const defSet = new Set(visibleIds);
-
-    const toAdd = visibleIds.filter((id) => !currentSet.has(id));
-    const toRemove = columnOrder.filter((id) => !defSet.has(id));
-
-    if (toAdd.length > 0 || toRemove.length > 0) {
-      const filtered = columnOrder.filter((id) => defSet.has(id));
-      setColumnOrder([...filtered, ...toAdd]);
-    }
-  }, [columns, columnOrder, setColumnOrder]);
-
   // Compute and set score ranges from data
   useEffect(() => {
     if (!data) return;
@@ -161,15 +135,7 @@ const EvaluationDatapointsTableContent = ({
       >
         <div className="flex flex-1 w-full space-x-2">
           <DataTableFilter columns={columnFilters} />
-          <EvalColumnsMenu
-            columnLabels={visibleColumns.map((column) => ({
-              id: column.id!,
-              label: typeof column.header === "string" ? column.header : column.id!,
-              ...(column.id!.startsWith("custom:") && {
-                onDelete: () => removeCustomColumn(column.id!.replace("custom:", "")),
-              }),
-            }))}
-          />
+          <EvalColumnsMenu />
           <DropdownMenu>
             <DropdownMenuTrigger asChild>
               <Button className="h-7 w-7" variant="outline" size="icon">
diff --git a/frontend/components/evaluation/evaluation-datapoints-table/index.tsx b/frontend/components/evaluation/evaluation-datapoints-table/index.tsx
index 0d3176b4b..28ed5951a 100644
--- a/frontend/components/evaluation/evaluation-datapoints-table/index.tsx
+++ b/frontend/components/evaluation/evaluation-datapoints-table/index.tsx
@@ -1,7 +1,7 @@
 import { type Row } from "@tanstack/react-table";
-import { useMemo } from "react";
+import { useShallow } from "zustand/react/shallow";
 
-import { useEvalStore } from "@/components/evaluation/store";
+import { selectVisibleColumns, useEvalStore } from "@/components/evaluation/store";
 import { DataTableStateProvider } from "@/components/ui/infinite-datatable/model/datatable-store";
 import { type EvalRow } from "@/lib/evaluation/types";
 
@@ -20,25 +20,16 @@ export interface EvaluationDatapointsTableProps {
   fetchNextPage: () => void;
 }
 
-const baseColumnOrder = ["status", "index", "data", "target", "metadata", "output", "duration", "cost"];
-
 const EvaluationDatapointsTable = (props: EvaluationDatapointsTableProps) => {
-  const { scores, isLoading } = props;
-  const customColumns = useEvalStore((s) => s.customColumns);
-  const defaultColumnOrder = useMemo(
-    () => [...baseColumnOrder, ...scores.map((s) => `score:${s}`), ...customColumns.map((cc) => `custom:${cc.name}`)],
-    [scores, customColumns]
-  );
+  const { isLoading } = props;
+  const visibleColumns = useEvalStore(useShallow(selectVisibleColumns));
 
-  // Delay mounting the store until scores are known, otherwise the store
-  // is created with an incomplete defaultColumnOrder and score columns
-  // won't be reorderable.
   if (isLoading) {
     return <EvalTableSkeleton />;
   }
 
   return (
-    <DataTableStateProvider storageKey="evaluation-datapoints-table" defaultColumnOrder={defaultColumnOrder}>
+    <DataTableStateProvider storageKey="evaluation-datapoints-table" columns={visibleColumns}>
       <EvaluationDatapointsTableContent {...props} />
     </DataTableStateProvider>
   );
diff --git a/frontend/components/evaluations/evaluations-groups-bar.tsx b/frontend/components/evaluations/evaluations-groups-bar.tsx
index 10355507e..5015051ba 100644
--- a/frontend/components/evaluations/evaluations-groups-bar.tsx
+++ b/frontend/components/evaluations/evaluations-groups-bar.tsx
@@ -11,14 +11,25 @@ import { InfiniteDataTable } from "../ui/infinite-datatable";
 
 type EvaluationGroup = { groupId: string; lastEvaluationCreatedAt: string };
 
-export const defaultEvaluationsGroupsBarColumnOrder = ["groupId", "lastEvaluationCreatedAt"];
+const columns: ColumnDef<EvaluationGroup>[] = [
+  {
+    id: "groupId",
+    header: "Group",
+    accessorFn: (row) => row.groupId,
+    size: 160,
+  },
+  {
+    id: "lastEvaluationCreatedAt",
+    header: "Last Evaluation",
+    accessorFn: (row) => row.lastEvaluationCreatedAt,
+    cell: ({ row }) => <ClientTimestampFormatter timestamp={row.original.lastEvaluationCreatedAt} />,
+    size: 160,
+  },
+];
 
 export default function EvaluationsGroupsBar() {
   return (
-    <DataTableStateProvider
-      storageKey="evaluations-groups-bar"
-      defaultColumnOrder={defaultEvaluationsGroupsBarColumnOrder}
-    >
+    <DataTableStateProvider storageKey="evaluations-groups-bar" columns={columns}>
       <EvaluationsGroupsBarContent />
     </DataTableStateProvider>
   );
@@ -43,22 +54,6 @@ function EvaluationsGroupsBarContent() {
     }
   }, [groups, groupId, router, projectId]);
 
-  const columns: ColumnDef<EvaluationGroup>[] = [
-    {
-      id: "groupId",
-      header: "Group",
-      accessorFn: (row) => row.groupId,
-      size: 160,
-    },
-    {
-      id: "lastEvaluationCreatedAt",
-      header: "Last Evaluation",
-      accessorFn: (row) => row.lastEvaluationCreatedAt,
-      cell: ({ row }) => <ClientTimestampFormatter timestamp={row.original.lastEvaluationCreatedAt} />,
-      size: 160,
-    },
-  ];
-
   return (
     <div className="max-w-80 flex flex-1 flex-col gap-2">
       <div className="flex overflow-hidden">
diff --git a/frontend/components/evaluations/evaluations.tsx b/frontend/components/evaluations/evaluations.tsx
index b005ac645..26c196efe 100644
--- a/frontend/components/evaluations/evaluations.tsx
+++ b/frontend/components/evaluations/evaluations.tsx
@@ -6,11 +6,11 @@ import React, { useCallback, useState } from "react";
 
 import SearchInput from "@/components/common/search-input";
 import ProgressionChart from "@/components/evaluations/progression-chart";
+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, useSelection } 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 JsonTooltip from "@/components/ui/json-tooltip.tsx";
@@ -59,15 +59,6 @@ const columns: ColumnDef<Evaluation>[] = [
   },
 ];
 
-export const defaultEvaluationsColumnOrder = [
-  "__row_selection",
-  "id",
-  "name",
-  "dataPointsCount",
-  "metadata",
-  "createdAt",
-];
-
 const filters: ColumnFilter[] = [
   {
     name: "ID",
@@ -95,7 +86,12 @@ const FETCH_SIZE = 50;
 
 export default function Evaluations() {
   return (
-    <DataTableStateProvider storageKey="evaluations-table" defaultColumnOrder={defaultEvaluationsColumnOrder}>
+    <DataTableStateProvider
+      storageKey="evaluations-table"
+      columns={columns}
+      enableRowSelection
+      lockedColumns={["__row_selection"]}
+    >
       <EvaluationsContent />
     </DataTableStateProvider>
   );
@@ -248,7 +244,6 @@ function EvaluationsContent() {
                 fetchNextPage={fetchNextPage}
                 state={{ rowSelection }}
                 onRowSelectionChange={onRowSelectionChange}
-                lockedColumns={["__row_selection"]}
                 selectionPanel={(selectedRowIds) => (
                   <div className="flex flex-col space-y-2">
                     <DeleteSelectedRows
@@ -261,13 +256,7 @@ function EvaluationsContent() {
               >
                 <div className="flex flex-1 w-full space-x-2">
                   <DataTableFilter columns={filters} />
-                  <ColumnsMenu
-                    lockedColumns={["__row_selection"]}
-                    columnLabels={columns.map((column) => ({
-                      id: column.id!,
-                      label: typeof column.header === "string" ? column.header : column.id!,
-                    }))}
-                  />
+                  <ColumnsMenu />
                   <SearchInput placeholder="Search evaluations by name..." />
                 </div>
                 <DataTableFilterList />
diff --git a/frontend/components/evaluators/evaluators-table.tsx b/frontend/components/evaluators/evaluators-table.tsx
index 4f8337245..dfa10c40d 100644
--- a/frontend/components/evaluators/evaluators-table.tsx
+++ b/frontend/components/evaluators/evaluators-table.tsx
@@ -4,9 +4,9 @@ import { type Row, type RowSelectionState } from "@tanstack/react-table";
 import { useCallback, useState } from "react";
 import useSWR from "swr";
 
+import { ColumnsMenu } from "@/components/ui/columns-menu";
 import DeleteSelectedRows from "@/components/ui/delete-selected-rows.tsx";
 import { InfiniteDataTable } from "@/components/ui/infinite-datatable";
-import ColumnsMenu from "@/components/ui/infinite-datatable/ui/columns-menu.tsx";
 import { type Evaluator } from "@/lib/evaluators/types";
 import { useToast } from "@/lib/hooks/use-toast";
 import { type PaginatedResponse } from "@/lib/types";
@@ -72,7 +72,6 @@ export default function EvaluatorsTable({ projectId, onRowClick }: EvaluatorsTab
       getRowId={(row) => row.id}
       onRowClick={onRowClick}
       enableRowSelection
-      lockedColumns={["__row_selection"]}
       state={{
         rowSelection,
       }}
@@ -87,13 +86,7 @@ export default function EvaluatorsTable({ projectId, onRowClick }: EvaluatorsTab
         </div>
       )}
     >
-      <ColumnsMenu
-        lockedColumns={["__row_selection"]}
-        columnLabels={columns.map((column) => ({
-          id: column.id!,
-          label: typeof column.header === "string" ? column.header : column.id!,
-        }))}
-      />
+      <ColumnsMenu />
     </InfiniteDataTable>
   );
 }
diff --git a/frontend/components/evaluators/evaluators.tsx b/frontend/components/evaluators/evaluators.tsx
index 78645efb1..76451e7e9 100644
--- a/frontend/components/evaluators/evaluators.tsx
+++ b/frontend/components/evaluators/evaluators.tsx
@@ -13,7 +13,7 @@ import { Button } from "@/components/ui/button";
 import { DataTableStateProvider } from "@/components/ui/infinite-datatable/model/datatable-store";
 import { type Evaluator } from "@/lib/evaluators/types";
 
-import { defaultEvaluatorsColumnOrder } from "./lib/consts";
+import { columns } from "./lib/consts";
 
 const manageEvaluatorSchema = z.object({
   id: z.string().optional(),
@@ -70,7 +70,12 @@ const Evaluators = () => {
             Evaluator
           </Button>
         </ManageEvaluatorSheet>
-        <DataTableStateProvider storageKey="evaluators-table" defaultColumnOrder={defaultEvaluatorsColumnOrder}>
+        <DataTableStateProvider
+          storageKey="evaluators-table"
+          columns={columns}
+          enableRowSelection
+          lockedColumns={["__row_selection"]}
+        >
           <EvaluatorsTable onRowClick={handleRowClick} projectId={projectId as string} />
         </DataTableStateProvider>
       </div>
diff --git a/frontend/components/evaluators/lib/consts.tsx b/frontend/components/evaluators/lib/consts.tsx
index 85ec84fe2..1f24a3dee 100644
--- a/frontend/components/evaluators/lib/consts.tsx
+++ b/frontend/components/evaluators/lib/consts.tsx
@@ -23,5 +23,3 @@ export const columns: ColumnDef<Evaluator>[] = [
     cell: ({ row }) => <ClientTimestampFormatter absolute timestamp={row.getValue("createdAt")} />,
   },
 ];
-
-export const defaultEvaluatorsColumnOrder = ["__row_selection", "name", "evaluatorType", "createdAt"];
diff --git a/frontend/components/playground/playground-history-table.tsx b/frontend/components/playground/playground-history-table.tsx
index b8fce098a..538bf488e 100644
--- a/frontend/components/playground/playground-history-table.tsx
+++ b/frontend/components/playground/playground-history-table.tsx
@@ -6,10 +6,10 @@ import { useCallback } from "react";
 
 import ClientTimestampFormatter from "@/components/client-timestamp-formatter";
 import SpanTypeIcon from "@/components/traces/span-type-icon";
+import { ColumnsMenu } from "@/components/ui/columns-menu";
 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 Mono from "@/components/ui/mono";
 import { useToast } from "@/lib/hooks/use-toast";
 import { type Trace } from "@/lib/traces/types";
@@ -126,18 +126,6 @@ const columns: ColumnDef<Trace, any>[] = [
   },
 ];
 
-export const defaultPlaygroundHistoryColumnOrder = [
-  "status",
-  "id",
-  "top_span_type",
-  "input",
-  "output",
-  "start_time",
-  "latency",
-  "cost",
-  "total_token_count",
-];
-
 interface PlaygroundHistoryTableProps {
   playgroundId: string;
   onRowClick?: (trace: Trace) => void;
@@ -148,11 +136,7 @@ const FETCH_SIZE = 50;
 
 export default function PlaygroundHistoryTable(props: PlaygroundHistoryTableProps) {
   return (
-    <DataTableStateProvider
-      storageKey="playground-history-table"
-      uniqueKey="id"
-      defaultColumnOrder={defaultPlaygroundHistoryColumnOrder}
-    >
+    <DataTableStateProvider storageKey="playground-history-table" uniqueKey="id" columns={columns}>
       <PlaygroundHistoryTableContent {...props} />
     </DataTableStateProvider>
   );
@@ -240,12 +224,7 @@ function PlaygroundHistoryTableContent({ playgroundId, onRowClick, onTraceSelect
       isLoading={isLoading}
       fetchNextPage={fetchNextPage}
     >
-      <ColumnsMenu
-        columnLabels={columns.map((column) => ({
-          id: column.id!,
-          label: typeof column.header === "string" ? column.header : column.id!,
-        }))}
-      />
+      <ColumnsMenu />
     </InfiniteDataTable>
   );
 }
diff --git a/frontend/components/playgrounds/playgrounds.tsx b/frontend/components/playgrounds/playgrounds.tsx
index e73d43191..d9167dd21 100644
--- a/frontend/components/playgrounds/playgrounds.tsx
+++ b/frontend/components/playgrounds/playgrounds.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 { 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 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";
@@ -52,8 +52,6 @@ const columns: ColumnDef<PlaygroundInfo>[] = [
   },
 ];
 
-export const defaultPlaygroundsColumnOrder = ["__row_selection", "id", "name", "createdAt"];
-
 const playgroundsTableFilters: ColumnFilter[] = [
   {
     name: "ID",
@@ -206,7 +204,6 @@ const PlaygroundsContent = () => {
             rowSelection,
           }}
           onRowSelectionChange={setRowSelection}
-          lockedColumns={["__row_selection"]}
           emptyRow={filter.length === 0 && !search ? EmptyRow : undefined}
           selectionPanel={(selectedRowIds) => (
             <div className="flex flex-col space-y-2">
@@ -240,13 +237,7 @@ const PlaygroundsContent = () => {
         >
           <div className="flex flex-1 w-full space-x-2 pt-1">
             <DataTableFilter columns={playgroundsTableFilters} />
-            <ColumnsMenu
-              columnLabels={columns.map((column) => ({
-                id: column.id!,
-                label: typeof column.header === "string" ? column.header : column.id!,
-              }))}
-              lockedColumns={["__row_selection"]}
-            />
+            <ColumnsMenu />
             <DataTableSearch className="mr-0.5" placeholder="Search by playground name..." />
           </div>
           <DataTableFilterList />
@@ -258,7 +249,12 @@ const PlaygroundsContent = () => {
 
 export default function Playgrounds() {
   return (
-    <DataTableStateProvider storageKey="playgrounds-table" defaultColumnOrder={defaultPlaygroundsColumnOrder}>
+    <DataTableStateProvider
+      storageKey="playgrounds-table"
+      columns={columns}
+      enableRowSelection
+      lockedColumns={["__row_selection"]}
+    >
       <PlaygroundsContent />
     </DataTableStateProvider>
   );
diff --git a/frontend/components/queues/queues.tsx b/frontend/components/queues/queues.tsx
index ed94b00fb..a438f4c4b 100644
--- a/frontend/components/queues/queues.tsx
+++ b/frontend/components/queues/queues.tsx
@@ -7,10 +7,10 @@ import { useCallback, useState } from "react";
 
 import ClientTimestampFormatter from "@/components/client-timestamp-formatter";
 import { Button } from "@/components/ui/button";
+import { ColumnsMenu } from "@/components/ui/columns-menu";
 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 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";
@@ -58,8 +58,6 @@ const columns: ColumnDef<LabelingQueue>[] = [
   },
 ];
 
-export const defaultQueuesColumnOrder = ["__row_selection", "id", "name", "count", "createdAt"];
-
 const queuesTableFilters: ColumnFilter[] = [
   {
     name: "ID",
@@ -221,7 +219,6 @@ const QueuesContent = () => {
             rowSelection,
           }}
           onRowSelectionChange={setRowSelection}
-          lockedColumns={["__row_selection"]}
           emptyRow={filter.length === 0 && !search ? EmptyRow : undefined}
           selectionPanel={(selectedRowIds) => (
             <div className="flex flex-col space-y-2">
@@ -255,13 +252,7 @@ const QueuesContent = () => {
         >
           <div className="flex flex-1 w-full space-x-2 pt-1">
             <DataTableFilter columns={queuesTableFilters} />
-            <ColumnsMenu
-              columnLabels={columns.map((column) => ({
-                id: column.id!,
-                label: typeof column.header === "string" ? column.header : column.id!,
-              }))}
-              lockedColumns={["__row_selection"]}
-            />
+            <ColumnsMenu />
             <DataTableSearch className="mr-0.5" placeholder="Search by queue name..." />
           </div>
           <DataTableFilterList />
@@ -273,7 +264,12 @@ const QueuesContent = () => {
 
 export default function Queues() {
   return (
-    <DataTableStateProvider storageKey="queues-table" defaultColumnOrder={defaultQueuesColumnOrder}>
+    <DataTableStateProvider
+      storageKey="queues-table"
+      columns={columns}
+      enableRowSelection
+      lockedColumns={["__row_selection"]}
+    >
       <QueuesContent />
     </DataTableStateProvider>
   );
diff --git a/frontend/components/signal/create-signal-job/index.tsx b/frontend/components/signal/create-signal-job/index.tsx
index cfd0f9753..e4ecd0379 100644
--- a/frontend/components/signal/create-signal-job/index.tsx
+++ b/frontend/components/signal/create-signal-job/index.tsx
@@ -10,18 +10,14 @@ import ConfirmSignalJobDialog from "@/components/signal/create-signal-job/confir
 import SelectionBanner from "@/components/signal/create-signal-job/selection-banner.tsx";
 import { useSignalStoreContext } from "@/components/signal/store.tsx";
 import TraceView from "@/components/traces/trace-view";
-import {
-  columns,
-  defaultTracesColumnOrder,
-  filters as tableFilters,
-} from "@/components/traces/traces-table/columns.tsx";
+import { columns, filters as tableFilters } from "@/components/traces/traces-table/columns.tsx";
 import { Button } from "@/components/ui/button.tsx";
+import { ColumnsMenu } from "@/components/ui/columns-menu";
 import DateRangeFilter from "@/components/ui/date-range-filter";
 import Header from "@/components/ui/header.tsx";
 import { InfiniteDataTable } from "@/components/ui/infinite-datatable";
 import { useInfiniteScroll, useSelection } from "@/components/ui/infinite-datatable/hooks";
 import { DataTableStateProvider } from "@/components/ui/infinite-datatable/model/datatable-store.tsx";
-import ColumnsMenu from "@/components/ui/infinite-datatable/ui/columns-menu.tsx";
 import RefreshButton from "@/components/ui/infinite-datatable/ui/refresh-button.tsx";
 import { useFeatureFlags } from "@/contexts/feature-flags-context";
 import type { Filter } from "@/lib/actions/common/filters.ts";
@@ -315,16 +311,9 @@ const CreateSignalJobContent = () => {
           }}
           onRowSelectionChange={onRowSelectionChange}
           getRowHref={getRowHref}
-          lockedColumns={["__row_selection", "status"]}
         >
           <div className="flex flex-1 w-full h-full gap-2">
-            <ColumnsMenu
-              lockedColumns={["__row_selection", "status"]}
-              columnLabels={columns.map((column) => ({
-                id: column.id!,
-                label: typeof column.header === "string" ? column.header : column.id!,
-              }))}
-            />
+            <ColumnsMenu />
             <DateRangeFilter mode="state" value={dateRange} onChange={setDateRange} />
             <RefreshButton onClick={refetch} variant="outline" />
           </div>
@@ -394,7 +383,7 @@ export default function CreateSignalJob({ traceId }: { traceId?: string }) {
   }, []);
 
   return (
-    <DataTableStateProvider defaultColumnOrder={["__row_selection", ...defaultTracesColumnOrder]}>
+    <DataTableStateProvider columns={columns} enableRowSelection lockedColumns={["__row_selection", "status"]}>
       <CreateSignalJobContent />
     </DataTableStateProvider>
   );
diff --git a/frontend/components/signal/events-table/index.tsx b/frontend/components/signal/events-table/index.tsx
index 0b20a4c16..3b1ad5c4b 100644
--- a/frontend/components/signal/events-table/index.tsx
+++ b/frontend/components/signal/events-table/index.tsx
@@ -11,12 +11,12 @@ import { useClusterId } from "@/components/signal/hooks/use-cluster-id";
 import { getFilterClusterIds, useSignalStoreContext } from "@/components/signal/store.tsx";
 import { type EventNavigationItem } from "@/components/signal/utils.ts";
 import { useTraceViewNavigation } from "@/components/traces/trace-view/navigation-context.tsx";
+import { ColumnsMenu } from "@/components/ui/columns-menu";
 import DateRangeFilter from "@/components/ui/date-range-filter";
 import { getDisplayRange, getTimeDifference } from "@/components/ui/date-range-filter/utils.ts";
 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";
 import DataTableFilter, { DataTableFilterList } from "@/components/ui/infinite-datatable/ui/datatable-filter";
 import { TableCell, TableRow } from "@/components/ui/table.tsx";
 import { UNCLUSTERED_ID } from "@/lib/actions/clusters";
@@ -216,12 +216,7 @@ function PureEventsTable() {
       >
         <div className="flex flex-1 w-full h-full gap-2">
           <DataTableFilter columns={filters} />
-          <ColumnsMenu
-            columnLabels={columns.map((column) => ({
-              id: column.id!,
-              label: typeof column.header === "string" ? column.header : column.id!,
-            }))}
-          />
+          <ColumnsMenu />
           <DateRangeFilter />
         </div>
         <ClusterBreadcrumbs />
@@ -235,10 +230,10 @@ function PureEventsTable() {
 export default function EventsTable() {
   const signal = useSignalStoreContext((state) => state.signal);
 
-  const { columnOrder } = useMemo(() => buildEventsColumns(signal.schemaFields), [signal.schemaFields]);
+  const { columns } = useMemo(() => buildEventsColumns(signal.schemaFields), [signal.schemaFields]);
 
   return (
-    <DataTableStateProvider storageKey={`events-table-${signal.id}`} uniqueKey="id" defaultColumnOrder={columnOrder}>
+    <DataTableStateProvider storageKey={`events-table-${signal.id}`} uniqueKey="id" columns={columns}>
       <PureEventsTable />
     </DataTableStateProvider>
   );
diff --git a/frontend/components/signal/jobs-table/index.tsx b/frontend/components/signal/jobs-table/index.tsx
index 5ed4bde1d..8734a3d58 100644
--- a/frontend/components/signal/jobs-table/index.tsx
+++ b/frontend/components/signal/jobs-table/index.tsx
@@ -10,9 +10,9 @@ import useSWR from "swr";
 import { type SignalJobRow, signalJobsColumns, signalJobsFilters } from "@/components/signal/jobs-table/columns.tsx";
 import { useSignalStoreContext } from "@/components/signal/store.tsx";
 import { Button } from "@/components/ui/button.tsx";
+import { ColumnsMenu } from "@/components/ui/columns-menu";
 import { InfiniteDataTable } from "@/components/ui/infinite-datatable";
 import { DataTableStateProvider } from "@/components/ui/infinite-datatable/model/datatable-store.tsx";
-import ColumnsMenu from "@/components/ui/infinite-datatable/ui/columns-menu.tsx";
 import FilterPopover, { FilterList } from "@/components/ui/infinite-datatable/ui/datatable-filter/ui";
 import { TableCell, TableRow } from "@/components/ui/table";
 import { type Filter } from "@/lib/actions/common/filters.ts";
@@ -100,7 +100,6 @@ const JobsTableContent = () => {
         columns={signalJobsColumns}
         data={jobs}
         getRowId={(job) => job.id}
-        lockedColumns={["id"]}
         hasMore={false}
         isFetching={isLoading}
         isLoading={isLoading}
@@ -111,13 +110,7 @@ const JobsTableContent = () => {
       >
         <div className="flex flex-1 w-full space-x-2">
           <FilterPopover columns={signalJobsFilters} filters={jobsFilters} onAddFilter={handleAddFilter} />
-          <ColumnsMenu
-            columnLabels={signalJobsColumns.map((column) => ({
-              id: column.id!,
-              label: typeof column.header === "string" ? column.header : column.id!,
-            }))}
-            lockedColumns={["id"]}
-          />
+          <ColumnsMenu />
         </div>
         <FilterList className="py-[3px] text-xs px-1" filters={jobsFilters} onRemoveFilter={handleRemoveFilter} />
       </InfiniteDataTable>
@@ -127,7 +120,7 @@ const JobsTableContent = () => {
 
 export default function SignalJobsTable() {
   return (
-    <DataTableStateProvider defaultColumnOrder={signalJobsColumns.map((c) => String(c.id))}>
+    <DataTableStateProvider columns={signalJobsColumns} lockedColumns={["id"]}>
       <JobsTableContent />
     </DataTableStateProvider>
   );
diff --git a/frontend/components/signal/runs-table/columns.tsx b/frontend/components/signal/runs-table/columns.tsx
index 41b1be4a2..a3ee4e5fe 100644
--- a/frontend/components/signal/runs-table/columns.tsx
+++ b/frontend/components/signal/runs-table/columns.tsx
@@ -94,8 +94,6 @@ export const getSignalRunsColumns = ({
   },
 ];
 
-export const defaultRunsColumnOrder = ["runId", "traceId", "eventId", "source", "mode", "status", "updatedAt"];
-
 export const signalRunsFilters: ColumnFilter[] = [
   {
     name: "Job ID",
diff --git a/frontend/components/signal/runs-table/index.tsx b/frontend/components/signal/runs-table/index.tsx
index 86bac7198..145ed9a1e 100644
--- a/frontend/components/signal/runs-table/index.tsx
+++ b/frontend/components/signal/runs-table/index.tsx
@@ -6,12 +6,12 @@ import { useParams, useRouter } from "next/navigation";
 import React, { useCallback, useState } from "react";
 
 import { useSignalStoreContext } from "@/components/signal/store";
+import { ColumnsMenu } from "@/components/ui/columns-menu";
 import DateRangeFilter from "@/components/ui/date-range-filter";
 import { getDisplayRange, getTimeDifference } from "@/components/ui/date-range-filter/utils.ts";
 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";
 import FilterPopover, { FilterList } from "@/components/ui/infinite-datatable/ui/datatable-filter/ui";
 import { TableCell, TableRow } from "@/components/ui/table.tsx";
 import { type Filter } from "@/lib/actions/common/filters";
@@ -19,7 +19,7 @@ import { Operator } from "@/lib/actions/common/operators";
 import { type SignalRunRow } from "@/lib/actions/signal-runs";
 import { useToast } from "@/lib/hooks/use-toast";
 
-import { defaultRunsColumnOrder, getSignalRunsColumns, signalRunsFilters } from "./columns";
+import { getSignalRunsColumns, signalRunsFilters } from "./columns";
 
 const FETCH_SIZE = 50;
 
@@ -164,12 +164,7 @@ function RunsTableContent() {
       >
         <div className="flex flex-1 w-full space-x-2">
           <FilterPopover columns={signalRunsFilters} filters={filter} onAddFilter={handleAddFilter} />
-          <ColumnsMenu
-            columnLabels={columns.map((column) => ({
-              id: column.id!,
-              label: typeof column.header === "string" ? column.header : column.id!,
-            }))}
-          />
+          <ColumnsMenu />
           <DateRangeFilter mode="state" value={dateRange} onChange={setDateRange} />
         </div>
         <FilterList className="py-[3px] text-xs px-1" filters={filter} onRemoveFilter={handleRemoveFilter} />
@@ -183,7 +178,7 @@ export default function SignalRunsTable() {
     <DataTableStateProvider
       storageKey="signal-runs-table"
       uniqueKey="runId"
-      defaultColumnOrder={defaultRunsColumnOrder}
+      columns={getSignalRunsColumns({ onJobNav: () => {} })}
     >
       <RunsTableContent />
     </DataTableStateProvider>
diff --git a/frontend/components/signal/triggers-table/columns.tsx b/frontend/components/signal/triggers-table/columns.tsx
index dd826f715..0e21e9971 100644
--- a/frontend/components/signal/triggers-table/columns.tsx
+++ b/frontend/components/signal/triggers-table/columns.tsx
@@ -61,8 +61,6 @@ export const getTriggersTableColumns = (): ColumnDef<TriggerRow>[] => [
   },
 ];
 
-export const defaultTriggersColumnOrder = ["filters", "mode", "createdAt"];
-
 export const triggersFilters: ColumnFilter[] = [
   {
     name: "Trigger ID",
diff --git a/frontend/components/signal/triggers-table/index.tsx b/frontend/components/signal/triggers-table/index.tsx
index 06ff9fcd0..c621284d7 100644
--- a/frontend/components/signal/triggers-table/index.tsx
+++ b/frontend/components/signal/triggers-table/index.tsx
@@ -8,17 +8,16 @@ import useSWR, { mutate } from "swr";
 
 import { useSignalStoreContext } from "@/components/signal/store.tsx";
 import {
-  defaultTriggersColumnOrder,
   getTriggersTableColumns,
   type TriggerRow,
   triggersFilters,
 } from "@/components/signal/triggers-table/columns.tsx";
 import ManageTriggerDialog from "@/components/signals/manage-trigger-dialog";
 import { Button } from "@/components/ui/button.tsx";
+import { ColumnsMenu } from "@/components/ui/columns-menu";
 import DeleteSelectedRows from "@/components/ui/delete-selected-rows";
 import { InfiniteDataTable } from "@/components/ui/infinite-datatable";
 import { DataTableStateProvider } from "@/components/ui/infinite-datatable/model/datatable-store.tsx";
-import ColumnsMenu from "@/components/ui/infinite-datatable/ui/columns-menu.tsx";
 import FilterPopover, { FilterList } from "@/components/ui/infinite-datatable/ui/datatable-filter/ui";
 import { TableCell, TableRow } from "@/components/ui/table";
 import { type Filter } from "@/lib/actions/common/filters.ts";
@@ -151,7 +150,6 @@ function TriggersTableContent() {
         columns={columns}
         data={triggers}
         getRowId={(trigger) => trigger.id}
-        lockedColumns={["__row_selection"]}
         hasMore={false}
         isFetching={isLoading}
         isLoading={isLoading}
@@ -167,13 +165,7 @@ function TriggersTableContent() {
       >
         <div className="flex flex-1 w-full space-x-2">
           <FilterPopover columns={triggersFilters} filters={storeTriggersFilters} onAddFilter={handleAddFilter} />
-          <ColumnsMenu
-            lockedColumns={["__row_selection"]}
-            columnLabels={columns.map((column) => ({
-              id: column.id!,
-              label: typeof column.header === "string" ? column.header : column.id!,
-            }))}
-          />
+          <ColumnsMenu />
         </div>
         <FilterList
           className="py-[3px] text-xs px-1"
@@ -187,7 +179,7 @@ function TriggersTableContent() {
 
 export default function TriggersTable() {
   return (
-    <DataTableStateProvider defaultColumnOrder={["__row_selection", ...defaultTriggersColumnOrder]}>
+    <DataTableStateProvider columns={getTriggersTableColumns()} enableRowSelection lockedColumns={["__row_selection"]}>
       <TriggersTableContent />
     </DataTableStateProvider>
   );
diff --git a/frontend/components/signals/index.tsx b/frontend/components/signals/index.tsx
index a8de7b515..0362526e9 100644
--- a/frontend/components/signals/index.tsx
+++ b/frontend/components/signals/index.tsx
@@ -30,7 +30,7 @@ const SIGNAL_QUICK_RANGES: DateRange[] = [
 
 export default function Signals() {
   return (
-    <DataTableStateProvider storageKey="signals-cards" uniqueKey="id" defaultColumnOrder={[]}>
+    <DataTableStateProvider storageKey="signals-cards" uniqueKey="id">
       <SignalsContent />
     </DataTableStateProvider>
   );
diff --git a/frontend/components/sql/editor-panel.tsx b/frontend/components/sql/editor-panel.tsx
index 746a91fe1..9b5a57375 100644
--- a/frontend/components/sql/editor-panel.tsx
+++ b/frontend/components/sql/editor-panel.tsx
@@ -265,7 +265,7 @@ export default function EditorPanel() {
             <div className="flex overflow-hidden h-full">
               {renderContent({
                 success: (
-                  <DataTableStateProvider>
+                  <DataTableStateProvider columns={columns}>
                     <InfiniteDataTable
                       className="w-full"
                       columns={columns}
diff --git a/frontend/components/traces/sessions-table/columns.tsx b/frontend/components/traces/sessions-table/columns.tsx
index d578361d9..fd6400f50 100644
--- a/frontend/components/traces/sessions-table/columns.tsx
+++ b/frontend/components/traces/sessions-table/columns.tsx
@@ -210,19 +210,3 @@ export const columns: ColumnDef<SessionRow, any>[] = [
     id: "tags",
   },
 ];
-
-export const defaultSessionsColumnOrder = [
-  "type",
-  "id",
-  "start_time",
-  "duration",
-  "input_cost",
-  "output_cost",
-  "total_cost",
-  "input_tokens",
-  "output_tokens",
-  "total_tokens",
-  "trace_count",
-  "user_id",
-  "tags",
-];
diff --git a/frontend/components/traces/sessions-table/index.tsx b/frontend/components/traces/sessions-table/index.tsx
index 93a5ddcfc..e5dafc156 100644
--- a/frontend/components/traces/sessions-table/index.tsx
+++ b/frontend/components/traces/sessions-table/index.tsx
@@ -6,14 +6,14 @@ import { useParams, usePathname, useRouter, useSearchParams } from "next/navigat
 import { useCallback, useEffect } from "react";
 
 import SearchInput from "@/components/common/search-input";
-import { columns, defaultSessionsColumnOrder, filters } from "@/components/traces/sessions-table/columns";
+import { columns, filters } from "@/components/traces/sessions-table/columns";
 import { useTraceViewNavigation } from "@/components/traces/trace-view/navigation-context";
 import { useTracesStoreContext } from "@/components/traces/traces-store";
+import { ColumnsMenu } from "@/components/ui/columns-menu";
 import DateRangeFilter from "@/components/ui/date-range-filter";
 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 DataTableFilter, { DataTableFilterList } from "@/components/ui/infinite-datatable/ui/datatable-filter";
 import RefreshButton from "@/components/ui/infinite-datatable/ui/refresh-button.tsx";
 import { useToast } from "@/lib/hooks/use-toast";
@@ -23,11 +23,7 @@ const FETCH_SIZE = 50;
 
 export default function SessionsTable() {
   return (
-    <DataTableStateProvider
-      storageKey="sessions-table"
-      uniqueKey="sessionId"
-      defaultColumnOrder={defaultSessionsColumnOrder}
-    >
+    <DataTableStateProvider storageKey="sessions-table" uniqueKey="sessionId" columns={columns}>
       <SessionsTableContent />
     </DataTableStateProvider>
   );
@@ -223,12 +219,7 @@ function SessionsTableContent() {
       >
         <div className="flex flex-1 w-full h-full gap-2">
           <DataTableFilter columns={filters} />
-          <ColumnsMenu
-            columnLabels={columns.map((column) => ({
-              id: column.id!,
-              label: typeof column.header === "string" ? column.header : column.id!,
-            }))}
-          />
+          <ColumnsMenu />
           <DateRangeFilter />
           <RefreshButton onClick={refetch} variant="outline" />
           <SearchInput placeholder="Search in sessions..." />
diff --git a/frontend/components/traces/spans-table/columns.tsx b/frontend/components/traces/spans-table/columns.tsx
index bffcb25d4..b6e2bcca3 100644
--- a/frontend/components/traces/spans-table/columns.tsx
+++ b/frontend/components/traces/spans-table/columns.tsx
@@ -279,19 +279,3 @@ export const columns: ColumnDef<SpanRow, any>[] = [
     id: "tags",
   },
 ];
-
-export const defaultSpansColumnOrder = [
-  "status",
-  "span_id",
-  "trace_id",
-  "span",
-  "path",
-  "input",
-  "output",
-  "start_time",
-  "duration",
-  "cost",
-  "tokens",
-  "model",
-  "tags",
-];
diff --git a/frontend/components/traces/spans-table/index.tsx b/frontend/components/traces/spans-table/index.tsx
index a0072ee64..6439dd1b0 100644
--- a/frontend/components/traces/spans-table/index.tsx
+++ b/frontend/components/traces/spans-table/index.tsx
@@ -5,14 +5,14 @@ import { useParams, usePathname, useRouter, useSearchParams } from "next/navigat
 import { useCallback, useEffect } from "react";
 
 import AdvancedSearch from "@/components/common/advanced-search";
-import { columns, defaultSpansColumnOrder, filters } from "@/components/traces/spans-table/columns";
+import { columns, filters } from "@/components/traces/spans-table/columns";
 import { useTraceViewNavigation } from "@/components/traces/trace-view/navigation-context";
 import { useTracesStoreContext } from "@/components/traces/traces-store";
+import { ColumnsMenu } from "@/components/ui/columns-menu";
 import DateRangeFilter from "@/components/ui/date-range-filter";
 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 DataTableFilter from "@/components/ui/infinite-datatable/ui/datatable-filter";
 import RefreshButton from "@/components/ui/infinite-datatable/ui/refresh-button.tsx";
 import { useToast } from "@/lib/hooks/use-toast";
@@ -22,7 +22,7 @@ const FETCH_SIZE = 50;
 
 export default function SpansTable() {
   return (
-    <DataTableStateProvider storageKey="spans-table" uniqueKey="spanId" defaultColumnOrder={defaultSpansColumnOrder}>
+    <DataTableStateProvider storageKey="spans-table" uniqueKey="spanId" columns={columns} lockedColumns={["status"]}>
       <SpansTableContent />
     </DataTableStateProvider>
   );
@@ -156,17 +156,10 @@ function SpansTableContent() {
         isFetching={isFetching}
         isLoading={isLoading}
         fetchNextPage={fetchNextPage}
-        lockedColumns={["status"]}
       >
         <div className="flex flex-1 w-full h-full gap-2">
           <DataTableFilter columns={filters} />
-          <ColumnsMenu
-            lockedColumns={["status"]}
-            columnLabels={columns.map((column) => ({
-              id: column.id!,
-              label: typeof column.header === "string" ? column.header : column.id!,
-            }))}
-          />
+          <ColumnsMenu />
           <DateRangeFilter />
           <RefreshButton onClick={refetch} variant="outline" />
         </div>
diff --git a/frontend/components/traces/trace-picker/columns.tsx b/frontend/components/traces/trace-picker/columns.tsx
index 5ee589fdd..df73c1273 100644
--- a/frontend/components/traces/trace-picker/columns.tsx
+++ b/frontend/components/traces/trace-picker/columns.tsx
@@ -101,6 +101,4 @@ export const tracePickerColumns: ColumnDef<TraceRow, any>[] = [
   },
 ];
 
-export const tracePickerColumnOrder = ["status", "top_span_type", "start_time", "duration", "total_tokens"];
-
 export const FETCH_SIZE = 30;
diff --git a/frontend/components/traces/trace-picker/index.tsx b/frontend/components/traces/trace-picker/index.tsx
index bb053e3a3..2da7cfec7 100644
--- a/frontend/components/traces/trace-picker/index.tsx
+++ b/frontend/components/traces/trace-picker/index.tsx
@@ -2,13 +2,13 @@
 
 import { DataTableStateProvider } from "@/components/ui/infinite-datatable/model/datatable-store";
 
-import { FETCH_SIZE, tracePickerColumnOrder } from "./columns";
+import { FETCH_SIZE, tracePickerColumns } from "./columns";
 import TracePickerContent, { type TracePickerProps } from "./trace-picker-content";
 
 export type { TracePickerProps };
 
 const TracePicker = (props: TracePickerProps) => (
-  <DataTableStateProvider defaultColumnOrder={tracePickerColumnOrder} pageSize={FETCH_SIZE}>
+  <DataTableStateProvider columns={tracePickerColumns} pageSize={FETCH_SIZE} lockedColumns={["status"]}>
     <TracePickerContent {...props} />
   </DataTableStateProvider>
 );
diff --git a/frontend/components/traces/trace-picker/trace-picker-content.tsx b/frontend/components/traces/trace-picker/trace-picker-content.tsx
index 5f813eeff..475262024 100644
--- a/frontend/components/traces/trace-picker/trace-picker-content.tsx
+++ b/frontend/components/traces/trace-picker/trace-picker-content.tsx
@@ -118,7 +118,6 @@ const TracePickerContent = ({
         isLoading={isLoading}
         fetchNextPage={fetchNextPage}
         estimatedRowHeight={36}
-        lockedColumns={["status"]}
       >
         <div className="flex gap-2 w-full items-center">
           <DateRangeFilter mode="state" value={dateRange} onChange={setDateRange} />
diff --git a/frontend/components/traces/traces-table/columns.tsx b/frontend/components/traces/traces-table/columns.tsx
index d51fc316e..5adbe4664 100644
--- a/frontend/components/traces/traces-table/columns.tsx
+++ b/frontend/components/traces/traces-table/columns.tsx
@@ -360,19 +360,3 @@ export const filters: ColumnFilter[] = [
     dataType: "string",
   },
 ];
-
-export const defaultTracesColumnOrder = [
-  "status",
-  "id",
-  "top_span_type",
-  "root_span_input",
-  "root_span_output",
-  "start_time",
-  "duration",
-  "cost",
-  "total_tokens",
-  "tags",
-  "metadata",
-  "session_id",
-  "user_id",
-];
diff --git a/frontend/components/traces/traces-table/index.tsx b/frontend/components/traces/traces-table/index.tsx
index 41711956e..0ae1c88a4 100644
--- a/frontend/components/traces/traces-table/index.tsx
+++ b/frontend/components/traces/traces-table/index.tsx
@@ -3,20 +3,19 @@ import { type Row } from "@tanstack/react-table";
 import { isEmpty, map } from "lodash";
 import { useParams, usePathname, useRouter, useSearchParams } from "next/navigation";
 import { useCallback, useEffect, useMemo, useRef } from "react";
-import { useStore } from "zustand";
 
 import { useTimeSeriesStatsUrl } from "@/components/charts/time-series-chart/use-time-series-stats-url";
 import AdvancedSearch from "@/components/common/advanced-search";
 import { useTraceViewNavigation } from "@/components/traces/trace-view/navigation-context";
 import TracesChart from "@/components/traces/traces-chart";
 import { useTracesStoreContext } from "@/components/traces/traces-store";
-import { defaultTracesColumnOrder, filters } from "@/components/traces/traces-table/columns";
+import { filters } from "@/components/traces/traces-table/columns";
 import TracesColumnsMenu from "@/components/traces/traces-table/traces-columns-menu";
 import { useTracesTableStore } from "@/components/traces/traces-table/traces-table-store";
 import DateRangeFilter from "@/components/ui/date-range-filter";
 import { InfiniteDataTable } from "@/components/ui/infinite-datatable";
 import { useInfiniteScroll } from "@/components/ui/infinite-datatable/hooks";
-import { DataTableStateProvider, useDataTableStore } from "@/components/ui/infinite-datatable/model/datatable-store";
+import { DataTableStateProvider } from "@/components/ui/infinite-datatable/model/datatable-store";
 import DataTableFilter from "@/components/ui/infinite-datatable/ui/datatable-filter";
 import RefreshButton from "@/components/ui/infinite-datatable/ui/refresh-button.tsx";
 import { Switch } from "@/components/ui/switch";
@@ -29,14 +28,10 @@ const FETCH_SIZE = 50;
 const DEFAULT_TARGET_BARS = 48;
 
 export default function TracesTable() {
-  const customColumns = useTracesTableStore((s) => s.customColumns);
-  const defaultColumnOrder = useMemo(
-    () => [...defaultTracesColumnOrder, ...customColumns.map((cc) => `custom:${cc.name}`)],
-    [customColumns]
-  );
+  const columnDefs = useTracesTableStore((s) => s.columnDefs);
 
   return (
-    <DataTableStateProvider storageKey="traces-table" defaultColumnOrder={defaultColumnOrder}>
+    <DataTableStateProvider storageKey="traces-table" columns={columnDefs} lockedColumns={["status"]}>
       <TracesTableContent />
     </DataTableStateProvider>
   );
@@ -85,7 +80,6 @@ function TracesTableContent() {
   // Initialize column defs (rebuild when store hydrates custom columns)
   const rebuildColumns = useTracesTableStore((s) => s.rebuildColumns);
   const columnDefs = useTracesTableStore((s) => s.columnDefs);
-  const removeCustomColumn = useTracesTableStore((s) => s.removeCustomColumn);
   const buildFetchParams = useTracesTableStore((s) => s.buildFetchParams);
 
   const customColumns = useTracesTableStore((s) => s.customColumns);
@@ -94,35 +88,8 @@ function TracesTableContent() {
     rebuildColumns();
   }, [customColumns, rebuildColumns]);
 
-  // SQL strings from column defs — only changes when columns structurally change.
-  // useInfiniteScroll uses JSON.stringify on deps, so identical SQL strings
-  // produce the same string → no spurious re-fetch.
   const columnSqls = useMemo(() => columnDefs.map((c) => c.meta?.sql).filter(Boolean), [columnDefs]);
 
-  // Sync datatable columnOrder with traces store columnDefs
-  const datatableStore = useDataTableStore();
-  const { columnOrder, setColumnOrder } = useStore(datatableStore, (s) => ({
-    columnOrder: s.columnOrder,
-    setColumnOrder: s.setColumnOrder,
-  }));
-
-  useEffect(() => {
-    // Skip sync before the store has hydrated columnDefs to avoid wiping saved column order.
-    if (columnDefs.length === 0) return;
-
-    const visibleIds = columnDefs.map((c) => c.id!);
-    const currentSet = new Set(columnOrder);
-    const defSet = new Set(visibleIds);
-
-    const toAdd = visibleIds.filter((id) => !currentSet.has(id));
-    const toRemove = columnOrder.filter((id) => !defSet.has(id));
-
-    if (toAdd.length > 0 || toRemove.length > 0) {
-      const filtered = columnOrder.filter((id) => defSet.has(id));
-      setColumnOrder([...filtered, ...toAdd]);
-    }
-  }, [columnDefs, columnOrder, setColumnOrder]);
-
   useEffect(() => {
     if (!chartContainerRef.current) return;
 
@@ -363,18 +330,6 @@ function TracesTableContent() {
     [searchParams, router, pathName]
   );
 
-  const columnLabels = useMemo(
-    () =>
-      columnDefs.map((column) => ({
-        id: column.id!,
-        label: typeof column.header === "string" ? column.header : column.id!,
-        ...(column.id!.startsWith("custom:") && {
-          onDelete: () => removeCustomColumn(column.id!.replace("custom:", "")),
-        }),
-      })),
-    [columnDefs, removeCustomColumn]
-  );
-
   return (
     <div className="flex flex-1 overflow-hidden px-4 pb-4">
       <InfiniteDataTable<TraceRow>
@@ -389,14 +344,13 @@ function TracesTableContent() {
         isLoading={isLoading}
         fetchNextPage={fetchNextPage}
         getRowHref={getRowHref}
-        lockedColumns={["status"]}
         sortBy={sortBy}
         sortDirection={sortDirection}
         onSort={handleSort}
       >
         <div className="flex flex-1 w-full h-full gap-2">
           <DataTableFilter columns={filters} />
-          <TracesColumnsMenu lockedColumns={["status"]} columnLabels={columnLabels} />
+          <TracesColumnsMenu />
           <DateRangeFilter />
           <RefreshButton onClick={handleRefresh} variant="outline" />
           <div className="flex items-center gap-2 px-2 border rounded-md bg-background h-7">
diff --git a/frontend/components/traces/traces-table/traces-columns-menu/index.tsx b/frontend/components/traces/traces-table/traces-columns-menu/index.tsx
index 432d1f69a..3f0bfa5ac 100644
--- a/frontend/components/traces/traces-table/traces-columns-menu/index.tsx
+++ b/frontend/components/traces/traces-table/traces-columns-menu/index.tsx
@@ -5,14 +5,10 @@ import { useMemo } from "react";
 import { useTracesTableStore } from "@/components/traces/traces-table/traces-table-store";
 import { type ColumnActions, ColumnsMenu, type CustomColumnPanelConfig } from "@/components/ui/columns-menu";
 
-interface TracesColumnsMenuProps {
-  lockedColumns?: string[];
-  columnLabels?: { id: string; label: string; onDelete?: () => void }[];
-}
-
-export default function TracesColumnsMenu({ lockedColumns = [], columnLabels = [] }: TracesColumnsMenuProps) {
+export default function TracesColumnsMenu() {
   const addCustomColumn = useTracesTableStore((s) => s.addCustomColumn);
   const updateCustomColumn = useTracesTableStore((s) => s.updateCustomColumn);
+  const removeCustomColumn = useTracesTableStore((s) => s.removeCustomColumn);
 
   const panelConfig = useMemo<CustomColumnPanelConfig>(
     () => ({
@@ -32,17 +28,11 @@ export default function TracesColumnsMenu({ lockedColumns = [], columnLabels = [
     () => ({
       addCustomColumn,
       updateCustomColumn,
+      removeCustomColumn,
       getColumnDef: (columnId) => useTracesTableStore.getState().columnDefs.find((c) => c.id === columnId),
     }),
-    [addCustomColumn, updateCustomColumn]
+    [addCustomColumn, updateCustomColumn, removeCustomColumn]
   );
 
-  return (
-    <ColumnsMenu
-      lockedColumns={lockedColumns}
-      columnLabels={columnLabels}
-      panelConfig={panelConfig}
-      columnActions={columnActions}
-    />
-  );
+  return <ColumnsMenu panelConfig={panelConfig} columnActions={columnActions} />;
 }
diff --git a/frontend/components/ui/columns-menu/columns-list-panel.tsx b/frontend/components/ui/columns-menu/columns-list-panel.tsx
index 925e85b4d..63e17c872 100644
--- a/frontend/components/ui/columns-menu/columns-list-panel.tsx
+++ b/frontend/components/ui/columns-menu/columns-list-panel.tsx
@@ -17,28 +17,29 @@ import { ColumnsMenuItem } from "./columns-menu-item";
 interface ColumnsListPanelProps {
   columnOrder: string[];
   columnVisibility: Record<string, boolean>;
-  columnLabels: { id: string; label: string; onDelete?: () => void }[];
+  columnLabelMap: Record<string, string>;
   lockedColumns: string[];
   onReorder: (newOrder: string[]) => void;
   onToggleVisibility: (columnId: string) => void;
   onReset: () => void;
-  onCustomColumnClick: () => void;
+  onCustomColumnClick?: () => void;
   onEditColumn?: (columnId: string) => void;
-  /** Whether to show the "Create column with SQL" button. Defaults to true. */
+  onDeleteColumn?: (columnId: string) => void;
   showCreateButton?: boolean;
 }
 
 export const ColumnsListPanel = ({
   columnOrder,
   columnVisibility,
-  columnLabels,
+  columnLabelMap,
   lockedColumns,
   onReorder,
   onToggleVisibility,
   onReset,
   onCustomColumnClick,
   onEditColumn,
-  showCreateButton = true,
+  onDeleteColumn,
+  showCreateButton = false,
 }: ColumnsListPanelProps) => {
   const sensors = useSensors(
     useSensor(PointerSensor),
@@ -78,17 +79,17 @@ export const ColumnsListPanel = ({
           >
             <SortableContext items={columnOrder} strategy={verticalListSortingStrategy}>
               {columnOrder.map((columnId) => {
-                const labelEntry = columnLabels?.find((col) => col.id === columnId);
+                const isCustom = columnId.startsWith("custom:");
                 return (
                   <ColumnsMenuItem
                     key={columnId}
                     id={columnId}
-                    label={labelEntry?.label || columnId}
+                    label={columnLabelMap[columnId] || columnId}
                     isVisible={columnVisibility[columnId] !== false}
                     isLocked={lockedColumns.includes(columnId)}
                     onToggleVisibility={onToggleVisibility}
-                    onDelete={labelEntry?.onDelete}
-                    onEdit={columnId.startsWith("custom:") && onEditColumn ? () => onEditColumn(columnId) : undefined}
+                    onDelete={isCustom && onDeleteColumn ? () => onDeleteColumn(columnId) : undefined}
+                    onEdit={isCustom && onEditColumn ? () => onEditColumn(columnId) : undefined}
                   />
                 );
               })}
@@ -105,7 +106,7 @@ export const ColumnsListPanel = ({
           <ListRestart className="w-3.5 h-3.5 text-secondary-foreground mr-2" />
           Reset columns
         </div>
-        {showCreateButton && (
+        {showCreateButton && onCustomColumnClick && (
           <div
             className="relative flex cursor-default select-none items-center rounded-sm px-2 py-1.5 text-xs transition-colors hover:bg-accent hover:text-accent-foreground"
             onClick={onCustomColumnClick}
diff --git a/frontend/components/ui/columns-menu/columns-menu.tsx b/frontend/components/ui/columns-menu/columns-menu.tsx
index dc745bbf2..20f7127b2 100644
--- a/frontend/components/ui/columns-menu/columns-menu.tsx
+++ b/frontend/components/ui/columns-menu/columns-menu.tsx
@@ -11,34 +11,35 @@ import { CustomColumnPanel } from "./custom-column-panel";
 import type { ColumnActions, CustomColumn, CustomColumnPanelConfig } from "./types";
 
 interface ColumnsMenuProps {
-  lockedColumns?: string[];
-  columnLabels?: { id: string; label: string; onDelete?: () => void }[];
-  /** Configuration for the custom column panel (schema, test query, etc.). */
-  panelConfig: CustomColumnPanelConfig;
-  /** Store actions for managing custom columns. */
-  columnActions: ColumnActions;
+  /** Configuration for the custom column panel (schema, test query, etc.). When omitted, the custom column UI is hidden. */
+  panelConfig?: CustomColumnPanelConfig;
+  /** Store actions for managing custom columns. Required when panelConfig is provided. */
+  columnActions?: ColumnActions;
   /** Whether to show the "Create column with SQL" button. Defaults to true. */
   showCreateButton?: boolean;
 }
 
-export default function ColumnsMenu({
-  lockedColumns = [],
-  columnLabels = [],
-  panelConfig,
-  columnActions,
-  showCreateButton = true,
-}: ColumnsMenuProps) {
+export default function ColumnsMenu({ panelConfig, columnActions, showCreateButton = true }: ColumnsMenuProps) {
   const store = useDataTableStore();
-  const { resetColumns, columnOrder, setColumnOrder, columnVisibility, setColumnVisibility } = useStore(
-    store,
-    (state) => ({
-      resetColumns: state.resetColumns,
-      columnOrder: state.columnOrder,
-      setColumnOrder: state.setColumnOrder,
-      columnVisibility: state.columnVisibility,
-      setColumnVisibility: state.setColumnVisibility,
-    })
-  );
+  const {
+    lockedColumns,
+    columnLabelMap,
+    resetColumns,
+    columnOrder,
+    setColumnOrder,
+    columnVisibility,
+    setColumnVisibility,
+  } = useStore(store, (state) => ({
+    lockedColumns: state.lockedColumns,
+    columnLabelMap: state.columnLabelMap,
+    resetColumns: state.resetColumns,
+    columnOrder: state.columnOrder,
+    setColumnOrder: state.setColumnOrder,
+    columnVisibility: state.columnVisibility,
+    setColumnVisibility: state.setColumnVisibility,
+  }));
+
+  const hasCustomColumns = !!panelConfig && !!columnActions;
 
   const [isOpen, setIsOpen] = useState(false);
   const [activePanel, setActivePanel] = useState<"list" | "form">("list");
@@ -54,6 +55,7 @@ export default function ColumnsMenu({
   }
 
   const handleEditColumn = (columnId: string) => {
+    if (!columnActions) return;
     const col = columnActions.getColumnDef(columnId);
     if (col?.meta?.isCustom) {
       setEditingColumn({
@@ -65,7 +67,13 @@ export default function ColumnsMenu({
     }
   };
 
+  const handleDeleteColumn = (columnId: string) => {
+    if (!columnActions) return;
+    columnActions.removeCustomColumn(columnId.replace("custom:", ""));
+  };
+
   const handleSave = (column: CustomColumn) => {
+    if (!columnActions) return;
     if (editingColumn) {
       columnActions.updateCustomColumn(editingColumn.name, column);
     } else {
@@ -95,7 +103,6 @@ export default function ColumnsMenu({
         align="start"
         onOpenAutoFocus={(e) => e.preventDefault()}
         onInteractOutside={(e) => {
-          // Prevent closing when interacting with CodeMirror autocomplete tooltips
           const target = e.target as HTMLElement | null;
           if (target?.closest(".cm-tooltip-autocomplete")) {
             e.preventDefault();
@@ -107,19 +114,24 @@ export default function ColumnsMenu({
             <ColumnsListPanel
               columnOrder={columnOrder}
               columnVisibility={columnVisibility}
-              columnLabels={columnLabels}
+              columnLabelMap={columnLabelMap}
               lockedColumns={lockedColumns}
               onReorder={setColumnOrder}
               onToggleVisibility={handleToggleVisibility}
               onReset={resetColumns}
-              onCustomColumnClick={() => {
-                setEditingColumn(null);
-                setActivePanel("form");
-              }}
-              onEditColumn={handleEditColumn}
-              showCreateButton={showCreateButton}
+              onCustomColumnClick={
+                hasCustomColumns
+                  ? () => {
+                      setEditingColumn(null);
+                      setActivePanel("form");
+                    }
+                  : undefined
+              }
+              onEditColumn={hasCustomColumns ? handleEditColumn : undefined}
+              onDeleteColumn={hasCustomColumns ? handleDeleteColumn : undefined}
+              showCreateButton={hasCustomColumns && showCreateButton}
             />
-          ) : (
+          ) : panelConfig ? (
             <CustomColumnPanel
               key={editingColumn?.name ?? "__new__"}
               onBack={() => {
@@ -130,7 +142,7 @@ export default function ColumnsMenu({
               editingColumn={editingColumn ?? undefined}
               config={panelConfig}
             />
-          )}
+          ) : null}
         </AnimatePresence>
       </PopoverContent>
     </Popover>
diff --git a/frontend/components/ui/columns-menu/types.ts b/frontend/components/ui/columns-menu/types.ts
index 08f439684..e9d96a424 100644
--- a/frontend/components/ui/columns-menu/types.ts
+++ b/frontend/components/ui/columns-menu/types.ts
@@ -32,6 +32,7 @@ export interface CustomColumnPanelConfig {
 export interface ColumnActions {
   addCustomColumn: (column: CustomColumn) => void;
   updateCustomColumn: (oldName: string, column: CustomColumn) => void;
+  removeCustomColumn: (name: string) => void;
   /** Return the ColumnDef for a given column ID, used for populating the edit form. */
   getColumnDef: (columnId: string) => ColumnDef<any> | undefined;
 }
diff --git a/frontend/components/ui/infinite-datatable/index.tsx b/frontend/components/ui/infinite-datatable/index.tsx
index 6c9fd118b..458b3ada8 100644
--- a/frontend/components/ui/infinite-datatable/index.tsx
+++ b/frontend/components/ui/infinite-datatable/index.tsx
@@ -50,7 +50,6 @@ export function InfiniteDataTable<TData extends RowData>({
   onRowClick,
   focusedRowId,
   selectionPanel,
-  lockedColumns = EMPTY_ARRAY as string[],
 
   // Styling
   className,
@@ -293,7 +292,6 @@ export function InfiniteDataTable<TData extends RowData>({
                 onHideColumn={(columnId) => {
                   setColumnVisibility({ ...columnVisibility, [columnId]: false });
                 }}
-                lockedColumns={lockedColumns}
               />
               <InfiniteDatatableBody
                 table={table}
diff --git a/frontend/components/ui/infinite-datatable/model/datatable-store.tsx b/frontend/components/ui/infinite-datatable/model/datatable-store.tsx
index 17b9b8ba4..5b16448e2 100644
--- a/frontend/components/ui/infinite-datatable/model/datatable-store.tsx
+++ b/frontend/components/ui/infinite-datatable/model/datatable-store.tsx
@@ -1,7 +1,8 @@
 "use client";
 
-import { intersection, pick, uniqBy } from "lodash";
-import { createContext, type ReactNode, useContext, useState } from "react";
+import { type ColumnDef } from "@tanstack/react-table";
+import { isEqual, uniqBy } from "lodash";
+import { createContext, type ReactNode, useContext, useEffect, useMemo, useState } from "react";
 import { createStore, type StoreApi } from "zustand";
 import { persist } from "zustand/middleware";
 
@@ -30,6 +31,9 @@ export interface InfiniteScrollActions<TData> {
 
 export interface SelectionState {
   selectedRows: Set<string>;
+  defaultColumnOrder: string[];
+  lockedColumns: string[];
+  columnLabelMap: Record<string, string>;
   columnVisibility: Record<string, boolean>;
   columnOrder: string[];
   columnSizing: Record<string, number>;
@@ -46,6 +50,7 @@ export interface SelectionActions {
   setColumnOrder: (order: string[]) => void;
   setColumnSizing: (sizing: Record<string, number>) => void;
   setDraggingColumnId: (columnId: string | null) => void;
+  reconcileColumns: (columnIds: string[]) => void;
   resetColumns: () => void;
   getStorageKey: () => string;
 }
@@ -58,8 +63,8 @@ type DataTableStore<TData> = InfiniteScrollState<TData> &
 function createDataTableStore<TData>(
   uniqueKey: string = "id",
   storageKey?: string,
-  defaultColumnOrder: string[] = [],
-  pageSize: number = 50
+  pageSize: number = 50,
+  lockedColumns: string[] = []
 ): StoreApi<DataTableStore<TData>> {
   const storeConfig = (
     set: StoreApi<DataTableStore<TData>>["setState"],
@@ -73,8 +78,11 @@ function createDataTableStore<TData>(
     uniqueKey,
     hasMore: true,
     pageSize,
+    defaultColumnOrder: [],
+    lockedColumns,
+    columnLabelMap: {},
     columnVisibility: {},
-    columnOrder: defaultColumnOrder,
+    columnOrder: [],
     columnSizing: {},
     draggingColumnId: null,
     setData: (updater) => set((state) => ({ data: updater(state.data) })),
@@ -87,12 +95,30 @@ function createDataTableStore<TData>(
     setColumnOrder: (order) => set({ columnOrder: order }),
     setColumnSizing: (sizing) => set({ columnSizing: sizing }),
     setDraggingColumnId: (columnId) => set({ draggingColumnId: columnId }),
-    resetColumns: () =>
+    reconcileColumns: (columnIds) => {
+      const state = get();
+      if (isEqual(state.defaultColumnOrder, columnIds)) return;
+
+      if (state.columnOrder.length === 0) {
+        set({ defaultColumnOrder: columnIds, columnOrder: columnIds });
+        return;
+      }
+
+      const idSet = new Set(columnIds);
+      const pruned = state.columnOrder.filter((id) => idSet.has(id));
+      const existingSet = new Set(pruned);
+      const added = columnIds.filter((id) => !existingSet.has(id));
+
+      set({ defaultColumnOrder: columnIds, columnOrder: [...pruned, ...added] });
+    },
+    resetColumns: () => {
+      const state = get();
       set({
         columnVisibility: {},
-        columnOrder: defaultColumnOrder,
+        columnOrder: state.defaultColumnOrder,
         columnSizing: {},
-      }),
+      });
+    },
     appendData: (items, count) =>
       set((state) => {
         const combined = [...state.data, ...items];
@@ -168,21 +194,16 @@ function createDataTableStore<TData>(
           columnOrder: state.columnOrder,
           columnSizing: state.columnSizing,
         }),
+        // Restore persisted state as-is; reconcileColumns (called by the provider) corrects it against the actual columns prop.
         merge: (persistedState, currentState) => {
           const persisted = persistedState as Partial<
             Pick<SelectionState, "columnVisibility" | "columnOrder" | "columnSizing">
           >;
-          const validColumns = intersection(persisted?.columnOrder ?? [], defaultColumnOrder);
-          const newColumns = defaultColumnOrder.filter((col) => !validColumns.includes(col));
-          const mergedColumnOrder = [...validColumns, ...newColumns];
-          const filteredColumnVisibility = pick(persisted?.columnVisibility ?? {}, defaultColumnOrder);
-          const filteredColumnSizing = pick(persisted?.columnSizing ?? {}, defaultColumnOrder);
-
           return {
             ...currentState,
-            columnVisibility: filteredColumnVisibility,
-            columnOrder: mergedColumnOrder,
-            columnSizing: filteredColumnSizing,
+            columnVisibility: persisted?.columnVisibility ?? {},
+            columnOrder: persisted?.columnOrder ?? [],
+            columnSizing: persisted?.columnSizing ?? {},
           };
         },
       })
@@ -201,7 +222,9 @@ export interface DataTableStateProviderProps {
   uniqueKey?: string;
   pageSize?: number;
   storageKey?: string;
-  defaultColumnOrder?: string[];
+  columns?: ColumnDef<any>[];
+  enableRowSelection?: boolean;
+  lockedColumns?: string[];
 }
 
 export function DataTableStateProvider<TData>({
@@ -209,9 +232,33 @@ export function DataTableStateProvider<TData>({
   storageKey,
   uniqueKey = "id",
   pageSize = 50,
-  defaultColumnOrder = [],
+  columns = [],
+  enableRowSelection = false,
+  lockedColumns = [],
 }: DataTableStateProviderProps) {
-  const [store] = useState(() => createDataTableStore<TData>(uniqueKey, storageKey, defaultColumnOrder, pageSize));
+  const [store] = useState(() => createDataTableStore<TData>(uniqueKey, storageKey, pageSize, lockedColumns));
+
+  const columnIds = useMemo(() => {
+    const ids = columns.map((c) => (c as ColumnDef<any> & { id?: string }).id).filter(Boolean) as string[];
+    return enableRowSelection ? ["__row_selection", ...ids] : ids;
+  }, [columns, enableRowSelection]);
+
+  const columnLabelMap = useMemo(() => {
+    const map: Record<string, string> = {};
+    for (const c of columns) {
+      const id = (c as ColumnDef<any> & { id?: string }).id;
+      if (!id) continue;
+      map[id] = typeof c.header === "string" ? c.header : id;
+    }
+    return map;
+  }, [columns]);
+
+  useEffect(() => {
+    if (columnIds.length > 0) {
+      store.getState().reconcileColumns(columnIds);
+      store.setState({ columnLabelMap });
+    }
+  }, [columnIds, columnLabelMap, store]);
 
   return <DataTableContext.Provider value={store}>{children}</DataTableContext.Provider>;
 }
diff --git a/frontend/components/ui/infinite-datatable/model/types.ts b/frontend/components/ui/infinite-datatable/model/types.ts
index 37696f096..c01f12721 100644
--- a/frontend/components/ui/infinite-datatable/model/types.ts
+++ b/frontend/components/ui/infinite-datatable/model/types.ts
@@ -8,12 +8,13 @@ export interface LoadMoreButtonProps {
   hasMore: boolean;
 }
 
-export interface InfiniteDataTableProps<TData extends RowData>
-  extends Omit<Partial<TableOptions<TData>>, "data" | "columns"> {
+export interface InfiniteDataTableProps<TData extends RowData> extends Omit<
+  Partial<TableOptions<TData>>,
+  "data" | "columns"
+> {
   data: TData[];
   columns: TableOptions<TData>["columns"];
 
-  lockedColumns?: string[];
   hasMore: boolean;
   isFetching: boolean;
   isLoading: boolean;
@@ -48,8 +49,6 @@ export interface InfiniteDataTableHeaderProps<TData extends RowData> {
   table: Table<TData>;
   columnOrder: string[];
   onHideColumn: (columnId: string) => void;
-
-  lockedColumns?: string[];
 }
 
 export interface InfiniteDataTableBodyProps<TData extends RowData> {
diff --git a/frontend/components/ui/infinite-datatable/ui/columns-menu-item.tsx b/frontend/components/ui/infinite-datatable/ui/columns-menu-item.tsx
deleted file mode 100644
index 613d644e2..000000000
--- a/frontend/components/ui/infinite-datatable/ui/columns-menu-item.tsx
+++ /dev/null
@@ -1,49 +0,0 @@
-import { useSortable } from "@dnd-kit/sortable";
-import { CSS } from "@dnd-kit/utilities";
-import { GripHorizontal } from "lucide-react";
-import React from "react";
-
-import { DropdownMenuItem } from "@/components/ui/dropdown-menu.tsx";
-import { Switch } from "@/components/ui/switch.tsx";
-import { cn } from "@/lib/utils.ts";
-
-interface ColumnsMenuItemProps {
-  id: string;
-  label: string;
-  isVisible: boolean;
-  isLocked: boolean;
-  onToggleVisibility: (columnId: string) => void;
-}
-
-export const ColumnsMenuItem = ({ id, label, isVisible, isLocked, onToggleVisibility }: ColumnsMenuItemProps) => {
-  const { attributes, listeners, setNodeRef, transform, transition } = useSortable({
-    id,
-    disabled: isLocked,
-  });
-  const style = {
-    transform: CSS.Transform.toString(transform),
-    transition,
-  };
-  return (
-    <DropdownMenuItem style={style} className={cn("flex items-center gap-4", isLocked && "hidden")}>
-      <div className="flex items-center gap-2">
-        <div
-          ref={setNodeRef}
-          {...(!isLocked && { ...attributes, ...listeners })}
-          className={cn("cursor-grab", isLocked && "cursor-not-allowed opacity-50")}
-        >
-          <GripHorizontal />
-        </div>
-        <span className={isLocked ? "text-muted-foreground mr-4" : ""}>{label}</span>
-      </div>
-
-      <Switch
-        className="ml-auto h-fit"
-        checked={isVisible}
-        disabled={isLocked}
-        onCheckedChange={() => !isLocked && onToggleVisibility(id)}
-        onClick={(e) => e.stopPropagation()}
-      />
-    </DropdownMenuItem>
-  );
-};
diff --git a/frontend/components/ui/infinite-datatable/ui/columns-menu.tsx b/frontend/components/ui/infinite-datatable/ui/columns-menu.tsx
deleted file mode 100644
index b5c26614b..000000000
--- a/frontend/components/ui/infinite-datatable/ui/columns-menu.tsx
+++ /dev/null
@@ -1,106 +0,0 @@
-import { closestCenter, DndContext, KeyboardSensor, PointerSensor, useSensor, useSensors } from "@dnd-kit/core";
-import { restrictToVerticalAxis } from "@dnd-kit/modifiers";
-import {
-  arrayMove,
-  SortableContext,
-  sortableKeyboardCoordinates,
-  verticalListSortingStrategy,
-} from "@dnd-kit/sortable";
-import { DropdownMenuSeparator } from "@radix-ui/react-dropdown-menu";
-import { ListRestart } from "lucide-react";
-import React from "react";
-import { useStore } from "zustand";
-
-import { Button } from "@/components/ui/button.tsx";
-import {
-  DropdownMenu,
-  DropdownMenuContent,
-  DropdownMenuGroup,
-  DropdownMenuItem,
-  DropdownMenuTrigger,
-} from "@/components/ui/dropdown-menu.tsx";
-import { useDataTableStore } from "@/components/ui/infinite-datatable/model/datatable-store.tsx";
-import { ColumnsMenuItem } from "@/components/ui/infinite-datatable/ui/columns-menu-item.tsx";
-
-interface ColumnsMenuProps {
-  lockedColumns?: string[];
-  columnLabels?: { id: string; label: string }[];
-}
-
-export default function ColumnsMenu({ lockedColumns = [], columnLabels = [] }: ColumnsMenuProps) {
-  const store = useDataTableStore();
-  const { resetColumns, columnOrder, setColumnOrder, columnVisibility, setColumnVisibility } = useStore(
-    store,
-    (state) => ({
-      resetColumns: state.resetColumns,
-      columnOrder: state.columnOrder,
-      setColumnOrder: state.setColumnOrder,
-      columnVisibility: state.columnVisibility,
-      setColumnVisibility: state.setColumnVisibility,
-    })
-  );
-  const sensors = useSensors(
-    useSensor(PointerSensor),
-    useSensor(KeyboardSensor, {
-      coordinateGetter: sortableKeyboardCoordinates,
-    })
-  );
-  function handleDragEnd(event: any) {
-    const { active, over } = event;
-
-    if (active && over && active.id !== over.id) {
-      const oldIndex = columnOrder.indexOf(active.id);
-      const newIndex = columnOrder.indexOf(over.id);
-      if (oldIndex !== -1 && newIndex !== -1) {
-        const newColumnOrder = arrayMove(columnOrder, oldIndex, newIndex);
-        setColumnOrder(newColumnOrder);
-      }
-    }
-  }
-
-  function handleToggleVisibility(columnId: string) {
-    if (lockedColumns.includes(columnId)) return;
-
-    setColumnVisibility({
-      ...columnVisibility,
-      [columnId]: !(columnVisibility[columnId] !== false),
-    });
-  }
-  return (
-    <DropdownMenu>
-      <DropdownMenuTrigger asChild>
-        <Button className="text-secondary-foreground focus:outline-0" icon="columns2" variant="outline">
-          Columns
-        </Button>
-      </DropdownMenuTrigger>
-      <DropdownMenuContent align="start" className="relative min-w-32">
-        <DropdownMenuGroup>
-          <DndContext
-            sensors={sensors}
-            collisionDetection={closestCenter}
-            onDragEnd={handleDragEnd}
-            modifiers={[restrictToVerticalAxis]}
-          >
-            <SortableContext items={columnOrder} strategy={verticalListSortingStrategy}>
-              {columnOrder.map((columnId) => (
-                <ColumnsMenuItem
-                  key={columnId}
-                  id={columnId}
-                  label={columnLabels?.find((col) => col.id === columnId)?.label || columnId}
-                  isVisible={columnVisibility[columnId] !== false}
-                  isLocked={lockedColumns.includes(columnId)}
-                  onToggleVisibility={handleToggleVisibility}
-                />
-              ))}
-            </SortableContext>
-          </DndContext>
-        </DropdownMenuGroup>
-        <DropdownMenuSeparator />
-        <DropdownMenuItem onClick={resetColumns}>
-          <ListRestart className="w-3.5 h-3.5 text-secondary-foreground" />
-          Reset columns
-        </DropdownMenuItem>
-      </DropdownMenuContent>
-    </DropdownMenu>
-  );
-}
diff --git a/frontend/components/ui/infinite-datatable/ui/header.tsx b/frontend/components/ui/infinite-datatable/ui/header.tsx
index db85db9a5..d288f761a 100644
--- a/frontend/components/ui/infinite-datatable/ui/header.tsx
+++ b/frontend/components/ui/infinite-datatable/ui/header.tsx
@@ -1,17 +1,22 @@
 import { horizontalListSortingStrategy, SortableContext } from "@dnd-kit/sortable";
 import { type RowData } from "@tanstack/react-table";
 import { forwardRef } from "react";
+import { useStore } from "zustand";
 
 import { TableHeader, TableRow } from "@/components/ui/table.tsx";
 
+import { useDataTableStore } from "../model/datatable-store.tsx";
 import { type InfiniteDataTableHeaderProps } from "../model/types.ts";
 import { InfiniteTableHead } from "./head.tsx";
 
 export const InfiniteDatatableHeader = forwardRef<HTMLTableSectionElement, InfiniteDataTableHeaderProps<RowData>>(
   function InfiniteDatatableHeader<TData extends RowData>(
-    { table, columnOrder, onHideColumn, lockedColumns }: InfiniteDataTableHeaderProps<TData>,
+    { table, columnOrder, onHideColumn }: InfiniteDataTableHeaderProps<TData>,
     ref: React.Ref<HTMLTableSectionElement>
   ) {
+    const store = useDataTableStore();
+    const lockedColumns = useStore(store, (s) => s.lockedColumns);
+
     return (
       <TableHeader
         ref={ref}

From a55c33ccd0db343420e37e65e1d02679c6039614 Mon Sep 17 00:00:00 2001
From: Olzhas Nurpeisov <olzhik1123@gmail.com>
Date: Mon, 23 Mar 2026 23:17:13 +0000
Subject: [PATCH 2/2] feat: final cleanup

---
 frontend/components/dataset/dataset.tsx       |   3 +-
 frontend/components/datasets/datasets.tsx     |   3 +-
 .../sidebar/runs-tab.tsx                      |   3 +-
 .../debugger-sessions/debugger-sessions.tsx   |   3 +-
 .../eval-columns-menu/eval-columns-menu.tsx   |  25 +-
 .../evaluation-datapoints-table-content.tsx   |  32 ++-
 .../evaluation-datapoints-table/index.tsx     |  16 +-
 frontend/components/evaluation/evaluation.tsx |  62 ++---
 frontend/components/evaluation/store.ts       | 111 +++------
 .../evaluations/evaluations-groups-bar.tsx    |   3 +-
 .../components/evaluations/evaluations.tsx    |   3 +-
 .../evaluators/evaluators-table.tsx           |   3 -
 frontend/components/evaluators/evaluators.tsx |   2 +-
 .../playground/playground-history-table.tsx   |   3 +-
 .../components/playgrounds/playgrounds.tsx    |   3 +-
 frontend/components/queues/queues.tsx         |   3 +-
 .../shared/evaluation/shared-evaluation.tsx   |  43 ++--
 .../signal/create-signal-job/index.tsx        |   5 +-
 .../components/signal/events-table/index.tsx  |   5 +-
 .../components/signal/jobs-table/index.tsx    |   3 +-
 .../components/signal/runs-table/index.tsx    |  17 +-
 .../signal/triggers-table/index.tsx           |   9 +-
 frontend/components/sql/editor-panel.tsx      |   3 +-
 .../traces/sessions-table/index.tsx           |   3 +-
 .../components/traces/spans-table/index.tsx   |   3 +-
 .../components/traces/trace-picker/index.tsx  |   2 +-
 .../trace-picker/trace-picker-content.tsx     |   3 +-
 .../traces/traces-table/columns.tsx           |   3 -
 .../components/traces/traces-table/index.tsx  |  54 ++---
 .../traces-columns-menu/index.tsx             |  26 +-
 .../traces/traces-table/traces-table-store.ts | 118 +++------
 .../ui/columns-menu/columns-menu.tsx          |  84 ++++---
 frontend/components/ui/columns-menu/index.ts  |   2 +-
 frontend/components/ui/columns-menu/types.ts  |  13 +-
 .../ui/infinite-datatable/index.tsx           |  11 +-
 .../model/datatable-store.tsx                 | 225 +++++++++++++-----
 .../ui/infinite-datatable/model/types.ts      |   3 +-
 37 files changed, 446 insertions(+), 467 deletions(-)

diff --git a/frontend/components/dataset/dataset.tsx b/frontend/components/dataset/dataset.tsx
index 9e35579e1..c34612f15 100644
--- a/frontend/components/dataset/dataset.tsx
+++ b/frontend/components/dataset/dataset.tsx
@@ -287,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}
@@ -349,7 +348,7 @@ export default function Dataset(props: DatasetProps) {
   return (
     <DataTableStateProvider
       storageKey="dataset-table"
-      columns={columns}
+      columnDefs={columns}
       enableRowSelection
       lockedColumns={["__row_selection"]}
     >
diff --git a/frontend/components/datasets/datasets.tsx b/frontend/components/datasets/datasets.tsx
index fefb43fd6..eaf5d3858 100644
--- a/frontend/components/datasets/datasets.tsx
+++ b/frontend/components/datasets/datasets.tsx
@@ -207,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}
@@ -245,7 +244,7 @@ export default function Datasets() {
   return (
     <DataTableStateProvider
       storageKey="datasets-table"
-      columns={columns}
+      columnDefs={columns}
       enableRowSelection
       lockedColumns={["__row_selection"]}
     >
diff --git a/frontend/components/debugger-sessions/debugger-session-view/sidebar/runs-tab.tsx b/frontend/components/debugger-sessions/debugger-session-view/sidebar/runs-tab.tsx
index 4439ab42d..933f442d2 100644
--- a/frontend/components/debugger-sessions/debugger-session-view/sidebar/runs-tab.tsx
+++ b/frontend/components/debugger-sessions/debugger-session-view/sidebar/runs-tab.tsx
@@ -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}
@@ -93,7 +92,7 @@ const RunsContent = () => {
 
 export default function RunsTab() {
   return (
-    <DataTableStateProvider columns={tracePickerColumns} pageSize={FETCH_SIZE} lockedColumns={["status"]}>
+    <DataTableStateProvider columnDefs={tracePickerColumns} pageSize={FETCH_SIZE} lockedColumns={["status"]}>
       <RunsContent />
     </DataTableStateProvider>
   );
diff --git a/frontend/components/debugger-sessions/debugger-sessions.tsx b/frontend/components/debugger-sessions/debugger-sessions.tsx
index 2854a9363..82568bf01 100644
--- a/frontend/components/debugger-sessions/debugger-sessions.tsx
+++ b/frontend/components/debugger-sessions/debugger-sessions.tsx
@@ -165,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}`}
@@ -190,7 +189,7 @@ function DebuggerSessionsContent() {
 
 export default function DebuggerSessions() {
   return (
-    <DataTableStateProvider storageKey="debugger-sessions-table" columns={columns}>
+    <DataTableStateProvider storageKey="debugger-sessions-table" columnDefs={columns}>
       <DebuggerSessionsContent />
     </DataTableStateProvider>
   );
diff --git a/frontend/components/evaluation/eval-columns-menu/eval-columns-menu.tsx b/frontend/components/evaluation/eval-columns-menu/eval-columns-menu.tsx
index d7114b76f..1711a295f 100644
--- a/frontend/components/evaluation/eval-columns-menu/eval-columns-menu.tsx
+++ b/frontend/components/evaluation/eval-columns-menu/eval-columns-menu.tsx
@@ -1,17 +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";
 
 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 removeCustomColumn = useEvalStore((s) => s.removeCustomColumn);
+  const store = useDataTableStore();
 
   const panelConfig = useMemo<CustomColumnPanelConfig>(
     () => ({
@@ -20,24 +19,16 @@ export default function EvalColumnsMenu() {
       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,
-      removeCustomColumn,
-      getColumnDef: (columnId) => useEvalStore.getState().columnDefs.find((c) => c.id === columnId),
-    }),
-    [addCustomColumn, updateCustomColumn, removeCustomColumn]
-  );
+  const getColumnDefs = useCallback(() => selectAllColumnDefs(store.getState()), [store]);
 
-  return <ColumnsMenu panelConfig={panelConfig} columnActions={columnActions} showCreateButton={!isShared} />;
+  return <ColumnsMenu panelConfig={panelConfig} getColumnDefs={getColumnDefs} showCreateButton={!isShared} />;
 }
diff --git a/frontend/components/evaluation/evaluation-datapoints-table/evaluation-datapoints-table-content.tsx b/frontend/components/evaluation/evaluation-datapoints-table/evaluation-datapoints-table-content.tsx
index 070593ef3..a1a13908a 100644
--- a/frontend/components/evaluation/evaluation-datapoints-table/evaluation-datapoints-table-content.tsx
+++ b/frontend/components/evaluation/evaluation-datapoints-table/evaluation-datapoints-table-content.tsx
@@ -1,11 +1,10 @@
 import { Settings as SettingsIcon } from "lucide-react";
 import { usePathname, useRouter, useSearchParams } from "next/navigation";
 import { useCallback, useEffect, useMemo } from "react";
-import { useShallow } from "zustand/react/shallow";
 
 import EvalColumnsMenu from "@/components/evaluation/eval-columns-menu";
 import SearchEvaluationInput from "@/components/evaluation/search-evaluation-input";
-import { selectVisibleColumns, useEvalStore } from "@/components/evaluation/store";
+import { useEvalStore } from "@/components/evaluation/store";
 import { Button } from "@/components/ui/button";
 import {
   DropdownMenu,
@@ -15,6 +14,10 @@ import {
   DropdownMenuTrigger,
 } from "@/components/ui/dropdown-menu";
 import { InfiniteDataTable } from "@/components/ui/infinite-datatable";
+import {
+  selectAllColumnDefs,
+  useDataTableStoreSelector,
+} from "@/components/ui/infinite-datatable/model/datatable-store";
 import DataTableFilter, { DataTableFilterList } from "@/components/ui/infinite-datatable/ui/datatable-filter";
 import { Switch } from "@/components/ui/switch";
 import { type EvalRow } from "@/lib/evaluation/types";
@@ -40,12 +43,20 @@ const EvaluationDatapointsTableContent = ({
   const sortBy = searchParams.get("sortBy") ?? undefined;
   const sortDirection = (searchParams.get("sortDirection")?.toLowerCase() ?? undefined) as "asc" | "desc" | undefined;
 
-  // Store state
-  const columns = useEvalStore((s) => s.columnDefs);
   const heatmapEnabled = useEvalStore((s) => s.heatmapEnabled);
   const setHeatmapEnabled = useEvalStore((s) => s.setHeatmapEnabled);
   const setScoreRanges = useEvalStore((s) => s.setScoreRanges);
-  // Compute and set score ranges from data
+  const isComparison = useEvalStore((s) => s.isComparison);
+
+  const allColumnDefs = useDataTableStoreSelector(selectAllColumnDefs);
+  const setColumnVisibility = useDataTableStoreSelector((s) => s.setColumnVisibility);
+  const columnVisibility = useDataTableStoreSelector((s) => s.columnVisibility);
+  useEffect(() => {
+    if (isComparison && columnVisibility["output"] !== false) {
+      setColumnVisibility({ ...columnVisibility, output: false });
+    }
+  }, [isComparison, columnVisibility, setColumnVisibility]);
+
   useEffect(() => {
     if (!data) return;
 
@@ -94,13 +105,9 @@ const EvaluationDatapointsTableContent = ({
     [searchParams, router, pathname]
   );
 
-  // Visible columns (hidden + output-in-comparison filtered out)
-  const visibleColumns = useEvalStore(useShallow(selectVisibleColumns));
-
-  // Derive filter definitions from column defs in the store
   const columnFilters = useMemo(
     () =>
-      columns
+      allColumnDefs
         .filter((c) => c.meta?.filterable)
         .map((c) => ({
           key: c.id!,
@@ -112,13 +119,12 @@ const EvaluationDatapointsTableContent = ({
                 ? ("number" as const)
                 : ("string" as const),
         })),
-    [columns]
+    [allColumnDefs]
   );
 
   return (
     <div className="flex overflow-hidden flex-1">
-      <InfiniteDataTable
-        columns={visibleColumns}
+      <InfiniteDataTable<EvalRow>
         data={data ?? []}
         hasMore={!searchParams.get("search") && hasMore}
         isFetching={isFetching}
diff --git a/frontend/components/evaluation/evaluation-datapoints-table/index.tsx b/frontend/components/evaluation/evaluation-datapoints-table/index.tsx
index 28ed5951a..c2b03c601 100644
--- a/frontend/components/evaluation/evaluation-datapoints-table/index.tsx
+++ b/frontend/components/evaluation/evaluation-datapoints-table/index.tsx
@@ -1,7 +1,7 @@
 import { type Row } from "@tanstack/react-table";
-import { useShallow } from "zustand/react/shallow";
+import { useMemo } from "react";
 
-import { selectVisibleColumns, useEvalStore } from "@/components/evaluation/store";
+import { buildEvalColumnDefs, buildEvalCustomColumnDef, useEvalStore } from "@/components/evaluation/store";
 import { DataTableStateProvider } from "@/components/ui/infinite-datatable/model/datatable-store";
 import { type EvalRow } from "@/lib/evaluation/types";
 
@@ -21,15 +21,21 @@ export interface EvaluationDatapointsTableProps {
 }
 
 const EvaluationDatapointsTable = (props: EvaluationDatapointsTableProps) => {
-  const { isLoading } = props;
-  const visibleColumns = useEvalStore(useShallow(selectVisibleColumns));
+  const { isLoading, scores } = props;
+  const isShared = useEvalStore((s) => s.isShared);
+
+  const columnDefs = useMemo(() => buildEvalColumnDefs(scores), [scores]);
 
   if (isLoading) {
     return <EvalTableSkeleton />;
   }
 
   return (
-    <DataTableStateProvider storageKey="evaluation-datapoints-table" columns={visibleColumns}>
+    <DataTableStateProvider
+      storageKey="evaluation-datapoints-table"
+      columnDefs={columnDefs}
+      buildCustomColumnDef={isShared ? undefined : buildEvalCustomColumnDef}
+    >
       <EvaluationDatapointsTableContent {...props} />
     </DataTableStateProvider>
   );
diff --git a/frontend/components/evaluation/evaluation.tsx b/frontend/components/evaluation/evaluation.tsx
index 5c8d08f4a..58dc7d740 100644
--- a/frontend/components/evaluation/evaluation.tsx
+++ b/frontend/components/evaluation/evaluation.tsx
@@ -11,7 +11,7 @@ import CompareChart from "@/components/evaluation/compare-chart";
 import EvaluationDatapointsTable from "@/components/evaluation/evaluation-datapoints-table";
 import EvaluationHeader from "@/components/evaluation/evaluation-header";
 import ScoreCard from "@/components/evaluation/score-card";
-import { useEvalStore } from "@/components/evaluation/store";
+import { buildEvalColumnDefs, useEvalStore } from "@/components/evaluation/store";
 import { useInfiniteScroll } from "@/components/ui/infinite-datatable/hooks";
 import { DataTableStateProvider } from "@/components/ui/infinite-datatable/model/datatable-store";
 import { Skeleton } from "@/components/ui/skeleton";
@@ -49,25 +49,20 @@ function EvaluationContent({ evaluations, evaluationId, evaluationName, initialT
     () => searchParams.get("datapointId") ?? undefined
   );
 
-  // Pagination state
   const pageSize = 50;
 
-  // Store
-  const rebuildColumns = useEvalStore((s) => s.rebuildColumns);
   const setIsComparison = useEvalStore((s) => s.setIsComparison);
   const setIsShared = useEvalStore((s) => s.setIsShared);
-  const columnDefs = useEvalStore((s) => s.columnDefs);
   const buildStatsParams = useEvalStore((s) => s.buildStatsParams);
   const buildFetchParams = useEvalStore((s) => s.buildFetchParams);
 
-  // Statistics URL (fetches all stats at once)
   const statsUrl = useMemo(() => {
     const base = `/api/projects/${params?.projectId}/evaluations/${evaluationId}/stats`;
-    const urlParams = buildStatsParams({ search, searchIn, filter, sortBy, sortDirection });
+    // Stats params don't need columns for the initial stats call (only for filter columns)
+    const urlParams = buildStatsParams({ search, searchIn, filter, sortBy, sortDirection }, []);
     const qs = urlParams.toString();
     return qs ? `${base}?${qs}` : base;
-    // columnDefs used internally in buildStatParams via store
-  }, [params?.projectId, evaluationId, search, searchIn, filter, sortBy, sortDirection, buildStatsParams, columnDefs]);
+  }, [params?.projectId, evaluationId, search, searchIn, filter, sortBy, sortDirection, buildStatsParams]);
 
   const { data: statsData, isLoading: isStatsLoading } = useSWR<{
     evaluation: EvaluationType;
@@ -76,15 +71,13 @@ function EvaluationContent({ evaluations, evaluationId, evaluationName, initialT
     scores: string[];
   }>(statsUrl, swrFetcher);
 
-  // Target statistics URL (if comparing)
   const targetStatsUrl = useMemo(() => {
     if (!targetId) return null;
     const base = `/api/projects/${params?.projectId}/evaluations/${targetId}/stats`;
-    const urlParams = buildStatsParams({ search, searchIn, filter, sortBy, sortDirection });
+    const urlParams = buildStatsParams({ search, searchIn, filter, sortBy, sortDirection }, []);
     const qs = urlParams.toString();
     return qs ? `${base}?${qs}` : base;
-    // columnDefs used internally in buildStatParams via store
-  }, [params.projectId, targetId, search, searchIn, filter, sortBy, sortDirection, buildStatsParams, columnDefs]);
+  }, [params.projectId, targetId, search, searchIn, filter, sortBy, sortDirection, buildStatsParams]);
 
   const { data: targetStatsData } = useSWR<{
     evaluation: EvaluationType;
@@ -95,28 +88,17 @@ function EvaluationContent({ evaluations, evaluationId, evaluationName, initialT
 
   const scores = useMemo(() => statsData?.scores ?? [], [statsData?.scores]);
 
-  // Sync comparison state from URL
+  const allColumnDefs = useMemo(() => buildEvalColumnDefs(scores), [scores]);
+
   useEffect(() => {
     setIsComparison(!!targetId);
   }, [targetId, setIsComparison]);
 
-  // Reset shared state — authenticated evals are not shared.
   useEffect(() => {
     setIsShared(false);
   }, [setIsShared]);
 
-  const customColumns = useEvalStore((s) => s.customColumns);
-
-  // Rebuild column defs when scores or custom columns change.
-  // This must run before useInfiniteScroll's effect (declaration order).
-  useEffect(() => {
-    rebuildColumns(scores);
-  }, [scores, customColumns, rebuildColumns]);
-
-  // SQL strings from column defs — only changes when columns structurally change.
-  // useInfiniteScroll uses JSON.stringify on deps, so identical SQL strings
-  // produce the same string → no spurious re-fetch.
-  const columnSqls = useMemo(() => columnDefs.map((c) => c.meta?.sql).filter(Boolean), [columnDefs]);
+  const columnSqls = useMemo(() => allColumnDefs.map((c) => c.meta?.sql).filter(Boolean), [allColumnDefs]);
 
   const onClose = useCallback(() => {
     setTraceId(undefined);
@@ -126,19 +108,21 @@ function EvaluationContent({ evaluations, evaluationId, evaluationName, initialT
     push(`${pathName}?${params}`);
   }, [searchParams, pathName, push]);
 
-  // Fetch function for datapoints — single query handles comparison via targetId
   const fetchDatapoints = useCallback(
     async (pageNumber: number) => {
-      const urlParams = buildFetchParams({
-        search,
-        searchIn,
-        filter,
-        sortBy,
-        sortDirection,
-        targetId,
-        pageNumber,
-        pageSize,
-      });
+      const urlParams = buildFetchParams(
+        {
+          search,
+          searchIn,
+          filter,
+          sortBy,
+          sortDirection,
+          targetId,
+          pageNumber,
+          pageSize,
+        },
+        allColumnDefs
+      );
 
       const url = `/api/projects/${params?.projectId}/evaluations/${evaluationId}?${urlParams.toString()}`;
       const response = await fetch(url);
@@ -160,10 +144,10 @@ function EvaluationContent({ evaluations, evaluationId, evaluationName, initialT
       sortDirection,
       targetId,
       buildFetchParams,
+      allColumnDefs,
     ]
   );
 
-  // Use infinite scroll hook — data is now EvalRow (Record<string, unknown>)
   const {
     data: allDatapoints,
     hasMore: hasMorePages,
diff --git a/frontend/components/evaluation/store.ts b/frontend/components/evaluation/store.ts
index eecfb77ff..45c63b6d7 100644
--- a/frontend/components/evaluation/store.ts
+++ b/frontend/components/evaluation/store.ts
@@ -3,7 +3,7 @@ import { create } from "zustand";
 import { persist } from "zustand/middleware";
 
 import { type ScoreRanges } from "@/components/evaluation/utils";
-import { type CustomColumn } from "@/components/ui/columns-menu";
+import { type CustomColumn } from "@/components/ui/infinite-datatable/model/datatable-store";
 import { type EvalQueryColumn } from "@/lib/actions/evaluation/query-builder";
 import { type EvalRow } from "@/lib/evaluation/types";
 
@@ -31,43 +31,53 @@ function toColumnsPayload(columnDefs: ColumnDef<EvalRow>[]): EvalQueryColumn[] {
     }));
 }
 
+export function buildEvalCustomColumnDef(cc: CustomColumn): ColumnDef<EvalRow> {
+  return {
+    id: `custom:${cc.name}`,
+    accessorFn: (row) => row[`custom:${cc.name}`],
+    cell: DataCell,
+    header: cc.name,
+    enableSorting: true,
+    meta: {
+      sql: cc.sql,
+      dataType: cc.dataType,
+      filterable: true,
+      comparable: true,
+      isCustom: true,
+    },
+  };
+}
+
+/** Build the full static + score column defs for a given set of score names. */
+export function buildEvalColumnDefs(scoreNames: string[]): ColumnDef<EvalRow>[] {
+  const scoreCols = scoreNames.map((name) => createScoreColumnDef(name));
+  return [...STATIC_COLUMNS, ...scoreCols];
+}
+
 interface EvalStoreState {
-  // Data
   scoreRanges: ScoreRanges;
   heatmapEnabled: boolean;
   isComparison: boolean;
   isShared: boolean;
-  columnDefs: ColumnDef<EvalRow>[];
-  customColumns: CustomColumn[];
-  lastScoreNames: string[];
 
-  // Actions
   setScoreRanges: (ranges: ScoreRanges) => void;
   setHeatmapEnabled: (enabled: boolean) => void;
   setIsComparison: (value: boolean) => void;
   setIsShared: (value: boolean) => void;
-  rebuildColumns: (scoreNames: string[]) => void;
-  addCustomColumn: (column: CustomColumn) => void;
-  updateCustomColumn: (oldName: string, column: CustomColumn) => void;
-  removeCustomColumn: (name: string) => void;
-  buildStatsParams: (raw: RawUrlParams) => URLSearchParams;
-  buildFetchParams: (raw: RawUrlParams & { pageNumber: number; pageSize: number }) => URLSearchParams;
+  buildStatsParams: (raw: RawUrlParams, allColumnDefs: ColumnDef<EvalRow>[]) => URLSearchParams;
+  buildFetchParams: (
+    raw: RawUrlParams & { pageNumber: number; pageSize: number },
+    allColumnDefs: ColumnDef<EvalRow>[]
+  ) => URLSearchParams;
 }
 
-/** Selector: visible columns, excluding output in comparison mode */
-export const selectVisibleColumns = (s: EvalStoreState): ColumnDef<EvalRow>[] =>
-  s.columnDefs.filter((c) => !c.meta?.hidden && !(s.isComparison && c.id === "output"));
-
 export const useEvalStore = create<EvalStoreState>()(
   persist(
-    (set, get) => ({
+    (set) => ({
       scoreRanges: {},
       heatmapEnabled: false,
       isComparison: false,
       isShared: false,
-      columnDefs: [],
-      customColumns: [],
-      lastScoreNames: [],
 
       setScoreRanges: (ranges) => set({ scoreRanges: ranges }),
 
@@ -77,58 +87,12 @@ export const useEvalStore = create<EvalStoreState>()(
 
       setIsShared: (value) => set({ isShared: value }),
 
-      rebuildColumns: (scoreNames) => {
-        const { customColumns, isShared } = get();
-        const scoreCols = scoreNames.map((name) => createScoreColumnDef(name));
-
-        // Don't include custom columns in shared evaluations to prevent
-        // browser-persisted custom SQL from being executed
-        const customCols: ColumnDef<EvalRow>[] = isShared
-          ? []
-          : customColumns.map((cc) => ({
-              id: `custom:${cc.name}`,
-              accessorFn: (row) => row[`custom:${cc.name}`],
-              cell: DataCell,
-              header: cc.name,
-              enableSorting: true,
-              meta: {
-                sql: cc.sql,
-                dataType: cc.dataType,
-                filterable: true,
-                comparable: true,
-                isCustom: true,
-              },
-            }));
-        set({ columnDefs: [...STATIC_COLUMNS, ...scoreCols, ...customCols], lastScoreNames: scoreNames });
-      },
-
-      addCustomColumn: (column) => {
-        const { customColumns } = get();
-        if (customColumns.some((cc) => cc.name === column.name)) return;
-        set({ customColumns: [...customColumns, column] });
-        get().rebuildColumns(get().lastScoreNames);
-      },
-
-      updateCustomColumn: (oldName, column) => {
-        const { customColumns } = get();
-        set({ customColumns: customColumns.map((cc) => (cc.name === oldName ? column : cc)) });
-        get().rebuildColumns(get().lastScoreNames);
-      },
-
-      removeCustomColumn: (name) => {
-        const { customColumns } = get();
-        set({ customColumns: customColumns.filter((cc) => cc.name !== name) });
-        get().rebuildColumns(get().lastScoreNames);
-      },
-
-      buildStatsParams: (raw) => {
-        const { columnDefs } = get();
+      buildStatsParams: (raw, allColumnDefs) => {
         const urlParams = new URLSearchParams();
         if (raw.search) urlParams.set("search", raw.search);
         raw.searchIn.forEach((v) => urlParams.append("searchIn", v));
         raw.filter.forEach((f) => urlParams.append("filter", f));
 
-        // Only send columns referenced by active filters (optimization for URL length)
         const parsedFilters = raw.filter
           .map((f) => {
             try {
@@ -140,15 +104,14 @@ export const useEvalStore = create<EvalStoreState>()(
           .filter(Boolean) as { column: string }[];
         if (parsedFilters.length > 0) {
           const filterIds = new Set(parsedFilters.map((f) => f.column));
-          const filterCols = columnDefs.filter((c) => filterIds.has(c.id!));
+          const filterCols = allColumnDefs.filter((c) => filterIds.has(c.id!));
           urlParams.set("columns", JSON.stringify(toColumnsPayload(filterCols)));
         }
 
         return urlParams;
       },
 
-      buildFetchParams: (raw) => {
-        const { columnDefs } = get();
+      buildFetchParams: (raw, allColumnDefs) => {
         const urlParams = new URLSearchParams();
         urlParams.set("pageNumber", raw.pageNumber.toString());
         urlParams.set("pageSize", raw.pageSize.toString());
@@ -156,13 +119,11 @@ export const useEvalStore = create<EvalStoreState>()(
         raw.searchIn.forEach((v) => urlParams.append("searchIn", v));
         raw.filter.forEach((f) => urlParams.append("filter", f));
 
-        // Full columns payload derived directly from column defs
-        urlParams.set("columns", JSON.stringify(toColumnsPayload(columnDefs)));
+        urlParams.set("columns", JSON.stringify(toColumnsPayload(allColumnDefs)));
 
-        // Sort — resolve SQL from column meta
         if (raw.sortBy) {
           urlParams.set("sortBy", raw.sortBy);
-          const col = columnDefs.find((c) => c.id === raw.sortBy);
+          const col = allColumnDefs.find((c) => c.id === raw.sortBy);
           if (col?.meta?.sql) urlParams.set("sortSql", col.meta.sql);
         }
         if (raw.sortDirection) urlParams.set("sortDirection", raw.sortDirection);
@@ -173,7 +134,7 @@ export const useEvalStore = create<EvalStoreState>()(
     }),
     {
       name: "evaluation-heatmap-enabled",
-      partialize: (state) => ({ heatmapEnabled: state.heatmapEnabled, customColumns: state.customColumns }),
+      partialize: (state) => ({ heatmapEnabled: state.heatmapEnabled }),
     }
   )
 );
diff --git a/frontend/components/evaluations/evaluations-groups-bar.tsx b/frontend/components/evaluations/evaluations-groups-bar.tsx
index 5015051ba..10c1d366e 100644
--- a/frontend/components/evaluations/evaluations-groups-bar.tsx
+++ b/frontend/components/evaluations/evaluations-groups-bar.tsx
@@ -29,7 +29,7 @@ const columns: ColumnDef<EvaluationGroup>[] = [
 
 export default function EvaluationsGroupsBar() {
   return (
-    <DataTableStateProvider storageKey="evaluations-groups-bar" columns={columns}>
+    <DataTableStateProvider storageKey="evaluations-groups-bar" columnDefs={columns}>
       <EvaluationsGroupsBarContent />
     </DataTableStateProvider>
   );
@@ -59,7 +59,6 @@ function EvaluationsGroupsBarContent() {
       <div className="flex overflow-hidden">
         <InfiniteDataTable<EvaluationGroup>
           className="w-full"
-          columns={columns}
           data={groups || []}
           getRowId={(row) => row.groupId}
           focusedRowId={groupId}
diff --git a/frontend/components/evaluations/evaluations.tsx b/frontend/components/evaluations/evaluations.tsx
index 26c196efe..301fbe66e 100644
--- a/frontend/components/evaluations/evaluations.tsx
+++ b/frontend/components/evaluations/evaluations.tsx
@@ -88,7 +88,7 @@ export default function Evaluations() {
   return (
     <DataTableStateProvider
       storageKey="evaluations-table"
-      columns={columns}
+      columnDefs={columns}
       enableRowSelection
       lockedColumns={["__row_selection"]}
     >
@@ -234,7 +234,6 @@ function EvaluationsContent() {
               <InfiniteDataTable<Evaluation>
                 className="w-full"
                 enableRowSelection
-                columns={columns}
                 data={evaluations}
                 getRowId={(evaluation) => evaluation.id}
                 getRowHref={(row) => `/project/${params?.projectId}/evaluations/${row.original.id}`}
diff --git a/frontend/components/evaluators/evaluators-table.tsx b/frontend/components/evaluators/evaluators-table.tsx
index dfa10c40d..01bcfa0a7 100644
--- a/frontend/components/evaluators/evaluators-table.tsx
+++ b/frontend/components/evaluators/evaluators-table.tsx
@@ -12,8 +12,6 @@ import { useToast } from "@/lib/hooks/use-toast";
 import { type PaginatedResponse } from "@/lib/types";
 import { swrFetcher } from "@/lib/utils";
 
-import { columns } from "./lib/consts";
-
 interface EvaluatorsTableProps {
   projectId: string;
   onRowClick: (row: Row<Evaluator>) => void;
@@ -63,7 +61,6 @@ export default function EvaluatorsTable({ projectId, onRowClick }: EvaluatorsTab
 
   return (
     <InfiniteDataTable
-      columns={columns}
       data={data?.items ?? []}
       hasMore={false}
       isFetching={false}
diff --git a/frontend/components/evaluators/evaluators.tsx b/frontend/components/evaluators/evaluators.tsx
index 76451e7e9..4bccdad75 100644
--- a/frontend/components/evaluators/evaluators.tsx
+++ b/frontend/components/evaluators/evaluators.tsx
@@ -72,7 +72,7 @@ const Evaluators = () => {
         </ManageEvaluatorSheet>
         <DataTableStateProvider
           storageKey="evaluators-table"
-          columns={columns}
+          columnDefs={columns}
           enableRowSelection
           lockedColumns={["__row_selection"]}
         >
diff --git a/frontend/components/playground/playground-history-table.tsx b/frontend/components/playground/playground-history-table.tsx
index 538bf488e..c6d0c3bcc 100644
--- a/frontend/components/playground/playground-history-table.tsx
+++ b/frontend/components/playground/playground-history-table.tsx
@@ -136,7 +136,7 @@ const FETCH_SIZE = 50;
 
 export default function PlaygroundHistoryTable(props: PlaygroundHistoryTableProps) {
   return (
-    <DataTableStateProvider storageKey="playground-history-table" uniqueKey="id" columns={columns}>
+    <DataTableStateProvider storageKey="playground-history-table" uniqueKey="id" columnDefs={columns}>
       <PlaygroundHistoryTableContent {...props} />
     </DataTableStateProvider>
   );
@@ -215,7 +215,6 @@ function PlaygroundHistoryTableContent({ playgroundId, onRowClick, onTraceSelect
   return (
     <InfiniteDataTable<Trace>
       className="w-full"
-      columns={columns}
       data={traces}
       getRowId={(trace) => trace.id}
       onRowClick={handleRowClick}
diff --git a/frontend/components/playgrounds/playgrounds.tsx b/frontend/components/playgrounds/playgrounds.tsx
index d9167dd21..ba49855cf 100644
--- a/frontend/components/playgrounds/playgrounds.tsx
+++ b/frontend/components/playgrounds/playgrounds.tsx
@@ -194,7 +194,6 @@ const PlaygroundsContent = () => {
           enableRowSelection={true}
           getRowHref={(row) => `/project/${projectId}/playgrounds/${row.original.id}`}
           getRowId={(row) => row.id}
-          columns={columns}
           data={playgrounds ?? []}
           hasMore={hasMore}
           isFetching={isFetching}
@@ -251,7 +250,7 @@ export default function Playgrounds() {
   return (
     <DataTableStateProvider
       storageKey="playgrounds-table"
-      columns={columns}
+      columnDefs={columns}
       enableRowSelection
       lockedColumns={["__row_selection"]}
     >
diff --git a/frontend/components/queues/queues.tsx b/frontend/components/queues/queues.tsx
index a438f4c4b..0234fbb57 100644
--- a/frontend/components/queues/queues.tsx
+++ b/frontend/components/queues/queues.tsx
@@ -209,7 +209,6 @@ const QueuesContent = () => {
           enableRowSelection={true}
           getRowHref={(row) => `/project/${projectId}/labeling-queues/${row.original.id}`}
           getRowId={(row: LabelingQueue) => row.id}
-          columns={columns}
           data={queues ?? []}
           hasMore={hasMore}
           isFetching={isFetching}
@@ -266,7 +265,7 @@ export default function Queues() {
   return (
     <DataTableStateProvider
       storageKey="queues-table"
-      columns={columns}
+      columnDefs={columns}
       enableRowSelection
       lockedColumns={["__row_selection"]}
     >
diff --git a/frontend/components/shared/evaluation/shared-evaluation.tsx b/frontend/components/shared/evaluation/shared-evaluation.tsx
index 4c9e909c8..698d7f7a1 100644
--- a/frontend/components/shared/evaluation/shared-evaluation.tsx
+++ b/frontend/components/shared/evaluation/shared-evaluation.tsx
@@ -12,7 +12,7 @@ import fullLogo from "@/assets/logo/logo.svg";
 import Chart from "@/components/evaluation/chart";
 import EvaluationDatapointsTable from "@/components/evaluation/evaluation-datapoints-table";
 import ScoreCard from "@/components/evaluation/score-card";
-import { useEvalStore } from "@/components/evaluation/store";
+import { buildEvalColumnDefs, useEvalStore } from "@/components/evaluation/store";
 import SharedEvalTraceView from "@/components/shared/evaluation/shared-eval-trace-view";
 import { useInfiniteScroll } from "@/components/ui/infinite-datatable/hooks";
 import { DataTableStateProvider } from "@/components/ui/infinite-datatable/model/datatable-store";
@@ -50,15 +50,11 @@ function SharedEvaluationContent({ evaluationId, evaluationName }: SharedEvaluat
 
   const pageSize = 50;
 
-  // Store actions
-  const rebuildColumns = useEvalStore((s) => s.rebuildColumns);
-  const columnDefs = useEvalStore((s) => s.columnDefs);
   const buildStatsParams = useEvalStore((s) => s.buildStatsParams);
   const buildFetchParams = useEvalStore((s) => s.buildFetchParams);
   const setIsComparison = useEvalStore((s) => s.setIsComparison);
   const setIsShared = useEvalStore((s) => s.setIsShared);
 
-  // Shared evals never have comparison mode — reset in case it persists from a previous page.
   useEffect(() => {
     setIsComparison(false);
     setIsShared(true);
@@ -66,11 +62,10 @@ function SharedEvaluationContent({ evaluationId, evaluationName }: SharedEvaluat
 
   const statsUrl = useMemo(() => {
     const base = `/api/shared/evals/${evaluationId}/stats`;
-    const urlParams = buildStatsParams({ search, searchIn, filter, sortBy, sortDirection });
+    const urlParams = buildStatsParams({ search, searchIn, filter, sortBy, sortDirection }, []);
     const qs = urlParams.toString();
     return qs ? `${base}?${qs}` : base;
-    // columnDefs used internally in buildStatParams via store
-  }, [evaluationId, search, searchIn, filter, sortBy, sortDirection, buildStatsParams, columnDefs]);
+  }, [evaluationId, search, searchIn, filter, sortBy, sortDirection, buildStatsParams]);
 
   const { data: statsData, isLoading: isStatsLoading } = useSWR<{
     evaluation: Evaluation;
@@ -81,13 +76,8 @@ function SharedEvaluationContent({ evaluationId, evaluationName }: SharedEvaluat
 
   const scores = useMemo(() => statsData?.scores ?? [], [statsData?.scores]);
 
-  // Rebuild column defs when scores change.
-  useEffect(() => {
-    rebuildColumns(scores);
-  }, [scores, rebuildColumns]);
-
-  // SQL strings from column defs — only changes when columns structurally change.
-  const columnSqls = useMemo(() => columnDefs.map((c) => c.meta?.sql).filter(Boolean), [columnDefs]);
+  const allColumnDefs = useMemo(() => buildEvalColumnDefs(scores), [scores]);
+  const columnSqls = useMemo(() => allColumnDefs.map((c) => c.meta?.sql).filter(Boolean), [allColumnDefs]);
 
   const onClose = useCallback(() => {
     setTraceId(undefined);
@@ -101,15 +91,18 @@ function SharedEvaluationContent({ evaluationId, evaluationName }: SharedEvaluat
 
   const fetchDatapoints = useCallback(
     async (pageNumber: number) => {
-      const urlParams = buildFetchParams({
-        search,
-        searchIn,
-        filter,
-        sortBy,
-        sortDirection,
-        pageNumber,
-        pageSize,
-      });
+      const urlParams = buildFetchParams(
+        {
+          search,
+          searchIn,
+          filter,
+          sortBy,
+          sortDirection,
+          pageNumber,
+          pageSize,
+        },
+        allColumnDefs
+      );
 
       const url = `/api/shared/evals/${evaluationId}?${urlParams.toString()}`;
       const response = await fetch(url);
@@ -120,7 +113,7 @@ function SharedEvaluationContent({ evaluationId, evaluationName }: SharedEvaluat
 
       return { items: data.results, count: 0 };
     },
-    [search, searchIn, filter, evaluationId, pageSize, sortBy, sortDirection, buildFetchParams]
+    [search, searchIn, filter, evaluationId, pageSize, sortBy, sortDirection, buildFetchParams, allColumnDefs]
   );
 
   const {
diff --git a/frontend/components/signal/create-signal-job/index.tsx b/frontend/components/signal/create-signal-job/index.tsx
index e4ecd0379..f286deb5b 100644
--- a/frontend/components/signal/create-signal-job/index.tsx
+++ b/frontend/components/signal/create-signal-job/index.tsx
@@ -10,7 +10,7 @@ import ConfirmSignalJobDialog from "@/components/signal/create-signal-job/confir
 import SelectionBanner from "@/components/signal/create-signal-job/selection-banner.tsx";
 import { useSignalStoreContext } from "@/components/signal/store.tsx";
 import TraceView from "@/components/traces/trace-view";
-import { columns, filters as tableFilters } from "@/components/traces/traces-table/columns.tsx";
+import { filters as tableFilters, STATIC_COLUMNS as columns } from "@/components/traces/traces-table/columns.tsx";
 import { Button } from "@/components/ui/button.tsx";
 import { ColumnsMenu } from "@/components/ui/columns-menu";
 import DateRangeFilter from "@/components/ui/date-range-filter";
@@ -295,7 +295,6 @@ const CreateSignalJobContent = () => {
       <div className="flex flex-1 overflow-hidden px-4 pb-4">
         <InfiniteDataTable<TraceRow>
           className="w-full"
-          columns={columns}
           data={traces}
           enableRowSelection
           getRowId={(trace) => trace.id}
@@ -383,7 +382,7 @@ export default function CreateSignalJob({ traceId }: { traceId?: string }) {
   }, []);
 
   return (
-    <DataTableStateProvider columns={columns} enableRowSelection lockedColumns={["__row_selection", "status"]}>
+    <DataTableStateProvider columnDefs={columns} enableRowSelection lockedColumns={["__row_selection", "status"]}>
       <CreateSignalJobContent />
     </DataTableStateProvider>
   );
diff --git a/frontend/components/signal/events-table/index.tsx b/frontend/components/signal/events-table/index.tsx
index 3b1ad5c4b..98cb7706d 100644
--- a/frontend/components/signal/events-table/index.tsx
+++ b/frontend/components/signal/events-table/index.tsx
@@ -72,7 +72,7 @@ function PureEventsTable() {
   const filterRaw = searchParams.getAll("filter");
   const filter = useMemo(() => filterRaw, [JSON.stringify(filterRaw)]);
 
-  const { columns, filters } = useMemo(() => buildEventsColumns(signal.schemaFields), [signal.schemaFields]);
+  const { filters } = useMemo(() => buildEventsColumns(signal.schemaFields), [signal.schemaFields]);
 
   const setTraceId = useSignalStoreContext((state) => state.setTraceId);
   const setSelectedEvent = useSignalStoreContext((state) => state.setSelectedEvent);
@@ -200,7 +200,6 @@ function PureEventsTable() {
     <div className="flex flex-1 overflow-hidden px-4 pb-4">
       <InfiniteDataTable<EventRow>
         className="w-full"
-        columns={columns}
         data={events}
         onRowClick={handleRowClick}
         getRowId={(row: EventRow) => row.id}
@@ -233,7 +232,7 @@ export default function EventsTable() {
   const { columns } = useMemo(() => buildEventsColumns(signal.schemaFields), [signal.schemaFields]);
 
   return (
-    <DataTableStateProvider storageKey={`events-table-${signal.id}`} uniqueKey="id" columns={columns}>
+    <DataTableStateProvider storageKey={`events-table-${signal.id}`} uniqueKey="id" columnDefs={columns}>
       <PureEventsTable />
     </DataTableStateProvider>
   );
diff --git a/frontend/components/signal/jobs-table/index.tsx b/frontend/components/signal/jobs-table/index.tsx
index 8734a3d58..66860e62e 100644
--- a/frontend/components/signal/jobs-table/index.tsx
+++ b/frontend/components/signal/jobs-table/index.tsx
@@ -97,7 +97,6 @@ const JobsTableContent = () => {
       </Link>
       <InfiniteDataTable<SignalJobRow>
         className="w-full"
-        columns={signalJobsColumns}
         data={jobs}
         getRowId={(job) => job.id}
         hasMore={false}
@@ -120,7 +119,7 @@ const JobsTableContent = () => {
 
 export default function SignalJobsTable() {
   return (
-    <DataTableStateProvider columns={signalJobsColumns} lockedColumns={["id"]}>
+    <DataTableStateProvider columnDefs={signalJobsColumns} lockedColumns={["id"]}>
       <JobsTableContent />
     </DataTableStateProvider>
   );
diff --git a/frontend/components/signal/runs-table/index.tsx b/frontend/components/signal/runs-table/index.tsx
index 145ed9a1e..8d7479b16 100644
--- a/frontend/components/signal/runs-table/index.tsx
+++ b/frontend/components/signal/runs-table/index.tsx
@@ -1,8 +1,7 @@
 "use client";
 
-import { type Row } from "@tanstack/react-table";
 import { isEqual } from "lodash";
-import { useParams, useRouter } from "next/navigation";
+import { useParams } from "next/navigation";
 import React, { useCallback, useState } from "react";
 
 import { useSignalStoreContext } from "@/components/signal/store";
@@ -15,7 +14,6 @@ import { DataTableStateProvider } from "@/components/ui/infinite-datatable/model
 import FilterPopover, { FilterList } from "@/components/ui/infinite-datatable/ui/datatable-filter/ui";
 import { TableCell, TableRow } from "@/components/ui/table.tsx";
 import { type Filter } from "@/lib/actions/common/filters";
-import { Operator } from "@/lib/actions/common/operators";
 import { type SignalRunRow } from "@/lib/actions/signal-runs";
 import { useToast } from "@/lib/hooks/use-toast";
 
@@ -55,13 +53,11 @@ const getEmptyRow = ({
 
 function RunsTableContent() {
   const { toast } = useToast();
-  const router = useRouter();
   const params = useParams<{ projectId: string; id: string }>();
-  const { signal, runsFilters, setRunsFilters, setJobsFilters } = useSignalStoreContext((state) => ({
+  const { signal, runsFilters, setRunsFilters } = useSignalStoreContext((state) => ({
     signal: state.signal,
     runsFilters: state.runsFilters,
     setRunsFilters: state.setRunsFilters,
-    setJobsFilters: state.setJobsFilters,
   }));
 
   const filter = runsFilters;
@@ -89,12 +85,6 @@ function RunsTableContent() {
     [setRunsFilters]
   );
 
-  const onJobNav = (row: Row<SignalRunRow>) => {
-    router.push(`/project/${params.projectId}/signals/${params.id}?tab=jobs`);
-    setJobsFilters([{ column: "job_id", operator: Operator.Eq, value: row.original.jobId }]);
-  };
-
-  const columns = getSignalRunsColumns({ onJobNav });
   const fetchRuns = useCallback(
     async (pageNumber: number) => {
       try {
@@ -153,7 +143,6 @@ function RunsTableContent() {
     <div className="flex flex-col gap-2 flex-1 overflow-hidden">
       <InfiniteDataTable<SignalRunRow>
         className="w-full"
-        columns={columns}
         data={runs}
         getRowId={(row: SignalRunRow) => row.runId}
         hasMore={hasMore}
@@ -178,7 +167,7 @@ export default function SignalRunsTable() {
     <DataTableStateProvider
       storageKey="signal-runs-table"
       uniqueKey="runId"
-      columns={getSignalRunsColumns({ onJobNav: () => {} })}
+      columnDefs={getSignalRunsColumns({ onJobNav: () => {} })}
     >
       <RunsTableContent />
     </DataTableStateProvider>
diff --git a/frontend/components/signal/triggers-table/index.tsx b/frontend/components/signal/triggers-table/index.tsx
index c621284d7..fd0a28207 100644
--- a/frontend/components/signal/triggers-table/index.tsx
+++ b/frontend/components/signal/triggers-table/index.tsx
@@ -66,8 +66,6 @@ function TriggersTableContent() {
 
   const triggers: TriggerRow[] = data?.items || [];
 
-  const columns = getTriggersTableColumns();
-
   const handleAddFilter = useCallback(
     (filter: Filter) => {
       setTriggersFilters((prev) => [...prev, filter]);
@@ -147,7 +145,6 @@ function TriggersTableContent() {
       </ManageTriggerDialog>
       <InfiniteDataTable<TriggerRow>
         className="w-full"
-        columns={columns}
         data={triggers}
         getRowId={(trigger) => trigger.id}
         hasMore={false}
@@ -179,7 +176,11 @@ function TriggersTableContent() {
 
 export default function TriggersTable() {
   return (
-    <DataTableStateProvider columns={getTriggersTableColumns()} enableRowSelection lockedColumns={["__row_selection"]}>
+    <DataTableStateProvider
+      columnDefs={getTriggersTableColumns()}
+      enableRowSelection
+      lockedColumns={["__row_selection"]}
+    >
       <TriggersTableContent />
     </DataTableStateProvider>
   );
diff --git a/frontend/components/sql/editor-panel.tsx b/frontend/components/sql/editor-panel.tsx
index 9b5a57375..e65909db6 100644
--- a/frontend/components/sql/editor-panel.tsx
+++ b/frontend/components/sql/editor-panel.tsx
@@ -265,10 +265,9 @@ export default function EditorPanel() {
             <div className="flex overflow-hidden h-full">
               {renderContent({
                 success: (
-                  <DataTableStateProvider columns={columns}>
+                  <DataTableStateProvider columnDefs={columns}>
                     <InfiniteDataTable
                       className="w-full"
-                      columns={columns}
                       data={results || []}
                       hasMore={false}
                       isFetching={false}
diff --git a/frontend/components/traces/sessions-table/index.tsx b/frontend/components/traces/sessions-table/index.tsx
index e5dafc156..4df766c64 100644
--- a/frontend/components/traces/sessions-table/index.tsx
+++ b/frontend/components/traces/sessions-table/index.tsx
@@ -23,7 +23,7 @@ const FETCH_SIZE = 50;
 
 export default function SessionsTable() {
   return (
-    <DataTableStateProvider storageKey="sessions-table" uniqueKey="sessionId" columns={columns}>
+    <DataTableStateProvider storageKey="sessions-table" uniqueKey="sessionId" columnDefs={columns}>
       <SessionsTableContent />
     </DataTableStateProvider>
   );
@@ -206,7 +206,6 @@ function SessionsTableContent() {
     <div className="flex overflow-hidden px-4 pb-6">
       <InfiniteDataTable<SessionRow>
         className="w-full"
-        columns={columns}
         data={sessions}
         getRowId={(session) => get(session, ["id"], session.sessionId)}
         onRowClick={handleRowClick}
diff --git a/frontend/components/traces/spans-table/index.tsx b/frontend/components/traces/spans-table/index.tsx
index 6439dd1b0..7ebb40dcc 100644
--- a/frontend/components/traces/spans-table/index.tsx
+++ b/frontend/components/traces/spans-table/index.tsx
@@ -22,7 +22,7 @@ const FETCH_SIZE = 50;
 
 export default function SpansTable() {
   return (
-    <DataTableStateProvider storageKey="spans-table" uniqueKey="spanId" columns={columns} lockedColumns={["status"]}>
+    <DataTableStateProvider storageKey="spans-table" uniqueKey="spanId" columnDefs={columns} lockedColumns={["status"]}>
       <SpansTableContent />
     </DataTableStateProvider>
   );
@@ -146,7 +146,6 @@ function SpansTableContent() {
     <div className="flex flex-1 overflow-hidden px-4 pb-6">
       <InfiniteDataTable<SpanRow>
         className="w-full"
-        columns={columns}
         data={spans}
         getRowId={(span) => span.spanId}
         onRowClick={handleRowClick}
diff --git a/frontend/components/traces/trace-picker/index.tsx b/frontend/components/traces/trace-picker/index.tsx
index 2da7cfec7..46b4a332c 100644
--- a/frontend/components/traces/trace-picker/index.tsx
+++ b/frontend/components/traces/trace-picker/index.tsx
@@ -8,7 +8,7 @@ import TracePickerContent, { type TracePickerProps } from "./trace-picker-conten
 export type { TracePickerProps };
 
 const TracePicker = (props: TracePickerProps) => (
-  <DataTableStateProvider columns={tracePickerColumns} pageSize={FETCH_SIZE} lockedColumns={["status"]}>
+  <DataTableStateProvider columnDefs={tracePickerColumns} pageSize={FETCH_SIZE} lockedColumns={["status"]}>
     <TracePickerContent {...props} />
   </DataTableStateProvider>
 );
diff --git a/frontend/components/traces/trace-picker/trace-picker-content.tsx b/frontend/components/traces/trace-picker/trace-picker-content.tsx
index 475262024..a327a7050 100644
--- a/frontend/components/traces/trace-picker/trace-picker-content.tsx
+++ b/frontend/components/traces/trace-picker/trace-picker-content.tsx
@@ -13,7 +13,7 @@ import RefreshButton from "@/components/ui/infinite-datatable/ui/refresh-button"
 import type { Filter } from "@/lib/actions/common/filters";
 import type { TraceRow } from "@/lib/traces/types";
 
-import { FETCH_SIZE, tracePickerColumns } from "./columns";
+import { FETCH_SIZE } from "./columns";
 
 export interface TracePickerProps {
   onTraceSelect: (trace: TraceRow) => void;
@@ -108,7 +108,6 @@ const TracePickerContent = ({
 
       <InfiniteDataTable<TraceRow>
         className="w-full flex-1"
-        columns={tracePickerColumns}
         data={traces}
         getRowId={(t) => t.id}
         onRowClick={handleRowClick}
diff --git a/frontend/components/traces/traces-table/columns.tsx b/frontend/components/traces/traces-table/columns.tsx
index 5adbe4664..bd48278a0 100644
--- a/frontend/components/traces/traces-table/columns.tsx
+++ b/frontend/components/traces/traces-table/columns.tsx
@@ -257,9 +257,6 @@ export const STATIC_COLUMNS: ColumnDef<TraceRow, any>[] = [
   },
 ];
 
-/** @deprecated Use STATIC_COLUMNS and useTracesTableStore().columnDefs instead */
-export const columns = STATIC_COLUMNS;
-
 export const filters: ColumnFilter[] = [
   {
     name: "ID",
diff --git a/frontend/components/traces/traces-table/index.tsx b/frontend/components/traces/traces-table/index.tsx
index 0ae1c88a4..42b3cc5ce 100644
--- a/frontend/components/traces/traces-table/index.tsx
+++ b/frontend/components/traces/traces-table/index.tsx
@@ -1,5 +1,5 @@
 "use client";
-import { type Row } from "@tanstack/react-table";
+import { type ColumnDef, type Row } from "@tanstack/react-table";
 import { isEmpty, map } from "lodash";
 import { useParams, usePathname, useRouter, useSearchParams } from "next/navigation";
 import { useCallback, useEffect, useMemo, useRef } from "react";
@@ -9,13 +9,17 @@ import AdvancedSearch from "@/components/common/advanced-search";
 import { useTraceViewNavigation } from "@/components/traces/trace-view/navigation-context";
 import TracesChart from "@/components/traces/traces-chart";
 import { useTracesStoreContext } from "@/components/traces/traces-store";
-import { filters } from "@/components/traces/traces-table/columns";
+import { filters, STATIC_COLUMNS } from "@/components/traces/traces-table/columns";
 import TracesColumnsMenu from "@/components/traces/traces-table/traces-columns-menu";
-import { useTracesTableStore } from "@/components/traces/traces-table/traces-table-store";
+import { buildCustomColumnDef, useTracesTableStore } from "@/components/traces/traces-table/traces-table-store";
 import DateRangeFilter from "@/components/ui/date-range-filter";
 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 {
+  DataTableStateProvider,
+  selectAllColumnDefs,
+  useDataTableStoreSelector,
+} from "@/components/ui/infinite-datatable/model/datatable-store";
 import DataTableFilter from "@/components/ui/infinite-datatable/ui/datatable-filter";
 import RefreshButton from "@/components/ui/infinite-datatable/ui/refresh-button.tsx";
 import { Switch } from "@/components/ui/switch";
@@ -28,10 +32,13 @@ const FETCH_SIZE = 50;
 const DEFAULT_TARGET_BARS = 48;
 
 export default function TracesTable() {
-  const columnDefs = useTracesTableStore((s) => s.columnDefs);
-
   return (
-    <DataTableStateProvider storageKey="traces-table" columns={columnDefs} lockedColumns={["status"]}>
+    <DataTableStateProvider
+      storageKey="traces-table"
+      columnDefs={STATIC_COLUMNS}
+      lockedColumns={["status"]}
+      buildCustomColumnDef={buildCustomColumnDef}
+    >
       <TracesTableContent />
     </DataTableStateProvider>
   );
@@ -77,18 +84,11 @@ function TracesTableContent() {
   const { setNavigationRefList } = useTraceViewNavigation();
   const isCurrentTimestampIncluded = !!pastHours || (!!endDate && new Date(endDate) >= new Date());
 
-  // Initialize column defs (rebuild when store hydrates custom columns)
-  const rebuildColumns = useTracesTableStore((s) => s.rebuildColumns);
-  const columnDefs = useTracesTableStore((s) => s.columnDefs);
   const buildFetchParams = useTracesTableStore((s) => s.buildFetchParams);
 
-  const customColumns = useTracesTableStore((s) => s.customColumns);
-
-  useEffect(() => {
-    rebuildColumns();
-  }, [customColumns, rebuildColumns]);
+  const allColumnDefs = useDataTableStoreSelector(selectAllColumnDefs) as ColumnDef<TraceRow>[];
 
-  const columnSqls = useMemo(() => columnDefs.map((c) => c.meta?.sql).filter(Boolean), [columnDefs]);
+  const columnSqls = useMemo(() => allColumnDefs.map((c) => c.meta?.sql).filter(Boolean), [allColumnDefs]);
 
   useEffect(() => {
     if (!chartContainerRef.current) return;
@@ -124,13 +124,16 @@ function TracesTableContent() {
   const fetchTraces = useCallback(
     async (pageNumber: number) => {
       try {
-        const urlParams = buildFetchParams({
-          pageNumber,
-          pageSize: FETCH_SIZE,
-          filter,
-          sortBy: sortBy ?? null,
-          sortDirection: sortDirection?.toUpperCase() as string | null,
-        });
+        const urlParams = buildFetchParams(
+          {
+            pageNumber,
+            pageSize: FETCH_SIZE,
+            filter,
+            sortBy: sortBy ?? null,
+            sortDirection: sortDirection?.toUpperCase() as string | null,
+          },
+          allColumnDefs
+        );
 
         if (pastHours != null) urlParams.set("pastHours", pastHours);
         if (startDate != null) urlParams.set("startDate", startDate);
@@ -174,6 +177,7 @@ function TracesTableContent() {
     },
     [
       buildFetchParams,
+      allColumnDefs,
       endDate,
       filter,
       pastHours,
@@ -234,15 +238,12 @@ function TracesTableContent() {
         const existingTraceIndex = currentTraces.findIndex((trace) => trace.id === traceData.id);
 
         if (existingTraceIndex !== -1) {
-          // Update existing trace
           const newTraces = [...currentTraces];
           newTraces[existingTraceIndex] = traceData;
           return newTraces;
         } else {
-          // New trace - insert at the beginning
           const newTraces = [traceData, ...currentTraces];
 
-          // Keep only the first FETCH_SIZE traces
           if (newTraces.length > FETCH_SIZE) {
             newTraces.splice(FETCH_SIZE);
           }
@@ -334,7 +335,6 @@ function TracesTableContent() {
     <div className="flex flex-1 overflow-hidden px-4 pb-4">
       <InfiniteDataTable<TraceRow>
         className="w-full"
-        columns={columnDefs.length > 0 ? columnDefs : []}
         data={traces}
         getRowId={(trace) => trace.id}
         onRowClick={handleRowClick}
diff --git a/frontend/components/traces/traces-table/traces-columns-menu/index.tsx b/frontend/components/traces/traces-table/traces-columns-menu/index.tsx
index 3f0bfa5ac..bacbe25ce 100644
--- a/frontend/components/traces/traces-table/traces-columns-menu/index.tsx
+++ b/frontend/components/traces/traces-table/traces-columns-menu/index.tsx
@@ -1,38 +1,28 @@
 "use client";
 
-import { useMemo } from "react";
+import { useCallback, useMemo } from "react";
 
-import { useTracesTableStore } from "@/components/traces/traces-table/traces-table-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";
 
 export default function TracesColumnsMenu() {
-  const addCustomColumn = useTracesTableStore((s) => s.addCustomColumn);
-  const updateCustomColumn = useTracesTableStore((s) => s.updateCustomColumn);
-  const removeCustomColumn = useTracesTableStore((s) => s.removeCustomColumn);
+  const store = useDataTableStore();
 
   const panelConfig = useMemo<CustomColumnPanelConfig>(
     () => ({
       schema: { tables: ["traces"] },
       generationMode: "trace-expression",
       buildTestQuery: (sql) => `SELECT ${sql} as \`test\` FROM traces LIMIT 1`,
-      getColumnDefs: () => useTracesTableStore.getState().columnDefs,
+      getColumnDefs: () => selectAllColumnDefs(store.getState()),
       namePlaceholder: "e.g. LLM span count",
       sqlPlaceholder: "e.g. total_tokens * total_cost",
       aiInputPlaceholder: "e.g. Calculate cost per token",
       sqlHint: "Expression is added as a column: SELECT <expr> FROM traces",
     }),
-    []
+    [store]
   );
 
-  const columnActions = useMemo<ColumnActions>(
-    () => ({
-      addCustomColumn,
-      updateCustomColumn,
-      removeCustomColumn,
-      getColumnDef: (columnId) => useTracesTableStore.getState().columnDefs.find((c) => c.id === columnId),
-    }),
-    [addCustomColumn, updateCustomColumn, removeCustomColumn]
-  );
+  const getColumnDefs = useCallback(() => selectAllColumnDefs(store.getState()), [store]);
 
-  return <ColumnsMenu panelConfig={panelConfig} columnActions={columnActions} />;
+  return <ColumnsMenu panelConfig={panelConfig} getColumnDefs={getColumnDefs} />;
 }
diff --git a/frontend/components/traces/traces-table/traces-table-store.ts b/frontend/components/traces/traces-table/traces-table-store.ts
index 6703d3c5b..7646e3d5b 100644
--- a/frontend/components/traces/traces-table/traces-table-store.ts
+++ b/frontend/components/traces/traces-table/traces-table-store.ts
@@ -1,12 +1,9 @@
 import { type ColumnDef } from "@tanstack/react-table";
 import { create } from "zustand";
-import { persist } from "zustand/middleware";
 
-import { type CustomColumn } from "@/components/ui/columns-menu";
+import { type CustomColumn } from "@/components/ui/infinite-datatable/model/datatable-store";
 import { type TraceRow } from "@/lib/traces/types";
 
-import { STATIC_COLUMNS } from "./columns";
-
 export interface TracesQueryColumn {
   id: string;
   sql: string;
@@ -31,85 +28,46 @@ function toColumnsPayload(columnDefs: ColumnDef<TraceRow>[]): TracesQueryColumn[
     }));
 }
 
-interface TracesTableStoreState {
-  columnDefs: ColumnDef<TraceRow>[];
-  customColumns: CustomColumn[];
-
-  rebuildColumns: () => void;
-  addCustomColumn: (column: CustomColumn) => void;
-  updateCustomColumn: (oldName: string, column: CustomColumn) => void;
-  removeCustomColumn: (name: string) => void;
-  buildFetchParams: (raw: RawUrlParams & { pageNumber: number; pageSize: number }) => URLSearchParams;
+export function buildCustomColumnDef(cc: CustomColumn): ColumnDef<TraceRow> {
+  return {
+    id: `custom:${cc.name}`,
+    accessorFn: (row) => (row as Record<string, unknown>)[`custom:${cc.name}`],
+    header: cc.name,
+    enableSorting: true,
+    meta: {
+      sql: cc.sql,
+      dataType: cc.dataType,
+      isCustom: true,
+    },
+  };
 }
 
-export const useTracesTableStore = create<TracesTableStoreState>()(
-  persist(
-    (set, get) => ({
-      columnDefs: [...STATIC_COLUMNS],
-      customColumns: [],
-
-      rebuildColumns: () => {
-        const { customColumns } = get();
-        const customCols: ColumnDef<TraceRow>[] = customColumns.map((cc) => ({
-          id: `custom:${cc.name}`,
-          accessorFn: (row) => (row as Record<string, unknown>)[`custom:${cc.name}`],
-          header: cc.name,
-          enableSorting: true,
-          meta: {
-            sql: cc.sql,
-            dataType: cc.dataType,
-            isCustom: true,
-          },
-        }));
-        set({ columnDefs: [...STATIC_COLUMNS, ...customCols] });
-      },
-
-      addCustomColumn: (column) => {
-        const { customColumns } = get();
-        if (customColumns.some((cc) => cc.name === column.name)) return;
-        set({ customColumns: [...customColumns, column] });
-        get().rebuildColumns();
-      },
-
-      updateCustomColumn: (oldName, column) => {
-        const { customColumns } = get();
-        set({ customColumns: customColumns.map((cc) => (cc.name === oldName ? column : cc)) });
-        get().rebuildColumns();
-      },
-
-      removeCustomColumn: (name) => {
-        const { customColumns } = get();
-        set({ customColumns: customColumns.filter((cc) => cc.name !== name) });
-        get().rebuildColumns();
-      },
-
-      buildFetchParams: (raw) => {
-        const { columnDefs } = get();
-        const urlParams = new URLSearchParams();
-        urlParams.set("pageNumber", raw.pageNumber.toString());
-        urlParams.set("pageSize", raw.pageSize.toString());
-        raw.filter.forEach((f) => urlParams.append("filter", f));
+interface TracesTableStoreState {
+  buildFetchParams: (
+    raw: RawUrlParams & { pageNumber: number; pageSize: number },
+    allColumnDefs: ColumnDef<TraceRow>[]
+  ) => URLSearchParams;
+}
 
-        // Send custom columns payload
-        const customCols = toColumnsPayload(columnDefs.filter((c) => c.meta?.isCustom));
-        if (customCols.length > 0) {
-          urlParams.set("customColumns", JSON.stringify(customCols));
-        }
+export const useTracesTableStore = create<TracesTableStoreState>()(() => ({
+  buildFetchParams: (raw, allColumnDefs) => {
+    const urlParams = new URLSearchParams();
+    urlParams.set("pageNumber", raw.pageNumber.toString());
+    urlParams.set("pageSize", raw.pageSize.toString());
+    raw.filter.forEach((f) => urlParams.append("filter", f));
 
-        // Sort — resolve SQL from column meta
-        if (raw.sortBy) {
-          urlParams.set("sortBy", raw.sortBy);
-          const col = columnDefs.find((c) => c.id === raw.sortBy);
-          if (col?.meta?.sql) urlParams.set("sortSql", col.meta.sql);
-        }
-        if (raw.sortDirection) urlParams.set("sortDirection", raw.sortDirection);
+    const customCols = toColumnsPayload(allColumnDefs.filter((c) => c.meta?.isCustom));
+    if (customCols.length > 0) {
+      urlParams.set("customColumns", JSON.stringify(customCols));
+    }
 
-        return urlParams;
-      },
-    }),
-    {
-      name: "traces-table-custom-columns",
-      partialize: (state) => ({ customColumns: state.customColumns }),
+    if (raw.sortBy) {
+      urlParams.set("sortBy", raw.sortBy);
+      const col = allColumnDefs.find((c) => c.id === raw.sortBy);
+      if (col?.meta?.sql) urlParams.set("sortSql", col.meta.sql);
     }
-  )
-);
+    if (raw.sortDirection) urlParams.set("sortDirection", raw.sortDirection);
+
+    return urlParams;
+  },
+}));
diff --git a/frontend/components/ui/columns-menu/columns-menu.tsx b/frontend/components/ui/columns-menu/columns-menu.tsx
index 20f7127b2..55b6be4ab 100644
--- a/frontend/components/ui/columns-menu/columns-menu.tsx
+++ b/frontend/components/ui/columns-menu/columns-menu.tsx
@@ -1,25 +1,25 @@
 import { AnimatePresence } from "framer-motion";
-import React, { useState } from "react";
+import React, { useCallback, useState } from "react";
 import { useStore } from "zustand";
 
 import { Button } from "@/components/ui/button.tsx";
-import { useDataTableStore } from "@/components/ui/infinite-datatable/model/datatable-store.tsx";
+import { type CustomColumn, useDataTableStore } from "@/components/ui/infinite-datatable/model/datatable-store.tsx";
 import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover.tsx";
 
 import { ColumnsListPanel } from "./columns-list-panel";
 import { CustomColumnPanel } from "./custom-column-panel";
-import type { ColumnActions, CustomColumn, CustomColumnPanelConfig } from "./types";
+import type { CustomColumnPanelConfig } from "./types";
 
 interface ColumnsMenuProps {
   /** Configuration for the custom column panel (schema, test query, etc.). When omitted, the custom column UI is hidden. */
   panelConfig?: CustomColumnPanelConfig;
-  /** Store actions for managing custom columns. Required when panelConfig is provided. */
-  columnActions?: ColumnActions;
   /** Whether to show the "Create column with SQL" button. Defaults to true. */
   showCreateButton?: boolean;
+  /** Return all current column defs for duplicate-name checking in the custom column panel. */
+  getColumnDefs?: () => import("@tanstack/react-table").ColumnDef<any>[];
 }
 
-export default function ColumnsMenu({ panelConfig, columnActions, showCreateButton = true }: ColumnsMenuProps) {
+export default function ColumnsMenu({ panelConfig, showCreateButton = true, getColumnDefs }: ColumnsMenuProps) {
   const store = useDataTableStore();
   const {
     lockedColumns,
@@ -29,6 +29,10 @@ export default function ColumnsMenu({ panelConfig, columnActions, showCreateButt
     setColumnOrder,
     columnVisibility,
     setColumnVisibility,
+    customColumns,
+    addCustomColumn,
+    updateCustomColumn,
+    removeCustomColumn,
   } = useStore(store, (state) => ({
     lockedColumns: state.lockedColumns,
     columnLabelMap: state.columnLabelMap,
@@ -37,9 +41,13 @@ export default function ColumnsMenu({ panelConfig, columnActions, showCreateButt
     setColumnOrder: state.setColumnOrder,
     columnVisibility: state.columnVisibility,
     setColumnVisibility: state.setColumnVisibility,
+    customColumns: state.customColumns,
+    addCustomColumn: state.addCustomColumn,
+    updateCustomColumn: state.updateCustomColumn,
+    removeCustomColumn: state.removeCustomColumn,
   }));
 
-  const hasCustomColumns = !!panelConfig && !!columnActions;
+  const hasCustomColumns = !!panelConfig;
 
   const [isOpen, setIsOpen] = useState(false);
   const [activePanel, setActivePanel] = useState<"list" | "form">("list");
@@ -54,34 +62,36 @@ export default function ColumnsMenu({ panelConfig, columnActions, showCreateButt
     });
   }
 
-  const handleEditColumn = (columnId: string) => {
-    if (!columnActions) return;
-    const col = columnActions.getColumnDef(columnId);
-    if (col?.meta?.isCustom) {
-      setEditingColumn({
-        name: col.header as string,
-        sql: col.meta.sql!,
-        dataType: col.meta.dataType as "string" | "number",
-      });
-      setActivePanel("form");
-    }
-  };
+  const handleEditColumn = useCallback(
+    (columnId: string) => {
+      const cc = customColumns.find((c) => `custom:${c.name}` === columnId);
+      if (cc) {
+        setEditingColumn(cc);
+        setActivePanel("form");
+      }
+    },
+    [customColumns]
+  );
 
-  const handleDeleteColumn = (columnId: string) => {
-    if (!columnActions) return;
-    columnActions.removeCustomColumn(columnId.replace("custom:", ""));
-  };
+  const handleDeleteColumn = useCallback(
+    (columnId: string) => {
+      removeCustomColumn(columnId.replace("custom:", ""));
+    },
+    [removeCustomColumn]
+  );
 
-  const handleSave = (column: CustomColumn) => {
-    if (!columnActions) return;
-    if (editingColumn) {
-      columnActions.updateCustomColumn(editingColumn.name, column);
-    } else {
-      columnActions.addCustomColumn(column);
-    }
-    setEditingColumn(null);
-    setActivePanel("list");
-  };
+  const handleSave = useCallback(
+    (column: CustomColumn) => {
+      if (editingColumn) {
+        updateCustomColumn(editingColumn.name, column);
+      } else {
+        addCustomColumn(column);
+      }
+      setEditingColumn(null);
+      setActivePanel("list");
+    },
+    [editingColumn, addCustomColumn, updateCustomColumn]
+  );
 
   const handleOpenChange = (open: boolean) => {
     setIsOpen(open);
@@ -91,6 +101,10 @@ export default function ColumnsMenu({ panelConfig, columnActions, showCreateButt
     }
   };
 
+  const panelConfigWithDefs = panelConfig
+    ? { ...panelConfig, getColumnDefs: getColumnDefs ?? panelConfig.getColumnDefs }
+    : undefined;
+
   return (
     <Popover open={isOpen} onOpenChange={handleOpenChange}>
       <PopoverTrigger asChild>
@@ -131,7 +145,7 @@ export default function ColumnsMenu({ panelConfig, columnActions, showCreateButt
               onDeleteColumn={hasCustomColumns ? handleDeleteColumn : undefined}
               showCreateButton={hasCustomColumns && showCreateButton}
             />
-          ) : panelConfig ? (
+          ) : panelConfigWithDefs ? (
             <CustomColumnPanel
               key={editingColumn?.name ?? "__new__"}
               onBack={() => {
@@ -140,7 +154,7 @@ export default function ColumnsMenu({ panelConfig, columnActions, showCreateButt
               }}
               onSave={handleSave}
               editingColumn={editingColumn ?? undefined}
-              config={panelConfig}
+              config={panelConfigWithDefs}
             />
           ) : null}
         </AnimatePresence>
diff --git a/frontend/components/ui/columns-menu/index.ts b/frontend/components/ui/columns-menu/index.ts
index 910afff9b..196a8a109 100644
--- a/frontend/components/ui/columns-menu/index.ts
+++ b/frontend/components/ui/columns-menu/index.ts
@@ -2,4 +2,4 @@ export { ColumnsListPanel } from "./columns-list-panel";
 export { default as ColumnsMenu } from "./columns-menu";
 export { ColumnsMenuItem } from "./columns-menu-item";
 export { CustomColumnPanel } from "./custom-column-panel";
-export type { ColumnActions, CustomColumn, CustomColumnPanelConfig } from "./types";
+export type { CustomColumn, CustomColumnPanelConfig } from "./types";
diff --git a/frontend/components/ui/columns-menu/types.ts b/frontend/components/ui/columns-menu/types.ts
index e9d96a424..d36b21a4e 100644
--- a/frontend/components/ui/columns-menu/types.ts
+++ b/frontend/components/ui/columns-menu/types.ts
@@ -3,8 +3,8 @@ import type { ColumnDef } from "@tanstack/react-table";
 import type { SQLSchemaConfig } from "@/components/sql/utils";
 import type { GenerationMode } from "@/lib/actions/sql";
 
-/** Shared custom column definition used by both evaluations and traces. */
-export type CustomColumn = { name: string; sql: string; dataType: "string" | "number" };
+/** @deprecated Import CustomColumn from @/components/ui/infinite-datatable/model/datatable-store instead. */
+export type { CustomColumn } from "@/components/ui/infinite-datatable/model/datatable-store";
 
 /** Configuration for the custom column panel to make it context-agnostic. */
 export interface CustomColumnPanelConfig {
@@ -27,12 +27,3 @@ export interface CustomColumnPanelConfig {
   /** Hint text shown below the SQL editor. */
   sqlHint?: string;
 }
-
-/** Callbacks for managing custom columns, injected by the consumer's store. */
-export interface ColumnActions {
-  addCustomColumn: (column: CustomColumn) => void;
-  updateCustomColumn: (oldName: string, column: CustomColumn) => void;
-  removeCustomColumn: (name: string) => void;
-  /** Return the ColumnDef for a given column ID, used for populating the edit form. */
-  getColumnDef: (columnId: string) => ColumnDef<any> | undefined;
-}
diff --git a/frontend/components/ui/infinite-datatable/index.tsx b/frontend/components/ui/infinite-datatable/index.tsx
index 458b3ada8..d9a66734c 100644
--- a/frontend/components/ui/infinite-datatable/index.tsx
+++ b/frontend/components/ui/infinite-datatable/index.tsx
@@ -15,6 +15,7 @@ import {
 import { restrictToHorizontalAxis } from "@dnd-kit/modifiers";
 import { arrayMove } from "@dnd-kit/sortable";
 import {
+  type ColumnDef,
   getCoreRowModel,
   getExpandedRowModel,
   type RowData,
@@ -30,7 +31,7 @@ import { Skeleton } from "@/components/ui/skeleton.tsx";
 import { Table } from "@/components/ui/table.tsx";
 import { cn } from "@/lib/utils.ts";
 
-import { useDataTableStore } from "./model/datatable-store.tsx";
+import { selectAllColumnDefs, useDataTableStore, useDataTableStoreSelector } from "./model/datatable-store.tsx";
 import { type InfiniteDataTableProps } from "./model/types.ts";
 import { InfiniteDatatableBody } from "./ui/body.tsx";
 import { InfiniteDatatableHeader } from "./ui/header.tsx";
@@ -65,7 +66,7 @@ export function InfiniteDataTable<TData extends RowData>({
   onSort,
 
   // TableOptions props
-  columns,
+  columns: columnsProp,
   data,
   state,
   enableRowSelection,
@@ -78,6 +79,11 @@ export function InfiniteDataTable<TData extends RowData>({
   ...tableOptions
 }: PropsWithChildren<InfiniteDataTableProps<TData>>) {
   const selectedRowIds = state?.rowSelection ? Object.keys(state.rowSelection) : [];
+
+  const store = useDataTableStore<TData>();
+  const storeColumns = useDataTableStoreSelector(selectAllColumnDefs) as ColumnDef<TData>[];
+  const columns = columnsProp ?? storeColumns;
+
   const finalColumns = useMemo(
     () => (enableRowSelection ? [createCheckboxColumn<TData>(), ...columns] : columns),
     [columns, enableRowSelection]
@@ -88,7 +94,6 @@ export function InfiniteDataTable<TData extends RowData>({
     [sortBy, sortDirection]
   );
 
-  const store = useDataTableStore();
   const {
     columnOrder,
     setColumnOrder,
diff --git a/frontend/components/ui/infinite-datatable/model/datatable-store.tsx b/frontend/components/ui/infinite-datatable/model/datatable-store.tsx
index 5b16448e2..8e1301be0 100644
--- a/frontend/components/ui/infinite-datatable/model/datatable-store.tsx
+++ b/frontend/components/ui/infinite-datatable/model/datatable-store.tsx
@@ -1,11 +1,17 @@
 "use client";
 
 import { type ColumnDef } from "@tanstack/react-table";
-import { isEqual, uniqBy } from "lodash";
-import { createContext, type ReactNode, useContext, useEffect, useMemo, useState } from "react";
-import { createStore, type StoreApi } from "zustand";
+import { uniqBy } from "lodash";
+import { createContext, type ReactNode, useContext, useState } from "react";
+import { createStore, type StoreApi, useStore } from "zustand";
 import { persist } from "zustand/middleware";
 
+export interface CustomColumn {
+  name: string;
+  sql: string;
+  dataType: "string" | "number";
+}
+
 export interface InfiniteScrollState<TData> {
   data: TData[];
   currentPage: number;
@@ -29,18 +35,20 @@ export interface InfiniteScrollActions<TData> {
   resetInfiniteScroll: () => void;
 }
 
-export interface SelectionState {
+export interface ColumnState {
   selectedRows: Set<string>;
-  defaultColumnOrder: string[];
   lockedColumns: string[];
   columnLabelMap: Record<string, string>;
   columnVisibility: Record<string, boolean>;
   columnOrder: string[];
   columnSizing: Record<string, number>;
   draggingColumnId: string | null;
+  customColumns: CustomColumn[];
+  staticColumnDefs: ColumnDef<any>[];
+  buildCustomColumnDef: ((cc: CustomColumn) => ColumnDef<any>) | null;
 }
 
-export interface SelectionActions {
+export interface ColumnActions {
   selectRow: (id: string) => void;
   deselectRow: (id: string) => void;
   toggleRow: (id: string) => void;
@@ -50,21 +58,65 @@ export interface SelectionActions {
   setColumnOrder: (order: string[]) => void;
   setColumnSizing: (sizing: Record<string, number>) => void;
   setDraggingColumnId: (columnId: string | null) => void;
-  reconcileColumns: (columnIds: string[]) => void;
   resetColumns: () => void;
   getStorageKey: () => string;
+  addCustomColumn: (column: CustomColumn) => void;
+  updateCustomColumn: (oldName: string, column: CustomColumn) => void;
+  removeCustomColumn: (name: string) => void;
+}
+
+type DataTableStore<TData> = InfiniteScrollState<TData> & InfiniteScrollActions<TData> & ColumnState & ColumnActions;
+
+/**
+ * Derive the combined column defs from static + custom.
+ * Exported so consumers can use it as a selector on the store.
+ */
+export function selectAllColumnDefs<TData>(state: DataTableStore<TData>): ColumnDef<TData>[] {
+  const { staticColumnDefs, customColumns, buildCustomColumnDef } = state;
+  if (!buildCustomColumnDef || customColumns.length === 0) {
+    return staticColumnDefs as ColumnDef<TData>[];
+  }
+  const customDefs = customColumns.map(buildCustomColumnDef);
+  return [...staticColumnDefs, ...customDefs] as ColumnDef<TData>[];
+}
+
+function buildColumnLabelMap(defs: ColumnDef<any>[]): Record<string, string> {
+  const map: Record<string, string> = {};
+  for (const c of defs) {
+    const id = (c as ColumnDef<any> & { id?: string }).id;
+    if (!id) continue;
+    map[id] = typeof c.header === "string" ? c.header : id;
+  }
+  return map;
 }
 
-type DataTableStore<TData> = InfiniteScrollState<TData> &
-  InfiniteScrollActions<TData> &
-  SelectionState &
-  SelectionActions;
+function buildColumnIds(defs: ColumnDef<any>[], enableRowSelection: boolean): string[] {
+  const ids = defs.map((c) => (c as ColumnDef<any> & { id?: string }).id).filter(Boolean) as string[];
+  return enableRowSelection ? ["__row_selection", ...ids] : ids;
+}
+
+/**
+ * Reconcile persisted columnOrder with the current set of column IDs.
+ * Returns the new columnOrder preserving user reordering where possible.
+ */
+function reconcileColumnOrder(currentOrder: string[], columnIds: string[]): string[] {
+  if (currentOrder.length === 0) return columnIds;
+
+  const idSet = new Set(columnIds);
+  const pruned = currentOrder.filter((id) => idSet.has(id));
+  const existingSet = new Set(pruned);
+  const added = columnIds.filter((id) => !existingSet.has(id));
+  return [...pruned, ...added];
+}
 
 function createDataTableStore<TData>(
   uniqueKey: string = "id",
   storageKey?: string,
   pageSize: number = 50,
-  lockedColumns: string[] = []
+  lockedColumns: string[] = [],
+  initialStaticColumnDefs: ColumnDef<any>[] = [],
+  initialBuildCustomColumnDef: ((cc: CustomColumn) => ColumnDef<any>) | null = null,
+  enableRowSelection: boolean = false
 ): StoreApi<DataTableStore<TData>> {
   const storeConfig = (
     set: StoreApi<DataTableStore<TData>>["setState"],
@@ -78,13 +130,15 @@ function createDataTableStore<TData>(
     uniqueKey,
     hasMore: true,
     pageSize,
-    defaultColumnOrder: [],
     lockedColumns,
     columnLabelMap: {},
     columnVisibility: {},
     columnOrder: [],
     columnSizing: {},
     draggingColumnId: null,
+    customColumns: [],
+    staticColumnDefs: initialStaticColumnDefs,
+    buildCustomColumnDef: initialBuildCustomColumnDef,
     setData: (updater) => set((state) => ({ data: updater(state.data) })),
     setCurrentPage: (currentPage) => set({ currentPage }),
     setIsFetching: (isFetching) => set({ isFetching }),
@@ -95,30 +149,69 @@ function createDataTableStore<TData>(
     setColumnOrder: (order) => set({ columnOrder: order }),
     setColumnSizing: (sizing) => set({ columnSizing: sizing }),
     setDraggingColumnId: (columnId) => set({ draggingColumnId: columnId }),
-    reconcileColumns: (columnIds) => {
-      const state = get();
-      if (isEqual(state.defaultColumnOrder, columnIds)) return;
-
-      if (state.columnOrder.length === 0) {
-        set({ defaultColumnOrder: columnIds, columnOrder: columnIds });
-        return;
-      }
-
-      const idSet = new Set(columnIds);
-      const pruned = state.columnOrder.filter((id) => idSet.has(id));
-      const existingSet = new Set(pruned);
-      const added = columnIds.filter((id) => !existingSet.has(id));
-
-      set({ defaultColumnOrder: columnIds, columnOrder: [...pruned, ...added] });
-    },
     resetColumns: () => {
       const state = get();
+      const allDefs = selectAllColumnDefs(state);
       set({
         columnVisibility: {},
-        columnOrder: state.defaultColumnOrder,
+        columnOrder: buildColumnIds(allDefs, enableRowSelection),
         columnSizing: {},
       });
     },
+    addCustomColumn: (column) => {
+      const { customColumns, columnOrder, buildCustomColumnDef: builder } = get();
+      if (customColumns.some((cc) => cc.name === column.name)) return;
+      const newCustomColumns = [...customColumns, column];
+      const columnId = `custom:${column.name}`;
+      const newState: Partial<DataTableStore<TData>> = {
+        customColumns: newCustomColumns,
+        columnOrder: [...columnOrder, columnId],
+      };
+      // Rebuild label map to include new column
+      if (builder) {
+        const allDefs = [...get().staticColumnDefs, ...newCustomColumns.map(builder)];
+        newState.columnLabelMap = buildColumnLabelMap(allDefs);
+      }
+      set(newState);
+    },
+    updateCustomColumn: (oldName, column) => {
+      const { customColumns, columnOrder, columnVisibility, buildCustomColumnDef: builder } = get();
+      const newCustomColumns = customColumns.map((cc) => (cc.name === oldName ? column : cc));
+      const oldId = `custom:${oldName}`;
+      const newId = `custom:${column.name}`;
+      const newState: Partial<DataTableStore<TData>> = { customColumns: newCustomColumns };
+      if (oldName !== column.name) {
+        newState.columnOrder = columnOrder.map((id) => (id === oldId ? newId : id));
+        const newVisibility = { ...columnVisibility };
+        if (oldId in newVisibility) {
+          newVisibility[newId] = newVisibility[oldId];
+          delete newVisibility[oldId];
+        }
+        newState.columnVisibility = newVisibility;
+      }
+      if (builder) {
+        const allDefs = [...get().staticColumnDefs, ...newCustomColumns.map(builder)];
+        newState.columnLabelMap = buildColumnLabelMap(allDefs);
+      }
+      set(newState);
+    },
+    removeCustomColumn: (name) => {
+      const { customColumns, columnVisibility, columnOrder, buildCustomColumnDef: builder } = get();
+      const columnId = `custom:${name}`;
+      const newVisibility = { ...columnVisibility };
+      delete newVisibility[columnId];
+      const newCustomColumns = customColumns.filter((cc) => cc.name !== name);
+      const newState: Partial<DataTableStore<TData>> = {
+        customColumns: newCustomColumns,
+        columnVisibility: newVisibility,
+        columnOrder: columnOrder.filter((id) => id !== columnId),
+      };
+      if (builder) {
+        const allDefs = [...get().staticColumnDefs, ...newCustomColumns.map(builder)];
+        newState.columnLabelMap = buildColumnLabelMap(allDefs);
+      }
+      set(newState);
+    },
     appendData: (items, count) =>
       set((state) => {
         const combined = [...state.data, ...items];
@@ -193,17 +286,18 @@ function createDataTableStore<TData>(
           columnVisibility: state.columnVisibility,
           columnOrder: state.columnOrder,
           columnSizing: state.columnSizing,
+          customColumns: state.customColumns,
         }),
-        // Restore persisted state as-is; reconcileColumns (called by the provider) corrects it against the actual columns prop.
         merge: (persistedState, currentState) => {
           const persisted = persistedState as Partial<
-            Pick<SelectionState, "columnVisibility" | "columnOrder" | "columnSizing">
+            Pick<ColumnState, "columnVisibility" | "columnOrder" | "columnSizing" | "customColumns">
           >;
           return {
             ...currentState,
             columnVisibility: persisted?.columnVisibility ?? {},
             columnOrder: persisted?.columnOrder ?? [],
             columnSizing: persisted?.columnSizing ?? {},
+            customColumns: persisted?.customColumns ?? [],
           };
         },
       })
@@ -222,9 +316,10 @@ export interface DataTableStateProviderProps {
   uniqueKey?: string;
   pageSize?: number;
   storageKey?: string;
-  columns?: ColumnDef<any>[];
+  columnDefs?: ColumnDef<any>[];
   enableRowSelection?: boolean;
   lockedColumns?: string[];
+  buildCustomColumnDef?: (cc: CustomColumn) => ColumnDef<any>;
 }
 
 export function DataTableStateProvider<TData>({
@@ -232,41 +327,57 @@ export function DataTableStateProvider<TData>({
   storageKey,
   uniqueKey = "id",
   pageSize = 50,
-  columns = [],
+  columnDefs = [],
   enableRowSelection = false,
   lockedColumns = [],
+  buildCustomColumnDef,
 }: DataTableStateProviderProps) {
-  const [store] = useState(() => createDataTableStore<TData>(uniqueKey, storageKey, pageSize, lockedColumns));
+  const [store] = useState(() => {
+    const s = createDataTableStore<TData>(
+      uniqueKey,
+      storageKey,
+      pageSize,
+      lockedColumns,
+      columnDefs,
+      buildCustomColumnDef ?? null,
+      enableRowSelection
+    );
 
-  const columnIds = useMemo(() => {
-    const ids = columns.map((c) => (c as ColumnDef<any> & { id?: string }).id).filter(Boolean) as string[];
-    return enableRowSelection ? ["__row_selection", ...ids] : ids;
-  }, [columns, enableRowSelection]);
+    // Synchronously reconcile persisted columnOrder with the column IDs
+    // from the initial defs + any persisted custom columns.
+    const state = s.getState();
+    const customDefs =
+      buildCustomColumnDef && state.customColumns.length > 0 ? state.customColumns.map(buildCustomColumnDef) : [];
+    const allDefs = [...columnDefs, ...customDefs];
+    const columnIds = buildColumnIds(allDefs, enableRowSelection);
+    const reconciledOrder = reconcileColumnOrder(state.columnOrder, columnIds);
 
-  const columnLabelMap = useMemo(() => {
-    const map: Record<string, string> = {};
-    for (const c of columns) {
-      const id = (c as ColumnDef<any> & { id?: string }).id;
-      if (!id) continue;
-      map[id] = typeof c.header === "string" ? c.header : id;
-    }
-    return map;
-  }, [columns]);
+    s.setState({
+      columnLabelMap: buildColumnLabelMap(allDefs),
+      columnOrder: reconciledOrder,
+    });
 
-  useEffect(() => {
-    if (columnIds.length > 0) {
-      store.getState().reconcileColumns(columnIds);
-      store.setState({ columnLabelMap });
-    }
-  }, [columnIds, columnLabelMap, store]);
+    return s;
+  });
 
   return <DataTableContext.Provider value={store}>{children}</DataTableContext.Provider>;
 }
 
-export function useDataTableStore<TData>() {
+function useDataTableContext() {
   const store = useContext(DataTableContext);
   if (!store) {
-    throw new Error("useDataTableStore must be used within DataTableStateProvider");
+    throw new Error("useDataTableStore / useDataTableStoreSelector must be used within DataTableStateProvider");
   }
-  return store as DataTableStoreApi<TData>;
+  return store;
+}
+
+export function useDataTableStore<TData>() {
+  return useDataTableContext() as DataTableStoreApi<TData>;
+}
+
+export function useDataTableStoreSelector<U>(
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  selector: (state: DataTableStore<any>) => U
+): U {
+  return useStore(useDataTableContext(), selector);
 }
diff --git a/frontend/components/ui/infinite-datatable/model/types.ts b/frontend/components/ui/infinite-datatable/model/types.ts
index c01f12721..19c84ae7f 100644
--- a/frontend/components/ui/infinite-datatable/model/types.ts
+++ b/frontend/components/ui/infinite-datatable/model/types.ts
@@ -13,7 +13,8 @@ export interface InfiniteDataTableProps<TData extends RowData> extends Omit<
   "data" | "columns"
 > {
   data: TData[];
-  columns: TableOptions<TData>["columns"];
+  /** Optional — when omitted, columns are read from the DataTableStore context. */
+  columns?: TableOptions<TData>["columns"];
 
   hasMore: boolean;
   isFetching: boolean;