From d18bad1933560cdd7d82c43672ef088140ae1ae7 Mon Sep 17 00:00:00 2001
From: Tim Sweeney <tim@wandb.com>
Date: Thu, 19 Dec 2024 15:24:38 -0800
Subject: [PATCH] Checkpoint: fast(er) initial load

---
 .../CompareEvaluationsPage.tsx                | 19 ++++++++++++++++---
 .../pages/CompareEvaluationsPage/ecpState.ts  |  6 +++---
 .../exampleCompareSectionUtil.ts              |  8 ++++----
 .../ExampleFilterSection.tsx                  |  4 ++--
 4 files changed, 25 insertions(+), 12 deletions(-)

diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CompareEvaluationsPage/CompareEvaluationsPage.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CompareEvaluationsPage/CompareEvaluationsPage.tsx
index 6679f3009a3..448a0619601 100644
--- a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CompareEvaluationsPage/CompareEvaluationsPage.tsx
+++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CompareEvaluationsPage/CompareEvaluationsPage.tsx
@@ -4,6 +4,7 @@
 
 import {Box} from '@material-ui/core';
 import {Alert} from '@mui/material';
+import {WaveLoader} from '@wandb/weave/components/Loaders/WaveLoader';
 import {Tailwind} from '@wandb/weave/components/Tailwind';
 import {maybePluralizeWord} from '@wandb/weave/core/util/string';
 import React, {FC, useCallback, useContext, useMemo, useState} from 'react';
@@ -178,10 +179,11 @@ const CompareEvaluationsPageInner: React.FC<{
   height: number;
 }> = props => {
   const {state, setSelectedMetrics} = useCompareEvaluationsState();
-  console.log(state);
   const showExampleFilter =
     Object.keys(state.summary.evaluationCalls).length === 2;
-  const showExamples = Object.keys(state.results?.resultRows ?? {}).length > 0;
+  const showExamples =
+    Object.keys(state.results.result?.resultRows ?? {}).length > 0;
+  const resultsLoading = state.results.loading;
   return (
     <Box
       sx={{
@@ -201,7 +203,18 @@ const CompareEvaluationsPageInner: React.FC<{
         <ComparisonDefinitionSection state={state} />
         <SummaryPlots state={state} setSelectedMetrics={setSelectedMetrics} />
         <ScorecardSection state={state} />
-        {showExamples ? (
+        {resultsLoading ? (
+          <Box
+            sx={{
+              width: '100%',
+              display: 'flex',
+              justifyContent: 'center',
+              alignItems: 'center',
+              height: '50px',
+            }}>
+            <WaveLoader size="small" />
+          </Box>
+        ) : showExamples ? (
           <>
             {showExampleFilter && <ExampleFilterSection state={state} />}
             <ResultExplorer state={state} height={props.height} />
diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CompareEvaluationsPage/ecpState.ts b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CompareEvaluationsPage/ecpState.ts
index f7f5cca0768..246876cb444 100644
--- a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CompareEvaluationsPage/ecpState.ts
+++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CompareEvaluationsPage/ecpState.ts
@@ -25,7 +25,7 @@ export type EvaluationComparisonState = {
   // The normalized data for the evaluations
   summary: EvaluationComparisonSummary;
   // The results of the evaluations
-  results: EvaluationComparisonResults | null;
+  results: Loadable<EvaluationComparisonResults>;
   // The dimensions to compare & filter results
   comparisonDimensions?: ComparisonDimensionsType;
   // The current digest which is in view
@@ -112,7 +112,7 @@ export const useEvaluationComparisonState = (
       loading: false,
       result: {
         summary: summaryData.result,
-        results: resultsData.result,
+        results: resultsData,
         comparisonDimensions: newComparisonDimensions,
         selectedInputDigest,
         selectedMetrics,
@@ -123,7 +123,7 @@ export const useEvaluationComparisonState = (
     summaryData.result,
     summaryData.loading,
     comparisonDimensions,
-    resultsData.result,
+    resultsData,
     selectedInputDigest,
     selectedMetrics,
     evaluationCallIds,
diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CompareEvaluationsPage/sections/ExampleCompareSection/exampleCompareSectionUtil.ts b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CompareEvaluationsPage/sections/ExampleCompareSection/exampleCompareSectionUtil.ts
index 76f9ceff775..0cfec0b6f69 100644
--- a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CompareEvaluationsPage/sections/ExampleCompareSection/exampleCompareSectionUtil.ts
+++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CompareEvaluationsPage/sections/ExampleCompareSection/exampleCompareSectionUtil.ts
@@ -124,13 +124,13 @@ export const useFilteredAggregateRows = (state: EvaluationComparisonState) => {
 
   const flattenedRows = useMemo(() => {
     const rows: FlattenedRow[] = [];
-    Object.entries(state.results?.resultRows ?? {}).forEach(
+    Object.entries(state.results.result?.resultRows ?? {}).forEach(
       ([rowDigest, rowCollection]) => {
         Object.values(rowCollection.evaluations).forEach(modelCollection => {
           Object.values(modelCollection.predictAndScores).forEach(
             predictAndScoreRes => {
               const datasetRow =
-                state.results?.inputs[predictAndScoreRes.rowDigest];
+                state.results.result?.inputs[predictAndScoreRes.rowDigest];
               if (datasetRow != null) {
                 const output = predictAndScoreRes._rawPredictTraceData?.output;
                 rows.push({
@@ -170,8 +170,8 @@ export const useFilteredAggregateRows = (state: EvaluationComparisonState) => {
     );
     return rows;
   }, [
-    state.results?.resultRows,
-    state.results?.inputs,
+    state.results.result?.resultRows,
+    state.results.result?.inputs,
     state.summary.scoreMetrics,
   ]);
 
diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CompareEvaluationsPage/sections/ExampleFilterSection/ExampleFilterSection.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CompareEvaluationsPage/sections/ExampleFilterSection/ExampleFilterSection.tsx
index 1c819c4ee16..13d590c35da 100644
--- a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CompareEvaluationsPage/sections/ExampleFilterSection/ExampleFilterSection.tsx
+++ b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CompareEvaluationsPage/sections/ExampleFilterSection/ExampleFilterSection.tsx
@@ -141,7 +141,7 @@ const SingleDimensionFilter: React.FC<{
       );
 
       if (baselineTargetDimension != null && compareTargetDimension != null) {
-        Object.entries(props.state.results?.resultRows ?? {}).forEach(
+        Object.entries(props.state.results.result?.resultRows ?? {}).forEach(
           ([digest, row]) => {
             const xVals: number[] = [];
             const yVals: number[] = [];
@@ -232,7 +232,7 @@ const SingleDimensionFilter: React.FC<{
     compareCallId,
     compositeMetricsMap,
     filteredDigest,
-    props.state.results?.resultRows,
+    props.state.results.result?.resultRows,
     targetDimension,
   ]);