-
Version
-
{objectVersionIndex}
+
+
+
+
Name
+
+
+ {objectName}
+ {objectVersions.loading ? (
+
+ ) : (
+
+ ({objectVersionCount} version
+ {objectVersionCount !== 1 ? 's' : ''})
+
+ )}
+
+
+
+
+
+
Version
+
{objectVersionIndex}
+
- {showDeleteButton && (
-
+
+ {isEditing ? (
+ renderEditingControls()
+ ) : (
+
+ )}
+ {showDeleteButton && !isEditing && (
-
- )}
+ )}
+
+ {renderPublishPopover()}
}
@@ -121,22 +366,16 @@ export const DatasetVersionPage: React.FC<{
label: 'Rows',
content: (
-
+
{data.loading ? (
) : (
-
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..3a3c63e02059
--- /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 = ({
+ params,
+ deleteRow,
+ restoreRow,
+ deleteAddedRow,
+ isDeleted,
+ isNew,
+}) => (
+
+
+);
+
+const ControlsColumnHeader: FC = ({onAddRow}) => (
+
+
+
+);
+
+export const EditableDatasetView: FC = ({
+ datasetObject,
+ isEditing,
+}) => {
+ const {useTableRowsQuery, useTableQueryStats} = useWFHooks();
+ const [sortBy, setSortBy] = useState([]);
+ const [sortModel, setSortModel] = useState([]);
+
+ 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({
+ 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([]);
+
+ 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>([]);
+ 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());
+
+ 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 = (
+ {rowLabel}} content={digestStr} />
+ );
+ return (
+
+ onClick(rowId)}>{rowSpan}
+
+ );
+ },
+ },
+ ...(isEditing
+ ? [
+ {
+ field: 'controls',
+ headerName: '',
+ width: 48,
+ sortable: false,
+ filterable: false,
+ editable: false,
+ renderHeader: () => (
+
+ ),
+ renderCell: (params: GridRenderCellParams) => (
+
+ ),
+ },
+ ]
+ : []),
+ ];
+
+ 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 (
+
+
+
+ );
+ }
+ return (
+
+ );
+ },
+ renderEditCell: (params: GridRenderEditCellParams) => (
+
+ ),
+ }));
+
+ return [...baseColumns, ...fieldColumns];
+ }, [
+ combinedRows,
+ editedCellsMap,
+ deleteRow,
+ restoreRow,
+ deletedRows,
+ deleteAddedRow,
+ handleAddRowsClick,
+ initialFields,
+ isEditing,
+ onClick,
+ ]);
+
+ return (
+
+ row.___weave?.id ?? row.id}
+ />
+
+ );
+};
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 d1a2543a63bc..71afee94a208 100644
--- a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/ObjectVersionPage.tsx
+++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/ObjectVersionPage.tsx
@@ -9,6 +9,7 @@ import {LoadingDots} from '../../../../LoadingDots';
import {Tailwind} from '../../../../Tailwind';
import {Tooltip} from '../../../../Tooltip';
import {useClosePeek} from '../context';
+import {DatasetEditProvider} from '../datasets/DatasetEditorContext';
import {DatasetVersionPage} from '../datasets/DatasetVersionPage';
import {NotFoundPanel} from '../NotFoundPanel';
import {CustomWeaveTypeProjectContext} from '../typeViews/CustomWeaveTypeDispatcher';
@@ -209,10 +210,13 @@ const ObjectVersionPageInner: React.FC<{
if (isDataset) {
return (
-
+
+
+
);
}