refactor(weave): dedicated component for dataset versions

wandb · Jan 7, 2025 · cdae1e8 · cdae1e8
1 parent 0152b4f
commit cdae1e8
Show file tree

Hide file tree

Showing 2 changed files with 152 additions and 9 deletions.
diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/datasets/DatasetVersionPage.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/datasets/DatasetVersionPage.tsx
@@ -0,0 +1,145 @@
+import Box from '@mui/material/Box';
+import React, {useMemo} from 'react';
+
+import {Icon} from '../../../../Icon';
+import {LoadingDots} from '../../../../LoadingDots';
+import {Tailwind} from '../../../../Tailwind';
+import {WeaveCHTableSourceRefContext} from '../pages/CallPage/DataTableView';
+import {ObjectViewerSection} from '../pages/CallPage/ObjectViewerSection';
+import {objectVersionText} from '../pages/common/Links';
+import {CenteredAnimatedLoader} from '../pages/common/Loader';
+import {
+  ScrollableTabContent,
+  SimplePageLayoutWithHeader,
+} from '../pages/common/SimplePageLayout';
+import {TabUseDataset} from '../pages/TabUseDataset';
+import {useWFHooks} from '../pages/wfReactInterface/context';
+import {objectVersionKeyToRefUri} from '../pages/wfReactInterface/utilities';
+import {ObjectVersionSchema} from '../pages/wfReactInterface/wfDataModelHooksInterface';
+import {CustomWeaveTypeProjectContext} from '../typeViews/CustomWeaveTypeDispatcher';
+
+export const DatasetVersionPage: React.FC<{
+  objectVersion: ObjectVersionSchema;
+}> = ({objectVersion}) => {
+  const {useRootObjectVersions, useRefsData} = useWFHooks();
+  const entityName = objectVersion.entity;
+  const projectName = objectVersion.project;
+  const objectName = objectVersion.objectId;
+  const objectVersionIndex = objectVersion.versionIndex;
+
+  const objectVersions = useRootObjectVersions(
+    entityName,
+    projectName,
+    {
+      objectIds: [objectName],
+    },
+    undefined,
+    true
+  );
+  const objectVersionCount = (objectVersions.result ?? []).length;
+  const refUri = objectVersionKeyToRefUri(objectVersion);
+
+  const data = useRefsData([refUri]);
+  const viewerData = useMemo(() => {
+    if (data.loading) {
+      return {};
+    }
+    return data.result?.[0] ?? {};
+  }, [data.loading, data.result]);
+
+  const viewerDataAsObject = useMemo(() => {
+    const dataIsPrimitive =
+      typeof viewerData !== 'object' ||
+      viewerData === null ||
+      Array.isArray(viewerData);
+    if (dataIsPrimitive) {
+      return {_result: viewerData};
+    }
+    return viewerData;
+  }, [viewerData]);
+
+  return (
+    <SimplePageLayoutWithHeader
+      title={
+        <Tailwind>
+          <div className="flex items-center gap-8">
+            <div className="flex h-22 w-22 items-center justify-center rounded-full bg-moon-300/[0.48] text-moon-600">
+              <Icon width={14} height={14} name="table" />
+            </div>
+            {objectVersionText(objectName, objectVersionIndex)}
+          </div>
+        </Tailwind>
+      }
+      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>
+              <div className="flex items-center">
+                <div className="group flex items-center font-semibold">
+                  <span>{objectName}</span>
+                  {objectVersions.loading ? (
+                    <LoadingDots />
+                  ) : (
+                    <span className="ml-[4px]">
+                      ({objectVersionCount} version
+                      {objectVersionCount !== 1 ? 's' : ''})
+                    </span>
+                  )}
+                </div>
+              </div>
+            </div>
+            <div className="block">
+              <p className="text-moon-500">Version</p>
+              <p>{objectVersionIndex}</p>
+            </div>
+          </div>
+        </Tailwind>
+      }
+      tabs={[
+        {
+          label: 'Rows',
+          content: (
+            <ScrollableTabContent sx={{p: 0}}>
+              <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
+                      />
+                    </CustomWeaveTypeProjectContext.Provider>
+                  </WeaveCHTableSourceRefContext.Provider>
+                )}
+              </Box>
+            </ScrollableTabContent>
+          ),
+        },
+        {
+          label: 'Use',
+          content: (
+            <ScrollableTabContent>
+              <Tailwind>
+                <TabUseDataset
+                  name={objectName}
+                  uri={refUri}
+                  versionIndex={objectVersionIndex}
+                />
+              </Tailwind>
+            </ScrollableTabContent>
+          ),
+        },
+      ]}
+    />
+  );
+};
diff --git 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 {DatasetVersionPage} from '../datasets/DatasetVersionPage';
 import {NotFoundPanel} from '../NotFoundPanel';
 import {CustomWeaveTypeProjectContext} from '../typeViews/CustomWeaveTypeDispatcher';
 import {WeaveCHTableSourceRefContext} from './CallPage/DataTableView';
@@ -27,7 +28,6 @@ import {
 } from './common/SimplePageLayout';
 import {EvaluationLeaderboardTab} from './LeaderboardTab';
 import {TabPrompt} from './TabPrompt';
-import {TabUseDataset} from './TabUseDataset';
 import {TabUseModel} from './TabUseModel';
 import {TabUseObject} from './TabUseObject';
 import {TabUsePrompt} from './TabUsePrompt';
@@ -198,6 +198,10 @@ const ObjectVersionPageInner: React.FC<{
     return <CenteredAnimatedLoader />;
   }
 
+  if (isDataset) {
+    return <DatasetVersionPage objectVersion={objectVersion} />;
+  }
+
   return (
     <SimplePageLayoutWithHeader
       title={
@@ -319,9 +323,9 @@ const ObjectVersionPageInner: React.FC<{
             ]
           : []),
         {
-          label: isDataset ? 'Rows' : 'Values',
+          label: 'Values',
           content: (
-            <ScrollableTabContent sx={isDataset ? {p: 0} : {}}>
+            <ScrollableTabContent>
               <Box
                 sx={{
                   flex: '0 0 auto',
@@ -359,12 +363,6 @@ const ObjectVersionPageInner: React.FC<{
                     projectName={projectName}
                     data={viewerDataAsObject}
                   />
-                ) : baseObjectClass === 'Dataset' ? (
-                  <TabUseDataset
-                    name={objectName}
-                    uri={refUri}
-                    versionIndex={objectVersionIndex}
-                  />
                 ) : baseObjectClass === 'Model' ? (
                   <TabUseModel
                     name={objectName}