refactor: extract countColumnRows helper + QueryModeError in fragment-do

- types.ts: new countColumnRows(columns) helper replaces 7 duplicated
  `columns[0]?.pages.reduce((s, p) => s + p.rowCount, 0) ?? 0` patterns
- local-executor, query-do, master-do: all 7 call sites use the helper
- fragment-do: 2 plain Error throws → QueryModeError (MEMORY_EXCEEDED, QUERY_FAILED)
- master-do: 1 missed plain Error → QueryModeError (SCHEMA_MISMATCH)

Author: teamchong

src/fragment-do.ts

@@ -11,6 +11,7 @@ import {
   buildKeptPageIndices,
 } from "./operators.js";
 import { mergeQueryResults } from "./merge.js";
+import { QueryModeError } from "./errors.js";
 import { resolveBucket } from "./bucket.js";
 import { concatQMCBBatches, decodeColumnarBatch, columnarBatchToRows } from "./columnar.js";
 import wasmModule from "./wasm-module.js";
@@ -187,12 +188,12 @@ export class FragmentDO extends DurableObject<Env> {
           pageInfos: keptPageIndices.map(pi => col.pages[pi]).filter(Boolean),
         }));
       if (!this.wasmEngine.registerColumns(fragTable, fragColEntries)) {
-        throw new Error(`WASM OOM: failed to register columns for fragment "${r2Key}"`);
+        throw new QueryModeError("MEMORY_EXCEEDED", `WASM OOM: failed to register columns for fragment "${r2Key}"`);
       }
 
       const fragQuery = { ...query, table: fragTable };
       const qmcb = this.wasmEngine.executeQueryColumnar(fragQuery);
-      if (!qmcb) throw new Error(`WASM query execution failed for fragment "${r2Key}"`);
+      if (!qmcb) throw new QueryModeError("QUERY_FAILED", `WASM query execution failed for fragment "${r2Key}"`);
       this.wasmEngine.clearTable(fragTable);
       allBatches.push(qmcb);
       totalWasmExecMs += Date.now() - wasmStart;

src/index.ts

@@ -1,7 +1,7 @@
 import { DataFrame, TableQuery } from "./client.js";
 import type { QueryDescriptor, QueryExecutor } from "./client.js";
 import type { AppendOptions, AppendResult, DropResult, ExplainResult, QueryResult, Row, QueryDORpc, MasterDORpc } from "./types.js";
-import { NULL_SENTINEL } from "./types.js";
+import { NULL_SENTINEL, countColumnRows } from "./types.js";
 import { LocalExecutor } from "./local-executor.js";
 import { createFromJSON, createFromCSV, createDemo } from "./convenience.js";
 import { sqlToDescriptor, buildSqlDataFrame } from "./sql/index.js";

src/local-executor.ts

@@ -7,7 +7,7 @@
  */
 import type { QueryDescriptor, QueryExecutor } from "./client.js";
 import type { AppendResult, ColumnMeta, DataType, DiffResult, ExplainResult, PageInfo, QueryResult, Row, TableMeta, DatasetMeta, VersionInfo } from "./types.js";
-import { queryReferencedColumns, queryCacheKey, NULL_SENTINEL } from "./types.js";
+import { queryReferencedColumns, queryCacheKey, countColumnRows, NULL_SENTINEL } from "./types.js";
 import { parseFooter, parseColumnMetaFromProtobuf, FOOTER_SIZE } from "./footer.js";
 import { parseManifest } from "./manifest.js";
 import { detectFormat, getParquetFooterLength, parseParquetFooter, parquetMetaToTableMeta } from "./parquet.js";
@@ -149,7 +149,7 @@ export class LocalExecutor implements QueryExecutor {
   async count(query: QueryDescriptor): Promise<number> {
     const meta = await this.getOrLoadMeta(query.table);
     if (query.filters.length === 0 && !query.filterGroups?.length && !query.aggregates?.length && !query.groupBy?.length && !query.distinct && !query.join && !query.vectorSearch && !query.setOperation && !query.subqueryIn && !query.computedColumns?.length) {
-      return meta.columns[0]?.pages.reduce((s, p) => s + p.rowCount, 0) ?? 0;
+      return countColumnRows(meta.columns);
     }
     // With filters: fall through to aggregate path
     const desc = { ...query, aggregates: [{ fn: "count" as const, column: "*" }] };
@@ -206,7 +206,7 @@ export class LocalExecutor implements QueryExecutor {
 
     const coalesced = coalesceRanges(ranges);
     const estimatedBytes = ranges.reduce((s, r) => s + r.length, 0);
-    const totalRows = columns[0]?.pages.reduce((s, p) => s + p.rowCount, 0) ?? 0;
+    const totalRows = countColumnRows(columns);
 
     // Detect format from file
     const format = await this.detectFileFormat(query.table);

src/master-do.ts

@@ -1,6 +1,6 @@
 import { DurableObject } from "cloudflare:workers";
 import type { ColumnMeta, Env, Footer, TableMeta, DatasetMeta, AppendResult, AppendOptions, DropResult } from "./types.js";
-import { NULL_SENTINEL } from "./types.js";
+import { NULL_SENTINEL, countColumnRows } from "./types.js";
 import { parseFooter, parseColumnMetaFromProtobuf, FOOTER_SIZE } from "./footer.js";
 import { parseManifest } from "./manifest.js";
 import { detectFormat, getParquetFooterLength, parseParquetFooter, parquetMetaToTableMeta } from "./parquet.js";
@@ -54,7 +54,7 @@ export class MasterDO extends DurableObject<Env> {
     if (!result) throw new QueryModeError("INVALID_FORMAT", `Failed to read footer for "${r2Key}"`);
 
     const tableName = r2Key.replace(/\.(lance|parquet)$/, "").split("/").pop() ?? r2Key;
-    const totalRows = result.columns[0]?.pages.reduce((s, p) => s + p.rowCount, 0) ?? 0;
+    const totalRows = countColumnRows(result.columns);
     const meta: TableMeta = {
       name: tableName, footer: result.parsed, format: result.format, columns: result.columns,
       totalRows, fileSize: result.fileSize, r2Key, updatedAt: Date.now(),
@@ -150,7 +150,7 @@ export class MasterDO extends DurableObject<Env> {
 
     // Convert row-major to column-major
     const columnArrays = rowsToColumnArrays(rows);
-    if (columnArrays.length === 0) throw new Error("No valid columns found");
+    if (columnArrays.length === 0) throw new QueryModeError("SCHEMA_MISMATCH", "No valid columns found in rows");
 
     // Build Lance fragment via WASM
     const fragmentBytes = wasm.buildFragment(columnArrays);

src/query-do.ts

@@ -1,6 +1,6 @@
 import { DurableObject } from "cloudflare:workers";
 import type { ColumnMeta, DataType, Env, ExplainResult, FilterOp, Footer, Row, TableMeta, DatasetMeta, IcebergDatasetMeta, QueryResult } from "./types.js";
-import { queryReferencedColumns, queryCacheKey, NULL_SENTINEL } from "./types.js";
+import { queryReferencedColumns, queryCacheKey, countColumnRows, NULL_SENTINEL } from "./types.js";
 import type { QueryDescriptor } from "./client.js";
 import { parseFooter, parseColumnMetaFromProtobuf } from "./footer.js";
 import { parseManifest, logicalTypeToDataType } from "./manifest.js";
@@ -589,7 +589,7 @@ export class QueryDO extends DurableObject<Env> {
             const columns = await this.readColumnMeta(r2Key, footer);
             meta = {
               name: table, footer, format: "lance", columns,
-              totalRows: columns[0]?.pages.reduce((s, p) => s + p.rowCount, 0) ?? 0,
+              totalRows: countColumnRows(columns),
               fileSize, r2Key, updatedAt,
             };
           }
@@ -639,7 +639,7 @@ export class QueryDO extends DurableObject<Env> {
     }
 
     // Use totalRows from broadcast if available and > 0, otherwise fall back to column computation
-    const computedRows = columns[0]?.pages.reduce((s, p) => s + p.rowCount, 0) ?? 0;
+    const computedRows = countColumnRows(columns);
     const totalRows = (body.totalRows != null && body.totalRows > 0) ? body.totalRows : computedRows;
 
     const meta: TableMeta = {
@@ -665,7 +665,7 @@ export class QueryDO extends DurableObject<Env> {
     if (query.filters.length === 0 && !query.filterGroups?.length && !query.aggregates?.length && !query.groupBy?.length && !query.distinct && !query.join && !query.vectorSearch && !query.setOperation && !query.subqueryIn && !query.computedColumns?.length) {
       const meta = this.footerCache.get(query.table)
         ?? (await this.loadTableFromR2(query.table)) ?? undefined;
-      if (meta) return meta.columns[0]?.pages.reduce((s, p) => s + p.rowCount, 0) ?? meta.totalRows;
+      if (meta) return countColumnRows(meta.columns) || meta.totalRows;
     }
     const countQuery = { ...query, aggregates: [{ fn: "count" as const, column: "*" }] };
     const result = await this.executeQuery(countQuery);
@@ -1455,7 +1455,7 @@ export class QueryDO extends DurableObject<Env> {
         const columns = await this.readColumnMeta(r2Key, footer);
         const meta: TableMeta = {
           name: tableName, footer, format: "lance", columns,
-          totalRows: columns[0]?.pages.reduce((s, p) => s + p.rowCount, 0) ?? 0,
+          totalRows: countColumnRows(columns),
           fileSize, r2Key, updatedAt: Date.now(),
         };
         this.footerCache.set(tableName, meta);

src/types.ts

@@ -57,6 +57,11 @@ export interface ColumnMeta {
   listDimension?: number;
 }
 
+/** Count total rows across all pages of a column set. */
+export function countColumnRows(columns: ColumnMeta[]): number {
+  return columns[0]?.pages.reduce((s, p) => s + p.rowCount, 0) ?? 0;
+}
+
 /** Parquet-specific encoding info attached to pages */
 export interface PageEncoding {
   encoding?: string;      // "PLAIN" | "RLE_DICTIONARY" | ...