From dd0841da0ee17a544f713e2222b2b6b72ed9b4e4 Mon Sep 17 00:00:00 2001
From: Ben Sherman <bcanfieldsherman@gmail.com>
Date: Tue, 7 Jan 2025 13:23:21 -0800
Subject: [PATCH] feat(weave): dataset editing UI

---
 .../Home/Browse3/datasets/CellRenderers.tsx   | 232 ++++++++
 .../Browse3/datasets/DatasetEditorContext.tsx | 132 +++++
 .../Browse3/datasets/DatasetVersionPage.tsx   | 349 ++++++++++--
 .../Browse3/datasets/EditableDatasetView.tsx  | 498 ++++++++++++++++++
 .../Browse3/pages/CallPage/DataTableView.tsx  |   4 +-
 .../Home/Browse3/pages/ObjectVersionPage.tsx  |   7 +-
 6 files changed, 1164 insertions(+), 58 deletions(-)
 create mode 100644 weave-js/src/components/PagePanelComponents/Home/Browse3/datasets/CellRenderers.tsx
 create mode 100644 weave-js/src/components/PagePanelComponents/Home/Browse3/datasets/DatasetEditorContext.tsx
 create mode 100644 weave-js/src/components/PagePanelComponents/Home/Browse3/datasets/EditableDatasetView.tsx

diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/datasets/CellRenderers.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/datasets/CellRenderers.tsx
new file mode 100644
index 000000000000..04f0731ae9b0
--- /dev/null
+++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/datasets/CellRenderers.tsx
@@ -0,0 +1,232 @@
+import {Box, TextField, Typography} from '@mui/material';
+import {Popover} from '@mui/material';
+import {
+  GridRenderCellParams,
+  GridRenderEditCellParams,
+} from '@mui/x-data-grid-pro';
+import {Button} from '@wandb/weave/components/Button';
+import {Icon} from '@wandb/weave/components/Icon';
+import React, {useState} from 'react';
+import {ResizableBox} from 'react-resizable';
+
+import {DraggableGrow, DraggableHandle} from '../../../../DraggablePopups';
+
+const cellViewingStyles = {
+  height: '100%',
+  width: '100%',
+  fontFamily: '"Source Sans Pro", sans-serif',
+  fontSize: '14px',
+  lineHeight: '1.5',
+  padding: '8px 12px',
+  display: 'flex',
+  alignItems: 'center',
+  transition: 'background-color 0.2s ease',
+};
+
+interface CellViewingRendererProps extends GridRenderCellParams {
+  isEdited?: boolean;
+  isDeleted?: boolean;
+  isNew?: boolean;
+}
+
+export const CellViewingRenderer: React.FC<CellViewingRendererProps> = ({
+  value,
+  isEdited = false,
+  isDeleted = false,
+  isNew = false,
+  api,
+  id,
+  field,
+}) => {
+  const [isHovered, setIsHovered] = useState(false);
+
+  const handleEditClick = (event: React.MouseEvent) => {
+    event.stopPropagation();
+    api.startCellEditMode({id, field});
+  };
+
+  const getBackgroundColor = () => {
+    if (isDeleted) {
+      return 'rgba(255, 0, 0, 0.1)';
+    }
+    if (isEdited) {
+      return 'rgba(0, 128, 128, 0.1)';
+    }
+    if (isNew) {
+      return 'rgba(0, 255, 0, 0.1)';
+    }
+    return 'transparent';
+  };
+
+  return (
+    <Box
+      onClick={handleEditClick}
+      onMouseEnter={() => setIsHovered(true)}
+      onMouseLeave={() => setIsHovered(false)}
+      sx={{
+        ...cellViewingStyles,
+        position: 'relative',
+        cursor: 'pointer',
+        backgroundColor: getBackgroundColor(),
+        opacity: isDeleted ? 0.5 : 1,
+        textDecoration: isDeleted ? 'line-through' : 'none',
+        '&:hover': {
+          backgroundColor: 'rgba(0, 0, 0, 0.04)',
+        },
+      }}>
+      <span style={{flex: 1}}>{value}</span>
+      {isHovered && (
+        <Box
+          sx={{
+            display: 'flex',
+            alignItems: 'center',
+            opacity: 0,
+            transition: 'opacity 0.2s ease',
+            cursor: 'pointer',
+            animation: 'fadeIn 0.2s ease forwards',
+            '@keyframes fadeIn': {
+              from: {opacity: 0},
+              to: {opacity: 0.5},
+            },
+            '&:hover': {
+              opacity: 0.8,
+            },
+          }}>
+          <Icon name="pencil-edit" height={14} width={14} />
+        </Box>
+      )}
+    </Box>
+  );
+};
+
+export const CellEditingRenderer: React.FC<
+  GridRenderEditCellParams
+> = params => {
+  const {id, value, field, hasFocus, api} = params;
+  const inputRef = React.useRef<HTMLTextAreaElement>(null);
+  const [anchorEl, setAnchorEl] = useState<HTMLDivElement | null>(null);
+
+  // Calculate initial position when edit mode starts
+  React.useLayoutEffect(() => {
+    const element = api.getCellElement(id, field);
+    if (element) {
+      setAnchorEl(element);
+    }
+  }, [api, id, field]);
+
+  React.useLayoutEffect(() => {
+    if (hasFocus) {
+      inputRef.current?.focus();
+    }
+  }, [hasFocus]);
+
+  const handleValueChange = (event: React.ChangeEvent<HTMLInputElement>) => {
+    const newValue = event.target.value;
+    api.setEditCellValue({id, field, value: newValue});
+  };
+
+  const handleKeyDown = (event: React.KeyboardEvent) => {
+    if (event.key === 'Enter' && !event.metaKey) {
+      event.stopPropagation();
+    }
+  };
+
+  return (
+    <>
+      <CellViewingRenderer {...params} />
+      <Popover
+        open={!!anchorEl}
+        anchorEl={anchorEl}
+        anchorOrigin={{
+          vertical: 'bottom',
+          horizontal: 'left',
+        }}
+        transformOrigin={{
+          vertical: 'top',
+          horizontal: 'left',
+        }}
+        onClose={() => setAnchorEl(null)}
+        TransitionComponent={DraggableGrow}
+        sx={{
+          '& .MuiPaper-root': {
+            backgroundColor: 'white',
+            boxShadow:
+              '0px 4px 20px rgba(0, 0, 0, 0.15), 0px 0px 40px rgba(0, 0, 0, 0.05)',
+            border: 'none',
+            borderRadius: '4px',
+            overflow: 'hidden',
+            padding: '0px',
+          },
+        }}>
+        <ResizableBox width={400} height={300}>
+          <Box
+            sx={{
+              height: '100%',
+              display: 'flex',
+              flexDirection: 'column',
+            }}>
+            <DraggableHandle>
+              <Box
+                sx={{
+                  cursor: 'grab',
+                  backgroundColor: 'white',
+                  display: 'flex',
+                  justifyContent: 'space-between',
+                  padding: '8px',
+                  borderBottom: '1px solid rgba(0, 0, 0, 0.05)',
+                }}>
+                <Typography
+                  variant="caption"
+                  sx={{
+                    fontFamily: '"Source Sans Pro", sans-serif',
+                    color: 'gray.500',
+                  }}>
+                  ⌘+Enter to close
+                </Typography>
+                <Button
+                  icon="close"
+                  size="small"
+                  variant="ghost"
+                  onClick={() => api.stopCellEditMode({id, field})}
+                />
+              </Box>
+            </DraggableHandle>
+            <Box
+              sx={{
+                flex: 1,
+                overflow: 'auto',
+                padding: '12px',
+              }}>
+              <TextField
+                inputRef={inputRef}
+                value={value}
+                onChange={handleValueChange}
+                onKeyDown={handleKeyDown}
+                fullWidth
+                multiline
+                autoFocus
+                sx={{
+                  '& .MuiInputBase-root': {
+                    fontFamily: '"Source Sans Pro", sans-serif',
+                    fontSize: '12px',
+                    border: 'none',
+                    backgroundColor: 'white',
+                  },
+                  '& .MuiInputBase-input': {
+                    padding: 0,
+                  },
+                  '& .MuiOutlinedInput-notchedOutline': {
+                    border: 'none',
+                  },
+                  '& textarea': {
+                    overflow: 'visible !important',
+                  },
+                }}
+              />
+            </Box>
+          </Box>
+        </ResizableBox>
+      </Popover>
+    </>
+  );
+};
diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/datasets/DatasetEditorContext.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/datasets/DatasetEditorContext.tsx
new file mode 100644
index 000000000000..5a0bd561e75c
--- /dev/null
+++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/datasets/DatasetEditorContext.tsx
@@ -0,0 +1,132 @@
+import React, {createContext, useCallback, useContext, useState} from 'react';
+
+interface DatasetRow {
+  // id: string | number;
+  [key: string]: any;
+  ___weave: {
+    id: string;
+    // The index must be set for rows of an existing dataset.
+    index?: number;
+    isNew?: boolean;
+  };
+}
+
+interface EditedCell {
+  [fieldName: string]: unknown;
+}
+
+interface DatasetEditContextType {
+  /** Map of edited cells, keyed by row absolute index */
+  editedCellsMap: Map<number, EditedCell>;
+  setEditedCellsMap: React.Dispatch<
+    React.SetStateAction<Map<number, EditedCell>>
+  >;
+  /** Map of complete edited rows, keyed by row absolute index */
+  editedRows: Map<number, DatasetRow>;
+  setEditedRows: React.Dispatch<React.SetStateAction<Map<number, DatasetRow>>>;
+  /** Callback to process row updates from the data grid */
+  processRowUpdate: (newRow: DatasetRow, oldRow: DatasetRow) => DatasetRow;
+  /** Array of row indices that have been marked for deletion */
+  deletedRows: number[];
+  setDeletedRows: React.Dispatch<React.SetStateAction<number[]>>;
+  /** Map of newly added rows, keyed by temporary row ID */
+  addedRows: Map<string, DatasetRow>;
+  setAddedRows: React.Dispatch<React.SetStateAction<Map<string, DatasetRow>>>;
+  /** Reset the context to its initial state */
+  reset: () => void;
+}
+
+export const DatasetEditContext = createContext<
+  DatasetEditContextType | undefined
+>(undefined);
+
+export const useDatasetEditContext = () => {
+  const context = useContext(DatasetEditContext);
+  if (!context) {
+    throw new Error(
+      'useDatasetEditContext must be used within a DatasetEditProvider'
+    );
+  }
+  return context;
+};
+
+interface DatasetEditProviderProps {
+  children: React.ReactNode;
+}
+
+export const DatasetEditProvider: React.FC<DatasetEditProviderProps> = ({
+  children,
+}) => {
+  const [editedCellsMap, setEditedCellsMap] = useState<Map<number, EditedCell>>(
+    new Map()
+  );
+  const [editedRows, setEditedRows] = useState<Map<number, DatasetRow>>(
+    new Map()
+  );
+  const [deletedRows, setDeletedRows] = useState<number[]>([]);
+  const [addedRows, setAddedRows] = useState<Map<string, DatasetRow>>(
+    new Map()
+  );
+
+  const processRowUpdate = useCallback(
+    (newRow: DatasetRow, oldRow: DatasetRow): DatasetRow => {
+      const changedField = Object.keys(newRow).find(
+        key => newRow[key] !== oldRow[key] && key !== 'id'
+      );
+
+      if (changedField) {
+        const rowKey = String(oldRow.___weave.id);
+        const rowIndex = oldRow.___weave.index;
+        if (oldRow.___weave.isNew) {
+          setAddedRows(prev => {
+            const updatedMap = new Map(prev);
+            updatedMap.set(rowKey, newRow);
+            return updatedMap;
+          });
+        } else {
+          setEditedCellsMap(prev => {
+            const existingEdits = prev.get(rowIndex!) || {};
+            const updatedMap = new Map(prev);
+            updatedMap.set(rowIndex!, {
+              ...existingEdits,
+              [changedField]: newRow[changedField],
+            });
+            return updatedMap;
+          });
+          setEditedRows(prev => {
+            const updatedMap = new Map(prev);
+            updatedMap.set(rowIndex!, newRow);
+            return updatedMap;
+          });
+        }
+      }
+      return newRow;
+    },
+    []
+  );
+
+  const reset = useCallback(() => {
+    setEditedCellsMap(new Map());
+    setEditedRows(new Map());
+    setDeletedRows([]);
+    setAddedRows(new Map());
+  }, []);
+
+  return (
+    <DatasetEditContext.Provider
+      value={{
+        editedCellsMap,
+        setEditedCellsMap,
+        editedRows,
+        setEditedRows,
+        processRowUpdate,
+        deletedRows,
+        setDeletedRows,
+        addedRows,
+        setAddedRows,
+        reset,
+      }}>
+      {children}
+    </DatasetEditContext.Provider>
+  );
+};
diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/datasets/DatasetVersionPage.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/datasets/DatasetVersionPage.tsx
index f3aed14e9efd..3e80255554f9 100644
--- a/weave-js/src/components/PagePanelComponents/Home/Browse3/datasets/DatasetVersionPage.tsx
+++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/datasets/DatasetVersionPage.tsx
@@ -1,13 +1,15 @@
-import Box from '@mui/material/Box';
-import React, {useMemo} from 'react';
+import {Box, Popover, Typography} from '@mui/material';
+import React, {useCallback, useMemo, useState} from 'react';
+import {Link, useHistory} from 'react-router-dom';
+import {toast} from 'react-toastify';
 
+import {Button} from '../../../../Button';
 import {Icon} from '../../../../Icon';
 import {LoadingDots} from '../../../../LoadingDots';
 import {Tailwind} from '../../../../Tailwind';
+import {useWeaveflowCurrentRouteContext} from '../context';
 import {WeaveCHTableSourceRefContext} from '../pages/CallPage/DataTableView';
-import {ObjectViewerSection} from '../pages/CallPage/ObjectViewerSection';
-import {objectVersionText} from '../pages/common/Links';
-import {ObjectVersionsLink} from '../pages/common/Links';
+import {ObjectVersionsLink, objectVersionText} from '../pages/common/Links';
 import {CenteredAnimatedLoader} from '../pages/common/Loader';
 import {
   ScrollableTabContent,
@@ -15,25 +17,60 @@ import {
 } from '../pages/common/SimplePageLayout';
 import {TabUseDataset} from '../pages/TabUseDataset';
 import {useWFHooks} from '../pages/wfReactInterface/context';
+import {TableUpdateSpec} from '../pages/wfReactInterface/traceServerClientTypes';
 import {objectVersionKeyToRefUri} from '../pages/wfReactInterface/utilities';
 import {ObjectVersionSchema} from '../pages/wfReactInterface/wfDataModelHooksInterface';
 import {CustomWeaveTypeProjectContext} from '../typeViews/CustomWeaveTypeDispatcher';
+import {useDatasetEditContext} from './DatasetEditorContext';
+import {EditableDatasetView} from './EditableDatasetView';
+
+const PUBLISHED_LINK_STYLES = {
+  color: 'rgb(94, 234, 212)',
+  textDecoration: 'none',
+  fontFamily: 'Inconsolata',
+  fontWeight: 600,
+} as const;
+
+const POPOVER_STYLES = {
+  '& .MuiPopover-paper': {
+    marginTop: '8px',
+    marginRight: '8px',
+  },
+} as const;
+
+const CODE_STYLES = {
+  fontFamily: 'Inconsolata',
+  fontWeight: 600,
+  backgroundColor: 'rgba(0, 0, 0, 0.04)',
+  padding: '2px 4px',
+  borderRadius: '4px',
+} as const;
 
 export const DatasetVersionPage: React.FC<{
   objectVersion: ObjectVersionSchema;
-}> = ({objectVersion}) => {
-  const {useRootObjectVersions, useRefsData} = useWFHooks();
+  refExtra?: string;
+}> = ({objectVersion, refExtra}) => {
+  const {editedCellsMap, editedRows, deletedRows, addedRows} =
+    useDatasetEditContext();
+  const router = useWeaveflowCurrentRouteContext();
+  const {useRootObjectVersions, useRefsData, useTableUpdate, useObjCreate} =
+    useWFHooks();
+
+  const [isEditing, setIsEditing] = useState(false);
+  const [publishAnchorEl, setPublishAnchorEl] = useState<HTMLElement | null>(
+    null
+  );
+
   const entityName = objectVersion.entity;
   const projectName = objectVersion.project;
   const objectName = objectVersion.objectId;
   const objectVersionIndex = objectVersion.versionIndex;
+  const projectId = `${entityName}/${projectName}`;
 
   const objectVersions = useRootObjectVersions(
     entityName,
     projectName,
-    {
-      objectIds: [objectName],
-    },
+    {objectIds: [objectName]},
     undefined,
     true
   );
@@ -53,12 +90,203 @@ export const DatasetVersionPage: React.FC<{
       typeof viewerData !== 'object' ||
       viewerData === null ||
       Array.isArray(viewerData);
-    if (dataIsPrimitive) {
-      return {_result: viewerData};
-    }
-    return viewerData;
+    return dataIsPrimitive ? {_result: viewerData} : viewerData;
   }, [viewerData]);
 
+  const originalTableDigest = viewerDataAsObject?.rows?.split('/').pop() ?? '';
+
+  const handleEditClick = useCallback(() => setIsEditing(true), []);
+  const handleCancelClick = useCallback(() => setIsEditing(false), []);
+  const handlePublishClose = () => setPublishAnchorEl(null);
+
+  const handlePublishClick = useCallback(
+    (event: React.MouseEvent<HTMLElement>) => {
+      if (
+        editedCellsMap.size > 0 ||
+        deletedRows.length > 0 ||
+        addedRows.size > 0
+      ) {
+        setPublishAnchorEl(event.currentTarget);
+      }
+    },
+    [editedCellsMap, deletedRows, addedRows]
+  );
+
+  const cleanRow = (row: any) => {
+    return Object.fromEntries(
+      Object.entries(row).filter(([key]) => !['___weave'].includes(key))
+    );
+  };
+
+  const convertEditsToTableUpdateSpec = useCallback(() => {
+    const updates: TableUpdateSpec[] = [];
+
+    editedRows.forEach((editedRow, rowIndex) => {
+      if (rowIndex !== undefined) {
+        updates.push({pop: {index: rowIndex}});
+        updates.push({
+          insert: {
+            index: rowIndex,
+            row: cleanRow(editedRow),
+          },
+        });
+      }
+    });
+
+    deletedRows
+      .sort((a, b) => b - a)
+      .forEach(rowIndex => {
+        updates.push({pop: {index: rowIndex}});
+      });
+
+    Array.from(addedRows.values())
+      .reverse()
+      .forEach(row => {
+        updates.push({
+          insert: {
+            index: 0,
+            row: cleanRow(row),
+          },
+        });
+      });
+
+    return updates;
+  }, [editedRows, deletedRows, addedRows]);
+
+  const tableUpdate = useTableUpdate();
+  const objCreate = useObjCreate();
+
+  const history = useHistory();
+
+  const handlePublish = useCallback(async () => {
+    setIsEditing(false);
+    setPublishAnchorEl(null);
+
+    const tableUpdateSpecs = convertEditsToTableUpdateSpec();
+    const tableUpdateResp = await tableUpdate(
+      projectId,
+      originalTableDigest,
+      tableUpdateSpecs
+    );
+    const tableRef = `weave:///${projectId}/table/${tableUpdateResp.digest}`;
+
+    const newObjVersion = await objCreate(projectId, objectName, {
+      ...objectVersion.val,
+      rows: tableRef,
+    });
+
+    const url = router.objectVersionUIUrl(
+      entityName,
+      projectName,
+      objectName,
+      newObjVersion,
+      undefined,
+      refExtra
+    );
+
+    toast(
+      <div style={{display: 'flex', alignItems: 'center', gap: '8px'}}>
+        <Icon name="checkmark" width={20} height={20} />
+        Published{' '}
+        <Link to={url} style={PUBLISHED_LINK_STYLES}>
+          {objectName}:v{objectVersionCount}
+        </Link>
+      </div>
+    );
+    history.push(url);
+  }, [
+    objectVersionCount,
+    history,
+    refExtra,
+    router,
+    objectName,
+    objectVersion.val,
+    convertEditsToTableUpdateSpec,
+    projectId,
+    objCreate,
+    tableUpdate,
+    originalTableDigest,
+    entityName,
+    projectName,
+  ]);
+
+  const renderEditingControls = () => (
+    <div className="flex gap-8">
+      <Typography
+        variant="body2"
+        sx={{
+          color: 'text.secondary',
+          display: 'flex',
+          alignItems: 'center',
+          fontSize: '14px',
+          fontFamily: 'Source Sans Pro',
+          gap: '4px',
+        }}>
+        <Icon name="pencil-edit" width={14} height={14} />
+        Editing dataset
+      </Typography>
+      <Button
+        title="Cancel"
+        variant="ghost"
+        size="medium"
+        icon="close"
+        onClick={handleCancelClick}>
+        Cancel
+      </Button>
+      <Button
+        title="Publish"
+        size="medium"
+        variant="primary"
+        icon="checkmark"
+        onClick={handlePublishClick}
+        disabled={
+          editedCellsMap.size === 0 &&
+          deletedRows.length === 0 &&
+          addedRows.size === 0
+        }>
+        Publish
+      </Button>
+    </div>
+  );
+
+  const renderPublishPopover = () => (
+    <Popover
+      open={Boolean(publishAnchorEl)}
+      anchorEl={publishAnchorEl}
+      onClose={handlePublishClose}
+      anchorOrigin={{vertical: 'bottom', horizontal: 'right'}}
+      transformOrigin={{vertical: 'top', horizontal: 'right'}}
+      sx={POPOVER_STYLES}>
+      <Box
+        sx={{
+          p: 3,
+          width: 300,
+          bgcolor: 'background.paper',
+          borderRadius: '8px',
+          '& .MuiTypography-root': {fontFamily: 'Source Sans Pro'},
+        }}>
+        <Typography sx={{mb: 2}}>
+          Publish changes to a new version of{' '}
+          <code style={CODE_STYLES}>{objectName}</code>?
+        </Typography>
+        <Box sx={{display: 'flex', justifyContent: 'flex-start', gap: 2}}>
+          <Button variant="ghost" onClick={handlePublishClose}>
+            Not yet
+          </Button>
+          <Button
+            icon="checkmark"
+            variant="primary"
+            onClick={() => {
+              handlePublish();
+              handlePublishClose();
+            }}>
+            Confirm
+          </Button>
+        </Box>
+      </Box>
+    </Popover>
+  );
+
   return (
     <SimplePageLayoutWithHeader
       title={
@@ -73,39 +301,56 @@ export const DatasetVersionPage: React.FC<{
       }
       headerContent={
         <Tailwind>
-          <div className="grid w-full auto-cols-max grid-flow-col gap-[16px] text-[14px]">
-            <div className="block">
-              <p className="text-moon-500">Name</p>
-              <ObjectVersionsLink
-                entity={entityName}
-                project={projectName}
-                filter={{objectName}}
-                versionCount={objectVersionCount}
-                neverPeek
-                variant="secondary">
-                <div className="group flex items-center font-semibold">
-                  <span>{objectName}</span>
-                  {objectVersions.loading ? (
-                    <LoadingDots />
-                  ) : (
-                    <span className="ml-[4px]">
-                      ({objectVersionCount} version
-                      {objectVersionCount !== 1 ? 's' : ''})
-                    </span>
-                  )}
-                  <Icon
-                    name="forward-next"
-                    width={16}
-                    height={16}
-                    className="ml-[2px] opacity-0 group-hover:opacity-100"
-                  />
-                </div>
-              </ObjectVersionsLink>
+          <div className="flex justify-between">
+            <div className="grid auto-cols-max grid-flow-col gap-[16px] text-[14px]">
+              <div className="block">
+                <p className="text-moon-500">Name</p>
+                <ObjectVersionsLink
+                  entity={entityName}
+                  project={projectName}
+                  filter={{objectName}}
+                  versionCount={objectVersionCount}
+                  neverPeek
+                  variant="secondary">
+                  <div className="group flex items-center font-semibold">
+                    <span>{objectName}</span>
+                    {objectVersions.loading ? (
+                      <LoadingDots />
+                    ) : (
+                      <span className="ml-[4px]">
+                        ({objectVersionCount} version
+                        {objectVersionCount !== 1 ? 's' : ''})
+                      </span>
+                    )}
+                    <Icon
+                      name="forward-next"
+                      width={16}
+                      height={16}
+                      className="ml-[2px] opacity-0 group-hover:opacity-100"
+                    />
+                  </div>
+                </ObjectVersionsLink>
+              </div>
+              <div className="block">
+                <p className="text-moon-500">Version</p>
+                <p>{objectVersionIndex}</p>
+              </div>
             </div>
-            <div className="block">
-              <p className="text-moon-500">Version</p>
-              <p>{objectVersionIndex}</p>
+            <div>
+              {isEditing ? (
+                renderEditingControls()
+              ) : (
+                <Button
+                  title="Edit"
+                  variant="secondary"
+                  size="medium"
+                  startIcon="pencil-edit"
+                  onClick={handleEditClick}>
+                  Edit dataset
+                </Button>
+              )}
             </div>
+            {renderPublishPopover()}
           </div>
         </Tailwind>
       }
@@ -114,22 +359,16 @@ export const DatasetVersionPage: React.FC<{
           label: 'Rows',
           content: (
             <ScrollableTabContent sx={{p: 0}}>
-              <Box
-                sx={{
-                  flex: '0 0 auto',
-                  height: '100%',
-                }}>
+              <Box sx={{flex: '0 0 auto', height: '100%'}}>
                 {data.loading ? (
                   <CenteredAnimatedLoader />
                 ) : (
                   <WeaveCHTableSourceRefContext.Provider value={refUri}>
                     <CustomWeaveTypeProjectContext.Provider
                       value={{entity: entityName, project: projectName}}>
-                      <ObjectViewerSection
-                        title=""
-                        data={viewerDataAsObject}
-                        noHide
-                        isExpanded
+                      <EditableDatasetView
+                        isEditing={isEditing}
+                        datasetObject={objectVersion.val}
                       />
                     </CustomWeaveTypeProjectContext.Provider>
                   </WeaveCHTableSourceRefContext.Provider>
diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/datasets/EditableDatasetView.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/datasets/EditableDatasetView.tsx
new file mode 100644
index 000000000000..cc16f001b13e
--- /dev/null
+++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/datasets/EditableDatasetView.tsx
@@ -0,0 +1,498 @@
+import {Box} from '@mui/material';
+import {
+  GridColDef,
+  GridPaginationModel,
+  GridRenderCellParams,
+  GridRenderEditCellParams,
+  GridRowModel,
+  GridSortModel,
+} from '@mui/x-data-grid-pro';
+import {A} from '@wandb/weave/common/util/links';
+import {Button} from '@wandb/weave/components/Button';
+import {RowId} from '@wandb/weave/components/PagePanelComponents/Home/Browse3/pages/CallPage/DataTableView';
+import {Tooltip} from '@wandb/weave/components/Tooltip';
+import React, {
+  FC,
+  useCallback,
+  useContext,
+  useEffect,
+  useMemo,
+  useState,
+} from 'react';
+import {useHistory} from 'react-router-dom';
+import {v4 as uuidv4} from 'uuid';
+
+import {isWeaveObjectRef, parseRef, parseRefMaybe} from '../../../../../react';
+import {flattenObjectPreservingWeaveTypes} from '../../Browse2/browse2Util';
+import {CellValue} from '../../Browse2/CellValue';
+import {useWeaveflowCurrentRouteContext} from '../context';
+import {WeaveCHTableSourceRefContext} from '../pages/CallPage/DataTableView';
+import {TABLE_ID_EDGE_NAME} from '../pages/wfReactInterface/constants';
+import {useWFHooks} from '../pages/wfReactInterface/context';
+import {SortBy} from '../pages/wfReactInterface/traceServerClientTypes';
+import {StyledDataGrid} from '../StyledDataGrid';
+import {CellEditingRenderer, CellViewingRenderer} from './CellRenderers';
+import {useDatasetEditContext} from './DatasetEditorContext';
+
+const ADDED_ROW_ID_PREFIX = 'new-';
+
+// Dataset object schema as it is stored in the database.
+interface DatasetObjectVal {
+  _type: 'Dataset';
+  name: string | null;
+  description: string | null;
+  rows: string;
+  _class_name: 'Dataset';
+  _bases: ['Object', 'BaseModel'];
+}
+
+interface EditableDataTableViewProps {
+  datasetObject: DatasetObjectVal;
+  isEditing: boolean;
+}
+
+interface ControlCellProps {
+  params: GridRenderCellParams;
+  deleteRow: (absoluteIndex: number) => void;
+  restoreRow: (absoluteIndex: number) => void;
+  deleteAddedRow: (rowId: string) => void;
+  isDeleted: boolean;
+  isNew: boolean;
+}
+
+interface ControlsColumnHeaderProps {
+  onAddRow: () => void;
+}
+
+const ControlCell: FC<ControlCellProps> = ({
+  params,
+  deleteRow,
+  restoreRow,
+  deleteAddedRow,
+  isDeleted,
+  isNew,
+}) => (
+  <Box
+    sx={{
+      display: 'flex',
+      alignItems: 'center',
+      justifyContent: 'center',
+      height: '100%',
+      width: '100%',
+      padding: '4px',
+    }}>
+    <Button
+      onClick={() => {
+        if (isNew) {
+          deleteAddedRow(params.row.___weave?.id);
+        } else if (isDeleted) {
+          restoreRow(params.row.___weave?.index);
+        } else {
+          deleteRow(params.row.___weave?.index);
+        }
+      }}
+      tooltip={isNew ? 'Remove' : isDeleted ? 'Restore' : 'Delete'}
+      icon={isNew ? 'close' : isDeleted ? 'undo' : 'delete'}
+      size="small"
+      variant="secondary"
+    />
+  </Box>
+);
+
+const ControlsColumnHeader: FC<ControlsColumnHeaderProps> = ({onAddRow}) => (
+  <Box
+    sx={{
+      margin: '2px',
+      display: 'flex',
+      justifyContent: 'center',
+      alignItems: 'center',
+      height: '100%',
+    }}>
+    <Button
+      icon="add-new"
+      onClick={onAddRow}
+      variant="ghost"
+      size="small"
+      tooltip="Add row"
+    />
+  </Box>
+);
+
+export const EditableDatasetView: FC<EditableDataTableViewProps> = ({
+  datasetObject,
+  isEditing,
+}) => {
+  const {useTableRowsQuery, useTableQueryStats} = useWFHooks();
+  const [sortBy, setSortBy] = useState<SortBy[]>([]);
+  const [sortModel, setSortModel] = useState<GridSortModel>([]);
+
+  const onSortModelChange = useCallback((model: GridSortModel) => {
+    setSortBy(
+      model.map(sort => ({
+        field: sort.field,
+        direction: sort.sort === 'asc' ? 'asc' : 'desc',
+      }))
+    );
+    setSortModel(model);
+  }, []);
+
+  const {
+    editedCellsMap,
+    processRowUpdate,
+    deletedRows,
+    setDeletedRows,
+    setAddedRows,
+    addedRows,
+  } = useDatasetEditContext();
+
+  const [paginationModel, setPaginationModel] = useState<GridPaginationModel>({
+    page: 0,
+    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();
+  const router = useWeaveflowCurrentRouteContext();
+  const onClick = useCallback(
+    val => {
+      const ref = parseRef(sharedRef!);
+      if (isWeaveObjectRef(ref)) {
+        const extra = 'attr/rows/' + TABLE_ID_EDGE_NAME + '/' + val;
+
+        const target = router.objectVersionUIUrl(
+          ref.entityName,
+          ref.projectName,
+          ref.artifactName,
+          ref.artifactVersion,
+          'obj',
+          extra
+        );
+        history.push(target);
+      }
+    },
+    [history, router, sharedRef]
+  );
+
+  const [initialFields, setInitialFields] = useState<string[]>([]);
+
+  const parsedRef = useMemo(
+    () => parseRefMaybe(datasetObject.rows),
+    [datasetObject.rows]
+  );
+
+  const lookupKey = useMemo(() => {
+    if (
+      parsedRef == null ||
+      !isWeaveObjectRef(parsedRef) ||
+      parsedRef.weaveKind !== 'table'
+    ) {
+      return null;
+    }
+    return {
+      entity: parsedRef.entityName,
+      project: parsedRef.projectName,
+      digest: parsedRef.artifactVersion,
+    };
+  }, [parsedRef]);
+
+  const numRowsQuery = useTableQueryStats(
+    lookupKey?.entity ?? '',
+    lookupKey?.project ?? '',
+    lookupKey?.digest ?? '',
+    {skip: lookupKey == null}
+  );
+
+  const numAddedRows = useMemo(
+    () => Array.from(addedRows.values()).length,
+    [addedRows]
+  );
+
+  const {numRowsToFetch, offset} = useMemo(() => {
+    const rowsToFetch =
+      numAddedRows <= paginationModel.page * paginationModel.pageSize
+        ? paginationModel.pageSize
+        : numAddedRows > (paginationModel.page + 1) * paginationModel.pageSize
+        ? 0
+        : paginationModel.pageSize - (numAddedRows % paginationModel.pageSize);
+
+    const offsetVal =
+      paginationModel.page * paginationModel.pageSize <= numAddedRows
+        ? 0
+        : paginationModel.page * paginationModel.pageSize - numAddedRows;
+
+    return {numRowsToFetch: rowsToFetch, offset: offsetVal};
+  }, [paginationModel, numAddedRows]);
+
+  const fetchQuery = useTableRowsQuery(
+    lookupKey?.entity ?? '',
+    lookupKey?.project ?? '',
+    lookupKey?.digest ?? '',
+    undefined,
+    numRowsToFetch,
+    offset,
+    sortBy,
+    {skip: lookupKey == null}
+  );
+
+  const [loadedRows, setLoadedRows] = useState<Array<{[key: string]: any}>>([]);
+  const [fetchQueryLoaded, setFetchQueryLoaded] = useState(false);
+
+  useEffect(() => {
+    if (!fetchQuery.loading) {
+      if (fetchQuery.result) {
+        setLoadedRows(fetchQuery.result.rows);
+      }
+      setFetchQueryLoaded(true);
+    }
+  }, [fetchQuery.loading, fetchQuery.result]);
+
+  const restoreRow = useCallback(
+    (absoluteIndex: number) => {
+      setDeletedRows(prev => prev.filter(index => index !== absoluteIndex));
+    },
+    [setDeletedRows]
+  );
+
+  const deleteRow = useCallback(
+    (absoluteIndex: number) => {
+      const rowKey = `${ADDED_ROW_ID_PREFIX}${absoluteIndex}`;
+      if (addedRows.has(rowKey)) {
+        setAddedRows(prev => {
+          const updatedMap = new Map(prev);
+          updatedMap.delete(rowKey);
+          return updatedMap;
+        });
+      } else {
+        setDeletedRows(prev => [...prev, absoluteIndex]);
+      }
+    },
+    [setDeletedRows, setAddedRows, addedRows]
+  );
+
+  const deleteAddedRow = useCallback(
+    (rowId: string) => {
+      setAddedRows(prev => {
+        const updatedMap = new Map(prev);
+        updatedMap.delete(rowId);
+        return updatedMap;
+      });
+    },
+    [setAddedRows]
+  );
+
+  const handleAddRowsClick = useCallback(() => {
+    setPaginationModel(prev => ({...prev, page: 0}));
+    setAddedRows(prev => {
+      const updatedMap = new Map(prev);
+      const newId = `${ADDED_ROW_ID_PREFIX}${uuidv4()}`;
+      console.log(initialFields);
+      const newRow = {
+        ___weave: {
+          id: newId,
+          isNew: true,
+        },
+        ...Object.fromEntries(initialFields.map(field => [field, ''])),
+      };
+      updatedMap.set(newId, newRow);
+      return updatedMap;
+    });
+  }, [setAddedRows, initialFields]);
+
+  const rows = useMemo(() => {
+    if (fetchQueryLoaded) {
+      return loadedRows.map((row, i) => {
+        const digest = row.digest;
+        const absoluteIndex =
+          i + paginationModel.pageSize * paginationModel.page;
+        const editedRow = editedCellsMap.get(absoluteIndex);
+        const value = flattenObjectPreservingWeaveTypes(row.val);
+        return {
+          ___weave: {
+            id: `${digest}_${absoluteIndex}`,
+            index: absoluteIndex,
+            isNew: false,
+          },
+          ...(editedRow ? {...value, ...editedRow} : value),
+        };
+      });
+    }
+    return [];
+  }, [loadedRows, fetchQueryLoaded, editedCellsMap, paginationModel]);
+
+  const combinedRows = useMemo(() => {
+    if (
+      !isEditing ||
+      numAddedRows <= paginationModel.page * paginationModel.pageSize
+    ) {
+      return rows;
+    }
+    const startIndex = paginationModel.page * paginationModel.pageSize;
+    const endIndex = startIndex + paginationModel.pageSize;
+    const displayedAddedRows = Array.from(addedRows.values()).slice(
+      startIndex,
+      endIndex
+    );
+    return [...displayedAddedRows, ...rows];
+  }, [rows, addedRows, numAddedRows, paginationModel, isEditing]);
+
+  const columns = useMemo(() => {
+    const allFields = combinedRows.reduce((acc, row) => {
+      Object.keys(row)
+        .filter(key => key !== '___weave')
+        .forEach(key => acc.add(key));
+      return acc;
+    }, new Set<string>());
+
+    if (initialFields.length === 0 && allFields.size > 0) {
+      setInitialFields(Array.from(allFields));
+    }
+
+    const baseColumns: GridColDef[] = [
+      {
+        field: '_row_click',
+        headerName: 'id',
+        sortable: false,
+        disableColumnMenu: true,
+        width: 50,
+        renderCell: params => {
+          const rowId = params.id as string;
+          if (isEditing && params.row.___weave?.isNew) {
+            return null;
+          }
+          const digestStr = rowId.split('_')[0];
+          const rowLabel = digestStr ? digestStr.slice(-4) : rowId;
+          const rowSpan = (
+            <Tooltip trigger={<RowId>{rowLabel}</RowId>} content={digestStr} />
+          );
+          return (
+            <Box sx={{marginLeft: '8px'}}>
+              <A onClick={() => onClick(rowId)}>{rowSpan}</A>
+            </Box>
+          );
+        },
+      },
+      ...(isEditing
+        ? [
+            {
+              field: 'controls',
+              headerName: '',
+              width: 48,
+              sortable: false,
+              filterable: false,
+              editable: false,
+              renderHeader: () => (
+                <ControlsColumnHeader onAddRow={handleAddRowsClick} />
+              ),
+              renderCell: (params: GridRenderCellParams) => (
+                <ControlCell
+                  params={params}
+                  deleteRow={deleteRow}
+                  deleteAddedRow={deleteAddedRow}
+                  restoreRow={restoreRow}
+                  isDeleted={deletedRows.includes(params.row.___weave?.index)}
+                  isNew={params.row.___weave?.isNew}
+                />
+              ),
+            },
+          ]
+        : []),
+    ];
+
+    const fieldColumns: GridColDef[] = Array.from(allFields).map(field => ({
+      field: field as string,
+      headerName: field as string,
+      flex: 1,
+      editable: isEditing,
+      sortable: !isEditing,
+      filterable: false,
+      renderCell: (params: GridRenderCellParams) => {
+        const editedRow = editedCellsMap.get(params.row.___weave?.index);
+        if (!isEditing) {
+          return (
+            <Box sx={{marginLeft: '8px'}}>
+              <CellValue value={params.value} />
+            </Box>
+          );
+        }
+        return (
+          <CellViewingRenderer
+            {...params}
+            isEdited={editedRow && editedRow[field as string] !== undefined}
+            isDeleted={deletedRows.includes(params.row.___weave?.index)}
+            isNew={params.row.___weave?.isNew}
+          />
+        );
+      },
+      renderEditCell: (params: GridRenderEditCellParams) => (
+        <CellEditingRenderer {...params} />
+      ),
+    }));
+
+    return [...baseColumns, ...fieldColumns];
+  }, [
+    combinedRows,
+    editedCellsMap,
+    deleteRow,
+    restoreRow,
+    deletedRows,
+    deleteAddedRow,
+    handleAddRowsClick,
+    initialFields,
+    isEditing,
+    onClick,
+  ]);
+
+  return (
+    <div style={{height: '100%', overflow: 'hidden'}}>
+      <StyledDataGrid
+        initialState={{
+          pinnedColumns: {
+            right: ['controls'],
+          },
+        }}
+        disableColumnMenu={true}
+        density="compact"
+        rows={combinedRows}
+        columns={columns}
+        sortingMode="server"
+        sortModel={sortModel}
+        onSortModelChange={onSortModelChange}
+        editMode="cell"
+        pagination
+        paginationMode="server"
+        paginationModel={paginationModel}
+        onPaginationModelChange={setPaginationModel}
+        rowCount={
+          (numRowsQuery.result?.count ?? 0) + (isEditing ? numAddedRows : 0)
+        }
+        disableMultipleColumnsSorting
+        loading={!fetchQueryLoaded}
+        disableRowSelectionOnClick
+        keepBorders={false}
+        pageSizeOptions={[]}
+        sx={{
+          border: 'none',
+          '& .MuiDataGrid-cell': {
+            padding: '0',
+          },
+          '& .MuiDataGrid-columnHeaders': {
+            borderBottom: '1px solid rgba(224, 224, 224, 1)',
+            marginBottom: '-1px', // offset the border
+          },
+        }}
+        processRowUpdate={processRowUpdate}
+        getRowId={(row: GridRowModel) => row.___weave?.id ?? row.id}
+      />
+    </div>
+  );
+};
diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/DataTableView.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/DataTableView.tsx
index 800b6a0fdf2b..ac8f56a05c46 100644
--- a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/DataTableView.tsx
+++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/DataTableView.tsx
@@ -48,7 +48,7 @@ import {TABLE_ID_EDGE_NAME} from '../wfReactInterface/constants';
 import {useWFHooks} from '../wfReactInterface/context';
 import {SortBy} from '../wfReactInterface/traceServerClientTypes';
 
-const RowId = styled.span`
+export const RowId = styled.span`
   font-family: 'Inconsolata', monospace;
 `;
 RowId.displayName = 'S.RowId';
@@ -238,7 +238,7 @@ export const WeaveCHTable: FC<{
   );
 };
 
-type DataTableServerSidePaginationControls = {
+export type DataTableServerSidePaginationControls = {
   paginationModel: GridPaginationModel;
   onPaginationModelChange: (model: GridPaginationModel) => void;
   totalRows: number;
diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/ObjectVersionPage.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/ObjectVersionPage.tsx
index 1b0bb491f1c0..007856acf066 100644
--- a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/ObjectVersionPage.tsx
+++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/ObjectVersionPage.tsx
@@ -7,6 +7,7 @@ import {Icon, IconName} from '../../../../Icon';
 import {LoadingDots} from '../../../../LoadingDots';
 import {Tailwind} from '../../../../Tailwind';
 import {Tooltip} from '../../../../Tooltip';
+import {DatasetEditProvider} from '../datasets/DatasetEditorContext';
 import {DatasetVersionPage} from '../datasets/DatasetVersionPage';
 import {NotFoundPanel} from '../NotFoundPanel';
 import {CustomWeaveTypeProjectContext} from '../typeViews/CustomWeaveTypeDispatcher';
@@ -199,7 +200,11 @@ const ObjectVersionPageInner: React.FC<{
   }
 
   if (isDataset) {
-    return <DatasetVersionPage objectVersion={objectVersion} />;
+    return (
+      <DatasetEditProvider>
+        <DatasetVersionPage objectVersion={objectVersion} refExtra={refExtra} />
+      </DatasetEditProvider>
+    );
   }
 
   return (