From 72b1cf66e06d52d99d9a5c39e076eddb2db5c6c9 Mon Sep 17 00:00:00 2001
From: Ben Sherman <ben@wandb.com>
Date: Tue, 4 Feb 2025 11:04:21 -0800
Subject: [PATCH] chore(weave): sort datasets while editing

---
 .../Browse3/datasets/EditableDatasetView.tsx  | 23 +++++--------------
 .../traceServerClientTypes.ts                 |  1 +
 2 files changed, 7 insertions(+), 17 deletions(-)

diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/datasets/EditableDatasetView.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/datasets/EditableDatasetView.tsx
index 504d6259a840..3098a0c0fbf9 100644
--- a/weave-js/src/components/PagePanelComponents/Home/Browse3/datasets/EditableDatasetView.tsx
+++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/datasets/EditableDatasetView.tsx
@@ -100,15 +100,6 @@ export const EditableDatasetView: FC<EditableDataTableViewProps> = ({
     pageSize: 50,
   });
 
-  // Reset sort model and pagination if we enter edit mode with sorting applied.
-  useEffect(() => {
-    if (isEditing && sortModel.length > 0) {
-      setPaginationModel({page: 0, pageSize: 50});
-      setSortModel([]);
-      setSortBy([]);
-    }
-  }, [isEditing, sortModel]);
-
   const sharedRef = useContext(WeaveCHTableSourceRefContext);
 
   const history = useHistory();
@@ -269,16 +260,14 @@ export const EditableDatasetView: FC<EditableDataTableViewProps> = ({
 
   const rows = useMemo(() => {
     if (fetchQueryLoaded) {
-      return loadedRows.map((row, i) => {
+      return loadedRows.map(row => {
         const digest = row.digest;
-        const absoluteIndex =
-          i + paginationModel.pageSize * paginationModel.page;
         const value = flattenObjectPreservingWeaveTypes(row.val);
-        const editedRow = editedRows.get(absoluteIndex);
+        const editedRow = editedRows.get(row.original_index);
         return {
           ___weave: {
-            id: `${digest}_${absoluteIndex}`,
-            index: absoluteIndex,
+            id: `${digest}_${row.original_index}`,
+            index: row.original_index,
             isNew: false,
             serverValue: value,
           },
@@ -287,7 +276,7 @@ export const EditableDatasetView: FC<EditableDataTableViewProps> = ({
       });
     }
     return [];
-  }, [loadedRows, fetchQueryLoaded, editedRows, paginationModel]);
+  }, [loadedRows, fetchQueryLoaded, editedRows]);
 
   const combinedRows = useMemo(() => {
     if (
@@ -404,7 +393,7 @@ export const EditableDatasetView: FC<EditableDataTableViewProps> = ({
       flex: columnWidths[field as string] ? undefined : 1,
       minWidth: 100,
       editable: isEditing,
-      sortable: !isEditing,
+      sortable: true,
       filterable: false,
       renderCell: (params: GridRenderCellParams) => {
         if (!isEditing) {
diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/wfReactInterface/traceServerClientTypes.ts b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/wfReactInterface/traceServerClientTypes.ts
index 9a72476ff8f4..25592bf4e6e1 100644
--- a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/wfReactInterface/traceServerClientTypes.ts
+++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/wfReactInterface/traceServerClientTypes.ts
@@ -297,6 +297,7 @@ export type TraceTableQueryRes = {
   rows: Array<{
     digest: string;
     val: any;
+    original_index?: number;
   }>;
 };