diff --git a/apps/opik-frontend/src/v2/pages/PlaygroundPage/PlaygroundOutputs/PlaygroundOutputTable/PlaygroundOutputCell.tsx b/apps/opik-frontend/src/v2/pages/PlaygroundPage/PlaygroundOutputs/PlaygroundOutputTable/PlaygroundOutputCell.tsx index 6edf74bb1d0..0563290bb68 100644 --- a/apps/opik-frontend/src/v2/pages/PlaygroundPage/PlaygroundOutputs/PlaygroundOutputTable/PlaygroundOutputCell.tsx +++ b/apps/opik-frontend/src/v2/pages/PlaygroundPage/PlaygroundOutputs/PlaygroundOutputTable/PlaygroundOutputCell.tsx @@ -4,11 +4,7 @@ import { ListTree } from "lucide-react"; import CellWrapper from "@/shared/DataTableCells/CellWrapper"; import { - useOutputLoadingByPromptDatasetItemId, - useOutputStaleStatusByPromptDatasetItemId, - useOutputValueByPromptDatasetItemId, - useSelectedRuleIdsByPromptDatasetItemId, - useTraceIdByPromptDatasetItemId, + useOutputByPromptDatasetItemId, useDatasetType, useExperimentIdByPromptId, } from "@/store/PlaygroundStore"; @@ -46,30 +42,23 @@ const PlaygroundOutputCell: React.FunctionComponent< const workspaceName = useAppStore((state) => state.activeWorkspaceName); - const value = useOutputValueByPromptDatasetItemId( - promptId, - originalRow.dataItemId, - ); - - const isLoading = useOutputLoadingByPromptDatasetItemId( - promptId, - originalRow.dataItemId, - ); - - const stale = useOutputStaleStatusByPromptDatasetItemId( - promptId, - originalRow.dataItemId, - ); - - const traceId = useTraceIdByPromptDatasetItemId( - promptId, - originalRow.dataItemId, - ); - - const selectedRuleIds = useSelectedRuleIdsByPromptDatasetItemId( + // One store subscription per cell, not five. Each of the fields below used to + // come from its own hook, and every one of those hooks runs the *same* + // selector, so a single streaming token re-ran it (dataset items x prompts x 5) + // times. Reading the output object once and picking the fields off it is + // equivalent — the defaults below are the ones those hooks applied — while + // cutting the per-token selector work by 5x. The selector returns the stored + // object reference, which only changes when this cell's own output changes, so + // unrelated updates still don't re-render this cell. + const output = useOutputByPromptDatasetItemId( promptId, originalRow.dataItemId, ); + const value = output?.value ?? null; + const isLoading = output?.isLoading ?? false; + const stale = output?.stale ?? false; + const traceId = output?.traceId ?? null; + const selectedRuleIds = output?.selectedRuleIds; const datasetType = useDatasetType(); const experimentId = useExperimentIdByPromptId(promptId); diff --git a/apps/opik-frontend/src/v2/pages/PlaygroundPage/useIncrementalDatasetHydration.ts b/apps/opik-frontend/src/v2/pages/PlaygroundPage/useIncrementalDatasetHydration.ts index 4a4161c7b72..8e6943d4054 100644 --- a/apps/opik-frontend/src/v2/pages/PlaygroundPage/useIncrementalDatasetHydration.ts +++ b/apps/opik-frontend/src/v2/pages/PlaygroundPage/useIncrementalDatasetHydration.ts @@ -1,6 +1,7 @@ -import { useEffect, useRef, useState } from "react"; +import { useEffect, useState } from "react"; import { DatasetItem } from "@/types/datasets"; import { useHydrateDatasetItemData } from "@/v2/pages/PlaygroundPage/useHydrateDatasetItemData"; +import { containsTruncatedMedia } from "@/lib/media"; export function useIncrementalDatasetHydration(datasetItems: DatasetItem[]): { hydratedItems: DatasetItem[]; @@ -9,10 +10,16 @@ export function useIncrementalDatasetHydration(datasetItems: DatasetItem[]): { const hydrateDatasetItemData = useHydrateDatasetItemData(); const [hydratedItems, setHydratedItems] = useState([]); const [isHydrating, setIsHydrating] = useState(false); - const cancelledRef = useRef(false); useEffect(() => { - cancelledRef.current = false; + // Scoped to this effect run rather than a shared ref. React runs the previous + // run's cleanup before this body, so an in-flight hydration from an earlier + // dataset observes its own `cancelled === true` and stops — even when this run + // returns early below and installs no cleanup of its own. A shared ref instead + // got reset here on every run, un-cancelling the previous run's loop; its + // continuation would then write the old dataset's data into the new array at a + // positional index. Also covers unmount. + let cancelled = false; if (datasetItems.length === 0) { setHydratedItems([]); @@ -21,19 +28,41 @@ export function useIncrementalDatasetHydration(datasetItems: DatasetItem[]): { } setHydratedItems(datasetItems); + + // Only items carrying truncated media need a round-trip — that is the sole + // condition under which hydrateDatasetItemData fetches anything. Selecting + // them up front means the loop below runs once per *media* item rather than + // once per row, and each pass rebuilt the whole array, so a 1000-row text + // dataset was doing ~1M array copies to arrive back at the data it started + // with. + const indexesToHydrate = datasetItems.reduce( + (acc, item, index) => { + if (containsTruncatedMedia(item.data)) { + acc.push(index); + } + return acc; + }, + [], + ); + + if (indexesToHydrate.length === 0) { + setIsHydrating(false); + return; + } + setIsHydrating(true); const hydrateItems = async () => { - for (let i = 0; i < datasetItems.length; i++) { - if (cancelledRef.current) return; + for (const index of indexesToHydrate) { + if (cancelled) return; - const hydratedData = await hydrateDatasetItemData(datasetItems[i]); + const hydratedData = await hydrateDatasetItemData(datasetItems[index]); - if (cancelledRef.current) return; + if (cancelled) return; setHydratedItems((prev) => prev.map((item, idx) => - idx === i ? { ...item, data: hydratedData } : item, + idx === index ? { ...item, data: hydratedData } : item, ), ); } @@ -44,7 +73,7 @@ export function useIncrementalDatasetHydration(datasetItems: DatasetItem[]): { hydrateItems(); return () => { - cancelledRef.current = true; + cancelled = true; }; }, [datasetItems, hydrateDatasetItemData]);