feat: wire multi-bucket sharding into all R2 paths + partition catalog for PB-scale

Multi-bucket sharding:
- Shared resolveBucket(env, r2Key) in src/bucket.ts — FNV-1a hash routing
- All 32 R2 call sites wired: query-do (23), fragment-do (2), master-do (12), worker (1)
- Reads, writes, list, head, delete all route through resolveBucket()
- 2-4x R2 rate limit increase with DATA_BUCKET_1/2/3 configured

Partition catalog:
- PartitionCatalog in src/partition-catalog.ts — O(1) lookup by partition key
- Maps partition values → fragment IDs (skips min/max evaluation)
- Auto-built when datasets load: picks best column by distinct value count
- Supports eq, in, neq, not_in for instant fragment pruning
- Two-phase pruning in executeMultiFragment: catalog → min/max → scan
- Stats exposed via diagnostics endpoint

Also: QueryExecutor.explain() now required, MaterializedExecutor has explain()

Author: teamchong

research/zig-engine-roadmap.md

@@ -146,11 +146,25 @@ Learnings from sibling Zig repos, prioritized by impact on QueryMode's WASM engi
 - `executeExplain()` reports `fragmentsSkipped` in explain output
 - `executeWithFragmentDOs()` accepts pruned fragment list (no wasted DO slots)
 
-### Multi-Bucket Sharding — FOUNDATION
+### Multi-Bucket Sharding — DONE
 - `Env` supports `DATA_BUCKET_1`, `DATA_BUCKET_2`, `DATA_BUCKET_3` (optional)
-- `resolveBucket(r2Key)` routes by FNV-1a hash of table prefix across all available buckets
-- Bypasses R2 per-bucket rate limits by distributing data across buckets
-- Remaining: wire `resolveBucket` into R2 read paths, add shard-aware ingest
+- `resolveBucket(env, r2Key)` in `src/bucket.ts` — shared utility, FNV-1a hash routing
+- Cached shard bucket array (computed once per DO lifetime)
+- **All 32 R2 call sites wired**: query-do.ts (23), fragment-do.ts (2), master-do.ts (12), worker.ts (1)
+- Reads, writes, list, head, delete all route through `resolveBucket()`
+- 2-4x R2 rate limit increase (10K→40K reads/s with 4 buckets)
+
+### Partition Catalog — DONE
+- `PartitionCatalog` in `src/partition-catalog.ts` — O(1) lookup by partition key values
+- Maps partition values → fragment IDs (no min/max evaluation needed)
+- Auto-built when datasets load: picks best column (most distinct values < 10K)
+- Supports `eq`, `in`, `neq`, `not_in` filters for instant fragment pruning
+- Two-phase pruning in `executeMultiFragment()`: catalog → min/max → scan
+- Stats exposed via diagnostics endpoint
+
+**Remaining:**
+- Partition-aware ingest (write data to correct shard bucket by partition key)
+- User-specified partition column (currently auto-detected)
 
 ## P3: Host Import Pattern for R2 I/O (from gitmode)
 

src/bucket.ts

@@ -0,0 +1,31 @@
+import type { Env } from "./types.js";
+
+/** Cached shard bucket array — computed once per DO lifetime. */
+let cachedAllBuckets: R2Bucket[] | null = null;
+
+/**
+ * Resolve R2 bucket for a given key. Routes by FNV-1a hash of the R2 key
+ * prefix (table name) across all available buckets.
+ *
+ * When DATA_BUCKET_1/2/3 are not configured, returns DATA_BUCKET (zero overhead).
+ * When configured, distributes tables across buckets for 2-4x rate limit increase.
+ */
+export function resolveBucket(env: Env, r2Key: string): R2Bucket {
+  if (!cachedAllBuckets) {
+    const shards = [env.DATA_BUCKET_1, env.DATA_BUCKET_2, env.DATA_BUCKET_3]
+      .filter((b): b is R2Bucket => !!b);
+    cachedAllBuckets = shards.length > 0
+      ? [env.DATA_BUCKET, ...shards]
+      : [env.DATA_BUCKET];
+  }
+
+  if (cachedAllBuckets.length === 1) return cachedAllBuckets[0];
+
+  const prefix = r2Key.split("/")[0] ?? r2Key;
+  let h = 0x811c9dc5;
+  for (let i = 0; i < prefix.length; i++) {
+    h ^= prefix.charCodeAt(i);
+    h = Math.imul(h, 0x01000193);
+  }
+  return cachedAllBuckets[(h >>> 0) % cachedAllBuckets.length];
+}

src/client.ts

@@ -2,6 +2,7 @@ import type {
   AggregateOp,
   AppendOptions,
   AppendResult,
+  DataType,
   DropResult,
   ExplainResult,
   FilterOp,
@@ -764,6 +765,14 @@ export class DataFrame<T extends Row = Row> {
           durationMs: 0,
         };
       },
+      async explain(_query: QueryDescriptor): Promise<ExplainResult> {
+        return {
+          table: "__from_operator__", format: "lance", totalRows: 0,
+          columns: [], pagesTotal: 0, pagesSkipped: 0, pagesScanned: 0,
+          estimatedBytes: 0, estimatedR2Reads: 0, fragments: 0,
+          filters: [], metaCached: false, estimatedRows: 0,
+        };
+      },
     };
     return new DataFrame<R>("__from_operator__", drainExecutor);
   }
@@ -898,7 +907,8 @@ export class MaterializedExecutor implements QueryExecutor {
     }
 
     // Apply aggregation
-    if (query.aggregates && query.aggregates.length > 0) {
+    const aggregates: AggregateOp[] | undefined = query.aggregates;
+    if (aggregates && aggregates.length > 0) {
       const groups = query.groupBy && query.groupBy.length > 0 ? query.groupBy : undefined;
       const buckets = new Map<string, Row[]>();
       for (const row of rows) {
@@ -911,7 +921,7 @@ export class MaterializedExecutor implements QueryExecutor {
       for (const [, bucket] of buckets) {
         const out: Row = {};
         if (groups) for (const g of groups) out[g] = bucket[0][g];
-        for (const agg of query.aggregates) {
+        for (const agg of aggregates) {
           const alias = agg.alias ?? `${agg.fn}_${agg.column}`;
           if (agg.fn === "count") {
             out[alias] = agg.column === "*" ? bucket.length : bucket.filter(r => r[agg.column] != null).length;

src/fragment-do.ts

@@ -9,6 +9,7 @@ import {
   type Operator, type RowBatch,
   buildEdgePipeline, drainPipeline,
 } from "./operators.js";
+import { resolveBucket } from "./bucket.js";
 import wasmModule from "./wasm-module.js";
 
 const R2_TIMEOUT_MS = 10_000;
@@ -123,7 +124,7 @@ export class FragmentDO extends DurableObject<Env> {
             withRetry(() =>
               withTimeout(
                 (async () => {
-                  const obj = await this.env.DATA_BUCKET.get(r2Key, {
+                  const obj = await resolveBucket(this.env, r2Key).get(r2Key, {
                     range: { offset: c.offset, length: c.length },
                   });
                   return obj ? { ...c, data: await obj.arrayBuffer() } : null;
@@ -179,7 +180,8 @@ export class FragmentDO extends DurableObject<Env> {
     let spillBytesRead: number | undefined;
 
     if (needsPipeline && allRows.length > 0) {
-      const spill = new R2SpillBackend(this.env.DATA_BUCKET, `__spill/${crypto.randomUUID()}`);
+      const spillBucket = fragments.length > 0 ? resolveBucket(this.env, fragments[0].r2Key) : resolveBucket(this.env, "");
+      const spill = new R2SpillBackend(spillBucket, `__spill/${crypto.randomUUID()}`);
       try {
         // Yield allRows in bounded batches to avoid holding 2x copies
         const BATCH_SIZE = 4096;

src/master-do.ts

@@ -5,6 +5,7 @@ import { parseManifest } from "./manifest.js";
 import { detectFormat, getParquetFooterLength, parseParquetFooter, parquetMetaToTableMeta } from "./parquet.js";
 import type { QueryDORpc } from "./types.js";
 import { instantiateWasm, type WasmEngine } from "./wasm-engine.js";
+import { resolveBucket } from "./bucket.js";
 import wasmModule from "./wasm-module.js";
 
 /** Master DO — single writer, reads footers, broadcasts invalidations. */
@@ -63,14 +64,14 @@ export class MasterDO extends DurableObject<Env> {
     const tableName = r2Prefix.replace(/\/$/, "").replace(/\.lance$/, "").split("/").pop() ?? r2Prefix;
 
     // Find latest manifest
-    const listed = await this.env.DATA_BUCKET.list({ prefix: `${r2Prefix}_versions/`, limit: 100 });
+    const listed = await resolveBucket(this.env, r2Prefix).list({ prefix: `${r2Prefix}_versions/`, limit: 100 });
     const manifestKeys = listed.objects
       .filter(o => o.key.endsWith(".manifest"))
       .sort((a, b) => a.key.localeCompare(b.key));
     if (manifestKeys.length === 0) throw new Error("No manifests found");
 
     const latestKey = manifestKeys[manifestKeys.length - 1].key;
-    const manifestObj = await this.env.DATA_BUCKET.get(latestKey);
+    const manifestObj = await resolveBucket(this.env, latestKey).get(latestKey);
     if (!manifestObj) throw new Error("Failed to read manifest");
 
     const manifest = parseManifest(await manifestObj.arrayBuffer());
@@ -172,15 +173,15 @@ export class MasterDO extends DurableObject<Env> {
     const dataR2Key = `${r2Prefix}${dataFilePath}`;
 
     // PUT data file to R2 (unique name = no conflict)
-    await this.env.DATA_BUCKET.put(dataR2Key, fragmentBytes);
+    await resolveBucket(this.env, dataR2Key).put(dataR2Key, fragmentBytes);
 
     // CAS loop for manifest update
     const MAX_RETRIES = 10;
     const latestKey = `${r2Prefix}_versions/_latest`;
 
     for (let attempt = 0; attempt < MAX_RETRIES; attempt++) {
       // Read current _latest with ETag
-      const latestObj = await this.env.DATA_BUCKET.get(latestKey);
+      const latestObj = await resolveBucket(this.env, latestKey).get(latestKey);
       let currentVersion = 0;
       let etag: string | undefined;
 
@@ -198,7 +199,7 @@ export class MasterDO extends DurableObject<Env> {
       let existingFragments: { id: number; filePath: string; physicalRows: number }[] = [];
       if (currentVersion > 0) {
         const manifestKey = `${r2Prefix}_versions/${currentVersion}.manifest`;
-        const manifestObj = await this.env.DATA_BUCKET.get(manifestKey);
+        const manifestObj = await resolveBucket(this.env, manifestKey).get(manifestKey);
         if (manifestObj) {
           const manifest = parseManifest(await manifestObj.arrayBuffer());
           if (manifest) existingFragments = manifest.fragments;
@@ -214,12 +215,12 @@ export class MasterDO extends DurableObject<Env> {
       // Write new manifest (protobuf-compatible binary)
       const manifestPayload = this.buildManifestBinary(newVersion, newFragments);
       const newManifestKey = `${r2Prefix}_versions/${newVersion}.manifest`;
-      await this.env.DATA_BUCKET.put(newManifestKey, manifestPayload);
+      await resolveBucket(this.env, newManifestKey).put(newManifestKey, manifestPayload);
 
       // CAS: write _latest with ETag condition
       try {
         const putOptions: R2PutOptions = etag ? { onlyIf: { etagMatches: etag } } : {};
-        const result = await this.env.DATA_BUCKET.put(latestKey, newVersionStr, putOptions);
+        const result = await resolveBucket(this.env, latestKey).put(latestKey, newVersionStr, putOptions);
 
         if (result === null && etag) {
           // 412 Precondition Failed — retry
@@ -346,13 +347,13 @@ export class MasterDO extends DurableObject<Env> {
     parsed?: Footer; raw: ArrayBuffer; fileSize: bigint; columns: ColumnMeta[];
     format: "lance" | "parquet";
   } | null> {
-    const head = await this.env.DATA_BUCKET.head(r2Key);
+    const head = await resolveBucket(this.env, r2Key).head(r2Key);
     if (!head) return null;
 
     const fileSize = BigInt(head.size);
     // Read last 40 bytes — enough for Lance footer or Parquet tail detection
     const tailSize = Math.min(Number(fileSize), FOOTER_SIZE);
-    const obj = await this.env.DATA_BUCKET.get(r2Key, {
+    const obj = await resolveBucket(this.env, r2Key).get(r2Key, {
       range: { offset: Number(fileSize) - tailSize, length: tailSize },
     });
     if (!obj) return null;
@@ -366,7 +367,7 @@ export class MasterDO extends DurableObject<Env> {
 
       // Fetch full Parquet Thrift footer
       const footerOffset = Number(fileSize) - footerLen - 8;
-      const footerObj = await this.env.DATA_BUCKET.get(r2Key, {
+      const footerObj = await resolveBucket(this.env, r2Key).get(r2Key, {
         range: { offset: footerOffset, length: footerLen },
       });
       if (!footerObj) return null;
@@ -386,7 +387,7 @@ export class MasterDO extends DurableObject<Env> {
     let columns: ColumnMeta[] = [];
     const metaLen = Number(parsed.columnMetaOffsetsStart) - Number(parsed.columnMetaStart);
     if (metaLen > 0) {
-      const metaObj = await this.env.DATA_BUCKET.get(r2Key, {
+      const metaObj = await resolveBucket(this.env, r2Key).get(r2Key, {
         range: { offset: Number(parsed.columnMetaStart), length: metaLen },
       });
       if (metaObj) columns = parseColumnMetaFromProtobuf(await metaObj.arrayBuffer(), parsed.numColumns);
@@ -456,7 +457,7 @@ export class MasterDO extends DurableObject<Env> {
       // List and delete all R2 objects under this prefix
       let cursor: string | undefined;
       do {
-        const listed = await this.env.DATA_BUCKET.list({
+        const listed = await resolveBucket(this.env, r2Prefix).list({
           prefix: r2Prefix,
           cursor,
           limit: 1000,
@@ -467,7 +468,7 @@ export class MasterDO extends DurableObject<Env> {
           bytesFreed += listed.objects.reduce((s, o) => s + o.size, 0);
           fragmentsDeleted += keys.length;
           // R2 delete supports up to 1000 keys per call
-          await this.env.DATA_BUCKET.delete(keys);
+          await resolveBucket(this.env, r2Prefix).delete(keys);
         }
 
         cursor = listed.truncated ? listed.cursor : undefined;