load the episode’s row range for v3 datasets (based on the file’s start index

src/app/[org]/[dataset]/[episode]/fetch-data.ts

@@ -748,7 +748,7 @@ async function loadEpisodeDataV3(
 
   try {
     const dataUrl = buildVersionedUrl(repoId, version, dataPath);
-    const arrayBuffer = await fetchParquetFile(dataUrl);
+    const parquetFile = await fetchParquetFile(dataUrl);
     const v3DataColumns = Array.from(
       new Set([
         "index",
@@ -779,32 +779,44 @@ async function loadEpisodeDataV3(
           .map(([key]) => key),
       ]),
     );
-    const fullData = await readParquetAsObjects(arrayBuffer, v3DataColumns);
-
     // Extract the episode-specific data slice
     const fromIndex = bigIntToNumber(episodeMetadata.dataset_from_index, 0);
-    const toIndex = bigIntToNumber(
-      episodeMetadata.dataset_to_index,
-      fullData.length,
-    );
+    let toIndex = bigIntToNumber(episodeMetadata.dataset_to_index, fromIndex);
+    if (toIndex <= fromIndex) {
+      toIndex = fromIndex + 1;
+    }
 
-    // Find the starting index of this parquet file by checking the first row's index
-    // This handles the case where episodes are split across multiple parquet files
-    let fileStartIndex = 0;
-    if (fullData.length > 0 && fullData[0].index !== undefined) {
-      fileStartIndex = Number(fullData[0].index);
+    let episodeRows: Record<string, unknown>[] = [];
+    let usedRowRange = false;
+
+    try {
+      const indexPreview = await readParquetAsObjects(parquetFile, ["index"], {
+        rowStart: 0,
+        rowEnd: 1,
+      });
+      const startIndexValue = indexPreview[0]?.index;
+      if (startIndexValue !== undefined && startIndexValue !== null) {
+        const fileStartIndex =
+          typeof startIndexValue === "bigint"
+            ? Number(startIndexValue)
+            : Number(startIndexValue);
+        const localFromIndex = Math.max(0, fromIndex - fileStartIndex);
+        const localToIndex = Math.max(localFromIndex, toIndex - fileStartIndex);
+        episodeRows = await readParquetAsObjects(parquetFile, v3DataColumns, {
+          rowStart: localFromIndex,
+          rowEnd: localToIndex,
+        });
+        usedRowRange = true;
+      }
+    } catch {
+      // Fall back to full reads if row-range selection fails.
     }
 
-    // Adjust indices to be relative to this file's starting position
-    const localFromIndex = Math.max(0, fromIndex - fileStartIndex);
-    const localToIndex = Math.min(fullData.length, toIndex - fileStartIndex);
+    if (!usedRowRange) {
+      episodeRows = await readParquetAsObjects(parquetFile, v3DataColumns);
+    }
 
-    const episodeData = sampleRowsFromRange(
-      fullData,
-      localFromIndex,
-      localToIndex,
-      MAX_EPISODE_POINTS,
-    );
+    const episodeData = evenlySampleArray(episodeRows, MAX_EPISODE_POINTS);
 
     if (episodeData.length === 0) {
       return {

src/utils/parquetUtils.ts

@@ -1,4 +1,10 @@
-import { parquetRead, parquetReadObjects } from "hyparquet";
+import {
+  asyncBufferFromUrl,
+  cachedAsyncBuffer,
+  parquetRead,
+  parquetReadObjects,
+  type AsyncBuffer,
+} from "hyparquet";
 
 export interface DatasetMetadata {
   codebase_version: string;
@@ -42,26 +48,36 @@ export function formatStringWithVars(
 }
 
 // Fetch and parse the Parquet file
-export async function fetchParquetFile(url: string): Promise<ArrayBuffer> {
-  const res = await fetch(url, { cache: "no-store" });
+type ParquetFile = ArrayBuffer | AsyncBuffer;
 
-  if (!res.ok) {
-    throw new Error(`Failed to fetch ${url}: ${res.status} ${res.statusText}`);
-  }
+const parquetFileCache = new Map<string, AsyncBuffer>();
 
-  return res.arrayBuffer();
+export async function fetchParquetFile(url: string): Promise<ParquetFile> {
+  const cached = parquetFileCache.get(url);
+  if (cached) return cached;
+
+  const file = await asyncBufferFromUrl({
+    url,
+    requestInit: { cache: "no-store" },
+  });
+  const wrapped = cachedAsyncBuffer(file);
+  parquetFileCache.set(url, wrapped);
+  return wrapped;
 }
 
 // Read specific columns from the Parquet file
 export async function readParquetColumn(
-  fileBuffer: ArrayBuffer,
+  fileBuffer: ParquetFile,
   columns: string[],
+  options?: { rowStart?: number; rowEnd?: number },
 ): Promise<unknown[][]> {
   return new Promise((resolve, reject) => {
     try {
       parquetRead({
         file: fileBuffer,
         columns: columns.length > 0 ? columns : undefined,
+        rowStart: options?.rowStart,
+        rowEnd: options?.rowEnd,
         onComplete: (data: unknown[][]) => {
           resolve(data);
         },
@@ -73,12 +89,15 @@ export async function readParquetColumn(
 }
 
 export async function readParquetAsObjects(
-  fileBuffer: ArrayBuffer,
+  fileBuffer: ParquetFile,
   columns: string[] = [],
+  options?: { rowStart?: number; rowEnd?: number },
 ): Promise<Record<string, unknown>[]> {
   return parquetReadObjects({
     file: fileBuffer,
     columns: columns.length > 0 ? columns : undefined,
+    rowStart: options?.rowStart,
+    rowEnd: options?.rowEnd,
   }) as Promise<Record<string, unknown>[]>;
 }